pred_complex.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "     code ch1 ch2  nres  len1  len2 dom_s1      dom_p1 dom_s2      dom_p2  \\\n",
      "0  4rk1_1   A   B  31.0   267   268  53822   PF13377.1  53822   PF13377.1   \n",
      "1  2gqn_1   D   B  53.0   392   392  53383  PF01053.15  53383  PF01053.15   \n",
      "2  2gqn_1   D   C  17.0   392   391  53383  PF01053.15  53383  PF01053.15   \n",
      "3  2gqn_1   A   C  52.0   391   391  53383  PF01053.15  53383  PF01053.15   \n",
      "4  2gqn_1   A   D  28.5   391   392  53383  PF01053.15  53383  PF01053.15   \n",
      "\n",
      "   ident  homo  \n",
      "0      1   1.0  \n",
      "1      1   1.0  \n",
      "2      1   1.0  \n",
      "3      1   1.0  \n",
      "4      1   1.0  \n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import networkx as nx\n",
    "#get known complexes\n",
    "datadir = '/home/dmoi/datasets/foldtree2/complexes/'\n",
    "complexdf = pd.read_csv(datadir+'contactDefinition.txt' , sep='\\t')\n",
    "print(complexdf.head())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "173593\n"
     ]
    }
   ],
   "source": [
    "import glob\n",
    "pdbs = glob.glob(datadir+'BU_all_renum/*.pdb')\n",
    "print(len(pdbs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "     code ch1 ch2  nres  len1  len2 dom_s1      dom_p1 dom_s2      dom_p2  \\\n",
      "0  4rk1_1   A   B  31.0   267   268  53822   PF13377.1  53822   PF13377.1   \n",
      "1  2gqn_1   D   B  53.0   392   392  53383  PF01053.15  53383  PF01053.15   \n",
      "2  2gqn_1   D   C  17.0   392   391  53383  PF01053.15  53383  PF01053.15   \n",
      "3  2gqn_1   A   C  52.0   391   391  53383  PF01053.15  53383  PF01053.15   \n",
      "4  2gqn_1   A   D  28.5   391   392  53383  PF01053.15  53383  PF01053.15   \n",
      "\n",
      "   ident  homo                                            pdbfile  \n",
      "0      1   1.0  /home/dmoi/datasets/foldtree2/complexes/BU_all...  \n",
      "1      1   1.0  /home/dmoi/datasets/foldtree2/complexes/BU_all...  \n",
      "2      1   1.0  /home/dmoi/datasets/foldtree2/complexes/BU_all...  \n",
      "3      1   1.0  /home/dmoi/datasets/foldtree2/complexes/BU_all...  \n",
      "4      1   1.0  /home/dmoi/datasets/foldtree2/complexes/BU_all...  \n"
     ]
    }
   ],
   "source": [
    "codes = { f.split('/')[-1].split('.')[0]:f for f in pdbs}\n",
    "complexdf['pdbfile'] = complexdf['code'].map(codes)\n",
    "print(complexdf.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "       code ch1 ch2  nres  len1  len2 dom_s1      dom_p1 dom_s2      dom_p2  \\\n",
      "292  1qxe_2   D   G   4.5   146   141  46458  PF00042.17  46458  PF00042.17   \n",
      "293  1qxe_2   E   G  10.0   141   141  46458  PF00042.17  46458  PF00042.17   \n",
      "294  1qxe_2   E   H   8.5   141   146  46458  PF00042.17  46458  PF00042.17   \n",
      "295  1qxe_2   B   C   8.5   146   141  46458  PF00042.17  46458  PF00042.17   \n",
      "296  1qxe_2   B   D   3.0   146   146  46458  PF00042.17  46458  PF00042.17   \n",
      "\n",
      "     ident  homo                                            pdbfile  \n",
      "292      0   1.0  /home/dmoi/datasets/foldtree2/complexes/BU_all...  \n",
      "293      1   1.0  /home/dmoi/datasets/foldtree2/complexes/BU_all...  \n",
      "294      0   1.0  /home/dmoi/datasets/foldtree2/complexes/BU_all...  \n",
      "295      0   1.0  /home/dmoi/datasets/foldtree2/complexes/BU_all...  \n",
      "296      1   1.0  /home/dmoi/datasets/foldtree2/complexes/BU_all...  \n",
      "8133\n"
     ]
    }
   ],
   "source": [
    "#sample 2000 codes\n",
    "import random\n",
    "sub = complexdf.code.unique().tolist()\n",
    "random.shuffle(sub)\n",
    "sub = sub[:2000]\n",
    "sub = complexdf[complexdf.code.isin(sub)]\n",
    "print(sub.head())\n",
    "print( len(sub) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "class ComplexDataset(Dataset):\n",
    "    def __init__(self, h5dataset  ):\n",
    "        super().__init__()\n",
    "        #keys should be the structures\n",
    "\n",
    "\n",
    "        self.h5dataset = h5dataset\n",
    "\n",
    "        if type(h5dataset) == str:\n",
    "            self.h5dataset = h5py.File(h5dataset, 'r')\n",
    "        self.structlist = list(self.h5dataset['structs'].keys())\n",
    "    def __len__(self):\n",
    "        return len(self.structlist)\n",
    "\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if type(idx) == str:\n",
    "            f = self.h5dataset['structs'][idx]\n",
    "        elif type(idx) == int:\n",
    "            f = self.h5dataset['structs'][self.structlist[idx]]\n",
    "        else:\n",
    "            raise 'use a structure filename or integer'\n",
    "        data = {}\n",
    "        hetero_data = HeteroData()\n",
    "        if 'node' in f.keys():\n",
    "            for node_type in f['node'].keys():\n",
    "                node_group = f['node'][node_type]\n",
    "                # Assuming 'x' exists\n",
    "                if 'x' in node_group.keys():\n",
    "                    hetero_data[node_type].x = torch.tensor(node_group['x'][:])\n",
    "        # Edge data\n",
    "        if 'edge' in f.keys():\n",
    "            for edge_name in f['edge'].keys():\n",
    "                edge_group = f['edge'][edge_name]\n",
    "                src_type, link_type, dst_type = edge_name.split('_')\n",
    "                edge_type = (src_type, link_type, dst_type)\n",
    "                # Assuming 'edge_index' exists\n",
    "                if 'edge_index' in edge_group.keys():\n",
    "                    hetero_data[edge_type].edge_index = torch.tensor(edge_group['edge_index'][:])\n",
    "                \n",
    "                # If there are edge attributes, load them too\n",
    "                if 'edge_attr' in edge_group.keys():\n",
    "                    hetero_data[edge_type].edge_attr = torch.tensor(edge_group['edge_attr'][:])\n",
    "        #return pytorch geometric heterograph\n",
    "        return hetero_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#iterate over the complexes\n",
    "\n",
    "#threshold for contact\n",
    "thresh = 10\n",
    "\n",
    "for code in tqdm.tqdm(sub.code.unique()):\n",
    "    \n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|█████████████████████████████▊                                            | 36180/89764 [10:45<15:56, 56.02it/s]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[6], line 6\u001b[0m\n\u001b[1;32m      4\u001b[0m graphs \u001b[38;5;241m=\u001b[39m {}\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m code \u001b[38;5;129;01min\u001b[39;00m tqdm\u001b[38;5;241m.\u001b[39mtqdm(complexdf\u001b[38;5;241m.\u001b[39mcode\u001b[38;5;241m.\u001b[39munique()):\n\u001b[0;32m----> 6\u001b[0m     sub \u001b[38;5;241m=\u001b[39m complexdf[\u001b[43mcomplexdf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcode\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m==\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mcode\u001b[49m]\n\u001b[1;32m      7\u001b[0m     chains \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mset\u001b[39m(sub\u001b[38;5;241m.\u001b[39mch1\u001b[38;5;241m.\u001b[39munique())\u001b[38;5;241m.\u001b[39munion( \u001b[38;5;28mset\u001b[39m( sub\u001b[38;5;241m.\u001b[39mch2\u001b[38;5;241m.\u001b[39munique() ))\n\u001b[1;32m      8\u001b[0m     \u001b[38;5;66;03m#add chains to graph\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/tf_gpu/lib/python3.8/site-packages/pandas/core/ops/common.py:81\u001b[0m, in \u001b[0;36m_unpack_zerodim_and_defer.<locals>.new_method\u001b[0;34m(self, other)\u001b[0m\n\u001b[1;32m     77\u001b[0m             \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mNotImplemented\u001b[39m\n\u001b[1;32m     79\u001b[0m other \u001b[38;5;241m=\u001b[39m item_from_zerodim(other)\n\u001b[0;32m---> 81\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mmethod\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mother\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/tf_gpu/lib/python3.8/site-packages/pandas/core/arraylike.py:40\u001b[0m, in \u001b[0;36mOpsMixin.__eq__\u001b[0;34m(self, other)\u001b[0m\n\u001b[1;32m     38\u001b[0m \u001b[38;5;129m@unpack_zerodim_and_defer\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__eq__\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     39\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__eq__\u001b[39m(\u001b[38;5;28mself\u001b[39m, other):\n\u001b[0;32m---> 40\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_cmp_method\u001b[49m\u001b[43m(\u001b[49m\u001b[43mother\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moperator\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43meq\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/tf_gpu/lib/python3.8/site-packages/pandas/core/series.py:6096\u001b[0m, in \u001b[0;36mSeries._cmp_method\u001b[0;34m(self, other, op)\u001b[0m\n\u001b[1;32m   6093\u001b[0m rvalues \u001b[38;5;241m=\u001b[39m extract_array(other, extract_numpy\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m, extract_range\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m   6095\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m np\u001b[38;5;241m.\u001b[39merrstate(\u001b[38;5;28mall\u001b[39m\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mignore\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m-> 6096\u001b[0m     res_values \u001b[38;5;241m=\u001b[39m \u001b[43mops\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcomparison_op\u001b[49m\u001b[43m(\u001b[49m\u001b[43mlvalues\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mrvalues\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mop\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   6098\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_construct_result(res_values, name\u001b[38;5;241m=\u001b[39mres_name)\n",
      "File \u001b[0;32m~/miniconda3/envs/tf_gpu/lib/python3.8/site-packages/pandas/core/ops/array_ops.py:293\u001b[0m, in \u001b[0;36mcomparison_op\u001b[0;34m(left, right, op)\u001b[0m\n\u001b[1;32m    290\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m invalid_comparison(lvalues, rvalues, op)\n\u001b[1;32m    292\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m is_object_dtype(lvalues\u001b[38;5;241m.\u001b[39mdtype) \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(rvalues, \u001b[38;5;28mstr\u001b[39m):\n\u001b[0;32m--> 293\u001b[0m     res_values \u001b[38;5;241m=\u001b[39m \u001b[43mcomp_method_OBJECT_ARRAY\u001b[49m\u001b[43m(\u001b[49m\u001b[43mop\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlvalues\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mrvalues\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    295\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    296\u001b[0m     res_values \u001b[38;5;241m=\u001b[39m _na_arithmetic_op(lvalues, rvalues, op, is_cmp\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n",
      "File \u001b[0;32m~/miniconda3/envs/tf_gpu/lib/python3.8/site-packages/pandas/core/ops/array_ops.py:82\u001b[0m, in \u001b[0;36mcomp_method_OBJECT_ARRAY\u001b[0;34m(op, x, y)\u001b[0m\n\u001b[1;32m     80\u001b[0m     result \u001b[38;5;241m=\u001b[39m libops\u001b[38;5;241m.\u001b[39mvec_compare(x\u001b[38;5;241m.\u001b[39mravel(), y\u001b[38;5;241m.\u001b[39mravel(), op)\n\u001b[1;32m     81\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m---> 82\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[43mlibops\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mscalar_compare\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mravel\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43my\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mop\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     83\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\u001b[38;5;241m.\u001b[39mreshape(x\u001b[38;5;241m.\u001b[39mshape)\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import tqdm\n",
    "import pickle\n",
    "chunksize = 10000\n",
    "graphs = {}\n",
    "for code in tqdm.tqdm(complexdf.code.unique()):\n",
    "    sub = complexdf[complexdf.code == code]\n",
    "    chains = set(sub.ch1.unique()).union( set( sub.ch2.unique() ))\n",
    "    #add chains to graph\n",
    "    G = nx.MultiGraph()\n",
    "    G.add_nodes_from(chains)\n",
    "    for i, row in sub.iterrows():\n",
    "        #homology link\n",
    "        if row.homo == 1:\n",
    "            #add homology type edge\n",
    "            ekey = G.add_edge(row.ch1, row.ch2, key='homology' )\n",
    "        #add contact type edge\n",
    "        ekey = G.add_edge(row.ch1, row.ch2, key='contact' )\n",
    "    graphs[code] = G\n",
    "    \n",
    "    \n",
    "\n",
    "with open('complexgraphs.pkl' , 'wb') as graphsout:\n",
    "    graphsout.write( pickle.dumps( graphs ) )\n",
    "\n",
    "    \n",
    "#embed all structures\n",
    "#use decoder and generate z vecs\n",
    "#use decoder sigmoid to get contact proba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.data import HeteroData\n",
    "#import sageconv\n",
    "from torch_geometric.nn import SAGEConv , Linear , FiLMConv , TransformerConv , FeaStConv , GATConv , GINConv , GatedGraphConv\n",
    "#import module dict and module list\n",
    "from torch.nn import ModuleDict, ModuleList , L1Loss\n",
    "from torch_geometric.nn import global_mean_pool\n",
    "#import negative sampling\n",
    "from torch_geometric.utils import negative_sampling\n",
    "\n",
    "\n",
    "EPS = 1e-10\n",
    "def recon_loss( z1: Tensor,z2: Tensor, pos_edge_index: Tensor , backbone:Tensor = None , decoder = None , poslossmod = 1 , neglossmod= 1) -> Tensor:\n",
    "    r\"\"\"Given latent variables :obj:`z`, computes the binary cross\n",
    "    entropy loss for positive edges :obj:`pos_edge_index` and negative\n",
    "    sampled edges.\n",
    "\n",
    "    Args:\n",
    "        z (torch.Tensor): The latent space :math:`\\mathbf{Z}`.\n",
    "        pos_edge_index (torch.Tensor): The positive edges to train against.\n",
    "        neg_edge_index (torch.Tensor, optional): The negative edges to\n",
    "            train against. If not given, uses negative sampling to\n",
    "            calculate negative edges. (default: :obj:`None`)\n",
    "    \"\"\"\n",
    "    \n",
    "    pos =decoder(z1,z2, pos_edge_index, { ( 'res','backbone','res'): backbone } )[1]\n",
    "    #turn pos edge index into a binary matrix\n",
    "    pos_loss = -torch.log( pos + EPS).mean()\n",
    "    neg_edge_index = negative_sampling(pos_edge_index, z.size(0))\n",
    "    neg = decoder(z1 , z2, neg_edge_index, { ( 'res','backbone','res'): backbone } )[1]\n",
    "    neg_loss = -torch.log( ( 1 - neg) + EPS ).mean()\n",
    "    return poslossmod*pos_loss + neglossmod*neg_loss\n",
    "\n",
    "#define loss for x reconstruction   \n",
    "def x_reconstruction_loss(x, recon_x):\n",
    "    \"\"\"\n",
    "    compute the loss over the node feature reconstruction.\n",
    "    \"\"\"\n",
    "    return F.l1_loss(recon_x, x)\n",
    "\n",
    "\n",
    "#amino acid onehot loss for x reconstruction\n",
    "def aa_reconstruction_loss(x, recon_x):\n",
    "    \"\"\"\n",
    "    compute the loss over the node feature reconstruction.\n",
    "    using categorical cross entropy\n",
    "    \"\"\"\n",
    "    x = torch.argmax(x, dim=1)\n",
    "    #recon_x = torch.argmax(recon_x, dim=1)\n",
    "    return F.cross_entropy(recon_x, x)\n",
    "\n",
    "def gaussian_loss(mu , logvar , beta= 1.5):\n",
    "    '''\n",
    "    \n",
    "    add beta to disentangle the features\n",
    "    \n",
    "    '''\n",
    "    kl_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())\n",
    "    return beta*kl_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class HeteroGAE_Pairwise_Decoder(torch.nn.Module):\n",
    "    def __init__(self, encoder_out_channels, xdim=20, hidden_channels={'res_backbone_res': [20, 20, 20]}, out_channels_hidden=20, nheads = 1 , Xdecoder_hidden=30, metadata={}, amino_mapper= None):\n",
    "        super(HeteroGAE_Decoder, self).__init__()\n",
    "        self.convs = torch.nn.ModuleList()\n",
    "        self.metadata = metadata\n",
    "        self.hidden_channels = hidden_channels\n",
    "        self.out_channels_hidden = out_channels_hidden\n",
    "        self.in_channels = encoder_out_channels\n",
    "        self.amino_acid_indices = amino_mapper\n",
    "        for i in range(len(self.hidden_channels[('res', 'backbone', 'res')])):\n",
    "            self.convs.append(\n",
    "                torch.nn.ModuleDict({\n",
    "                    '_'.join(edge_type): SAGEConv(self.in_channels if i == 0 else self.hidden_channels[edge_type][i-1], self.hidden_channels[edge_type][i]  )\n",
    "                    for edge_type in [('res', 'backbone', 'res')]\n",
    "                })\n",
    "            )\n",
    "\n",
    "        self.lin = Linear(hidden_channels[('res', 'backbone', 'res')][-1], self.out_channels_hidden)\n",
    "        \n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "\n",
    "        self.decoder = torch.nn.Sequential(\n",
    "            torch.nn.Linear( self.in_channels + self.out_channels_hidden , Xdecoder_hidden),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(Xdecoder_hidden, Xdecoder_hidden),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(Xdecoder_hidden, Xdecoder_hidden),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(Xdecoder_hidden, Xdecoder_hidden),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(Xdecoder_hidden, xdim),\n",
    "            torch.nn.LogSoftmax(dim=1)         )\n",
    "        \n",
    "\n",
    "    def forward(self, z1, z2, , edge_index, backbones, **kwargs):\n",
    "        #copy z for later concatenation\n",
    "        xrs = []\n",
    "        zs = []\n",
    "        for z in [z1,z2]:\n",
    "            inz = z\n",
    "            for layer in self.convs:\n",
    "                for edge_type, conv in layer.items():\n",
    "                    z = conv(z, backbones[tuple(edge_type.split('_'))])\n",
    "                    z = F.relu(z)\n",
    "            z = self.lin(z)\n",
    "            x_r = self.decoder(  torch.cat( [inz,  z] , axis = 1) )\n",
    "            xrs.append(x_r)\n",
    "            zs.append(z)\n",
    "\n",
    "        sim_matrix = (zs[0][edge_index[0]] * zs[1][edge_index[1]]).sum(dim=1)\n",
    "        edge_probs = self.sigmoid(sim_matrix)\n",
    "        #plddt_r = self.plddt_decoder(z)\n",
    "        #return x_r, plddt_r,  edge_probs\n",
    "        return x_r[0],x_r[1]  edge_probs\n",
    "\n",
    "    def x_to_amino_acid_sequence(self, x_r):\n",
    "        \"\"\"\n",
    "        Converts the reconstructed 20-dimensional matrix to a sequence of amino acids.\n",
    "\n",
    "        Args:\n",
    "            x_r (Tensor): Reconstructed 20-dimensional tensor.\n",
    "\n",
    "        Returns:\n",
    "            str: A string representing the sequence of amino acids.\n",
    "        \"\"\"\n",
    "        # Find the index of the maximum value in each row to get the predicted amino acid\n",
    "        indices = torch.argmax(x_r, dim=1)\n",
    "        \n",
    "        # Convert indices to amino acids\n",
    "        amino_acid_sequence = ''.join(self.amino_mapper[idx.item()] for idx in indices)\n",
    "        \n",
    "        return amino_acid_sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load pretrained encoder\n",
    "encoder = torch.load(\"encoder.pt\")\n",
    "encoder.eval()\n",
    "\n",
    "#create new decoder\n",
    "decoder = HeteroGAE_Pairwise_Decoder(encoder_out_channels = encoder.out_channels , \n",
    "                            hidden_channels={ ( 'res','backbone','res'):[ 300 ] * 5  } , \n",
    "                            out_channels_hidden= 150 , metadata=metadata , amino_mapper = aaindex , Xdecoder_hidden=100 )\n",
    "\n",
    "\n",
    "encoder_save = 'encoder_mk2_aa_50_AAq_transformermk2'\n",
    "decoder_save = 'decoder_mk2_aa_50_AAq_transformermk2'\n",
    "\n",
    "betafactor = 2\n",
    "#put encoder and decoder on the device\n",
    "encoder = encoder.to(device)\n",
    "decoder = decoder.to(device)\n",
    "# Create a DataLoader for training\n",
    "total_loss_x = 0\n",
    "total_loss_edge = 0\n",
    "total_vq=0\n",
    "total_kl = 0\n",
    "total_plddt=0\n",
    "# Training loop\n",
    "if os.path.exists(encoder_save):\n",
    "    encoder.load_state_dict(torch.load(encoder_save))\n",
    "    decoder.load_state_dict(torch.load(decoder_save))\n",
    "    print(\"loaded encoder and decoder\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "    train_loader = DataLoader(struct_dat, batch_size=40, shuffle=True)\n",
    "\n",
    "    optimizer = torch.optim.Adam( list(decoder.parameters()), lr=0.001)\n",
    "    encoder.train()\n",
    "    decoder.train()\n",
    "    \n",
    "    xlosses = []\n",
    "    edgelosses = []\n",
    "    vqlosses = []\n",
    "    plddtlosses = []\n",
    "\n",
    "    edgeweight = 1\n",
    "    xweight = 2\n",
    "    vqweight = 2\n",
    "    plddtweight = 1\n",
    "\n",
    "    for epoch in range(1000):\n",
    "        for data in tqdm.tqdm(train_loader):\n",
    "            data = data.to(device)\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            z1 = encoder(data['res'].x, data['AA'].x , data.edge_index_dict)\n",
    "            z2 = encoder(data['res'].x, data['AA'].x , data.edge_index_dict)\n",
    "            \n",
    "            #add positional encoding to give to the decoder\n",
    "            edgeloss = recon_loss(z1,z2, , data.edge_index_dict[( 'res','contactPoints','res')]\n",
    "                                  , data.edge_index_dict[( 'res','backbone','res')], decoder)\n",
    "            \n",
    "            recon_x , edge_probs = decoder(z, data.edge_index_dict[( 'res','contactPoints','res')] , data.edge_index_dict )        \n",
    "    \n",
    "            xloss = aa_reconstruction_loss(data['AA'].x, recon_x)\n",
    "            #plddtloss = x_reconstruction_loss(data['plddt'].x, recon_plddt)\n",
    "            loss = xweight*xloss + edgeweight*edgeloss + vqweight*vqloss #+ plddtloss\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            total_loss_edge += edgeloss.item()\n",
    "            total_loss_x += xloss.item()\n",
    "            total_vq += vqloss.item()\n",
    "            #total_plddt += plddtloss.item()\n",
    "\n",
    "        if epoch % 100 == 0 :\n",
    "            #save model\n",
    "            torch.save(encoder.state_dict(), encoder_save)\n",
    "            torch.save(decoder.state_dict(), decoder_save)\n",
    "        \"\"\"\n",
    "        for loss in [( total_loss_x , xlosses , xweight ), (total_loss_edge, edgelosses, edgeweight), ( total_vq, vqlosses, vqweight ) ]:\n",
    "            loss[1].append(loss[0])\n",
    "            #calculate the mean delta loss for past 10 epochs\n",
    "            if len(loss[1]) > 10:\n",
    "                loss[1].pop(0)\n",
    "                mean_loss = np.mean(loss[0:5])\n",
    "                #calculate the delta loss for the last 5 epochs\n",
    "                delta_loss = np.mean(loss[1][-5:])\n",
    "                delta_loss = delta_loss- mean_loss\n",
    "                if delta_loss > 0:\n",
    "                    loss[2] = loss[2]*2\n",
    "                else:\n",
    "                    loss[2] = loss[2]/1.5\n",
    "                loss[2] = np.clip(loss[2], 1e-5, 1e5)\n",
    "        \"\"\"    \n",
    "        print(f'Epoch {epoch}, AALoss: {total_loss_x:.4f}, Edge Loss: {total_loss_edge:.4f}, vq Loss: {total_vq:.4f}') #, plddt Loss: {total_plddt:.4f}')\n",
    "        total_loss_x = 0\n",
    "        total_loss_edge = 0\n",
    "        total_vq = 0\n",
    "        #total_plddt = 0\n",
    "    torch.save(encoder.state_dict(), encoder_save)\n",
    "    torch.save(decoder.state_dict(), decoder_save)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#brainstorming to get complex graphs...\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.nn import TransformerEncoder, TransformerEncoderLayer\n",
    "from torch.nn import TransformerDecoder, TransformerDecoderLayer\n",
    "\n",
    "class GraphGenerationModel(nn.Module):\n",
    "    def __init__(self, input_dim, hidden_dim, output_dim, nhead, num_encoder_layers, num_decoder_layers, num_classes):\n",
    "        super(GraphGenerationModel, self).__init__()\n",
    "        self.encoder_layer = TransformerEncoderLayer(d_model=input_dim, nhead=nhead)\n",
    "        self.transformer_encoder = TransformerEncoder(self.encoder_layer, num_layers=num_encoder_layers)\n",
    "        \n",
    "        self.decoder_layer = TransformerDecoderLayer(d_model=output_dim, nhead=nhead)\n",
    "        self.transformer_decoder = TransformerDecoder(self.decoder_layer, num_layers=num_decoder_layers)\n",
    "        \n",
    "        self.node_mlp = nn.Linear(output_dim, num_classes + 1)  # +1 for EOS token\n",
    "        self.edge_mlp = nn.Linear(2 * output_dim, 1)\n",
    "        self.embedding = nn.Embedding(num_classes + 1, output_dim)\n",
    "        \n",
    "    def encode(self, x):\n",
    "        return self.transformer_encoder(x)\n",
    "    \n",
    "    def decode(self, z, memory):\n",
    "        return self.transformer_decoder(z, memory)\n",
    "    \n",
    "    def forward(self, x, max_length=20):\n",
    "        # Encode the input sequence using the transformer encoder\n",
    "        memory = self.encode(x)  # Shape: [batch_size, seq_len, input_dim]\n",
    "        \n",
    "        # Start with an initial input (e.g., a start token)\n",
    "        start_token = torch.zeros((memory.size(0), 1, memory.size(-1)), device=x.device)\n",
    "        z = start_token\n",
    "        \n",
    "        node_logits = []\n",
    "        for _ in range(max_length):\n",
    "            z = self.decode(z, memory)\n",
    "            node_logit = self.node_mlp(z[:, -1, :])  # Take the last token's output\n",
    "            node_logits.append(node_logit)\n",
    "            \n",
    "            # Get the next input token (embedding of the predicted class)\n",
    "            next_token = node_logit.argmax(dim=-1).unsqueeze(1)\n",
    "            next_token_embedding = self.embedding(next_token)\n",
    "            z = torch.cat([z, next_token_embedding], dim=1)\n",
    "            \n",
    "            if (next_token == num_classes).all():  # EOS token\n",
    "                break\n",
    "        \n",
    "        node_logits = torch.stack(node_logits, dim=1)\n",
    "        \n",
    "        # Decode edges (fully connected example, customize as needed)\n",
    "        num_nodes = node_logits.size(1)\n",
    "        edge_index = torch.combinations(torch.arange(num_nodes), r=2).t().contiguous()\n",
    "        edge_features = torch.cat([node_logits[:, edge_index[0]], node_logits[:, edge_index[1]]], dim=-1)\n",
    "        edge_logits = self.edge_mlp(edge_features).squeeze(-1)\n",
    "        edge_probs = torch.sigmoid(edge_logits)\n",
    "        \n",
    "        return node_logits, edge_index, edge_probs\n",
    "\n",
    "# Example usage\n",
    "input_dim = 10   # Dimension of input embeddings\n",
    "hidden_dim = 32  # Dimension of hidden layer\n",
    "output_dim = 16  # Dimension of node features\n",
    "nhead = 2        # Number of heads in the multi-head attention\n",
    "num_encoder_layers = 2  # Number of transformer encoder layers\n",
    "num_decoder_layers = 2  # Number of transformer decoder layers\n",
    "num_classes = 5  # Number of node categories (excluding EOS token)\n",
    "\n",
    "model = GraphGenerationModel(input_dim, hidden_dim, output_dim, nhead, num_encoder_layers, num_decoder_layers, num_classes)\n",
    "\n",
    "# Example input: batch of sequences (batch_size, seq_len, input_dim)\n",
    "x = torch.randn(5, 10, input_dim)  # Batch of 5 sequences, each of length 10\n",
    "\n",
    "node_logits, edge_index, edge_probs = model(x)\n",
    "print(\"Node logits:\\n\", node_logits)\n",
    "print(\"Edge index:\\n\", edge_index)\n",
    "print(\"Edge probabilities:\\n\", edge_probs)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "def train(model, dataloader, epochs=10):\n",
    "    optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "    edge_loss_fn = nn.BCELoss()  # Loss function for edge probabilities\n",
    "    node_loss_fn = nn.NLLLoss()  # Loss function for node classification\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        model.train()\n",
    "        total_edge_loss = 0\n",
    "        total_node_loss = 0\n",
    "\n",
    "        for sequences, graphs, node_labels in dataloader:\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            node_logits, edge_index, edge_probs = model(sequences)\n",
    "\n",
    "            # Flatten the edge probabilities and create a target tensor\n",
    "            target_edge_probs = torch.cat([graph.edge_index for graph in graphs], dim=1)\n",
    "            target = torch.ones(target_edge_probs.size(1))  # Example target, customize as needed\n",
    "\n",
    "            # Compute edge loss\n",
    "            edge_loss = edge_loss_fn(edge_probs, target)\n",
    "\n",
    "            # Compute node classification loss\n",
    "            node_labels = node_labels.view(-1)\n",
    "            node_classes = node_logits.view(-1, num_classes + 1)  # +1 for EOS token\n",
    "            node_loss = node_loss_fn(node_classes, node_labels)\n",
    "\n",
    "            # Total loss\n",
    "            loss = edge_loss + node_loss\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            total_edge_loss += edge_loss.item()\n",
    "            total_node_loss += node_loss.item()\n",
    "\n",
    "        print(f\"Epoch {epoch+1}, Edge Loss: {total_edge_loss / len(dataloader)}, Node Loss: {total_node_loss / len(dataloader)}\")\n",
    "\n",
    "# Train the model\n",
    "train(model, dataloader)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}