virtual_screening.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Virtual Screening"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.stats import ranksums"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read in the three sets of generated embeddings using the protein structures from the three lowest affinity binders, first for the decoys, then for the binders."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "B1MDI3_decoys = torch.load('checkpoints/quickbind_default/virtual_screening/B1MDI3_predictionsw_single_rep.pt')\n",
    "B1MDI3_decoys_1 = torch.load('checkpoints/quickbind_default/virtual_screening/decoy_B1MDI3_predictions_1-w-single-rep.pt')\n",
    "B1MDI3_decoys_2 = torch.load('checkpoints/quickbind_default/virtual_screening/decoy_B1MDI3_predictions_2-w-single-rep.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "P56817_decoys = torch.load('checkpoints/quickbind_default/virtual_screening/P56817_predictionsw_single_rep.pt')\n",
    "P56817_decoys_1 = torch.load('checkpoints/quickbind_default/virtual_screening/decoy_P56817_predictions_1-w-single-rep.pt')\n",
    "P56817_decoys_2 = torch.load('checkpoints/quickbind_default/virtual_screening/decoy_P56817_predictions_2-w-single-rep.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "P17931_decoys = torch.load('checkpoints/quickbind_default/virtual_screening/P17931_predictionsw_single_rep.pt')\n",
    "P17931_decoys_1 = torch.load('checkpoints/quickbind_default/virtual_screening/decoy_P17931_predictions_1-w-single-rep.pt')\n",
    "P17931_decoys_2 = torch.load('checkpoints/quickbind_default/virtual_screening/decoy_P17931_predictions_2-w-single-rep.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "Q8ULI9_decoys = torch.load('checkpoints/quickbind_default/virtual_screening/Q8ULI9_predictionsw_single_rep.pt')\n",
    "Q8ULI9_decoys_1 = torch.load('checkpoints/quickbind_default/virtual_screening/decoy_Q8ULI9_predictions_1-w-single-rep.pt')\n",
    "Q8ULI9_decoys_2 = torch.load('checkpoints/quickbind_default/virtual_screening/decoy_Q8ULI9_predictions_2-w-single-rep.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "P01116_decoys = torch.load('checkpoints/quickbind_default/virtual_screening/P01116_predictionsw_single_rep.pt')\n",
    "P01116_decoys_1 = torch.load('checkpoints/quickbind_default/virtual_screening/decoy_P01116_predictions_1-w-single-rep.pt')\n",
    "P01116_decoys_2 = torch.load('checkpoints/quickbind_default/virtual_screening/decoy_P01116_predictions_2-w-single-rep.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "B1MDI3_binders = [\n",
    "    '6qqt', '6qrf', '6qqw', '6qre', '6qrc', '6qqv', '6qrg', '6qr1', '6qqq', '6qqu', '6qr2', '6qra', '6qr4', '6qr3',\n",
    "    '6qr9', '6qqz', '6qr0', '6qrd', '6qr7'\n",
    "]\n",
    "B1MDI3_binders_out = torch.load('checkpoints/quickbind_default/virtual_screening/true_B1MDI3_predictions_w_single_rep.pt')\n",
    "B1MDI3_binders_out_1 = torch.load('checkpoints/quickbind_default/virtual_screening/true_B1MDI3_predictions_1-w-single-rep.pt')\n",
    "B1MDI3_binders_out_2 = torch.load('checkpoints/quickbind_default/virtual_screening/true_B1MDI3_predictions_2-w-single-rep.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "P56817_binders = [\n",
    "    '6uvp', '6uvv', '6uvy', '6uwp', '6nw3', '6e3z', '6nv7', '6uwv', '6nv9', '6od6', '6jt3', '6jsg',\n",
    "    '6jsn', '6jsf', '6jse', '6pz4'\n",
    "]\n",
    "P56817_binders_out = torch.load('checkpoints/quickbind_default/virtual_screening/true_P56817_predictions_w_single_rep.pt')\n",
    "P56817_binders_out_1 = torch.load('checkpoints/quickbind_default/virtual_screening/true_P56817_predictions_1-w-single-rep.pt')\n",
    "P56817_binders_out_2 = torch.load('checkpoints/quickbind_default/virtual_screening/true_P56817_predictions_2-w-single-rep.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "P17931_binders = [\n",
    "    '6qlt', '6qlq', '6i75', '6qlu', '6qln', '6i78', '6qlr', '6i77', '6qlo', '6i76', '6qlp', '6qls',\n",
    "    '6i74', '6qge', '6qgf'\n",
    "]\n",
    "P17931_binders_out = torch.load('checkpoints/quickbind_default/virtual_screening/true_P17931_predictions_w_single_rep.pt')\n",
    "P17931_binders_out_1 = torch.load('checkpoints/quickbind_default/virtual_screening/true_P17931_predictions_1-w-single-rep.pt')\n",
    "P17931_binders_out_2 = torch.load('checkpoints/quickbind_default/virtual_screening/true_P17931_predictions_2-w-single-rep.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "Q8ULI9_binders = [\n",
    "    '6oy0', '6oxy', '6oxz', '6oy2', '6oxt', '6oxs', '6oxx', '6oxu', '6oy1', '6oxr', '6oxw', '6oxv',\n",
    "    '6oxp', '6oxq'\n",
    "]\n",
    "Q8ULI9_binders_out = torch.load('checkpoints/quickbind_default/virtual_screening/true_Q8ULI9_predictions_w_single_rep.pt')\n",
    "Q8ULI9_binders_out_1 = torch.load('checkpoints/quickbind_default/virtual_screening/true_Q8ULI9_predictions_1-w-single-rep.pt')\n",
    "Q8ULI9_binders_out_2 = torch.load('checkpoints/quickbind_default/virtual_screening/true_Q8ULI9_predictions_2-w-single-rep.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "P01116_binders = [\n",
    "    '6quw', '6quv', '6gj5', '6v5l', '6gj6', '6p8z', '6pgo', '6p8x', '6gj8', '6gj7', '6p8y', '6pgp',\n",
    "    '6oim'\n",
    "]\n",
    "P01116_binders_out = torch.load('checkpoints/quickbind_default/virtual_screening/true_P01116_predictions_w_single_rep.pt')\n",
    "P01116_binders_out_1 = torch.load('checkpoints/quickbind_default/virtual_screening/true_P01116_predictions_1-w-single-rep.pt')\n",
    "P01116_binders_out_2 = torch.load('checkpoints/quickbind_default/virtual_screening/true_P01116_predictions_2-w-single-rep.pt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Predict binding affinities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "from train_binding_affinity import BindingAffinityPredictor, BindingAffinityData, get_predictions\n",
    "from torch.utils.data import DataLoader\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = BindingAffinityPredictor(64)\n",
    "model.load_state_dict(torch.load(\n",
    "    'checkpoints/quickbind_default/binding_affinity_prediction/ckpt_seed42.pt',\n",
    "))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('binding_affinity_dict.pkl', 'rb') as f:\n",
    "    binding_affinity_dict = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_proteins = torch.load('data/processed/timesplit_test/rec_input_proc_ids.pt')\n",
    "test_protein_names = [p['complex_names'] for p in test_proteins]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "lists = []\n",
    "lists_1 = []\n",
    "lists_2 = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "B1MDI3_binders_input = [s.squeeze() for s in B1MDI3_binders_out['s_pre_struct']]\n",
    "B1MDI3_binders_ordered = [n for n in test_protein_names if n in B1MDI3_binders]\n",
    "B1MDI3_binders_dataset = BindingAffinityData(B1MDI3_binders_input, B1MDI3_binders_ordered, binding_affinity_dict)\n",
    "B1MDI3_binders_loader = DataLoader(B1MDI3_binders_dataset, batch_size=1, shuffle=False)\n",
    "predictions_binders, _ = get_predictions(model, B1MDI3_binders_loader)\n",
    "\n",
    "B1MDI3_decoys_input = [s.squeeze() for s in B1MDI3_decoys['s_pre_struct']]\n",
    "B1MDI3_decoys_dataset = BindingAffinityData(B1MDI3_decoys_input, B1MDI3_decoys['names'], binding_affinity_dict)\n",
    "B1MDI3_decoys_loader = DataLoader(B1MDI3_decoys_dataset, batch_size=1, shuffle=False)\n",
    "predictions_decoys, _ = get_predictions(model, B1MDI3_decoys_loader)\n",
    "\n",
    "lists.append([predictions_binders, predictions_decoys])\n",
    "\n",
    "######################\n",
    "\n",
    "B1MDI3_binders_input_1 = [s.squeeze() for s in B1MDI3_binders_out_1['s_pre_struct']]\n",
    "B1MDI3_binders_dataset = BindingAffinityData(B1MDI3_binders_input_1, B1MDI3_binders_ordered, binding_affinity_dict)\n",
    "B1MDI3_binders_loader = DataLoader(B1MDI3_binders_dataset, batch_size=1, shuffle=False)\n",
    "predictions_binders, _ = get_predictions(model, B1MDI3_binders_loader)\n",
    "\n",
    "B1MDI3_decoys_input_1 = [s.squeeze() for s in B1MDI3_decoys_1['s_pre_struct']]\n",
    "B1MDI3_decoys_dataset = BindingAffinityData(B1MDI3_decoys_input_1, B1MDI3_decoys_1['names'], binding_affinity_dict)\n",
    "B1MDI3_decoys_loader = DataLoader(B1MDI3_decoys_dataset, batch_size=1, shuffle=False)\n",
    "predictions_decoys, _ = get_predictions(model, B1MDI3_decoys_loader)\n",
    "\n",
    "lists_1.append([predictions_binders, predictions_decoys])\n",
    "\n",
    "######################\n",
    "\n",
    "B1MDI3_binders_input_2 = [s.squeeze() for s in B1MDI3_binders_out_2['s_pre_struct']]\n",
    "B1MDI3_binders_dataset = BindingAffinityData(B1MDI3_binders_input_2, B1MDI3_binders_ordered, binding_affinity_dict)\n",
    "B1MDI3_binders_loader = DataLoader(B1MDI3_binders_dataset, batch_size=1, shuffle=False)\n",
    "\n",
    "predictions_binders, _ = get_predictions(model, B1MDI3_binders_loader)\n",
    "\n",
    "B1MDI3_decoys_input_2 = [s.squeeze() for s in B1MDI3_decoys_2['s_pre_struct']]\n",
    "B1MDI3_decoys_dataset = BindingAffinityData(B1MDI3_decoys_input_2, B1MDI3_decoys_2['names'], binding_affinity_dict)\n",
    "B1MDI3_decoys_loader = DataLoader(B1MDI3_decoys_dataset, batch_size=1, shuffle=False)\n",
    "predictions_decoys, _ = get_predictions(model, B1MDI3_decoys_loader)\n",
    "\n",
    "lists_2.append([predictions_binders, predictions_decoys])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P56817_binders_input = [s.squeeze() for s in P56817_binders_out['s_pre_struct']]\n",
    "P56817_binders_ordered = [n for n in test_protein_names if n in P56817_binders]\n",
    "P56817_binders_dataset = BindingAffinityData(P56817_binders_input, P56817_binders_ordered, binding_affinity_dict)\n",
    "P56817_binders_loader = DataLoader(P56817_binders_dataset, batch_size=1, shuffle=False)\n",
    "predictions_binders, _ = get_predictions(model, P56817_binders_loader)\n",
    "\n",
    "P56817_decoys_input = [s.squeeze() for s in P56817_decoys['s_pre_struct']]\n",
    "P56817_decoys_dataset = BindingAffinityData(P56817_decoys_input, P56817_decoys['names'], binding_affinity_dict)\n",
    "P56817_decoys_loader = DataLoader(P56817_decoys_dataset, batch_size=1, shuffle=False)\n",
    "predictions_decoys, _ = get_predictions(model, P56817_decoys_loader)\n",
    "\n",
    "lists.append([predictions_binders, predictions_decoys])\n",
    "\n",
    "######################\n",
    "\n",
    "P56817_binders_input_1 = [s.squeeze() for s in P56817_binders_out_1['s_pre_struct']]\n",
    "P56817_binders_dataset = BindingAffinityData(P56817_binders_input_1, P56817_binders_ordered, binding_affinity_dict)\n",
    "P56817_binders_loader = DataLoader(P56817_binders_dataset, batch_size=1, shuffle=False)\n",
    "predictions_binders, _ = get_predictions(model, P56817_binders_loader)\n",
    "\n",
    "P56817_decoys_input_1 = [s.squeeze() for s in P56817_decoys_1['s_pre_struct']]\n",
    "P56817_decoys_dataset = BindingAffinityData(P56817_decoys_input_1, P56817_decoys_1['names'], binding_affinity_dict)\n",
    "P56817_decoys_loader = DataLoader(P56817_decoys_dataset, batch_size=1, shuffle=False)\n",
    "predictions_decoys, _ = get_predictions(model, P56817_decoys_loader)\n",
    "\n",
    "lists_1.append([predictions_binders, predictions_decoys])\n",
    "\n",
    "######################\n",
    "\n",
    "P56817_binders_input_2 = [s.squeeze() for s in P56817_binders_out_2['s_pre_struct']]\n",
    "P56817_binders_dataset = BindingAffinityData(P56817_binders_input_2, P56817_binders_ordered, binding_affinity_dict)\n",
    "P56817_binders_loader = DataLoader(P56817_binders_dataset, batch_size=1, shuffle=False)\n",
    "predictions_binders, _ = get_predictions(model, P56817_binders_loader)\n",
    "\n",
    "P56817_decoys_input_2 = [s.squeeze() for s in P56817_decoys_2['s_pre_struct']]\n",
    "P56817_decoys_dataset = BindingAffinityData(P56817_decoys_input_2, P56817_decoys_2['names'], binding_affinity_dict)\n",
    "P56817_decoys_loader = DataLoader(P56817_decoys_dataset, batch_size=1, shuffle=False)\n",
    "predictions_decoys, _ = get_predictions(model, P56817_decoys_loader)\n",
    "\n",
    "lists_2.append([predictions_binders, predictions_decoys])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P17931_binders_input = [s.squeeze() for s in P17931_binders_out['s_pre_struct']]\n",
    "P17931_binders_ordered = [n for n in test_protein_names if n in P17931_binders]\n",
    "P17931_binders_dataset = BindingAffinityData(P17931_binders_input, P17931_binders_ordered, binding_affinity_dict)\n",
    "P17931_binders_loader = DataLoader(P17931_binders_dataset, batch_size=1, shuffle=False)\n",
    "predictions_binders, _ = get_predictions(model, P17931_binders_loader)\n",
    "\n",
    "P17931_decoys_input = [s.squeeze() for s in P17931_decoys['s_pre_struct']]\n",
    "P17931_decoys_dataset = BindingAffinityData(P17931_decoys_input, P17931_decoys['names'], binding_affinity_dict)\n",
    "P17931_decoys_loader = DataLoader(P17931_decoys_dataset, batch_size=1, shuffle=False)\n",
    "predictions_decoys, _ = get_predictions(model, P17931_decoys_loader)\n",
    "\n",
    "lists.append([predictions_binders, predictions_decoys])\n",
    "\n",
    "######################\n",
    "\n",
    "P17931_binders_input_1 = [s.squeeze() for s in P17931_binders_out_1['s_pre_struct']]\n",
    "P17931_binders_dataset = BindingAffinityData(P17931_binders_input_1, P17931_binders_ordered, binding_affinity_dict)\n",
    "P17931_binders_loader = DataLoader(P17931_binders_dataset, batch_size=1, shuffle=False)\n",
    "predictions_binders, _ = get_predictions(model, P17931_binders_loader)\n",
    "\n",
    "P17931_decoys_input_1 = [s.squeeze() for s in P17931_decoys_1['s_pre_struct']]\n",
    "P17931_decoys_dataset = BindingAffinityData(P17931_decoys_input_1, P17931_decoys_1['names'], binding_affinity_dict)\n",
    "P17931_decoys_loader = DataLoader(P17931_decoys_dataset, batch_size=1, shuffle=False)\n",
    "predictions_decoys, _ = get_predictions(model, P17931_decoys_loader)\n",
    "\n",
    "lists_1.append([predictions_binders, predictions_decoys])\n",
    "\n",
    "######################\n",
    "\n",
    "P17931_binders_input_2 = [s.squeeze() for s in P17931_binders_out_2['s_pre_struct']]\n",
    "P17931_binders_dataset = BindingAffinityData(P17931_binders_input_2, P17931_binders_ordered, binding_affinity_dict)\n",
    "P17931_binders_loader = DataLoader(P17931_binders_dataset, batch_size=1, shuffle=False)\n",
    "predictions_binders, _ = get_predictions(model, P17931_binders_loader)\n",
    "\n",
    "P17931_decoys_input_2 = [s.squeeze() for s in P17931_decoys_2['s_pre_struct']]\n",
    "P17931_decoys_dataset = BindingAffinityData(P17931_decoys_input_2, P17931_decoys_2['names'], binding_affinity_dict)\n",
    "P17931_decoys_loader = DataLoader(P17931_decoys_dataset, batch_size=1, shuffle=False)\n",
    "predictions_decoys, _ = get_predictions(model, P17931_decoys_loader)\n",
    "\n",
    "lists_2.append([predictions_binders, predictions_decoys])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Q8ULI9_binders_input = [s.squeeze() for s in Q8ULI9_binders_out['s_pre_struct']]\n",
    "Q8ULI9_binders_ordered = [n for n in test_protein_names if n in Q8ULI9_binders]\n",
    "Q8ULI9_binders_dataset = BindingAffinityData(Q8ULI9_binders_input, Q8ULI9_binders_ordered, binding_affinity_dict)\n",
    "Q8ULI9_binders_loader = DataLoader(Q8ULI9_binders_dataset, batch_size=1, shuffle=False)\n",
    "predictions_binders, _ = get_predictions(model, Q8ULI9_binders_loader)\n",
    "\n",
    "Q8ULI9_decoys_input = [s.squeeze() for s in Q8ULI9_decoys['s_pre_struct']]\n",
    "Q8ULI9_decoys_dataset = BindingAffinityData(Q8ULI9_decoys_input, Q8ULI9_decoys['names'], binding_affinity_dict)\n",
    "Q8ULI9_decoys_loader = DataLoader(Q8ULI9_decoys_dataset, batch_size=1, shuffle=False)\n",
    "predictions_decoys, _ = get_predictions(model, Q8ULI9_decoys_loader)\n",
    "\n",
    "lists.append([predictions_binders, predictions_decoys])\n",
    "\n",
    "######################\n",
    "\n",
    "Q8ULI9_binders_input_1 = [s.squeeze() for s in Q8ULI9_binders_out_1['s_pre_struct']]\n",
    "Q8ULI9_binders_dataset = BindingAffinityData(Q8ULI9_binders_input_1, Q8ULI9_binders_ordered, binding_affinity_dict)\n",
    "Q8ULI9_binders_loader = DataLoader(Q8ULI9_binders_dataset, batch_size=1, shuffle=False)\n",
    "predictions_binders, _ = get_predictions(model, Q8ULI9_binders_loader)\n",
    "\n",
    "Q8ULI9_decoys_input_1 = [s.squeeze() for s in Q8ULI9_decoys_1['s_pre_struct']]\n",
    "Q8ULI9_decoys_dataset = BindingAffinityData(Q8ULI9_decoys_input_1, Q8ULI9_decoys_1['names'], binding_affinity_dict)\n",
    "Q8ULI9_decoys_loader = DataLoader(Q8ULI9_decoys_dataset, batch_size=1, shuffle=False)\n",
    "predictions_decoys, _ = get_predictions(model, Q8ULI9_decoys_loader)\n",
    "\n",
    "lists_1.append([predictions_binders, predictions_decoys])\n",
    "\n",
    "######################\n",
    "\n",
    "Q8ULI9_binders_input_2 = [s.squeeze() for s in Q8ULI9_binders_out_2['s_pre_struct']]\n",
    "Q8ULI9_binders_dataset = BindingAffinityData(Q8ULI9_binders_input_2, Q8ULI9_binders_ordered, binding_affinity_dict)\n",
    "Q8ULI9_binders_loader = DataLoader(Q8ULI9_binders_dataset, batch_size=1, shuffle=False)\n",
    "predictions_binders, _ = get_predictions(model, Q8ULI9_binders_loader)\n",
    "\n",
    "Q8ULI9_decoys_input_2 = [s.squeeze() for s in Q8ULI9_decoys_2['s_pre_struct']]\n",
    "Q8ULI9_decoys_dataset = BindingAffinityData(Q8ULI9_decoys_input_2, Q8ULI9_decoys_2['names'], binding_affinity_dict)\n",
    "Q8ULI9_decoys_loader = DataLoader(Q8ULI9_decoys_dataset, batch_size=1, shuffle=False)\n",
    "predictions_decoys, _ = get_predictions(model, Q8ULI9_decoys_loader)\n",
    "\n",
    "lists_2.append([predictions_binders, predictions_decoys])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P01116_binders_input = [s.squeeze() for s in P01116_binders_out['s_pre_struct']]\n",
    "P01116_binders_ordered = [n for n in test_protein_names if n in P01116_binders]\n",
    "P01116_binders_dataset = BindingAffinityData(P01116_binders_input, P01116_binders_ordered, binding_affinity_dict)\n",
    "P01116_binders_loader = DataLoader(P01116_binders_dataset, batch_size=1, shuffle=False)\n",
    "predictions_binders, _ = get_predictions(model, P01116_binders_loader)\n",
    "\n",
    "P01116_decoys_input = [s.squeeze() for s in P01116_decoys['s_pre_struct']]\n",
    "P01116_decoys_dataset = BindingAffinityData(P01116_decoys_input, P01116_decoys['names'], binding_affinity_dict)\n",
    "P01116_decoys_loader = DataLoader(P01116_decoys_dataset, batch_size=1, shuffle=False)\n",
    "predictions_decoys, _ = get_predictions(model, P01116_decoys_loader)\n",
    "\n",
    "lists.append([predictions_binders, predictions_decoys])\n",
    "\n",
    "######################\n",
    "\n",
    "P01116_binders_input_1 = [s.squeeze() for s in P01116_binders_out_1['s_pre_struct']]\n",
    "P01116_binders_dataset = BindingAffinityData(P01116_binders_input_1, P01116_binders_ordered, binding_affinity_dict)\n",
    "P01116_binders_loader = DataLoader(P01116_binders_dataset, batch_size=1, shuffle=False)\n",
    "predictions_binders, _ = get_predictions(model, P01116_binders_loader)\n",
    "\n",
    "P01116_decoys_input_1 = [s.squeeze() for s in P01116_decoys_1['s_pre_struct']]\n",
    "P01116_decoys_dataset = BindingAffinityData(P01116_decoys_input_1, P01116_decoys_1['names'], binding_affinity_dict)\n",
    "P01116_decoys_loader = DataLoader(P01116_decoys_dataset, batch_size=1, shuffle=False)\n",
    "predictions_decoys, _ = get_predictions(model, P01116_decoys_loader)\n",
    "\n",
    "lists_1.append([predictions_binders, predictions_decoys])\n",
    "\n",
    "######################\n",
    "\n",
    "P01116_binders_input_2 = [s.squeeze() for s in P01116_binders_out_2['s_pre_struct']]\n",
    "P01116_binders_dataset = BindingAffinityData(P01116_binders_input_2, P01116_binders_ordered, binding_affinity_dict)\n",
    "P01116_binders_loader = DataLoader(P01116_binders_dataset, batch_size=1, shuffle=False)\n",
    "predictions_binders, _ = get_predictions(model, P01116_binders_loader)\n",
    "\n",
    "P01116_decoys_input_2 = [s.squeeze() for s in P01116_decoys_2['s_pre_struct']]\n",
    "P01116_decoys_dataset = BindingAffinityData(P01116_decoys_input_2, P01116_decoys_2['names'], binding_affinity_dict)\n",
    "P01116_decoys_loader = DataLoader(P01116_decoys_dataset, batch_size=1, shuffle=False)\n",
    "predictions_decoys, _ = get_predictions(model, P01116_decoys_loader)\n",
    "\n",
    "lists_2.append([predictions_binders, predictions_decoys])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 5, figsize=(10, 5))\n",
    "sns.set(style=\"ticks\")\n",
    "names = ['B1MDI3', 'P56817', 'P17931', 'Q8ULI9', 'P01116']\n",
    "\n",
    "for i, set_of_lists in enumerate(lists):\n",
    "    name = names[i]\n",
    "    statistic, p_value = ranksums(set_of_lists[0], set_of_lists[1], alternative='greater')\n",
    "\n",
    "    sns.boxplot(data=[set_of_lists[0], set_of_lists[1]], ax=axs[i])\n",
    "    axs[i].set_xticks([0, 1])\n",
    "    axs[i].set_xticklabels(['Binders', 'Decoys'])\n",
    "    axs[i].set_title(f'{name}\\np-value: {p_value:.4f}')\n",
    "\n",
    "axs[0].set_ylabel('Predicted Binding Affinity')\n",
    "plt.tight_layout()\n",
    "plt.savefig('binding_affinity_screen.jpg', bbox_inches='tight', format='jpg')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 5, figsize=(10, 5))\n",
    "sns.set(style=\"ticks\")\n",
    "names = ['B1MDI3', 'P56817', 'P17931', 'Q8ULI9', 'P01116']\n",
    "\n",
    "for i, set_of_lists in enumerate(lists_1):\n",
    "    name = names[i]\n",
    "    statistic, p_value = ranksums(set_of_lists[0], set_of_lists[1], alternative='greater')\n",
    "\n",
    "    sns.boxplot(data=[set_of_lists[0], set_of_lists[1]], ax=axs[i])\n",
    "    axs[i].set_xticks([0, 1])\n",
    "    axs[i].set_xticklabels(['Binders', 'Decoys'])\n",
    "    axs[i].set_title(f'{name}\\np-value: {p_value:.4f}')\n",
    "\n",
    "axs[0].set_ylabel('Predicted Binding Affinity')\n",
    "plt.tight_layout()\n",
    "plt.savefig('binding_affinity_screen_1.jpg', bbox_inches='tight', format='jpg')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 5, figsize=(10, 5))\n",
    "sns.set(style=\"ticks\")\n",
    "names = ['B1MDI3', 'P56817', 'P17931', 'Q8ULI9', 'P01116']\n",
    "\n",
    "for i, set_of_lists in enumerate(lists_2):\n",
    "    name = names[i]\n",
    "    statistic, p_value = ranksums(set_of_lists[0], set_of_lists[1], alternative='greater')\n",
    "\n",
    "    sns.boxplot(data=[set_of_lists[0], set_of_lists[1]], ax=axs[i])\n",
    "    axs[i].set_xticks([0, 1])\n",
    "    axs[i].set_xticklabels(['Binders', 'Decoys'])\n",
    "    axs[i].set_title(f'{name}\\np-value: {p_value:.4f}')\n",
    "\n",
    "axs[0].set_ylabel('Predicted Binding Affinity')\n",
    "plt.tight_layout()\n",
    "plt.savefig('binding_affinity_screen_2.jpg', bbox_inches='tight', format='jpg')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cross-docking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from tqdm import tqdm\n",
    "from commons.utils import read_molecule, reorder_atoms, get_symmetry_rmsd, rigid_transform_Kabsch_3D\n",
    "from copy import deepcopy\n",
    "from rdkit.Chem import RemoveHs\n",
    "from rdkit.Geometry import Point3D\n",
    "from inference import print_results\n",
    "import os\n",
    "import numpy as np\n",
    "from Bio.Align import PairwiseAligner "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# functions to find indices to remove from sequences\n",
    "def find_extra_indices(reference, input_str):\n",
    "    extra_indices = []\n",
    "    i_ref = 0\n",
    "    i_input = 0\n",
    "\n",
    "    while i_ref < len(reference) and i_input < len(input_str):\n",
    "        if reference[i_ref] == input_str[i_input]:\n",
    "            i_ref += 1\n",
    "            i_input += 1\n",
    "        else:\n",
    "            extra_indices.append(i_input)\n",
    "            i_input += 1\n",
    "\n",
    "    # If there are extra characters at the end of input_str\n",
    "    while i_input < len(input_str):\n",
    "        extra_indices.append(i_input)\n",
    "        i_input += 1\n",
    "\n",
    "    return extra_indices\n",
    "\n",
    "def find_indices_to_remove(seq1, seq2):\n",
    "    # Perform sequence alignment\n",
    "    aligner = PairwiseAligner()\n",
    "    alignments = aligner.align(seqA=seq1, seqB=seq2)\n",
    "\n",
    "    # Get the aligned sequences\n",
    "    aligned_seq1 = alignments[0][0]\n",
    "    aligned_seq2 = alignments[0][1]\n",
    "    \n",
    "    # Find indices of differing amino acids\n",
    "    differing_indices = [i for i, (a1, a2) in enumerate(zip(aligned_seq1, aligned_seq2)) if a1 != a2]\n",
    "\n",
    "    aligned_seq1_clean = ''.join([s for i,s in enumerate(aligned_seq1) if i not in differing_indices])\n",
    "    aligned_seq2_clean = ''.join([s for i,s in enumerate(aligned_seq2) if i not in differing_indices])\n",
    "\n",
    "    assert len(aligned_seq1_clean) <= len(seq1)\n",
    "    assert len(aligned_seq2_clean) <= len(seq2)\n",
    "\n",
    "    indices_to_remove_seq1 = find_extra_indices(aligned_seq1_clean, seq1)\n",
    "    indices_to_remove_seq2 = find_extra_indices(aligned_seq2_clean, seq2)\n",
    "\n",
    "    assert len(seq1) - len(indices_to_remove_seq1) == len(seq2) - len(indices_to_remove_seq2)\n",
    "    \n",
    "    return indices_to_remove_seq1, indices_to_remove_seq2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_proteins = torch.load('data/processed/timesplit_test/rec_input_proc_ids.pt')\n",
    "test_protein_names = [p['complex_names'] for p in test_proteins]\n",
    "test_c_alphas = [p['c_alpha_coords'] for p in test_proteins]\n",
    "test_sequences = [p['sequence'] for p in test_proteins]\n",
    "outputs = torch.load('checkpoints/quickbind_default/predictions-w-single-rep.pt')\n",
    "\n",
    "coms = []\n",
    "for rec in test_proteins:\n",
    "    c_alpha_coords = rec['c_alpha_coords']\n",
    "    if c_alpha_coords.shape[0] > 2000: # inference is currently limited to complexes with less than 2000 residues\n",
    "        continue\n",
    "    c = torch.mean(c_alpha_coords, dim=0)\n",
    "    coms.append(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_crossdocking(binders_list, uid, ref_name, idx=None):\n",
    "    binders_out = torch.load(\n",
    "        f'virtual_screening/true_{uid}_predictions{f\"_{idx}\" if idx else \"\"}-w-single-rep.pt'\n",
    "    )\n",
    "    ref_index = test_protein_names.index(ref_name)\n",
    "    binders_ordered = [n for n in test_protein_names if n in binders_list]\n",
    "    true_c_alphas_ordered = [c for c, n in zip(test_c_alphas, test_protein_names) if n in binders_list]\n",
    "    true_sequences = [s for s, n in zip(test_sequences, test_protein_names) if n in binders_list]\n",
    "    true_coms = [c for c, n in zip(coms, test_protein_names) if n in binders_list]\n",
    "    input_c_alphas = test_c_alphas[ref_index]\n",
    "    input_sequence = test_sequences[ref_index]\n",
    "    input_com = coms[ref_index]\n",
    "    targets = [t for t, n in zip(outputs['targets'], outputs['names']) if n in binders_ordered]\n",
    "\n",
    "    kabsch_rmsds, rmsds, centroid_distances, protein_rmsds = [], [], [], []\n",
    "    for prediction, target, name, sequence, true_c_alphas, com in tqdm(zip(\n",
    "        binders_out['predictions'], targets, binders_ordered, true_sequences, true_c_alphas_ordered, true_coms),\n",
    "        desc='Evaluating model predictions', total = len(binders_out['predictions'])\n",
    "    ):\n",
    "        if name == ref_name:\n",
    "            continue\n",
    "        assert prediction.shape == target.shape\n",
    "        lig = read_molecule(os.path.join('data/PDBBind/', name, f'{name}_ligand.mol2'), remove_hs=True)\n",
    "        if lig is None:\n",
    "            lig = read_molecule(os.path.join('data/PDBBind/', name, f'{name}_ligand.sdf'), remove_hs=True)\n",
    "        lig = RemoveHs(lig)\n",
    "        lig = reorder_atoms(lig)\n",
    "    \n",
    "        lig_pred = deepcopy(lig)\n",
    "        conf = lig_pred.GetConformer()\n",
    "        prediction = prediction.squeeze().cpu().numpy()\n",
    "        for i in range(lig_pred.GetNumAtoms()):\n",
    "            x, y, z = prediction[i]\n",
    "            conf.SetAtomPosition(i, Point3D(float(x), float(y), float(z)))\n",
    "        coords_pred = lig_pred.GetConformer().GetPositions()\n",
    "    \n",
    "        # transform prediction based on alignment of protein C alpha atoms\n",
    "        indices_input, indices_true = find_indices_to_remove(input_sequence, sequence)\n",
    "        keep_input = torch.tensor([i for i in range(len(input_sequence)) if i not in indices_input])\n",
    "        keep_true = torch.tensor([i for i in range(len(sequence)) if i not in indices_true])\n",
    "        updated_input = torch.index_select(input_c_alphas, 0, keep_input).numpy() - input_com.numpy()\n",
    "        updated_true = torch.index_select(true_c_alphas, 0, keep_true).numpy() - com.numpy()\n",
    "        R, t = rigid_transform_Kabsch_3D(updated_input.T, updated_true.T)\n",
    "    \n",
    "        coords_pred = (R @ (coords_pred).T).T + t.squeeze()\n",
    "        transformed_input = (R @ (updated_input).T).T + t.squeeze()\n",
    "        protein_rmsds.append(np.sqrt((((transformed_input - updated_true)** 2)*3).mean()))\n",
    "    \n",
    "        lig_true = deepcopy(lig)\n",
    "        conf_true = lig_true.GetConformer()\n",
    "        for i in range(lig_true.GetNumAtoms()):\n",
    "            x, y, z = target.squeeze()[i]\n",
    "            conf_true.SetAtomPosition(i, Point3D(float(x), float(y), float(z)))\n",
    "        coords_native = lig_true.GetConformer().GetPositions()\n",
    "    \n",
    "        try:\n",
    "            rmsd = get_symmetry_rmsd(lig_true, coords_native, coords_pred, lig_pred)\n",
    "        except Exception as e:\n",
    "            print(\"Using non corrected RMSD because of the error:\", e)\n",
    "            rmsd = np.sqrt(np.sum((coords_pred - coords_native) ** 2, axis=1).mean())\n",
    "        centroid_distance = np.linalg.norm(coords_native.mean(axis=0) - coords_pred.mean(axis=0))\n",
    "        R, t = rigid_transform_Kabsch_3D(coords_pred.T, coords_native.T)\n",
    "        moved_coords = (R @ (coords_pred).T).T + t.squeeze()\n",
    "        kabsch_rmsd = np.sqrt(np.sum((moved_coords - coords_native) ** 2, axis=1).mean())\n",
    "        kabsch_rmsds.append(kabsch_rmsd)\n",
    "        rmsds.append(rmsd)\n",
    "        centroid_distances.append(centroid_distance)\n",
    "    \n",
    "    print(uid)\n",
    "    kabsch_rmsds = np.array(kabsch_rmsds)\n",
    "    print_results(\n",
    "        rmsds, centroid_distances, False\n",
    "    )\n",
    "    print(f'Mean Kabsch RMSD: {kabsch_rmsds.mean().__round__(2)} +- {kabsch_rmsds.std().__round__(2)}')\n",
    "    print(f'Median Kabsch RMSD: {np.median(kabsch_rmsds).__round__(2)} +- {kabsch_rmsds.std().__round__(2)}')\n",
    "    print('Kabsch RMSD percentiles: ', np.percentile(kabsch_rmsds, [25, 50, 75]).round(2))\n",
    "    print('Average protein RMSD', np.mean(protein_rmsds))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions: 100%|██████████| 19/19 [00:00<00:00, 140.75it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B1MDI3\n",
      "----------------------------------------------------------------------------------------------------\n",
      "|                              Test statistics (excl. hydrogen atoms)                              |\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Mean RMSD: 7.36 +- 2.02\n",
      "RMSD percentiles:  [6.   6.75 8.64]\n",
      "% RMSD below 2: 0.0%\n",
      "% RMSD below 5: 5.56%\n",
      "Mean centroid distance: 5.21 +- 2.67\n",
      "Centroid percentiles:  [3.53 4.65 6.2 ]\n",
      "% centroid distances below 2: 11.11%\n",
      "% centroid distances below 5: 55.56%\n",
      "Mean Kabsch RMSD: 1.56 +- 0.34\n",
      "Median Kabsch RMSD: 1.5 +- 0.34\n",
      "Kabsch RMSD percentiles:  [1.33 1.5  1.79]\n",
      "Average protein RMSD 0.2936543735706902\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "B1MDI3_binders = [\n",
    "    '6qqt', '6qrf', '6qqw', '6qre', '6qrc', '6qqv', '6qrg', '6qr1', '6qqq', '6qqu', '6qr2', '6qra', '6qr4', '6qr3',\n",
    "    '6qr9', '6qqz', '6qr0', '6qrd', '6qr7'\n",
    "]\n",
    "evaluate_crossdocking(B1MDI3_binders, 'B1MDI3', '6qqt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions: 100%|██████████| 19/19 [00:00<00:00, 137.22it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B1MDI3\n",
      "----------------------------------------------------------------------------------------------------\n",
      "|                              Test statistics (excl. hydrogen atoms)                              |\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Mean RMSD: 11.35 +- 2.4\n",
      "RMSD percentiles:  [10.41 11.99 13.06]\n",
      "% RMSD below 2: 0.0%\n",
      "% RMSD below 5: 0.0%\n",
      "Mean centroid distance: 10.35 +- 3.06\n",
      "Centroid percentiles:  [ 9.13 11.44 12.74]\n",
      "% centroid distances below 2: 0.0%\n",
      "% centroid distances below 5: 5.56%\n",
      "Mean Kabsch RMSD: 1.63 +- 0.41\n",
      "Median Kabsch RMSD: 1.5 +- 0.41\n",
      "Kabsch RMSD percentiles:  [1.43 1.5  1.99]\n",
      "Average protein RMSD 0.4883690880740943\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "B1MDI3_binders = [\n",
    "    '6qqt', '6qrf', '6qqw', '6qre', '6qrc', '6qqv', '6qrg', '6qr1', '6qqq', '6qqu', '6qr2', '6qra', '6qr4', '6qr3',\n",
    "    '6qr9', '6qqz', '6qr0', '6qrd', '6qr7'\n",
    "]\n",
    "evaluate_crossdocking(B1MDI3_binders, 'B1MDI3', '6qrf', idx=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions: 100%|██████████| 19/19 [00:00<00:00, 97.92it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B1MDI3\n",
      "----------------------------------------------------------------------------------------------------\n",
      "|                              Test statistics (excl. hydrogen atoms)                              |\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Mean RMSD: 9.96 +- 2.95\n",
      "RMSD percentiles:  [ 7.23  9.89 12.89]\n",
      "% RMSD below 2: 0.0%\n",
      "% RMSD below 5: 0.0%\n",
      "Mean centroid distance: 8.4 +- 3.66\n",
      "Centroid percentiles:  [ 5.07  8.4  12.03]\n",
      "% centroid distances below 2: 0.0%\n",
      "% centroid distances below 5: 27.78%\n",
      "Mean Kabsch RMSD: 1.61 +- 0.38\n",
      "Median Kabsch RMSD: 1.6 +- 0.38\n",
      "Kabsch RMSD percentiles:  [1.35 1.6  1.87]\n",
      "Average protein RMSD 0.3002542470502205\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "B1MDI3_binders = [\n",
    "    '6qqt', '6qrf', '6qqw', '6qre', '6qrc', '6qqv', '6qrg', '6qr1', '6qqq', '6qqu', '6qr2', '6qra', '6qr4', '6qr3',\n",
    "    '6qr9', '6qqz', '6qr0', '6qrd', '6qr7'\n",
    "]\n",
    "evaluate_crossdocking(B1MDI3_binders, 'B1MDI3', '6qqw', idx=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions: 100%|██████████| 16/16 [00:00<00:00, 83.94it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "P56817\n",
      "----------------------------------------------------------------------------------------------------\n",
      "|                              Test statistics (excl. hydrogen atoms)                              |\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Mean RMSD: 2.13 +- 1.57\n",
      "RMSD percentiles:  [1.5  1.58 1.98]\n",
      "% RMSD below 2: 73.33%\n",
      "% RMSD below 5: 93.33%\n",
      "Mean centroid distance: 1.08 +- 0.51\n",
      "Centroid percentiles:  [0.69 0.91 1.22]\n",
      "% centroid distances below 2: 93.33%\n",
      "% centroid distances below 5: 100.0%\n",
      "Mean Kabsch RMSD: 1.12 +- 0.39\n",
      "Median Kabsch RMSD: 1.19 +- 0.39\n",
      "Kabsch RMSD percentiles:  [0.9  1.19 1.35]\n",
      "Average protein RMSD 1.0027781721654772\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "P56817_binders = [\n",
    "    '6uvp', '6uvv', '6uvy', '6uwp', '6nw3', '6e3z', '6nv7', '6uwv', '6nv9', '6od6', '6jt3', '6jsg',\n",
    "    '6jsn', '6jsf', '6jse', '6pz4'\n",
    "]\n",
    "evaluate_crossdocking(P56817_binders, 'P56817', '6uvp')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions: 100%|██████████| 16/16 [00:00<00:00, 85.16it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "P56817\n",
      "----------------------------------------------------------------------------------------------------\n",
      "|                              Test statistics (excl. hydrogen atoms)                              |\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Mean RMSD: 2.09 +- 1.52\n",
      "RMSD percentiles:  [1.41 1.52 2.06]\n",
      "% RMSD below 2: 73.33%\n",
      "% RMSD below 5: 93.33%\n",
      "Mean centroid distance: 0.93 +- 0.37\n",
      "Centroid percentiles:  [0.72 0.85 0.98]\n",
      "% centroid distances below 2: 100.0%\n",
      "% centroid distances below 5: 100.0%\n",
      "Mean Kabsch RMSD: 1.18 +- 0.4\n",
      "Median Kabsch RMSD: 1.17 +- 0.4\n",
      "Kabsch RMSD percentiles:  [0.92 1.17 1.44]\n",
      "Average protein RMSD 0.9673100870998292\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "P56817_binders = [\n",
    "    '6uvp', '6uvv', '6uvy', '6uwp', '6nw3', '6e3z', '6nv7', '6uwv', '6nv9', '6od6', '6jt3', '6jsg',\n",
    "    '6jsn', '6jsf', '6jse', '6pz4'\n",
    "]\n",
    "evaluate_crossdocking(P56817_binders, 'P56817', '6uvv', idx=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions: 100%|██████████| 16/16 [00:00<00:00, 95.69it/s] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "P56817\n",
      "----------------------------------------------------------------------------------------------------\n",
      "|                              Test statistics (excl. hydrogen atoms)                              |\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Mean RMSD: 2.15 +- 1.64\n",
      "RMSD percentiles:  [1.44 1.54 2.09]\n",
      "% RMSD below 2: 73.33%\n",
      "% RMSD below 5: 93.33%\n",
      "Mean centroid distance: 0.9 +- 0.41\n",
      "Centroid percentiles:  [0.6  0.83 0.94]\n",
      "% centroid distances below 2: 100.0%\n",
      "% centroid distances below 5: 100.0%\n",
      "Mean Kabsch RMSD: 1.17 +- 0.41\n",
      "Median Kabsch RMSD: 1.18 +- 0.41\n",
      "Kabsch RMSD percentiles:  [0.89 1.18 1.4 ]\n",
      "Average protein RMSD 0.9086242446514984\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "P56817_binders = [\n",
    "    '6uvp', '6uvv', '6uvy', '6uwp', '6nw3', '6e3z', '6nv7', '6uwv', '6nv9', '6od6', '6jt3', '6jsg',\n",
    "    '6jsn', '6jsf', '6jse', '6pz4'\n",
    "]\n",
    "evaluate_crossdocking(P56817_binders, 'P56817', '6uvy', idx=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions: 100%|██████████| 15/15 [00:00<00:00, 89.94it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "P17931\n",
      "----------------------------------------------------------------------------------------------------\n",
      "|                              Test statistics (excl. hydrogen atoms)                              |\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Mean RMSD: 2.19 +- 2.47\n",
      "RMSD percentiles:  [1.33 1.46 1.72]\n",
      "% RMSD below 2: 85.71%\n",
      "% RMSD below 5: 92.86%\n",
      "Mean centroid distance: 0.88 +- 0.32\n",
      "Centroid percentiles:  [0.64 0.82 0.99]\n",
      "% centroid distances below 2: 100.0%\n",
      "% centroid distances below 5: 100.0%\n",
      "Mean Kabsch RMSD: 1.22 +- 0.47\n",
      "Median Kabsch RMSD: 1.04 +- 0.47\n",
      "Kabsch RMSD percentiles:  [1.   1.04 1.22]\n",
      "Average protein RMSD 0.16096047345873035\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "P17931_binders = [\n",
    "    '6qlt', '6qlq', '6i75', '6qlu', '6qln', '6i78', '6qlr', '6i77', '6qlo', '6i76', '6qlp', '6qls',\n",
    "    '6i74', '6qge', '6qgf'\n",
    "]\n",
    "evaluate_crossdocking(P17931_binders, 'P17931', '6qlt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions: 100%|██████████| 15/15 [00:00<00:00, 91.47it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "P17931\n",
      "----------------------------------------------------------------------------------------------------\n",
      "|                              Test statistics (excl. hydrogen atoms)                              |\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Mean RMSD: 2.25 +- 2.44\n",
      "RMSD percentiles:  [1.44 1.53 1.73]\n",
      "% RMSD below 2: 85.71%\n",
      "% RMSD below 5: 92.86%\n",
      "Mean centroid distance: 0.96 +- 0.25\n",
      "Centroid percentiles:  [0.76 0.94 1.07]\n",
      "% centroid distances below 2: 100.0%\n",
      "% centroid distances below 5: 100.0%\n",
      "Mean Kabsch RMSD: 1.22 +- 0.47\n",
      "Median Kabsch RMSD: 1.05 +- 0.47\n",
      "Kabsch RMSD percentiles:  [1.   1.05 1.21]\n",
      "Average protein RMSD 0.166106119457376\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "P17931_binders = [\n",
    "    '6qlt', '6qlq', '6i75', '6qlu', '6qln', '6i78', '6qlr', '6i77', '6qlo', '6i76', '6qlp', '6qls',\n",
    "    '6i74', '6qge', '6qgf'\n",
    "]\n",
    "evaluate_crossdocking(P17931_binders, 'P17931', '6qlq', idx=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions: 100%|██████████| 15/15 [00:00<00:00, 94.57it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "P17931\n",
      "----------------------------------------------------------------------------------------------------\n",
      "|                              Test statistics (excl. hydrogen atoms)                              |\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Mean RMSD: 2.24 +- 2.46\n",
      "RMSD percentiles:  [1.4  1.5  1.75]\n",
      "% RMSD below 2: 85.71%\n",
      "% RMSD below 5: 92.86%\n",
      "Mean centroid distance: 0.89 +- 0.29\n",
      "Centroid percentiles:  [0.71 0.86 1.02]\n",
      "% centroid distances below 2: 100.0%\n",
      "% centroid distances below 5: 100.0%\n",
      "Mean Kabsch RMSD: 1.2 +- 0.48\n",
      "Median Kabsch RMSD: 1.03 +- 0.48\n",
      "Kabsch RMSD percentiles:  [0.99 1.03 1.17]\n",
      "Average protein RMSD 0.31159448404717477\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "P17931_binders = [\n",
    "    '6qlt', '6qlq', '6i75', '6qlu', '6qln', '6i78', '6qlr', '6i77', '6qlo', '6i76', '6qlp', '6qls',\n",
    "    '6i74', '6qge', '6qgf'\n",
    "]\n",
    "evaluate_crossdocking(P17931_binders, 'P17931', '6i75', idx=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions:   0%|          | 0/14 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions: 100%|██████████| 14/14 [00:00<00:00, 76.31it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Q8ULI9\n",
      "----------------------------------------------------------------------------------------------------\n",
      "|                              Test statistics (excl. hydrogen atoms)                              |\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Mean RMSD: 0.99 +- 0.14\n",
      "RMSD percentiles:  [0.92 0.97 1.07]\n",
      "% RMSD below 2: 100.0%\n",
      "% RMSD below 5: 100.0%\n",
      "Mean centroid distance: 0.19 +- 0.07\n",
      "Centroid percentiles:  [0.12 0.16 0.24]\n",
      "% centroid distances below 2: 100.0%\n",
      "% centroid distances below 5: 100.0%\n",
      "Mean Kabsch RMSD: 0.94 +- 0.14\n",
      "Median Kabsch RMSD: 0.94 +- 0.14\n",
      "Kabsch RMSD percentiles:  [0.89 0.94 0.99]\n",
      "Average protein RMSD 0.14522072243775053\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "Q8ULI9_binders = [\n",
    "    '6oy0', '6oxy', '6oxz', '6oy2', '6oxt', '6oxs', '6oxx', '6oxu', '6oy1', '6oxr', '6oxw', '6oxv',\n",
    "    '6oxp', '6oxq'\n",
    "]\n",
    "evaluate_crossdocking(Q8ULI9_binders, 'Q8ULI9', '6oy0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions: 100%|██████████| 14/14 [00:00<00:00, 73.56it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Q8ULI9\n",
      "----------------------------------------------------------------------------------------------------\n",
      "|                              Test statistics (excl. hydrogen atoms)                              |\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Mean RMSD: 0.99 +- 0.14\n",
      "RMSD percentiles:  [0.91 0.97 1.06]\n",
      "% RMSD below 2: 100.0%\n",
      "% RMSD below 5: 100.0%\n",
      "Mean centroid distance: 0.2 +- 0.1\n",
      "Centroid percentiles:  [0.11 0.18 0.27]\n",
      "% centroid distances below 2: 100.0%\n",
      "% centroid distances below 5: 100.0%\n",
      "Mean Kabsch RMSD: 0.94 +- 0.13\n",
      "Median Kabsch RMSD: 0.94 +- 0.13\n",
      "Kabsch RMSD percentiles:  [0.9  0.94 0.98]\n",
      "Average protein RMSD 0.14797126746950764\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "Q8ULI9_binders = [\n",
    "    '6oy0', '6oxy', '6oxz', '6oy2', '6oxt', '6oxs', '6oxx', '6oxu', '6oy1', '6oxr', '6oxw', '6oxv',\n",
    "    '6oxp', '6oxq'\n",
    "]\n",
    "evaluate_crossdocking(Q8ULI9_binders, 'Q8ULI9', '6oxy', idx=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions: 100%|██████████| 14/14 [00:00<00:00, 84.27it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Q8ULI9\n",
      "----------------------------------------------------------------------------------------------------\n",
      "|                              Test statistics (excl. hydrogen atoms)                              |\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Mean RMSD: 1.0 +- 0.15\n",
      "RMSD percentiles:  [0.93 1.01 1.06]\n",
      "% RMSD below 2: 100.0%\n",
      "% RMSD below 5: 100.0%\n",
      "Mean centroid distance: 0.22 +- 0.09\n",
      "Centroid percentiles:  [0.17 0.23 0.27]\n",
      "% centroid distances below 2: 100.0%\n",
      "% centroid distances below 5: 100.0%\n",
      "Mean Kabsch RMSD: 0.95 +- 0.14\n",
      "Median Kabsch RMSD: 0.95 +- 0.14\n",
      "Kabsch RMSD percentiles:  [0.86 0.95 0.99]\n",
      "Average protein RMSD 0.16909177256970379\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "Q8ULI9_binders = [\n",
    "    '6oy0', '6oxy', '6oxz', '6oy2', '6oxt', '6oxs', '6oxx', '6oxu', '6oy1', '6oxr', '6oxw', '6oxv',\n",
    "    '6oxp', '6oxq'\n",
    "]\n",
    "evaluate_crossdocking(Q8ULI9_binders, 'Q8ULI9', '6oxz', idx=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions: 100%|██████████| 13/13 [00:00<00:00, 75.51it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "P01116\n",
      "----------------------------------------------------------------------------------------------------\n",
      "|                              Test statistics (excl. hydrogen atoms)                              |\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Mean RMSD: 13.0 +- 5.41\n",
      "RMSD percentiles:  [ 8.21 15.32 17.76]\n",
      "% RMSD below 2: 0.0%\n",
      "% RMSD below 5: 8.33%\n",
      "Mean centroid distance: 11.27 +- 5.92\n",
      "Centroid percentiles:  [ 4.71 14.1  16.68]\n",
      "% centroid distances below 2: 0.0%\n",
      "% centroid distances below 5: 33.33%\n",
      "Mean Kabsch RMSD: 1.97 +- 0.43\n",
      "Median Kabsch RMSD: 1.91 +- 0.43\n",
      "Kabsch RMSD percentiles:  [1.63 1.91 2.38]\n",
      "Average protein RMSD 9.57859230944588\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "P01116_binders = [\n",
    "    '6quw', '6quv', '6gj5', '6v5l', '6gj6', '6p8z', '6pgo', '6p8x', '6gj8', '6gj7', '6p8y', '6pgp',\n",
    "    '6oim'\n",
    "]\n",
    "evaluate_crossdocking(P01116_binders, 'P01116', '6quw')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions:   0%|          | 0/13 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions: 100%|██████████| 13/13 [00:00<00:00, 83.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "P01116\n",
      "----------------------------------------------------------------------------------------------------\n",
      "|                              Test statistics (excl. hydrogen atoms)                              |\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Mean RMSD: 12.99 +- 5.4\n",
      "RMSD percentiles:  [ 7.99 15.23 17.78]\n",
      "% RMSD below 2: 0.0%\n",
      "% RMSD below 5: 8.33%\n",
      "Mean centroid distance: 11.29 +- 5.87\n",
      "Centroid percentiles:  [ 4.94 14.08 16.64]\n",
      "% centroid distances below 2: 0.0%\n",
      "% centroid distances below 5: 25.0%\n",
      "Mean Kabsch RMSD: 2.03 +- 0.39\n",
      "Median Kabsch RMSD: 1.94 +- 0.39\n",
      "Kabsch RMSD percentiles:  [1.81 1.94 2.42]\n",
      "Average protein RMSD 9.574189161626935\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "P01116_binders = [\n",
    "    '6quw', '6quv', '6gj5', '6v5l', '6gj6', '6p8z', '6pgo', '6p8x', '6gj8', '6gj7', '6p8y', '6pgp',\n",
    "    '6oim'\n",
    "]\n",
    "evaluate_crossdocking(P01116_binders, 'P01116', '6quv', idx=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating model predictions: 100%|██████████| 13/13 [00:00<00:00, 104.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "P01116\n",
      "----------------------------------------------------------------------------------------------------\n",
      "|                              Test statistics (excl. hydrogen atoms)                              |\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Mean RMSD: 12.46 +- 6.27\n",
      "RMSD percentiles:  [ 4.95 14.69 18.41]\n",
      "% RMSD below 2: 0.0%\n",
      "% RMSD below 5: 33.33%\n",
      "Mean centroid distance: 10.74 +- 6.69\n",
      "Centroid percentiles:  [ 2.4  13.77 16.81]\n",
      "% centroid distances below 2: 16.67%\n",
      "% centroid distances below 5: 33.33%\n",
      "Mean Kabsch RMSD: 1.92 +- 0.64\n",
      "Median Kabsch RMSD: 1.69 +- 0.64\n",
      "Kabsch RMSD percentiles:  [1.42 1.69 2.19]\n",
      "Average protein RMSD 3.855804703990517\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "P01116_binders = [\n",
    "    '6quw', '6quv', '6gj5', '6v5l', '6gj6', '6p8z', '6pgo', '6p8x', '6gj8', '6gj7', '6p8y', '6pgp',\n",
    "    '6oim'\n",
    "]\n",
    "evaluate_crossdocking(P01116_binders, 'P01116', '6gj5', idx=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "alphabind_new",
   "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.10.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}