count_geomol_failures.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "243efd0e",
   "metadata": {},
   "source": [
    "# Count how many input SMILES cannot be processed by GeoMol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "1aec1229",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.8.12 (default, Oct 13 2021, 06:42:42) \n",
      "[Clang 13.0.0 (clang-1300.0.29.3)]\n",
      "1.7.0\n",
      "1.6.3\n",
      "1.19.2\n",
      "0.7.0\n",
      "2021.09.4\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "print (sys.version)  # Should be 3.8, otherwise brew install python@3.8\n",
    "\n",
    "## Make sure to have the correct package versions :\n",
    "# Should be torch-geometric 1.6.3. and pytorch to 1.7.0, otherwise:\n",
    "# python3.8  -m pip install --upgrade --force-reinstall gast==0.3.3 grpcio~=1.32.0 typing-extensions~=3.7.4 h5py~=2.10.0 numpy~=1.19.2 six~=1.15.0 torch==1.7.0 torchvision==0.8.1 torchaudio==0.7.0  torch-scatter==2.0.7 torch-sparse==0.6.9 torch-cluster==1.5.9 torch-spline-conv==1.2.1 torch-geometric==1.6.3 -f https://data.pyg.org/whl/torch-1.7.0+cpu.html\n",
    "# (from https://github.com/pytorch/pytorch/issues/47354)\n",
    "\n",
    "# Install rdkit:\n",
    "# !python3.8 -m pip install rdkit-pypi\n",
    "# !python3.8 -m pip install --upgrade --force-reinstall numpy==1.19.2 pot==0.7.0\n",
    "\n",
    "\n",
    "import torch \n",
    "print(torch.__version__)\n",
    "import torch_geometric as tg\n",
    "print(tg.__version__)\n",
    "\n",
    "\n",
    "import numpy as np\n",
    "print(np.__version__)\n",
    "\n",
    "\n",
    "import ot\n",
    "print(ot.__version__)\n",
    "\n",
    "import rdkit\n",
    "print(rdkit.__version__)\n",
    "\n",
    "import os\n",
    "import json\n",
    "from rdkit import Chem, Geometry\n",
    "from rdkit.Chem import AllChem\n",
    "import pickle\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import random\n",
    "import yaml\n",
    "\n",
    "from model.model import GeoMol\n",
    "from model.featurization import featurize_mol_from_smiles\n",
    "from torch_geometric.data import Batch\n",
    "from model.inference import construct_conformers\n",
    "\n",
    "import random\n",
    "\n",
    "import pickle\n",
    "\n",
    "#!python3.8 -m pip install ipywidgets py3Dmol\n",
    "from ipywidgets import interact, fixed, IntSlider\n",
    "import ipywidgets\n",
    "import py3Dmol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7889c3a4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<script>\n",
       "code_show_err=false; \n",
       "function code_toggle_err() {\n",
       " if (code_show_err){\n",
       " $('div.output_stderr').hide();\n",
       " } else {\n",
       " $('div.output_stderr').show();\n",
       " }\n",
       " code_show_err = !code_show_err\n",
       "} \n",
       "$( document ).ready(code_toggle_err);\n",
       "</script>\n",
       "To toggle on/off output_stderr, click <a href=\"javascript:code_toggle_err()\">here</a>."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import HTML\n",
    "HTML('''<script>\n",
    "code_show_err=false; \n",
    "function code_toggle_err() {\n",
    " if (code_show_err){\n",
    " $('div.output_stderr').hide();\n",
    " } else {\n",
    " $('div.output_stderr').show();\n",
    " }\n",
    " code_show_err = !code_show_err\n",
    "} \n",
    "$( document ).ready(code_toggle_err);\n",
    "</script>\n",
    "To toggle on/off output_stderr, click <a href=\"javascript:code_toggle_err()\">here</a>.''')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "749ae87d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "f = open('summary_qm9.json')\n",
    "data = json.load(f)\n",
    "print(len(data.keys()))\n",
    "worked = 0\n",
    "processed = 0\n",
    "for smi in data.keys():\n",
    "    processed += 1\n",
    "    tg_data = featurize_mol_from_smiles(smi, dataset='qm9')\n",
    "    if not tg_data:\n",
    "#         print(f'failed to featurize SMILES: {smi}')\n",
    "        continue\n",
    "    else:\n",
    "        worked += 1\n",
    "        \n",
    "print(processed, worked)\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "0a57b56c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Same code as in visualize_confs.ipynb\n",
    "\n",
    "def show_mol(mol, view, grid):\n",
    "    mb = Chem.MolToMolBlock(mol)\n",
    "    view.removeAllModels(viewer=grid)\n",
    "    view.addModel(mb,'sdf', viewer=grid)\n",
    "    view.setStyle({'model':0},{'stick': {}}, viewer=grid)\n",
    "    view.zoomTo(viewer=grid)\n",
    "    return view\n",
    "\n",
    "def view_single(mol):\n",
    "    view = py3Dmol.view(width=600, height=600, linked=False, viewergrid=(1,1))\n",
    "    show_mol(mol, view, grid=(0, 0))\n",
    "    return view\n",
    "\n",
    "def MolTo3DView(mol, size=(600, 600), style=\"stick\", surface=False, opacity=0.5, confId=0):\n",
    "    \"\"\"Draw molecule in 3D\n",
    "    \n",
    "    Args:\n",
    "    ----\n",
    "        mol: rdMol, molecule to show\n",
    "        size: tuple(int, int), canvas size\n",
    "        style: str, type of drawing molecule\n",
    "               style can be 'line', 'stick', 'sphere', 'carton'\n",
    "        surface, bool, display SAS\n",
    "        opacity, float, opacity of surface, range 0.0-1.0\n",
    "    Return:\n",
    "    ----\n",
    "        viewer: py3Dmol.view, a class for constructing embedded 3Dmol.js views in ipython notebooks.\n",
    "    \"\"\"\n",
    "    assert style in ('line', 'stick', 'sphere', 'carton')\n",
    "    mblock = Chem.MolToMolBlock(mol[confId])\n",
    "    viewer = py3Dmol.view(width=size[0], height=size[1])\n",
    "    viewer.addModel(mblock, 'mol')\n",
    "    viewer.setStyle({style:{}})\n",
    "    if surface:\n",
    "        viewer.addSurface(py3Dmol.SAS, {'opacity': opacity})\n",
    "    viewer.zoomTo()\n",
    "    return viewer\n",
    "\n",
    "def conf_viewer(idx, mol):\n",
    "    return MolTo3DView(mol, confId=idx).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78eb8af3",
   "metadata": {},
   "source": [
    "## Code taken from the generate_confs.py. We ask GeoMol to generate one conformer for each input SMILES and check how often it fails.\n",
    "\n",
    "## NOTE: to reproduce numbers in our paper, one needs additionally to run scripts/clean_smiles.py and to sample two times the number of ground truth conformers per each mol. Otherwise, the code below would still work, just the final metric values would be worse. The file clean_smiles.py accounts for inconsistent molecules in the true dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "912e97a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def process_dataset(dataset='qm9'):\n",
    "    trained_model_dir = 'trained_models/' + dataset + '/'\n",
    "    mmff = False\n",
    "\n",
    "    with open(f'{trained_model_dir}/model_parameters.yml') as f:\n",
    "        model_parameters = yaml.full_load(f)\n",
    "    model = GeoMol(**model_parameters)\n",
    "\n",
    "    state_dict = torch.load(f'{trained_model_dir}/best_model.pt', map_location=torch.device('cpu'))\n",
    "    model.load_state_dict(state_dict, strict=True)\n",
    "    model.eval()\n",
    "\n",
    "    f = open('summary_' + dataset + '.json')\n",
    "    data = json.load(f)\n",
    "    print(len(data.keys()))\n",
    "    notworked = 0\n",
    "    processed = 0\n",
    "\n",
    "    all_dataset_smiles = list(data.keys())\n",
    "    \n",
    "    random.shuffle(all_dataset_smiles)\n",
    "    \n",
    "    conformer_dict = {}\n",
    "    for smi in all_dataset_smiles[:10000]:\n",
    "        processed += 1\n",
    "\n",
    "        # create data object (skip smiles rdkit can't handle)\n",
    "        tg_data = featurize_mol_from_smiles(smi, dataset=dataset)\n",
    "        if not tg_data:\n",
    "            notworked += 1\n",
    "            continue\n",
    "\n",
    "        # generate model predictions\n",
    "        data = Batch.from_data_list([tg_data])\n",
    "        model(data, inference=True, n_model_confs=1)\n",
    "\n",
    "        # set coords\n",
    "        n_atoms = tg_data.x.size(0)\n",
    "        model_coords = construct_conformers(data, model)\n",
    "        mols = []\n",
    "        for x in model_coords.split(1, dim=1):\n",
    "            mol = Chem.AddHs(Chem.MolFromSmiles(smi))\n",
    "            coords = x.squeeze(1).double().cpu().detach().numpy()\n",
    "            mol.AddConformer(Chem.Conformer(n_atoms), assignId=True)\n",
    "            for i in range(n_atoms):\n",
    "                mol.GetConformer(0).SetAtomPosition(i, Geometry.Point3D(coords[i, 0], coords[i, 1], coords[i, 2]))\n",
    "\n",
    "            if mmff:\n",
    "                try:\n",
    "                    AllChem.MMFFOptimizeMoleculeConfs(mol, mmffVariant='MMFF94s')\n",
    "                except Exception as e:\n",
    "                    pass\n",
    "            mols.append(mol)\n",
    "\n",
    "        if len(mols) == 0:\n",
    "            notworked += 1\n",
    "            continue\n",
    "\n",
    "        conformer_dict[smi] = mols\n",
    "\n",
    "\n",
    "#         if processed % 10000 == 0:\n",
    "#             print(notworked, processed)\n",
    "    return conformer_dict\n",
    "\n",
    "conformer_dict = process_dataset('qm9')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "217e3d59",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using 10000 random SMILES from QM9, we get   num SMILES with successfully generated conformers  9938\n"
     ]
    }
   ],
   "source": [
    "print('Using 10000 random SMILES from QM9, we get ',\n",
    "      ' num SMILES with successfully generated conformers ', len(conformer_dict.keys()))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "cf9faf74",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "304466\n",
      "Using 10000 random SMILES from DRUGS, we get   num SMILES with successfully generated conformers  9999\n"
     ]
    }
   ],
   "source": [
    "conformer_dict = process_dataset('drugs')\n",
    "print('Using 10000 random SMILES from DRUGS, we get ',\n",
    "      ' num SMILES with successfully generated conformers ', len(conformer_dict.keys()))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "63f3aad8",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a671644138f14d909234ed7fb83c583e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=0, description='idx', max=0), Output()), _dom_classes=('widget-interact'…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8ec488b2c9ee4693bd747a44ed190e35",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=0, description='idx', max=0), Output()), _dom_classes=('widget-interact'…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a895775a991a4c538c4e7d1f94744756",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=0, description='idx', max=0), Output()), _dom_classes=('widget-interact'…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Show the conformer of one random SMILES\n",
    "mols = conformer_dict[list(conformer_dict.keys())[0]]\n",
    "interact(conf_viewer, idx=ipywidgets.IntSlider(min=0, max=len(mols)-1, step=1), mol=fixed(mols));\n",
    "\n",
    "mols = conformer_dict[list(conformer_dict.keys())[1]]\n",
    "interact(conf_viewer, idx=ipywidgets.IntSlider(min=0, max=len(mols)-1, step=1), mol=fixed(mols));\n",
    "\n",
    "mols = conformer_dict[list(conformer_dict.keys())[2]]\n",
    "interact(conf_viewer, idx=ipywidgets.IntSlider(min=0, max=len(mols)-1, step=1), mol=fixed(mols));\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e262a80e",
   "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}