GGNN.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "edges = torch.tensor(([[[0,1,0,0],\n",
    "                        [1,0,1,0],\n",
    "                        [0,1,0,1],\n",
    "                        [0,0,1,0]]]))\n",
    "nodes = torch.tensor(([[[0,0,0,1,1],\n",
    "                        [0,0,1,1,0],\n",
    "                        [0,1,0,0,1],\n",
    "                        [1,0,0,0,1\n",
    "                        ]]]))\n",
    "edges.shape,nodes.shape\n",
    "edges = edges.view(1,4,4,1)\n",
    "#single graph with multiple nodes and edges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((tensor([0, 0, 0, 0, 0, 0]),\n",
       "  tensor([0, 1, 1, 2, 2, 3]),\n",
       "  tensor([1, 0, 2, 1, 3, 2])),\n",
       " (tensor([0, 0, 0, 0]), tensor([0, 1, 2, 3])))"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adjacency.nonzero(as_tuple=True),adjacency.sum(-1).nonzero(as_tuple=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "adj is  torch.Size([1, 4, 4])\n",
      "nodes various ares (tensor([0, 0, 0, 0]), tensor([0, 1, 2, 3]))\n"
     ]
    }
   ],
   "source": [
    "adjacency = torch.sum(edges, dim=3)\n",
    "print(\"adj is \",adjacency.shape)\n",
    "(edge_batch_batch_idc,\n",
    "edge_batch_node_idc,\n",
    "edge_batch_nghb_idc) = adjacency.nonzero(as_tuple=True)\n",
    "(node_batch_batch_idc,\n",
    "node_batch_node_idc) = adjacency.sum(-1).nonzero(as_tuple=True)\n",
    "print(\"nodes various ares\",(node_batch_batch_idc,node_batch_node_idc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ True, False, False, False, False, False],\n",
       "         [False,  True,  True, False, False, False],\n",
       "         [False, False, False,  True,  True, False],\n",
       "         [False, False, False, False, False,  True]]),\n",
       " tensor([[True, True, True, True, True, True],\n",
       "         [True, True, True, True, True, True],\n",
       "         [True, True, True, True, True, True],\n",
       "         [True, True, True, True, True, True]]))"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "same_batch = node_batch_batch_idc.view(-1,1)== edge_batch_batch_idc\n",
    "same_node  = node_batch_node_idc.view(-1, 1) == edge_batch_node_idc\n",
    "message_summation_matrix = (same_batch * same_node).float()\n",
    "same_node,same_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "edge_batch_edges  torch.Size([6, 1]) hello  torch.Size([1, 4, 4, 1])\n"
     ]
    }
   ],
   "source": [
    "edge_batch_edges = edges[edge_batch_batch_idc, edge_batch_node_idc, edge_batch_nghb_idc, :]\n",
    "print(\"edge_batch_edges \",edge_batch_edges.shape,'hello ',edges.shape)\n",
    "hidden_node_features=100\n",
    "hidden_nodes = torch.zeros(nodes.shape[0],\n",
    "                            nodes.shape[1],\n",
    "                            hidden_node_features,\n",
    "                            device='cuda')\n",
    "hidden_nodes[:nodes.shape[0], :nodes.shape[1], :nodes.shape[2]] = nodes.clone()#padding upto 13 nodes,with features to 100....we have only 9 features here\n",
    "node_batch_nodes = hidden_nodes[node_batch_batch_idc, node_batch_node_idc, :]#picking out the same batch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 4, 5])"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nodes.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "in message pass  torch.Size([6, 100]) nodes  torch.Size([6, 100])\n"
     ]
    }
   ],
   "source": [
    "message_passes = 4\n",
    "for _ in range(message_passes):\n",
    "            edge_batch_nodes = hidden_nodes[edge_batch_batch_idc, edge_batch_node_idc, :]#getting hi  13*13*100\n",
    "\n",
    "            edge_batch_nghbs = hidden_nodes[edge_batch_batch_idc, edge_batch_nghb_idc, :]#getting neighs(hj)  13*13*1#why its 1\n",
    "            print(\"in message pass \",edge_batch_nghbs.shape, \"nodes \",edge_batch_nodes.shape)\n",
    "            break\n",
    "            message_terms    = message_terms(edge_batch_nodes,  \n",
    "                                                  edge_batch_nghbs,\n",
    "                                                  edge_batch_edges)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([6, 1])"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "edge_batch_edges.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_batch_edges = edge_batch_edges.view(6,1)#and in bond case it will be something like.... 6,4,1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "The size of tensor a (6) must match the size of tensor b (100) at non-singleton dimension 2",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[1;32m~\\AppData\\Local\\Temp/ipykernel_19488/2125836865.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      4\u001b[0m \u001b[0mouts\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m[\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      5\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 6\u001b[1;33m     \u001b[0mouts\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0medge_batch_edges\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m,\u001b[0m\u001b[0mi\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0mmsg_nns\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mi\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mnode_neighbours_v\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m,\u001b[0m\u001b[0mi\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;31m#here we upper part we pull the edges type bw i,j. it should be a matrix like [1,4] but in our case we use single matrix only\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      7\u001b[0m \u001b[0moutput\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0msum\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mouts\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      8\u001b[0m \u001b[0mmessages\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmatmul\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmessage_summation_matrix\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0moutput\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mRuntimeError\u001b[0m: The size of tensor a (6) must match the size of tensor b (100) at non-singleton dimension 2"
     ]
    }
   ],
   "source": [
    "edges_v               = edge_batch_edges.view(-1,1, 1)\n",
    "node_neighbours_v     = edges_v * edge_batch_nghbs.view(-1,1,100).cpu()\n",
    "node_neighbours_v.shape\n",
    "outs = []\n",
    "for i in range(1):\n",
    "    outs.append(edge_batch_edges[:,i,:]*msg_nns[i](node_neighbours_v[:,i:,])) #here we upper part we pull the edges type bw i,j. it should be a matrix like [1,4] but in our case we use single matrix only\n",
    "output = sum(outs)\n",
    "messages = torch.matmul(message_summation_matrix, output)\n",
    "print(messages.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "def s2v_message_passing(nodes,neighs,edges):\n",
    "    edges_v = edges.view(-1,1,1)#it gives value for edges.....type of edge..size->(6,4)->in molecule with 6 edges kind\n",
    "    neighs = edges_v*neighs.view(-1,1,1)#multiplying each by this number\n",
    "    outs = []\n",
    "    for i in range(1):\n",
    "        outs.append(nodes[:,i,:]*msg_nns[i](edges_v[:,i,:]))\n",
    "    output = sum(outs)\n",
    "    return sum\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([4, 100]), torch.Size([6, 6, 100]), torch.Size([6, 100]))"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "node_batch_nodes.shape,outs[0].shape,edge_batch_nghbs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "ename": "IndexError",
     "evalue": "too many indices for tensor of dimension 2",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[1;32m~\\AppData\\Local\\Temp/ipykernel_19488/3789738656.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[1;32m----> 1\u001b[1;33m \u001b[0mnode_batch_nodes\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;33m:\u001b[0m\u001b[0mmessages\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m:\u001b[0m\u001b[0mmessages\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mmessages\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;31mIndexError\u001b[0m: too many indices for tensor of dimension 2"
     ]
    }
   ],
   "source": [
    "node_batch_nodes[:messages.shape[0],:messages.shape[1],:] = messages\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([4, 6]), torch.Size([6, 6, 100]))"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "message_summation_matrix.shape,output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([6, 4, 100])\n"
     ]
    }
   ],
   "source": [
    "messages = torch.matmul(message_summation_matrix, output)\n",
    "print(messages.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([6, 6, 100])"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([6, 100])"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "edge_batch_nghbs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([6, 1, 100])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "edge_batch_nghbs.view(-1,1,100).cpu().shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "msg_nns = []\n",
    "for _ in range(4):#number of edge features\n",
    "            msg_nns.append(\n",
    "                MLP(\n",
    "                    in_features=100,\n",
    "                    hidden_layer_sizes=[250] * 4,\n",
    "                    out_features=100,#message size\n",
    "                    dropout_p=0,\n",
    "                )\n",
    "            )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from collections import namedtuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class MLP(torch.nn.Module):\n",
    "    def __init__(self, in_features : int, hidden_layer_sizes : list, out_features : int,\n",
    "                 dropout_p : float) -> None:\n",
    "        super().__init__()\n",
    "\n",
    "        activation_function = torch.nn.SELU\n",
    "        # create list of all layer feature sizes\n",
    "        fs = [in_features, *hidden_layer_sizes, out_features]\n",
    "        # create list of linear_blocks\n",
    "        layers = [self._linear_block(in_f, out_f,\n",
    "                                     activation_function,\n",
    "                                     dropout_p)\n",
    "                  for in_f, out_f in zip(fs, fs[1:])]\n",
    "        # concatenate modules in all sequentials in layers list\n",
    "        layers = [module for sq in layers for module in sq.children()]\n",
    "\n",
    "        # add modules to sequential container\n",
    "        self.seq = torch.nn.Sequential(*layers)\n",
    "\n",
    "    def _linear_block(self, in_f : int, out_f : int, activation : torch.nn.Module,\n",
    "                      dropout_p : float) -> torch.nn.Sequential:\n",
    "        \n",
    "        # bias must be used in most MLPs in our models to learn from empty graphs\n",
    "        linear = torch.nn.Linear(in_f, out_f, bias=True)\n",
    "        torch.nn.init.xavier_uniform_(linear.weight)\n",
    "        return torch.nn.Sequential(linear, activation(), torch.nn.AlphaDropout(dropout_p))\n",
    "\n",
    "    def forward(self, layers_input : torch.nn.Sequential) -> torch.nn.Sequential:\n",
    "        return self.seq(layers_input)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "hyperparameters = {\n",
    "        \"enn_depth\"            : 4,\n",
    "        \"enn_dropout_p\"        : 0.0,\n",
    "        \"enn_hidden_dim\"       : 250,\n",
    "        \"mlp1_depth\"           : 4,\n",
    "        \"mlp1_dropout_p\"       : 0.0,\n",
    "        \"mlp1_hidden_dim\"      : 500,\n",
    "        \"mlp2_depth\"           : 4,\n",
    "        \"mlp2_dropout_p\"       : 0.0,\n",
    "        \"mlp2_hidden_dim\"      : 500,\n",
    "        \"gather_att_depth\"     : 4,\n",
    "        \"gather_att_dropout_p\" : 0.0,\n",
    "        \"gather_att_hidden_dim\": 250,\n",
    "        \"gather_emb_depth\"     : 4,\n",
    "        \"gather_emb_dropout_p\" : 0.0,\n",
    "        \"gather_emb_hidden_dim\": 250,\n",
    "        \"gather_width\"         : 100,\n",
    "        \"hidden_node_features\" : 100,\n",
    "        \"message_passes\"       : 3,\n",
    "        \"message_size\"         : 100,\n",
    "    }\n",
    "import json\n",
    "class dotdict(dict):\n",
    "    \"\"\"dot.notation access to dictionary attributes\"\"\"\n",
    "    __getattr__ = dict.get\n",
    "    __setattr__ = dict.__setitem__\n",
    "    __delattr__ = dict.__delitem__\n",
    "constants = dotdict(hyperparameters)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "<rdkit.Chem.rdmolfiles.SmilesMolSupplier object at 0x0000021A8DC75CC0>\n",
      "<ashishcode.MolecularGraph object at 0x0000021AAC78EEB0>\n"
     ]
    }
   ],
   "source": [
    "from ashishcode import load_molecule,MolecularGraph,params\n",
    "path = \"data\\pre-training\\gdb13_1K\\Train.smi\"\n",
    "molecule_set = load_molecule(path)\n",
    "print(molecule_set)\n",
    "for mol in molecule_set:\n",
    "    molecule = MolecularGraph(mol)\n",
    "    print(molecule)\n",
    "    break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class SummationMPNN(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    Abstract `SummationMPNN` class. Specific models using this class are\n",
    "    defined in `mpnn.py`; these are MNN, S2V, and GGNN.\n",
    "    \"\"\"\n",
    "    def __init__(self, constants : namedtuple):\n",
    "\n",
    "        super().__init__()\n",
    "\n",
    "        self.hidden_node_features = constants.hidden_node_features\n",
    "        self.edge_features        = constants.n_edge_features\n",
    "        self.message_size         = constants.message_size\n",
    "        self.message_passes       = constants.message_passes\n",
    "        self.constants            = constants\n",
    "\n",
    "    def forward(self, nodes : torch.Tensor, edges : torch.Tensor) -> None:\n",
    "        adjacency = torch.sum(edges, dim=3)\n",
    "        # **note: \"idc\" == \"indices\", \"nghb{s}\" == \"neighbour(s)\"\n",
    "        (edge_batch_batch_idc,\n",
    "         edge_batch_node_idc,\n",
    "         edge_batch_nghb_idc) = adjacency.nonzero(as_tuple=True)\n",
    "        #print(\"sizes are edge_batch_node_idc \",edge_batch_node_idc)\n",
    "\n",
    "        (node_batch_batch_idc, node_batch_node_idc) = adjacency.sum(-1).nonzero(as_tuple=True)\n",
    "\n",
    "        same_batch = node_batch_batch_idc.view(-1, 1) == edge_batch_batch_idc\n",
    "        same_node  = node_batch_node_idc.view(-1, 1) == edge_batch_node_idc\n",
    "\n",
    "        # element ij of `message_summation_matrix` is 1 if `edge_batch_edges[j]`\n",
    "        # is connected with `node_batch_nodes[i]`, else 0\n",
    "        message_summation_matrix = (same_batch * same_node).float()\n",
    "\n",
    "        edge_batch_edges = edges[edge_batch_batch_idc, edge_batch_node_idc, edge_batch_nghb_idc, :]\n",
    "\n",
    "        # pad up the hidden nodes\n",
    "        hidden_nodes = torch.zeros(nodes.shape[0],\n",
    "                                   nodes.shape[1],\n",
    "                                   self.hidden_node_features,\n",
    "                                   device=self.constants.device)\n",
    "        hidden_nodes[:nodes.shape[0], :nodes.shape[1], :nodes.shape[2]] = nodes.clone()\n",
    "        hidden_nodes_original = hidden_nodes.clone()\n",
    "        node_batch_nodes = hidden_nodes[node_batch_batch_idc, node_batch_node_idc, :]\n",
    "        \n",
    "\n",
    "        for _ in range(self.message_passes):\n",
    "            edge_batch_nodes = hidden_nodes[edge_batch_batch_idc, edge_batch_node_idc, :]\n",
    "\n",
    "            edge_batch_nghbs = hidden_nodes[edge_batch_batch_idc, edge_batch_nghb_idc, :]\n",
    "\n",
    "            #print(\"hello ji \",edge_batch_nghbs.shape,edge_batch_nodes.shape,hidden_nodes.shape)\n",
    "\n",
    "            message_terms = self.message_terms(edge_batch_nodes,\n",
    "                                                  edge_batch_nghbs,\n",
    "                                                  edge_batch_edges)\n",
    "\n",
    "            if len(message_terms.size()) == 1:  # if a single graph in batch\n",
    "                message_terms = message_terms.unsqueeze(0)\n",
    "\n",
    "            # the summation in eq. 1 of the NMPQC paper happens here\n",
    "            messages = torch.matmul(message_summation_matrix, message_terms)\n",
    "\n",
    "            node_batch_nodes = self.update(node_batch_nodes, messages)\n",
    "            hidden_nodes[node_batch_batch_idc, node_batch_node_idc, :] = node_batch_nodes.clone() #updated the hidden states\n",
    "        node_mask = adjacency.sum(-1) != 0\n",
    "        output    = self.readout(hidden_nodes, hidden_nodes_original, node_mask)\n",
    "        return output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from graphinvent.gnn.modules import GlobalReadout \n",
    "class GGNN(SummationMPNN):\n",
    "    def __init__(self,constants) -> None:\n",
    "        super().__init__(constants)\n",
    "        self.constants       = constants\n",
    "        message_weights      = torch.Tensor(self.constants.message_size,\n",
    "                                            self.constants.hidden_node_features,\n",
    "                                            4)#edge features\n",
    "        if False:#\"cuda\" == \"cuda\":\n",
    "            message_weights = message_weights.to(\"cuda\", non_blocking=True)\n",
    "        \n",
    "    \n",
    "        self.message_weights = torch.nn.Parameter(message_weights)\n",
    "\n",
    "        self.gru             = torch.nn.GRUCell(\n",
    "            input_size=self.constants.message_size,\n",
    "            hidden_size=self.constants.hidden_node_features,\n",
    "            bias=True\n",
    "        )\n",
    "        \n",
    "        self.msg_nns = []\n",
    "        for _ in range(4):#number of edge features\n",
    "                    self.msg_nns.append(\n",
    "                        MLP(\n",
    "                            in_features=100,\n",
    "                            hidden_layer_sizes=[250] * 4,\n",
    "                            out_features=100,#message size\n",
    "                            dropout_p=0,\n",
    "                        )\n",
    "                    )\n",
    "\n",
    "        self.gather = GraphGather(node_features=100, hidden_node_features=100, out_features=100, \n",
    "                                    att_depth = constants.gather_att_depth,\n",
    "                                    att_hidden_dim=constants.gather_att_hidden_dim,\n",
    "                                    att_dropout_p=0.4,\n",
    "                                    emb_depth=constants.gather_emb_depth, emb_hidden_dim=constants.gather_emb_hidden_dim,\n",
    "                                  emb_dropout_p=0.0, big_positive=0.1)\n",
    "        \n",
    "        self.APDReadout = GlobalReadout()\n",
    "\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self) -> None:\n",
    "        import math\n",
    "        stdev = 1.0 / math.sqrt(self.message_weights.size(1))\n",
    "        self.message_weights.data.uniform_(-stdev, stdev)\n",
    "\n",
    "    def message_terms(self, nodes : torch.Tensor, node_neighbours : torch.Tensor,\n",
    "                        edges : torch.Tensor) -> torch.Tensor:\n",
    "        \n",
    "        edges_v = edges.view(-1,4,1)#it gives value for edges.....type of edge..size->(6,4)->in molecule with 6 edges kind\n",
    "        outs = []\n",
    "        for i in range(4):\n",
    "            # print(\"in last loop \",edges_v[:,i,:].shape,node_neighbours.shape)\n",
    "            outs.append(edges_v[:,i,:]*self.msg_nns[i](node_neighbours))#check node neigh element wise help\n",
    "        output = sum(outs)\n",
    "        return output\n",
    "\n",
    "        \n",
    "    #torch.broadcast_to(x, (3, 3))\n",
    "    def update(self, nodes : torch.Tensor, messages : torch.Tensor) -> torch.Tensor:\n",
    "        return self.gru(messages, nodes)\n",
    "\n",
    "    def readout(self, hidden_nodes : torch.Tensor, input_nodes : torch.Tensor,\n",
    "                node_mask : torch.Tensor) -> torch.Tensor:\n",
    "        # graph_embeddings = torch.sum(hidden_nodes, dim=1)\n",
    "        # print(\"input to network \",hidden_nodes.shape,input_nodes.shape)\n",
    "        output_gather =  self.gather(input_nodes,hidden_nodes,node_mask)\n",
    "        output        = self.APDReadout(hidden_nodes,output_gather)\n",
    "        return output\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape of feat in global readout is  torch.Size([1, 10, 100])\n",
      "output of final net is , torch.Size([1, 23, 100])\n"
     ]
    }
   ],
   "source": [
    "network = GGNN(constants)\n",
    "node,edge = molecule.get_graph_state()\n",
    "nodes,edges = torch.Tensor(node).view((1,10,9)),torch.Tensor(edge).view((1,10,10,4))\n",
    "print(\"output of final net is ,\",network(nodes,edges).shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class GraphGather(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    GGNN readout function.\n",
    "    \"\"\"\n",
    "    def __init__(self, node_features : int, hidden_node_features : int,\n",
    "                 out_features : int, att_depth : int, att_hidden_dim : int,\n",
    "                 att_dropout_p : float, emb_depth : int, emb_hidden_dim : int,\n",
    "                 emb_dropout_p : float, big_positive : float) -> None:\n",
    "\n",
    "        super().__init__()\n",
    "\n",
    "        self.big_positive = big_positive\n",
    "\n",
    "        self.att_nn = MLP(\n",
    "            in_features=node_features + hidden_node_features,\n",
    "            hidden_layer_sizes=[att_hidden_dim] * att_depth,\n",
    "            out_features=out_features,\n",
    "            dropout_p=att_dropout_p\n",
    "        )\n",
    "\n",
    "        self.emb_nn = MLP(\n",
    "            in_features=hidden_node_features,\n",
    "            hidden_layer_sizes=[emb_hidden_dim] * emb_depth,\n",
    "            out_features=out_features,\n",
    "            dropout_p=emb_dropout_p\n",
    "        )\n",
    "\n",
    "    def forward(self, hidden_nodes : torch.Tensor, input_nodes : torch.Tensor,\n",
    "                node_mask : torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Defines forward pass.\n",
    "        \"\"\"\n",
    "        Softmax     = torch.nn.Softmax(dim=1)\n",
    "\n",
    "        cat         = torch.cat((hidden_nodes, input_nodes), dim=2)\n",
    "        #energy_mask = (node_mask == 0).float() * self.big_positive\n",
    "        #energies    = self.att_nn(cat) - energy_mask.unsqueeze(-1)\n",
    "        #attention   = Softmax(energies)\n",
    "        attention = self.att_nn(cat)#activation is SELU-->done automatically inside\n",
    "        embedding   = self.emb_nn(hidden_nodes)#this is mlp-a-->\n",
    "\n",
    "        return torch.sum(attention * embedding, dim=1)\n",
    "class GlobalReadout(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.mlp1 = MLP(in_features=constants.message_size,\n",
    "                  hidden_layer_sizes=[constants.mlp1_hidden_dim]*constants.mlp1_depth,\n",
    "                  out_features=constants.message_size,\n",
    "                  dropout_p=0.0)\n",
    "        self.mlp2 = MLP(in_features=constants.message_size,\n",
    "                  hidden_layer_sizes=[constants.mlp2_hidden_dim]*constants.mlp2_depth,\n",
    "                  out_features=constants.message_size,\n",
    "                  dropout_p=0.0)\n",
    "        \n",
    "        self.mlp3 = MLP(in_features=constants.message_size,\n",
    "                  hidden_layer_sizes=[constants.mlp1_hidden_dim]*constants.mlp1_depth,\n",
    "                  out_features=constants.message_size,\n",
    "                  dropout_p=0.0)\n",
    "        self.mlp4 = MLP(in_features=2*constants.message_size,\n",
    "                  hidden_layer_sizes=[constants.mlp2_hidden_dim]*constants.mlp2_depth,\n",
    "                  out_features=constants.message_size,\n",
    "                  dropout_p=0.0)\n",
    "        self.mlpt = MLP(in_features=constants.message_size,\n",
    "                  hidden_layer_sizes=[constants.mlp1_hidden_dim]*constants.mlp1_depth,\n",
    "                  out_features=constants.message_size,\n",
    "                  dropout_p=0.0)\n",
    "          \n",
    "    def forward(self,features,g):\n",
    "        print(\"shape of feat in global readout is \",features.shape)\n",
    "        #g= torch.sum(features,dim=1)\n",
    "        # print(\"graph embedding shapes \",g.shape)\n",
    "        g = g.view(1,1,100)\n",
    "        #g = torch.broadcast_to(g, (1,10, 100))\n",
    "        # print(\"api is \",g.shape)\n",
    "        fadd1 = self.mlp1(features)\n",
    "        fconn1 = self.mlp2(features)  \n",
    "\n",
    "        # print(\"dims fadd ip\",fconn1.shape,g.shape,torch.cat([fadd1,g],dim=1).shape)\n",
    "        fadd = self.mlp3(torch.cat([fadd1,g],dim=1)).unsqueeze(dim=1)\n",
    "        fconn = self.mlp3(torch.cat([fconn1,g],dim=1)).unsqueeze(dim=1)\n",
    "\n",
    "        fterm = self.mlpt(g)\n",
    "        cat = torch.cat((fadd.squeeze(dim=1), fconn.squeeze(dim=1), fterm), dim=1)\n",
    "        # print(\"final shape \",cat.shape)\n",
    "        return cat\n",
    "        #apd = self.Softmax()....from original code its removed\n",
    "\n",
    "    \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "5a7aeb117152383494077b830ed8f2bcff9728640e9e954d18ee6388f442a456"
  },
  "kernelspec": {
   "display_name": "Python 3.9.9 64-bit (windows store)",
   "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.9.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}