model_inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6c6bf22e-55e9-42b3-a53e-e148749c14ef",
   "metadata": {},
   "source": "# MapDiff sequence inference from a pdb file "
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "from hydra import initialize, compose\n",
    "from model.egnn_pytorch.egnn_net import EGNN_NET\n",
    "from model.ipa.ipa_net import IPANetPredictor\n",
    "from model.prior_diff import Prior_Diff\n",
    "from utils import enable_dropout\n",
    "from dataloader.collator import CollatorDiff\n",
    "from data.generate_graph_cath import pdb2graph, get_processed_graph, amino_acids_type\n",
    "from tqdm import tqdm"
   ],
   "id": "c162835a1a606be9",
   "execution_count": null,
   "outputs": []
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# Initialize config and load trained model weights",
   "id": "6b294692739302fd"
  },
  {
   "cell_type": "code",
   "id": "bae5c616-e8bb-4b68-a574-0c3d73d4bd06",
   "metadata": {},
   "source": [
    "with initialize(version_base=None, config_path=\"conf\"):\n",
    "    cfg = compose(config_name=\"inference\")\n",
    "\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# load trained model\n",
    "egnn = EGNN_NET(input_feat_dim=cfg.model.input_feat_dim, hidden_channels=cfg.model.hidden_dim,\n",
    "                edge_attr_dim=cfg.model.edge_attr_dim,\n",
    "                dropout=cfg.model.drop_out, n_layers=cfg.model.depth, update_edge=cfg.model.update_edge,\n",
    "                norm_coors=cfg.model.norm_coors, update_coors=cfg.model.update_coors,\n",
    "                update_global=cfg.model.update_global, embedding=cfg.model.embedding,\n",
    "                embedding_dim=cfg.model.embedding_dim, norm_feat=cfg.model.norm_feat, embed_ss=cfg.model.embed_ss)\n",
    "\n",
    "ipa = IPANetPredictor(dropout=cfg.model.ipa_drop_out)\n",
    "model = Prior_Diff(egnn, ipa, timesteps=cfg.diffusion.timesteps,\n",
    "                   objective=cfg.diffusion.objective,\n",
    "                   noise_type=cfg.diffusion.noise_type, sample_method=cfg.diffusion.sample_method,\n",
    "                   min_mask_ratio=cfg.mask_prior.min_mask_ratio,\n",
    "                   dev_mask_ratio=cfg.mask_prior.dev_mask_ratio,\n",
    "                   marginal_dist_path=cfg.dataset.marginal_train_dir).to(device)\n",
    "\n",
    "checkpoint = torch.load(cfg.test_model.path)\n",
    "model.load_state_dict(checkpoint['model'], strict=True)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "b141d4ad-9704-4cfc-86fd-35ef1e7b0d38",
   "metadata": {},
   "source": "# MapDiff sequence inference"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# load data\n",
    "pdb_dir = \"data/sample_pdb/\"\n",
    "pdb_files = os.listdir(pdb_dir)\n",
    "\n",
    "for pdb_file in tqdm(pdb_files):\n",
    "    graph = get_processed_graph(pdb2graph(pdb_dir + pdb_file))\n",
    "    collator = CollatorDiff()\n",
    "    g_batch, ipa_batch = collator([graph])\n",
    "    g_batch, ipa_batch = g_batch.to(device), ipa_batch.to(device)\n",
    "\n",
    "    # predict\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        ens_logits = []\n",
    "        enable_dropout(model)\n",
    "        for _ in range(cfg.diffusion.ensemble_num):\n",
    "            logits, sample_graph = model.mc_ddim_sample(g_batch, ipa_batch, diverse=True, step=cfg.diffusion.ddim_steps)\n",
    "            ens_logits.append(logits)\n",
    "        ens_logits_tensor = torch.stack(ens_logits)\n",
    "        mean_sample_logits = ens_logits_tensor.mean(dim=0).cpu()\n",
    "        true_label = g_batch.x.cpu()\n",
    "        true_sample_seq = ''.join([amino_acids_type[i] for i in true_label.argmax(dim=1).tolist()])\n",
    "        pred_sample_seq = ''.join([amino_acids_type[i] for i in mean_sample_logits.argmax(dim=1).tolist()])\n",
    "    \n",
    "        ll_fullseq = F.cross_entropy(mean_sample_logits, true_label, reduction='mean').item()\n",
    "        perplexity = np.exp(ll_fullseq)\n",
    "        sample_recovery = (mean_sample_logits.argmax(dim=1) == true_label.argmax(dim=1)).sum() / true_label.shape[0]\n",
    "        \n",
    "        print(f'PDB file: {pdb_file}')\n",
    "        print(f'Sequence length: {len(pred_sample_seq)}')\n",
    "        print(f'Sample perplexity: {perplexity}')\n",
    "        print(f'Sample recovery rate {sample_recovery:.4f}')\n",
    "        print(f'Pred sequence: {pred_sample_seq}')\n",
    "        print(f'True sequence: {true_sample_seq}')"
   ],
   "id": "a06013d968d648ae",
   "outputs": [],
   "execution_count": null
  }
 ],
 "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.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}