GCNN_batching.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from rlvs.molecule_world.datasets import DataStore\n",
    "from rlvs.molecule_world.complex import Complex\n",
    "from rlvs.molecule_world.helper_functions import *\n",
    "\n",
    "DataStore.init(crop=False)\n",
    "complexes = []\n",
    "ligands = []\n",
    "for i in range(2):\n",
    "    protein, ligand = DataStore.DATA[i]\n",
    "    print(protein.n_atoms, ligand.n_atoms)\n",
    "    ligands.append(ligand)\n",
    "    complex = Complex(protein, ligand)\n",
    "    complexes.append(complex)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from deepchem.feat.mol_graphs import MultiConvMol, ConvMol\n",
    "\n",
    "def mols_to_inputs(mols):\n",
    "    multiConvMol = ConvMol.agglomerate_mols(mols)\n",
    "    n_samples = np.array([len(mols)])\n",
    "    inputs = [multiConvMol.get_atom_features(), multiConvMol.deg_slice,\n",
    "                np.array(multiConvMol.membership), n_samples]\n",
    "    for i in range(1, len(multiConvMol.get_deg_adjacency_lists())):\n",
    "        inputs.append(multiConvMol.get_deg_adjacency_lists()[i])\n",
    "    return inputs\n",
    "    \n",
    "def get_gc_inputs(inputs):\n",
    "    atom_features = np.expand_dims(inputs[0], axis=0)\n",
    "    degree_slice = np.expand_dims(tf.cast(inputs[1], dtype=tf.int32), axis=0)\n",
    "    membership = np.expand_dims(tf.cast(inputs[2], dtype=tf.int32), axis=0)\n",
    "    n_samples = np.expand_dims(tf.cast(inputs[3], dtype=tf.int32), axis=0)\n",
    "    deg_adjs = [np.expand_dims(tf.cast(deg_adj, dtype=tf.int32), axis=0) for deg_adj in inputs[4:]]\n",
    "\n",
    "    in_layer = atom_features\n",
    "\n",
    "    gc_in = [in_layer, degree_slice, membership, n_samples] + deg_adjs\n",
    "    \n",
    "    return gc_in"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from rlvs.network.graph_layer import GraphConv, GraphGather, GraphPool\n",
    "from tensorflow.keras.layers import Dense, Input, add\n",
    "from tensorflow.keras.models import Model\n",
    "\n",
    "import tensorflow as tf\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _create_molecule_network(jj=0):\n",
    "\n",
    "    features_input = Input(shape=(None, 18,), name=f\"critic_Feature_{jj}\", batch_size=1) \n",
    "    degree_slice_input = Input(shape=(11,2), dtype=tf.int32, name=f\"critic_Degree_slice_{jj}\", batch_size=1)\n",
    "    membership = Input(shape=(None,), dtype=tf.int32, name=f'membership_{jj}', batch_size=1)\n",
    "    n_samples = Input(shape=(1,), dtype=tf.int32, name=f'n_samples_{jj}', batch_size=1)\n",
    "    deg_adjs_input = [Input(shape=(None,None,), dtype=tf.int32, name=f\"critic_deg_adjs_{jj}_{i}\", batch_size=1) for i in  range(10)]\n",
    "\n",
    "    input_states = [features_input, degree_slice_input, membership, n_samples] + deg_adjs_input\n",
    "    graph_layer = GraphConv(out_channel=64, activation_fn=tf.nn.relu)(input_states)\n",
    "\n",
    "    graph_pool_in = [graph_layer, degree_slice_input, membership, n_samples] + deg_adjs_input\n",
    "    graph_pool = GraphPool()(graph_pool_in)\n",
    "    dense_layer = Dense(128, activation=tf.nn.relu)(graph_pool)\n",
    "\n",
    "    return input_states, GraphGather(activation_fn=tf.nn.tanh)([dense_layer, membership, n_samples])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ip_1, graph_gather_layer_1 = _create_molecule_network(1)\n",
    "ip_2, graph_gather_layer_2 = _create_molecule_network(2)\n",
    "mol1_model = Model(inputs=ip_1, outputs=graph_gather_layer_1)\n",
    "mol2_model = Model(inputs=ip_2, outputs=graph_gather_layer_2)\n",
    "\n",
    "combination_layer = add([mol1_model.output, mol2_model.output])\n",
    "combined_dense_layer = Dense(64, activation=tf.nn.relu)(combination_layer)\n",
    "conv_model_1 = Model([ip_1, ip_2], combined_dense_layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ligands = [c.ligand for c in complexes]\n",
    "proteins = [c.protein for c in complexes]\n",
    "\n",
    "protein_batch  = get_gc_inputs(mols_to_inputs(proteins))\n",
    "ligand_batch  = get_gc_inputs(mols_to_inputs(ligands))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_model_1.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_model_1([protein_batch, ligand_batch])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert( (atom_features[0, np.argwhere(membership==1)[:,1], :] == ligands[1].atom_features ).all())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.6.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}