3DGNN.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import argparse\n",
    "from tqdm.auto import tqdm\n",
    "import torch\n",
    "from torch.nn.utils import clip_grad_norm_\n",
    "# import torch_geometric\n",
    "# assert not torch_geometric.__version__.startswith('2'), 'Please use torch_geometric lower than version 2.0.0'\n",
    "from torch_geometric.loader import DataLoader\n",
    "\n",
    "from models.surfgen import SurfGen\n",
    "from utils.datasets import *\n",
    "from utils.transforms import *\n",
    "from utils.misc import *\n",
    "from utils.train import *\n",
    "from utils.datasets.surfdata import SurfGenDataset\n",
    "from time import time\n",
    "from utils.train import get_model_loss\n",
    "\n",
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument('--config', type=str, default='/home/haotian/molecules_confs/Protein_test/Pocket2Mol-main/configs/train.yml')\n",
    "parser.add_argument('--device', type=str, default='cpu')\n",
    "parser.add_argument('--logdir', type=str, default='/home/haotian/molecules_confs/Protein_test/Pocket2Mol-main/logs')\n",
    "args = parser.parse_args([])\n",
    "base_path = '/home/haotian/molecules_confs/Protein_test/SurfGen'\n",
    "args.config = os.path.join(base_path, 'configs/train_surf.yml')\n",
    "args.logdir = os.path.join(base_path, 'logs')\n",
    "config = load_config(args.config)\n",
    "config_name = os.path.basename(args.config)[:os.path.basename(args.config).rfind('.')]\n",
    "seed_all(config.train.seed)\n",
    "config.dataset.path = os.path.join(base_path, 'data/crossdocked_pocket10')\n",
    "config.dataset.split = os.path.join(base_path, 'data/split_by_name.pt')\n",
    "log_dir = get_new_log_dir(args.logdir, prefix=config_name)\n",
    "ckpt_dir = os.path.join(log_dir, 'checkpoints')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "protein_featurizer = FeaturizeProteinAtom()\n",
    "ligand_featurizer = FeaturizeLigandAtom()\n",
    "masking = get_mask(config.train.transform.mask)\n",
    "composer = AtomComposer(protein_featurizer.feature_dim, ligand_featurizer.feature_dim, config.model.encoder.knn)\n",
    "edge_sampler = EdgeSample(config.train.transform.edgesampler)\n",
    "cfg_ctr = config.train.transform.contrastive\n",
    "contrastive_sampler = ContrastiveSample(cfg_ctr.num_real, cfg_ctr.num_fake, cfg_ctr.pos_real_std, cfg_ctr.pos_fake_std, config.model.field.knn)\n",
    "transform = Compose([\n",
    "    RefineData(),\n",
    "    LigandCountNeighbors(),\n",
    "    protein_featurizer,\n",
    "    ligand_featurizer,\n",
    "    masking,\n",
    "    composer,\n",
    "\n",
    "    FocalBuilder(),\n",
    "    edge_sampler,\n",
    "    contrastive_sampler,\n",
    "])\n",
    "\n",
    "dataset, subsets = get_dataset(\n",
    "    config = config.dataset,\n",
    "    transform = transform,\n",
    ")\n",
    "dataset, subsets = get_dataset(\n",
    "    config = config.dataset,\n",
    "    transform = transform,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset, subsets = get_dataset(\n",
    "    config = config.dataset,\n",
    "    transform = transform,\n",
    ")\n",
    "train_set, val_set = subsets['train'], subsets['test']\n",
    "follow_batch = []\n",
    "collate_exclude_keys = ['ligand_nbh_list']\n",
    "train_iterator = inf_iterator(DataLoader(\n",
    "    train_set, \n",
    "    batch_size = config.train.batch_size, \n",
    "    shuffle = True,\n",
    "    num_workers = config.train.num_workers,\n",
    "    pin_memory = config.train.pin_memory,\n",
    "    follow_batch = follow_batch,\n",
    "    exclude_keys = collate_exclude_keys,\n",
    "))\n",
    "val_loader = DataLoader(val_set, config.train.batch_size, shuffle=False, follow_batch=follow_batch, exclude_keys = collate_exclude_keys,)\n",
    "train_loader = DataLoader(train_set, config.train.batch_size, shuffle=False,  exclude_keys = collate_exclude_keys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = val_set[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "from models.invariant import VNLinear,GVPerceptronVN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [],
   "source": [
    "pos = data.ligand_pos\n",
    "nbh_list = data.ligand_nbh_list \n",
    "node_feat = data.ligand_atom_feature_full.float()\n",
    "edge_index = data.ligand_bond_index\n",
    "edge_type = data.ligand_bond_type\n",
    "edge_feat = F.one_hot(edge_type-1, num_classes=3).float()\n",
    "edge_vec = (pos[edge_index[0]] - pos[edge_index[1]]).unsqueeze(1)\n",
    "num_atoms = pos.shape[0]\n",
    "max_dim = len(max(list(nbh_list.values()),key=len))\n",
    "local_coords = 0.1 * torch.ones([num_atoms,max_dim, 3])\n",
    "masker = torch.zeros([num_atoms,max_dim,3])\n",
    "for atom_idx in range(num_atoms):\n",
    "    local_coord = pos[atom_idx] - pos[nbh_list[atom_idx]]\n",
    "    local_coords[atom_idx,:local_coord.shape[0],:] = local_coord\n",
    "    masker[atom_idx,:local_coord.shape[0],:] = torch.ones([local_coord.shape[0],3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [],
   "source": [
    "mol = Chem.MolFromMol2File('2z3h_A_rec_1wn6_bst_lig_tt_docked_3.mol2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.chem import fragmentize_mol, remove_dummys_mol\n",
    "from torch_scatter import scatter_add"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [],
   "source": [
    "fragmentations = fragmentize_mol(mol)\n",
    "fragmentation = fragmentations[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "metadata": {},
   "outputs": [],
   "source": [
    "combine_mols = remove_dummys_mol(fragmentation[1])[0]\n",
    "frags = Chem.GetMolFrags(combine_mols,asMols=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "metadata": {},
   "outputs": [],
   "source": [
    "frag = remove_dummys_mol(frags[1])[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "metadata": {},
   "outputs": [],
   "source": [
    "motif_idx1 = torch.tensor(mol.GetSubstructMatch(frags[0]))\n",
    "motif_idx2 = torch.tensor(mol.GetSubstructMatch(frag))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_node_vec_dim = max_dim\n",
    "input_node_sca_dim = 13\n",
    "input_edge_vec_dim = 1\n",
    "input_edge_sca_dim = 3\n",
    "out_dim = 16\n",
    "\n",
    "node_vec_net = VNLinear(input_node_vec_dim,out_dim)\n",
    "node_sca_net = nn.Linear(input_node_sca_dim, out_dim)\n",
    "edge_vec_net = VNLinear(input_edge_vec_dim, out_dim)\n",
    "edge_sca_net = nn.Linear(input_edge_sca_dim, out_dim)\n",
    "node_net = nn.Linear(input_node_sca_dim, out_dim)\n",
    "edge_net = nn.Linear(input_edge_sca_dim, out_dim)\n",
    "edge_mapper = GVPerceptronVN(out_dim,out_dim,out_dim,out_dim)\n",
    "node_mapper = GVPerceptronVN(out_dim,out_dim,out_dim,out_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_raw = edge_index[0]\n",
    "node_vec_hid = node_vec_net(local_coords) \n",
    "node_sca_hid = node_sca_net(node_feat)\n",
    "edge_vec_hid = edge_vec_net(edge_vec)\n",
    "edge_sca_hid = edge_sca_net(edge_feat)\n",
    "\n",
    "msg_sca_j = node_sca_hid[edge_raw]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([13, 10, 14, 11, 12,  6,  7,  8,  9])"
      ]
     },
     "execution_count": 147,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "motif_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([31, 13])"
      ]
     },
     "execution_count": 148,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "node_feat.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Tensor"
      ]
     },
     "execution_count": 163,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(motif_idx[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Tensor"
      ]
     },
     "execution_count": 161,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(torch.ones_like(motif_idx)[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "metadata": {},
   "outputs": [],
   "source": [
    "### input motif index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 178,
   "metadata": {},
   "outputs": [],
   "source": [
    "motifs = [motif_idx1,motif_idx2]\n",
    "motif_index = torch.zeros(num_atoms).long()\n",
    "for idx,motif in enumerate(motifs):\n",
    "    motif_index[motif] = torch.ones_like(motif) * idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 180,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[13.,  5.,  4.,  0.,  0.,  0.,  0., 22., 45., 51., 39.,  6.,  0.],\n",
       "        [ 4.,  3.,  2.,  0.,  0.,  0.,  0.,  9., 19., 23., 15.,  4.,  0.]])"
      ]
     },
     "execution_count": 180,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "### motif pooling techniques\n",
    "scatter_add(node_feat, motif_index, dim=0)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.13 ('SurfGen')",
   "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.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "58c8be33f25b1934ef57e2589968d87ff83eb50288ed379ce5e63e82eedc880a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}