plots.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "007ab727-e54f-4116-bedc-24b5fe32db85",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit import Chem\n",
    "from rdkit.Chem.Draw import IPythonConsole, rdMolDraw2D\n",
    "from rdkit.Chem import AllChem, DataStructs, Descriptors, PandasTools\n",
    "\n",
    "import os\n",
    "from os.path import join, getsize\n",
    "\n",
    "from useful_rdkit_utils import mol2numpy_fp\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "from sklearn.manifold import TSNE\n",
    "\n",
    "from umap import UMAP\n",
    "\n",
    "from itertools import product\n",
    "\n",
    "import concurrent.futures\n",
    "\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from chemplot import Plotter\n",
    "\n",
    "import numba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1372700-f497-457b-93ed-871e68a99914",
   "metadata": {},
   "outputs": [],
   "source": [
    "current_dir = os.path.dirname(os.path.abspath('plots.ipynb'))\n",
    "\n",
    "# reset the working directory to the notebook's location\n",
    "print(current_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dedc2109-3f92-4b02-976d-68d63753c1c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a relative path to the CSV file\n",
    "oracle_csv_path = os.path.join(current_dir, 'cs_49k.csv') # precomputed cnnaffinity as oracle, for 49k mols\n",
    "oracle_df = pd.read_csv(oracle_csv_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08103881-24b5-4653-a8c1-6985647f035b",
   "metadata": {},
   "source": [
    "Functions for data analysis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52573167-1928-46b0-a4dd-2d543551a4f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_test_df(directory):\n",
    "    \"\"\"\n",
    "    Concatenate all selection.csv files from each cycle directory within the specified directory.\n",
    "\n",
    "    Parameters:\n",
    "    - directory: str, the path to the directory containing cycle folders.\n",
    "\n",
    "    Returns:\n",
    "    - all_selections_df: DataFrame, containing all the data with an additional 'cycle' column.\n",
    "    \"\"\"\n",
    "    # Initialize an empty DataFrame\n",
    "    all_selections_df = pd.DataFrame()\n",
    "\n",
    "    # List all cycle directories in the specified directory\n",
    "    cycle_dirs = [d for d in os.listdir(directory) if\n",
    "                  os.path.isdir(os.path.join(directory, d)) and d.startswith('cycle_')]\n",
    "\n",
    "    # Iterate over each directory and append the data to the DataFrame\n",
    "    for cycle_dir in cycle_dirs:\n",
    "        cycle_path = os.path.join(directory, cycle_dir, 'selection.csv')\n",
    "        if os.path.exists(cycle_path):\n",
    "            # Read the selection.csv file\n",
    "            cycle_df = pd.read_csv(cycle_path)\n",
    "            print(cycle_df.columns)\n",
    "            new_name = 'cnnaffinity'\n",
    "            if 'combo1' in cycle_df.columns:\n",
    "                cycle_df.rename(columns={'combo1': new_name}, inplace=True)\n",
    "            \n",
    "            if 'plip' in cycle_df.columns:\n",
    "                cycle_df.rename(columns={'plip': new_name}, inplace=True)\n",
    "            # Add a 'cycle' column to keep track of the source\n",
    "            cycle_number = int(cycle_dir.split('_')[1])  # Extract cycle number from the directory name\n",
    "            cycle_df['cycle'] = cycle_number\n",
    "            cycle_df['cnnaffinity'] = cycle_df['cnnaffinity'].abs()\n",
    "            # cycle_df['cnnaffinity'].astype('float').dtypes\n",
    "\n",
    "            # Append the DataFrame to the main DataFrame\n",
    "            all_selections_df = pd.concat([all_selections_df, cycle_df], ignore_index=True)\n",
    "    all_selections_df['expt'] = directory\n",
    "\n",
    "    print(all_selections_df.columns)\n",
    "    return all_selections_df\n",
    "\n",
    "\n",
    "def calculate_performance_metrics(data_df, oracle_df, top_percentage):\n",
    "    \"\"\"\n",
    "    Calculate performance metrics for a given dataframe.\n",
    "\n",
    "    Parameters:\n",
    "    - cycle_data: DataFrame for a specific cycle containing 'is_active' column.\n",
    "    - oracle_df: DataFrame of the oracle data containing 'cnnaffinity' column.\n",
    "    - threshold: Float, the threshold for determining active compounds.\n",
    "\n",
    "    Returns:\n",
    "    - A dictionary containing TP, FP, TN, FN, recall, precision, and F1 score for the cycle.\n",
    "    \"\"\"\n",
    "    # calculate true positives (TP), false positives (FP), and false negatives (FN)\n",
    "    threshold = oracle_df['cnnaffinity'].quantile(1 - top_percentage)\n",
    "    TP = data_df['is_active'].sum()  #!! just sum of all the actives, correctly identified as positive via oracle\n",
    "    FN = oracle_df[oracle_df['cnnaffinity'].astype('float') >= float(threshold)].shape[0] - TP  # all the actives from the oracle, minus the ones we picked i.e. were active but we didnt pick them\n",
    "    FP = data_df.shape[0] - TP  #!! num mols we picked when we shouldn't\n",
    "    TN = oracle_df.shape[0] - TP - FP - FN  # only 4 options TP/TN/FP/FN; so total - (TP-FP-FN) = TN\n",
    "\n",
    "    # ensure the counts add up to the total number of samples\n",
    "    assert TP + FN + FP + TN == len(oracle_df), \"Count mismatch\"\n",
    "\n",
    "    recall = TP / (TP + FN) if TP + FN > 0 else 0\n",
    "    precision = TP / (TP + FP) if TP + FP > 0 else 0\n",
    "\n",
    "    f1_score = 2 * precision * recall / (precision + recall) if precision + recall > 0 else 0 #check\n",
    "\n",
    "    print(f\"True Positive (TP): {TP}\")\n",
    "    print(f\"False Negative (FN): {FN}\")\n",
    "    print(f\"False Positive (FP): {FP}\")\n",
    "    print(f\"True Negative (TN): {TN}\")\n",
    "\n",
    "    # Check for Count Mismatch\n",
    "    print(f\"Total Samples: {len(oracle_df)}\")\n",
    "    print(f\"Sum of Counts (TP+FN+FP+TN): {TP + FN + FP + TN}\")\n",
    "\n",
    "    print(f\"Recall: {recall}\")\n",
    "    print(f\"Precision: {precision}\")\n",
    "    print(f\"F1 Score: {f1_score}\")\n",
    "\n",
    "    return {\n",
    "        'expt': [data_df['expt'].iloc[0].replace('_generated', '')], #[data_df.cycle.iloc[0]], #\n",
    "        'TP': [TP],\n",
    "        'FP': [FP],\n",
    "        'TN': [TN],\n",
    "        'FN': [FN],\n",
    "        'recall': [recall],\n",
    "        'precision': [precision],\n",
    "        'f1_score': [f1_score]\n",
    "    }\n",
    "\n",
    "\n",
    "def find_active_mols(df, top_percentage, oracle_df):\n",
    "    \"\"\"\n",
    "    Assign activity to chosen molecules from oracle\n",
    "\n",
    "    Parameters:\n",
    "    - df: DataFrame containing 'cnnaffinity'\n",
    "    - top_percentage: float, the top percentage of 'cnnaffinity' values to consider as 'active'.\n",
    "\n",
    "    Returns:\n",
    "    - metrics_df: DataFrame containing 'is_active' column, where active is defined as above top_percentage score in the oracle\n",
    "    \"\"\"\n",
    "    threshold = oracle_df['cnnaffinity'].quantile(1 - top_percentage)\n",
    "    df['is_active'] = df['cnnaffinity'] >= threshold\n",
    "\n",
    "    print(f'Threshold for {top_percentage} is {threshold}.')\n",
    "    print(f'Number of active mols for {df.expt.iloc[0]} is {df.is_active.sum()}.')\n",
    "\n",
    "\n",
    "    return df\n",
    "\n",
    "def gather_data(data_dir):\n",
    "    dir_list = []\n",
    "    for root, dirs, files in os.walk('.'):\n",
    "        dir_list.append(dirs) \n",
    "    return dir_list[0]\n",
    "\n",
    "\n",
    "def get_params(exp_list):\n",
    "    search_strs = [text.split('_') for text in exp_list]\n",
    "    flat_list = [item for sublist in search_strs for item in sublist]\n",
    "    param_list = list(set(flat_list))\n",
    "    return param_list\n",
    "\n",
    "def gen_rep_data(data_dir, percent=0.02, trunc=True):\n",
    "    exp_list = gather_data(data_dir)\n",
    "    df_list = [create_test_df(exp) for exp in exp_list]\n",
    "\n",
    "    \n",
    "    if trunc:\n",
    "        df_sorted = [df.sort_values(by=['cycle'])[:2500] for df in df_list]\n",
    "    else:\n",
    "        \n",
    "        df_sorted = [df.sort_values(by=['cycle']) for df in df_list]\n",
    "    [len(df) for df in df_sorted]\n",
    "    df_list = df_sorted\n",
    "    \n",
    "    dat = {percent : [find_active_mols(df, float(percent), oracle_df) for df in df_list]}\n",
    "    \n",
    "    \n",
    "    metrics = [calculate_performance_metrics(df, oracle_df, float(percent)) for df in dat[percent]]\n",
    "    \n",
    "    metrics = [pd.DataFrame(dict) for dict in metrics]\n",
    "    metric_df = pd.concat(metrics)\n",
    "    return df_list, metric_df\n",
    "\n",
    "\n",
    "def visualize_data(search_string, final_df):\n",
    "    # Filter DataFrame using Search String\n",
    "    search_df = final_df[final_df['expt'].str.contains(search_string)]\n",
    "    \n",
    "    # Melt the DataFrame for Visualisation\n",
    "    metrics_melted = pd.melt(search_df, id_vars='expt', value_vars=['f1_score', 'recall'], var_name='metric', value_name='score')\n",
    "    \n",
    "    # Plot Data\n",
    "    fig, ax1 = plt.subplots()\n",
    "    custom_palette = sns.color_palette(\"colorblind\", n_colors=len(metrics_melted['metric'].unique()))\n",
    "    g = sns.barplot(x='expt', y='score', hue='metric', data=metrics_melted, ax=ax1, palette=custom_palette, hatch='/')\n",
    "    \n",
    "    # Customize Plot\n",
    "    plt.xticks(rotation=90)\n",
    "    plt.title(f'{search_string} - Top {percent}')  # Note: 'percent' is not defined in the provided code snippet.\n",
    "    ax1.set_ylabel('Score')\n",
    "    ax1.set_xlabel('Cycle size')\n",
    "    plt.legend(loc=(1.04, 0))\n",
    "\n",
    "    sns.set_style(\"darkgrid\", {\"axes.facecolor\": \".9\"})\n",
    "    plt.savefig(f'{search_string}_{percent}.png', bbox_inches=\"tight\")\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def create_violin_plot(df_combined):\n",
    "    \"\"\"\n",
    "    Load data from a df and create a violin plot.\n",
    "\n",
    "    Parameters:\n",
    "    - directory_path: Path to the directory containing SDF files.\n",
    "    - plot_title: Title for the violin plot.\n",
    "\n",
    "    Returns:\n",
    "    None.\n",
    "    \"\"\"\n",
    "    \n",
    "    \n",
    "    df_combined['cycle'] = df_combined['cycle'].astype('int')\n",
    "    df_combined['cnnaffinity'] = df_combined['cnnaffinity'].astype('float')\n",
    "    plot_title = df_combined['expt'].iloc[0]\n",
    "\n",
    "    # Create a violin plot of sf1 vs cycle\n",
    "    sns.violinplot(x='cycle', y='cnnaffinity', data=df_combined)\n",
    "    #plt.title(plot_title)\n",
    "    plt.xlabel('Cycle')\n",
    "    plt.ylabel('Predicted pK')\n",
    "    plt.xticks(range(min(df_combined['cycle']), max(df_combined['cycle']) + 1, 2))\n",
    "    plt.savefig(f'/home/cree/code/gal/cs50k/{plot_title}')\n",
    "    plt.show()\n",
    "    return df_combined\n",
    "\n",
    "def plot_metric_over_cycles(df, metric, oracle_df, percent):\n",
    "    # Check if the metric is valid\n",
    "    if metric not in ['recall', 'precision', 'accuracy', 'f1_score']:  # Add other valid metrics if needed\n",
    "        raise ValueError(\"Invalid metric specified.\")\n",
    "    expt = df['expt'].iloc[0]\n",
    "    # Split df into a list of DataFrames based on the 'cycle' column\n",
    "    max_cycle = df['cycle'].max()\n",
    "    df_list = [df[df['cycle'] <= i] for i in range(1, max_cycle + 1)]\n",
    "\n",
    "    # Calculate performance metrics for each DataFrame\n",
    "    data = [calculate_performance_metrics(df, oracle_df, percent) for df in df_list]\n",
    "\n",
    "    # Extracting the specified metric values from the data\n",
    "    metric_values = [entry[metric][0] for entry in data]  # Adjust indexing if necessary\n",
    "\n",
    "    # Plotting\n",
    "    plt.figure(figsize=(10, 6))\n",
    "    plt.plot(range(1, len(metric_values) + 1), metric_values, marker='o')\n",
    "    plt.title(f'{metric.capitalize()} Values Over Different Cycles')\n",
    "    plt.xlabel('Cycle')\n",
    "    plt.ylabel(metric.capitalize())\n",
    "    plt.grid(True)\n",
    "    plt.xticks(range(1, len(metric_values) + 1))\n",
    "    \n",
    "    # Save the plot\n",
    "    plt.savefig(f'{expt}_{metric}_{percent}.png')\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def visualize_avg_data(search_terms, df_list):\n",
    "    # Combine all DataFrames in the list\n",
    "    combined_df = pd.concat(df_list)\n",
    "    \n",
    "    # Convert search_terms to a list if it's a single string\n",
    "    if isinstance(search_terms, str):\n",
    "        search_terms = [search_terms]\n",
    "    \n",
    "    # Apply AND logic in filtering: all terms must be present\n",
    "    mask = combined_df['expt'].apply(lambda x: all(term in x for term in search_terms))\n",
    "    search_df = combined_df[mask]\n",
    "        # Calculate Mean and Standard Deviation\n",
    "    mean_df = search_df.groupby('expt')[['f1_score', 'recall']].mean().reset_index()\n",
    "    std_df = search_df.groupby('expt')[['f1_score', 'recall']].sem().reset_index()\n",
    "    \n",
    "    # Melt the DataFrames for Visualization\n",
    "    mean_melted = pd.melt(mean_df, id_vars='expt', value_vars=['f1_score', 'recall'], var_name='metric', value_name='mean_score')\n",
    "    std_melted = pd.melt(std_df, id_vars='expt', value_vars=['f1_score', 'recall'], var_name='metric', value_name='std_dev')\n",
    "    \n",
    "    # Merge Mean and Standard Deviation Data\n",
    "    merged_df = pd.merge(mean_melted, std_melted, on=['expt', 'metric'])\n",
    "    print(merged_df)\n",
    "    # Plot Data\n",
    "    fig, ax1 = plt.subplots()\n",
    "    custom_palette = sns.color_palette(\"colorblind\", n_colors=len(merged_df['metric'].unique()))\n",
    "    g = sns.barplot(x='expt', y='mean_score', hue='metric', data=merged_df, ax=ax1, palette=custom_palette, hatch='/', )\n",
    "    \n",
    "    # Add Error Bars\n",
    "    for i, bar in enumerate(g.patches):\n",
    "        bar.set_width(0.3)\n",
    "        hue_index = i % len(merged_df['metric'].unique())\n",
    "        x = bar.get_x() #+ bar.get_width() #* hue_index\n",
    "        y = bar.get_height()\n",
    "        #error = merged_df.iloc[i]['std_dev']\n",
    "        #ax1.errorbar(x + bar.get_width()/2, y, yerr=error, fmt='none', color='black', capsize=3)\n",
    "\n",
    "    # Customize Plot\n",
    "    plt.xticks(rotation=90)\n",
    "    ax1.set_xticklabels(ax1.get_xticklabels(), fontsize=10)\n",
    "    plt.title(f'{search_terms[0]} Metrics')\n",
    "    ax1.set_ylabel('Score')\n",
    "    ax1.set_xlabel('Experiment')\n",
    "    plt.legend(loc=(1.04, 0))\n",
    "    plt.xticks(rotation=90)\n",
    "    #custom_xtick_labels = ['false', 'true',]\n",
    "    #ax1.set_xticklabels(custom_xtick_labels, fontsize=10, rotation=90)\n",
    "    plt.title(f'{search_terms[0]} Metrics')\n",
    "    sns.set_style(\"darkgrid\", {\"axes.facecolor\": \".9\"})\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46d5d8d1-5d58-4fc4-bf29-a174d73eaea0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def smi2svg(smi):\n",
    "    mol = Chem.MolFromSmiles(smi)\n",
    "    d2d = rdMolDraw2D.MolDraw2DSVG(200, 100)\n",
    "    d2d.DrawMolecule(mol)\n",
    "    d2d.FinishDrawing()\n",
    "    return d2d.GetDrawingText()\n",
    "    \n",
    "\n",
    "def umap(df, nbits=2048):\n",
    "    \"\"\"\n",
    "    Compute UMAP projections for molecular data.\n",
    "\n",
    "    Parameters:\n",
    "    - df: Dataframe containing a 'ROMol' column with molecular data.\n",
    "\n",
    "    Returns:\n",
    "    - res: UMAP reduced dimensionality output.\n",
    "    \"\"\"\n",
    "    # Compute Morgan Fingerprints\n",
    "    df['fp'] = df['ROMol'].apply(lambda x: AllChem.GetMorganFingerprintAsBitVect(x, 2, nBits=nbits))\n",
    "    df['svg'] = df['Smiles'].apply(lambda x: smi2svg(x))\n",
    "    # Tanimoto Distance function\n",
    "    def tanimoto_dist(a, b):\n",
    "        dotprod = np.dot(a, b)\n",
    "        tc = dotprod / (np.sum(a) + np.sum(b) - dotprod)\n",
    "        return 1.0 - tc\n",
    "\n",
    "    # UMAP dimensionality reduction\n",
    "    fps = df['fp'].apply(lambda fp: np.array(fp)).tolist()\n",
    "    from umap import UMAP\n",
    "    reducer = UMAP(metric=tanimoto_dist)\n",
    "    res = reducer.fit_transform(fps)\n",
    "    \n",
    "    return res\n",
    "\n",
    "\n",
    "def cluster_data(df, res, min_samples, min_cluster_size):\n",
    "    # Apply HDBSCAN clustering on the UMAP results\n",
    "    clusterer = hdbscan.HDBSCAN(min_samples=min_samples, min_cluster_size=min_cluster_size, cluster_selection_method='leaf')\n",
    "    cluster_labels = clusterer.fit_predict(res)\n",
    "\n",
    "    # Add the cluster labels to the original DataFrame\n",
    "    df['cluster'] = cluster_labels\n",
    "    return df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d10362f0-f1c5-47b2-93aa-c31df76cd355",
   "metadata": {},
   "source": [
    "### Reproduce plots for analysis of the chemical space of the oracle dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89520ba6-2c52-4291-bfd0-ef55cecded37",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit.Chem.Descriptors import ExactMolWt\n",
    "\n",
    "\n",
    "# MW distribution of the oracle\n",
    "oracle_df['sf1'] = oracle_df['cnnaffinity']\n",
    "oracle_df['ROMol'] = oracle_df['Smiles'].apply(lambda x: Chem.MolFromSmiles(x))\n",
    "oracle_df['MW'] = oracle_df['ROMol'].apply(lambda x: ExactMolWt(x))\n",
    "\n",
    "\n",
    "# plotting the histogram with adjustments\n",
    "plt.hist(oracle_df['MW'], color='lightgrey', edgecolor='black', bins=50)\n",
    "\n",
    "# removing gridlines\n",
    "plt.grid(False)\n",
    "\n",
    "# setting x-axis label to 'MW'\n",
    "plt.xlabel('MW')\n",
    "plt.savefig('oracle_mw')\n",
    "# simple histogram of the oracle\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de9ffb50-9803-474d-85f8-c3497a8781a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "oracle_csv_path = os.path.join(current_dir, 'cs_49k.csv')\n",
    "\n",
    "oracle_df = pd.read_csv(oracle_csv_path)\n",
    "# Identify duplicates based on 'Smiles' and 'cnnaffinity'\n",
    "duplicates = oracle_df[oracle_df.duplicated(subset=['cnnaffinity'], keep=False)]\n",
    "duplicates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35657856-affa-4dec-8fe3-06156958de9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add additional mpro inhibitors from prospective search (just to build umap, not to plot):\n",
    "sdf = 'onebyone_it14_over6cnnaffinity.sdf'\n",
    "enamine_df = PandasTools.LoadSDF(sdf)\n",
    "enamine_df = enamine_df.rename(columns={'filename': 'Smiles',})\n",
    "enamine_df "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51014d04-d231-4999-bb32-980311ee994d",
   "metadata": {},
   "outputs": [],
   "source": [
    "oracle_df = oracle_df.reset_index(drop=True,)\n",
    "enamine_df = enamine_df.reset_index(drop=True,)\n",
    "enamine_oracle_df = pd.concat([oracle_df, enamine_df])\n",
    "enamine_oracle_df = enamine_oracle_df[['Smiles','cnnaffinity', 'enamine_id']]\n",
    "enamine_oracle_df.fillna(0,inplace=True)\n",
    "enamine_oracle_df.cnnaffinity = enamine_oracle_df.cnnaffinity.astype('float').abs()\n",
    "enamine_oracle_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fdd3988-b558-42b2-9d73-4f9ce9916f57",
   "metadata": {},
   "source": [
    "## Fig 2b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f309f1d-7071-4654-bffc-0e7657780515",
   "metadata": {},
   "outputs": [],
   "source": [
    "# check distribution of predicted pK for oracle dataset\n",
    "oracle_df = oracle_df.drop_duplicates(subset=['Smiles', 'cnnaffinity'], keep='first')\n",
    "oracle_df.hist('cnnaffinity', bins=175, color='grey')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bc137c8-98d5-4dbc-9d54-dabb8fe45b79",
   "metadata": {},
   "source": [
    "## Fig 2c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8276b760-7e6f-44ff-85fe-e7f30dfe452d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import BoundaryNorm, ListedColormap\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "\n",
    "# Custom color map and normalization\n",
    "my_colors = ['#c7c7c7', 'gold', 'orange', 'red', 'darkred']\n",
    "my_cmap = ListedColormap(my_colors)\n",
    "bounds = [3, 4.5, 5, 5.5, 6]\n",
    "my_norm = BoundaryNorm(bounds, ncolors=len(my_colors))\n",
    "\n",
    "def plot_umap(df, title, size=100):\n",
    "    sns.set_style('white')\n",
    "    plt.figure(figsize=(10, 8))\n",
    "\n",
    "    # Use custom colormap and norm\n",
    "    points = plt.scatter(x=df['UMAP-1'], y=df['UMAP-2'], c=df['cnnaffinity'], cmap=my_cmap, norm=my_norm, s=size)\n",
    "    \n",
    "    # Creating color bar and legend\n",
    "    cbar = plt.colorbar(points, spacing='proportional', ticks=bounds, shrink=0.6, aspect=30)\n",
    "    cbar.ax.tick_params(labelsize=16)\n",
    "    cbar.set_label('Predicted pK', fontsize=16) \n",
    "    #plt.title(f'Chemical space UMAP - {title}', fontsize=16)\n",
    "    plt.xlabel('UMAP-1', fontsize=16)\n",
    "    plt.ylabel('UMAP-2', fontsize=16)\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n",
    "    \n",
    "    #plt.legend(title='CNNaffinity', title_fontsize=14, fontsize=14, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)\n",
    "    plt.savefig(title, bbox_inches='tight')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3abf3676-8937-4961-9bf2-eaae7c39bb9c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# calculate umap for oracle df\n",
    "cp = Plotter.from_smiles(enamine_oracle_df[\"Smiles\"],)\n",
    "res = cp.umap()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69988044-15fb-463c-bf4d-a5369dbffa6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot oracle data onto umap:\n",
    "oracle_df = oracle_df.reset_index(drop=True,)\n",
    "res = res.reset_index(drop=True,)\n",
    "df = pd.concat([oracle_df, res], axis=1)\n",
    "df['cnnaffinity'] = df['cnnaffinity'].astype(float)\n",
    "umap_df = df[df['cnnaffinity'] > 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5020c261-b78a-4b8a-8cfe-67693357bb0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# note that umap is a stochastic algorithm, so different runs may show variations:\n",
    "plot_umap(umap_df, '50K', size=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b87c879-e9f0-4d01-9d8b-07afa18107a9",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "## Fig 3 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dd1c7f2-9d78-4043-a34a-57cbc3ae60e2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# some examples of calculation of recall and F1 score for different active learning hyperparameters:\n",
    "\n",
    "os.chdir(current_dir)\n",
    "base_dir = os.path.dirname(os.path.abspath('plots.ipynb'))\n",
    "df_list = []\n",
    "percent=0.02\n",
    "# create 5 dfs in a loop with different directories for each rep\n",
    "dfs = {}\n",
    "for i in [1,2,3,4,5]:\n",
    "    data_dir = f'{base_dir}/rep_{i}'  # change directory for each rep\n",
    "    os.chdir(data_dir)\n",
    "    print(data_dir)\n",
    "    df_name = f'rep{i}_df'  # naming each df as rep1_df, rep2_df, etc.\n",
    "    expt_df, dfs[df_name] = gen_rep_data(data_dir, percent, trunc=False)\n",
    "    df_list.append(expt_df)\n",
    "\n",
    "os.chdir(current_dir)\n",
    "# Access dataframes: dfs['rep1_df'], dfs['rep2_df'], ..., dfs['rep5_df']\n",
    "dfs['rep3_df']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ea288f9-ed96-4dc8-9ecb-b2470658fedb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# average f1/recall values over 5 independent runs\n",
    "# dfs is a dictionary of all reps, accessed via their keys, repX_df\n",
    "# ignore error\n",
    "visualize_avg_data([\"0.1\", \"gp\",], [dfs[key] for key in ['rep1_df', 'rep2_df', 'rep3_df', 'rep4_df', 'rep5_df']])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebb1e5c0-6379-4cda-99c2-9c3029eb55f9",
   "metadata": {},
   "source": [
    "## Fig 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9bc72ff-36fa-4550-94cf-b7917bfb8dea",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_cycles(df):\n",
    "    cyc_last = df['cycle'].max()\n",
    "    first_df = df[df['cycle'] == 1]\n",
    "    last_df = df[df['cycle'] == cyc_last]\n",
    "    return {'first' : first_df, 'last' : last_df}\n",
    "    \n",
    "def create_test_df(directory, feature_column):\n",
    "    \"\"\"\n",
    "    Concatenate all selection.csv files from each cycle directory within the specified directory.\n",
    "\n",
    "    Parameters:\n",
    "    - directory: str, the path to the directory containing cycle folders.\n",
    "\n",
    "    Returns:\n",
    "    - all_selections_df: DataFrame, containing all the data with an additional 'cycle' column.\n",
    "    \"\"\"\n",
    "    # Initialize an empty DataFrame\n",
    "    all_selections_df = pd.DataFrame()\n",
    "\n",
    "    # List all cycle directories in the specified directory\n",
    "    cycle_dirs = [d for d in os.listdir(directory) if\n",
    "                  os.path.isdir(os.path.join(directory, d)) and d.startswith('cycle_')]\n",
    "\n",
    "    # Iterate over each directory and append the data to the DataFrame\n",
    "    for cycle_dir in cycle_dirs:\n",
    "        cycle_path = os.path.join(directory, cycle_dir, 'selection.csv')\n",
    "        if os.path.exists(cycle_path):\n",
    "            # Read the selection.csv file\n",
    "            cycle_df = pd.read_csv(cycle_path)\n",
    "\n",
    "            # Add a 'cycle' column to keep track of the source\n",
    "            cycle_number = int(cycle_dir.split('_')[1])  # Extract cycle number from the directory name\n",
    "            cycle_df['cycle'] = cycle_number\n",
    "            cycle_df[feature_column] = cycle_df[feature_column].abs()\n",
    "            # cycle_df['cnnaffinity'].astype('float').dtypes\n",
    "\n",
    "            # Append the DataFrame to the main DataFrame\n",
    "            all_selections_df = pd.concat([all_selections_df, cycle_df], ignore_index=True)\n",
    "    all_selections_df['expt'] = directory\n",
    "    return all_selections_df\n",
    "\n",
    "import pandas as pd\n",
    "from rdkit import Chem, DataStructs\n",
    "from rdkit.Chem import AllChem, PandasTools\n",
    "import itertools\n",
    "\n",
    "dfs = [] # list to store each DataFrame\n",
    "\n",
    "# generate and store dfs\n",
    "for i in range(1, 6):\n",
    "    path = f'{base_dir}/rep_{i}/gp_200_UCB_True_10'\n",
    "    df = create_test_df(path, 'cnnaffinity')\n",
    "    # optionally, add source identifier before appending\n",
    "    df['source'] = f'rep_{i}'\n",
    "    dfs.append(df)\n",
    "\n",
    "# concatenate all dfs\n",
    "concatenated_df = pd.concat(dfs, ignore_index=True)\n",
    "concatenated_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17ab3d02-8b03-4a14-aae3-2f4db28ee19d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_cycles(df):\n",
    "    cyc_last = df['cycle'].max()\n",
    "    first_df = df[df['cycle'] == 1]\n",
    "    last_df = df[df['cycle'] == cyc_last]\n",
    "    return {'first' : first_df, 'last' : last_df}\n",
    "\n",
    "\n",
    "\n",
    "cycle_dict = extract_cycles(concatenated_df)\n",
    "\n",
    "def fid_filter(mask_df):\n",
    "    mask = umap_df['fid'].isin(mask_df['fid'])\n",
    "    filtered_df = umap_df[mask]\n",
    "    return filtered_df\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import BoundaryNorm, ListedColormap\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "\n",
    "# Custom color map and normalization\n",
    "my_colors = ['#c7c7c7', 'gold', 'orange', 'red', 'darkred']\n",
    "my_cmap = ListedColormap(my_colors)\n",
    "bounds = [3, 4.5, 5, 5.5, 6]\n",
    "my_norm = BoundaryNorm(bounds, ncolors=len(my_colors))\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def plot_umap(df, title, size=100):\n",
    "    sns.set_style('white')\n",
    "    plt.figure(figsize=(10, 8))\n",
    "\n",
    "    # Use custom colormap and norm\n",
    "    points = plt.scatter(x=df['UMAP-1'], y=df['UMAP-2'], c=df['cnnaffinity'], cmap=my_cmap, norm=my_norm, s=size)\n",
    "    \n",
    "    # Creating color bar and legend\n",
    "    cbar = plt.colorbar(points, spacing='proportional', ticks=bounds, shrink=0.6, aspect=30)\n",
    "    cbar.ax.tick_params(labelsize=16)\n",
    "    cbar.set_label('Predicted pK', fontsize=16) \n",
    "    #plt.title(f'Chemical space UMAP - {title}', fontsize=16)\n",
    "    plt.xlabel('UMAP-1', fontsize=16)\n",
    "    plt.ylabel('UMAP-2', fontsize=16)\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n",
    "\n",
    "    #plt.legend(title='CNNaffinity', title_fontsize=14, fontsize=14, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)\n",
    "    plt.savefig(title, bbox_inches='tight')\n",
    "    plt.show()\n",
    "    \n",
    "cycles = ['1','13']\n",
    "for i, df in enumerate(cycle_dict.values()):\n",
    "    number_of_zeros = (df['cnnaffinity'] == 0).sum()\n",
    "    mean_value = df['cnnaffinity'].mean()\n",
    "    #print(f'mean cnnaffinity: {mean_value}, # of 0s: {number_of_zeros}')\n",
    "    filtered_df = fid_filter(df)\n",
    "    #print(filtered_df)\n",
    "    #cycle = str(filtered_df['cycle'].iloc[0])\n",
    "    #print(cycle)\n",
    "    filtered_df['cnnaffinity'] = filtered_df['cnnaffinity'].astype(float)\n",
    "    filtered_df = filtered_df[filtered_df['cnnaffinity'] > 0]\n",
    "    plot_umap(filtered_df, f'{cycles[i]}',)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f684525c-84e2-4f6a-b215-cbc0ebc51d31",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f4e65a9-c4c0-4b1d-ab08-575f812ef5f1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fe8058e-281f-4afd-8f7f-8e9a9f09d9ef",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "6a8b43ee-5f47-47d3-8ec2-3665a7449868",
   "metadata": {},
   "source": [
    "## Fig 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c324ee8-9419-45e6-a272-35fa8212c056",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def create_violin_plot(df_combined, title, hline=None):\n",
    "    \"\"\"\n",
    "    Load data from a df and create a violin plot with optional horizontal line.\n",
    "    \n",
    "    Parameters:\n",
    "    - df_combined: DataFrame containing the data.\n",
    "    - hline: Optional; y-value for a horizontal line across the plot.\n",
    "    \n",
    "    Returns:\n",
    "    None.\n",
    "    \"\"\"\n",
    "    \n",
    "    df_combined['cycle'] = df_combined['cycle'].astype('int')\n",
    "    df_combined['cnnaffinity'] = df_combined['cnnaffinity'].astype('float')\n",
    "    plot_title = df_combined['expt'].iloc[0]\n",
    "\n",
    "    # create violin plot without black borders and with specified inner\n",
    "    sns.violinplot(x='cycle', y='cnnaffinity', data=df_combined, linewidth=0, edgecolor='none')\n",
    "    #ax.set_facecolor('white')\n",
    "    # add optional horizontal line\n",
    "    if hline is not None:\n",
    "        plt.axhline(y=hline, color='black', linewidth=2)  # increase thickness of hline\n",
    "    \n",
    "    plt.xlabel('Cycle')\n",
    "    plt.ylabel('Predicted pK')\n",
    "    \n",
    "    plt.xticks(range(min(df_combined['cycle']), max(df_combined['cycle']) + 1, 2))\n",
    "    plt.savefig(f'{title}')\n",
    "    plt.show()\n",
    "    return df_combined\n",
    "\n",
    "def create_test_df(directory):\n",
    "    \"\"\"\n",
    "    Concatenate all selection.csv files from each cycle directory within the specified directory.\n",
    "\n",
    "    Parameters:\n",
    "    - directory: str, the path to the directory containing cycle folders.\n",
    "\n",
    "    Returns:\n",
    "    - all_selections_df: DataFrame, containing all the data with an additional 'cycle' column.\n",
    "    \"\"\"\n",
    "    # Initialize an empty DataFrame\n",
    "    all_selections_df = pd.DataFrame()\n",
    "\n",
    "    # List all cycle directories in the specified directory\n",
    "    cycle_dirs = [d for d in os.listdir(directory) if\n",
    "                  os.path.isdir(os.path.join(directory, d)) and d.startswith('cycle_')]\n",
    "\n",
    "    # Iterate over each directory and append the data to the DataFrame\n",
    "    for cycle_dir in cycle_dirs:\n",
    "        cycle_path = os.path.join(directory, cycle_dir, 'selection.csv')\n",
    "        if os.path.exists(cycle_path):\n",
    "            # Read the selection.csv file\n",
    "            cycle_df = pd.read_csv(cycle_path)\n",
    "            print(cycle_df.columns)\n",
    "            new_name = 'cnnaffinity'\n",
    "            if 'combo1' in cycle_df.columns:\n",
    "                cycle_df.rename(columns={'combo1': new_name}, inplace=True)\n",
    "            \n",
    "            if 'plip' in cycle_df.columns:\n",
    "                cycle_df.rename(columns={'plip': new_name}, inplace=True)\n",
    "            # Add a 'cycle' column to keep track of the source\n",
    "            cycle_number = int(cycle_dir.split('_')[1])  # Extract cycle number from the directory name\n",
    "            cycle_df['cycle'] = cycle_number\n",
    "            cycle_df['cnnaffinity'] = cycle_df['cnnaffinity'].abs()\n",
    "            # cycle_df['cnnaffinity'].astype('float').dtypes\n",
    "\n",
    "            # Append the DataFrame to the main DataFrame\n",
    "            all_selections_df = pd.concat([all_selections_df, cycle_df], ignore_index=True)\n",
    "    all_selections_df['expt'] = directory\n",
    "\n",
    "    print(all_selections_df.columns)\n",
    "    return all_selections_df\n",
    "    \n",
    "os.chdir(current_dir)\n",
    "base_dir = os.path.dirname(os.path.abspath('plots.ipynb'))\n",
    "\n",
    "dirs = ['mpro-al-pK-beta01','mpro-al-pK-beta10', 'mpro-al-plip', 'mpro-al-cs']\n",
    "\n",
    "for d in dirs:\n",
    "    combined_df = create_test_df(f'{base_dir}/{d}/generated')\n",
    "    \n",
    "    combined_df = combined_df.rename(columns={'Cycle': 'cycle'})\n",
    "    create_violin_plot(combined_df, hline=None, title=f'{base_dir}/{d}.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "538feb1d-d97f-47e5-a162-52ec45c2a224",
   "metadata": {},
   "source": [
    "## AL regression vs CNNaffinity plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "118e4a13-9e3f-4e36-8d79-89a0749c3b56",
   "metadata": {},
   "outputs": [],
   "source": [
    "# al regression model vs cnnaffinity for a random rep & parameter\n",
    "os.chdir(current_dir)\n",
    "base_dir = os.path.dirname(os.path.abspath('plots.ipynb'))\n",
    "os.chdir(f'{base_dir}/rep_3/gp_300_UCB_True_10/')\n",
    "\n",
    "import os\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def create_test_df(directory):\n",
    "    \"\"\"\n",
    "    Concatenate all selection.csv files from each cycle directory within the specified directory.\n",
    "\n",
    "    Parameters:\n",
    "    - directory: str, the path to the directory containing cycle folders.\n",
    "\n",
    "    Returns:\n",
    "    - all_selections_df: DataFrame, containing all the data with an additional 'cycle' column.\n",
    "    \"\"\"\n",
    "    # Initialize an empty DataFrame\n",
    "    all_selections_df = []\n",
    "    # List all cycle directories in the specified directory\n",
    "    cycle_dirs = [d for d in os.listdir(directory) if\n",
    "                  os.path.isdir(os.path.join(directory, d)) and d.startswith('cycle_') and not d.endswith('0001')]\n",
    "    cycle_dirs = sorted(cycle_dirs, key=lambda x: int(x.split('_')[1]))\n",
    "\n",
    "\n",
    "\n",
    "    print(cycle_dirs)\n",
    "    # Iterate over each directory and append the data to the DataFrame\n",
    "    for cycle_dir in cycle_dirs:\n",
    "        cycle_path = os.path.join(directory, cycle_dir, 'virtual_library_with_predictions.csv')\n",
    "        #print(cycle_path)\n",
    "        if os.path.exists(cycle_path):\n",
    "            # Read the selection.csv file\n",
    "            cycle_df = pd.read_csv(cycle_path)\n",
    "\n",
    "            # Add a 'cycle' column to keep track of the source\n",
    "            cycle_number = int(cycle_dir.split('_')[1])  # Extract cycle number from the directory name\n",
    "            cycle_df['cycle'] = cycle_number\n",
    "            cycle_df['regression'] = cycle_df['regression'].abs()\n",
    "            # cycle_df['cnnaffinity'].astype('float').dtypes\n",
    "\n",
    "            # Append the DataFrame to the main DataFrame\n",
    "            all_selections_df.append(cycle_df)\n",
    "            #all_selections_df['expt] = directory\n",
    "    return all_selections_df\n",
    "\n",
    "oracle_csv_path = f'{base_dir}/cs_49k.csv'\n",
    "oracle_df = pd.read_csv(oracle_csv_path)\n",
    "\n",
    "directory = '.'\n",
    "cycle_dirs = [d for d in os.listdir(directory) if os.path.isdir(os.path.join(directory, d))]\n",
    "cycle_dirs\n",
    "\n",
    "a = create_test_df('.')\n",
    "#len(a)\n",
    "\n",
    "def calc_rmse(df_list):\n",
    "    rmse_list = []\n",
    "    print(f'calculating rmses for {len(df_list)} cycles')\n",
    "    for i, df in enumerate(df_list):\n",
    "        # Merge with oracle_df on 'Smiles'\n",
    "        merged_df = pd.merge(df, oracle_df, how='outer', on='Smiles')\n",
    "        # Calculate RMSE\n",
    "        rmse = mean_squared_error(merged_df['cnnaffinity_y'], merged_df['regression'], squared=False)\n",
    "        rmse_list.append(rmse)\n",
    "        print(f'rmse for cycle {i} is {rmse}')\n",
    "        print(f'df:\\n {df}')\n",
    "        #print(f'cycle number: {len(rmse_list)}')\n",
    "    return rmse_list\n",
    "\n",
    "#a[0].reset_index(\n",
    "b = pd.merge(a[-1],oracle_df, how='outer', on='Smiles')\n",
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39ce402f-54c5-44b2-8427-e647e2674527",
   "metadata": {},
   "outputs": [],
   "source": [
    "merged_df_last = b\n",
    "\n",
    "# Calculate RMSE\n",
    "# Prepare the data\n",
    "y = merged_df_last['regression']\n",
    "x = merged_df_last['cnnaffinity_y']\n",
    "\n",
    "# Create a scatter plot\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.scatter(x, y, alpha=0.1)  # alpha for transparency in case of overlapping points\n",
    "\n",
    "# Adding labels and title\n",
    "plt.xlabel('CNNaffinity')\n",
    "plt.ylabel('Regression Model Pred. CNNaffinity')\n",
    "#plt.title('Scatter Plot of RMSE vs cnnaffinity_y')\n",
    "\n",
    "# Display the plot\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15483e8c-41e8-48a8-b3fc-da9e4d01b046",
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Prepare the data\n",
    "x = merged_df_last['cnnaffinity_y']\n",
    "y = merged_df_last['regression']\n",
    "\n",
    "# Create a density plot\n",
    "plt.figure(figsize=(10, 6))\n",
    "sns.kdeplot(x=x, y=y, cmap=\"Reds\", fill=True, bw_adjust=0.5)\n",
    "\n",
    "# Adding labels\n",
    "plt.xlabel('Oracle Predicted pK')\n",
    "plt.ylabel('Regression Predicted pK')\n",
    "\n",
    "# Display the plot\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c8b21ae-ce6f-4aa6-abe2-184032f8a049",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ceda824-4167-4073-be0b-1c6914a7f1d4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3daae24c-a6f5-45d6-b5a3-48d5a157e79b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48e88fed-f57d-4aa6-a182-1b81c9abf247",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fegrow_al_data-v2",
   "language": "python",
   "name": "fegrow_al_data-v2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}