notebooks/AF2_model_docking.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **🔥🔥DiffBindFR**\n",
    "<div class=\"alert alert-info\"> Diffusion model based flexible protein-ligand docking\n",
    "</div>\n",
    "\n",
    "for more details see our [Preprint](https://arxiv.org/abs/2311.15201)\n",
    "\n",
    "Here, we conduct a demo by using DiffBindFR to redock ligand (fetched by PDB ID: 2ZEC) into the predefined pocket of **AlphaFold2 modelled structure** (Uniprot ID: Q15661)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ec9cd73f6b314ce7b4a09e61ab80afee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import os, sys, glob\n",
    "from pathlib import Path\n",
    "home = os.path.join(os.getcwd(), '../')\n",
    "home = os.path.abspath(home)\n",
    "if home not in sys.path:\n",
    "    sys.path.insert(0, home)\n",
    "home = Path(home)\n",
    "import pandas as pd\n",
    "from rdkit import Chem\n",
    "import MDAnalysis as mda\n",
    "import nglview as nv\n",
    "from nglview.color import ColormakerRegistry\n",
    "from DiffBindFR import common\n",
    "from DiffBindFR.evaluation import get_traj_id, export_xtc\n",
    "from DiffBindFR.app.predict import runner\n",
    "from DiffBindFR.utils import (\n",
    "    pair_spatial_metrics, \n",
    "    PDBPocketResidues, \n",
    "    to_complex_block,\n",
    "    read_molblock,\n",
    "    update_mol_pose,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_path = home / 'examples' / 'AF2'\n",
    "holo = example_path / '2zec.pdb'\n",
    "crystal_ligand = example_path / 'ligand.sdf'\n",
    "af2 = example_path / 'Q15661_AF2.pdb'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### PoseView of Holo structure\n",
    "\n",
    "Here, pocket residues within 5 angstrom of crystal ligand are visualized (colored by <font color='red'>red</font>)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8e09574522e44d9296c1d576f56ca767",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "NGLWidget()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pocket_buffer = 5\n",
    "holo_pocket = PDBPocketResidues.RDmolPocketResidues(\n",
    "    str(holo), str(crystal_ligand),\n",
    ")\n",
    "view = holo_pocket.visualize_pocket(pocket_buffer)\n",
    "view._remote_call('setSize', target='Widget', args=['','600px'])\n",
    "view"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Pocket Conformation visual inspection\n",
    "\n",
    "Compare the AF2 modeled pocket conformation (<font color='yellow'>yellow cartoon</font> and <font color='blue'>blue sticks</font>) with the crystal structure in advance.\n",
    "\n",
    "We could get the knowledge:\n",
    "\n",
    "- AF2 modeled structure has holo-like backbone with CA RMSD = 0.32 A\n",
    "- There are significant differences in pocket side chain conformation with sc-RMSD = 1.24 A, mainly from A:218:ASP, A:219:SER, A:221:GLN, A:244:TRP, A:246:GLU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9a3a66a5d2f94d769e09dab3d1ab748c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "NGLWidget()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "view = holo_pocket.compare(str(af2))\n",
    "view._remote_call('setSize', target='Widget', args=['','600px'])\n",
    "view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pocket CA RMSD within 5A of ligand: 0.32\n",
      "pocket side chain heavy atoms RMSD within 5A of ligand: 1.24\n"
     ]
    }
   ],
   "source": [
    "# Quantitative comparison of pocket conformation between af2 structure and holo\n",
    "holo_chainid = 'A' # see the receptor chain ID in poseview\n",
    "results_df = pair_spatial_metrics(\n",
    "    str(holo), str(crystal_ligand), str(af2),\n",
    "    holo_chainid, 'A', # af2 chain id is A as it is monomer prediction\n",
    "    bs_cutoff = pocket_buffer,\n",
    ")\n",
    "ca_rmsd = results_df.iloc[0].mean_ca_rmsd\n",
    "sc_rmsd = results_df.iloc[0].mean_sc_rmsd\n",
    "print('pocket CA RMSD within 5A of ligand:', round(ca_rmsd, 2))\n",
    "print('pocket side chain heavy atoms RMSD within 5A of ligand:', round(sc_rmsd, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pocket CA RMSD of A:218:ASP,A:219:SER,A:221:GLN,A:244:TRP,A:246:GLU: 0.23\n",
      "pocket side chain heavy atoms RMSD of A:218:ASP,A:219:SER,A:221:GLN,A:244:TRP,A:246:GLU: 1.78\n"
     ]
    }
   ],
   "source": [
    "# let's see the key residues\n",
    "flexible_residues = 'A:218:ASP,A:219:SER,A:221:GLN,A:244:TRP,A:246:GLU'\n",
    "holo_chainid = 'A' # see the receptor chain ID in poseview\n",
    "# reverse the input as the residue is on af2 structure and crystal_ligand here is useless\n",
    "results_df = pair_spatial_metrics(\n",
    "    str(af2), str(crystal_ligand), str(holo), \n",
    "    'A', holo_chainid, \n",
    "    bs_res_str = flexible_residues.split(','),\n",
    ")\n",
    "ca_rmsd = results_df.iloc[0].mean_ca_rmsd\n",
    "sc_rmsd = results_df.iloc[0].mean_sc_rmsd\n",
    "print(f'pocket CA RMSD of {flexible_residues}:', round(ca_rmsd, 2))\n",
    "print(f'pocket side chain heavy atoms RMSD of {flexible_residues}:', round(sc_rmsd, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run demo\n",
    "\n",
    "So here, we would like to use **DiffBindFR** to perform flexible docking and dock the ligand into pocket and refine the side chain conformation so that the refined structure is close to holo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment_name = 'Q15661'\n",
    "export_dir = 'demo_af2_docking'\n",
    "export_dir = os.path.abspath(export_dir)\n",
    "seed = 888 # 8888 has better ligand rmsd but not sc RMSD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " ____  _  __  __ ____  _           _ _____ ____  \n",
      "|  _ \\(_)/ _|/ _| __ )(_)_ __   __| |  ___|  _ \\ \n",
      "| | | | | |_| |_|  _ \\| | '_ \\ / _` | |_  | |_) |\n",
      "| |_| | |  _|  _| |_) | | | | | (_| |  _| |  _ < \n",
      "|____/|_|_| |_| |____/|_|_| |_|\\__,_|_|   |_| \\_\\ \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-03-13 22:17:38,637 - DiffBindFR - INFO - `use_rank_shift` = False,Set random seed to 888, deterministic: False\n",
      "2024-03-13 22:17:38,672 - DiffBindFR - INFO - Start to prepare job (experiment name: Q15661).\n",
      "2024-03-13 22:17:42,885 - DiffBindFR - INFO - dock Status: Prep task is Done!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Use Background Generator supported dataloader.\n",
      "Initializing diffusion model...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-03-13 22:17:46,597 - DiffBindFR - INFO - load checkpoint from local path: /data01/zhujintao/projects/home/research/DiffBindFR_OpenSource/DiffBindFR/weights/diffbindfr_paper.pth\n",
      "2024-03-13 22:17:48,655 - DiffBindFR - INFO - Running model inference...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>] 40/40, 0.5 task/s, elapsed: 76s, ETA:     0s"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-03-13 22:19:04,791 - DiffBindFR - INFO - Model inference is done which tasks 76.13553285598755s\n",
      "2024-03-13 22:19:04,792 - DiffBindFR - INFO - Export model results to path: /data01/zhujintao/projects/home/research/DiffBindFR_OpenSource/notebooks/demo_af2_docking/Q15661/results/model_output.pt\n",
      "2024-03-13 22:19:04,849 - DiffBindFR - INFO - Export binding structures....\n",
      "2024-03-13 22:19:36,899 - DiffBindFR - INFO - Binding structure export is completed.\n",
      "2024-03-13 22:19:36,902 - DiffBindFR - INFO - Start to binding conformation enrichment analysis...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-------------------------------------------------------------------------------------+\n",
      "|                             \u001b[5;36mDiffBindFR Model evaluations\u001b[0m                            |\n",
      "+-----------------------------------------------------------------------+-------------+\n",
      "| \u001b[31mMetric\u001b[0m                                                                | \u001b[34mPerformance\u001b[0m |\n",
      "+-----------------------------------------------------------------------+-------------+\n",
      "| Total poses                                                           | 1x40=40     |\n",
      "| l-rmsd gold standard                                                  | 2.0         |\n",
      "| centroid gold standard                                                | 1.0         |\n",
      "| chi1_15 gold standard                                                 | 0.75        |\n",
      "| sc-rmsd gold standard                                                 | 1.0         |\n",
      "|                                                                       |             |\n",
      "| All poses l-rmsd mean-std                                             | 4.65-4.55   |\n",
      "| All poses l-rmsd median                                               | 2.73        |\n",
      "| All poses l-rmsd count < 1.0A                                         | 1           |\n",
      "| All poses l-rmsd fraction < 1.0A                                      | 2.5%        |\n",
      "| All poses l-rmsd count < 2.0A                                         | 15          |\n",
      "| All poses l-rmsd fraction < 2.0A                                      | 37.5%       |\n",
      "| All poses l-rmsd count < 2.5A                                         | 18          |\n",
      "| All poses l-rmsd fraction < 2.5A                                      | 45.0%       |\n",
      "| All poses l-rmsd count < 3A                                           | 22          |\n",
      "| All poses l-rmsd fraction < 3A                                        | 55.0%       |\n",
      "| All poses l-rmsd count < 4.0A                                         | 26          |\n",
      "| All poses l-rmsd fraction < 4.0A                                      | 65.0%       |\n",
      "| Each complex l-rmsd < 2.0A SR mean-std                                | 0.38-0.0    |\n",
      "| Each complex l-rmsd < 2.0A SR median                                  | 0.38        |\n",
      "| top 1 l-rmsd all < 2.0A complex fraction                              | 100.0%      |\n",
      "| top 1 l-rmsd mapped centroid all < 1.0A complex fraction              | 100.0%      |\n",
      "| top 1 l-rmsd mapped chi1 < 15 degree SR all < 75.0% complex fraction  | 100.0%      |\n",
      "| top 1 l-rmsd mapped sc-rmsd all < 1.0A complex fraction               | 100.0%      |\n",
      "| top 3 l-rmsd all < 2.0A complex fraction                              | 100.0%      |\n",
      "| top 3 l-rmsd mapped centroid all < 1.0A complex fraction              | 100.0%      |\n",
      "| top 3 l-rmsd mapped chi1 < 15 degree SR all < 75.0% complex fraction  | 100.0%      |\n",
      "| top 3 l-rmsd mapped sc-rmsd all < 1.0A complex fraction               | 100.0%      |\n",
      "| top 5 l-rmsd all < 2.0A complex fraction                              | 100.0%      |\n",
      "| top 5 l-rmsd mapped centroid all < 1.0A complex fraction              | 100.0%      |\n",
      "| top 5 l-rmsd mapped chi1 < 15 degree SR all < 75.0% complex fraction  | 100.0%      |\n",
      "| top 5 l-rmsd mapped sc-rmsd all < 1.0A complex fraction               | 100.0%      |\n",
      "| top 10 l-rmsd all < 2.0A complex fraction                             | 100.0%      |\n",
      "| top 10 l-rmsd mapped centroid all < 1.0A complex fraction             | 0.0%        |\n",
      "| top 10 l-rmsd mapped chi1 < 15 degree SR all < 75.0% complex fraction | 100.0%      |\n",
      "| top 10 l-rmsd mapped sc-rmsd all < 1.0A complex fraction              | 100.0%      |\n",
      "|                                                                       |             |\n",
      "| All poses centroid mean-std                                           | 3.58-4.02   |\n",
      "| All poses centroid median                                             | 1.77        |\n",
      "| All poses centroid distance count < 0.5A                              | 1           |\n",
      "| All poses centroid fraction < 0.5A                                    | 2.5%        |\n",
      "| All poses centroid distance count < 1.0A                              | 9           |\n",
      "| All poses centroid fraction < 1.0A                                    | 22.5%       |\n",
      "| All poses centroid distance count < 2.0A                              | 23          |\n",
      "| All poses centroid fraction < 2.0A                                    | 57.5%       |\n",
      "| Each complex centroid < 1.0A SR mean-std                              | 0.22-0.0    |\n",
      "| Each complex centroid < 1.0A SR median                                | 0.22        |\n",
      "| top 1 centroid all < 1.0 angstrom complex fraction                    | 100.0%      |\n",
      "| top 3 centroid all < 1.0 angstrom complex fraction                    | 100.0%      |\n",
      "| top 5 centroid all < 1.0 angstrom complex fraction                    | 100.0%      |\n",
      "| top 10 centroid all < 1.0 angstrom complex fraction                   | 0.0%        |\n",
      "|                                                                       |             |\n",
      "| All conf chi1 15 SR mean-std                                          | 0.81-0.05   |\n",
      "| All conf chi1 15 SR median                                            | 0.81        |\n",
      "| All conf chi1 15 SR count > 80.0%                                     | 26          |\n",
      "| All conf chi1 15 SR fraction > 80.0%                                  | 65.0%       |\n",
      "| All conf chi1 15 SR count > 70.0%                                     | 40          |\n",
      "| All conf chi1 15 SR fraction > 70.0%                                  | 100.0%      |\n",
      "| All conf chi1 15 SR count > 60.0%                                     | 40          |\n",
      "| All conf chi1 15 SR fraction > 60.0%                                  | 100.0%      |\n",
      "| All conf chi1 15 SR count > 50.0%                                     | 40          |\n",
      "| All conf chi1 15 SR fraction > 50.0%                                  | 100.0%      |\n",
      "| Each complex chi1 15 SR > 75.0% SR mean-std                           | 0.88-0.0    |\n",
      "| Each complex chi1 15 SR > 75.0% SR median                             | 0.88        |\n",
      "| top 1 chi1 15 SR all > 75.0% complex fraction                         | 100.0%      |\n",
      "| top 3 chi1 15 SR all > 75.0% complex fraction                         | 100.0%      |\n",
      "| top 5 chi1 15 SR all > 75.0% complex fraction                         | 100.0%      |\n",
      "| top 10 chi1 15 SR all > 75.0% complex fraction                        | 100.0%      |\n",
      "|                                                                       |             |\n",
      "| All conf sc-rmsd mean-std                                             | 0.76-0.12   |\n",
      "| All conf sc-rmsd median                                               | 0.78        |\n",
      "| All conf sc-rmsd count < 0.5A                                         | 0           |\n",
      "| All conf sc-rmsd fraction < 0.5A                                      | 0.0%        |\n",
      "| All conf sc-rmsd count < 1.0A                                         | 40          |\n",
      "| All conf sc-rmsd fraction < 1.0A                                      | 100.0%      |\n",
      "| All conf sc-rmsd count < 1.5A                                         | 40          |\n",
      "| All conf sc-rmsd fraction < 1.5A                                      | 100.0%      |\n",
      "| All conf sc-rmsd count < 2.0A                                         | 40          |\n",
      "| All conf sc-rmsd fraction < 2.0A                                      | 100.0%      |\n",
      "| Each complex sc-rmsd < 1.0A SR mean-std                               | 1.0-0.0     |\n",
      "| Each complex sc-rmsd < 1.0A SR median                                 | 1.0         |\n",
      "| top 1 sc-rmsd all < 1.0A complex fraction                             | 100.0%      |\n",
      "| top 3 sc-rmsd all < 1.0A complex fraction                             | 100.0%      |\n",
      "| top 5 sc-rmsd all < 1.0A complex fraction                             | 100.0%      |\n",
      "| top 10 sc-rmsd all < 1.0A complex fraction                            | 100.0%      |\n",
      "+-----------------------------------------------------------------------+-------------+\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-03-13 22:19:36,926 - DiffBindFR - INFO - Start to correct error...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO: Pandarallel will run on 1 workers.\n",
      "INFO: Pandarallel will use Memory file system to transfer data between the main process and workers.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-03-13 22:19:52,701 - DiffBindFR - INFO - Error correction is completed.\n",
      "2024-03-13 22:19:53,948 - DiffBindFR - INFO - Start to MDN based pose scoring...\n",
      "Processing 40 jobs, chunk 50000 per iter, total 1 iters...: 100%|██████████| 1/1 [00:05<00:00,  5.39s/it]\n",
      "2024-03-13 22:20:00,014 - DiffBindFR - INFO - Load scoring model...\n",
      "100%|██████████| 3/3 [00:00<00:00, 14.51it/s]\n",
      "2024-03-13 22:20:00,429 - DiffBindFR - INFO - Model Scoring is Done!\n",
      "2024-03-13 22:20:00,469 - DiffBindFR - INFO - MDN based pose scoring is completed.\n",
      "2024-03-13 22:20:00,471 - DiffBindFR - INFO - Clean up datasets and model_out.pt.\n",
      "2024-03-13 22:20:00,475 - DiffBindFR - INFO - DiffBindFR docking is Done!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of DiffBindFR perfect poses: 1\n",
      "The DiffBindFR perfect selected L-RMSD success rate: 100.0 %\n",
      "The DiffBindFR perfect selected sc-RMSD success rate: 100.0 %\n",
      "The DiffBindFR perfect selected centroid success rate: 100.0 %\n",
      "The DiffBindFR perfect selected chi1_15 success rate: 100.0 %\n",
      "\n",
      "Use error corrected poses...\n",
      "Total number of DiffBindFR-Smina perfect poses: 1\n",
      "The DiffBindFR-Smina perfect selected success rate: 100.0 %\n",
      "The DiffBindFR-Smina perfect selected sc-rmsd success rate: 100.0 %\n",
      "The DiffBindFR-Smina perfect selected centroid success rate: 100.0 %\n",
      "The DiffBindFR-Smina perfect selected chi1_15 success rate: 100.0 %\n",
      "\n",
      "Total number of DiffBindFR-Smina top1 poses: 1\n",
      "The DiffBindFR-Smina top1 selected success rate: 100.0 %\n",
      "The DiffBindFR-Smina top1 selected sc-rmsd success rate: 100.0 %\n",
      "The DiffBindFR-Smina top1 selected centroid success rate: 100.0 %\n",
      "The DiffBindFR-Smina top1 selected chi1_15 success rate: 100.0 %\n",
      "\n",
      "Total number of DiffBindFR-MDN top1 poses: 1\n",
      "The DiffBindFR-MDN top1 l-rmsd success rate: 100.0 %\n",
      "The DiffBindFR-MDN top1 sc-rmsd success rate: 100.0 %\n",
      "The DiffBindFR-MDN top1 centroid success rate: 0.0 %\n",
      "The DiffBindFR-MDN top1 chi1_15 success rate: 100.0 %\n",
      "Report perfect and top 1 performance is done!\n"
     ]
    }
   ],
   "source": [
    "# input parameters in jupyter using argparse\n",
    "parser = common.parse_args()\n",
    "args = parser.parse_args(\n",
    "    [\n",
    "        '-l', str(crystal_ligand),\n",
    "        '-p', str(af2),\n",
    "        '-o', export_dir,\n",
    "        '-np', '40', \n",
    "        '-gpu', '0', \n",
    "        '-cpu', '1', \n",
    "        '-bs', '16', \n",
    "        '-eval', '-rp', # here we automatically evaluate the redock performance\n",
    "        '-cl', \n",
    "        '-st',\n",
    "        '-n', experiment_name,\n",
    "        '--seed', str(seed),\n",
    "    ]\n",
    ")\n",
    "args.cfg_options = None\n",
    "job_df = common.make_inference_jobs(args)\n",
    "runner(job_df, args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_dir = os.path.join(export_dir, experiment_name, 'results')\n",
    "smina_top1 = os.path.join(results_dir, 'results_ec_smina_top1.csv')\n",
    "smina_top1 = pd.read_csv(smina_top1)\n",
    "smina_top1 = smina_top1.iloc[0]\n",
    "smina_top1_protein = smina_top1.protein_pdb\n",
    "smina_top1_pose = smina_top1.docked_lig\n",
    "mdn_top1 = os.path.join(results_dir, 'results_ec_mdn_top1.csv')\n",
    "mdn_top1 = pd.read_csv(mdn_top1)\n",
    "mdn_top1 = mdn_top1.iloc[0]\n",
    "mdn_top1_protein = mdn_top1.protein_pdb\n",
    "mdn_top1_pose = mdn_top1.docked_lig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Smina top1 prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "031821e561fb480fb3629ceb0e86e044",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "NGLWidget()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "holo_pocket = PDBPocketResidues.RDmolPocketResidues(\n",
    "    str(holo), str(crystal_ligand),\n",
    ")\n",
    "view = holo_pocket.visualize_pocket(pocket_buffer)\n",
    "view = holo_pocket.compare(smina_top1_protein, ligand_sdf = smina_top1_pose)\n",
    "view._remote_call('setSize', target='Widget', args=['','600px'])\n",
    "view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DiffBindFR-Smina\n",
      "ligand RMSD: 0.87\n",
      "pocket side chain heavy atoms RMSD within 5A of ligand: 1.18\n",
      "pocket side chain heavy atoms RMSD of A:218:ASP,A:219:SER,A:221:GLN,A:244:TRP,A:246:GLU: 0.97\n"
     ]
    }
   ],
   "source": [
    "print('DiffBindFR-Smina')\n",
    "print(f'ligand RMSD:', round(smina_top1['l-rmsd_ec'], 2))\n",
    "\n",
    "results_df = pair_spatial_metrics(\n",
    "    str(holo), str(crystal_ligand), str(smina_top1_protein),\n",
    "    holo_chainid, 'A', # af2 chain id is A as it is monomer prediction\n",
    "    bs_cutoff = pocket_buffer,\n",
    ")\n",
    "sc_rmsd = results_df.iloc[0].mean_sc_rmsd\n",
    "print('pocket side chain heavy atoms RMSD within 5A of ligand:', round(sc_rmsd, 2))\n",
    "\n",
    "results_df = pair_spatial_metrics(\n",
    "    str(smina_top1_protein), str(crystal_ligand), str(holo), \n",
    "    'A', holo_chainid, \n",
    "    bs_res_str = flexible_residues.split(','),\n",
    ")\n",
    "sc_rmsd = results_df.iloc[0].mean_sc_rmsd\n",
    "print(f'pocket side chain heavy atoms RMSD of {flexible_residues}:', round(sc_rmsd, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DiffBindFR-MDN\n",
      "ligand RMSD: 1.82\n",
      "pocket side chain heavy atoms RMSD within 5A of ligand: 1.47\n",
      "pocket side chain heavy atoms RMSD of A:218:ASP,A:219:SER,A:221:GLN,A:244:TRP,A:246:GLU: 1.42\n"
     ]
    }
   ],
   "source": [
    "print('DiffBindFR-MDN')\n",
    "print(f'ligand RMSD:', round(mdn_top1['l-rmsd_ec'], 2))\n",
    "\n",
    "results_df = pair_spatial_metrics(\n",
    "    str(holo), str(crystal_ligand), str(mdn_top1_protein),\n",
    "    holo_chainid, 'A', # af2 chain id is A as it is monomer prediction\n",
    "    bs_cutoff = pocket_buffer,\n",
    ")\n",
    "sc_rmsd = results_df.iloc[0].mean_sc_rmsd\n",
    "print('pocket side chain heavy atoms RMSD within 5A of ligand:', round(sc_rmsd, 2))\n",
    "\n",
    "results_df = pair_spatial_metrics(\n",
    "    str(mdn_top1_protein), str(crystal_ligand), str(holo), \n",
    "    'A', holo_chainid, \n",
    "    bs_res_str = flexible_residues.split(','),\n",
    ")\n",
    "sc_rmsd = results_df.iloc[0].mean_sc_rmsd\n",
    "print(f'pocket side chain heavy atoms RMSD of {flexible_residues}:', round(sc_rmsd, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Make trajectory movie"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "':A and ( 218 or 219 or 221 or 244 or 246 )'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cm = ColormakerRegistry\n",
    "cm.add_scheme_func('lig_atomwise','''\n",
    " this.atomColor = function (atom) {\n",
    "     if (atom.element == \"C\") {\n",
    "       return 0x7272e6 // C\n",
    "     } else if (atom.element == \"H\") {\n",
    "       return 0xecf0f1\n",
    "     } else if (atom.element == \"S\") {\n",
    "       return 0xf1c40f\n",
    "     } else if (atom.element == \"N\") {\n",
    "       return 0x2d2de1\n",
    "     } else if (atom.element == \"O\") {\n",
    "       return 0xff5252\n",
    "     }\n",
    " }\n",
    "''')\n",
    "cm.add_scheme_func('prot_atomwise','''\n",
    " this.atomColor = function (atom) {\n",
    "     if (atom.element == \"C\") {\n",
    "       return 0xf9f902 // C\n",
    "     } else if (atom.element == \"H\") {\n",
    "       return 0xecf0f1\n",
    "     } else if (atom.element == \"S\") {\n",
    "       return 0xf1c40f\n",
    "     } else if (atom.element == \"N\") {\n",
    "       return 0x2d2de1\n",
    "     } else if (atom.element == \"O\") {\n",
    "       return 0xff5252\n",
    "     }\n",
    " }\n",
    "''')\n",
    "\n",
    "def add_ec_to_xtc(\n",
    "    sample_dir: str,\n",
    "    topology: str,\n",
    "    new_name: str = 'new_prl_traj.xtc',\n",
    ") -> str:\n",
    "    pdb_final = os.path.join(sample_dir, 'prot_final.pdb')\n",
    "    lig_final = os.path.join(sample_dir, 'lig_final_ec.sdf')\n",
    "    lig_final_mol = Chem.SDMolSupplier(lig_final)[0]\n",
    "    lig_final_mol = Chem.MolFromPDBBlock(Chem.MolToPDBBlock(lig_final_mol))\n",
    "\n",
    "    traj_dir = os.path.join(sample_dir, 'prl_traj')\n",
    "    trajs = list(Path(traj_dir).glob('prl_*.pdb'))\n",
    "    assert len(trajs) > 0, 'please export trajectory when you run DiffBindFR sampling by turn on -st.'\n",
    "    ids = []\n",
    "    for traj in trajs:\n",
    "        stem = traj.stem\n",
    "        traj_id = get_traj_id(stem)\n",
    "        ids.append(traj_id)\n",
    "    max_id = max(ids)\n",
    "    final_id = max_id + 1\n",
    "    final_traj_path = os.path.join(traj_dir, f'prl_{final_id}.pdb')\n",
    "    seed_traj_path = trajs[0]\n",
    "\n",
    "    mol_seed_block = read_molblock(seed_traj_path)\n",
    "    mol_seed = Chem.MolFromPDBBlock(mol_seed_block) # use mol_seed topology to export PDB block\n",
    "    lig_final_mol = update_mol_pose(mol_seed, lig_final_mol)\n",
    "\n",
    "    trajectory = os.path.join(sample_dir, new_name)\n",
    "    p_pdbblock = Path(pdb_final).read_text()\n",
    "    l_pdbblock = Chem.MolToPDBBlock(lig_final_mol)\n",
    "    try:\n",
    "        complex_pdb_block = to_complex_block(p_pdbblock, l_pdbblock, final_traj_path)\n",
    "        export_xtc(\n",
    "            topology,\n",
    "            traj_dir,\n",
    "            trajectory,\n",
    "        )\n",
    "    finally:\n",
    "        if os.path.exists(final_traj_path):\n",
    "            os.remove(final_traj_path) # avoid increment by multiple run\n",
    "    return trajectory\n",
    "\n",
    "def show_nv_traj(\n",
    "    sample_dir: str,\n",
    "    repr_sel: str,\n",
    "    add_ec_to_xtc_flag = True,\n",
    "):\n",
    "    topology = os.path.join(sample_dir, '../prl_topol.pdb')\n",
    "\n",
    "    if add_ec_to_xtc_flag:\n",
    "        # add ec ligand into xtc\n",
    "        trajectory = add_ec_to_xtc(sample_dir, topology)\n",
    "    else:\n",
    "        trajectory = os.path.join(sample_dir, 'prl_traj.xtc')\n",
    "\n",
    "    u = mda.Universe(topology, trajectory)\n",
    "    system = u.select_atoms('all')\n",
    "    t = nv.MDAnalysisTrajectory(system)\n",
    "    w = nv.NGLWidget(t)\n",
    "    w.clear_representations()\n",
    "    w.add_cartoon(colorScheme = 'sstruc')\n",
    "    w.add_representation(\n",
    "        repr_type='ball+stick', \n",
    "        selection='[UNL]', # ligand resname\n",
    "        color_scheme = 'lig_atomwise'\n",
    "    )\n",
    "    w.add_representation('licorice', selection=repr_sel, color_scheme='prot_atomwise')\n",
    "\n",
    "    if add_ec_to_xtc_flag:\n",
    "        os.remove(trajectory)\n",
    "    \n",
    "    return w\n",
    "\n",
    "# make nglview selection expression\n",
    "flex_residue_list = flexible_residues.split(',')\n",
    "flex_resnumber = [x.split(':')[1] for x in flex_residue_list]\n",
    "flex_resnumber = ':A and ' + '( ' + ' or '.join(flex_resnumber) + ' )'\n",
    "flex_resnumber"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0cc2c2ecdd014ee29b1a0ee32383d694",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "NGLWidget(max_frame=20)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sample_dir = os.path.dirname(smina_top1_protein)\n",
    "w = show_nv_traj(sample_dir, flex_resnumber, True)\n",
    "w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e4601b06865747a6947f46bb42ccea28",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "NGLWidget(max_frame=20)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sample_dir = os.path.dirname(mdn_top1_protein)\n",
    "w = show_nv_traj(sample_dir, flex_resnumber, True)\n",
    "w"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🎉🎉End\n",
    "\n",
    "Thanks for your interest in DiffBindFR. We are still working hard to further improve performance and extend it to other applications.\n",
    "\n",
    "If you have any question, feel free to open a [github issue](https://github.com/HBioquant/DiffBindFR/issues) or reach out to me: [zhujt@stu.pku.edu.cn](zhujt@stu.pku.edu.cn)\n",
    "\n",
    "👋👋👋"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "quantum",
   "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}