{

 "cells": [

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "%load_ext autoreload\n",

    "%autoreload 2\n",

    "\n",

    "\n",

    "import numpy as np\n",

    "import rdkit\n",

    "from rdkit import Chem\n",

    "\n",

    "import h5py, ast, pickle\n",

    "\n",

    "# Occupy a GPU for the model to be loaded \n",

    "%env CUDA_DEVICE_ORDER=PCI_BUS_ID\n",

    "# GPU ID, if occupied change to an available GPU ID listed under !nvidia-smi\n",

    "%env CUDA_VISIBLE_DEVICES=2 \n",

    "\n",

    "from ddc_pub import ddc_v3 as ddc"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "# Load model"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "# Import existing (trained) model\n",

    "# Ignore any warning(s) about training configuration or non-seriazable keyword arguments\n",

    "model_name = \"models/heteroencoder_model\"\n",

    "model = ddc.DDC(model_name=model_name)"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "# Load data from dataset"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "dataset_name = \"datasets/CHEMBL25_TEST.h5\"\n",

    "npoints = 1000\n",

    "\n",

    "dataset = h5py.File(dataset_name, \"r\")\n",

    "mols    = dataset[\"mols\"][:]\n",

    "# Select random npoints\n",

    "mols_in = mols[np.random.choice(len(mols), npoints, replace=False)]\n",

    "dataset.close()\n",

    "\n",

    "# Get the SMILES behind the binary mols\n",

    "smiles_in = [Chem.MolToSmiles(Chem.Mol(mol)) for mol in mols_in]"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "# Alternatively, use your own SMILES"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "# Input SMILES to auto-encode\n",

    "smiles_in = ['Cc1cccn2c(CN(C)C3CCCc4ccccc43)c(C(=O)N3CCOCC3)nc12',\n",

    "             'COC(=O)NN=C(c1ccc(O)cc1)C1C(=O)N(C)C(=O)N(C)C1=O',\n",

    "             'CCc1cc(CC)nc(OCCCn2c3c(c4cc(-c5nc(C)no5)ccc42)CC(F)(F)CC3)n1',\n",

    "             'Cc1ccc2c(C(=O)Nc3ccccc3)c(SSc3c(C(=O)Nc4ccccc4)c4ccc(C)cc4n3C)n(C)c2c1',\n",

    "             'Cc1cccc(-c2ccccc2)c1Oc1nc(O)nc(NCc2ccc3occc3c2)n1',\n",

    "             'Cn1nnnc1SCC(=O)NN=Cc1ccc(Cl)cc1',\n",

    "             'COc1cccc(NS(=O)(=O)c2ccc(OC)c(OC)c2)c1',\n",

    "             'COc1ccc(OC)c(S(=O)(=O)n2nc(C)cc2C)c1',\n",

    "             'NCCCn1cc(C2=C(c3ccncc3)C(=O)NC2=O)c2ccccc21',\n",

    "             'CN(C)C(=O)N1CCN(C(c2ccc(Cl)cc2)c2cccnc2)CC1']\n",

    "\n",

    "# MUST convert SMILES to binary mols for the model to accept them (it re-converts them to SMILES internally)\n",

    "mols_in = [Chem.rdchem.Mol.ToBinary(Chem.MolFromSmiles(smiles)) for smiles in smiles_in]"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "# Encode the binary mols into their latent representations\n",

    "latent = model.transform(model.vectorize(mols_in))"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "# Convert back to SMILES\n",

    "smiles_out = []\n",

    "for lat in latent:   \n",

    "    smiles, _ = model.predict(lat, temp=0)\n",

    "    smiles_out.append(smiles)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "# To compare the results, convert smiles_out to CANONICAL\n",

    "for idx, smiles in enumerate(smiles_out):\n",

    "    mol = Chem.MolFromSmiles(smiles)\n",

    "    if mol:\n",

    "        smiles_out[idx] = Chem.MolToSmiles(mol, canonical=True)\n",

    "    else:\n",

    "        smiles_out[idx] = \"INVALID\""

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "scrolled": true

   },

   "outputs": [],

   "source": [

    "smiles_in"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "scrolled": true

   },

   "outputs": [],

   "source": [

    "smiles_out"

   ]

  }

 ],

 "metadata": {

  "kernelspec": {

   "display_name": "ddc",

   "language": "python",

   "name": "ddc"

  },

  "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.6.7"

  }

 },

 "nbformat": 4,

 "nbformat_minor": 2

}