Uni_Dock+OpenBPMD.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/pablo-arantes/Cloud-Bind/blob/main/Uni_Dock%2BOpenBPMD.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pj2BWZxUDbEE"
      },
      "source": [
        "# **Hi there!**\n",
        "\n",
        "This is a Jupyter notebook for running molecular docking calculations with Uni-Dock, a GPU-accelerated molecular docking program, and for running OpenBPMD, to evaluate the ligand pose stability using metadynamics. OpenBPMD is powered by the OpenMM simulation engine and uses a revised scoring function.\n",
        "\n",
        "The main goal of this notebook is to demonstrate how to harness the power of cloud-computing to perform drug binding structure prediction in a cheap and yet feasible fashion.\n",
        "\n",
        "---\n",
        "\n",
        " **This notebook is NOT a standard protocol for docking calculations and MD simulations!** It is just simple docking pipeline illustrating each step of a docking and MD protocol.\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "**Bugs**\n",
        "- If you encounter any bugs, please report the issue to https://github.com/pablo-arantes/Cloud-Bind/issues\n",
        "\n",
        "**Acknowledgments**\n",
        "- We would like to thank the [Uni-Dock](https://github.com/dptech-corp/Uni-Dock) team for doing an excellent job open sourcing the software.\n",
        "- We would like to thank the [Roitberg](https://roitberg.chem.ufl.edu/) team for developing the fantastic [TorchANI](https://github.com/aiqm/torchani).\n",
        "- We would like to thank [@ruiz_moreno_aj](https://twitter.com/ruiz_moreno_aj) for his work on [Jupyter Dock](https://github.com/AngelRuizMoreno/Jupyter_Dock)\n",
        "- We would like to thank the ChemosimLab ([@ChemosimLab](https://twitter.com/ChemosimLab)) team for their incredible [ProLIF](https://prolif.readthedocs.io/en/latest/index.html#) (Protein-Ligand Interaction Fingerprints) tool.\n",
        "- We would like to thank the [OpenBPMD](https://github.com/Gervasiolab/OpenBPMD) team for their open source implementation of binding pose metadynamics (BPMD).\n",
        "- Also, credit to [David Koes](https://github.com/dkoes) for his awesome [py3Dmol](https://3dmol.csb.pitt.edu/) plugin.\n",
        "- Finally, we would like to thank [Making it rain](https://github.com/pablo-arantes/making-it-rain) team, **Pablo R. Arantes** ([@pablitoarantes](https://twitter.com/pablitoarantes)), **Marcelo D. Polêto** ([@mdpoleto](https://twitter.com/mdpoleto)), **Conrado Pedebos** ([@ConradoPedebos](https://twitter.com/ConradoPedebos)) and **Rodrigo Ligabue-Braun** ([@ligabue_braun](https://twitter.com/ligabue_braun)), for their amazing work.\n",
        "- A Cloud-Bind by **Pablo R. Arantes** ([@pablitoarantes](https://twitter.com/pablitoarantes))\n",
        "\n",
        "- For related notebooks see: [Cloud-Bind](https://github.com/pablo-arantes/Cloud-Bind)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Install Conda Colab**\n",
        "#@markdown It will restart the kernel (session), don't worry.\n",
        "!pip install -q condacolab\n",
        "import condacolab\n",
        "condacolab.install()"
      ],
      "metadata": {
        "id": "k5xm80qcWGWd",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wH1oMiVUlxO_",
        "cellView": "form"
      },
      "source": [
        "#@title **Install dependencies**\n",
        "#@markdown It will take a few minutes, please, drink a coffee and wait. ;-)\n",
        "# install dependencies\n",
        "%%capture\n",
        "import sys\n",
        "import tarfile\n",
        "import os\n",
        "import subprocess\n",
        "import sys\n",
        "\n",
        "commands = [\n",
        "    \"rm -rf /usr/local/conda-meta/pinned\",\n",
        "    \"pip -q install py3Dmol\",\n",
        "    \"pip install git+https://github.com/pablo-arantes/biopandas\",\n",
        "    \"pip install bio\",\n",
        "    \"mamba install openmm=7.7.0\",\n",
        "    \"mamba install openmmforcefields -c conda-forge -y\",\n",
        "    \"pip install prolif==1.1.0\",\n",
        "    \"mamba install pytorch torchvision -c pytorch\",\n",
        "    \"pip install torchani\",\n",
        "    \"pip install ase\",\n",
        "    \"mamba install -c conda-forge pdbfixer -y\",\n",
        "    \"mamba install -c conda-forge parmed -y\",\n",
        "    \"mamba install -c conda-forge openbabel -y\",\n",
        "    \"pip install --upgrade MDAnalysis\",\n",
        "    \"wget https://github.com/dptech-corp/Uni-Dock/releases/download/1.0.0/unidock\",\n",
        "    \"chmod +x unidock\",\n",
        "    \"wget https://github.com/gnina/gnina/releases/download/v1.0.3/gnina\",\n",
        "    \"chmod +x gnina\",\n",
        "    \"wget https://github.com/rdk/p2rank/releases/download/2.4/p2rank_2.4.tar.gz\"\n",
        "]\n",
        "\n",
        "for cmd in commands:\n",
        "    subprocess.run(cmd, shell=True)\n",
        "\n",
        "file = tarfile.open('p2rank_2.4.tar.gz')\n",
        "file.extractall('/content/')\n",
        "file.close()\n",
        "os.remove('p2rank_2.4.tar.gz')\n",
        "\n",
        "#load dependencies\n",
        "from openmm import app, unit\n",
        "from openmm.app import HBonds, NoCutoff, PDBFile\n",
        "from openff.toolkit.topology import Molecule, Topology\n",
        "from openff.toolkit.typing.engines.smirnoff import ForceField\n",
        "from openff.toolkit.utils import get_data_file_path\n",
        "import parmed as pmd\n",
        "from biopandas.pdb import PandasPdb\n",
        "import openmm as mm\n",
        "from openmm import *\n",
        "from openmm.app import *\n",
        "from openmm.unit import *\n",
        "import urllib.request\n",
        "import numpy as np\n",
        "import MDAnalysis as mda\n",
        "import py3Dmol\n",
        "import pytraj as pt\n",
        "import platform\n",
        "import scipy.cluster.hierarchy\n",
        "from scipy.spatial.distance import squareform\n",
        "import scipy.stats as stats\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "from scipy.interpolate import griddata\n",
        "import seaborn as sb\n",
        "from statistics import mean, stdev\n",
        "from pytraj import matrix\n",
        "from matplotlib import colors\n",
        "from IPython.display import set_matplotlib_formats"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fDQnAKJLFxtt"
      },
      "source": [
        "## Using Google Drive to store simulation data\n",
        "\n",
        "Google Colab does not allow users to keep data on their computing nodes. However, we can use Google Drive to read, write, and store our simulations files. Therefore, we suggest to you to:\n",
        "\n",
        "1.   Create a folder in your own Google Drive and copy the necessary input files there.\n",
        "2.   Copy the path of your created directory. We will use it below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Lm7Akepv_vl-",
        "cellView": "form"
      },
      "source": [
        "#@title ### **Import Google Drive**\n",
        "#@markdown Click in the \"Run\" buttom to make your Google Drive accessible.\n",
        "from google.colab import drive\n",
        "\n",
        "drive.mount('/content/drive', force_remount=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "n96Vyb3mMemN"
      },
      "source": [
        "#@title **Check if you correctly allocated GPU nodes**\n",
        "\n",
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Select the Runtime > \"Change runtime type\" menu to enable a GPU accelerator, ')\n",
        "  print('and then re-execute this cell.')\n",
        "else:\n",
        "  print(gpu_info)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BZwl66HTGI7v"
      },
      "source": [
        "\n",
        "---\n",
        "# **Loading the necessary input files**\n",
        "\n",
        "At this point, we should have all libraries and dependencies installed.\n",
        "\n",
        "**Important**: Make sure the PDB file points to the correct structure.\n",
        "\n",
        "Below, you should provide the names of all input files and the pathway of your Google Drive folder containing them.\n",
        "\n",
        "**Please, don't use spaces in the files and folders names, i.e. MyDrive/protein_ligand and so on.**"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Please, provide the necessary input files below for receptor**:\n",
        "#@markdown **Important:** Run the cell to prepare your receptor and select your reference residue for the construction of an optimal box size for the docking calculations.\n",
        "from openmm.app.pdbfile import PDBFile\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "import os\n",
        "from Bio.PDB import PDBParser, PDBIO, Select\n",
        "from Bio.PDB import is_aa\n",
        "import pandas as pd\n",
        "from pdbfixer import PDBFixer\n",
        "from openbabel import pybel\n",
        "\n",
        "Google_Drive_Path = '/content/' #@param {type:\"string\"}\n",
        "workDir = Google_Drive_Path\n",
        "\n",
        "workDir2 = os.path.join(workDir)\n",
        "workDir_check = os.path.exists(workDir2)\n",
        "if workDir_check == False:\n",
        "  os.mkdir(workDir2)\n",
        "else:\n",
        "  pass\n",
        "\n",
        "if os.path.exists(os.path.join(workDir, \"name_residues.txt\")):\n",
        "  os.remove(os.path.join(workDir, \"name_residues.txt\"))\n",
        "  os.remove(os.path.join(workDir,\"name_residues_receptor.txt\"))\n",
        "else:\n",
        "  pass\n",
        "\n",
        "temp = os.path.join(workDir, \"temp.pdb\")\n",
        "receptor = os.path.join(workDir, \"receptor.pdb\")\n",
        "receptor_pdbqt = os.path.join(workDir, \"receptor.pdbqt\")\n",
        "ligand = os.path.join(workDir, \"ligand.pdbqt\")\n",
        "\n",
        "\n",
        "Query_PDB_ID = '3HTB' #@param {type:\"string\"}\n",
        "\n",
        "pdbfn = Query_PDB_ID + \".pdb\"\n",
        "url = 'https://files.rcsb.org/download/' + pdbfn\n",
        "outfnm = os.path.join(workDir, pdbfn)\n",
        "urllib.request.urlretrieve(url, outfnm)\n",
        "\n",
        "\n",
        "ppdb = PandasPdb().read_pdb(outfnm)\n",
        "ppdb.df['ATOM'] = ppdb.df['ATOM']\n",
        "ppdb.df['HETATM'] = ppdb.df['HETATM'][ppdb.df['HETATM']['residue_name'] != 'HOH']\n",
        "ppdb.to_pdb(path=temp, records=['ATOM', 'HETATM'], gz=False, append_newline=True)\n",
        "\n",
        "#prepare receptor\n",
        "ppdb = PandasPdb().read_pdb(outfnm)\n",
        "ppdb.df['ATOM'] = ppdb.df['ATOM']\n",
        "ppdb.df['HETATM'] = ppdb.df['HETATM'][ppdb.df['HETATM']['residue_name'] != 'HOH']\n",
        "ppdb.df['ATOM'] = ppdb.df['ATOM'][ppdb.df['ATOM']['atom_name'] != 'OXT']\n",
        "ppdb.df['ATOM']= ppdb.df['ATOM'][ppdb.df['ATOM']['element_symbol'] != 'H']\n",
        "ppdb.to_pdb(path=receptor, records=['ATOM', 'HETATM'], gz=False, append_newline=True)\n",
        "\n",
        "fixer = PDBFixer(filename=receptor)\n",
        "fixer.removeHeterogens()\n",
        "fixer.findMissingResidues()\n",
        "fixer.findMissingAtoms()\n",
        "fixer.addMissingAtoms()\n",
        "fixer.addMissingHydrogens(pH=7.4)\n",
        "PDBFile.writeFile(fixer.topology, fixer.positions, open(receptor, 'w'))\n",
        "\n",
        "path = '/content/'\n",
        "\n",
        "\n",
        "def is_het(residue):\n",
        "    res = residue.id[0]\n",
        "    return res != \" \" and res != \"W\"\n",
        "\n",
        "def aa(residue):\n",
        "    res = residue.id[0]\n",
        "    return res != \"W\"\n",
        "\n",
        "\n",
        "class ResidueSelect(Select):\n",
        "    def __init__(self, chain, residue):\n",
        "        self.chain = chain\n",
        "        self.residue = residue\n",
        "\n",
        "    def accept_chain(self, chain):\n",
        "        return chain.id == self.chain.id\n",
        "\n",
        "    def accept_residue(self, residue):\n",
        "        return residue == self.residue and aa(residue)\n",
        "\n",
        "def extract_ligands(path):\n",
        "    pdb = PDBParser().get_structure(temp, temp)\n",
        "    io = PDBIO()\n",
        "    io.set_structure(pdb)\n",
        "    i = 1\n",
        "    name_residues = []\n",
        "    for model in pdb:\n",
        "      for chain in model:\n",
        "        for residue in chain:\n",
        "          if not aa(residue):\n",
        "            continue\n",
        "          # print(f\"{chain[i].resname} {i}\")\n",
        "          name_residues.append(residue)\n",
        "          print((f\"saving {residue}\"), file=open(os.path.join(workDir, \"name_residues.txt\"), \"a\",))\n",
        "          i += 1\n",
        "\n",
        "extract_ligands(path)\n",
        "\n",
        "def extract_ligands2(path):\n",
        "    pdb = PDBParser().get_structure(receptor, receptor)\n",
        "    io = PDBIO()\n",
        "    io.set_structure(pdb)\n",
        "    i2 = 1\n",
        "    name_residues2 = []\n",
        "    for model in pdb:\n",
        "      for chain in model:\n",
        "        for residue in chain:\n",
        "          if not aa(residue):\n",
        "            continue\n",
        "          # print(f\"{chain[i].resname} {i}\")\n",
        "          name_residues2.append(residue)\n",
        "          print((f\"saving {residue}\"), file=open(os.path.join(workDir, \"name_residues_receptor.txt\"), \"a\",))\n",
        "          i2 += 1\n",
        "\n",
        "extract_ligands2(path)\n",
        "\n",
        "# mol= [m for m in pybel.readfile(filename=receptor, format='pdb')][0]\n",
        "# out=pybel.Outputfile(filename=receptor_pdbqt,format='pdbqt',overwrite=True)\n",
        "# pybel.Molecule()\n",
        "# out.write(mol)\n",
        "# out.close()\n",
        "\n",
        "os.system(\"obabel -i pdb \" + str(receptor) + \" -o pdbqt -O \" + str(receptor_pdbqt) + \" -xr --partialcharge\")\n",
        "\n",
        "dataset = pd.read_csv(os.path.join(workDir, 'name_residues.txt'), delimiter = \" \", header=None)\n",
        "df = pd.DataFrame(dataset)\n",
        "df = df.iloc[:, [2]]\n",
        "new = df.to_numpy()\n",
        "\n",
        "dataset2 = pd.read_csv(os.path.join(workDir, 'name_residues_receptor.txt'), delimiter = \" \", header=None)\n",
        "df2 = pd.DataFrame(dataset2)\n",
        "df2 = df2.iloc[:, [2]]\n",
        "new2 = df2.to_numpy()\n",
        "\n",
        "b = 1\n",
        "res_number = []\n",
        "for j in new2:\n",
        "  res_number.append(b)\n",
        "  b += 1\n",
        "\n",
        "print(\"Residue\" + \" - \"  + \"Number\" )\n",
        "a = 1\n",
        "for j in new:\n",
        "  print(', '.join(j) + \" - \"  + str(a))\n",
        "  a += 1"
      ],
      "metadata": {
        "id": "NxmwlxdxlyiD",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Predict ligand-binding pockets from your protein structure using P2Rank**:\n",
        "#@markdown **P2Rank** is a stand-alone command line program that predicts ligand-binding pockets from a protein structure. It achieves high prediction success rates without relying on an external software for computation of complex features or on a database of known protein-ligand templates.\n",
        "#@markdown P2Rank makes predictions by scoring and clustering points on the protein's solvent accessible surface. Ligandability score of individual points is determined by a machine learning based model trained on the dataset of known protein-ligand complexes. For more details see [here](https://github.com/rdk/p2rank).\n",
        "\n",
        "import subprocess\n",
        "import csv\n",
        "\n",
        "output_p2rank = os.path.join(workDir, \"output_p2rank\")\n",
        "p2rank = \"/content/p2rank_2.4/prank predict -f \" + str(receptor) + \" -o \" + str(output_p2rank)\n",
        "original_stdout = sys.stdout\n",
        "with open('p2rank.sh', 'w') as f:\n",
        "  sys.stdout = f\n",
        "  print(p2rank)\n",
        "  sys.stdout = original_stdout\n",
        "subprocess.run([\"chmod 700 p2rank.sh\"], shell=True)\n",
        "subprocess.run([\"./p2rank.sh\"], shell=True,)\n",
        "\n",
        "with open(os.path.join(workDir, \"output_p2rank/receptor.pdb_predictions.csv\"), 'r') as file:\n",
        "  csvreader = csv.reader(file)\n",
        "  residue = []\n",
        "  score = []\n",
        "  center_x = []\n",
        "  center_y = []\n",
        "  center_z = []\n",
        "  for row in csvreader:\n",
        "    residue.append(row[9:10])\n",
        "    score.append(row[2:3])\n",
        "    center_x.append(row[6:7])\n",
        "    center_y.append(row[7:8])\n",
        "    center_z.append(row[8:9])\n",
        "\n",
        "for i in range(1,len(residue)):\n",
        "  file = str((residue[i])[0]).split()\n",
        "  score_end = str((score[i])[0]).split()\n",
        "  center_x_end = str((center_x[i])[0]).split()\n",
        "  center_y_end = str((center_y[i])[0]).split()\n",
        "  center_z_end = str((center_z[i])[0]).split()\n",
        "  print(\"Pocket \" + str(i))\n",
        "  print(\"Score = \" + score_end[0])\n",
        "  final_residues = []\n",
        "  for i in range(0,len(file)):\n",
        "    test = file[i]\n",
        "    final_residues.append(int(test[2:]))\n",
        "  print(\"Selected Residues = \" + str(final_residues))\n",
        "  print(\"Center x = \"+ str(center_x_end[0]), \"Center y = \"+ str(center_y_end[0]), \"Center z = \"+ str(center_z_end[0]) + \"\\n\")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "6zS9eKySxq13"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Please, provide the pocket or residue number for the selection**:\n",
        "#@markdown **Important:** The selected pocket or residues will be used as a reference for the construction of an optimal box size for the ligand during the docking. If you want to select more than one residue, please, use comma to separte the numbers (i.e. 147,150,155,160). **Please, DO NOT USE SPACES BETWEEN THEM.**\n",
        "\n",
        "import re\n",
        "import csv\n",
        "\n",
        "if os.path.exists(os.path.join(workDir, \"name_residue.txt\")):\n",
        "  os.remove(os.path.join(workDir, \"name_residue.txt\"))\n",
        "else:\n",
        "  pass\n",
        "\n",
        "# Python code to convert string to list\n",
        "def Convert(string):\n",
        "\tli = list(string.split(\",\"))\n",
        "\treturn li\n",
        "\n",
        "def extract_ligands(path,residues):\n",
        "    pdb = PDBParser().get_structure(temp, temp)\n",
        "    io = PDBIO()\n",
        "    io.set_structure(pdb)\n",
        "    i = 1\n",
        "    name_residues = []\n",
        "    for model in pdb:\n",
        "      for chain in model:\n",
        "        for residue in chain:\n",
        "          if not aa(residue):\n",
        "            continue\n",
        "          if i == int(residues):\n",
        "            # print(residues)\n",
        "            print((f\"saving {residue}\"), file=open(os.path.join(workDir, \"name_residue.txt\"), \"a\",))\n",
        "            io.save(f\"res_{i}_certo.pdb\", ResidueSelect(chain, residue))\n",
        "          i += 1\n",
        "\n",
        "Selection = \"Pocket\" #@param [\"Pocket\", \"Residues\"]\n",
        "\n",
        "number = '2' #@param {type:\"string\"}\n",
        "\n",
        "if Selection == \"Pocket\":\n",
        "  file = str((residue[int(number)])[0]).split()\n",
        "  score_end = str((score[int(number)])[0]).split()\n",
        "  center_x_end = str((center_x[int(number)])[0]).split()\n",
        "  center_y_end = str((center_y[int(number)])[0]).split()\n",
        "  center_z_end = str((center_z[int(number)])[0]).split()\n",
        "  center_x_gnina = float(center_x_end[0])\n",
        "  center_y_gnina = float(center_y_end[0])\n",
        "  center_z_gnina = float(center_z_end[0])\n",
        "  print(\"Pocket \" + str(number))\n",
        "  print(\"Score = \" + score_end[0])\n",
        "  print(\"Center x = \"+ str(center_x_end[0]), \"Center y = \"+ str(center_y_end[0]), \"Center z = \"+ str(center_z_end[0]) + \"\\n\")\n",
        "  final_residues = []\n",
        "  for i in range(0,len(file)):\n",
        "    test = file[i]\n",
        "    final_residues.append(int(test[2:]))\n",
        "  residues_num = final_residues\n",
        "else:\n",
        "  residues_num = Convert(number)\n",
        "\n",
        "filenames=[]\n",
        "for k in range(0, len(residues_num)):\n",
        "  extract_ligands(path, residues_num[k])\n",
        "  filenames.append(f\"res_{residues_num[k]}_certo.pdb\")\n",
        "\n",
        "\n",
        "with open('selection_merge.pdb', 'w') as outfile:\n",
        "    for fname in filenames:\n",
        "        with open(fname) as infile:\n",
        "            for line in infile:\n",
        "                outfile.write(line)\n",
        "\n",
        "# reading each line from original\n",
        "# text file\n",
        "file1 = open('/content/selection_merge.pdb', 'r')\n",
        "file2 = open('/content/selection_merge_end.pdb','w')\n",
        "\n",
        "for line in file1.readlines():\n",
        "\n",
        "    # reading all lines that begin\n",
        "    # with \"TextGenerator\"\n",
        "    x = re.findall(\"^END\", line)\n",
        "\n",
        "    if not x:\n",
        "        file2.write(line)\n",
        "\n",
        "# close and save the files\n",
        "file1.close()\n",
        "file2.close()\n",
        "\n",
        "dataset = pd.read_csv(os.path.join(workDir, \"name_residue.txt\"), delimiter = \" \", header=None)\n",
        "df = pd.DataFrame(dataset)\n",
        "df = df.iloc[:, [2]]\n",
        "new = df.to_numpy()\n",
        "\n",
        "print(\"Selected Residue\" + \" - \"  + \"Number\" )\n",
        "for j, i in zip(new, range(0, len(residues_num))):\n",
        "# for j in new:\n",
        "  print(', '.join(j) + \" - \"  + str(residues_num[i]))\n",
        "res_box = '/content/selection_merge_end.pdb'"
      ],
      "metadata": {
        "cellView": "form",
        "id": "8knPKMYiNcYE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Setting the grid box**:\n",
        "#@markdown Now, we will use our ViewProtGrid to visualize the protein, binding site residues and a grid box of variable size and position that we can manipulate using a slider.\n",
        "\n",
        "#These definitions will enable loading our protein and then\n",
        "#drawing a box with a given size and centroid on the cartesian space\n",
        "#This box will enable us to set up the system coordinates for the simulation\n",
        "#\n",
        "#HINT: The active site of the HIV-2 protease is near the beta strands in green\n",
        "#\n",
        "#ACKNOWLEDGE: This script is largely based on the one created by Jose Manuel\n",
        "#Napoles Duarte, Physics Teacher at the Chemical Sciences Faculty of the\n",
        "#Autonomous University of Chihuahua (https://github.com/napoles-uach)\n",
        "#\n",
        "#First, we define the grid box\n",
        "def definegrid(object,bxi,byi,bzi,bxf,byf,bzf):\n",
        "  object.addBox({'center':{'x':bxi,'y':byi,'z':bzi},'dimensions': {'w':bxf,'h':byf,'d':bzf},'color':'blue','opacity': 0.6})\n",
        "\n",
        "#Next, we define how the protein will be shown in py3Dmol\n",
        "#Note that we are also adding a style representation for active site residues\n",
        "def viewprot(object,prot_PDBfile):\n",
        "  mol1 = open(prot_PDBfile, 'r').read()\n",
        "  object.addModel(mol1,'pdb')\n",
        "  object.setStyle({'cartoon': {'color':'white'}})\n",
        "\n",
        "# def viewresid(res_box)\n",
        "#   mol2 = (open(res_box,'r').read(),format='mol2')\n",
        "\n",
        "#Lastly, we combine the box grid and protein into a single viewer\n",
        "def viewprotgrid(prot_PDBfile,resids,bxi,byi,bzi,bxf=10,byf=10,bzf=10):\n",
        "  mol_view = py3Dmol.view(1000,600)\n",
        "  definegrid(mol_view,bxi,byi,bzi,bxf,byf,bzf)\n",
        "  viewprot(mol_view,prot_PDBfile)\n",
        "  mol_view.addModel(open(resids,'r').read(),format='mol2')\n",
        "  ref_m = mol_view.getModel()\n",
        "  ref_m.setStyle({},{'stick':{'colorscheme':'greenCarbon','radius':0.2}})\n",
        "  mol_view.setBackgroundColor('0xffffff')\n",
        "  mol_view.zoomTo()\n",
        "  mol_view.show()\n",
        "\n",
        "#@markdown X coordinate of the center (Angstrom):\n",
        "centerX = 23 #@param {type:\"slider\", min:-100, max:100, step:1}\n",
        "#@markdown Y coordinate of the center (Angstrom):\n",
        "centerY = -25 #@param {type:\"slider\", min:-100, max:100, step:1}\n",
        "#@markdown Z coordinate of the center (Angstrom):\n",
        "centerZ = -3 #@param {type:\"slider\", min:-100, max:100, step:1}\n",
        "#@markdown Size in the X dimension (Angstrom):\n",
        "sizeX = 10 #@param {type:\"slider\", min:0, max:30, step:1}\n",
        "#@markdown Size in the Y dimension (Angstrom):\n",
        "sizeY = 10 #@param {type:\"slider\", min:0, max:30, step:1}\n",
        "#@markdown Size in the Z dimension (Angstrom):\n",
        "sizeZ = 10 #@param {type:\"slider\", min:0, max:30, step:1}\n",
        "\n",
        "viewprotgrid(receptor,res_box,centerX,centerY,centerZ,sizeX,sizeY,sizeZ)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "ak4T6zB-R-w4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Please, provide the necessary input files for the ligand**:\n",
        "\n",
        "#@markdown Type the smiles or filename (PDB and MOL formats) of your molecule. **Ex: C=CC(=O)OC, molecule.pdb or molecule.mol**\n",
        "\n",
        "#@markdown Just remind you that if you want to use pdb or mol file, you should first upload the file here in Colab or in your Google Drive.\n",
        "\n",
        "#@markdown If you don't know the exactly smiles, please, check at https://pubchem.ncbi.nlm.nih.gov/\n",
        "\n",
        "from rdkit import Chem\n",
        "from rdkit.Chem import AllChem\n",
        "from rdkit.Chem import Draw\n",
        "from rdkit.Chem import rdMolTransforms\n",
        "from rdkit.Chem.Draw import rdMolDraw2D\n",
        "from rdkit.Chem import rdDepictor\n",
        "from IPython.display import SVG\n",
        "import ipywidgets as widgets\n",
        "import rdkit\n",
        "from rdkit.Chem.Draw import IPythonConsole\n",
        "AllChem.SetPreferCoordGen(True)\n",
        "from IPython.display import Image\n",
        "from openbabel import pybel\n",
        "import matplotlib.image as mpimg\n",
        "\n",
        "import os\n",
        "\n",
        "import py3Dmol\n",
        "\n",
        "\n",
        "Type = \"smiles\" #@param [\"smiles\", \"pdb\", \"mol\"]\n",
        "\n",
        "smiles_or_filename = 'CCCC1=CC=CC=C1O' #@param {type:\"string\"}\n",
        "\n",
        "if Type == \"smiles\":\n",
        "  Smiles = smiles_or_filename\n",
        "  smiles_fig = Chem.MolFromSmiles(Smiles)\n",
        "  hmol = Chem.AddHs(smiles_fig)\n",
        "  AllChem.EmbedMolecule(hmol)\n",
        "  hmol.GetConformer(0)\n",
        "  mp = AllChem.MMFFGetMoleculeProperties(hmol)\n",
        "  ff = AllChem.MMFFGetMoleculeForceField(hmol, mp)\n",
        "  # Optimize\n",
        "  AllChem.OptimizeMoleculeConfs(hmol, ff, numThreads=1, maxIters=1000)\n",
        "  AllChem.MolToMolFile(hmol, (os.path.join(workDir, \"ligand.mol\")))\n",
        "  AllChem.MolToPDBFile(hmol, (os.path.join(workDir, \"ligand.pdb\")))\n",
        "elif Type == \"pdb\":\n",
        "  pdb_name = os.path.join(workDir, smiles_or_filename)\n",
        "  mol= [m for m in pybel.readfile(filename=pdb_name, format='pdb')][0]\n",
        "  out=pybel.Outputfile(filename='mol.mol',format='mol',overwrite=True)\n",
        "  out.write(mol)\n",
        "  out.close()\n",
        "  mol = Chem.MolFromMolFile('mol.mol')\n",
        "  Smiles = Chem.MolToSmiles(mol)\n",
        "  smiles_fig = Chem.MolFromSmiles(Smiles)\n",
        "  hmol = Chem.AddHs(smiles_fig)\n",
        "  AllChem.EmbedMolecule(hmol)\n",
        "  hmol.GetConformer(0)\n",
        "  mp = AllChem.MMFFGetMoleculeProperties(hmol)\n",
        "  ff = AllChem.MMFFGetMoleculeForceField(hmol, mp)\n",
        "  # Optimize\n",
        "  AllChem.OptimizeMoleculeConfs(hmol, ff, numThreads=1, maxIters=1000)\n",
        "  AllChem.MolToMolFile(hmol, (os.path.join(workDir, \"ligand.mol\")))\n",
        "  AllChem.MolToPDBFile(hmol, (os.path.join(workDir, \"ligand.pdb\")))\n",
        "else:\n",
        "  mol_name = os.path.join(workDir, smiles_or_filename)\n",
        "  mol = Chem.MolFromMolFile(mol_name)\n",
        "  Smiles = Chem.MolToSmiles(mol)\n",
        "  smiles_fig = Chem.MolFromSmiles(Smiles)\n",
        "  hmol = Chem.AddHs(smiles_fig)\n",
        "  AllChem.EmbedMolecule(hmol)\n",
        "  hmol.GetConformer(0)\n",
        "  mp = AllChem.MMFFGetMoleculeProperties(hmol)\n",
        "  ff = AllChem.MMFFGetMoleculeForceField(hmol, mp)\n",
        "  # Optimize\n",
        "  AllChem.OptimizeMoleculeConfs(hmol, ff, numThreads=1, maxIters=1000)\n",
        "  AllChem.MolToMolFile(hmol, (os.path.join(workDir, \"ligand.mol\")))\n",
        "  AllChem.MolToPDBFile(hmol, (os.path.join(workDir, \"ligand.pdb\")))\n",
        "\n",
        "print(\"Smiles: \" + str(Smiles))\n",
        "smi = Draw.MolToFile(smiles_fig, size=(600, 600), filename=os.path.join(workDir,str(Smiles) + '.png'))\n",
        "img = mpimg.imread(os.path.join(workDir,str(Smiles) + '.png'))\n",
        "plt.figure(figsize = (8,8))\n",
        "imgplot = plt.imshow(img)\n",
        "plt.axis('off')\n",
        "plt.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "KLwXCHI30fqz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from typing import List\n",
        "from ase import Atoms\n",
        "from ase.lattice.cubic import Diamond\n",
        "from ase.md.langevin import Langevin\n",
        "from ase.optimize import BFGS\n",
        "from ase import io\n",
        "from ase.io import read, write\n",
        "from ase import units\n",
        "from ase.constraints import ExternalForce, FixInternals\n",
        "import torch\n",
        "import torchani\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "from openbabel import pybel\n",
        "from torchani.units import HARTREE_TO_KCALMOL\n",
        "\n",
        "\n",
        "#@title **Ligand geometry optimization using TorchANI**:\n",
        "\n",
        "#@markdown Geometry optimization for the ligand 3D structure, using ANI-1x, ANI-1ccx or ANI-2x as the optimizing engine.\n",
        "\n",
        "#@markdown If you want to know more about **TorchANI**, please, check at https://aiqm.github.io/torchani/\n",
        "\n",
        "model_name = \"ANI-2x\" #@param [\"ANI-1x\", \"ANI-1ccx\", \"ANI-2x\"]\n",
        "\n",
        "#@markdown Convergence threshold for geometry optimization:\n",
        "\n",
        "opt_tol = 0.0001 #@param {type:\"slider\", min:0.0001, max:0.01, step:0.0001}\n",
        "\n",
        "\n",
        "\n",
        "# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "device = torch.device('cpu')\n",
        "\n",
        "if model_name == \"ANI-2x\":\n",
        "  model = torchani.models.ANI2x(periodic_table_index=True).to(device)\n",
        "  calculator = torchani.models.ANI2x().ase()\n",
        "  print(\"Model = ANI2x\")\n",
        "elif model_name == \"ANI-1ccx\":\n",
        "  model = torchani.models.ANI1ccx(periodic_table_index=True).to(device)\n",
        "  calculator = torchani.models.ANI1ccx().ase()\n",
        "  print(\"Model = ANI1ccx\")\n",
        "elif model_name == \"ANI-1x\":\n",
        "  model = torchani.models.ANI1x(periodic_table_index=True).to(device)\n",
        "  calculator = torchani.models.ANI1x().ase()\n",
        "  print(\"Model = ANI1x\")\n",
        "else:\n",
        "  pass\n",
        "\n",
        "def mol2arr(mols, device=device):\n",
        "    coordinates = []\n",
        "    spices = []\n",
        "    for mol in mols:\n",
        "        pos = mol.GetConformer().GetPositions().tolist()\n",
        "        atomnums = [a.GetAtomicNum() for a in mol.GetAtoms()]\n",
        "        coordinates.append(pos)\n",
        "        spices.append(atomnums)\n",
        "    coordinates = torch.tensor(coordinates,\n",
        "                               requires_grad=True,\n",
        "                               device=device)\n",
        "    species = torch.tensor(spices, device=device)\n",
        "    return coordinates, species\n",
        "\n",
        "mol_deg = AllChem.MolFromMolFile ((os.path.join(workDir, \"ligand.mol\")), removeHs=False)\n",
        "mol = io.read(os.path.join(workDir, \"ligand.mol\"))\n",
        "coordinates, species = mol2arr([mol_deg], device)\n",
        "tensor1 = coordinates.detach().numpy()\n",
        "atoms = Atoms(mol, positions=tensor1[0])\n",
        "atoms.center(vacuum=3.0)\n",
        "atoms.set_calculator(calculator)\n",
        "print(\"Begin Geometry Optimization \")\n",
        "opt = BFGS(atoms)\n",
        "opt.run(fmax=opt_tol)\n",
        "# print()\n",
        "write((os.path.join(workDir, \"ligand_min.xyz\")), format=\"xyz\", images=atoms)\n",
        "\n",
        "atomic_symbols = []\n",
        "xyz_coordinates = []\n",
        "\n",
        "with open((os.path.join(workDir, \"ligand_min.xyz\")), \"r\") as file:\n",
        "  for line_number,line in enumerate(file):\n",
        "      if line_number == 0:\n",
        "          num_atoms = int(line)\n",
        "      elif line_number == 1:\n",
        "          comment = line # might have useful information\n",
        "      else:\n",
        "          atomic_symbol, x, y, z = line.split()\n",
        "          atomic_symbols.append(atomic_symbol)\n",
        "          xyz_coordinates.append([float(x),float(y),float(z)])\n",
        "\n",
        "from rdkit.Geometry import Point3D\n",
        "conf = mol_deg.GetConformer()\n",
        "\n",
        "for i in range(mol_deg.GetNumAtoms()):\n",
        "  x,y,z = xyz_coordinates[i]\n",
        "  conf.SetAtomPosition(i,Point3D(x,y,z))\n",
        "AllChem.MolToMolFile(mol_deg, (os.path.join(workDir, \"ligand_min.mol\")))\n",
        "AllChem.MolToPDBFile(mol_deg, (os.path.join(workDir, \"ligand_min.pdb\")))\n",
        "#convert to sdf format\n",
        "# mol= [m for m in pybel.readfile(filename=os.path.join(workDir, \"ligand_min.mol\"), format='mol')][0]\n",
        "# out=pybel.Outputfile(filename=ligand,format='pdbqt',overwrite=True)\n",
        "# out.Molecule()\n",
        "# out.write(mol)\n",
        "# out.close()\n",
        "ligand_mol=os.path.join(workDir, \"ligand_min.mol\")\n",
        "os.system(\"obabel -i mol \" + str(ligand_mol) + \" -o pdbqt -O \" + str(ligand) + \" -xh --partialcharge\")\n",
        "\n",
        "#TorchANI Energies\n",
        "mol_deg = AllChem.MolFromMolFile ((os.path.join(workDir, \"ligand_min.mol\")), removeHs=False)\n",
        "coordinates, species = mol2arr([mol_deg], device)\n",
        "energy = model((species, coordinates)).energies\n",
        "derivative = torch.autograd.grad(energy.sum(), coordinates)[0]\n",
        "force = -derivative\n",
        "print('Energy:', energy.item())\n",
        "# print('Force:', force.squeeze())"
      ],
      "metadata": {
        "cellView": "form",
        "id": "hTcrO_yCzDEb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Ligand Visualization**:\n",
        "#@markdown Now the ligand has been optimized, it would be recomended to visualize and check the ligand.\n",
        "\n",
        "view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',)\n",
        "view.removeAllModels()\n",
        "view.setViewStyle({'style':'outline','color':'black','width':0.05})\n",
        "\n",
        "# view.addModel(open(receptor,'r').read(),format='pdb')\n",
        "# Prot=view.getModel()\n",
        "# Prot.setStyle({'cartoon':{'arrows':True, 'tubes':True, 'style':'oval', 'color':'white'}})\n",
        "# view.addSurface(py3Dmol.VDW,{'opacity':0.6,'color':'white'})\n",
        "\n",
        "view.addModel(open(os.path.join(workDir, \"ligand_min.mol\"),'r').read(),format='pdbqt')\n",
        "ref_m = view.getModel()\n",
        "ref_m.setStyle({},{'stick':{'colorscheme':'greenCarbon','radius':0.2}})\n",
        "\n",
        "view.zoomTo()\n",
        "view.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "0GZCAGa9_Y3k"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Parameters for the docking calculation:**\n",
        "\n",
        "#@markdown Explicit random seed:\n",
        "seed = \"0\" #@param {type:\"string\"}\n",
        "\n",
        "scoring = \"vina\" #@param [\"vina\", \"vinardo\"]\n",
        "\n",
        "\n",
        "#@markdown **Advanced options** `--search_mode` is the recommended setting of `--exhaustiveness` and `--max_step`, with three combinations called `fast`, `balance`, and `detail`.\n",
        "\n",
        "#@markdown - `fast` mode: `--exhaustiveness 128` & `--max_step 20`\n",
        "#@markdown - `balance` mode: `--exhaustiveness 384` & `--max_step 40`\n",
        "#@markdown - `detail` mode: `--exhaustiveness 512` & `--max_step 40`\n",
        "\n",
        "#@markdown The larger `--exhaustiveness` and `--max_step`, the higher the computational complexity, the higher the accuracy, but the larger the computational cost.\n",
        "\n",
        "search_mode = \"balance\" #@param [\"fast\", \"balance\", \"detail\"]\n",
        "exhaustiveness = 384 #@param {type:\"slider\", min:8, max:512, step:8}\n",
        "max_step = 40 #@param {type:\"slider\", min:0, max:40, step:5}\n",
        "\n",
        "#@markdown Maximum number of binding modes to generate:\n",
        "num_modes = 10 #@param {type:\"slider\", min:1, max:10, step:1}\n",
        "\n",
        "import locale\n",
        "def getpreferredencoding(do_setlocale = True):\n",
        "    return \"UTF-8\"\n",
        "locale.getpreferredencoding = getpreferredencoding\n",
        "\n",
        "unidock = \"./unidock --receptor \" + str(receptor_pdbqt) + \" --gpu_batch \" + str(ligand) + \" --search_mode \" + str(search_mode) + \" --scoring \" + str(scoring) + \" --center_x \" + str(centerX) + \" --center_y \" + str(centerY) + \" --center_z \" + str(centerZ) + \" --size_x \" + str(sizeX) + \" --size_y \" + str(sizeY) + \" --size_z \" + str(sizeZ) + \" --num_modes \" + str(num_modes) + \" --dir \" + str(workDir) + \" --seed \" + str(seed) + \" --max_step \" + str(max_step) + \" --exhaustiveness \" + str(exhaustiveness)\n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "with open('unidock.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(unidock)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "print(unidock)\n",
        "!chmod 700 unidock.sh 2>&1 1>/dev/null\n",
        "!bash unidock.sh\n",
        "\n",
        "ligand_out=os.path.join(workDir, \"ligand_out.pdbqt\")\n",
        "ligand_out_sdf=os.path.join(workDir, \"ligand_out.sdf\")\n",
        "os.system(\"obabel -i pdbqt \" + str(ligand_out) + \" -o sdf -O \" + str(ligand_out_sdf) + \" -xh\")\n",
        "\n",
        "\n",
        "import gzip\n",
        "v = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',)\n",
        "v.setViewStyle({'style':'outline','color':'black','width':0.05})\n",
        "v.addModel(open(receptor).read())\n",
        "v.setStyle({'cartoon':{},'stick':{'colorscheme':'white','radius':.1}})\n",
        "v.addModel(open(res_box).read())\n",
        "v.setStyle({'model':1},{'stick':{'colorscheme':'dimgrayCarbon','radius':.175}})\n",
        "v.addModelsAsFrames(open(ligand_out_sdf,'rt').read())\n",
        "v.setStyle({'model':2},{'stick':{'colorscheme':'greenCarbon'}})\n",
        "v.animate({'interval':1000})\n",
        "v.zoomTo({'model':1})\n",
        "v.rotate(90)"
      ],
      "metadata": {
        "id": "wyzlRC-sBY3J",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Best pose selection:**\n",
        "\n",
        "mode_number = \"1\" #@param [\"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\", \"9\", \"10\"]\n",
        "\n",
        "view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',)\n",
        "view.removeAllModels()\n",
        "view.setViewStyle({'style':'outline','color':'black','width':0.05})\n",
        "\n",
        "view.addModel(open(receptor,'r').read(),'pdb')\n",
        "Prot=view.getModel()\n",
        "# Prot.setStyle({'cartoon':{'arrows':True, 'tubes':True, 'style':'oval', 'color':'white'},'stick':{'radius':.1}})\n",
        "Prot.setStyle({'cartoon':{'arrows':True, 'tubes':True, 'style':'oval', 'color':'white'}})\n",
        "view.addSurface(py3Dmol.VDW,{'opacity':0.2,'color':'white'})\n",
        "\n",
        "\n",
        "view.addModel(open(res_box,'r').read(),'pdb')\n",
        "ref_m = view.getModel()\n",
        "ref_m.setStyle({},{'stick':{'colorscheme':'grayCarbon','radius':0.2}})\n",
        "\n",
        "\n",
        "results=Chem.SDMolSupplier(ligand_out_sdf,False)\n",
        "p=Chem.MolToMolBlock(results[(int(mode_number)-1)],False)\n",
        "p2=Chem.MolToMolFile(results[(int(mode_number)-1)],(os.path.join(workDir, str(int(mode_number)) + \"_pose.sdf\")))\n",
        "\n",
        "print('Reference residues: Gray | Uni-Dock Pose: Green')\n",
        "view.addModel(p,'mol')\n",
        "x = view.getModel()\n",
        "x.setStyle({},{'stick':{'colorscheme':'greenCarbon','radius':0.2}})\n",
        "\n",
        "view.zoomTo()\n",
        "view.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "3GggfJfLIVe8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import parmed as pmd\n",
        "from openmmforcefields.generators import SMIRNOFFTemplateGenerator\n",
        "from openmmforcefields.generators import GAFFTemplateGenerator\n",
        "\n",
        "# Imports from the toolkit\n",
        "import openff.toolkit\n",
        "from openff.toolkit.topology import Molecule, Topology\n",
        "from openff.toolkit.typing.engines.smirnoff import ForceField\n",
        "# Imports from dependencie\n",
        "from openff.units import unit\n",
        "from openff.units.openmm import to_openmm, from_openmm\n",
        "try:\n",
        "    import openmm\n",
        "    from openmm import unit\n",
        "except ImportError:\n",
        "    from simtk import openmm, unit\n",
        "\n",
        "#@title **Parameters to generate the topology:**\n",
        "\n",
        "#@markdown **Parameters to generate the protein topology:**\n",
        "\n",
        "Force_field = \"AMBER99SB\" #@param [\"AMBER14SB\", \"AMBER99\", \"AMBER99SB\"]\n",
        "Water_type = \"TIP3P\" #@param [\"TIP3P\", \"SPC/E\"]\n",
        "\n",
        "#@markdown Padding distance (angstrons) determines the largest size of the solute along any axis (x, y, or z). It then creates a cubic box of width (solute size)+2*(padding). The above line guarantees that no part of the solute comes closer than 1 nm to any edge of the box.\n",
        "\n",
        "Padding_distance = 4 #@param {type:\"slider\", min:1, max:10, step:1}\n",
        "\n",
        "#@markdown **ATTENTION**: Give the concentration in Molar units:\n",
        "\n",
        "Ions = \"NaCl\" #@param [\"NaCl\", \"KCl\" ]\n",
        "Concentration = \"0.15\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown **Function to add missing hydrogen atoms based on pH:**\n",
        "\n",
        "pH = \"7.4\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown **Parameters to generate the ligand topology:**\n",
        "\n",
        "Ligand_Force_field = \"gaff-2.11\" #@param ['gaff-1.4', 'gaff-1.8', 'gaff-1.81', 'gaff-2.1', 'gaff-2.11', 'smirnoff99Frosst-1.1.0', 'smirnoff99Frosst-1.0.0', 'smirnoff99Frosst-1.0.9', 'smirnoff99Frosst-1.0.1', 'smirnoff99Frosst-1.0.6', 'smirnoff99Frosst-1.0.5', 'smirnoff99Frosst-1.0.7', 'smirnoff99Frosst-1.0.8', 'smirnoff99Frosst-1.0.2', 'smirnoff99Frosst-1.0.4', 'smirnoff99Frosst-1.0.3', 'openff-1.0.0-RC1', 'openff-1.0.1', 'openff-1.3.1-alpha.1', 'openff-1.0.0', 'openff-1.2.0', 'openff-2.0.0-rc.2', 'openff-2.0.0-rc.1', 'openff-2.0.0', 'openff-1.1.0', 'openff-1.3.1', 'openff-1.1.1', 'openff-1.0.0-RC2', 'openff-1.3.0', 'openff-1.2.1']\n",
        "\n",
        "lig_end = (os.path.join(workDir, str(int(mode_number)) + \"_ligand.sdf\"))\n",
        "\n",
        "results=Chem.SDMolSupplier(ligand_out_sdf,True)\n",
        "p=Chem.MolToMolFile(results[(int(mode_number)-1)],(os.path.join(workDir, str(int(mode_number)) + \"_ligand.sdf\")))\n",
        "\n",
        "# Load a molecule from a SDF file\n",
        "ligand = Molecule.from_file(os.path.join(workDir, str(int(mode_number)) + \"_ligand.sdf\"))\n",
        "# Molecule stores both the co-ordinates of the atoms and their bond graph\n",
        "ligand_positions = ligand.conformers[0]\n",
        "ligand_topology = ligand.to_topology()\n",
        "\n",
        "# if Force_field == \"CHARMM36\" and Water_type == \"TIP3P\":\n",
        "#   water = \"charmm36/water.xml\"\n",
        "#   water_model = \"tip3p\"\n",
        "#   ff = \"charmm36.xml\"\n",
        "if Force_field == \"AMBER14SB\" and Water_type == \"TIP3P\":\n",
        "  water = \"amber14/tip3p.xml\"\n",
        "  water_model = \"tip3p\"\n",
        "  ff = \"amber14/protein.ff14SB.xml\"\n",
        "elif Force_field == \"AMBER99SB\" and Water_type == \"TIP3P\":\n",
        "  water = \"tip3p.xml\"\n",
        "  water_model = \"tip3p\"\n",
        "  ff = \"amber99sbildn.xml\"\n",
        "elif Force_field == \"AMBER99\" and Water_type == \"TIP3P\":\n",
        "  water = \"tip3p.xml\"\n",
        "  water_model = \"tip3p\"\n",
        "  ff = \"amber99sb.xml\"\n",
        "# elif Force_field == \"CHARMM36\" and Water_type == \"SPC/E\":\n",
        "#   water = \"charmm36/spce.xml\"\n",
        "#   water_model = \"spce\"\n",
        "#   ff = \"charmm36.xml\"\n",
        "elif Force_field == \"AMBER14SB\" and Water_type == \"SPC/E\":\n",
        "  water = \"amber14/spce.xml\"\n",
        "  water_model = \"spce\"\n",
        "  ff = \"amber14/protein.ff14SB.xml\"\n",
        "elif Force_field == \"AMBER99SB\" and Water_type == \"SPC/E\":\n",
        "  water = \"spce.xml\"\n",
        "  water_model = \"spce\"\n",
        "  ff = \"amber99sbildn.xml\"\n",
        "elif Force_field == \"AMBER99\" and Water_type == \"SPC/E\":\n",
        "  water = \"spce.xml\"\n",
        "  water_model = \"spce\"\n",
        "  ff = \"amber99sb.xml\"\n",
        "else:\n",
        "  pass\n",
        "\n",
        "# Load protein and water force field parameters\n",
        "omm_forcefield = openmm.app.ForceField(ff, water)\n",
        "\n",
        "small_ff_options = ['gaff-1.4', 'gaff-1.8', 'gaff-1.81', 'gaff-2.1', 'gaff-2.11', 'smirnoff99Frosst-1.1.0', 'smirnoff99Frosst-1.0.0', 'smirnoff99Frosst-1.0.9', 'smirnoff99Frosst-1.0.1',\n",
        "            'smirnoff99Frosst-1.0.6', 'smirnoff99Frosst-1.0.5', 'smirnoff99Frosst-1.0.7', 'smirnoff99Frosst-1.0.8', 'smirnoff99Frosst-1.0.2', 'smirnoff99Frosst-1.0.4',\n",
        "            'smirnoff99Frosst-1.0.3', 'openff-1.0.0-RC1', 'openff-1.0.1', 'openff-1.3.1-alpha.1', 'openff-1.0.0', 'openff-1.2.0', 'openff-2.0.0-rc.2', 'openff-2.0.0-rc.1',\n",
        "            'openff-2.0.0', 'openff-1.1.0', 'openff-1.3.1', 'openff-1.1.1', 'openff-1.0.0-RC2', 'openff-1.3.0', 'openff-1.2.1']\n",
        "\n",
        "type_generator_options = ('GAFF', 'GAFF', 'GAFF', 'GAFF', 'GAFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF',\n",
        "                          'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF', 'SMIRNOFF')\n",
        "\n",
        "small_ff_include = [\"'gaff-1.4'\", \"'gaff-1.8'\", \"'gaff-1.81'\", \"'gaff-2.1'\", \"'gaff-2.11'\", \"'smirnoff99Frosst-1.1.0'\", \"'smirnoff99Frosst-1.0.0'\", \"'smirnoff99Frosst-1.0.9'\", \"'smirnoff99Frosst-1.0.1'\",\n",
        "            \"'smirnoff99Frosst-1.0.6'\", \"'smirnoff99Frosst-1.0.5'\", \"'smirnoff99Frosst-1.0.7'\", \"'smirnoff99Frosst-1.0.8'\", \"'smirnoff99Frosst-1.0.2'\", \"'smirnoff99Frosst-1.0.4'\",\n",
        "            \"'smirnoff99Frosst-1.0.3'\", \"'openff-1.0.0-RC1'\", \"'openff-1.0.1'\", \"'openff-1.3.1-alpha.1'\", \"'openff-1.0.0'\", \"'openff-1.2.0'\", \"'openff-2.0.0-rc.2'\", \"'openff-2.0.0-rc.1'\",\n",
        "            \"'openff-2.0.0'\", \"'openff-1.1.0'\", \"'openff-1.3.1'\", \"'openff-1.1.1'\", \"'openff-1.0.0-RC2'\", \"'openff-1.3.0'\", \"'openff-1.2.1'\"]\n",
        "\n",
        "\n",
        "small_ff_id = small_ff_options.index(Ligand_Force_field)\n",
        "type_generator = type_generator_options[small_ff_id]\n",
        "small_ff = small_ff_include[small_ff_id]\n",
        "\n",
        "# Teach OpenMM about the ligand molecule and force field\n",
        "ligand_generator = type_generator + \"TemplateGenerator(forcefield=\" + small_ff + \", molecules=[ligand])\"\n",
        "type_generator_options = eval(ligand_generator)\n",
        "omm_forcefield.registerTemplateGenerator(type_generator_options.generator)\n",
        "\n",
        "# Load the receptor structure into Modeller\n",
        "pdb = openmm.app.PDBFile(receptor)\n",
        "modeller = openmm.app.Modeller(pdb.topology, pdb.positions)\n",
        "\n",
        "# Add the ligand\n",
        "# Under the hood, this line uses the OpenFF Toolkit to generate new parameters for the ligand\n",
        "modeller.add(ligand_topology.to_openmm(), to_openmm(ligand_positions))\n",
        "\n",
        "modeller.addHydrogens(omm_forcefield, pH=float(pH))\n",
        "\n",
        "if Ions == \"NaCl\":\n",
        "  positive_ion = 'Na+'\n",
        "else:\n",
        "  positive_ion = 'K+'\n",
        "\n",
        "# solvate it in 0.15 M NaCl solution\n",
        "modeller.addSolvent(\n",
        "    omm_forcefield,\n",
        "    model=water_model,\n",
        "    padding=int(Padding_distance) * unit.angstrom,\n",
        "    ionicStrength=float(Concentration) * unit.molar,\n",
        "    positiveIon=positive_ion, negativeIon='Cl-',\n",
        ")\n",
        "\n",
        "# Retrieve the OpenMM Topology, which stores the atoms and connectivity\n",
        "topology = modeller.getTopology()\n",
        "\n",
        "# Get the initial positions\n",
        "# The box is about 75 angstroms per side, so add (30, 30, 30) to center the protein\n",
        "positions = modeller.getPositions() + np.array([30, 30, 30]) * unit.angstrom\n",
        "\n",
        "#Export PDB file\n",
        "system_pdb = os.path.join(workDir, 'system.pdb')\n",
        "PDBFile.writeFile(topology, positions, open(system_pdb, 'w'))\n",
        "\n",
        "# ParmEd's GROMACS exporter can't handle constraints from openmm, so we need a variant for export without them\n",
        "export_system = omm_forcefield.createSystem(\n",
        "    modeller.topology,\n",
        "    # constraints=None,\n",
        "    # rigidWater=False,\n",
        "    flexibleConstraints=True\n",
        ")\n",
        "\n",
        "# Combine the topology, system and positions into a ParmEd Structure\n",
        "pmd_complex_struct = pmd.openmm.load_topology(topology, export_system, positions)\n",
        "\n",
        "# Export AMBER files.\n",
        "amber_prmtop = os.path.join(workDir, 'system_amber.prmtop')\n",
        "amber_inpcrd = os.path.join(workDir, 'system_amber.inpcrd')\n",
        "pmd_complex_struct.save(amber_prmtop, overwrite=True)\n",
        "pmd_complex_struct.save(amber_inpcrd, overwrite=True)\n",
        "pprm = pmd.load_file(amber_prmtop, amber_inpcrd)\n",
        "mprm_from_parmed = mda.Universe(pprm)\n",
        "\n",
        "# Export GROMACS files.\n",
        "gromacs_top = os.path.join(workDir, 'system_gromacs.top')\n",
        "gromacs_gro = os.path.join(workDir, 'system_gromacs.gro')\n",
        "pmd_complex_struct.save(gromacs_top, overwrite=True)\n",
        "pmd_complex_struct.save(gromacs_gro, overwrite=True)\n",
        "\n",
        "pdb_check = os.path.exists(system_pdb)\n",
        "top_amber = os.path.exists(amber_prmtop)\n",
        "inpcrd_amber = os.path.exists(amber_inpcrd)\n",
        "top_gromacs = os.path.exists(gromacs_top)\n",
        "gro_gromacs = os.path.exists(gromacs_gro)\n",
        "\n",
        "if pdb_check == True and top_amber == True and inpcrd_amber == True and top_gromacs == True and gro_gromacs == True:\n",
        "  print(\"Successfully generated topology! :-)\")\n",
        "else:\n",
        "  print(\"ERROR: Check your inputs! \")"
      ],
      "metadata": {
        "id": "Kf88JFKUJYc_",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C8kKR7bpI86W"
      },
      "source": [
        "## Let's take a look on our simulation box:"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Show 3D structure**\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "import py3Dmol\n",
        "\n",
        "color = \"gray\" #@param [\"gray\", \"rainbow\"]\n",
        "show_sidechains = False #@param {type:\"boolean\"}\n",
        "show_mainchains = False #@param {type:\"boolean\"}\n",
        "show_ligand = True #@param {type:\"boolean\"}\n",
        "show_box = True #@param {type:\"boolean\"}\n",
        "box_opacity = 0.6 #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
        "\n",
        "\n",
        "def show_pdb(show_sidechains=False, show_mainchains=False, show_ligand = False, show_box = False, color=\"rainbow\"):\n",
        "  view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',)\n",
        "  view.addModel(open(system_pdb,'r').read(),'pdb')\n",
        "\n",
        "  if color == \"gray\":\n",
        "    view.setStyle({'cartoon':{}})\n",
        "  elif color == \"rainbow\":\n",
        "    view.setStyle({'cartoon': {'color':'spectrum'}})\n",
        "\n",
        "  if show_sidechains:\n",
        "    BB = ['C','O','N']\n",
        "    view.addStyle({'and':[{'resn':[\"GLY\",\"PRO\"],'invert':True},{'atom':BB,'invert':True}]},\n",
        "                        {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "    view.addStyle({'and':[{'resn':\"GLY\"},{'atom':'CA'}]},\n",
        "                        {'sphere':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "    view.addStyle({'and':[{'resn':\"PRO\"},{'atom':['C','O'],'invert':True}]},\n",
        "                        {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "  if show_mainchains:\n",
        "    BB = ['C','O','N','CA']\n",
        "    view.addStyle({'atom':BB},{'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "\n",
        "  if show_box:\n",
        "    view.addSurface(py3Dmol.SAS, {'opacity': box_opacity, 'color':'white'})\n",
        "\n",
        "  if show_ligand:\n",
        "    HP = ['UNK']\n",
        "    view.addStyle({'and':[{'resn':HP}]},\n",
        "                       {'stick':{'colorscheme':'greenCarbon','radius':0.3}})\n",
        "    view.setViewStyle({'style':'outline','color':'black','width':0.1})\n",
        "\n",
        "  view.zoomTo()\n",
        "  return view\n",
        "\n",
        "\n",
        "show_pdb(show_sidechains, show_mainchains, show_ligand, show_box, color).show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "Pr1eoCsHxsWl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **View and check the Ligand Interaction Network (LigPlot)**\n",
        "#@markdown This diagram is interactive and allows moving around the residues, as well as clicking the legend to toggle the display of specific residues types or interactions. The diagram will be saved as an HTML file (initial.html).\n",
        "\n",
        "import MDAnalysis as mda\n",
        "import prolif as plf\n",
        "import numpy as np\n",
        "import os\n",
        "from prolif.plotting.network import LigNetwork\n",
        "\n",
        "# load topology\n",
        "u = mda.Universe(amber_prmtop, system_pdb)\n",
        "lig = u.select_atoms(\"not protein and not (resname HOH) and not (resname CL) and not (resname NA) and not (resname K)\")\n",
        "prot = u.select_atoms(\"protein\")\n",
        "\n",
        "# create RDKit-like molecules for visualisation\n",
        "lmol = plf.Molecule.from_mda(lig)\n",
        "pmol = plf.Molecule.from_mda(prot)\n",
        "\n",
        "fp = plf.Fingerprint()\n",
        "fp.run(u.trajectory[::10], lig, prot)\n",
        "df = fp.to_dataframe(return_atoms=True)\n",
        "\n",
        "net = LigNetwork.from_ifp(df, lmol,\n",
        "                          # replace with `kind=\"frame\", frame=0` for the other depiction\n",
        "                          kind=\"aggregate\", frame=0,\n",
        "                          rotation=270)\n",
        "net.save(os.path.join(workDir, \"initial.html\"))\n",
        "net.display()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "J7YMpWZ8pIiK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Open Binding Pose Metadynamics**: evaluating ligand pose stability using metadynamics\n",
        "\n",
        "#@markdown **OpenBPMD** is an open source implementation of binding pose metadynamics (BPMD). Based on work by [Clark et al, 2016](https://doi.org/10.1021/acs.jctc.6b00201). Primarily built for reranking docked poses out of a list of candidates generated by a docking program. OpenBPMD biases ligands using metadynamics and evaluates how stable they are during the simulation.\n",
        "\n",
        "# OpenMM\n",
        "from openmm.app.metadynamics import *\n",
        "\n",
        "# The rest\n",
        "import argparse\n",
        "import numpy as np\n",
        "import MDAnalysis as mda\n",
        "from MDAnalysis.analysis import rms, contacts\n",
        "import mdtraj as md\n",
        "import pandas as pd\n",
        "import parmed as pmd\n",
        "import glob\n",
        "import os\n",
        "from sys import stdout, exit, stderr\n",
        "\n",
        "Google_Drive_Path = '/content/' #@param {type:\"string\"}\n",
        "workDir = Google_Drive_Path\n",
        "\n",
        "#@markdown Name of the output directory:\n",
        "output = \"ligand_pose\" #@param {type:\"string\"}\n",
        "output = os.path.join(workDir, output)\n",
        "#@markdown Number of repeat OpenBPMD simulations to run in series:\n",
        "\n",
        "nreps = \"1\" #@param [\"1\",\"2\",\"3\",\"4\",\"5\",\"6\",\"7\",\"8\",\"9\",\"10\"]\n",
        "\n",
        "#@markdown Size of the metadynamical hill, in kcal/mol:\n",
        "\n",
        "hill_height = 0.3 #@param {type:\"slider\", min:0.1, max:1, step:0.1}\n",
        "#Parameters\n",
        "lig_resname = 'UNK' # Residue name of the ligand in the structure/parameter file.\n",
        "nreps = int(nreps)\n",
        "parameters = os.path.join(workDir, 'system_amber.prmtop')\n",
        "structure = os.path.join(workDir, 'system_amber.inpcrd')\n",
        "continue_simulation = \"no\" #@param [\"yes\", \"no\"]\n",
        "\n",
        "\n",
        "\n",
        "def main():\n",
        "\n",
        "    if not os.path.isdir(f'{output}'):\n",
        "      os.mkdir(f'{output}')\n",
        "\n",
        "    coords = AmberInpcrdFile(structure)\n",
        "    parm = AmberPrmtopFile(parameters)\n",
        "\n",
        "    # Minimize\n",
        "    min_file_name = 'minimized_system.pdb'\n",
        "    if not os.path.isfile(os.path.join(output,min_file_name)):\n",
        "        print(\"Minimizing...\")\n",
        "        #min_pos = minimize(parm, coords.positions, output)\n",
        "        minimize(parm, coords.positions, output, min_file_name)\n",
        "    min_pdb = os.path.join(output,min_file_name)\n",
        "\n",
        "    # Equilibrate\n",
        "    eq_file_name = 'equil_system.pdb'\n",
        "    if not os.path.isfile(os.path.join(output,eq_file_name)):\n",
        "        print(\"Equilibrating...\")\n",
        "        equilibrate(min_pdb, parm, output, eq_file_name)\n",
        "    eq_pdb = os.path.join(output,eq_file_name)\n",
        "    cent_eq_pdb = os.path.join(output,'centred_'+eq_file_name)\n",
        "    if os.path.isfile(eq_pdb) and not os.path.isfile(cent_eq_pdb):\n",
        "        mdtraj_top = parameters\n",
        "        mdu = md.load(eq_pdb, top=mdtraj_top)\n",
        "        mdu.image_molecules()\n",
        "        mdu.save_pdb(cent_eq_pdb)\n",
        "\n",
        "    # Run NREPS number of production simulations\n",
        "    for idx in range(1, (nreps+1)):\n",
        "        rep_dir = os.path.join(output,f'rep_{idx}')\n",
        "        if not os.path.isdir(rep_dir):\n",
        "            os.mkdir(rep_dir)\n",
        "        print(\"Running BPMD simulation \" + str(idx))\n",
        "        if os.path.isfile(os.path.join(rep_dir,'bpm_results.csv')):\n",
        "            continue\n",
        "\n",
        "        produce(output, idx, lig_resname, eq_pdb, parm, parameters,\n",
        "                structure, hill_height)\n",
        "\n",
        "        trj_name = os.path.join(rep_dir,'trj.dcd')\n",
        "\n",
        "        PoseScoreArr = get_pose_score(cent_eq_pdb, trj_name, lig_resname)\n",
        "\n",
        "        ContactScoreArr = get_contact_score(cent_eq_pdb, trj_name, lig_resname)\n",
        "\n",
        "        # Calculate the CompScore at every frame\n",
        "        CompScoreArr = np.zeros(99)\n",
        "        for index in range(ContactScoreArr.shape[0]):\n",
        "            ContactScore, PoseScore = ContactScoreArr[index], PoseScoreArr[index]\n",
        "            CompScore = PoseScore - 5 * ContactScore\n",
        "            CompScoreArr[index] = CompScore\n",
        "\n",
        "        Scores = np.stack((CompScoreArr, PoseScoreArr, ContactScoreArr), axis=-1)\n",
        "\n",
        "        # Save a DataFrame to CSV\n",
        "        df = pd.DataFrame(Scores, columns=['CompScore', 'PoseScore',\n",
        "                                           'ContactScore'])\n",
        "        df.to_csv(os.path.join(rep_dir,'bpm_results.csv'), index=False)\n",
        "\n",
        "    collect_results(output, output)\n",
        "\n",
        "    return None\n",
        "\n",
        "\n",
        "def get_contact_score(structure_file, trajectory_file, lig_resname):\n",
        "    \"\"\"A function the gets the ContactScore from an OpenBPMD trajectory.\n",
        "\n",
        "    Parameters\n",
        "    ----------\n",
        "    structure_file : str\n",
        "        The name of the centred equilibrated system PDB file that\n",
        "        was used to start the OpenBPMD simulation.\n",
        "    trajectory_file : str\n",
        "        The name of the OpenBPMD trajectory file.\n",
        "    lig_resname : str\n",
        "        Residue name of the ligand that was biased.\n",
        "\n",
        "    Returns\n",
        "    -------\n",
        "    contact_scores : np.array\n",
        "        ContactScore for every frame of the trajectory.\n",
        "    \"\"\"\n",
        "    u = mda.Universe(structure_file, trajectory_file)\n",
        "\n",
        "    sel_donor = f\"resname {lig_resname} and not name *H*\"\n",
        "    sel_acceptor = f\"protein and not name H* and \\\n",
        "                     around 5 resname {lig_resname}\"\n",
        "\n",
        "    # reference groups (first frame of the trajectory, but you could also use\n",
        "    # a separate PDB, eg crystal structure)\n",
        "    a_donors = u.select_atoms(sel_donor)\n",
        "    a_acceptors = u.select_atoms(sel_acceptor)\n",
        "\n",
        "    cont_analysis = contacts.Contacts(u, select=(sel_donor, sel_acceptor),\n",
        "                                      refgroup=(a_donors, a_acceptors),\n",
        "                                      radius=3.5)\n",
        "\n",
        "    cont_analysis.run()\n",
        "    # print number of average contacts in the first ns\n",
        "    # NOTE - hard coded number of frames (100 per traj)\n",
        "    frame_idx_first_ns = int(len(cont_analysis.timeseries)/10)\n",
        "    first_ns_mean = np.mean(cont_analysis.timeseries[1:frame_idx_first_ns, 1])\n",
        "    if first_ns_mean == 0:\n",
        "        normed_contacts = cont_analysis.timeseries[1:, 1]\n",
        "    else:\n",
        "        normed_contacts = cont_analysis.timeseries[1:, 1]/first_ns_mean\n",
        "    contact_scores = np.where(normed_contacts > 1, 1, normed_contacts)\n",
        "\n",
        "    return contact_scores\n",
        "\n",
        "\n",
        "def get_pose_score(structure_file, trajectory_file, lig_resname):\n",
        "    \"\"\"A function the gets the PoseScore (ligand RMSD) from an OpenBPMD\n",
        "    trajectory.\n",
        "\n",
        "    Parameters\n",
        "    ----------\n",
        "    'structure_file : str\n",
        "        The name of the centred equilibrated system\n",
        "        PDB file that was used to start the OpenBPMD simulation.\n",
        "    trajectory_file : str\n",
        "        The name of the OpenBPMD trajectory file.\n",
        "    lig_resname : str\n",
        "        Residue name of the ligand that was biased.\n",
        "\n",
        "    Returns\n",
        "    -------\n",
        "    pose_scores : np.array\n",
        "        PoseScore for every frame of the trajectory.\n",
        "    \"\"\"\n",
        "    # Load a MDA universe with the trajectory\n",
        "    u = mda.Universe(structure_file, trajectory_file)\n",
        "    # Align each frame using the backbone as reference\n",
        "    # Calculate the RMSD of ligand heavy atoms\n",
        "    r = rms.RMSD(u, select='backbone',\n",
        "                 groupselections=[f'resname {lig_resname} and not name H*'],\n",
        "                 ref_frame=0).run()\n",
        "    # Get the PoseScores as np.array\n",
        "    pose_scores = r.rmsd[1:, -1]\n",
        "\n",
        "    return pose_scores\n",
        "\n",
        "\n",
        "def minimize(parm, input_positions, out_dir, min_file_name):\n",
        "    \"\"\"An energy minimization function down with an energy tolerance\n",
        "    of 10 kJ/mol.\n",
        "\n",
        "    Parameters\n",
        "    ----------\n",
        "    parm : Parmed or OpenMM parameter file object\n",
        "        Used to create the OpenMM System object.\n",
        "    input_positions : OpenMM Quantity\n",
        "        3D coordinates of the equilibrated system.\n",
        "    out_dir : str\n",
        "        Directory to write the outputs.\n",
        "    min_file_name : str\n",
        "        Name of the minimized PDB file to write.\n",
        "    \"\"\"\n",
        "    system = omm_forcefield.createSystem(modeller.topology, nonbondedMethod=PME, nonbondedCutoff=1*nanometers, constraints=HBonds,)\n",
        "\n",
        "    # Define platform properties\n",
        "    platform = Platform.getPlatformByName('CUDA')\n",
        "    properties = {'CudaPrecision': 'mixed'}\n",
        "\n",
        "    # Set up the simulation parameters\n",
        "    # Langevin integrator at 300 K w/ 1 ps^-1 friction coefficient\n",
        "    # and a 2-fs timestep\n",
        "    # NOTE - no dynamics performed, but required for setting up\n",
        "    # the OpenMM system.\n",
        "    integrator = LangevinIntegrator(300*kelvin, 1/picosecond,\n",
        "                                    0.002*picoseconds)\n",
        "    simulation = Simulation(topology, system, integrator, platform,\n",
        "                            properties)\n",
        "    simulation.context.setPositions(input_positions)\n",
        "\n",
        "    # Minimize the system - no predefined number of steps\n",
        "    simulation.minimizeEnergy()\n",
        "\n",
        "    # Write out the minimized system to use w/ MDAnalysis\n",
        "    positions = simulation.context.getState(getPositions=True).getPositions()\n",
        "    out_file = os.path.join(out_dir,min_file_name)\n",
        "    PDBFile.writeFile(simulation.topology, positions,\n",
        "                      open(out_file, 'w'))\n",
        "\n",
        "    return None\n",
        "\n",
        "\n",
        "def equilibrate(min_pdb, parm, out_dir, eq_file_name):\n",
        "    \"\"\"A function that does a 500 ps NVT equilibration with position\n",
        "    restraints, with a 5 kcal/mol/A**2 harmonic constant on solute heavy\n",
        "    atoms, using a 2 fs timestep.\n",
        "\n",
        "    Parameters\n",
        "    ----------\n",
        "    min_pdb : str\n",
        "        Name of the minimized PDB file.\n",
        "    parm : Parmed or OpenMM parameter file object\n",
        "        Used to create the OpenMM System object.\n",
        "    out_dir : str\n",
        "        Directory to write the outputs to.\n",
        "    eq_file_name : str\n",
        "        Name of the equilibrated PDB file to write.\n",
        "    \"\"\"\n",
        "    # Get the solute heavy atom indices to use\n",
        "    # for defining position restraints during equilibration\n",
        "    universe = mda.Universe(min_pdb,\n",
        "                            format='XPDB', in_memory=True)\n",
        "    solute_heavy_atom_idx = universe.select_atoms('not resname WAT and\\\n",
        "                                                   not resname SOL and\\\n",
        "                                                   not resname HOH and\\\n",
        "                                                   not resname CL and \\\n",
        "                                                   not resname NA and \\\n",
        "                                                   not resname K and \\\n",
        "                                                   not name H*').indices\n",
        "    # Necessary conversion to int from numpy.int64,\n",
        "    # b/c it breaks OpenMM C++ function\n",
        "    solute_heavy_atom_idx = [int(idx) for idx in solute_heavy_atom_idx]\n",
        "\n",
        "    # Add the restraints.\n",
        "    # We add a dummy atoms with no mass, which are therefore unaffected by\n",
        "    # any kind of scaling done by barostat (if used). And the atoms are\n",
        "    # harmonically restrained to the dummy atom. We have to redefine the\n",
        "    # system, b/c we're adding new particles and this would clash with\n",
        "    # modeller.topology.\n",
        "    system = omm_forcefield.createSystem(modeller.topology, nonbondedMethod=PME, nonbondedCutoff=1*nanometers, constraints=HBonds)\n",
        "\n",
        "    # Add the harmonic restraints on the positions\n",
        "    # of specified atoms\n",
        "    restraint = HarmonicBondForce()\n",
        "    restraint.setUsesPeriodicBoundaryConditions(True)\n",
        "    system.addForce(restraint)\n",
        "    nonbonded = [force for force in system.getForces()\n",
        "                 if isinstance(force, NonbondedForce)][0]\n",
        "    dummyIndex = []\n",
        "    input_positions = PDBFile(min_pdb).getPositions()\n",
        "    positions = input_positions\n",
        "    # Go through the indices of all atoms that will be restrained\n",
        "    for i in solute_heavy_atom_idx:\n",
        "        j = system.addParticle(0)\n",
        "        # ... and add a dummy/ghost atom next to it\n",
        "        nonbonded.addParticle(0, 1, 0)\n",
        "        # ... that won't interact with the restrained atom\n",
        "        nonbonded.addException(i, j, 0, 1, 0)\n",
        "        # ... but will be have a harmonic restraint ('bond')\n",
        "        # between the two atoms\n",
        "        restraint.addBond(i, j, 0*nanometers,\n",
        "                          5*kilocalories_per_mole/angstrom**2)\n",
        "        dummyIndex.append(j)\n",
        "        input_positions.append(positions[i])\n",
        "\n",
        "    integrator = LangevinIntegrator(300*kelvin, 1/picosecond,\n",
        "                                    0.002*picoseconds)\n",
        "    platform = Platform.getPlatformByName('CUDA')\n",
        "    properties = {'CudaPrecision': 'mixed'}\n",
        "    sim = Simulation(topology, system, integrator,\n",
        "                     platform, properties)\n",
        "    sim.context.setPositions(input_positions)\n",
        "    integrator.step(250000)  # run 500 ps of equilibration\n",
        "    all_positions = sim.context.getState(\n",
        "        getPositions=True, enforcePeriodicBox=True).getPositions()\n",
        "    # we don't want to write the dummy atoms, so we only\n",
        "    # write the positions of atoms up to the first dummy atom index\n",
        "    relevant_positions = all_positions[:dummyIndex[0]]\n",
        "    out_file = os.path.join(out_dir,eq_file_name)\n",
        "    PDBFile.writeFile(sim.topology, relevant_positions,\n",
        "                      open(out_file, 'w'))\n",
        "\n",
        "    return None\n",
        "\n",
        "\n",
        "def produce(out_dir, idx, lig_resname, eq_pdb, parm, parm_file,\n",
        "            coords_file, set_hill_height):\n",
        "    \"\"\"An OpenBPMD production simulation function. Ligand RMSD is biased with\n",
        "    metadynamics. The integrator uses a 4 fs time step and\n",
        "    runs for 10 ns, writing a frame every 100 ps.\n",
        "\n",
        "    Writes a 'trj.dcd', 'COLVAR.npy', 'bias_*.npy' and 'sim_log.csv' files\n",
        "    during the metadynamics simulation in the '{out_dir}/rep_{idx}' directory.\n",
        "    After the simulation is done, it analyses the trajectories and writes a\n",
        "    'bpm_results.csv' file with time-resolved PoseScore and ContactScore.\n",
        "\n",
        "    Parameters\n",
        "    ----------\n",
        "    out_dir : str\n",
        "        Directory where your equilibration PDBs and 'rep_*' dirs are at.\n",
        "    idx : int\n",
        "        Current replica index.\n",
        "    lig_resname : str\n",
        "        Residue name of the ligand.\n",
        "    eq_pdb : str\n",
        "        Name of the PDB for equilibrated system.\n",
        "    parm : Parmed or OpenMM parameter file object\n",
        "        Used to create the OpenMM System object.\n",
        "    parm_file : str\n",
        "        The name of the parameter or topology file of the system.\n",
        "    coords_file : str\n",
        "        The name of the coordinate file of the system.\n",
        "    set_hill_height : float\n",
        "        Metadynamic hill height, in kcal/mol.\n",
        "\n",
        "    \"\"\"\n",
        "    # First, assign the replica directory to which we'll write the files\n",
        "    write_dir = os.path.join(out_dir,f'rep_{idx}')\n",
        "    # Get the anchor atoms by ...\n",
        "    universe = mda.Universe(eq_pdb,\n",
        "                            format='XPDB', in_memory=True)\n",
        "    # ... finding the protein's COM ...\n",
        "    prot_com = universe.select_atoms('protein').center_of_mass()\n",
        "    x, y, z = prot_com[0], prot_com[1], prot_com[2]\n",
        "    # ... and taking the heavy backbone atoms within 5A of the COM\n",
        "    sel_str = f'point {x} {y} {z} 5 and backbone and not name H*'\n",
        "    anchor_atoms = universe.select_atoms(sel_str)\n",
        "    # ... or 10 angstrom\n",
        "    if len(anchor_atoms) == 0:\n",
        "        sel_str = f'point {x} {y} {z} 10 and backbone and not name H*'\n",
        "        anchor_atoms = universe.select_atoms(sel_str)\n",
        "\n",
        "    anchor_atom_idx = anchor_atoms.indices.tolist()\n",
        "\n",
        "    # Get indices of ligand heavy atoms\n",
        "    lig = universe.select_atoms(f'resname {lig_resname} and not name H*')\n",
        "\n",
        "    lig_ha_idx = lig.indices.tolist()\n",
        "\n",
        "    # Set up the system to run metadynamics\n",
        "    if continue_simulation == \"yes\":\n",
        "      topology = PDBFile(eq_pdb).getTopology()\n",
        "      with open(os.path.join(workDir, 'system_openMM.xml')) as input:\n",
        "        system = XmlSerializer.deserialize(input.read())\n",
        "    else:\n",
        "      topology = modeller.getTopology()\n",
        "      system = omm_forcefield.createSystem(modeller.topology, nonbondedMethod=PME, nonbondedCutoff=1*nanometers, constraints=HBonds, hydrogenMass=4*amu)\n",
        "      openMM_system = os.path.join(workDir, 'system_openMM.xml')\n",
        "      openMM_system_check = os.path.exists(openMM_system)\n",
        "      if openMM_system_check == False:\n",
        "        with open(os.path.join(workDir, 'system_openMM.xml'), 'w') as output:\n",
        "          output.write(XmlSerializer.serialize(system))\n",
        "      else:\n",
        "        pass\n",
        "\n",
        "    # get the atom positions for the system from the equilibrated\n",
        "    # system\n",
        "    input_positions = PDBFile(eq_pdb).getPositions()\n",
        "\n",
        "    # Add an 'empty' flat-bottom restraint to fix the issue with PBC.\n",
        "    # Without one, RMSDForce object fails to account for PBC.\n",
        "    k = 0*kilojoules_per_mole  # NOTE - 0 kJ/mol constant\n",
        "    upper_wall = 10.00*nanometer\n",
        "    fb_eq = '(k/2)*max(distance(g1,g2) - upper_wall, 0)^2'\n",
        "    upper_wall_rest = CustomCentroidBondForce(2, fb_eq)\n",
        "    upper_wall_rest.addGroup(lig_ha_idx)\n",
        "    upper_wall_rest.addGroup(anchor_atom_idx)\n",
        "    upper_wall_rest.addBond([0, 1])\n",
        "    upper_wall_rest.addGlobalParameter('k', k)\n",
        "    upper_wall_rest.addGlobalParameter('upper_wall', upper_wall)\n",
        "    upper_wall_rest.setUsesPeriodicBoundaryConditions(True)\n",
        "    system.addForce(upper_wall_rest)\n",
        "\n",
        "    alignment_indices = lig_ha_idx + anchor_atom_idx\n",
        "\n",
        "    rmsd = RMSDForce(input_positions, alignment_indices)\n",
        "    # Set up the typical metadynamics parameters\n",
        "    grid_min, grid_max = 0.0, 1.0  # nm\n",
        "    hill_height = set_hill_height*kilocalories_per_mole\n",
        "    hill_width = 0.002  # nm, also known as sigma\n",
        "\n",
        "    grid_width = hill_width / 5\n",
        "    # 'grid' here refers to the number of grid points\n",
        "    grid = int(abs(grid_min - grid_max) / grid_width)\n",
        "\n",
        "    rmsd_cv = BiasVariable(rmsd, grid_min, grid_max, hill_width,\n",
        "                           False, gridWidth=grid)\n",
        "\n",
        "    # define the metadynamics object\n",
        "    # deposit bias every 1 ps, BF = 4, write bias every ns\n",
        "    meta = Metadynamics(system, [rmsd_cv], 300.0*kelvin, 4.0, hill_height,\n",
        "                        250, biasDir=write_dir,\n",
        "                        saveFrequency=250000)\n",
        "\n",
        "    # Set up and run metadynamics\n",
        "    integrator = LangevinIntegrator(300*kelvin, 1.0/picosecond,\n",
        "                                    0.004*picoseconds)\n",
        "    platform = Platform.getPlatformByName('CUDA')\n",
        "    properties = {'CudaPrecision': 'mixed'}\n",
        "\n",
        "    simulation = Simulation(topology, system, integrator, platform,\n",
        "                            properties)\n",
        "    simulation.context.setPositions(input_positions)\n",
        "\n",
        "    trj_name = os.path.join(write_dir,'trj.dcd')\n",
        "\n",
        "    sim_time = 10  # ns\n",
        "    steps = 250000 * sim_time\n",
        "\n",
        "    simulation.reporters.append(DCDReporter(trj_name, 25000))  # every 100 ps\n",
        "    simulation.reporters.append(StateDataReporter(stdout, 25000,\n",
        "                                step=True, speed=True, progress=True,\n",
        "                                totalSteps=steps, remainingTime=True,\n",
        "                                separator='\\t\\t'))\n",
        "    simulation.reporters.append(StateDataReporter(\n",
        "                                os.path.join(write_dir,'sim_log.csv'), 250000,\n",
        "                                step=True, temperature=True, progress=True,\n",
        "                                remainingTime=True, speed=True,\n",
        "                                totalSteps=steps, separator=','))  # every 1 ns\n",
        "\n",
        "    print(\"\\n> Simulating \" + str(steps) + \" steps...\")\n",
        "\n",
        "    colvar_array = np.array([meta.getCollectiveVariables(simulation)])\n",
        "    for i in range(0, int(steps), 500):\n",
        "        if i % 25000 == 0:\n",
        "            # log the stored COLVAR every 100ps\n",
        "            np.save(os.path.join(write_dir,'COLVAR.npy'), colvar_array)\n",
        "        meta.step(simulation, 500)\n",
        "        current_cvs = meta.getCollectiveVariables(simulation)\n",
        "        # record the CVs every 2 ps\n",
        "        colvar_array = np.append(colvar_array, [current_cvs], axis=0)\n",
        "    np.save(os.path.join(write_dir,'COLVAR.npy'), colvar_array)\n",
        "    # center everything using MDTraj, to fix any PBC imaging issues\n",
        "    # mdtraj can't use GMX TOP, so we have to specify the GRO file instead\n",
        "    mdtraj_top = parm_file\n",
        "    mdu = md.load(trj_name, top=mdtraj_top)\n",
        "    mdu.image_molecules()\n",
        "    mdu.save(trj_name)\n",
        "\n",
        "    return None\n",
        "\n",
        "\n",
        "def collect_results(in_dir, out_dir):\n",
        "    \"\"\"A function that collects the time-resolved BPM results,\n",
        "    takes the scores from last 2 ns of the simulation, averages them\n",
        "    and writes that average as the final score for a given pose.\n",
        "\n",
        "    Writes a 'results.csv' file in 'out_dir' directory.\n",
        "\n",
        "    Parameters\n",
        "    ----------\n",
        "    in_dir : str\n",
        "        Directory with 'rep_*' directories.\n",
        "    out_dir : str\n",
        "        Directory where the 'results.csv' file will be written\n",
        "    \"\"\"\n",
        "    compList = []\n",
        "    contactList = []\n",
        "    poseList = []\n",
        "    # find how many repeats have been run\n",
        "    glob_str = os.path.join(in_dir,'rep_*')\n",
        "    nreps = len(glob.glob(glob_str))\n",
        "    for idx in range(1, (nreps+1)):\n",
        "        f = os.path.join(in_dir,f'rep_{idx}','bpm_results.csv')\n",
        "        df = pd.read_csv(f)\n",
        "        # Since we only want last 2 ns, get the index of\n",
        "        # the last 20% of the data points\n",
        "        last_2ns_idx = round(len(df['CompScore'].values)/5)  # round up\n",
        "        compList.append(df['CompScore'].values[-last_2ns_idx:])\n",
        "        contactList.append(df['ContactScore'].values[-last_2ns_idx:])\n",
        "        poseList.append(df['PoseScore'].values[-last_2ns_idx:])\n",
        "\n",
        "    # Get the means of the last 2 ns\n",
        "    meanCompScore = np.mean(compList)\n",
        "    meanPoseScore = np.mean(poseList)\n",
        "    meanContact = np.mean(contactList)\n",
        "    # Get the standard deviation of the final 2 ns\n",
        "    meanCompScore_std = np.std(compList)\n",
        "    meanPoseScore_std = np.std(poseList)\n",
        "    meanContact_std = np.std(contactList)\n",
        "    # Format it the Pandas way\n",
        "    d = {'CompScore': [meanCompScore], 'CompScoreSD': [meanCompScore_std],\n",
        "         'PoseScore': [meanPoseScore], 'PoseScoreSD': [meanPoseScore_std],\n",
        "         'ContactScore': [meanContact], 'ContactScoreSD': [meanContact_std]}\n",
        "\n",
        "    results_df = pd.DataFrame(data=d)\n",
        "    results_df = results_df.round(3)\n",
        "    results_df.to_csv(os.path.join(out_dir,'results.csv'), index=False)\n",
        "\n",
        "main()"
      ],
      "metadata": {
        "id": "b7GrsMTGTV__",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "##  **Analysing BPMD simulations**\n",
        "\n",
        "Understanding the BPMD results:\n",
        "\n",
        "- The lower (more negative) the CompScore, the more likely a given pose to be the correct pose. 'Correct' here means a binding pose that is similar to a pose observed in an experimentally determined structure.\n",
        "\n",
        "- PoseScore, being a more objective measure than ContactScore, should be given a slighlty higher weight when the CompScores of two poses are very similar.\n",
        "\n",
        "- The standard deviations of the scores might seem really high when compared to the scores themselves. Short metadynamics simulations are known to be very noisy and this is expected. Poses with lower standard deviation also tend to be have a lower RMSD to the known pose. Therefore, standard deviations is a useful indicator of confidence.\n",
        "\n",
        "- CompScores should only be compared between poses of the same ligand, not for comparing stability of different ligands."
      ],
      "metadata": {
        "id": "Tt0X7UGx2oYh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title  **Plot the time-wise BPMD scores for each replica:**\n",
        "#@markdown Select the replica:\n",
        "\n",
        "replica = \"1\" #@param [\"1\",\"2\",\"3\",\"4\",\"5\",\"6\",\"7\",\"8\",\"9\",\"10\"]\n",
        "\n",
        "#@markdown **Important:** Probably, the scores will be very noisy. This is typical of short metadynamics simulations that haven't yet converged the free energy surface. However, we're not interested in the free energy. We're only want to test the stability of the ligand in the binding pose.\n",
        "if int(replica) > int(nreps):\n",
        "  print(\"You have less replicas than requested.\")\n",
        "  exit()\n",
        "else:\n",
        "  pass\n",
        "\n",
        "f = os.path.join(workDir,output,'rep_' + str(int(replica)),'bpm_results.csv')\n",
        "df = pd.read_csv(f)\n",
        "time_sequence = np.linspace(0,10,99)\n",
        "\n",
        "plt.title('CompScore from ' + str(replica))\n",
        "plt.plot(time_sequence,df['CompScore'])\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "plt.xlabel('time(ns)')\n",
        "plt.ylabel('CompScore')\n",
        "plt.ylim(-5,5)\n",
        "plt.show()\n",
        "plt.savefig(os.path.join(workDir, \"CompScore_\" + str(replica) + \"_BPMD.png\"), dpi=600, bbox_inches='tight')\n",
        "raw_data=pd.DataFrame(df['CompScore'])\n",
        "raw_data.to_csv(os.path.join(workDir, \"CompScore_\" + str(replica) + \"_BPMD.csv\"))\n",
        "\n",
        "plt.title('PoseScore from 1 ' + str(replica))\n",
        "plt.plot(time_sequence,df['PoseScore'],color='darkorange')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "plt.xlabel('time(ns)')\n",
        "plt.ylabel('PoseScore')\n",
        "plt.ylim(0,5)\n",
        "plt.show()\n",
        "plt.savefig(os.path.join(workDir, \"PoseScore_\" + str(replica) + \"_BPMD.png\"), dpi=600, bbox_inches='tight')\n",
        "raw_data=pd.DataFrame(df['PoseScore'])\n",
        "raw_data.to_csv(os.path.join(workDir, \"PoseScore_\" + str(replica) + \"_BPMD.csv\"))\n",
        "\n",
        "plt.title('ContactScore from 1 ' + str(replica))\n",
        "plt.plot(time_sequence,df['ContactScore'],color='green')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "plt.xlabel('time(ns)')\n",
        "plt.ylabel('ContactScore')\n",
        "plt.ylim(-0.1,1.1)\n",
        "plt.show()\n",
        "plt.savefig(os.path.join(workDir, \"ContactScore_\" + str(replica) + \"_BPMD.png\"), dpi=600, bbox_inches='tight')\n",
        "raw_data=pd.DataFrame(df['ContactScore'])\n",
        "raw_data.to_csv(os.path.join(workDir, \"ContactScore_\" + str(replica) + \"_BPMD.csv\"))"
      ],
      "metadata": {
        "cellView": "form",
        "id": "pF4pReqAv0qi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title  **Plot the time-wise BPMD scores average over multiple replicas :**\n",
        "#@markdown **Important:** To have more confidence in the stability scores, we run multiple repeat simulations. In the final PoseScore and ContactScore, after N repeats, we take the mean of the score over the last 2 ns, which also helps with the noise.\n",
        "\n",
        "# We'll store the results from 10 repeats in a 10x99 matrix.\n",
        "CompScores = np.zeros((10,99))\n",
        "PoseScores = np.zeros((10,99))\n",
        "ContactScores = np.zeros((10,99))\n",
        "\n",
        "# Fill those matrices with the scores from each repeat\n",
        "for idx in range(1,(nreps+1)):\n",
        "    f = os.path.join(workDir,output,f'rep_{idx}','bpm_results.csv')\n",
        "    df = pd.read_csv(f)\n",
        "    CompScores[idx] = df['CompScore']\n",
        "    PoseScores[idx] = df['PoseScore']\n",
        "    ContactScores[idx] = df['ContactScore']\n",
        "\n",
        "# Average out the scores from all of the repeats,\n",
        "# giving a mean of the scores at each frame of the trajectory\n",
        "averagedCompScore = np.array([ np.mean(CompScores[:,i]) for i in range(0,99) ])\n",
        "averagedPoseScore = [ np.mean(PoseScores[:,i]) for i in range(0,99) ]\n",
        "averagedContactScore = [ np.mean(ContactScores[:,i]) for i in range(0,99) ]\n",
        "# Get the standard deviation for the CompScore\n",
        "CompScore_stddev = np.array([ np.std(CompScores[:,i]) for i in range(0,99) ])\n",
        "# An array of time steps for plotting the x axis\n",
        "time_sequence = np.linspace(0,10,99)\n",
        "\n",
        "plt.title('CompScore from ' + str(nreps) + ' repeats')\n",
        "plt.plot(time_sequence,averagedCompScore, color='blue')\n",
        "# Visualise the standard deviation of CompScore at each frame\n",
        "plt.fill_between(time_sequence, averagedCompScore-CompScore_stddev, averagedCompScore+CompScore_stddev,\n",
        "                 color='blue', alpha=0.3, lw=0)\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "plt.xlabel('time(ns)')\n",
        "plt.ylabel('CompScore')\n",
        "plt.ylim(-5,5)\n",
        "plt.show()\n",
        "plt.savefig(os.path.join(workDir, \"CompScore_mean_BPMD.png\"), dpi=600, bbox_inches='tight')\n",
        "raw_data=pd.DataFrame(averagedCompScore)\n",
        "raw_data.to_csv(os.path.join(workDir, \"CompScore_mean_BPMD.csv\"))\n",
        "\n",
        "plt.title('PoseScore from ' + str(nreps) + ' repeats')\n",
        "time_sequence = np.linspace(0,10,99)\n",
        "plt.plot(time_sequence,averagedPoseScore,color='darkorange')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "plt.xlabel('time(ns)')\n",
        "plt.ylabel('PoseScore')\n",
        "plt.ylim(0,5)\n",
        "plt.show()\n",
        "plt.savefig(os.path.join(workDir, \"PoseScore_mean_BPMD.png\"), dpi=600, bbox_inches='tight')\n",
        "raw_data=pd.DataFrame(averagedPoseScore)\n",
        "raw_data.to_csv(os.path.join(workDir, \"PoseScore_mean_BPMD.csv\"))\n",
        "\n",
        "plt.title('ContactScore from ' + str(nreps) + ' repeats')\n",
        "time_sequence = np.linspace(0,10,99)\n",
        "plt.plot(time_sequence,averagedContactScore,color='green')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "plt.xlabel('time(ns)')\n",
        "plt.ylabel('ContactScore')\n",
        "plt.ylim(-0.1,1.1)\n",
        "plt.show()\n",
        "plt.savefig(os.path.join(workDir, \"ContactScore_mean_BPMD.png\"), dpi=600, bbox_inches='tight')\n",
        "raw_data=pd.DataFrame(averagedContactScore)\n",
        "raw_data.to_csv(os.path.join(workDir, \"ContactScore_mean_BPMD.csv\"))\n",
        "\n",
        "compList = []\n",
        "contactList = []\n",
        "poseList = []\n",
        "# Find how many repeats have been run\n",
        "glob_str = os.path.join(workDir,output,'rep_*')\n",
        "nreps = len(glob.glob(glob_str))\n",
        "for idx in range(1, (nreps+1)):\n",
        "    f = os.path.join(workDir,output,f'rep_{idx}','bpm_results.csv')\n",
        "    df = pd.read_csv(f)\n",
        "    # Since we only want last 2 ns, get the index of\n",
        "    # the last 20% of the data points\n",
        "    last_2ns_idx = round(len(df['CompScore'].values)/5)  # round up\n",
        "    compList.append(df['CompScore'].values[-last_2ns_idx:])\n",
        "    contactList.append(df['ContactScore'].values[-last_2ns_idx:])\n",
        "    poseList.append(df['PoseScore'].values[-last_2ns_idx:])\n",
        "\n",
        "# Get the means of the last 2 ns\n",
        "meanCompScore = np.mean(compList)\n",
        "meanPoseScore = np.mean(poseList)\n",
        "meanContact = np.mean(contactList)\n",
        "# Get the standard deviation of the final 2 ns\n",
        "meanCompScore_std = np.std(compList)\n",
        "meanPoseScore_std = np.std(poseList)\n",
        "meanContact_std = np.std(contactList)\n",
        "# Format it the Pandas way\n",
        "d = {'CompScore': [meanCompScore], 'CompScoreSD': [meanCompScore_std],\n",
        "     'PoseScore': [meanPoseScore], 'PoseScoreSD': [meanPoseScore_std],\n",
        "     'ContactScore': [meanContact], 'ContactScoreSD': [meanContact_std]}\n",
        "\n",
        "results_df = pd.DataFrame(data=d)\n",
        "results_df = results_df.round(3)\n",
        "print(\"In order to get a single number that evaluates a given pose/ligand, we take the scores of the final 2 ns.\")\n",
        "results_df"
      ],
      "metadata": {
        "id": "0w0UVaCUpFrk",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import MDAnalysis as mda\n",
        "from MDAnalysis.analysis import align, rms\n",
        "import subprocess\n",
        "import pandas as pd\n",
        "from rdkit import Chem\n",
        "from rdkit.Chem import rdFMCS,AllChem, Draw, PandasTools\n",
        "import seaborn as sns\n",
        "#@title **Evaluate the poses over the ensemble with the pre-trained GNINA models**\n",
        "#@markdown Pose (CNNscore) and binding affinity (Affinity Score and CNNaffinity) predictions with GNINA models.\n",
        "\n",
        "# @markdown Minimize ligands already positioned in a binding site:\n",
        "minimization = \"Yes\" #@param [\"Yes\",\"No\"]\n",
        "Skip = 10 #@param {type:\"slider\", min:1, max:100, step:1}\n",
        "pdb_dir = os.path.join(workDir,'PDBs/')\n",
        "pdb_dir_check = os.path.exists(pdb_dir)\n",
        "if pdb_dir_check == False:\n",
        "  os.mkdir(pdb_dir)\n",
        "else:\n",
        "  pass\n",
        "\n",
        "flist = []\n",
        "pdb = os.path.join(output,f\"centred_equil_system.pdb\")\n",
        "ref = os.path.join(str(output) + '/rep_' + str(1) + '/trj.dcd')\n",
        "for i in range(1, int(nreps+1)):\n",
        "  flist.append(os.path.join(str(output) + '/rep_' + str(i) + '/trj.dcd'))\n",
        "\n",
        "u1 = mda.Universe(pdb, flist)\n",
        "u2 = mda.Universe(pdb, ref)\n",
        "u2.trajectory[0] # set u2 to first frame\n",
        "\n",
        "align.AlignTraj(u1, u2, select='name CA', in_memory=True).run()\n",
        "\n",
        "#Save a trajectory with protein + ligand in PDB format\n",
        "protein_ligand = u1.select_atoms(\"not (resname HOH) and not (resname NA) and not (resname CL) and not (resname K)\")\n",
        "with mda.Writer(os.path.join(output,f'protein_ligand_GNINA.dcd'), protein_ligand.n_atoms) as W:\n",
        "  for ts in u1.trajectory[0:len(u1.trajectory):int(Skip)]:\n",
        "    W.write(protein_ligand)\n",
        "\n",
        "if minimization == 'Yes':\n",
        "  minimize = \"--minimize\"\n",
        "else:\n",
        "  minimize = \"--score_only\"\n",
        "\n",
        "affinity = []\n",
        "CNNscore = []\n",
        "CNNaffinity = []\n",
        "protein_atoms = u1.select_atoms(\"protein\")\n",
        "ligand_atoms = u1.select_atoms(\"resname UNK\")\n",
        "i = 0\n",
        "for ts in u1.trajectory[0:len(u1.trajectory):int(Skip)]:\n",
        "  if i > -1:\n",
        "    with mda.Writer(os.path.join(pdb_dir,'protein.pdb'), protein_atoms.n_atoms) as W:\n",
        "      W.write(protein_atoms)\n",
        "    with mda.Writer(os.path.join(pdb_dir,'ligand.pdb'), ligand_atoms.n_atoms) as W:\n",
        "      W.write(ligand_atoms)\n",
        "  gnina_score = \"./gnina -r \" + str(pdb_dir) + \"protein.pdb\" + \" -l \" +  str(pdb_dir) + \"ligand.pdb \" + str(minimize) + \" -o score_gnina.sdf\"\n",
        "  original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "  with open('gnina_score.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(gnina_score)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "  print(\"Running frame \"+ str(i))\n",
        "  subprocess.run([\"chmod 700 gnina_score.sh\"], shell=True)\n",
        "  subprocess.run([\"./gnina_score.sh\"], shell=True)\n",
        "  result=Chem.SDMolSupplier(\"score_gnina.sdf\",True)\n",
        "  affinity.append(float(result[0].GetProp('minimizedAffinity')))\n",
        "  CNNscore.append(float(result[0].GetProp('CNNscore')))\n",
        "  CNNaffinity.append(float(result[0].GetProp('CNNaffinity')))\n",
        "  i = i + 1\n",
        "\n",
        "raw_data=pd.DataFrame(affinity)\n",
        "raw_data.to_csv(os.path.join(output,f\"GNINA_affinity.csv\"))\n",
        "raw_data=pd.DataFrame(CNNscore)\n",
        "raw_data.to_csv(os.path.join(output,f\"GNINA_CNNscore.csv\"))\n",
        "raw_data=pd.DataFrame(CNNaffinity)\n",
        "raw_data.to_csv(os.path.join(output,f\"GNINA_CNNaffinity.csv\"))"
      ],
      "metadata": {
        "cellView": "form",
        "id": "lIIacmKcOI3_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Plot the Affinity Score, CNNscore and CNNaffinity, calculated with GNINA models**\n",
        "\n",
        "\n",
        "time_sequence = np.linspace(0,int(len(u1.trajectory)/int(Skip)),int(len(u1.trajectory)/int(Skip)))\n",
        "\n",
        "plt.title('Affinity Score')\n",
        "plt.plot(time_sequence,affinity)\n",
        "plt.xlim(0, time_sequence[-1])\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "plt.xlabel('Frames')\n",
        "plt.ylabel('Affinity (kcal/mol)')\n",
        "plt.show()\n",
        "plt.savefig(os.path.join(workDir, \"Affinity_Score_BPMD.png\"), dpi=600, bbox_inches='tight')\n",
        "raw_data=pd.DataFrame(affinity)\n",
        "raw_data.to_csv(os.path.join(workDir, \"Affinity_Score_BPMD.csv\"))\n",
        "\n",
        "plt.title('CNNscore')\n",
        "plt.plot(time_sequence,CNNscore,color='darkorange')\n",
        "plt.xlim(0, time_sequence[-1])\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "plt.xlabel('Frames')\n",
        "plt.ylabel('CNNscore')\n",
        "plt.show()\n",
        "plt.savefig(os.path.join(workDir, \"CNNscore_BPMD.png\"), dpi=600, bbox_inches='tight')\n",
        "raw_data=pd.DataFrame(CNNscore)\n",
        "raw_data.to_csv(os.path.join(workDir, \"CNNscore_BPMD.csv\"))\n",
        "\n",
        "plt.title('CNNaffinity')\n",
        "plt.plot(time_sequence,CNNaffinity,color='green')\n",
        "plt.xlim(0, time_sequence[-1])\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "plt.xlabel('Frames')\n",
        "plt.ylabel('CNNaffinity')\n",
        "plt.show()\n",
        "plt.savefig(os.path.join(workDir, \"CNNaffinity_BPMD.png\"), dpi=600, bbox_inches='tight')\n",
        "raw_data=pd.DataFrame(CNNaffinity)\n",
        "raw_data.to_csv(os.path.join(workDir, \"CNNaffinity_BPMD.csv\"))"
      ],
      "metadata": {
        "id": "S_jZB4wXraR9",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Best pose selection from BPMD simulations:**\n",
        "\n",
        "Frame = \"1\" #@param {type:\"string\"}\n",
        "\n",
        "if int(Frame) > (len(u1.trajectory)/int(Skip)-int(1)):\n",
        "  print(\"You didn't save enough frames in your trajectory. Please, run the previous cell again with a lower skip.\")\n",
        "  exit()\n",
        "else:\n",
        "  pass\n",
        "\n",
        "protein_ligand = u1.select_atoms(\"not (resname HOH) and not (resname NA) and not (resname CL) and not (resname K)\")\n",
        "with mda.Writer(os.path.join(output,f'protein_ligand_' + str(Frame) + '.pdb'), protein_ligand.n_atoms) as W:\n",
        "  for ts in u1.trajectory[int(Frame):(int(Frame)+1):1]:\n",
        "    W.write(protein_ligand)\n",
        "\n",
        "\n",
        "view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',)\n",
        "view.removeAllModels()\n",
        "view.setViewStyle({'style':'outline','color':'black','width':0.05})\n",
        "\n",
        "view.addModel(open(os.path.join(output,f'protein_ligand_' + str(Frame) + '.pdb'),'r').read(),'pdb')\n",
        "Prot=view.getModel()\n",
        "# Prot.setStyle({'cartoon':{'arrows':True, 'tubes':True, 'style':'oval', 'color':'white'},'stick':{'radius':.1}})\n",
        "Prot.setStyle({'cartoon':{'arrows':True, 'tubes':True, 'style':'oval', 'color':'white'}})\n",
        "view.addSurface(py3Dmol.VDW,{'opacity':0.2,'color':'white'})\n",
        "\n",
        "print ('Affinity: {}'.format(affinity[(int(Frame)-1)]))\n",
        "print ('CNN Score: {}'.format(CNNscore[(int(Frame)-1)]))\n",
        "print ('CNN Affinity: {}'.format(CNNaffinity[(int(Frame)-1)]))\n",
        "\n",
        "HP = ['UNK']\n",
        "view.addStyle({'and':[{'resn':HP}]},\n",
        "              {'stick':{'colorscheme':'greenCarbon','radius':0.3}})\n",
        "view.setViewStyle({'style':'outline','color':'black','width':0.1})\n",
        "\n",
        "\n",
        "view.zoomTo()\n",
        "view.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "3pdCUueu9wtU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **View and check the Ligand Interaction Network (LigPlot) on the best pose**\n",
        "#@markdown This diagram is interactive and allows moving around the residues, as well as clicking the legend to toggle the display of specific residues types or interactions. The diagram will be saved as an HTML file (frame_number.html).\n",
        "Frame = \"1\" #@param {type:\"string\"}\n",
        "\n",
        "if int(Frame) > (len(u1.trajectory)/int(Skip)-int(1)):\n",
        "  print(\"You didn't save enough frames in your trajectory. Please, run the previous cell again with a lower skip.\")\n",
        "  exit()\n",
        "else:\n",
        "  pass\n",
        "\n",
        "import MDAnalysis as mda\n",
        "import prolif as plf\n",
        "import numpy as np\n",
        "import os\n",
        "from prolif.plotting.network import LigNetwork\n",
        "\n",
        "\n",
        "protein_ligand = u1.select_atoms(\"all\")\n",
        "with mda.Writer(os.path.join(output,f'protein_ligand_' + str(Frame) + '.pdb'), protein_ligand.n_atoms) as W:\n",
        "  for ts in u1.trajectory[int(Frame):(int(Frame)+1):1]:\n",
        "    W.write(protein_ligand)\n",
        "\n",
        "# load topology\n",
        "u = mda.Universe(os.path.join(workDir,\"system_amber.prmtop\"), os.path.join(output,f'protein_ligand_' + str(Frame) + '.pdb'))\n",
        "lig = u.select_atoms(\"not protein and not (resname HOH) and not (resname CL) and not (resname NA) and not (resname K)\")\n",
        "prot = u.select_atoms(\"protein\")\n",
        "\n",
        "# create RDKit-like molecules for visualisation\n",
        "lmol = plf.Molecule.from_mda(lig)\n",
        "pmol = plf.Molecule.from_mda(prot)\n",
        "\n",
        "fp = plf.Fingerprint()\n",
        "fp.run(u.trajectory[::5], lig, prot)\n",
        "df = fp.to_dataframe(return_atoms=True)\n",
        "\n",
        "net = LigNetwork.from_ifp(df, lmol,\n",
        "                          # replace with `kind=\"frame\", frame=0` for the other depiction\n",
        "                          kind=\"aggregate\", frame=0,\n",
        "                          rotation=270)\n",
        "net.save(os.path.join(workDir, \"frame_\" + str(Frame) + \".html\"))\n",
        "net.display()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "Fn4rgXoa2R3D"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Compute RMSD of the ligand during BPMD simulations**\n",
        "#@markdown **Provide output file names below:**\n",
        "import MDAnalysis.analysis.rms\n",
        "Output_name = 'rmsd_lig' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "u = mda.Universe(os.path.join(output,f\"centred_equil_system.pdb\"), flist)\n",
        "ref = mda.Universe(os.path.join(output,f\"centred_equil_system.pdb\"), os.path.join(str(output) + '/rep_' + str(1) + '/trj.dcd'))\n",
        "ref.trajectory[0] # set u2 to first frame\n",
        "\n",
        "\n",
        "#Save a trajectory with protein + ligand in PDB format\n",
        "R = MDAnalysis.analysis.rms.RMSD(u, ref,\n",
        "           select=\"resname UNK and not (name H*)\",\n",
        "           groupselections=[\"resname UNK and not (name H*)\"])\n",
        "R.run()\n",
        "\n",
        "# Plotting:\n",
        "rmsd = R.rmsd.T   # transpose makes it easier for plotting\n",
        "time = rmsd[0]/10\n",
        "\n",
        "ax = plt.plot(time, rmsd[2], alpha=0.6, color = 'blue', linewidth = 1.0)\n",
        "# plt.xlim(0, time[-1])\n",
        "# plt.ylim(2, 6)\n",
        "\n",
        "plt.xlabel(\"Time (ns)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel(\"RMSD [$\\AA$]\", fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "raw_data=pd.DataFrame(rmsd[2])\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "metadata": {
        "cellView": "form",
        "id": "K7rz94vbEtNk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZHyMpikjuaLT",
        "cellView": "form"
      },
      "source": [
        "#@title **Plot RMSD as a ditribution**\n",
        "\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = 'rmsd_dist' #@param {type:\"string\"}\n",
        "\n",
        "ax = sb.kdeplot(rmsd[2], color=\"blue\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('RMSD [$\\AA$]', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks([])\n",
        "plt.ylabel('')\n",
        "ax.spines['top'].set_visible(False)\n",
        "ax.spines['right'].set_visible(False)\n",
        "ax.spines['bottom'].set_visible(True)\n",
        "ax.spines['left'].set_visible(False)\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}