notebooks/02_scaffold_split_with_custom_estimators.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "beaedbca-4728-443d-a8fc-1d79f9fee27a",
   "metadata": {},
   "source": [
    "# Custom Sklearn estimators for cheminformatics: Scaffold splits example\n",
    "\n",
    "Molpipeline provides some custom Sklearn-like estimators solving common cheminformatics tasks. The estimators comply the [Sklearn estimator API](https://scikit-learn.org/stable/developers/develop.html) and can be used in pipelines.\n",
    "\n",
    "This notebook shows how to use the MurckoScaffoldClustering estimator for generating scaffold splits for molecular machine learning. It applies the widely used [Murcko-type decomposition](https://www.rdkit.org/docs/GettingStartedInPython.html#murcko-decomposition) to a molecule data sets. From the decomposition a clustering is generated which then can be directly used with Sklearn's group-based splitters for cluster cross-validation.\n",
    "\n",
    "This is a simple example notebook using dummy data to illustrate the usage of custom estimators for cheminformatics like MurckoScaffoldClustering. Please look at the advanced notebooks for more detailed examples. \n",
    "\n",
    "[**Scaffold clustering**](#estimators)\n",
    "* Murcko-scaffolds and generic scaffolds with MurckoScaffoldClustering estimator\n",
    "\n",
    "[**Putting it together**](#fullexample)\n",
    "* Train and evaluate a classfier with MolPipeline\n",
    "* Cross validation evaluation with scaffold-split\n",
    "    * Combine MurckoScaffoldClustering with Sklearn's GroupKFold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fd299f4f-3cb5-419e-a992-d3ce90a2ec2f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to find the pandas get_adjustment() function to patch\n",
      "Failed to patch pandas - PandasTools will have limited functionality\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "from molpipeline import Pipeline\n",
    "from molpipeline.any2mol import AutoToMol\n",
    "from molpipeline.mol2mol import ElementFilter\n",
    "from molpipeline.mol2any import MolToMorganFP\n",
    "from molpipeline.estimators import MurckoScaffoldClustering\n",
    "\n",
    "import pandas as pd\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import PandasTools\n",
    "\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.metrics import roc_auc_score\n",
    "from sklearn.model_selection import GroupKFold"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97508e86-2d13-4a79-808c-dcf53f927156",
   "metadata": {},
   "source": [
    "## Scaffold clustering <a class=\"anchor\" id=\"estimators\"></a>\n",
    "\n",
    "MolPipeline implements custom Sklearn estimators for standard molecular machine learning tasks. For example, we created a MurckoScaffoldClustering estimator which can be used like a normal Sklearn clustering estimator.   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2f90d155-80bc-4f2b-9f6e-d44405f83f56",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 1.,  0.,  1., nan, nan, nan])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scaffold_smiles = [\n",
    "    \"Nc1ccccc1\",\n",
    "    \"Cc1cc(Oc2nccc(CCC)c2)ccc1\",\n",
    "    \"c1ccccc1\",\n",
    "]\n",
    "linear_smiles = [\"CC\", \"CCC\", \"CCCN\"]\n",
    "\n",
    "# run the scaffold clustering\n",
    "scaffold_clustering = MurckoScaffoldClustering(\n",
    "    n_jobs=1, linear_molecules_strategy=\"ignore\"\n",
    ")\n",
    "scaffold_clustering.fit_predict(scaffold_smiles + linear_smiles)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa305965-28d6-4b5b-9fb7-df54fd9b7453",
   "metadata": {},
   "source": [
    "The cluster labels of the above scaffold clustering assigns nan to linear molecules. This is because we used linear_molecules_strategy=\"ignore\". Instead we can also use the \"own_cluster\" strategy which groups all linear molecules in a new cluster:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ac6270f0-b656-4552-8891-bf26b28c5bf7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1., 0., 1., 2., 2., 2.])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scaffold_clustering = MurckoScaffoldClustering(\n",
    "    n_jobs=1, linear_molecules_strategy=\"own_cluster\"\n",
    ")\n",
    "scaffold_clustering.fit_predict(scaffold_smiles + linear_smiles)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "112b7f01-4da8-48ae-a3cd-4d9014e082dc",
   "metadata": {},
   "source": [
    "In addition, instead of using the basic Murcko scaffolds for clustering we can cluster with the generic scaffolds. Generic scaffolds are Murcko scaffolds but with all atom elements set to carbons. We can do this by setting make_generic=True like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "be53c55b-4754-441f-86f2-dc2a3ce6631c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1., 0., 1., 2., 2., 2., 1.])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scaffold_clustering = MurckoScaffoldClustering(\n",
    "    n_jobs=1, linear_molecules_strategy=\"own_cluster\", make_generic=True\n",
    ")\n",
    "scaffold_clustering.fit_predict(scaffold_smiles + linear_smiles + [\"c1ncccc1\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bb2c70b-34ae-49c3-9044-6abb61f839d2",
   "metadata": {},
   "source": [
    "## Cross validation with scaffold split <a class=\"anchor\" id=\"fullexample\"></a>\n",
    "\n",
    "A cross validation with scaffold splits is straightfoward to implement with MurckoScaffoldClustering. We can simply combine the generated scaffold clusters with Sklearn's group-based splitters, like [GroupKFold](https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.GroupKFold.html), [StratifiedGroupKFold](https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.StratifiedGroupKFold.html), [LeaveOneGroupOut](https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.LeaveOneGroupOut.html#sklearn.model_selection.LeaveOneGroupOut), etc."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e3d67a4-7588-4b01-bef3-c140e26b3f46",
   "metadata": {},
   "source": [
    "Let's setup some data and do the clustering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7d64781a-919f-4958-9d60-935e5927594a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to patch pandas - unable to change molecule rendering\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>has_nitrogen_label</th>\n",
       "      <th>murcko_clusters</th>\n",
       "      <th>Molecule</th>\n",
       "      <th>Murcko_SMILES</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>smiles</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Nc1ccccc1</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>&lt;rdkit.Chem.rdchem.Mol object at 0x7f8bdbf40580&gt;</td>\n",
       "      <td>c1ccccc1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Cc1cc(Oc2nccc(CCC)c2)ccc1</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>&lt;rdkit.Chem.rdchem.Mol object at 0x7f8bdbf403c0&gt;</td>\n",
       "      <td>c1ccc(Oc2ccccn2)cc1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>c1ccccc1</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>&lt;rdkit.Chem.rdchem.Mol object at 0x7f8bdbf40510&gt;</td>\n",
       "      <td>c1ccccc1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>CCCCN</th>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>&lt;rdkit.Chem.rdchem.Mol object at 0x7f8bdbf405f0&gt;</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>CCC</th>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>&lt;rdkit.Chem.rdchem.Mol object at 0x7f8bdbf40660&gt;</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>CCO</th>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>&lt;rdkit.Chem.rdchem.Mol object at 0x7f8bdbf409e0&gt;</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Oc1ccccc1</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>&lt;rdkit.Chem.rdchem.Mol object at 0x7f8bdbf40a50&gt;</td>\n",
       "      <td>c1ccccc1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Oc1ccc(N)cc1</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>&lt;rdkit.Chem.rdchem.Mol object at 0x7f8bdbf40ac0&gt;</td>\n",
       "      <td>c1ccccc1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                           has_nitrogen_label  murcko_clusters  \\\n",
       "smiles                                                           \n",
       "Nc1ccccc1                                   1                1   \n",
       "Cc1cc(Oc2nccc(CCC)c2)ccc1                   1                0   \n",
       "c1ccccc1                                    0                1   \n",
       "CCCCN                                       1                2   \n",
       "CCC                                         0                2   \n",
       "CCO                                         0                2   \n",
       "Oc1ccccc1                                   0                1   \n",
       "Oc1ccc(N)cc1                                1                1   \n",
       "\n",
       "                                                                   Molecule  \\\n",
       "smiles                                                                        \n",
       "Nc1ccccc1                  <rdkit.Chem.rdchem.Mol object at 0x7f8bdbf40580>   \n",
       "Cc1cc(Oc2nccc(CCC)c2)ccc1  <rdkit.Chem.rdchem.Mol object at 0x7f8bdbf403c0>   \n",
       "c1ccccc1                   <rdkit.Chem.rdchem.Mol object at 0x7f8bdbf40510>   \n",
       "CCCCN                      <rdkit.Chem.rdchem.Mol object at 0x7f8bdbf405f0>   \n",
       "CCC                        <rdkit.Chem.rdchem.Mol object at 0x7f8bdbf40660>   \n",
       "CCO                        <rdkit.Chem.rdchem.Mol object at 0x7f8bdbf409e0>   \n",
       "Oc1ccccc1                  <rdkit.Chem.rdchem.Mol object at 0x7f8bdbf40a50>   \n",
       "Oc1ccc(N)cc1               <rdkit.Chem.rdchem.Mol object at 0x7f8bdbf40ac0>   \n",
       "\n",
       "                                 Murcko_SMILES  \n",
       "smiles                                          \n",
       "Nc1ccccc1                             c1ccccc1  \n",
       "Cc1cc(Oc2nccc(CCC)c2)ccc1  c1ccc(Oc2ccccn2)cc1  \n",
       "c1ccccc1                              c1ccccc1  \n",
       "CCCCN                                           \n",
       "CCC                                             \n",
       "CCO                                             \n",
       "Oc1ccccc1                             c1ccccc1  \n",
       "Oc1ccc(N)cc1                          c1ccccc1  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# a list of dummy smiles\n",
    "smiles_data = np.array(\n",
    "    [\n",
    "        \"Nc1ccccc1\",\n",
    "        \"Cc1cc(Oc2nccc(CCC)c2)ccc1\",\n",
    "        \"c1ccccc1\",\n",
    "        \"CCCCN\",\n",
    "        \"CCC\",\n",
    "        \"CCO\",\n",
    "        \"Oc1ccccc1\",\n",
    "        \"Oc1ccc(N)cc1\",\n",
    "    ]\n",
    ")\n",
    "# a simple dummy target variable y that indicates whether the molecule contains a nitrogen (1=has N, 0=no N)\n",
    "has_nitrogen_label = np.array([1, 1, 0, 1, 0, 0, 0, 1])\n",
    "\n",
    "# we cluster the molecules by their murcko scaffold for our cross validation split\n",
    "scaffold_clustering = MurckoScaffoldClustering(\n",
    "    n_jobs=1, linear_molecules_strategy=\"own_cluster\"\n",
    ")\n",
    "groups = scaffold_clustering.fit_predict(smiles_data)\n",
    "\n",
    "# let's look at the data in a nice dataframe\n",
    "df = pd.DataFrame(\n",
    "    {\n",
    "        \"smiles\": smiles_data,\n",
    "        \"has_nitrogen_label\": has_nitrogen_label,\n",
    "        \"murcko_clusters\": groups.astype(int),\n",
    "    }\n",
    ")\n",
    "PandasTools.AddMoleculeColumnToFrame(df, \"smiles\", \"Molecule\")\n",
    "PandasTools.AddMurckoToFrame(df, molCol=\"Molecule\")\n",
    "df.set_index(\"smiles\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf4dad46-7698-4f36-affc-86f13ba82559",
   "metadata": {},
   "source": [
    "Now we can run a cross validation using GroupKFold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "91646c3c-a173-401e-8cf4-071b84b38994",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 0: Nc1ccccc1 prediction=0.31 label=1\n",
      "fold 0: c1ccccc1 prediction=0.20 label=0\n",
      "fold 0: Oc1ccccc1 prediction=0.22 label=0\n",
      "fold 0: Oc1ccc(N)cc1 prediction=0.27 label=1\n",
      "fold 0: test ROC AUC score: 1.0\n",
      "-------------------\n",
      "fold 1: Cc1cc(Oc2nccc(CCC)c2)ccc1 prediction=0.29 label=1\n",
      "fold 1: CCCCN prediction=0.30 label=1\n",
      "fold 1: CCC prediction=0.24 label=0\n",
      "fold 1: CCO prediction=0.27 label=0\n",
      "fold 1: test ROC AUC score: 1.0\n",
      "-------------------\n"
     ]
    }
   ],
   "source": [
    "# setup a splitter that handles the cluster split for us\n",
    "grouper = GroupKFold(n_splits=2)\n",
    "grouper.random_state = 67056\n",
    "\n",
    "for i, (train, test) in enumerate(\n",
    "    grouper.split(smiles_data, has_nitrogen_label, groups=groups)\n",
    "):\n",
    "\n",
    "    # setup the pipeline\n",
    "    pipeline = Pipeline(\n",
    "        [(\"auto2mol\", AutoToMol())]\n",
    "        + [\n",
    "            (\n",
    "                \"morgan\",\n",
    "                MolToMorganFP(n_bits=1024, radius=2),\n",
    "            ),\n",
    "            (\"RandomForestClassifier\", RandomForestClassifier(random_state=67056)),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    # fit the pipeline to the training data\n",
    "    pipeline.fit(X=smiles_data[train], y=has_nitrogen_label[train])\n",
    "\n",
    "    # evaluate the pipeline on the test set\n",
    "    predictions = pipeline.predict_proba(\n",
    "        X=smiles_data[test],\n",
    "    )\n",
    "\n",
    "    # print the performance for predicting the presence of nitrogens on the test set\n",
    "    for smi, pred, label in zip(\n",
    "        smiles_data[test], predictions[:, 1], has_nitrogen_label[test]\n",
    "    ):\n",
    "        print(f\"fold {i}:\", smi, f\"prediction={pred:.2f}\", f\"label={label}\")\n",
    "    print(\n",
    "        f\"fold {i}:\",\n",
    "        \"test ROC AUC score:\",\n",
    "        roc_auc_score(has_nitrogen_label[test], predictions[:, 1]),\n",
    "    )\n",
    "    print(\"-------------------\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a8f2b31-3614-4740-a528-39b0fba4f4a7",
   "metadata": {},
   "source": [
    "The results above show that in this dummy example the presence of a nitrogen can be learned and predicted with a ROC AUC score of 1.0 on the test sets of the scaffold split.   "
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}