03_virtual_screening_solution.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3b484b21-037f-4538-a86c-fa330ab70b9e",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "source": [
    "# Virtual screening"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ce37a10-c5a0-47df-91dc-357bbf51a8b3",
   "metadata": {},
   "source": [
    "Notebook inspired by [Tutorial for the Teach-Discover-Treat (TDT) Competition 2014](https://github.com/sriniker/TDT-tutorial-2014/tree/master) by Sereina Riniker and Gregory Landrum."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb8cae79-c856-4a14-bc70-9a5ca268de93",
   "metadata": {},
   "source": [
    "Install dependencies if necessary (uncomment and run):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dca19722-a442-4d24-8d61-77ab7274bc5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbc7644c-69cd-4c48-9191-7b082211e8a0",
   "metadata": {},
   "source": [
    "## 1. Primary screening data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a4b6f14-7a6c-4cff-97ef-4e14e3880219",
   "metadata": {},
   "source": [
    "We will focus on screening data for *Plasmodium falciparum*, parasite that causes malaria. This disease causes over 600 thousand deaths anually, mostly among African children. First effective malaria vaccine was approved by WHO in 2023, meeting their goal of >75% efficacy (reduction in risk of getting sick). Subsequently, it was approved and distributed in Ghana and Nigeria. However, research in this area is far from over.\n",
    "\n",
    "Our dataset comes from high-throughput screening (HTS) campaigns from around 2010-2012. It was used for [TDT competition](http://www.tdtproject.org/challenge-1---malaria-hts.html). Note that this data contains imprecise labels, e.g. molecules can be false positives. For this dataset we also have a confirmatory (secondary) screen with precise labels, which we will use later.\n",
    "\n",
    "In HTS, biological activity is measured as a real value, e.g. [EC50](https://en.wikipedia.org/wiki/EC50) or [IC50](https://en.wikipedia.org/wiki/IC50), and typically standardized to z-score. Due to inherent measurement inaccuracy, it is preferred to threshold it and perform binary classification, i.e. active/non-active. Sometimes an \"ambiguous\" class is also introduced, to denote highly uncertain molecules, but here we will ignore all such compounds for simplicity.\n",
    "\n",
    "For details and papers, see:\n",
    "- [TDT challenge: Malaria High-throughput Screen](http://www.tdtproject.org/challenge-1---malaria-hts.html)\n",
    "- [\"Chemical genetics of Plasmodium falciparum\" W. A. Guiguemde et al.](https://www.nature.com/articles/nature09099)\n",
    "- [\"Global Phenotypic Screening for Antimalarials\" W. A. Guiguemde et al.](https://www.sciencedirect.com/science/article/pii/S1074552112000178?via%3Dihub)\n",
    "\n",
    "First, we will load the data. It is slightly simplified, compared to the original TDT data, and compressed in Parquet format (original CSV is really big). `label` means whether activity in HTS was high enough to mark a molecule as a hit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dae96658-cea9-4438-8cc4-cd7c6bfb8a42",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "df = pd.read_parquet(\"data/malaria_hts_train.parquet\")\n",
    "print(f\"Data size: {len(df)}\")\n",
    "\n",
    "df[\"label\"].value_counts().plot.bar(title=\"Activity labels distribution\")\n",
    "plt.show()\n",
    "\n",
    "df = df[df[\"label\"] != \"ambiguous\"]\n",
    "df[\"label\"] = df[\"label\"].map({\"false\": 0, \"true\": 1})\n",
    "\n",
    "smiles_list = df[\"SMILES\"].tolist()\n",
    "labels = df[\"label\"].values\n",
    "\n",
    "display(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77ebe4c2-089c-43ff-93ae-592f82d84f97",
   "metadata": {},
   "source": [
    "### Exercise 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0dea69f3-dde2-4141-b716-1e49b31b54d4",
   "metadata": {},
   "source": [
    "We need to prepare our data for further processing. We will use `preprocessing` module in *scikit-fingerprints* for this - [documentation](https://scikit-fingerprints.github.io/scikit-fingerprints/modules/preprocessing.html).\n",
    "\n",
    "Transform SMILES strings to RDKit `Mol` objects with `MolFromSmilesTransformer`. Use only RDKit-valid molecules with `valid_only` argument and `.transform_x_y` method.\n",
    "\n",
    "`batch_size=1000` and `verbose=True` are recommended to see the progress bar.\n",
    "\n",
    "Check how many molecules were removed during this process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "873472c3-3546-4638-9c0c-e0a9ecde2df0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from skfp.preprocessing import MolFromSmilesTransformer, MolStandardizer\n",
    "\n",
    "\n",
    "mol_from_smiles = MolFromSmilesTransformer(valid_only=True, n_jobs=-1, batch_size=1000, verbose=True)\n",
    "mols, y = mol_from_smiles.transform_x_y(smiles_list, labels)\n",
    "\n",
    "print(f\"Removed molecules: {len(smiles_list) - len(mols)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f9165b4-f42b-452e-8c94-dc30c186b054",
   "metadata": {},
   "source": [
    "## 2. Filtering"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ae7b6d7-af71-4daf-a103-cddf26620470",
   "metadata": {},
   "source": [
    "First step in virtual screening (VS) workflows is filtering, removing molecules with unwanted features. *scikit-fingerprints* implements them in `filters` module. Both property-based and structural-based filters share the same interface with `.transform()` and .`transform_x_y()` methods, that return only molecules fulfilling the filter rules.\n",
    "\n",
    "We can chain filters one after another, to apply many one after another. For example, to use all three PAINS filter sets, we can use it many times. Another useful downstream filter is ZINC druglike filter, used by ZINC database, which curates commercially available chemical compounds for virtual screening in medicinal chemistry.\n",
    "\n",
    "However, scikit-learn pipelines don't support using `.transform_x_y()` currently, so we will need to apply them manually. Note that filters can be applied in any order - resulting molecules have to pass all of them. As such, if we focus on efficiency, putting the most restrictive filters first can be useful.\n",
    "\n",
    "For ZINC filter definitions, see e.g. [scikit-fingerprints documentation](https://scikit-fingerprints.github.io/scikit-fingerprints/modules/generated/skfp.filters.ZINCDruglikeFilter.html).\n",
    "\n",
    "References:\n",
    "- [\"New Substructure Filters for Removal of Pan Assay Interference Compounds (PAINS) from Screening Libraries and for Their Exclusion in Bioassays\" J. Baell, G. Holloway](https://doi.org/10.1021/jm901137j)\n",
    "- [\"ZINC − A Free Database of Commercially Available Compounds for Virtual Screening\" J.Irwin, B. Shoichet](https://doi.org/10.1021/ci049714+)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1dd42cc-d99a-442a-a198-a89bcb4290cc",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from skfp.filters import PAINSFilter, ZINCDruglikeFilter\n",
    "\n",
    "\n",
    "verbosity_args = dict(batch_size=1000, n_jobs=-1, verbose=True)\n",
    "\n",
    "filt_pains_a = PAINSFilter(variant=\"A\", **verbosity_args)\n",
    "filt_pains_b = PAINSFilter(variant=\"B\", **verbosity_args)\n",
    "filt_pains_c = PAINSFilter(variant=\"C\", **verbosity_args)\n",
    "filt_zinc = ZINCDruglikeFilter(**verbosity_args)\n",
    "\n",
    "print(f\"Molecules before filtering: {len(mols)}\")\n",
    "\n",
    "mols_filtered, y_filtered = filt_pains_a.transform_x_y(mols, y)\n",
    "print(f\"Molecules after PAINS A: {len(mols_filtered)}\")\n",
    "\n",
    "mols_filtered, y_filtered = filt_pains_b.transform_x_y(mols_filtered, y_filtered)\n",
    "print(f\"Molecules after PAINS B: {len(mols_filtered)}\")\n",
    "\n",
    "mols_filtered, y_filtered = filt_pains_c.transform_x_y(mols_filtered, y_filtered)\n",
    "print(f\"Molecules after PAINS C: {len(mols_filtered)}\")\n",
    "\n",
    "mols_filtered, y_filtered = filt_zinc.transform_x_y(mols_filtered, y_filtered)\n",
    "print(f\"Molecules after ZINC druglike: {len(mols_filtered)}\")\n",
    "\n",
    "y_active = y_filtered.sum()\n",
    "y_inactive = len(y_filtered) - y_active\n",
    "print(f\"Final inactive molecules: {y_inactive}\")\n",
    "print(f\"Final active molecules: {y_active}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe4fa382-7a65-49e8-b62d-e75fb47aa04f",
   "metadata": {},
   "source": [
    "### Exercise 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36756adf-cf92-4ecf-a03e-febad5de8913",
   "metadata": {},
   "source": [
    "A common criticism of molecular filters is that they are too conservative, and limit the explored chemical space too much. The most famous Lipinski's Rule of 5 filter is often shown as an example of this. It was originally designed to keep orally bioavailable drugs.\n",
    "\n",
    "Try out the Lipinski's Rule of 5 filter, and two alternatives proposed in literature:\n",
    "- Beyond Rule of 5\n",
    "- Rule of Veber\n",
    "\n",
    "You will find them in [scikit-fingerprints documentation](https://scikit-fingerprints.github.io/scikit-fingerprints/modules/filters.html).\n",
    "\n",
    "Filter the initial set of ~295 thousand compounds. Use filters separately (not as a pipeline) and compare how many non-actives and actives did they leave. Was there any big difference in speed?\n",
    "\n",
    "**Note:** do not overwrite the variables `mols_filtered` and `y_filtered` from above, otherwise you will need to recompute them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "863b3586-eabb-4ce7-8c6c-02f9574f6dd2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from skfp.filters import BeyondRo5Filter, LipinskiFilter, RuleOfVeberFilter\n",
    "\n",
    "\n",
    "for filt_cls, filt_name in [\n",
    "    (LipinskiFilter, \"Lipinski\"),\n",
    "    (BeyondRo5Filter, \"Beyond Rule of 5\"),\n",
    "    (RuleOfVeberFilter, \"Rule of Veber\"),\n",
    "]:\n",
    "    filt = filt_cls(**verbosity_args)\n",
    "    mols_filtered_2, y_filtered_2 = filt.transform_x_y(mols, y)\n",
    "    y_active_2 = y_filtered_2.sum()\n",
    "    y_inactive_2 = len(labels) - y_active_2\n",
    "    print(\n",
    "        f\"Molecules after {filt_name}:\",\n",
    "        f\"{len(mols_filtered_2)},\",\n",
    "        f\"active: {y_active_2},\",\n",
    "        f\"inactive: {y_inactive_2}\"\n",
    "    )\n",
    "    print()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81b00647-8df1-41a1-a290-1cae309f1732",
   "metadata": {},
   "source": [
    "### What next?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6fbaa2d-210c-4940-8804-6225bb1d3d57",
   "metadata": {},
   "source": [
    "At this point, we typically would get a chemist involved to look over rejected active molecules. Some that look interesting to experts could be rescued, to be fixed in later parts, e.g. their toxicity reduced through structural changes. However, here we focus on a fully automated virtual screening workflow.\n",
    "\n",
    "Then, we can do two things:\n",
    "1. Confirm actives in confirmatory screen\n",
    "2. Perform virtual screening (VS)\n",
    "\n",
    "We will do the latter, since we focus on the virtual screening workflow. If you are interested, see the bonus exercise for confirmatory screen exercise."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0ec32b5-1b63-4837-9a4d-fa3f4b034b0f",
   "metadata": {},
   "source": [
    "## Ligand-based virtual screening"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffd686c3-a597-4e6c-8730-fef4f326ea4d",
   "metadata": {},
   "source": [
    "Ligand-based VS uses 2D molecular structures, and is a much cheaper and scalable alternative to the 3D structure-based VS. We use molecular fingerprints to quantify similarity to active molecules, and those most similar are assumed to also be active. This can be seen as either a binary classification (active/non-active), or similarity searching, i.e. finding molecules most similar to the active ones.\n",
    "\n",
    "Here, we will focus on classification approach, which has a few advantages:\n",
    "- typically great results\n",
    "- quite noise-resistant\n",
    "- fast and scalable\n",
    "\n",
    "Note that the classification is often highly imbalanced, so we definitely need to use class weighting or other techniques.\n",
    "\n",
    "Steps:\n",
    "1. Train-test split\n",
    "2. Make a pipeline (fingerprint + classifier), tune it and train on training set.\n",
    "3. Calculate a score (e.g. probability) for each test molecule.\n",
    "4. Rank test compounds by their score (sort descending), and evaluate this ranking with chosen metric(s).\n",
    "5. Repeat steps 1-4 many times (bootstrap), e.g. 50.\n",
    "\n",
    "We want the method to recognize structural and functional properties of actives, putting true actives from test set first. In practice, if the test set was a large library of commercially available compounds, we would take first e.g. 100 and test them in the lab. Therefore, we need early enrichment - highly probably actives as close to the beginning as possible. Metrics like enrichment factor (EF) take this into consideration. Due to high imbalance and variance, bootstrap helps with evaluation.\n",
    "\n",
    "**Additional notes**\n",
    "\n",
    "The workflow above assumes no hyperparameter tuning for fingerprints. This is quite common, but of course we could perform such tuning. It \n",
    "would require slight changes, mainly creating a validation set (perhaps with bootstrap or cross-validation), and calculating final fingerprints after tuning.\n",
    "\n",
    "The similarity search approach would work as follows. After train-test set, for each test molecule we compute similarity to actives from the train set, and take the highest similarity. This is called MAX fusion. This is the score, and we sort compounds in the order of descending similarity.\n",
    "\n",
    "Further reading:\n",
    "- [\"Open-source platform to benchmark fingerprints for ligand-based virtual screening\" S. Riniker, G. Landrum](https://doi.org/10.1186/1758-2946-5-26)\n",
    "- [\"Heterogeneous Classifier Fusion for Ligand-Based Virtual Screening: Or, How Decision Making by Committee Can Be a Good Thing\" S. Riniker et al.](https://doi.org/10.1021/ci400466r)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bdacfb1-3698-46c8-b3af-b3ebad9c0887",
   "metadata": {},
   "source": [
    "If we don't tune fingerprint hyperparameters, we can do this once before the entire training. Even if we do tune them, due to bootstrapping precomputing the actual matrices and caching them at this point can be useful. For simplicity, we will just use default settings of AtomPair fingerprint, using scikit-fingerprints and `AtomPairFingerprint` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f44cfdf-2b66-4db4-a87e-57d50ab6e232",
   "metadata": {},
   "outputs": [],
   "source": [
    "from skfp.fingerprints import AtomPairFingerprint\n",
    "\n",
    "\n",
    "fp_ap = AtomPairFingerprint(batch_size=1000, n_jobs=-1, verbose=True)\n",
    "X_ap = fp_ap.transform(mols_filtered)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29ae9bea-240c-404c-8666-7f9ab2e3d559",
   "metadata": {},
   "source": [
    "Now we need 2 functions:\n",
    "1. Classifier training, which also gets scores (probabilities) for test compounds\n",
    "2. Scores ranking and evaluation.\n",
    "\n",
    "Popular classifiers are logistic regression and tree-based ensembles like Random Forest or boosting. We will use LightGBM, since it's very fast."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "300533cc-a9fd-450f-b768-67f32f2d4d65",
   "metadata": {},
   "source": [
    "## Exercise 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "185371bd-1fb1-48d6-91e5-00e550185e66",
   "metadata": {},
   "source": [
    "Implement `get_vs_scores` function.\n",
    "\n",
    "1. Train LightGBM classifier ([documentation](https://lightgbm.readthedocs.io/en/latest/pythonapi/lightgbm.LGBMClassifier.html)) with settings:\n",
    "   - `class_weight=\"balanced\"`\n",
    "   - `n_jobs=-1`\n",
    "   - `random_state=0`\n",
    "   - `verbose=-1` - to turn off unnecessary logging\n",
    "3. Predict probability of positive class (active) for test compounds. Those are the molecule scores.\n",
    "4. Calculate metrics using scikit-learn and scikit-fingerprints ([documentation](https://scikit-fingerprints.github.io/scikit-fingerprints/modules/metrics.html)): AUROC and enrichment factor at 5% (EF5%).\n",
    "5. Return calculated metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82854e67-e46f-4ffc-8142-6e098ae3bf56",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from lightgbm import LGBMClassifier\n",
    "from sklearn.metrics import roc_auc_score\n",
    "from skfp.metrics import enrichment_factor\n",
    "\n",
    "\n",
    "def perform_vs_experiment(\n",
    "    X_train: np.ndarray,\n",
    "    X_test: np.ndarray,\n",
    "    y_train: np.ndarray,\n",
    "    y_test: np.ndarray,\n",
    ") -> tuple[float, float]:\n",
    "    clf = LGBMClassifier(n_jobs=-1, random_state=0, verbose=-1)\n",
    "    clf.fit(X_train, y_train)\n",
    "    y_pred_proba = clf.predict_proba(X_test)[:, 1]\n",
    "\n",
    "    auroc = roc_auc_score(y_test, y_pred_proba)\n",
    "    ef = enrichment_factor(y_test, y_pred_proba, fraction=0.05)\n",
    "\n",
    "    return auroc, ef\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a2d1952-c7dd-4a33-a289-d52c96591035",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    X_ap, y_filtered, test_size=0.2, random_state=0, stratify=y_filtered\n",
    ")\n",
    "\n",
    "auroc, ef = perform_vs_experiment(X_train, X_test, y_train, y_test)\n",
    "print(f\"AUROC: {auroc:.2%}\")\n",
    "print(f\"Enrichment factor at 5% (EF5%): {ef:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b42eb0ee-3a09-4861-ad46-36171bd81306",
   "metadata": {},
   "source": [
    "Since we have so few actives, single train-test split is not enough. We should perform bootstrapping, i.e. repeat the experiment many times with different splits, and average the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "222e31f0-0e1a-4848-9688-4d689ba39da5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "def perform_boostrap_vs_experiment(\n",
    "    X_fp: np.ndarray, y: np.ndarray, n_repetitions: int = 10\n",
    ") -> None:\n",
    "    auroc_values = []\n",
    "    ef_values = []\n",
    "    for i in tqdm(range(n_repetitions), total=n_repetitions):\n",
    "        X_train, X_test, y_train, y_test = train_test_split(\n",
    "            X_fp, y, test_size=0.2, random_state=i, stratify=y_filtered\n",
    "        )\n",
    "        auroc, ef = perform_vs_experiment(X_train, X_test, y_train, y_test)\n",
    "        auroc_values.append(auroc)\n",
    "        ef_values.append(ef)\n",
    "\n",
    "    auroc_mean = 100 * np.mean(auroc_values)\n",
    "    auroc_std = 100 * np.std(auroc_values)\n",
    "    ef_mean = np.mean(ef_values)\n",
    "    ef_std = np.std(ef_values)\n",
    "    \n",
    "    print(f\"AUROC: {auroc_mean:.2f} +- {auroc_std:.2f}\")\n",
    "    print(f\"Enrichment factor at 5% (EF5%): {ef_mean:.2f} +- {ef_std:.2f}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84ee4747-f3ce-4f33-9bd0-0955ef8eb678",
   "metadata": {},
   "outputs": [],
   "source": [
    "perform_boostrap_vs_experiment(X_ap, y_filtered)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04939966-7b85-4207-a98b-a0aaeb9cd9aa",
   "metadata": {},
   "source": [
    "AUROC can be easily interpreted. However, is this enrichment factor good or not? It has minimal value 0, random classifier would get value 1, but the maximum value has a slightly more complicated formula.\n",
    "\n",
    "Mark fraction as $X$, number of actives $n$, and test set size $N$. Then, the maximal value is $1/X$ if $X \\geq n/N$, and $N/n$ otherwise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4ea7897-21d0-4854-920f-0565875b6d40",
   "metadata": {},
   "outputs": [],
   "source": [
    "fraction = 0.05\n",
    "n = sum(y_test)\n",
    "N = len(y_test)\n",
    "\n",
    "max_ef = 1 / fraction if fraction >= n/N else N/b\n",
    "\n",
    "print(f\"Maximal EF5% value: {max_ef:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4b1a46e-e0b1-462f-a18f-d938c7bee344",
   "metadata": {},
   "source": [
    "## Exercise 4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ce3d7a1-9d52-498b-a960-436a52f80ece",
   "metadata": {},
   "source": [
    "Perform the virtual screening experiment with other fingerprints:\n",
    "- ECFP4\n",
    "- RDKit fingerprint\n",
    "- MACCS\n",
    "- Laggner\n",
    "\n",
    "Use scikit-fingerprints [documentation](https://scikit-fingerprints.github.io/scikit-fingerprints/modules/fingerprints.html) as needed. Which fingeprint got the best result?\n",
    "\n",
    "If you have the time, you can also try out the count variants."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2933507-f74e-4e2e-bfab-8dcd0999ed09",
   "metadata": {},
   "outputs": [],
   "source": [
    "from skfp.fingerprints import ECFPFingerprint, RDKitFingerprint, MACCSFingerprint, LaggnerFingerprint\n",
    "\n",
    "\n",
    "for fp_cls, fp_name in [\n",
    "    (ECFPFingerprint, \"ECFP\"),\n",
    "    (RDKitFingerprint, \"RDKit\"),\n",
    "    (MACCSFingerprint, \"MACCS\"),\n",
    "    (LaggnerFingerprint, \"Laggner\"),\n",
    "]:\n",
    "    fp = fp_cls(batch_size=1000, n_jobs=-1, verbose=True)\n",
    "    X_fp = fp.transform(mols_filtered)\n",
    "    print(f\"{fp_name} fingerprint\")\n",
    "    perform_boostrap_vs_experiment(X_fp, y_filtered)\n",
    "    print()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eccdbe9d-83ac-4f74-ba4c-b809db2b82fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from skfp.fingerprints import ECFPFingerprint, RDKitFingerprint, MACCSFingerprint, LaggnerFingerprint\n",
    "\n",
    "\n",
    "for fp_cls, fp_name in [\n",
    "    (ECFPFingerprint, \"ECFP\"),\n",
    "    (RDKitFingerprint, \"RDKit\"),\n",
    "    (MACCSFingerprint, \"MACCS\"),\n",
    "    (LaggnerFingerprint, \"Laggner\"),\n",
    "]:\n",
    "    fp = fp_cls(count=True, batch_size=1000, n_jobs=-1, verbose=True)\n",
    "    X_fp = fp.transform(mols_filtered)\n",
    "    print(f\"{fp_name} count fingerprint\")\n",
    "    perform_boostrap_vs_experiment(X_fp, y_filtered)\n",
    "    print()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cc30666-521f-4922-8ad8-f278fb1052a6",
   "metadata": {},
   "source": [
    "Now, that we have the final model selected, real-world deployment steps would be:\n",
    "1. Retrain model on all available data\n",
    "2. Download a large-scale virtual screening collection, e.g. [Enamine Screening Collection](https://enamine.net/compound-collections/screening-collection) of ~4.4 million compounds\n",
    "3. Filter and predict probabilities, selecting most likely molecules\n",
    "4. Select diverse compounds\n",
    "5. Perform confirmatory screening"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6087fbb-d518-42c8-bd99-282a8c83e804",
   "metadata": {},
   "source": [
    "## Bonus exercise"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e95b2da5-1719-42f7-a167-148f4f8c4da5",
   "metadata": {},
   "source": [
    "We got about 1000 compounds after filtering. Testing them thoroughly in the lab may be too expensive, so we want to pick a subset for purely economic reasons, e.g. 500.\n",
    "\n",
    "Some of those may turn active, but toxic, expensive to synthesize, or be rejected for other reasons. Therefore, we want diverse molecules, to have many alternatives. This is the problem of **maximum diversity picking (MDP)**. It is typically approximated with heuristics, such as clustering.\n",
    "\n",
    "Molecules typically use ECFP vectorization. We first pick diverse cluster centers (centroid molecules), e.g. with Butina (sphere exclusion) clustering - a density-based clustering frequently used for molecules. Alternatively, MaxMin algorithm can be used, which tries to maximize the sum of Tanimoto distances between molecules. Each molecule is then assigned to its nearest cluster based on Tanimoto similarity to the cluster centroid. We can pick molecules from each cluster until we have as many molecules as we can test.\n",
    "\n",
    "1. Implement Butina (sphere exclusion / leader following) clustering and cluster the active molecules after filtering. [This tutorial](https://greglandrum.github.io/rdkit-blog/posts/2020-11-18-sphere-exclusion-clustering.html) will be useful.\n",
    "2. Plot molecules from cluster centers - are they diverse? Do their general structure (scaffolds) look different?\n",
    "3. Pick 1 molecule from each cluster randomly in a loop, to get 500 molecules.\n",
    "4. Calculate pairwise Tanimoto similarity between resulting molecule set to quantitatively assess their diversity. Plot the distribution. Are they generally below 0.65, commonly used as similarity threshold?\n",
    "\n",
    "The resulting list would be submitted for synthesis in a confirmatory screen. [TDT data](http://www.tdtproject.org/challenge-1---malaria-hts.html) provides EC50 values for molecules that were tested by them in confirmatory screen, and you can also compare our molecules (based on SMILES) with those. See [this tutorial](https://github.com/sriniker/TDT-tutorial-2014/blob/master/TDT_challenge_tutorial.ipynb) for details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80590d7d-f8fe-404a-8123-619f07fd0348",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit.Chem import rdMolDescriptors\n",
    "from rdkit.Chem.rdFingerprintGenerator import GetMorganGenerator\n",
    "from rdkit.SimDivFilters.rdSimDivPickers import LeaderPicker\n",
    "\n",
    "\n",
    "fp_gen = GetMorganGenerator(radius=2)\n",
    "fps = fp_gen.GetFingerprints(mols_filtered, numThreads=6)\n",
    "fps = list(fps)\n",
    "\n",
    "picker = LeaderPicker()\n",
    "centroid_idxs = picker.LazyBitVectorPick(fps, len(fps), threshold=0.65)\n",
    "centroid_idxs = list(centroid_idxs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a8497e9-ac29-4891-b703-68cd42c8aef0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "\n",
    "import numpy as np\n",
    "from rdkit.DataStructs import BulkTanimotoSimilarity\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "def assign_points_to_clusters(centroid_idxs: list, fps: list) -> list[list]:\n",
    "    # initialize clusters with centers\n",
    "    clusters = {\n",
    "        cluster_idx: [mol_idx] for cluster_idx, mol_idx in enumerate(centroid_idxs)\n",
    "    }\n",
    "    centroid_fps = [fps[i] for i, fp in enumerate(centroid_idxs) if i in centroid_idxs]\n",
    "    \n",
    "    for idx, fp in tqdm(enumerate(fps), total=len(fps)):\n",
    "        # skip cluster centers, they are already assigned\n",
    "        if idx in cluster_centers:\n",
    "            continue\n",
    "\n",
    "        sims = BulkTanimotoSimilarity(fp, centroid_fps)\n",
    "        most_sim_cluster_idx = np.argmax(sims)\n",
    "        clusters[most_sim_cluster_idx].append(idx)\n",
    "\n",
    "    return clusters\n",
    "\n",
    "\n",
    "clustering = assign_points_to_clusters(centroid_idxs, fps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62e01248-7a9e-4457-b757-bcd85f7d7c9d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from rdkit.Chem.Draw import MolsToGridImage\n",
    "\n",
    "centrid_mols = [mols_filtered[idx] for idx in centroid_idxs]\n",
    "MolsToGridImage(centrid_mols[:16], molsPerRow=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3473ed3a-b30d-49be-9076-feff34f98a73",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "picks = []\n",
    "while len(picks) < 500:\n",
    "    for cluster_idx, mol_idxs in clustering.items():\n",
    "        mol_idx = random.choice(mol_idxs)\n",
    "        if mol_idx not in picks:\n",
    "            picks.append(mol_idx)\n",
    "\n",
    "picked_mols = [mols_filtered[idx] for idx in picks]\n",
    "picked_mols[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c73aa1b-038f-4ccb-ac29-da73b914efbd",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "picked_fps = [fps[idx] for idx in picks]\n",
    "\n",
    "all_sims = []\n",
    "for idx, fp in tqdm(enumerate(picked_fps), total=len(picked_fps)):\n",
    "    mol_sims = BulkTanimotoSimilarity(fp, picked_fps)\n",
    "    \n",
    "    # remove self-similarity equal to 1.0\n",
    "    del mol_sims[idx]\n",
    "    \n",
    "    all_sims.extend(mol_sims)\n",
    "\n",
    "pd.Series(all_sims).plot.hist(title=\"Pairwise Tanimoto similarities distribution\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fac266b-f5d0-41a9-afa6-1362e4f257db",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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
}