2_CRLV_biased_G__3.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "547ce239",
   "metadata": {},
   "source": [
    "## Imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "above-exhibit",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "import unidecode\n",
    "import random\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from models.charRNN import make_charRNN, get_random_batch, generate\n",
    "from utils import n_chars, check_validity, check_novelty, strsmis2listsmis, list2txt, get_props"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "86245df7",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62a0c0c7",
   "metadata": {},
   "source": [
    "## Data preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "freelance-fundamental",
   "metadata": {},
   "outputs": [],
   "source": [
    "file = unidecode.unidecode(open('all_data/s_100_str_+1M.txt').read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "innovative-forty",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<COc1ccc2[C@@H]3[C@H](COc2c1)C(C)(C)OC4=C3C(=O)C(=O)C5=C4OC(C)(C)[C@@H]6COc7cc(OC)ccc7[C@H]56><C[S+]'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "file[0:100]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b1e301b",
   "metadata": {},
   "source": [
    "## Make the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2b15bbab",
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 0.0005"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "spectacular-italian",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'G'\n",
    "rnn, optimizer, criterion = make_charRNN(n_chars=56, hidden_size=512, num_layers=3,\n",
    "                                        lr=lr, pretrained_file=model_name) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5d0ba093",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RNN(\n",
      "  (embed): Embedding(56, 30)\n",
      "  (gru): GRU(30, 512, num_layers=3, batch_first=True)\n",
      "  (fc): Linear(in_features=512, out_features=56, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "print(rnn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7bcc61e",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dcf4eaf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "## --- for reproducibility of the results ---\n",
    "seeds = [3, 0, 3]\n",
    "\n",
    "torch.manual_seed(seeds[0])\n",
    "np.random.seed(seeds[1])\n",
    "random.seed(seeds[2])\n",
    "\n",
    "obj = '_3_'\n",
    "k=20  \n",
    "chunk_len=128  \n",
    "n_epochs = 800\n",
    "\n",
    "trained_model = 'models/pretrained/CRLV_' +obj+str(lr)+'_'+str(n_epochs)+'_'+str(k)+'_'+str(chunk_len)+'_'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "smart-validation",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(k=20, chunk_len=128, n_epochs=800, print_every=10):\n",
    "    print(\"=> Starting training...\")\n",
    "    Ls = []\n",
    "    prev_best_loss = 1000000\n",
    "    for epoch in range(1, n_epochs + 1):            \n",
    "        inp, target = get_random_batch(file, chunk_len)\n",
    "        hidden = rnn.init_hidden()\n",
    "\n",
    "        rnn.zero_grad()\n",
    "        loss = 0\n",
    "        inp = inp.to(device)\n",
    "        target = target.to(device)\n",
    "        \n",
    "        for i in range(chunk_len): \n",
    "            output, hidden = rnn(inp[:, i], hidden)\n",
    "            loss += criterion(output, target[:,i])\n",
    "\n",
    "#-------- the following lines are for conditional loss ---------------: \n",
    "        n = 1\n",
    "        for i in range(k):\n",
    "            smi = generate(rnn, initial_str='<', predict_len=102, temperature=0.50)\n",
    "            prop1, prop2, prop3, prop4, prop5 = get_props(smi, c=1)\n",
    "            if prop1<=3 and prop2<=480 and prop3<=3 and prop4<=3 and prop5<=3:\n",
    "                n = n*2\n",
    "            \n",
    "        loss = loss / n\n",
    "#---------------------------------------------------------------------------------------\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        loss = loss.item()/(chunk_len)\n",
    "        \n",
    "        if prev_best_loss > loss:                     # save trained model given by the best epoch\n",
    "            prev_best_loss = loss\n",
    "            best_model = trained_model + 'best_.pth'\n",
    "            torch.save(rnn.state_dict(), best_model)\n",
    "        \n",
    "        Ls.append(loss)\n",
    "        if epoch % print_every == 0:  \n",
    "            print(f'Loss: {loss} epoch: {epoch}')   \n",
    "        \n",
    "#     last_model = trained_model + 'last_.pth'            # save trained model given by the last epoch\n",
    "#     torch.save(rnn.state_dict(), last_model)               \n",
    "    return Ls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "bcf08041",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "Losses = train(k=k, chunk_len=chunk_len, n_epochs=n_epochs, print_every=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "72ec700d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def moving_average(values, window):\n",
    "    weights = np.repeat(1.0, window) / window\n",
    "    return np.convolve(values, weights, 'valid')\n",
    "\n",
    "Losses_ = moving_average(Losses, 20)\n",
    "plt.plot(Losses_)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "focused-initial",
   "metadata": {},
   "source": [
    "## 0. generate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db358e59",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model_name = 'CRLV_3'  #use our pretrained model, or the one you just trained:\n",
    "model_name = trained_model.replace('models/pretrained/', '') + 'best_'  #or +'last'\n",
    "model_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35e4b559",
   "metadata": {},
   "outputs": [],
   "source": [
    "rnn, optimizer, criterion = make_charRNN(n_chars=56, hidden_size=512, num_layers=3,\n",
    "                                        lr=lr, pretrained_file=model_name) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8e09cde4",
   "metadata": {},
   "outputs": [],
   "source": [
    "CRLV_temp = 0.5 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "certain-booth",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#-----------------------generate:\n",
    "n_ep = 10000\n",
    "all_smis = []\n",
    "for ep in range(n_ep):\n",
    "    length = 102 \n",
    "    ic = '<'\n",
    "    smi = generate(rnn, initial_str=ic, predict_len=length, temperature=CRLV_temp)\n",
    "    all_smis.append(smi)\n",
    "    print(\"iteration:\", ep)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0251adc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "### save the list of generated smiles if you want\n",
    "# filename = model_name + '_all_smis_.txt'\n",
    "# path = 'gen_smis/' + filename\n",
    "# list2txt(path, mylist=all_smis)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ac389cc",
   "metadata": {},
   "source": [
    "## 1. How many among all_smis are valid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a6ed67b7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "val_smis = []\n",
    "for i, s in enumerate(all_smis):\n",
    "    v = check_validity(s)\n",
    "    if v == 1:\n",
    "        val_smis.append(s)\n",
    "        print(\"mol at:\", i, \" : verified\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "eab16342",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"There is \", len(val_smis), \"valid molecule among the\", len(all_smis), \"generated\")\n",
    "print(\"meaning a percent of: \", 100 * len(val_smis)/len(all_smis), \"%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "642f6393",
   "metadata": {},
   "source": [
    "## 2. How many among val_smis are novel:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "6dd3a15d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1498669"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "smis_list = strsmis2listsmis(file)\n",
    "len(smis_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "8ce90768",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "val_and_nov_smis = []\n",
    "for i, s in enumerate(val_smis):\n",
    "    nv = check_novelty(s, smis_list)\n",
    "    if nv == 1:\n",
    "        val_and_nov_smis.append(s)\n",
    "        print(\"mol at:\", i, \" : verified\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "201322f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"There is \", len(val_and_nov_smis), \"valid & novel molecule among the\", len(all_smis), \"generated\")\n",
    "print(\"meaning a percent of: \", 100 * len(val_and_nov_smis)/ len(val_smis), \"%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a046a943",
   "metadata": {},
   "source": [
    "## 3. How many among valid_and_novel_smis are unique \n",
    "* meaning, they don't have duplicates/were not repeated in the generated molecules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f27de21d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_uniqueness(smis_list):\n",
    "    return list(set(smis_list))\n",
    "\n",
    "val_nov_and_unique_smis = check_uniqueness(val_and_nov_smis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "500d88e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"There is \", len(val_nov_and_unique_smis), \"valid, novel and unique molecule among the\", len(all_smis), \"generated\")\n",
    "print(\"meaning, a percent of : \", 100 * len(val_nov_and_unique_smis)/ len(val_and_nov_smis), \"%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99e04490",
   "metadata": {},
   "source": [
    "## 4. Diversity (intDiv) of the valid, novel, and unique smiles:\n",
    "* using the moses library https://github.com/molecularsets/moses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "116ddb85",
   "metadata": {},
   "outputs": [],
   "source": [
    "import moses\n",
    "moses.metrics.internal_diversity(val_nov_and_unique_smis)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95f9b4ac",
   "metadata": {},
   "source": [
    "## 5. How many have the desired chemical property:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "494afddb",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "succ = []\n",
    "for i, smi in enumerate(val_nov_and_unique_smis):\n",
    "    prop1, prop2, prop3, prop4, prop5 = get_props(smi, c=1)\n",
    "    if prop1<=3 and prop2<=480 and prop3<=3 and prop4<=3 and prop5<=3:\n",
    "        succ.append(smi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "5d1242c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('there is', len(succ), 'molecule with the desired property')\n",
    "print(\"meaning a percent of\", 100 * len(succ)/len(val_nov_and_unique_smis), \"% of smiles among the valid, \\\n",
    "novel and unique ones with the desired property\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3464be3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}