Notebooks/intro_demo_of_DESMILES.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*start this notebook server on a machine with a GPU; optionally, use a modern rdkit version* "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image\n",
    "import os\n",
    "from desmiles.config import DATA_DIR\n",
    "\n",
    "fig_dir = os.path.join(DATA_DIR, 'notebooks', 'Figures')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The high level idea behind DESMILES is that if we can learn to generate small molecules from a reduced small molecule representation that has been very successful in modelling structure activity relationships, then we will be able to generate useful molecules for a variety of practical tasks in drug discovery.  Furthermore, if this representation learning is encoding the chemical similarity, then we will be able to easily generate chemically similar molecules starting from any molecule.  The following two images from the DESMILES publication show the outline of the model and the representation of a slice of chemical space.  The notebook below demonstrates the basic functionality of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = Image(filename=f\"{fig_dir}/deep learn chem space (desmiles)__extended data fig 5__3__mcgillen__2019__.png\")\n",
    "display(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "img = Image(filename=f\"{fig_dir}/deep learn chem space (desmiles)__extended data fig 1__1__maragakis__2019__.png\")\n",
    "display(img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demo of DESMILES"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next couple of cell define some high level code for generating and displaying new molecules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "import argparse\n",
    "import tempfile\n",
    "import multiprocessing\n",
    "import functools\n",
    "import subprocess\n",
    "from pathlib import Path\n",
    "\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import scipy\n",
    "\n",
    "\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem.AllChem import GetMorganFingerprintAsBitVect\n",
    "\n",
    "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\n",
    "from desmiles.models import get_fp_to_embedding_model, get_embedded_fp_to_smiles_model\n",
    "from desmiles.utils import load_old_pretrained_desmiles, load_pretrained_desmiles\n",
    "from desmiles.utils import accuracy4\n",
    "from desmiles.utils import smiles_idx_to_string\n",
    "from desmiles.learner import OriginalFastaiOneCycleScheduler, Learner\n",
    "from desmiles.decoding.astar import AstarTreeParallelHybrid as AstarTree\n",
    "from desmiles.config import DATA_DIR\n",
    "\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "from IPython.display import SVG, display, HTML\n",
    "import rdkit.Chem\n",
    "from rdkit.Chem import Draw,rdMolDescriptors,AllChem, rdDepictor\n",
    "from rdkit.Chem.Draw import IPythonConsole, rdMolDraw2D\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "table_from_to_template='<table border=1> <tr> <td>{}</td> <td> --> </td> <td>{}</td> </tr> </table>'\n",
    "table_4_template='<table border=1 width={}> <tr> <td>{}</td> <td>{}</td> <td>{}</td> <td>{}</td> </tr> </table>'\n",
    "\n",
    "\n",
    "def canon_smiles(x):\n",
    "    s = ''\n",
    "    try: \n",
    "        s = Chem.CanonSmiles(x, useChiral=True)\n",
    "    except:\n",
    "        pass\n",
    "    return s\n",
    "\n",
    "\n",
    "def get_itos_8k():\n",
    "    return np.load(os.path.join(DATA_DIR, 'pretrained', \"itos.npy\"))\n",
    "\n",
    "\n",
    "def vec_to_smiles(idx_vec_inp, itos):\n",
    "    \"\"\"Return a SMILES string from an index vector (deals with reversal)\"\"\"\n",
    "    ##HACK TO WORK WITH NEWER VERSION 2020-06-08\n",
    "    if idx_vec_inp[0] == 3:\n",
    "        idx_vec = idx_vec_inp[1:]\n",
    "    else:\n",
    "        idx_vec = idx_vec_inp\n",
    "    ##\n",
    "    if idx_vec[0] == 1:  # SMILES string is in fwd direction\n",
    "        return ''.join(itos[x] for x in idx_vec if x > 3)\n",
    "    if idx_vec[0] == 2:  # SMILES string is in bwd direction\n",
    "        #despot.Print(\"decoder: bwd direction\")\n",
    "        return ''.join(itos[x] for x in idx_vec[::-1] if x > 3)\n",
    "    else: # don't know how to deal with it---do your best\n",
    "        print(\"decoder received an invalid start to the SMILES\", idx_vec)\n",
    "        return ''.join(itos[x] for x in idx_vec if x > 3)\n",
    "\n",
    "    \n",
    "def smiles_to_fingerprint(smiles_str, sparse=False, as_tensor=True):\n",
    "    \"Return the desmiles fp\"\n",
    "    rdmol = Chem.MolFromSmiles(smiles_str)\n",
    "    fp = np.concatenate([\n",
    "        np.asarray(GetMorganFingerprintAsBitVect(rdmol, 2, useChirality=True), dtype=np.uint8),\n",
    "        np.asarray(GetMorganFingerprintAsBitVect(rdmol, 3, useChirality=True), dtype=np.uint8)])\n",
    "    if sparse:\n",
    "        return scipy.sparse.csr_matrix(fp)\n",
    "    if as_tensor:\n",
    "        import torch\n",
    "        device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "        return torch.tensor(fp.astype(np.float32)).to(device)\n",
    "    return fp\n",
    "\n",
    "\n",
    "def barcode_fp(fp, width=8, height=0.5):\n",
    "    fig = plt.figure()\n",
    "    ax2 = fig.add_axes([0, 0, width, height], xticks=[], yticks=[])\n",
    "    barprops = dict(aspect='auto', cmap=plt.cm.binary, interpolation='nearest')\n",
    "    return ax2.imshow(fp.reshape((1, -1)), **barprops)\n",
    "    \n",
    "    \n",
    "def moltosvg(rdkit_mol, size_x=450, size_y=150):\n",
    "    try:\n",
    "        rdkit_mol.GetAtomWithIdx(0).GetExplicitValence()\n",
    "    except RuntimeError:\n",
    "        rdkit_mol.UpdatePropertyCache(False)\n",
    "    try:\n",
    "        mc_mol = rdMolDraw2D.PrepareMolForDrawing(rdkit_mol, kekulize=True)\n",
    "    except ValueError:  # <- can happen on a kekulization failure                                                                                                                             \n",
    "        mc_mol = rdMolDraw2D.PrepareMolForDrawing(rdkit_mol, kekulize=False)\n",
    "    drawer = rdMolDraw2D.MolDraw2DSVG(size_x, size_y)\n",
    "    drawer.DrawMolecule(mc_mol)\n",
    "    drawer.FinishDrawing()\n",
    "    svg = drawer.GetDrawingText()\n",
    "    # It seems that the svg renderer used doesn't quite hit the spec.\n",
    "    # Here are some fixes to make it work in the notebook, although I think\n",
    "    # the underlying issue needs to be resolved at the generation step\n",
    "    return svg.replace('svg:','')\n",
    "\n",
    "\n",
    "def displayTable4SMILES(smiles, size_x=225, size_y=150, width=980):\n",
    "    assert(len(smiles)==4)\n",
    "    svgs = map(lambda x: moltosvg(Chem.MolFromSmiles(x), size_x=size_x, size_y=size_y), smiles)\n",
    "    display(HTML(table_4_template.format(width, *svgs)))\n",
    "    \n",
    "    \n",
    "def procSMILES(sm):\n",
    "    m = Chem.MolFromSmiles(sm)\n",
    "    AllChem.Compute2DCoords(m)\n",
    "    return m\n",
    "\n",
    "\n",
    "def imageOfMols(smiles_list, molsPerRow=4, subImgSize=(240,200), labels=None):\n",
    "    mols = [procSMILES(sm) for sm in smiles_list]\n",
    "    if labels is not None:\n",
    "        labels = [str(x) for x in labels]\n",
    "    img = Draw.MolsToGridImage(mols, molsPerRow=molsPerRow, subImgSize=subImgSize, useSVG=True, legends=labels)\n",
    "    return img\n",
    "\n",
    "\n",
    "def imageOfMolsLabels(smiles_labels_list, molsPerRow=5, subImgSize=(200,200)):\n",
    "    mols = [procSMILES(sm[0]) for sm in smiles_labels_list]\n",
    "    labels = [str(sm[1]) for sm in smiles_labels_list]\n",
    "    img = Draw.MolsToGridImage(mols, molsPerRow=molsPerRow, subImgSize=subImgSize, useSVG=True, legends=labels)\n",
    "    return img\n",
    "\n",
    "\n",
    "from itertools import zip_longest\n",
    "def grouper(iterable, n, fillvalue=None):\n",
    "    args = [iter(iterable)] * n\n",
    "    return zip_longest(*args, fillvalue=fillvalue)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#from astar_purge import get_astar_tree, isLeafNode\n",
    "def get_most_probable_smiles(emb_fp, emfp_to_smiles, max_branches=50, num_expand=2000):\n",
    "    with torch.no_grad():\n",
    "        astar_tree = AstarTree(emb_fp, \n",
    "                               emfp_to_smiles, \n",
    "                               max_branches=max_branches, num_expand=num_expand)\n",
    "        score, smile_idx = next(astar_tree)\n",
    "        return smiles_idx_to_string(smile_idx)\n",
    "\n",
    "\n",
    "\n",
    "def get_upto_n_most_probable_valid_smiles(n, emb_fp, emfp_to_smiles, max_branches=50, num_expand=2000):\n",
    "    with torch.no_grad():\n",
    "        astar_tree = AstarTree(emb_fp, #.astype(np.float32), \n",
    "                                    emfp_to_smiles, \n",
    "                                    max_branches=max_branches, num_expand=num_expand)\n",
    "        for _ in range(n):\n",
    "            smile = smiles_idx_to_string(next(astar_tree)[1])\n",
    "            if Chem.MolFromSmiles(smile):\n",
    "                yield smile\n",
    "            else:\n",
    "                yield \"\"\n",
    "\n",
    "                \n",
    "def get_first_n_most_probable_valid_smiles(n, emb_fp, emfp_to_smiles, \n",
    "                                           max_branches=50, max_search=200,\n",
    "                                           num_expand=2000,\n",
    "                                           verbose=False):\n",
    "    results = set()\n",
    "    with torch.no_grad():\n",
    "        astar_tree = AstarTree(emb_fp, \n",
    "                               emfp_to_smiles, \n",
    "                               max_branches=max_branches, num_expand=num_expand)\n",
    "        for i in range(max_search):\n",
    "            if len(results) == n:\n",
    "                break\n",
    "            score, smile = next(astar_tree)\n",
    "            smile = smiles_idx_to_string(smile)\n",
    "            if verbose:\n",
    "                print(score, smile)\n",
    "            if (smile is not None) and Chem.MolFromSmiles(smile):\n",
    "                smile = Chem.CanonSmiles(smile)\n",
    "                results.add(smile)\n",
    "                yield smile\n",
    "            else:\n",
    "                yield \"\"\n",
    "\n",
    "                \n",
    "def set_of_upto_n_most_probable_valid_smiles(n, emb_fp, embfp_to_smiles, max_branches=50, num_expand=2000):\n",
    "    result = list(get_upto_n_most_probable_valid_smiles(n, emb_fp, embfp_to_smiles, \n",
    "                                                        max_branches, num_expand))\n",
    "    return set([Chem.CanonSmiles(x) for x in result if x != \"\"])\n",
    "\n",
    "    \n",
    "def get_most_probable_valid_smiles(emb_fp, emfp_to_smiles, max_branches=50, num_expand=2000):\n",
    "    counter = 0\n",
    "    with torch.no_grad():\n",
    "        astar_tree = AstarTree(emb_fp, emfp_to_smiles,\n",
    "                                   max_branches=max_branches,\n",
    "                              num_expand=num_expand)\n",
    "        while counter < 5:\n",
    "            smile = smiles_idx_to_string(next(astar_tree)[1])\n",
    "            if Chem.MolFromSmiles(smile):\n",
    "                return smile\n",
    "            counter += 1\n",
    "        return smile\n",
    "    \n",
    "    \n",
    "def dedup(seq):\n",
    "    seen = set()\n",
    "    seen_add = seen.add\n",
    "    return [x for x in seq if not (x in seen or seen_add(x))]\n",
    "\n",
    "\n",
    "def clean(seq):\n",
    "    return [x for x in seq if x]\n",
    "\n",
    "\n",
    "def dedup_clean(seq):\n",
    "    return dedup(clean(seq))\n",
    "\n",
    "\n",
    "def to_numpy_int(f1):\n",
    "    return f1.cpu().numpy().astype(np.int32) if isinstance(f1, torch.Tensor) else f1\n",
    "\n",
    "\n",
    "def tanimoto(f1, f2):\n",
    "    s1 = to_numpy_int(f1)\n",
    "    s2 = to_numpy_int(f2)\n",
    "    return np.sum(s1 & s2) / np.sum(s1 | s2)\n",
    "\n",
    "\n",
    "def displayTable4labels(labels, width=980):\n",
    "    display(HTML(table_4_template.format(width, *labels)))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <div id=\"load_data\">Load some data </div>\n",
    "\n",
    "First load the encoded version of the training set and the validation set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "itos = get_itos_8k()\n",
    "vec_to_smiles_8k = partial(vec_to_smiles, itos=itos)\n",
    "\n",
    "def random_smiles_enc8k(enc_table, n=2):\n",
    "    idx = np.random.randint(0, len(enc_table), n)\n",
    "    return [(vec_to_smiles_8k(enc_table[i]), enc_table[i][enc_table[i]>0]) for i in idx]\n",
    "\n",
    "def random_smiles(enc_table, n=10):\n",
    "    return [s for s,e in random_smiles_enc8k(enc_table, n)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "training_smiles_enc8k = np.load(os.path.join(DATA_DIR, 'pretrained', 'training.enc8000.npy'))\n",
    "val2_smiles_enc8k = np.load(os.path.join(DATA_DIR, 'pretrained', 'val2.enc8000.npy'))\n",
    "v2samples = list(pd.read_csv(os.path.join(DATA_DIR, 'notebooks', \"fast_val2_molecules.csv\"))[\"SMILES\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print('\\n'.join(random_smiles(training_smiles_enc8k, 5)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "len(training_smiles_enc8k)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hit Ctrl-Enter on the following cell several times to explore random samples from the training set.  \n",
    "Change \"training_smiles_enc8k\" to \"val2_smiles_enc8k\" to explore the validation set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import rdkit.Chem.Descriptors\n",
    "smiles = random_smiles(training_smiles_enc8k, 4) # pick training_smiles... or val2_smiles...\n",
    "displayTable4SMILES(smiles)\n",
    "for x in smiles: \n",
    "    m = Chem.MolFromSmiles(x)\n",
    "    print( m.GetNumAtoms(), m.GetNumAtoms(onlyExplicit=False), \n",
    "          Chem.Descriptors.NumAromaticRings(m), \n",
    "          np.round(Chem.Descriptors.TPSA(m), 2),\n",
    "          np.round(Chem.Descriptors.MolWt(m), 2), \n",
    "          np.round(Chem.Descriptors.MolLogP(m), 2) )\n",
    "for x in smiles: barcode_fp(smiles_to_fingerprint(x), height=0.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <div id=\"subSMILES\"> subSMILES and DESMILES </div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each molecule is made out of up to 26 subSMILES (byte-pair encoded symbols).\n",
    "These are represented as integer numbers in the encoded tables, and 'itos' converts them to strings.\n",
    "The following example decomposes a random molecule from the training set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_smiles_enc8k.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "rsamples = random_smiles_enc8k(training_smiles_enc8k, 1)\n",
    "print(rsamples)\n",
    "display(imageOfMolsLabels(rsamples, subImgSize=(600,300), molsPerRow=1))\n",
    "barcode_fp(smiles_to_fingerprint(rsamples[0][0]));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rsamples[0][0], [(x, itos[x]) for s, e in rsamples for x in e]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load a pretrained model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load a model that was trained on molecules from both the validation and the training set. This would be the model to use in future applications of DESMILES.  The model parameters came as the result of the hyperoptimization discussed in the publication."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "get_model = desmiles.utils.load_old_pretrained_desmiles\n",
    "fp_to_smiles_5layer = get_model(os.path.join(DATA_DIR, 'pretrained', 'train_val1_val2','model_2000_400_2000_5'))\n",
    "rmodel = RecurrentDESMILES(fp_to_smiles_5layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rmodel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <div id=\"first_tests\">First tests: create a fragment from its fingerprint</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The simplest application of the model is to create a molecule from its fingerprint.\n",
    "Here is the example of a fragment that is outside of the original library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "scrolled": true,
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "fragment = 'Nc1ncc(C(F)(F)F)cc1F'  \n",
    "display(imageOfMols([fragment], labels=[fragment]))\n",
    "fp = smiles_to_fingerprint(fragment,  as_tensor=True)\n",
    "barcode_fp(fp);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fp, fp.size(), set(fp.cpu().numpy()), sum(fp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can invert this fingerprint to generate a small molecule.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.hycell": false,
    "ein.tags": "worksheet-0",
    "scrolled": true,
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "smiles = get_most_probable_smiles(fp, rmodel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "imageOfMols([smiles])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <div id=\"collection\"> Generate a collection of fragments with Astar </div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Often one wants to generate a whole bunch of variations of a single molecule.\n",
    "\n",
    "The example below shows variations of this simple fragment outside the training/validation set, \n",
    "together with a measure of the fingerprint similarity (higher is better; 1.0 is perfect match of fingerprints.)\n",
    "\n",
    "Sometimes the model will go through a number of invalid intermediate attempts, before finding the next example.\n",
    "The function get_n_most_probable_valid_smiles will only return the valid molecules.\n",
    "The parameter max_branches limits the search; for very complicated molecules, \n",
    "the search might be exhausted before the optimal molecules get returned, \n",
    "so higher values of max_branches might get \"better\" molecules \n",
    "but the search will take up more GPU memory.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "%%capture --no-stdout --no-display\n",
    "smiles = dedup_clean([Chem.CanonSmiles(x) \n",
    "                       for x in get_first_n_most_probable_valid_smiles(8, fp, rmodel, \n",
    "                                                                      max_branches=100)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = np.round([tanimoto(fp, smiles_to_fingerprint(x)) for x in smiles], 2)\n",
    "for x in smiles: barcode_fp(smiles_to_fingerprint(x), height=0.2)\n",
    "display(imageOfMols(smiles, molsPerRow=4, subImgSize=(240, 200), labels=labels));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check out some more complex and random molecules (subset of validation 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below are some complicated molecules and their top 3 variants according to DEMSMILES.  For this demonstration we picked as inputs a subset of the validation molecules that decoded rather quickly.  \n",
    "\n",
    "Please click Ctrl-Enter on the next cell a couple of times until you see some interesting molecules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "rsamples = list(np.random.choice(v2samples, 4))\n",
    "print(rsamples)\n",
    "displayTable4SMILES(rsamples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "%%capture --no-stdout --no-display\n",
    "for s in rsamples:\n",
    "    newfp = smiles_to_fingerprint(s)\n",
    "    newsmiles = dedup_clean([Chem.CanonSmiles(x) for x in \n",
    "                         get_first_n_most_probable_valid_smiles(4,newfp,rmodel, max_branches=200)])\n",
    "    newsmiles.extend([\"\", \"\", \"\", \"\"])\n",
    "    newsmiles = newsmiles[:4]\n",
    "    displayTable4SMILES(newsmiles)\n",
    "    labels = [np.round(tanimoto(newfp, smiles_to_fingerprint(s)), 2) for s in newsmiles]\n",
    "    displayTable4labels(labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <div id=\"Perturbations\"> Perturbations of a molecule</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's add a little noise to the fingerprints of our little fragment by turning some random bits on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "smiles = []\n",
    "extra_fp_on = [30, 40, 50, 60]\n",
    "torch.random.manual_seed(314)\n",
    "for num_bits_on in extra_fp_on:\n",
    "    random_indices = torch.randint(0, fp.size().numel(), torch.Size([num_bits_on]))\n",
    "    fp_add = torch.zeros_like(fp)\n",
    "    fp_add[random_indices] = 1\n",
    "    fp_pert = fp + fp_add\n",
    "    s = get_most_probable_valid_smiles(fp_pert, rmodel, max_branches=100)\n",
    "    smiles.append(s)\n",
    "display(imageOfMols(smiles, labels=extra_fp_on))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <div id=\"intro_to_algebra\">Intro to algebra of molecules</div>\n",
    "We can \"add\" two molecules by mixing their fingerprints \n",
    "(or by mixing their embeddings, or other internal layers)\n",
    "Even though the model is highly nonlinear, the \"addition\" is often intuitive, \n",
    "for example when the model is able to combine the fingerprints \n",
    "and create a molecule that matches both inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fragment2 = \"c1cccnc1N2CCCCC2\"\n",
    "\n",
    "imageOfMols([fragment, fragment2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "%%capture --no-stdout --no-display\n",
    "fps = [smiles_to_fingerprint(x, as_tensor=False) for x in [fragment, fragment2]]\n",
    "ftarget = (fps[0] | fps[1])\n",
    "ftarget = torch.Tensor(ftarget).cuda()\n",
    "smiles = dedup_clean(get_first_n_most_probable_valid_smiles(6, ftarget, \n",
    "                                                           rmodel))\n",
    "labels = [tanimoto(smiles_to_fingerprint(x), ftarget) for x in smiles]\n",
    "labels = [\"A\", \"B\"] + [str(x) for x in np.round(labels, 2)]\n",
    "display(imageOfMols([fragment, fragment2, *smiles], labels=labels))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get rid of the tri-fluoromethyl group"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "fragment3 = \"c1ccccc1C(F)(F)F\"\n",
    "imageOfMols([fragment, fragment2, fragment3], labels=[\"A\", \"+  B\", \"-  C\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "%%capture --no-stdout --no-display\n",
    "fps = [smiles_to_fingerprint(x, as_tensor=False) for x in [fragment, fragment2, fragment3]]\n",
    "f3target = (fps[0] | fps[1]) - fps[2]\n",
    "f3target = np.clip(f3target, 0, 1)\n",
    "f3target = torch.Tensor(f3target).cuda()\n",
    "smiles = dedup_clean(get_first_n_most_probable_valid_smiles(5, f3target, rmodel))\n",
    "labels = [tanimoto(smiles_to_fingerprint(x), f3target) for x in smiles]\n",
    "labels = [\"A\", \"+ B\", \"- C\"] + [str(x) for x in np.round(labels, 2)]\n",
    "display(imageOfMols([fragment, fragment2, fragment3, *smiles], labels=labels))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <div id=\"fine_tuning\">Fine tuning applications</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A promising way to generate new molecules is to finetune the DESMILES model to improve the outputs using training inputs from matched pairs.  This somewhat more complicated application of the model was described in the publication, and was used to generate the figures below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "img = Image(filename=f\"{fig_dir}/deep learn chem space (desmiles)__extended data fig 5__2__nisonoff__2019__.png\")\n",
    "display(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "img = Image(filename=f\"{fig_dir}/deep learn chem space (desmiles)__extended data fig 4__2__nisonoff__2019__.png\")\n",
    "display(img)"
   ]
  }
 ],
 "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": "FPEmbedding_Paul.ipynb"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}