Notebooks/overview_of_DESMILES.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "use_gpu = True\n",
    "if not use_gpu:\n",
    "    import os\n",
    "    os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"-1\"\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "import sys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "import desmiles\n",
    "from desmiles.data import Vocab, FpSmilesList, DesmilesLoader, DataBunch\n",
    "from desmiles.learner import desmiles_model_learner\n",
    "from desmiles.models import Desmiles, RecurrentDESMILES, get_fp_to_embedding_model, get_embedded_fp_to_smiles_model\n",
    "from desmiles.utils import load_old_pretrained_desmiles, load_pretrained_desmiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import torch\n",
    "import pandas as pd\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from desmiles.config import DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "model_fn = Path(DATA_DIR) / 'pretrained/model_2000_400_2000_5.h5'\n",
    "architecture = {'fp_emb_sz': 2000, 'emb_sz': 400, 'nh': 2000, 'nl': 5, 'clip':0.3, 'alpha':2., 'beta':1.}\n",
    "# load fastai learner\n",
    "learner = load_old_pretrained_desmiles(model_fn, return_learner=True, **architecture)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "from desmiles.learner import OriginalFastaiOneCycleScheduler, Learner\n",
    "\n",
    "# generate training data\n",
    "n=1000\n",
    "sigma=0.1\n",
    "# learn function y = x**2 + noise\n",
    "x = np.linspace(-1,1, n)\n",
    "y = x**2 + (np.random.randn(n) * sigma)\n",
    "x_t = torch.tensor(x, dtype=torch.float).unsqueeze(1)\n",
    "y_t = torch.tensor(y, dtype=torch.float).unsqueeze(1)\n",
    "\n",
    "# create databunch\n",
    "trn_ds = torch.utils.data.TensorDataset(x_t, y_t)\n",
    "val_ds = torch.utils.data.TensorDataset(x_t, y_t)\n",
    "trn_loader = torch.utils.data.DataLoader(trn_ds, batch_size=10, shuffle=True)\n",
    "val_loader = torch.utils.data.DataLoader(val_ds, batch_size=10, shuffle=False)\n",
    "db = DataBunch(trn_loader, val_loader)\n",
    "\n",
    "# train model \n",
    "model = torch.nn.Sequential(torch.nn.Linear(1,100), torch.nn.ReLU(), torch.nn.Linear(100,1), torch.nn.ReLU())\n",
    "learner = Learner(db, model, loss_func=torch.nn.functional.mse_loss)\n",
    "div_factor=10\n",
    "# Use the old fastai one cycle training policy\n",
    "one_cycle_linear_cb = OriginalFastaiOneCycleScheduler(learner, 0.002, div_factor=div_factor)\n",
    "learner.fit(5, callbacks=[one_cycle_linear_cb])\n",
    "learner.recorder.plot_lr()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "# The current fastai library uses the following for their one cycle training policy\n",
    "learner = Learner(db, model, loss_func=torch.nn.functional.mse_loss)\n",
    "learner.fit_one_cycle(5, 0.002)\n",
    "learner.recorder.plot_lr()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "import scipy.sparse\n",
    "import os\n",
    "MYDATA=os.path.join(DATA_DIR, 'notebooks')\n",
    "    \n",
    "trn_smiles = np.load(os.path.join(MYDATA, 'training.enc8000.split_0.npy'))\n",
    "trn_fps = scipy.sparse.load_npz(os.path.join(MYDATA, 'training_fp.split_0.npz'))\n",
    "\n",
    "val_smiles = np.load(os.path.join(MYDATA, 'validation.enc8000.npy'))\n",
    "val_fps = scipy.sparse.load_npz(os.path.join(MYDATA,'validation_fp.npz'))\n",
    "\n",
    "itos_fn=os.path.join(DATA_DIR, 'pretrained', 'id.dec8000')\n",
    "itos = [s.strip() for i,s in enumerate(open(itos_fn, encoding='utf-8'))]\n",
    "vocab = Vocab(itos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "#Let's train DESMILSE on 1% of 1/4 the data\n",
    "\n",
    "num_trn_smiles = trn_smiles.shape[0] \n",
    "trn_inds = np.random.permutation(np.arange(num_trn_smiles))\n",
    "num_to_keep = int(num_trn_smiles*0.01)\n",
    "trn_inds = trn_inds[:num_to_keep]\n",
    "\n",
    "num_val_smiles = val_smiles.shape[0] \n",
    "val_inds = np.random.permutation(np.arange(num_val_smiles))\n",
    "num_to_keep = int(num_val_smiles*0.01)\n",
    "val_inds = val_inds[:num_to_keep]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "# create data bunch\n",
    "bs=200\n",
    "trn_ds = FpSmilesList(trn_smiles[trn_inds], trn_fps[trn_inds], vocab)\n",
    "val_ds = FpSmilesList(val_smiles[val_inds], val_fps[val_inds], vocab)\n",
    "trn_dl = DesmilesLoader(trn_ds, bs=bs, vocab=vocab)\n",
    "val_dl = DesmilesLoader(val_ds, bs=bs, vocab=vocab)\n",
    "db = DataBunch(trn_dl, val_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "scrolled": true,
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "from desmiles.utils import accuracy4\n",
    "architecture = {'fp_emb_sz': 200, 'emb_sz': 200, 'nh': 200, 'nl': 1}\n",
    "regularization = {'clip':0.3, 'alpha':2., 'beta':1.}\n",
    "\n",
    "# Training parameters\n",
    "max_lr = 0.001\n",
    "div_factor = 10.\n",
    "\n",
    "# 1) Create learner object\n",
    "learner = desmiles_model_learner(db, **architecture, **regularization)\n",
    "\n",
    "learner.metrics = [accuracy4]\n",
    "# 2) Specify training schedule\n",
    "one_cycle_linear_cb = OriginalFastaiOneCycleScheduler(learner, max_lr, div_factor=div_factor)\n",
    "# 3) Train \n",
    "learner.fit(50, callbacks=[one_cycle_linear_cb])\n",
    "# 4) Save model\n",
    "#learner.save('model_1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "learner.recorder.plot_lr()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "from desmiles.utils import decoder, image_of_mols\n",
    "from functools import partial\n",
    "\n",
    "# helper function to map from one-hot-encoded vector to smiles string\n",
    "decoder = partial(decoder, itos=itos)\n",
    "def smiles_idx_to_string(smiles_idx, decoder=decoder):\n",
    "    return decoder(smiles_idx[smiles_idx > 0].tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "from desmiles.decoding.astar import AstarTreeParallelHybrid as AstarTree\n",
    "\n",
    "#learner.load('model_1')\n",
    "\n",
    "# Lets see if we at least fit our training set a bit\n",
    "(smiles_idx, fps, lengths), y = next(iter(trn_dl))\n",
    "test_smiles_idx = smiles_idx[:,-1]\n",
    "test_fp = fps[-1]\n",
    "test_smiles = smiles_idx_to_string(test_smiles_idx)\n",
    "image_of_mols([test_smiles])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "scrolled": true,
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "model = learner.model\n",
    "model.eval()\n",
    "rdesmiles = RecurrentDESMILES(model)\n",
    "\n",
    "astar = AstarTree(test_fp.unsqueeze(0).to('cuda'), rdesmiles, num_expand=100)\n",
    "neg_log_prob, smiles_idx =  next(astar)\n",
    "smiles = smiles_idx_to_string(smiles_idx)\n",
    "image_of_mols([smiles, test_smiles])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from desmiles.decoding.astar import AstarTreeParallelHybrid as AstarTree\n",
    "\n",
    "#learner.load('model_1')\n",
    "\n",
    "# Lets see if we at least fit our training set a bit\n",
    "(smiles_idx, fps, lengths), y = next(iter(val_dl))\n",
    "test_smiles_idx = smiles_idx[:,-1]\n",
    "test_fp = fps[-1]\n",
    "test_smiles = smiles_idx_to_string(test_smiles_idx)\n",
    "image_of_mols([test_smiles])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model = learner.model\n",
    "model.eval()\n",
    "rdesmiles = RecurrentDESMILES(model)\n",
    "\n",
    "astar = AstarTree(test_fp.unsqueeze(0).to('cuda'), rdesmiles, num_expand=100)\n",
    "neg_log_prob, smiles_idx =  next(astar)\n",
    "smiles = smiles_idx_to_string(smiles_idx)\n",
    "image_of_mols([smiles, test_smiles])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "model_fn = Path(DATA_DIR) / 'pretrained/model_2000_400_2000_5.h5'\n",
    "architecture = {'fp_emb_sz': 2000, 'emb_sz': 400, 'nh': 2000, 'nl': 5, 'clip':0.3, 'alpha':2., 'beta':1.}\n",
    "learner = load_old_pretrained_desmiles(model_fn, return_learner=True, **architecture)\n",
    "model = learner.model\n",
    "model.eval()\n",
    "# make a RecurrentDESMILES model\n",
    "model = RecurrentDESMILES(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "from desmiles.utils import smiles_to_fingerprint\n",
    "validation_smiles = [s.strip() for s in open(os.path.join(DATA_DIR, 'pretrained', 'validation_smiles_10k.smi'))]\n",
    "inds = np.random.permutation(np.arange(len(validation_smiles)))\n",
    "i = 0\n",
    "smiles_to_invert = validation_smiles[inds[i]]\n",
    "fp = smiles_to_fingerprint(smiles_to_invert, as_tensor=True)\n",
    "image_of_mols([smiles_to_invert])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "astar = AstarTree(fp.unsqueeze(0), model, num_expand=100)\n",
    "nlp, smiles_idx = next(astar)\n",
    "smiles = smiles_idx_to_string(smiles_idx)\n",
    "image_of_mols([smiles, smiles_to_invert])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "# model size\n",
    "np.sum([np.prod(p.shape) for p in model.parameters()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "# Lets use the fast variant of A* to get 100 top solution\n",
    "astar = AstarTree(fp.unsqueeze(0), model, num_expand=1000, max_branches=5000)\n",
    "from collections import defaultdict\n",
    "scores = defaultdict(float)\n",
    "all_leaf_nodes = []\n",
    "for _ in range(1000):\n",
    "    nlp, smiles_idx = next(astar)\n",
    "    smiles = smiles_idx_to_string(smiles_idx)\n",
    "    print(smiles,  np.exp(-nlp))\n",
    "    scores[smiles] += np.exp(-nlp)\n",
    "    all_leaf_nodes.append(smiles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "sorted(scores.items(), key=lambda x: -x[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Don't forget to regenerate the DRD2 dataset if you haven't already done it.\n",
    "To do so, please run DESMILES/tests/download_drd2_dataset.sh <PATH/TO>/DESMILES/data/notebooks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "from drd2 import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "bs=200\n",
    "original_smile, train_fp, train_enc = load_training_data(raise_prob=True)\n",
    "db = create_databunch(train_fp, train_enc, itos_fn, bs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "model_fn = Path(os.path.join(DATA_DIR, 'pretrained', 'model_2000_400_2000_5.h5'))\n",
    "learner = load_old_pretrained_desmiles(model_fn, return_learner=True)\n",
    "learner.metrics = [accuracy4]\n",
    "learner.data = db\n",
    "\n",
    "num_epochs = 5\n",
    "max_lr = 0.001\n",
    "div_factor = 7\n",
    "\n",
    "one_cycle_linear_cb = OriginalFastaiOneCycleScheduler(learner, max_lr, div_factor=div_factor)\n",
    "learner.fit(num_epochs, callbacks=[one_cycle_linear_cb])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "(val_smiles_idx, val_fps, _), _ = next(iter(db.valid_dl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "model = learner.model\n",
    "model.eval()\n",
    "model = RecurrentDESMILES(model)\n",
    "astar = AstarTree(val_fps[0].unsqueeze(0), model, num_expand=100)\n",
    "all_leaf_nodes = []\n",
    "for _ in range(100):\n",
    "    nlp, smiles_idx = next(astar)\n",
    "    smiles = smiles_idx_to_string(smiles_idx)\n",
    "    print(smiles)\n",
    "    all_leaf_nodes.append(smiles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.7.12"
  },
  "name": "DESMILES.ipynb"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}