01_rdkit_intro_solution.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3b484b21-037f-4538-a86c-fa330ab70b9e",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "source": [
    "# RDKit introduction"
   ]
  },
  {
   "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": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## 1. Loading and visualizing molecules"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a4b6f14-7a6c-4cff-97ef-4e14e3880219",
   "metadata": {},
   "source": [
    "SMILES strings are typically used for exchanging molecules. Let's load an example molecule - [ibuprofen](https://en.wikipedia.org/wiki/Ibuprofen) - with RDKit.\n",
    "\n",
    "Things to note about RDKit:\n",
    "- main module is `rdkit.Chem`, with most common functions directly available\n",
    "- it uses CamelCase, following C++ conventions\n",
    "- `MolFromSmiles` automatically performs basic sanitization and cleanup checks ([details](https://www.rdkit.org/docs/RDKit_Book.html#molecular-sanitization))\n",
    "- Jupyter automatically renders molecules from the last line"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dae96658-cea9-4438-8cc4-cd7c6bfb8a42",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit.Chem import MolFromSmiles\n",
    "\n",
    "smiles = \"CC(C)Cc1ccc(cc1)[C@@H](C)C(=O)O\"\n",
    "ibuprofen_mol = MolFromSmiles(smiles)\n",
    "ibuprofen_mol"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83f46161-d333-44fe-97e8-fb5cd11f3fd5",
   "metadata": {},
   "source": [
    "Note that SMILES is not unique. Each software deterministically outputs its own **canonical SMILES**, but they differ between programs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d94479d-74ed-4f28-9e9e-cd17d3e468f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit.Chem import MolToSmiles\n",
    "\n",
    "MolToSmiles(ibuprofen_mol)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd14f7d9-be32-4a46-8b65-20e29ecf2cd2",
   "metadata": {},
   "source": [
    "If you need more control over drawing, or are using regular Python scripts, `rdkit.Chem.Draw` contains necessary functions. It operates on `pillow` images, which are sometimes much less convenient than `matplotlib` plots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3deade90-ce2f-421a-88f3-83169bd824f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit.Chem.Draw import MolToImage\n",
    "\n",
    "\n",
    "img = MolToImage(ibuprofen_mol)\n",
    "print(type(img))\n",
    "ibuprofen_mol"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77ebe4c2-089c-43ff-93ae-592f82d84f97",
   "metadata": {},
   "source": [
    "### Exercise 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0dea69f3-dde2-4141-b716-1e49b31b54d4",
   "metadata": {},
   "source": [
    "Read [caffeine molecule](https://en.wikipedia.org/wiki/Caffeine) from Wikipedia and plot it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "873472c3-3546-4638-9c0c-e0a9ecde2df0",
   "metadata": {},
   "outputs": [],
   "source": [
    "caffeine_mol = MolFromSmiles(\"CN1C=NC2=C1C(=O)N(C(=O)N2C)C\")\n",
    "caffeine_mol"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12b5934d-4029-4819-a758-48dfb210e090",
   "metadata": {},
   "source": [
    "Note that RDKit by default uses **implicit hydrogens** where possible. To use them explicitly, e.g. for drawing, we use `AddHs` function. It returns new molecule with all hydrogens explicitly present."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8f74e60-9c58-4df6-a251-56beca716757",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit.Chem import AddHs\n",
    "\n",
    "AddHs(ibuprofen_mol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c38bd985-0ebe-4c2c-bdf7-f9a05b0006bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "AddHs(caffeine_mol)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e15b2b8b-0a70-4458-900e-0029684720b7",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## 2. Atoms, bonds, rings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c99ece33-3c3a-4988-9582-922a34ea6995",
   "metadata": {},
   "source": [
    "Molecules (`Mol` class) are built of atoms (`Atom` class) and bonds (`Bond` class). They are numbered, but everything on graphs is **permutation invariant**, i.e. this order is artificial, and purely for convenience.\n",
    "\n",
    "Rings are simply sequences of atoms and bonds. Information about them is in `RingInfo` object for each molecule.\n",
    "\n",
    "All are available in `rdkit.Chem.rdchem` ([docs](https://www.rdkit.org/docs/source/rdkit.Chem.rdchem.html)), but also in `rdkit.Chem`.\n",
    "\n",
    "RDKit is very explicit about getting many objects. No typical Python convenience for iteration or nice printing by default!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0ed051e-12df-4ab6-85fc-22fcb49c80d9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for atom in ibuprofen_mol.GetAtoms():\n",
    "    print(atom)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c97548a-85f5-49fc-950f-89cd3bca6ff4",
   "metadata": {},
   "outputs": [],
   "source": [
    "for atom in ibuprofen_mol.GetAtoms():\n",
    "    print(f\"{atom.GetIdx():2}\", atom.GetAtomicNum(), atom.GetSymbol())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce236623-896d-4b92-9b20-e62e8e0c349a",
   "metadata": {},
   "outputs": [],
   "source": [
    "for bond in ibuprofen_mol.GetBonds():\n",
    "    begin_atom = bond.GetBeginAtom()\n",
    "    end_atom = bond.GetEndAtom()\n",
    "    print(\n",
    "        f\"{bond.GetIdx():2}\",\n",
    "        begin_atom.GetSymbol(),\n",
    "        end_atom.GetSymbol(),\n",
    "        f\"{bond.GetBondType()}\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b6920ab-91bf-4bb0-8698-27b05fd90908",
   "metadata": {},
   "source": [
    "There is no easy way to know which atom exactly has a given index. To plot this, we have to explicitly use `SetAtomMapNum` and assign the atom its own index as a number."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78b4e2fd-85db-4895-9c0d-2e530c9d78ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "for atom in ibuprofen_mol.GetAtoms():\n",
    "    atom.SetAtomMapNum(atom.GetIdx())\n",
    "\n",
    "ibuprofen_mol"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb0c0cd9-0ef4-47e6-9986-9e6e6349b7fd",
   "metadata": {},
   "source": [
    "### Exercise 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9de779c1-cf82-4351-8323-ab5fb23cc7ac",
   "metadata": {},
   "source": [
    "Counts of elements and bond types are arguably the simplest features available in chemistry. We used it e.g. in [MOLTOP model](https://arxiv.org/abs/2407.12136).\n",
    "\n",
    "Implement `get_simple_element_counts` function, which returns tuple with 3 integers: counts of carbons, oxygens, and nitrogens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d26cc36-22e7-45ea-a5fe-f12f0c1d098a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit.Chem import Mol\n",
    "\n",
    "\n",
    "def get_simple_element_counts(mol: Mol) -> tuple[int, int, int]:\n",
    "    carbons = sum(atom.GetSymbol() == \"C\" for atom in mol.GetAtoms())\n",
    "    oxygens = sum(atom.GetSymbol() == \"O\" for atom in mol.GetAtoms())\n",
    "    nitrogens = sum(atom.GetSymbol() == \"N\" for atom in mol.GetAtoms())\n",
    "\n",
    "    return carbons, oxygens, nitrogens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ffa4cc9-01af-4042-8f0a-b2310a8e0ba0",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert get_simple_element_counts(ibuprofen_mol) == (13, 2, 0)\n",
    "assert get_simple_element_counts(caffeine_mol) == (8, 2, 4)\n",
    "\n",
    "carbons, oxygens, nitrogens = get_simple_element_counts(ibuprofen_mol)\n",
    "print(f\"Ibuprofen | carbons: {carbons}, oxygens: {oxygens}, nitrogens: {nitrogens}\")\n",
    "\n",
    "carbons, oxygens, nitrogens = get_simple_element_counts(caffeine_mol)\n",
    "print(f\"Caffeine  | carbons: {carbons}, oxygens: {oxygens}, nitrogens: {nitrogens}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c918c5a8-1adc-42d5-885b-05f81224b11f",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## 3. Rings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fddc8d53-0751-4823-aa5f-71a6df72c2a6",
   "metadata": {},
   "source": [
    "Rings information can be obtained with `.GetRingInfo()`. Note that its `.AtomRings()` method returns atom indexes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b80beaa7-5cbc-45d6-980a-218dcd846725",
   "metadata": {},
   "outputs": [],
   "source": [
    "ri = ibuprofen_mol.GetRingInfo()\n",
    "\n",
    "print(f\"Number of rings: {ri.NumRings()}\")\n",
    "print(f\"Rings: {ri.AtomRings()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4de21996-4156-44ca-82c0-a9193eef2ff1",
   "metadata": {},
   "source": [
    "### Exercise 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f9a6da8-ded5-46ea-9159-d8a8335e6cd0",
   "metadata": {},
   "source": [
    "Implement function `get_num_aromatic_rings`, which counts the **aromatic rings** in a molecule. Ring is aromatic if it consists only of aromatic atoms (or bonds, it's equivalent). Both atoms and bonds have `.GetIsAromatic()` method, which checks aromaticity. Molecules have `.GetAtomWithIdx()` method to get atom by index.\n",
    "\n",
    "Test it on ibuprofen and caffeine molecules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42433f20-27d8-4e96-a2e1-a0a295a34f53",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Iterable\n",
    "\n",
    "\n",
    "def get_num_aromatic_rings(mol: Mol) -> int:\n",
    "    ri = mol.GetRingInfo()\n",
    "    counter = 0\n",
    "    for ring in ri.AtomRings():\n",
    "        counter += all(\n",
    "            mol.GetAtomWithIdx(atom_idx).GetIsAromatic() for atom_idx in ring\n",
    "        )\n",
    "\n",
    "    return counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97107a3e-c50c-4a6b-8374-688c5646d180",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert get_num_aromatic_rings(ibuprofen_mol) == 1\n",
    "assert get_num_aromatic_rings(caffeine_mol) == 2\n",
    "\n",
    "print(f\"Ibuprofen aromatic rings: {get_num_aromatic_rings(ibuprofen_mol)}\")\n",
    "print(f\"Caffeine aromatic rings: {get_num_aromatic_rings(caffeine_mol)}\")\n",
    "\n",
    "print(\"Solution is correct!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cc84865-2090-4a3a-ae73-f9d54778ef3d",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## 4. Molecular descriptors"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fea4cac4-1898-4391-a1de-b81ec1b36d0a",
   "metadata": {},
   "source": [
    "To use computational methods on molecules, we need to compute some numerical properties. In particular, **descriptors** are particular properties of molecule, for example:\n",
    "- element counts\n",
    "- bond type counts\n",
    "- atomic masses\n",
    "- electric charges\n",
    "- various topological indexes\n",
    "\n",
    "They are very interpretable features, and often quite fast to compute. Using many well-selected descriptors can result in powerful models for predicting more complex properties.\n",
    "\n",
    "In RDKit, they are in a few different places, e.g. in `rdkit.Chem.Descriptors`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "493a7085-ddbf-4134-a7b1-95f778e6d2c1",
   "metadata": {},
   "source": [
    "**Molecular mass** (or weight) is the sum of masses of all atoms in a molecule, typically counted in [daltons (Da)](https://en.wikipedia.org/wiki/Dalton_(unit)). Too large molecules are more rarely used, since it's hard for them to go through membranes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f19f78e2-6c78-449f-a4b6-34a6376d9b73",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit.Chem.Descriptors import MolWt\n",
    "\n",
    "\n",
    "MolWt(ibuprofen_mol)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98d6f878-dd6b-4ea6-a5c4-fdbcabf27824",
   "metadata": {},
   "source": [
    "**Partition coefficient P** measures **lipophilicity** of molecule, i.e. how well it dissolves into fats (lipids). If molecule is lipophilic, it is also hydrophobic, i.e. does not dissolve well in water. This is also very important, since it strongly affects how easily the drug can reach its intended target in the body, and how well it will work when it reaches its target. Fortunately, it can be shown that we can quite well approximated by simple atomic contributions. Each element has its own contribution (positive or negative), and we sum them for a molecule. For details, see:\n",
    "> Wildman, S.A. and Crippen, G.M. (1999) \"Prediction of Physicochemical Parameters by Atomic Contribution\" Journal of Chemical Information and Computer Sciences, 39, 868-873. ([link](http://dx.doi.org/10.1021/ci990307l))\n",
    "\n",
    "Other references:\n",
    "- [LogP - Making Sense of the Value](https://www.acdlabs.com/wp-content/uploads/download/app/physchem/making_sense.pdf)\n",
    "- [Understanding Lipinski’s Rule of 5 and the Role of LogP Value in Drug Design and Development](https://www.sailife.com/understanding-lipinskis-rule-of-5-and-the-role-of-logp-value-in-drug-design-and-development/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64acc2cc-360c-4975-994e-7bc18fb78abe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit.Chem.Crippen import MolLogP\n",
    "\n",
    "\n",
    "MolLogP(ibuprofen_mol)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1362db7d-8a22-4783-bb40-50682463dc99",
   "metadata": {},
   "source": [
    "**Topological indices** describe the topology, or shape of the molecule, by simple numerical statistics. There are a lot, since they are closely tied to mathematics and graph theory, and molecules are graphs. It turns out that simple statistics of e.g. shortest paths of a molecule can be quite powerful in representing its properties as a single number (or a few such numbers). We used similar approaches in [LTP](https://arxiv.org/abs/2305.00724) and [MOLTOP](https://arxiv.org/abs/2407.12136) models.\n",
    "\n",
    "Arguably the first computational approach in chemistry, using descriptors and statistics, was **Wiener index**:\n",
    "> Wiener, H. (1947), \"Structural determination of paraffin boiling points\", Journal of the American Chemical Society, 1 (69): 17–20\n",
    "\n",
    "It is defined as the sum of lengths of shortest paths between molecule atoms:\n",
    "$$\n",
    "W(G) = \\frac{1}{2} \\sum_{i=1}^N \\sum_{j=1}^N dist(v_i, v_j)\n",
    "$$\n",
    "\n",
    "Length of shortest path is the minimal number of bonds to go through between atoms $i$ and $j$. Note we have to divide by 2 due to symmetry - here we sum both $dist(v_i, v_j)$ and $dist(v_j, v_i)$ for simplicity.\n",
    "\n",
    "It turns out to be well correlated with many physico-chemical properties of alkanes (paraffins), e.g. boiling point (original Wiener's work), melting point, liquid density and viscocity, and more.\n",
    "\n",
    "Other references:\n",
    "- [Wikipedia page](https://en.wikipedia.org/wiki/Wiener_index)\n",
    "- [Wolfram MathWorld page](https://mathworld.wolfram.com/WienerIndex.html)\n",
    "\n",
    "Unfortunately, it is not implemented in RDKit."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c0777d3-b9e7-41f9-b8ed-8f80fcadd2a4",
   "metadata": {},
   "source": [
    "### Exercise 4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68d7f8f0-8c15-486b-b788-485828732ca1",
   "metadata": {},
   "source": [
    "Finish the implementation of `wiener_index()` function, which computes the Wiener index value for a molecule.\n",
    "\n",
    "`GetDistanceMatrix()` function in RDKit returns a matrix of shortest path distances between atoms as a NumPy array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "306f236a-0a06-4d84-9ebb-077eb949f6d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from rdkit.Chem import GetDistanceMatrix\n",
    "\n",
    "\n",
    "def wiener_index(mol: Mol) -> int:\n",
    "    distances = GetDistanceMatrix(mol)\n",
    "    return int(distances.sum() / 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa2f656d-3293-44ce-88e3-c39c67ef64f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert np.isclose(wiener_index(ibuprofen_mol), 404)\n",
    "assert np.isclose(wiener_index(caffeine_mol), 258)\n",
    "\n",
    "print(f\"Ibuprofen Wiener index: {wiener_index(ibuprofen_mol)}\")\n",
    "print(f\"Caffeine Wiener index: {wiener_index(caffeine_mol)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d0f6271-f8c8-4ca8-b356-5ecff94c1a5f",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## 5. SMARTS patterns matching"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2f94872-c0db-4de7-85c5-3eebb43fc77c",
   "metadata": {},
   "source": [
    "SMARTS is like regular expressions (regex / regexp) for molecules. It allows us to define patterns in strings to detect substructures, e.g. methyl group, or aromatic ring of size N. They aren't very fast, but highly expressive and concise.\n",
    "\n",
    "SMARTS matches, called **substructures**, are one of the most commonly used ML features, e.g. in searching and property prediction.\n",
    "\n",
    "Example references:\n",
    "- [Daylight - SMARTS - A Language for Describing Molecular Patterns](https://www.daylight.com/dayhtml/doc/theory/theory.smarts.html)\n",
    "- [Daylight - SMARTS examples](https://www.daylight.com/dayhtml_tutorials/languages/smarts/smarts_examples.html)\n",
    "- [Basic SMARTS patterns by R. Hanson](https://jcheminf.biomedcentral.com/articles/10.1186/s13321-016-0160-4/tables/6)\n",
    "- [SMARTS.plus visualizer](https://smarts.plus/)\n",
    "\n",
    "In RDKit, they are `Mol` objects, created as `MolFromSmarts`. They typically are \"templates\" with wildcards, which can match any atoms. To match them, `Mol` has methods:\n",
    "- `GetSubstructMatch()` - gets single match as list of atom indexes ([docs](https://www.rdkit.org/docs/source/rdkit.Chem.rdchem.html#rdkit.Chem.rdchem.Mol.GetSubstructMatch))\n",
    "- `GetSubstructMatches()` - gets a list of matches aslists of atoms' indexes ([docs](https://www.rdkit.org/docs/source/rdkit.Chem.rdchem.html#rdkit.Chem.rdchem.Mol.GetSubstructMatches))\n",
    "- `HasSubstructMatch()` - boolean whether we have at least 1 match ([docs](https://www.rdkit.org/docs/source/rdkit.Chem.rdchem.html#rdkit.Chem.rdchem.Mol.HasSubstructMatch))\n",
    "\n",
    "Let's check 6-membered aromatic ring."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efd851cc-240d-4ea2-872a-999845f722b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit.Chem import MolFromSmarts\n",
    "\n",
    "\n",
    "ring_6_smarts = MolFromSmarts(\"a1aaaaa1\")\n",
    "ring_6_smarts"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83ccb8aa-25ea-4cb0-b6ec-c81659158085",
   "metadata": {},
   "source": [
    "To draw methyl group $CH_3$, it's useful to explicitly draw all hydrogens. `AddHs` also works on molecules created from SMARTS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3007e46c-d637-4bfe-99e9-2d74e27b8e25",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit.Chem import AddHs\n",
    "\n",
    "methyl_smarts = MolFromSmarts(\"[CH3]\")\n",
    "\n",
    "print(\"Methyl group, implicit hydrogens:\")\n",
    "display(MolToImage(methyl_smarts))\n",
    "\n",
    "print(\"Methyl group, explicit hydrogens:\")\n",
    "explicit_methyl_smarts = AddHs(methyl_smarts)\n",
    "display(MolToImage(explicit_methyl_smarts))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "373ddd5b-6ba7-4302-a38b-1cd329598a79",
   "metadata": {},
   "source": [
    "### Exercise 5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2565774-ad1b-41bb-a22f-1a7207b1eb87",
   "metadata": {},
   "source": [
    "Implement function `count_methyl_groups()`, which returns a number of methyl groups in a molecule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c81f8a1-5476-4af7-b329-41e84dc52fa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_methyl_groups(mol: Mol) -> int:\n",
    "    methyl_smarts = MolFromSmarts(\"[CH3]\")\n",
    "    return len(mol.GetSubstructMatches(methyl_smarts))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c602f172-fd9a-488b-ab29-e76aa051f812",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert count_methyl_groups(ibuprofen_mol) == 3\n",
    "assert count_methyl_groups(caffeine_mol) == 3\n",
    "\n",
    "print(f\"Ibuprofen has {count_methyl_groups(ibuprofen_mol)} methyl groups\")\n",
    "print(f\"Caffeine has {count_methyl_groups(caffeine_mol)} methyl groups\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6087fbb-d518-42c8-bd99-282a8c83e804",
   "metadata": {},
   "source": [
    "## Bonus exercise"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e95b2da5-1719-42f7-a167-148f4f8c4da5",
   "metadata": {},
   "source": [
    "Wiener index was originally used to predict melting point of alkanes (paraffins), which are crucial elements of fuels. However, he used a very small dataset. One of the largest and highest quality datasets for melting point of molecules is [Jean-Claude Bradley Double Plus Good Melting Point Dataset](https://figshare.com/articles/dataset/Jean_Claude_Bradley_Double_Plus_Good_Highly_Curated_and_Validated_Melting_Point_Dataset/1031638?file=1503991). It uses data from multiple data sources and measurements, not only for alkanes.\n",
    "\n",
    "1. Download the dataset, read the CSV with Pandas. Check the initial number of molecules.\n",
    "2. Get SMILES from `smiles` column and transform them to molecules.\n",
    "3. Get labels from `mpC` column - melting points in Celcius degrees.\n",
    "4. Transform SMILES to molecules, keeping only molecules and labels that can be successfully processed by RDKit. Note that `None` is returned in case of failure.\n",
    "5. Keep only alkanes, by using SMARTS `\"[CX4]\"`. If a molecule has such a match, it's an alkane. Note that you also need to filter labels. Check number of molecules after this step.\n",
    "6. Visualize `y` on a histogram. Print minimal, maximal and median values. This will help you understand the target values.\n",
    "7. Wiener also suggested *polarity number* in his paper, which is number of shortest paths with length exactly 3. Implement `polarity_number()` function, which computes this feature.\n",
    "8. For each molecule, extract features using functions from above:\n",
    "   - Wiener index\n",
    "   - polarity number\n",
    "   - number of aromatic rings\n",
    "   - molecular weight\n",
    "   - logP\n",
    "9. Split molecules with random split into train-test in 75-25% proportions with scikit-learn.\n",
    "10. Train a simple linear regression. Calculate and print MAE and $R^2$ on training and testing dataset.\n",
    "11. Many papers about predicting physico-chemical properties of molecules use polynomial regression. Add polynomial features with degree 3 (`PolynomialFeatures` in scikit-learn) and re-run the regression. Compare MAE and $R^2$ values to plain linear regression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35cee301-9c69-4016-ac1d-50685a0ad406",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install openpyxl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80590d7d-f8fe-404a-8123-619f07fd0348",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from skfp.preprocessing import MolFromSmilesTransformer\n",
    "\n",
    "\n",
    "df = pd.read_excel(\"data/BradleyDoublePlusGoodMeltingPointDataset.xlsx\")\n",
    "smiles_list = df[\"smiles\"]\n",
    "y = df[\"mpC\"]\n",
    "print(\"Initial number of molecules:\", len(df))\n",
    "\n",
    "mols = [MolFromSmiles(smiles) for smiles in smiles_list]\n",
    "valid_mask = [mol is not None for mol in mols]\n",
    "mols = np.array(mols)[valid_mask]\n",
    "y = np.array(y)[valid_mask]\n",
    "print(\"Number of RDKit-valid molecules:\", len(mols))\n",
    "\n",
    "alkane_smarts = MolFromSmarts(\"[CX4]\")\n",
    "alkane_mask = [mol.HasSubstructMatch(alkane_smarts) for mol in mols]\n",
    "mols = mols[alkane_mask]\n",
    "y = y[alkane_mask]\n",
    "print(\"Number of alkanes:\", len(mols))\n",
    "\n",
    "pd.Series(y).plot.hist(title=\"Melting points distribution\")\n",
    "plt.show()\n",
    "print(f\"Targets: min {y.min():.2f}, max {y.max():.2f}, median {np.median(y):.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00ac25b3-c826-4581-8e25-9e092eef0d5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def polarity_number(mol: Mol) -> int:\n",
    "    distances = GetDistanceMatrix(mol).astype(int)\n",
    "    return np.sum(distances == 3) // 2\n",
    "\n",
    "\n",
    "X = np.array(\n",
    "    [\n",
    "        [\n",
    "            wiener_index(mol),\n",
    "            polarity_number(mol),\n",
    "            get_num_aromatic_rings(mol),\n",
    "            MolWt(mol),\n",
    "            MolLogP(mol),\n",
    "        ]\n",
    "        for mol in mols\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c376cef-c6e0-46b1-a279-290b09be4681",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LinearRegression\n",
    "from sklearn.metrics import (\n",
    "    mean_absolute_error,\n",
    "    r2_score,\n",
    ")\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn.preprocessing import PolynomialFeatures\n",
    "\n",
    "\n",
    "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=0)\n",
    "\n",
    "lin_reg = LinearRegression()\n",
    "lin_reg.fit(X_train, y_train)\n",
    "y_pred = lin_reg.predict(X_test)\n",
    "print(f\"Linear regression MAE: {mean_absolute_error(y_test, y_pred):.2f}\")\n",
    "print(f\"Linear regression R^2: {r2_score(y_test, y_pred):.2f}\")\n",
    "\n",
    "print()\n",
    "\n",
    "poly_reg = make_pipeline(PolynomialFeatures(degree=3), LinearRegression())\n",
    "poly_reg.fit(X_train, y_train)\n",
    "y_pred = poly_reg.predict(X_test)\n",
    "print(f\"Polynomial regression MAE: {mean_absolute_error(y_test, y_pred):.2f}\")\n",
    "print(f\"Polynomial regression R^2: {r2_score(y_test, y_pred):.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "866a0dca-2e2d-4e58-8630-6b5eeca0c5aa",
   "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
}