graph_convolution.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from rlvs.molecule_world.datasets import DataStore\n",
    "\n",
    "import tensorflow as tf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.2.0'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rlvs.molecule_world.helper_functions import *\n",
    "file_ligand = \"data/pafnucy_data/complexes/4mrw/4mrw_ligand.mol2\"\n",
    "file_protein = \"data/pafnucy_data/complexes/4mrw/4mrw_pocket.mol2\"\n",
    "\n",
    "f = Featurizer()\n",
    "\n",
    "obmol_ligand = read_to_OB(filename=file_ligand, filetype=\"mol2\")\n",
    "nodes_ligand, canon_adj_list_ligand = f.get_mol_features(obmol=obmol_ligand, molecule_type=1, bond_verbose=0)\n",
    "ligand = Molecule(atom_features=nodes_ligand, canon_adj_list=canon_adj_list_ligand)\n",
    "\n",
    "obmol_protein = read_to_OB(filename=file_protein, filetype=\"mol2\")\n",
    "nodes_protein, canon_adj_list_protein = f.get_mol_features(obmol=obmol_protein, molecule_type=-1, bond_verbose=0)\n",
    "protein = Molecule(atom_features=nodes_protein, canon_adj_list=canon_adj_list_protein)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from rdkit import Chem\n",
    "# from deepchem.feat import ConvMolFeaturizer\n",
    "\n",
    "# smiles = ['COC(C)(C)CCCC(C)CC=CC(C)=CC(=O)OC(C)C',\n",
    "#           'CCOC(=O)CC',\n",
    "#           'CSc1nc(NC(C)C)nc(NC(C)C)n1',\n",
    "#           'CC(C#C)N(C)C(=O)Nc1ccc(Cl)cc1',\n",
    "#           'Cc1cc2ccccc2cc1C']\n",
    "\n",
    "# mols = [Chem.MolFromSmiles(s) for s in smiles]\n",
    "# featurizer = ConvMolFeaturizer()\n",
    "# mols = featurizer.featurize(mols)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((550, 18), (18, 18))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mol1 = protein\n",
    "mol2 = ligand\n",
    "mol1.get_atom_features().shape, mol2.get_atom_features().shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare tensor inputs for convolution layer\n",
    "\n",
    "TODO: Feature normalization?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = []\n",
    "for mol in [mol1, mol2]:\n",
    "    atom_features = np.expand_dims(mol.get_atom_features(), axis=0)\n",
    "    degree_slice = np.expand_dims((mol.deg_slice.astype(dtype='int32')), axis=0)\n",
    "    deg_adjs = [np.expand_dims(deg_adj.astype(dtype='int32'), axis=0) for deg_adj in mol.get_deg_adjacency_lists()[1:]]\n",
    "\n",
    "    gc_in = [atom_features, degree_slice] + deg_adjs\n",
    "    \n",
    "    inputs.append(gc_in)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define model (inputs -> protein-ligand combined feature vector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model_2\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "input_1 (InputLayer)            [(1, None, 18)]      0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_2 (InputLayer)            [(1, 11, 2)]         0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_3 (InputLayer)            [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_4 (InputLayer)            [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_5 (InputLayer)            [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_6 (InputLayer)            [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_7 (InputLayer)            [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_8 (InputLayer)            [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_9 (InputLayer)            [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_10 (InputLayer)           [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_11 (InputLayer)           [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_12 (InputLayer)           [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_13 (InputLayer)           [(1, None, 18)]      0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_14 (InputLayer)           [(1, 11, 2)]         0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_15 (InputLayer)           [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_16 (InputLayer)           [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_17 (InputLayer)           [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_18 (InputLayer)           [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_19 (InputLayer)           [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_20 (InputLayer)           [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_21 (InputLayer)           [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_22 (InputLayer)           [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_23 (InputLayer)           [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_24 (InputLayer)           [(1, None, None)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "graph_conv (GraphConv)          (None, 64)           25536       input_1[0][0]                    \n",
      "                                                                 input_2[0][0]                    \n",
      "                                                                 input_3[0][0]                    \n",
      "                                                                 input_4[0][0]                    \n",
      "                                                                 input_5[0][0]                    \n",
      "                                                                 input_6[0][0]                    \n",
      "                                                                 input_7[0][0]                    \n",
      "                                                                 input_8[0][0]                    \n",
      "                                                                 input_9[0][0]                    \n",
      "                                                                 input_10[0][0]                   \n",
      "                                                                 input_11[0][0]                   \n",
      "                                                                 input_12[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "graph_conv_1 (GraphConv)        (None, 64)           25536       input_13[0][0]                   \n",
      "                                                                 input_14[0][0]                   \n",
      "                                                                 input_15[0][0]                   \n",
      "                                                                 input_16[0][0]                   \n",
      "                                                                 input_17[0][0]                   \n",
      "                                                                 input_18[0][0]                   \n",
      "                                                                 input_19[0][0]                   \n",
      "                                                                 input_20[0][0]                   \n",
      "                                                                 input_21[0][0]                   \n",
      "                                                                 input_22[0][0]                   \n",
      "                                                                 input_23[0][0]                   \n",
      "                                                                 input_24[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "graph_pool (GraphPool)          (None, 64)           0           graph_conv[0][0]                 \n",
      "                                                                 input_2[0][0]                    \n",
      "                                                                 input_3[0][0]                    \n",
      "                                                                 input_4[0][0]                    \n",
      "                                                                 input_5[0][0]                    \n",
      "                                                                 input_6[0][0]                    \n",
      "                                                                 input_7[0][0]                    \n",
      "                                                                 input_8[0][0]                    \n",
      "                                                                 input_9[0][0]                    \n",
      "                                                                 input_10[0][0]                   \n",
      "                                                                 input_11[0][0]                   \n",
      "                                                                 input_12[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "graph_pool_1 (GraphPool)        (None, 64)           0           graph_conv_1[0][0]               \n",
      "                                                                 input_14[0][0]                   \n",
      "                                                                 input_15[0][0]                   \n",
      "                                                                 input_16[0][0]                   \n",
      "                                                                 input_17[0][0]                   \n",
      "                                                                 input_18[0][0]                   \n",
      "                                                                 input_19[0][0]                   \n",
      "                                                                 input_20[0][0]                   \n",
      "                                                                 input_21[0][0]                   \n",
      "                                                                 input_22[0][0]                   \n",
      "                                                                 input_23[0][0]                   \n",
      "                                                                 input_24[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense (Dense)                   (None, 128)          8320        graph_pool[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "dense_1 (Dense)                 (None, 128)          8320        graph_pool_1[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "graph_gather (GraphGather)      (1, 256)             0           dense[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "graph_gather_1 (GraphGather)    (1, 256)             0           dense_1[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "add (Add)                       (1, 256)             0           graph_gather[0][0]               \n",
      "                                                                 graph_gather_1[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "dense_2 (Dense)                 (1, 64)              16448       add[0][0]                        \n",
      "==================================================================================================\n",
      "Total params: 84,160\n",
      "Trainable params: 84,160\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "from rlvs.network.graph_layer import GraphConv, GraphPool, GraphGather\n",
    "from tensorflow.keras.layers import Input, Dense, add\n",
    "from tensorflow.keras.models import Model\n",
    "\n",
    "# first molecule inputs\n",
    "features_input_1 = Input(shape=(None,mol.n_feat,), batch_size=1)\n",
    "degree_slice_input_1 = Input(shape=(11,2), dtype=tf.int32, batch_size=1) #degree per atom\n",
    "deg_adjs_input_1 = []\n",
    "for i in range(10):\n",
    "    deg_adjs_input_1.append(Input(shape=(None,None,), dtype=tf.int32, batch_size=1))\n",
    "ip_1 = [features_input_1, degree_slice_input_1] + deg_adjs_input_1\n",
    "\n",
    "# first molecule convolution\n",
    "graph_layer_1 = GraphConv(out_channel=64, activation_fn=tf.nn.relu)(ip_1)\n",
    "gp_in_1 = [graph_layer_1, degree_slice_input_1] + deg_adjs_input_1\n",
    "\n",
    "# first molecule pooling (analogous to max pooling)\n",
    "graph_pool_1 = GraphPool()(gp_in_1)\n",
    "dense_layer_1 = Dense(128, activation=\"relu\")(graph_pool_1)\n",
    "\n",
    "# first molecule gather node level features into molecule level features\n",
    "graph_gather_layer_1 = GraphGather(activation_fn=tf.nn.relu)(dense_layer_1)\n",
    "\n",
    "mol1_model = Model(inputs=ip_1, outputs=graph_gather_layer_1)\n",
    "\n",
    "# second molecule inputs\n",
    "features_input_2 = Input(shape=(None,mol.n_feat,), batch_size=1)\n",
    "degree_slice_input_2 = Input(shape=(11,2), dtype=tf.int32, batch_size=1)\n",
    "deg_adjs_input_2 = []\n",
    "for i in range(10):\n",
    "    deg_adjs_input_2.append(Input(shape=(None,None,), dtype=tf.int32, batch_size=1))\n",
    "ip_2 = [features_input_2, degree_slice_input_2] + deg_adjs_input_2\n",
    "\n",
    "# second molecule convolution\n",
    "graph_layer_2 = GraphConv(out_channel=64, activation_fn=tf.nn.relu)(ip_2)\n",
    "gp_in_2 = [graph_layer_2, degree_slice_input_2] + deg_adjs_input_2\n",
    "\n",
    "# second molecule pooling (analogous to max pooling)\n",
    "graph_pool_2 = GraphPool()(gp_in_2)\n",
    "dense_layer_2 = Dense(128, activation=\"relu\")(graph_pool_2)\n",
    "\n",
    "# second molecule gather node level features into molecule level features\n",
    "graph_gather_layer_2 = GraphGather(activation_fn=tf.nn.relu)(dense_layer_2)\n",
    "\n",
    "mol2_model = Model(inputs=ip_2, outputs=graph_gather_layer_2)\n",
    "\n",
    "# combine molecule1 and molecule2 feature vectors\n",
    "#TODO: is \"add\" the correct operation?\n",
    "combination_layer = add([mol1_model.output, mol2_model.output])\n",
    "combined_dense_layer = Dense(64, activation=\"relu\")(combination_layer)\n",
    "\n",
    "# define full model\n",
    "m = Model([ip_1, ip_2], combined_dense_layer)\n",
    "m.compile()\n",
    "m.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Actor Network Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "18"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mol.n_feat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([[[ 9.5850e+00,  6.3040e+00,  4.8043e+01, ...,  0.0000e+00,\n",
       "           2.0000e+00, -1.0000e+00],\n",
       "         [ 1.5418e+01,  1.4069e+01,  4.5875e+01, ...,  1.0000e+00,\n",
       "           2.7190e-01, -1.0000e+00],\n",
       "         [ 1.3540e+01,  1.2143e+01,  4.8911e+01, ...,  0.0000e+00,\n",
       "          -5.6790e-01, -1.0000e+00],\n",
       "         ...,\n",
       "         [-2.9130e+00,  1.1315e+01,  4.6582e+01, ...,  0.0000e+00,\n",
       "          -1.8250e-01, -1.0000e+00],\n",
       "         [-3.8290e+00,  2.2978e+01,  4.4985e+01, ...,  1.0000e+00,\n",
       "          -2.7500e-02, -1.0000e+00],\n",
       "         [-3.4360e+00,  2.1593e+01,  4.4447e+01, ...,  0.0000e+00,\n",
       "          -5.0000e-03, -1.0000e+00]]]),\n",
       " array([[[  0,   1],\n",
       "         [  1, 240],\n",
       "         [241,  28],\n",
       "         [269,  96],\n",
       "         [365,  64],\n",
       "         [  0,   0],\n",
       "         [  0,   0],\n",
       "         [  0,   0],\n",
       "         [  0,   0],\n",
       "         [  0,   0],\n",
       "         [  0,   0]]], dtype=int32),\n",
       " array([[[241],\n",
       "         [269],\n",
       "         [242],\n",
       "         [365],\n",
       "         [366],\n",
       "         [366],\n",
       "         [271],\n",
       "         [272],\n",
       "         [273],\n",
       "         [274],\n",
       "         [276],\n",
       "         [277],\n",
       "         [367],\n",
       "         [368],\n",
       "         [368],\n",
       "         [280],\n",
       "         [281],\n",
       "         [279],\n",
       "         [282],\n",
       "         [283],\n",
       "         [369],\n",
       "         [370],\n",
       "         [370],\n",
       "         [286],\n",
       "         [287],\n",
       "         [285],\n",
       "         [288],\n",
       "         [289],\n",
       "         [371],\n",
       "         [372],\n",
       "         [372],\n",
       "         [373],\n",
       "         [374],\n",
       "         [374],\n",
       "         [374],\n",
       "         [375],\n",
       "         [375],\n",
       "         [375],\n",
       "         [290],\n",
       "         [291],\n",
       "         [292],\n",
       "         [292],\n",
       "         [376],\n",
       "         [377],\n",
       "         [377],\n",
       "         [293],\n",
       "         [294],\n",
       "         [378],\n",
       "         [379],\n",
       "         [379],\n",
       "         [380],\n",
       "         [381],\n",
       "         [381],\n",
       "         [381],\n",
       "         [382],\n",
       "         [382],\n",
       "         [382],\n",
       "         [295],\n",
       "         [296],\n",
       "         [383],\n",
       "         [384],\n",
       "         [384],\n",
       "         [245],\n",
       "         [297],\n",
       "         [298],\n",
       "         [246],\n",
       "         [385],\n",
       "         [386],\n",
       "         [386],\n",
       "         [299],\n",
       "         [300],\n",
       "         [387],\n",
       "         [388],\n",
       "         [389],\n",
       "         [389],\n",
       "         [389],\n",
       "         [390],\n",
       "         [390],\n",
       "         [390],\n",
       "         [301],\n",
       "         [302],\n",
       "         [247],\n",
       "         [391],\n",
       "         [392],\n",
       "         [393],\n",
       "         [393],\n",
       "         [393],\n",
       "         [303],\n",
       "         [304],\n",
       "         [248],\n",
       "         [394],\n",
       "         [395],\n",
       "         [396],\n",
       "         [396],\n",
       "         [396],\n",
       "         [305],\n",
       "         [306],\n",
       "         [397],\n",
       "         [398],\n",
       "         [398],\n",
       "         [398],\n",
       "         [307],\n",
       "         [308],\n",
       "         [399],\n",
       "         [400],\n",
       "         [401],\n",
       "         [401],\n",
       "         [403],\n",
       "         [403],\n",
       "         [403],\n",
       "         [402],\n",
       "         [402],\n",
       "         [402],\n",
       "         [309],\n",
       "         [310],\n",
       "         [249],\n",
       "         [404],\n",
       "         [405],\n",
       "         [405],\n",
       "         [312],\n",
       "         [313],\n",
       "         [314],\n",
       "         [315],\n",
       "         [317],\n",
       "         [318],\n",
       "         [406],\n",
       "         [407],\n",
       "         [407],\n",
       "         [320],\n",
       "         [321],\n",
       "         [322],\n",
       "         [323],\n",
       "         [324],\n",
       "         [325],\n",
       "         [326],\n",
       "         [408],\n",
       "         [409],\n",
       "         [409],\n",
       "         [410],\n",
       "         [410],\n",
       "         [411],\n",
       "         [411],\n",
       "         [411],\n",
       "         [327],\n",
       "         [328],\n",
       "         [412],\n",
       "         [412],\n",
       "         [329],\n",
       "         [330],\n",
       "         [331],\n",
       "         [332],\n",
       "         [332],\n",
       "         [413],\n",
       "         [414],\n",
       "         [414],\n",
       "         [415],\n",
       "         [415],\n",
       "         [333],\n",
       "         [334],\n",
       "         [416],\n",
       "         [416],\n",
       "         [335],\n",
       "         [336],\n",
       "         [417],\n",
       "         [418],\n",
       "         [418],\n",
       "         [338],\n",
       "         [339],\n",
       "         [340],\n",
       "         [341],\n",
       "         [342],\n",
       "         [343],\n",
       "         [344],\n",
       "         [251],\n",
       "         [419],\n",
       "         [420],\n",
       "         [420],\n",
       "         [346],\n",
       "         [347],\n",
       "         [348],\n",
       "         [349],\n",
       "         [351],\n",
       "         [352],\n",
       "         [421],\n",
       "         [422],\n",
       "         [423],\n",
       "         [423],\n",
       "         [423],\n",
       "         [424],\n",
       "         [424],\n",
       "         [424],\n",
       "         [353],\n",
       "         [354],\n",
       "         [425],\n",
       "         [426],\n",
       "         [426],\n",
       "         [426],\n",
       "         [355],\n",
       "         [252],\n",
       "         [359],\n",
       "         [427],\n",
       "         [428],\n",
       "         [428],\n",
       "         [357],\n",
       "         [361],\n",
       "         [362],\n",
       "         [363],\n",
       "         [364],\n",
       "         [253],\n",
       "         [253],\n",
       "         [254],\n",
       "         [254],\n",
       "         [255],\n",
       "         [255],\n",
       "         [256],\n",
       "         [256],\n",
       "         [257],\n",
       "         [257],\n",
       "         [258],\n",
       "         [258],\n",
       "         [259],\n",
       "         [259],\n",
       "         [260],\n",
       "         [260],\n",
       "         [261],\n",
       "         [261],\n",
       "         [262],\n",
       "         [262],\n",
       "         [263],\n",
       "         [263],\n",
       "         [264],\n",
       "         [264],\n",
       "         [265],\n",
       "         [265],\n",
       "         [266],\n",
       "         [266],\n",
       "         [267],\n",
       "         [267],\n",
       "         [268],\n",
       "         [268]]], dtype=int32),\n",
       " array([[[365,   1],\n",
       "         [275,   3],\n",
       "         [280, 281],\n",
       "         [286, 287],\n",
       "         [384,  63],\n",
       "         [386,  66],\n",
       "         [392,  82],\n",
       "         [395,  90],\n",
       "         [316, 116],\n",
       "         [410, 411],\n",
       "         [350, 174],\n",
       "         [427, 199],\n",
       "         [209, 210],\n",
       "         [211, 212],\n",
       "         [213, 214],\n",
       "         [215, 216],\n",
       "         [217, 218],\n",
       "         [219, 220],\n",
       "         [221, 222],\n",
       "         [223, 224],\n",
       "         [225, 226],\n",
       "         [227, 228],\n",
       "         [229, 230],\n",
       "         [231, 232],\n",
       "         [233, 234],\n",
       "         [235, 236],\n",
       "         [237, 238],\n",
       "         [239, 240]]], dtype=int32),\n",
       " array([[[365,   2, 276],\n",
       "         [366, 271, 272],\n",
       "         [270, 273,   7],\n",
       "         [270, 274,   8],\n",
       "         [271, 275,   9],\n",
       "         [272, 275,  10],\n",
       "         [273, 274, 242],\n",
       "         [367,  11, 269],\n",
       "         [367,  12, 282],\n",
       "         [368, 280, 279],\n",
       "         [278, 281,  18],\n",
       "         [278, 243,  16],\n",
       "         [279, 243,  17],\n",
       "         [369,  19, 277],\n",
       "         [369,  20, 288],\n",
       "         [370, 286, 285],\n",
       "         [284, 287,  26],\n",
       "         [284, 244,  24],\n",
       "         [285, 244,  25],\n",
       "         [ 27, 371, 283],\n",
       "         [371,  28, 290],\n",
       "         [ 39, 376, 289],\n",
       "         [376,  40, 293],\n",
       "         [377,  41,  42],\n",
       "         [ 46, 378, 291],\n",
       "         [378,  47, 295],\n",
       "         [ 58, 383, 294],\n",
       "         [383,  59, 297],\n",
       "         [ 64, 385, 296],\n",
       "         [385,  65, 299],\n",
       "         [ 70, 387, 298],\n",
       "         [387,  71, 301],\n",
       "         [ 80, 391, 300],\n",
       "         [391,  81, 303],\n",
       "         [ 88, 394, 302],\n",
       "         [394,  89, 305],\n",
       "         [ 96, 397, 304],\n",
       "         [397,  97, 307],\n",
       "         [102, 399, 306],\n",
       "         [399, 103, 309],\n",
       "         [114, 404, 308],\n",
       "         [404, 115, 317],\n",
       "         [405, 312, 313],\n",
       "         [311, 314, 120],\n",
       "         [311, 315, 121],\n",
       "         [312, 316, 122],\n",
       "         [313, 316, 123],\n",
       "         [314, 315, 249],\n",
       "         [124, 406, 310],\n",
       "         [406, 125, 325],\n",
       "         [407, 320, 321],\n",
       "         [319, 322, 129],\n",
       "         [319, 323, 130],\n",
       "         [320, 324, 131],\n",
       "         [321, 324, 132],\n",
       "         [322, 323, 133],\n",
       "         [134, 408, 318],\n",
       "         [408, 135, 327],\n",
       "         [144, 412, 326],\n",
       "         [412, 145, 329],\n",
       "         [148, 413, 328],\n",
       "         [413, 149, 333],\n",
       "         [415, 150, 332],\n",
       "         [331, 151, 152],\n",
       "         [158, 416, 330],\n",
       "         [416, 159, 335],\n",
       "         [162, 417, 334],\n",
       "         [417, 163, 343],\n",
       "         [418, 338, 339],\n",
       "         [337, 340, 167],\n",
       "         [337, 341, 168],\n",
       "         [338, 342, 169],\n",
       "         [339, 342, 170],\n",
       "         [340, 341, 171],\n",
       "         [172, 419, 336],\n",
       "         [419, 173, 351],\n",
       "         [420, 346, 347],\n",
       "         [345, 348, 178],\n",
       "         [345, 349, 179],\n",
       "         [346, 350, 180],\n",
       "         [347, 350, 181],\n",
       "         [348, 349, 251],\n",
       "         [182, 421, 344],\n",
       "         [421, 183, 353],\n",
       "         [192, 425, 352],\n",
       "         [425, 193, 355],\n",
       "         [198, 427, 354],\n",
       "         [428, 358, 357],\n",
       "         [356, 359, 204],\n",
       "         [356, 360, 361],\n",
       "         [357, 200, 360],\n",
       "         [358, 359, 362],\n",
       "         [358, 363, 205],\n",
       "         [360, 364, 206],\n",
       "         [361, 364, 207],\n",
       "         [362, 363, 208]]], dtype=int32),\n",
       " array([[[241, 269, 366,   4],\n",
       "         [365, 270,   5,   6],\n",
       "         [276, 368, 277,  13],\n",
       "         [367, 278,  14,  15],\n",
       "         [282, 370, 283,  21],\n",
       "         [369, 284,  22,  23],\n",
       "         [288, 289, 372,  29],\n",
       "         [371, 373,  30,  31],\n",
       "         [372, 374, 375,  32],\n",
       "         [373,  33,  34,  35],\n",
       "         [373,  36,  37,  38],\n",
       "         [290, 377, 291,  43],\n",
       "         [376, 292,  44,  45],\n",
       "         [293, 294, 379,  48],\n",
       "         [378, 380,  49,  50],\n",
       "         [379, 381, 382,  51],\n",
       "         [380,  52,  53,  54],\n",
       "         [380,  55,  56,  57],\n",
       "         [295, 384, 296,  60],\n",
       "         [383, 245,  61,  62],\n",
       "         [297, 298, 386,  67],\n",
       "         [385, 246,  68,  69],\n",
       "         [299, 388, 300,  72],\n",
       "         [387, 389, 390,  73],\n",
       "         [388,  74,  75,  76],\n",
       "         [388,  77,  78,  79],\n",
       "         [301, 392, 302,  83],\n",
       "         [391, 393, 247,  84],\n",
       "         [392,  85,  86,  87],\n",
       "         [303, 395, 304,  91],\n",
       "         [394, 396, 248,  92],\n",
       "         [395,  93,  94,  95],\n",
       "         [305, 398, 306,  98],\n",
       "         [397,  99, 100, 101],\n",
       "         [307, 400, 308, 104],\n",
       "         [399, 402, 401, 105],\n",
       "         [400, 403, 106, 107],\n",
       "         [400, 111, 112, 113],\n",
       "         [401, 108, 109, 110],\n",
       "         [309, 405, 310, 117],\n",
       "         [404, 311, 118, 119],\n",
       "         [317, 407, 318, 126],\n",
       "         [406, 319, 127, 128],\n",
       "         [325, 409, 326, 136],\n",
       "         [408, 410, 137, 138],\n",
       "         [409, 250, 139, 140],\n",
       "         [250, 141, 142, 143],\n",
       "         [327, 328, 146, 147],\n",
       "         [329, 414, 330, 153],\n",
       "         [413, 415, 154, 155],\n",
       "         [414, 331, 156, 157],\n",
       "         [333, 334, 160, 161],\n",
       "         [335, 418, 336, 164],\n",
       "         [417, 337, 165, 166],\n",
       "         [343, 420, 344, 175],\n",
       "         [419, 345, 176, 177],\n",
       "         [351, 422, 352, 184],\n",
       "         [421, 423, 424, 185],\n",
       "         [422, 186, 187, 188],\n",
       "         [422, 189, 190, 191],\n",
       "         [353, 426, 354, 194],\n",
       "         [425, 195, 196, 197],\n",
       "         [355, 252, 428, 201],\n",
       "         [427, 356, 202, 203]]], dtype=int32),\n",
       " array([], shape=(1, 0, 5), dtype=int32),\n",
       " array([], shape=(1, 0, 6), dtype=int32),\n",
       " array([], shape=(1, 0, 7), dtype=int32),\n",
       " array([], shape=(1, 0, 8), dtype=int32),\n",
       " array([], shape=(1, 0, 9), dtype=int32),\n",
       " array([], shape=(1, 0, 10), dtype=int32)]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rlvs.network import ActorGNN\n",
    "from rlvs.molecule_world.env import GraphEnv\n",
    "env = GraphEnv()\n",
    "actor = ActorGNN(\n",
    "            18,\n",
    "            6,\n",
    "            0.00005,\n",
    "            0.001\n",
    "        )\n",
    "inputs = env.reset()\n",
    "actor.actor(inputs)\n",
    "inputs[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Forward run of test data through model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "o = m(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1, 64), dtype=float32, numpy=\n",
       "array([[  545.26935,     0.     ,     0.     ,  3791.2192 ,  2360.0024 ,\n",
       "         9495.732  ,  1209.4473 ,  4813.2646 ,     0.     ,   789.11066,\n",
       "          684.5331 ,   826.39307,  7852.6924 ,  3336.8691 ,     0.     ,\n",
       "            0.     ,  5453.1475 ,     0.     ,     0.     ,  2616.227  ,\n",
       "            0.     ,  4508.199  ,  3329.6023 ,  6843.436  ,     0.     ,\n",
       "         2942.189  ,     0.     ,     0.     ,   386.92505,     0.     ,\n",
       "         5114.6855 ,  5607.7563 , 10091.157  ,   925.6692 ,     0.     ,\n",
       "         3731.1091 ,     0.     ,  8693.792  ,     0.     ,  1674.274  ,\n",
       "            0.     ,  1856.3945 ,   466.98758,     0.     ,     0.     ,\n",
       "         2826.5874 ,  6300.3926 ,     0.     ,  1114.1597 ,  4391.5474 ,\n",
       "            0.     ,  3172.9285 ,     0.     ,     0.     ,     0.     ,\n",
       "         2555.6826 ,     0.     ,     0.     ,     0.     ,     0.     ,\n",
       "            0.     ,  1317.4073 ,  2840.6704 ,     0.     ]],\n",
       "      dtype=float32)>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[  545.26935,     0.     ,     0.     ,  3791.2192 ,  2360.0024 ,\n",
       "         9495.732  ,  1209.4473 ,  4813.2646 ,     0.     ,   789.11066,\n",
       "          684.5331 ,   826.39307,  7852.6924 ,  3336.8691 ,     0.     ,\n",
       "            0.     ,  5453.1475 ,     0.     ,     0.     ,  2616.227  ,\n",
       "            0.     ,  4508.199  ,  3329.6023 ,  6843.436  ,     0.     ,\n",
       "         2942.189  ,     0.     ,     0.     ,   386.92505,     0.     ,\n",
       "         5114.6855 ,  5607.7563 , 10091.157  ,   925.6692 ,     0.     ,\n",
       "         3731.1091 ,     0.     ,  8693.792  ,     0.     ,  1674.274  ,\n",
       "            0.     ,  1856.3945 ,   466.98758,     0.     ,     0.     ,\n",
       "         2826.5874 ,  6300.3926 ,     0.     ,  1114.1597 ,  4391.5474 ,\n",
       "            0.     ,  3172.9285 ,     0.     ,     0.     ,     0.     ,\n",
       "         2555.6826 ,     0.     ,     0.     ,     0.     ,     0.     ,\n",
       "            0.     ,  1317.4073 ,  2840.6704 ,     0.     ]],\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.predict(inputs)"
   ]
  }
 ],
 "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.6.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}