2_RL_biased_G__4.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ae439c53",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "painted-irrigation",
   "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.PG_agent import PolicyGradientAgent\n",
    "from models.charRNN import generate\n",
    "from utils import n_chars, all_chars, check_validity, check_novelty, strsmis2listsmis, list2txt, get_props"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "handmade-specification",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1e237b0",
   "metadata": {},
   "source": [
    "## Load the general generator:\n",
    "* please see the file PG_agent.py inside the models folder for the code behind the developed model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d4931958",
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 0.00005   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "federal-disney",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name =  'G'   \n",
    "agent = PolicyGradientAgent(n_chars=n_chars, hidden_size=512, num_layers=3, \n",
    "                            output_size=n_chars, lr=lr, policy_file=model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "accurate-suite",
   "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(agent.policy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b48d6025",
   "metadata": {},
   "source": [
    "# Training function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a555ba16",
   "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='_4_'\n",
    "\n",
    "k = 20  \n",
    "n_epochs = 800\n",
    "\n",
    "trained_model = 'models/pretrained/RL_' +obj+str(lr)+'_'+str(n_epochs)+'_'+str(k)+ '_'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "annual-gathering",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def train_REINFORCE(n_epochs, k=1): \n",
    "    scores = []\n",
    "    smis = []\n",
    "    prev_high_r = 0\n",
    "    score = 0\n",
    "    avg_score = 0\n",
    "    for i in range(n_epochs*k):\n",
    "        done = False\n",
    "        ic = '<' \n",
    "        observation = all_chars.index(ic)\n",
    "        step = 0\n",
    "        length = 102  #max_length\n",
    "        smi = ''\n",
    "        while not done:\n",
    "            action = agent.choose_action(observation) \n",
    "            pred_char = all_chars[action]\n",
    "            if pred_char=='>' or step==length:  #if any of the terminal conditions is reached ----\n",
    "                prop1, prop2, prop3, prop4, prop5 = get_props(smi, c=1)\n",
    "                arr, alr, oh, cooh, coor, nh2, rval = get_props(smi, c=2)\n",
    "\n",
    "                cycle, grp, rv = 0, 0, 0        # -------------------------------   \n",
    "                if arr==2 or alr==1: \n",
    "                    cycle = 1\n",
    "                if oh>=1 or cooh>=1 or coor>=1 or nh2>=1:\n",
    "                    grp = 1\n",
    "                if 0.05<rval<0.5:\n",
    "                    rv = 1\n",
    "                                               #----------------------------------\n",
    "                if cycle==1 and grp==1 and rv==1 and prop1<=3 and prop2<=480 and prop3<=3 and prop4<=3 and prop5<=3: \n",
    "                    reward = 10           \n",
    "    \n",
    "                score += reward\n",
    "                agent.store_rewards(reward)\n",
    "                done = True\n",
    "            \n",
    "            else:                              #if terminal conditions are not reached yet -----\n",
    "                smi += pred_char\n",
    "                reward = 0                     #step reward = 0 \n",
    "                score += reward\n",
    "                step += 1\n",
    "                ic = pred_char \n",
    "                observation = all_chars.index(ic)\n",
    "            agent.store_rewards(reward)\n",
    "        \n",
    "        if (i+1) % k == 0:  #learn each k episodes \n",
    "            agent.learn()\n",
    "            avg_score = score/(k*10)\n",
    "            scores.append(avg_score)\n",
    "            print(\"episod: \", i, \"episodic reward: \", avg_score)\n",
    "            score = 0\n",
    "            \n",
    "        if prev_high_r < avg_score:  #save trained model given by the best epoch\n",
    "            print('--- improvment made ---')\n",
    "            prev_high_r = avg_score\n",
    "            best_model = trained_model + 'best_.pth'\n",
    "            agent.save_policy(best_model)\n",
    "\n",
    "    last_model = trained_model + 'last_.pth'   #save trained model given by the last epoch\n",
    "    agent.save_policy(last_model) \n",
    "    return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "facial-aurora",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "scores = train_REINFORCE(n_epochs, k=k)   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "f2e38c0a",
   "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",
    "smooth_score = moving_average(scores, 5)\n",
    "plt.plot(smooth_score)\n",
    "# plt.savefig('gen_smis/RL_obj4.png', dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d455d744",
   "metadata": {},
   "outputs": [],
   "source": [
    "# np.savetxt('RL_scores_4_.txt', scores, delimiter=',')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe9f4f95",
   "metadata": {},
   "source": [
    "## 0. generate: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "bc0f4604",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model_name = 'RL_4'  #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": 26,
   "id": "d0ebdb9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = PolicyGradientAgent(n_chars=n_chars, hidden_size=512, num_layers=3, \n",
    "                            output_size=n_chars, policy_file= model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "d3aca9ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "rnn = agent.policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "4d8a6877",
   "metadata": {},
   "outputs": [],
   "source": [
    "RL_temp = 0.5   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "793170fd",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "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=RL_temp)\n",
    "    all_smis.append(smi)\n",
    "    print(\"iteration:\", ep)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "eaebc03b",
   "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": "32fc1660",
   "metadata": {},
   "source": [
    "## 1. How many among all_smis are valid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "2668d2e8",
   "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": 48,
   "id": "ebc8c74d",
   "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": "ae2e9d7f",
   "metadata": {},
   "source": [
    "## 2. How many among val_smis are novel:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "2ae7ecb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "file = unidecode.unidecode(open('all_data/s_100_str_+1M.txt').read())\n",
    "smis_list = strsmis2listsmis(file)\n",
    "len(smis_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "59252d2a",
   "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": 43,
   "id": "d845edd4",
   "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": "2c40509a",
   "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": 35,
   "id": "fb856e90",
   "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": 42,
   "id": "e1da45b4",
   "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": "87efd22f",
   "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": 41,
   "id": "8ec715e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import moses\n",
    "moses.metrics.internal_diversity(val_nov_and_unique_smis)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01357e49",
   "metadata": {},
   "source": [
    "## 5. How many have the desired chemical property:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "5a664490",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "succ = []\n",
    "for i, smi in enumerate(val_nov_and_unique_smis):\n",
    "    arr, alr, oh, cooh, coor, nh2, rval = get_props(smi, c=2)                                                        \n",
    "    cycle, grp, rv = 0, 0, 0 \n",
    "    if arr==2 or alr==1: \n",
    "        cycle = 1\n",
    "    if oh>=1 or cooh>=1 or coor>=1 or nh2>=1:\n",
    "        grp = 1\n",
    "    if 0.05<rval<0.5:\n",
    "        rv = 1\n",
    "    \n",
    "    prop1, prop2, prop3, prop4, prop5 = get_props(smi, c=1)\n",
    "        \n",
    "    if cycle==1 and grp==1 and rv==1 and prop1<=3 and prop2<=480 and prop3<=3 and prop4<=3 and prop5<=3:\n",
    "        succ.append(smi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "6f97e5a1",
   "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": "a77e2e95",
   "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
}