Uni_Dock+MD.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%2BMD.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 Molecular Dynamics (MD) simulations using OpenMM engine.\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 git+https://github.com/amber-md/pytraj\",\n",
        "    \"pip install bio\",\n",
        "    \"mamba install openmmforcefields -c conda-forge -y\",\n",
        "    \"mamba install openmm\",\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": "IxfNxnIjcoUN"
      },
      "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": "markdown",
      "metadata": {
        "id": "n85MrAO7M7uQ"
      },
      "source": [
        "---\n",
        "---\n",
        "# **Equilibrating the simulation box**\n",
        "\n",
        "Proper MD equilibration protocol is designed to equilibrate both temperature and pressure throughout the simulation box while preserving the protein experimental conformation. In addition, we also allow the solvent to accomodate around the protein, creating proper solvation layers.\n",
        "\n",
        "Below, we will set up the MD equilibration parameters, such as temperature, pressure and the desired simulation time. We will define the force constant used to restraint protein heavy-atoms in place and the frequency at which we want to save atomic coordinates in a trajectory file (.dcd).\n",
        "\n",
        "After you are done, you can run the next 2 cells to equilibrate your system."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8x9Qp_dbr9HP",
        "cellView": "form"
      },
      "source": [
        "#@title ### **Parameters for MD Equilibration protocol:**\n",
        "\n",
        "# remove whitespaces\n",
        "Jobname = 'prot_lig_equil' #@param {type:\"string\"}\n",
        "\n",
        "Minimization_steps = \"10000\" #@param [\"1000\", \"5000\", \"10000\", \"20000\", \"50000\", \"100000\"]\n",
        "\n",
        "#@markdown Simulation time (in nanoseconds) and integration time (in femtoseconds):\n",
        "Time = \"0.5\" #@param {type:\"string\"}\n",
        "stride_time_eq = Time\n",
        "Integration_timestep = \"2\" #@param [\"0.5\", \"1\", \"2\", \"3\", \"4\"]\n",
        "dt_eq = Integration_timestep\n",
        "\n",
        "#@markdown Temperature (in Kelvin) and Pressure (in bar)\n",
        "Temperature = 298 #@param {type:\"string\"}\n",
        "temperature_eq = Temperature\n",
        "Pressure = 1 #@param {type:\"string\"}\n",
        "pressure_eq = Pressure\n",
        "\n",
        "#@markdown Position restraints force constant (in kJ/mol):\n",
        "Force_constant = 700 #@param {type:\"slider\", min:0, max:2000, step:100}\n",
        "\n",
        "#@markdown Frequency to write the trajectory file (in picoseconds):\n",
        "\n",
        "Write_the_trajectory = \"10\" #@param [\"10\", \"100\", \"200\", \"500\", \"1000\"]\n",
        "write_the_trajectory_eq = Write_the_trajectory\n",
        "#@markdown Frequency to write the log file (in picoseconds):\n",
        "\n",
        "Write_the_log = \"10\" #@param [\"10\", \"100\", \"200\", \"500\", \"1000\"]\n",
        "write_the_log_eq = Write_the_log\n",
        "\n",
        "\n",
        "#@markdown ---\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Runs an Equilibration MD simulation (NPT ensemble)**\n",
        "#@markdown Now, let's equilibrate our system!\n",
        "\n",
        "###########################################\n",
        "\n",
        "from sys import stdout, exit, stderr\n",
        "import os, math, fnmatch\n",
        "\n",
        "#############################################\n",
        "# Defining MD simulation parameters\n",
        "\n",
        "jobname = os.path.join(workDir, Jobname)\n",
        "pdbfile = system_pdb\n",
        "\n",
        "time_ps = float(Time)*1000\n",
        "simulation_time = float(time_ps)*picosecond\t\t# in ps\n",
        "dt = int(dt_eq)*femtosecond\n",
        "temperature = float(temperature_eq)*kelvin\n",
        "savcrd_freq = int(write_the_trajectory_eq)*picosecond\n",
        "print_freq  = int(write_the_log_eq)*picosecond\n",
        "\n",
        "pressure\t= float(pressure_eq)*bar\n",
        "\n",
        "restraint_fc = int(Force_constant) # kJ/mol\n",
        "\n",
        "nsteps  = int(simulation_time.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nprint  = int(print_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nsavcrd = int(savcrd_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "\n",
        "#############################################\n",
        "# Defining functions to use below:\n",
        "def backup_old_log(pattern, string):\n",
        "\tresult = []\n",
        "\tfor root, dirs, files in os.walk(\"./\"):\n",
        "\t\tfor name in files:\n",
        "\t\t\tif fnmatch.fnmatch(name, pattern):\n",
        "\n",
        "\t\t\t\ttry:\n",
        "\t\t\t\t\tnumber = int(name[-2])\n",
        "\t\t\t\t\tavail = isinstance(number, int)\n",
        "\t\t\t\t\t#print(name,avail)\n",
        "\t\t\t\t\tif avail == True:\n",
        "\t\t\t\t\t\tresult.append(number)\n",
        "\t\t\t\texcept:\n",
        "\t\t\t\t\tpass\n",
        "\n",
        "\tif len(result) > 0:\n",
        "\t\tmaxnumber = max(result)\n",
        "\telse:\n",
        "\t\tmaxnumber = 0\n",
        "\n",
        "\tbackup_file = \"\\#\" + string + \".\" + str(maxnumber + 1) + \"#\"\n",
        "\tos.system(\"mv \" + string + \" \" + backup_file)\n",
        "\treturn backup_file\n",
        "\n",
        "def restraints(system, crd, fc, restraint_array):\n",
        "\n",
        "\tboxlx = system.getDefaultPeriodicBoxVectors()[0][0].value_in_unit(nanometers)\n",
        "\tboxly = system.getDefaultPeriodicBoxVectors()[1][1].value_in_unit(nanometers)\n",
        "\tboxlz = system.getDefaultPeriodicBoxVectors()[2][2].value_in_unit(nanometers)\n",
        "\n",
        "\tif fc > 0:\n",
        "\t\t# positional restraints for all heavy-atoms\n",
        "\t\tposresPROT = CustomExternalForce('k*periodicdistance(x, y, z, x0, y0, z0)^2;')\n",
        "\t\tposresPROT.addPerParticleParameter('k')\n",
        "\t\tposresPROT.addPerParticleParameter('x0')\n",
        "\t\tposresPROT.addPerParticleParameter('y0')\n",
        "\t\tposresPROT.addPerParticleParameter('z0')\n",
        "\n",
        "\t\tfor atom1 in restraint_array:\n",
        "\t\t\tatom1 = int(atom1)\n",
        "\n",
        "\t\t\txpos  = crd.positions[atom1].value_in_unit(nanometers)[0]\n",
        "\t\t\typos  = crd.positions[atom1].value_in_unit(nanometers)[1]\n",
        "\t\t\tzpos  = crd.positions[atom1].value_in_unit(nanometers)[2]\n",
        "\n",
        "\t\t\tposresPROT.addParticle(atom1, [fc, xpos, ypos, zpos])\n",
        "\n",
        "\t\tsystem.addForce(posresPROT)\n",
        "\n",
        "\treturn system\n",
        "##############################################\n",
        "\n",
        "#############################################\n",
        "print(\"\\n> Simulation details:\\n\")\n",
        "print(\"\\tJob name = \" + jobname)\n",
        "print(\"\\tPDB file = \" + str(pdbfile))\n",
        "\n",
        "print(\"\\n\\tSimulation_time = \" + str(simulation_time))\n",
        "print(\"\\tIntegration timestep = \" + str(dt))\n",
        "print(\"\\tTotal number of steps = \" +  str(nsteps))\n",
        "\n",
        "print(\"\\n\\tSave coordinates each \" + str(savcrd_freq))\n",
        "print(\"\\tPrint in log file each \" + str(print_freq))\n",
        "\n",
        "print(\"\\n\\tTemperature = \" + str(temperature))\n",
        "print(\"\\tPressure = \" + str(pressure))\n",
        "#############################################\n",
        "\n",
        "print(\"\\n> Setting the system:\\n\")\n",
        "\n",
        "print(\"\\t- Reading topology and structure file...\")\n",
        "prmtop = pmd.load_file(amber_prmtop)\n",
        "inpcrd = AmberInpcrdFile(amber_inpcrd)\n",
        "\n",
        "\n",
        "print(\"\\t- Creating system and setting parameters...\")\n",
        "nonbondedMethod = PME\n",
        "nonbondedCutoff = 1.0*nanometers\n",
        "ewaldErrorTolerance = 0.0005\n",
        "constraints = HBonds\n",
        "rigidWater = True\n",
        "constraintTolerance = 0.000001\n",
        "friction = 1.0\n",
        "system = omm_forcefield.createSystem(modeller.topology, nonbondedMethod=nonbondedMethod, nonbondedCutoff=nonbondedCutoff,\n",
        "                                          constraints=constraints, rigidWater=rigidWater, ewaldErrorTolerance=ewaldErrorTolerance)\n",
        "\n",
        "print(\"\\t- Applying restraints. Force Constant = \" + str(Force_constant) + \"kJ/mol\")\n",
        "pt_system = pt.iterload(amber_inpcrd, amber_prmtop)\n",
        "pt_topology = pt_system.top\n",
        "restraint_array = pt.select_atoms('!(:H*) & !(:HOH) & !(:NA) & !(:CL) & !(:K)', pt_topology)\n",
        "\n",
        "system = restraints(system, inpcrd, restraint_fc, restraint_array)\n",
        "\n",
        "print(\"\\t- Setting barostat...\")\n",
        "system.addForce(MonteCarloBarostat(pressure, temperature))\n",
        "\n",
        "print(\"\\t- Setting integrator...\")\n",
        "integrator = LangevinIntegrator(temperature, friction, dt)\n",
        "integrator.setConstraintTolerance(constraintTolerance)\n",
        "simulation = Simulation(topology, system, integrator)\n",
        "simulation.context.setPositions(inpcrd.positions)\n",
        "if inpcrd.boxVectors is not None:\n",
        "\tsimulation.context.setPeriodicBoxVectors(*inpcrd.boxVectors)\n",
        "\n",
        "print(\"\\t- Energy minimization: \" + str(Minimization_steps) + \" steps\")\n",
        "simulation.minimizeEnergy(tolerance=10*kilojoule/mole/nanometer, maxIterations=int(Minimization_steps))\n",
        "\n",
        "print(\"\\t-> Potential Energy = \" + str(simulation.context.getState(getEnergy=True).getPotentialEnergy()))\n",
        "\n",
        "print(\"\\t- Setting initial velocities...\")\n",
        "simulation.context.setVelocitiesToTemperature(temperature)\n",
        "\n",
        "#############################################\n",
        "# Running Equilibration on NPT ensemble\n",
        "\n",
        "dcd_file = jobname + \".dcd\"\n",
        "log_file = jobname + \".log\"\n",
        "rst_file = jobname + \".rst\"\n",
        "prv_rst_file = jobname + \".rst\"\n",
        "pdb_file = jobname + \".pdb\"\n",
        "\n",
        "# Creating a trajectory file and reporters\n",
        "dcd = DCDReporter(dcd_file, nsavcrd)\n",
        "firstdcdstep = (nsteps) + nsavcrd\n",
        "dcd._dcd = DCDFile(dcd._out, topology, simulation.integrator.getStepSize(), firstdcdstep, nsavcrd) # charmm doesn't like first step to be 0\n",
        "\n",
        "simulation.reporters.append(dcd)\n",
        "simulation.reporters.append(StateDataReporter(stdout, nprint, step=True, speed=True, progress=True, totalSteps=nsteps, remainingTime=True, separator='\\t\\t'))\n",
        "simulation.reporters.append(StateDataReporter(log_file, nprint, step=True, kineticEnergy=True, potentialEnergy=True, totalEnergy=True, temperature=True, volume=True, speed=True))\n",
        "\n",
        "print(\"\\n> Simulating \" + str(nsteps) + \" steps...\")\n",
        "simulation.step(nsteps)\n",
        "\n",
        "simulation.reporters.clear() # remove all reporters so the next iteration don't trigger them.\n",
        "\n",
        "\n",
        "##################################\n",
        "# Writing last frame information of stride\n",
        "print(\"\\n> Writing state file (\" + str(rst_file) + \")...\")\n",
        "state = simulation.context.getState( getPositions=True, getVelocities=True )\n",
        "with open(rst_file, 'w') as f:\n",
        "\tf.write(XmlSerializer.serialize(state))\n",
        "\n",
        "last_frame = int(nsteps/nsavcrd)\n",
        "print(\"> Writing coordinate file (\" + str(pdb_file) + \", frame = \" + str(last_frame) + \")...\")\n",
        "positions = simulation.context.getState(getPositions=True).getPositions()\n",
        "PDBFile.writeFile(simulation.topology, positions, open(pdb_file, 'w'))\n",
        "\n",
        "print(\"\\n> Finished!\\n\")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "mBpeYxUYTdDt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "---\n",
        "---\n",
        "# **Running a Production MD simulation**\n",
        "\n",
        "Finally, we will proceed with the Production simulation itself using the equilibrated system coordinates as input structure.\n",
        "\n",
        "Note that we will use here a *.rst state file* , which contains atomic velocities and positions from the last frame of the equilibration simulation, guaranteeing that our production simulation begins from a thermodynamically equilibrated system.\n",
        "\n",
        "Another important information here is the **Number_of_strides** and the **Stride_Time**. In this notebook, we simulate a defined number of *strides*, so the **simulation time = Number_of_strides*Stride_Time**. For example, we can simulate 100ns by setting *Number_of_strides=10* and *Stride_Time=10 ns*.\n",
        "\n",
        "**Important: at the end of the Production simulation, we concatenate all strides to create a complete trajectory file which can be visualized and analyzed**\n",
        "\n",
        "The idea behind this approach is to make use of the intermitent 12h/24h period that Google Colab allows us to use its GPUs."
      ],
      "metadata": {
        "id": "TEa6b6t3SscK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ### **Provide input file names below:**\n",
        "\n",
        "Equilibrated_PDB = 'prot_lig_equil.pdb' #@param {type:\"string\"}\n",
        "State_file = 'prot_lig_equil.rst' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown ### **Parameters for MD Production protocol:**\n",
        "\n",
        "\n",
        "# remove whitespaces\n",
        "Jobname = 'prot_lig_prod' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Simulation time (in nanoseconds), number of strides (integers) and integration timestep (in femtoseconds):\n",
        "Stride_Time = \"1\" #@param {type:\"string\"}\n",
        "stride_time_prod = Stride_Time\n",
        "Number_of_strides = \"1\" #@param {type:\"string\"}\n",
        "nstride = Number_of_strides\n",
        "Integration_timestep = \"2\" #@param [\"0.5\", \"1\", \"2\", \"3\", \"4\"]\n",
        "dt_prod = Integration_timestep\n",
        "\n",
        "#@markdown Temperature (in Kelvin) and Pressure (in bar)\n",
        "Temperature = 298 #@param {type:\"string\"}\n",
        "temperature_prod = Temperature\n",
        "Pressure = 1 #@param {type:\"string\"}\n",
        "pressure_prod = Pressure\n",
        "\n",
        "#@markdown Frequency to write the trajectory file (in picoseconds):\n",
        "Write_the_trajectory = \"10\" #@param [\"10\", \"100\", \"200\", \"500\", \"1000\"]\n",
        "write_the_trajectory_prod = Write_the_trajectory\n",
        "\n",
        "#@markdown Frequency to write the log file (in picoseconds):\n",
        "Write_the_log = \"10\" #@param [\"10\", \"100\", \"200\", \"500\", \"1000\"]\n",
        "write_the_log_prod = Write_the_log\n",
        "\n",
        "#@markdown ---"
      ],
      "metadata": {
        "cellView": "form",
        "id": "nwSCUV-OStuF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Runs a Production MD simulation (NPT ensemble) after equilibration**\n",
        "#\n",
        "###########################################\n",
        "from sys import stdout, exit, stderr\n",
        "import os, math, fnmatch\n",
        "\n",
        "#############################################\n",
        "# Defining MD simulation parameters\n",
        "\n",
        "jobname = os.path.join(workDir, str(Jobname))\n",
        "pdbfile = os.path.join(workDir, Equilibrated_PDB)\n",
        "equil_rst_file = os.path.join(workDir, State_file)\n",
        "\n",
        "\n",
        "stride_time_ps = float(stride_time_prod)*1000\n",
        "stride_time = float(stride_time_ps)*picosecond\n",
        "nstride = int(Number_of_strides)\n",
        "dt = int(dt_prod)*femtosecond\n",
        "temperature = float(temperature_prod)*kelvin\n",
        "savcrd_freq = int(write_the_trajectory_prod)*picosecond\n",
        "print_freq  = int(write_the_log_prod)*picosecond\n",
        "\n",
        "pressure\t= float(pressure_prod)*bar\n",
        "\n",
        "simulation_time = stride_time*nstride\n",
        "nsteps  = int(stride_time.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nprint  = int(print_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nsavcrd = int(savcrd_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "firststride = 1 # must be integer\n",
        "#############################################\n",
        "# Defining functions to use below:\n",
        "def backup_old_log(pattern, string):\n",
        "\tresult = []\n",
        "\tfor root, dirs, files in os.walk(\"./\"):\n",
        "\t\tfor name in files:\n",
        "\t\t\tif fnmatch.fnmatch(name, pattern):\n",
        "\n",
        "\t\t\t\ttry:\n",
        "\t\t\t\t\tnumber = int(name[-2])\n",
        "\t\t\t\t\tavail = isinstance(number, int)\n",
        "\t\t\t\t\t#print(name,avail)\n",
        "\t\t\t\t\tif avail == True:\n",
        "\t\t\t\t\t\tresult.append(number)\n",
        "\t\t\t\texcept:\n",
        "\t\t\t\t\tpass\n",
        "\n",
        "\tif len(result) > 0:\n",
        "\t\tmaxnumber = max(result)\n",
        "\telse:\n",
        "\t\tmaxnumber = 0\n",
        "\n",
        "\tbackup_file = \"\\#\" + string + \".\" + str(maxnumber + 1) + \"#\"\n",
        "\tos.system(\"mv \" + string + \" \" + backup_file)\n",
        "\treturn backup_file\n",
        "##############################################\n",
        "\n",
        "#############################################\n",
        "print(\"\\n> Simulation details:\\n\")\n",
        "print(\"\\tJob name = \" + jobname)\n",
        "print(\"\\tPDB file = \" + str(pdbfile))\n",
        "\n",
        "print(\"\\n\\tSimulation_time = \" + str(stride_time*nstride))\n",
        "print(\"\\tIntegration timestep = \" + str(dt))\n",
        "print(\"\\tTotal number of steps = \" +  str(nsteps*nstride))\n",
        "print(\"\\tNumber of strides = \" + str(nstride) + \" (\" + str(stride_time) + \" in each stride)\")\n",
        "\n",
        "print(\"\\n\\tSave coordinates each \" + str(savcrd_freq))\n",
        "print(\"\\tSave checkpoint each \" + str(savcrd_freq))\n",
        "print(\"\\tPrint in log file each \" + str(print_freq))\n",
        "\n",
        "print(\"\\n\\tTemperature = \" + str(temperature))\n",
        "print(\"\\tPressure = \" + str(pressure))\n",
        "#############################################\n",
        "\n",
        "print(\"\\n> Setting the system:\\n\")\n",
        "\n",
        "\n",
        "print(\"\\t- Reading topology and structure file...\")\n",
        "prmtop = pmd.load_file(amber_prmtop)\n",
        "inpcrd = AmberInpcrdFile(amber_inpcrd)\n",
        "\n",
        "print(\"\\t- Creating system and setting parameters...\")\n",
        "nonbondedMethod = PME\n",
        "nonbondedCutoff = 1.0*nanometers\n",
        "ewaldErrorTolerance = 0.0005\n",
        "constraints = HBonds\n",
        "rigidWater = True\n",
        "constraintTolerance = 0.000001\n",
        "friction = 1.0\n",
        "system = omm_forcefield.createSystem(modeller.topology,nonbondedMethod=nonbondedMethod, nonbondedCutoff=nonbondedCutoff,\n",
        "                                          constraints=constraints, rigidWater=rigidWater, ewaldErrorTolerance=ewaldErrorTolerance)\n",
        "\n",
        "print(\"\\t- Setting barostat...\")\n",
        "system.addForce(MonteCarloBarostat(pressure, temperature))\n",
        "\n",
        "print(\"\\t- Setting integrator...\")\n",
        "integrator = LangevinIntegrator(temperature, friction, dt)\n",
        "integrator.setConstraintTolerance(constraintTolerance)\n",
        "simulation = Simulation(topology, system, integrator)\n",
        "simulation.context.setPositions(inpcrd.positions)\n",
        "if inpcrd.boxVectors is not None:\n",
        "\tsimulation.context.setPeriodicBoxVectors(*inpcrd.boxVectors)\n",
        "\n",
        "#############################################\n",
        "# Opening a loop of extension NSTRIDE to simulate the entire STRIDE_TIME*NSTRIDE\n",
        "for n in range(1, nstride + 1):\n",
        "\n",
        "\tprint(\"\\n\\n>>> Simulating Stride #\" + str(n) + \" <<<\")\n",
        "\n",
        "\tdcd_file = jobname + \"_\" + str(n) + \".dcd\"\n",
        "\tlog_file = jobname + \"_\" + str(n) + \".log\"\n",
        "\trst_file = jobname + \"_\" + str(n) + \".rst\"\n",
        "\tprv_rst_file = jobname + \"_\" + str(n-1) + \".rst\"\n",
        "\tpdb_file = jobname + \"_\" + str(n) + \".pdb\"\n",
        "\n",
        "\tif os.path.exists(rst_file):\n",
        "\t\tprint(\"> Stride #\" + str(n) + \" finished (\" + rst_file + \" present). Moving to next stride... <\")\n",
        "\t\tcontinue\n",
        "\n",
        "\tif n == 1:\n",
        "\t\tprint(\"\\n> Loading previous state from equilibration > \" + equil_rst_file + \" <\")\n",
        "\t\twith open(equil_rst_file, 'r') as f:\n",
        "\t\t\tsimulation.context.setState(XmlSerializer.deserialize(f.read()))\n",
        "\t\t\tcurrstep = int((n-1)*nsteps)\n",
        "\t\t\tcurrtime = currstep*dt.in_units_of(picosecond)\n",
        "\t\t\tsimulation.currentStep = currstep\n",
        "\t\t\tsimulation.context.setTime(currtime)\n",
        "\t\t\tprint(\"> Current time: \" + str(currtime) + \" (Step = \" + str(currstep) + \")\")\n",
        "\n",
        "\telse:\n",
        "\t\tprint(\"> Loading previous state from > \" + prv_rst_file + \" <\")\n",
        "\t\twith open(prv_rst_file, 'r') as f:\n",
        "\t\t\tsimulation.context.setState(XmlSerializer.deserialize(f.read()))\n",
        "\t\t\tcurrstep = int((n-1)*nsteps)\n",
        "\t\t\tcurrtime = currstep*dt.in_units_of(picosecond)\n",
        "\t\t\tsimulation.currentStep = currstep\n",
        "\t\t\tsimulation.context.setTime(currtime)\n",
        "\t\t\tprint(\"> Current time: \" + str(currtime) + \" (Step = \" + str(currstep) + \")\")\n",
        "\n",
        "\n",
        "\tdcd = DCDReporter(dcd_file, nsavcrd)\n",
        "\tfirstdcdstep = (currstep) + nsavcrd\n",
        "\tdcd._dcd = DCDFile(dcd._out, topology, simulation.integrator.getStepSize(), firstdcdstep, nsavcrd) # first step should not be 0\n",
        "\n",
        "\tsimulation.reporters.append(dcd)\n",
        "\tsimulation.reporters.append(StateDataReporter(stdout, nprint, step=True, speed=True, progress=True, totalSteps=(nsteps*nstride), remainingTime=True, separator='\\t\\t'))\n",
        "\tsimulation.reporters.append(StateDataReporter(log_file, nprint, step=True, kineticEnergy=True, potentialEnergy=True, totalEnergy=True, temperature=True, volume=True, speed=True))\n",
        "\n",
        "\tprint(\"\\n> Simulating \" + str(nsteps) + \" steps... (Stride #\" + str(n) + \")\")\n",
        "\tsimulation.step(nsteps)\n",
        "\n",
        "\tsimulation.reporters.clear() # remove all reporters so the next iteration don't trigger them.\n",
        "\n",
        "\n",
        "\t##################################\n",
        "\t# Writing last frame information of stride\n",
        "\tprint(\"\\n> Writing state file (\" + str(rst_file) + \")...\")\n",
        "\tstate = simulation.context.getState( getPositions=True, getVelocities=True )\n",
        "\twith open(rst_file, 'w') as f:\n",
        "\t\tf.write(XmlSerializer.serialize(state))\n",
        "\n",
        "\tlast_frame = int(nsteps/nsavcrd)\n",
        "\tprint(\"> Writing coordinate file (\" + str(pdb_file) + \", frame = \" + str(last_frame) + \")...\")\n",
        "\tpositions = simulation.context.getState(getPositions=True).getPositions()\n",
        "\tPDBFile.writeFile(simulation.topology, positions, open(pdb_file, 'w'))\n",
        "\n",
        "print(\"\\n> Finished!\\n\")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "XOLGNAnNTBUx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Concatenate and align the trajectory**\n",
        "import MDAnalysis as mda\n",
        "from MDAnalysis.analysis import align, rms\n",
        "\n",
        "Google_Drive_Path = '/content/' #@param {type:\"string\"}\n",
        "workDir = Google_Drive_Path\n",
        "Equilibrated_PDB = 'prot_lig_equil.pdb' #@param {type:\"string\"}\n",
        "Jobname = \"prot_lig_prod\" #@param {type: \"string\"}\n",
        "Skip = \"1\" #@param [\"1\", \"2\", \"5\", \"10\", \"20\", \"50\"]\n",
        "stride_traj = Skip\n",
        "Output_format = \"dcd\" #@param [\"dcd\", \"pdb\", \"trr\", \"xtc\"]\n",
        "first_stride = \"1\" #@param {type:\"string\"}\n",
        "Number_of_strides = \"1\" #@param {type:\"string\"}\n",
        "nstride = int(Number_of_strides)\n",
        "stride_time = \"1\" #@param {type:\"string\"}\n",
        "trajectory_saved_frequency = \"10\" #@param [\"10\", \"100\", \"200\", \"500\", \"1000\"]\n",
        "traj_save_freq = trajectory_saved_frequency\n",
        "Remove_waters = \"yes\" #@param [\"yes\", \"no\"]\n",
        "# stride_id_as_ref_for_alignment = \"1\" #@param {type: \"string\"}\n",
        "output_prefix = first_stride+\"-\"+str(int(first_stride)+nstride-1)\n",
        "\n",
        "stride_time_ps = float(stride_time)*1000\n",
        "simulation_time_analysis = stride_time_ps*nstride\n",
        "simulation_ns = float(stride_time)*int(Number_of_strides)\n",
        "number_frames = int(simulation_time_analysis)/int(traj_save_freq)\n",
        "number_frames_analysis = number_frames/int(stride_traj)\n",
        "\n",
        "\n",
        "# traj_end = os.path.join(workDir, str(Jobname) + \"_all.dcd\")\n",
        "nw_dcd = os.path.join(workDir, str(Jobname) + output_prefix + \"_nw.\" + str(Output_format))\n",
        "nw_pdb = os.path.join(workDir, str(Jobname) +  \"_nw.pdb\")\n",
        "whole_pdb = os.path.join(workDir, str(Jobname) +  \"_whole.pdb\")\n",
        "# traj_end2 = os.path.join(workDir, str(Jobname) + \"_all.\" + str(Output_format))\n",
        "whole_dcd = os.path.join(workDir, str(Jobname) + output_prefix + \"_whole.\" + str(Output_format))\n",
        "template =  os.path.join(workDir, str(Jobname) + '_%s.dcd')\n",
        "pdb = os.path.join(workDir, Equilibrated_PDB)\n",
        "\n",
        "flist = [template % str(i) for i in range(int(first_stride), int(first_stride) + nstride)]\n",
        "ref = [template % int(1)]\n",
        "\n",
        "pprm = pmd.load_file(amber_prmtop, amber_inpcrd)\n",
        "mprm_from_parmed = mda.Universe(pprm)\n",
        "u = mprm_from_parmed.select_atoms('not (resname HOH)')\n",
        "top_nw = u.convert_to('PARMED')\n",
        "\n",
        "u1 = mda.Universe(pdb, flist)\n",
        "u2 = mda.Universe(pdb, ref)\n",
        "\n",
        "u2.trajectory[0] # set u2 to first frame\n",
        "\n",
        "# print(rms.rmsd(u1.select_atoms('name CA').positions, u2.select_atoms('name CA').positions, superposition=False))\n",
        "\n",
        "align.AlignTraj(u1, u2, select='name CA', in_memory=True).run()\n",
        "\n",
        "nw = u1.select_atoms(\"not (resname HOH)\")\n",
        "if Remove_waters == \"yes\":\n",
        "  with mda.Writer(nw_dcd, nw.n_atoms) as W:\n",
        "    for ts in u1.trajectory[::int(Skip)]:\n",
        "        W.write(nw, )\n",
        "  not_waters = u2.select_atoms(\"not (resname HOH)\")\n",
        "  not_waters.write(nw_pdb)\n",
        "  traj_dcd_check = os.path.exists(nw_dcd)\n",
        "  topology_mda = top_nw\n",
        "  traj = nw_dcd\n",
        "  pdb_ref = nw_pdb\n",
        "else:\n",
        "  with mda.Writer(whole_dcd, u1.select_atoms(\"all\").n_atoms) as W:\n",
        "    for ts in u1.trajectory[::int(Skip)]:\n",
        "        W.write(u1.select_atoms(\"all\"))\n",
        "  whole = u2.select_atoms(\"all\")\n",
        "  whole.write(whole_pdb)\n",
        "  traj_dcd_check = os.path.exists(whole_dcd)\n",
        "  topology_mda = pprm\n",
        "  traj = whole_dcd\n",
        "  pdb_ref = whole_pdb\n",
        "\n",
        "traj_load = pt.load(traj, pdb_ref)\n",
        "print(traj_load)\n",
        "\n",
        "if traj_dcd_check == True:\n",
        "  print(\"Trajectory concatenated successfully! :-)\")\n",
        "else:\n",
        "  print(\"ERROR: Check your inputs! \")"
      ],
      "metadata": {
        "id": "3NMA7K6OSjmH",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x_awbwdZnp6Q",
        "cellView": "form"
      },
      "source": [
        "#@title **Load, view and check the trajectory**\n",
        "#@markdown This will take a few minutes. Another coffee would be great. :-)\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "# !rm *.pdb 2> /dev/null\n",
        "\n",
        "#py3dmol functions\n",
        "class Atom(dict):\n",
        "  def __init__(self, line):\n",
        "    self[\"type\"] = line[0:6].strip()\n",
        "    self[\"idx\"] = line[6:11].strip()\n",
        "    self[\"name\"] = line[12:16].strip()\n",
        "    self[\"resname\"] = line[17:20].strip()\n",
        "    self[\"resid\"] = int(int(line[22:26]))\n",
        "    self[\"x\"] = float(line[30:38])\n",
        "    self[\"y\"] = float(line[38:46])\n",
        "    self[\"z\"] = float(line[46:54])\n",
        "    self[\"sym\"] = line[76:78].strip()\n",
        "\n",
        "  def __str__(self):\n",
        "    line = list(\" \" * 80)\n",
        "    line[0:6] = self[\"type\"].ljust(6)\n",
        "    line[6:11] = self[\"idx\"].ljust(5)\n",
        "    line[12:16] = self[\"name\"].ljust(4)\n",
        "    line[17:20] = self[\"resname\"].ljust(3)\n",
        "    line[22:26] = str(self[\"resid\"]).ljust(4)\n",
        "    line[30:38] = str(self[\"x\"]).rjust(8)\n",
        "    line[38:46] = str(self[\"y\"]).rjust(8)\n",
        "    line[46:54] = str(self[\"z\"]).rjust(8)\n",
        "    line[76:78] = self[\"sym\"].rjust(2)\n",
        "    return \"\".join(line) + \"\\n\"\n",
        "\n",
        "class Molecule(list):\n",
        "  def __init__(self, file):\n",
        "    for line in file:\n",
        "      if \"ATOM\" in line or \"HETATM\" in line:\n",
        "        self.append(Atom(line))\n",
        "\n",
        "    def __str__(self):\n",
        "      outstr = \"\"\n",
        "      for at in self:\n",
        "        outstr += str(at)\n",
        "      return outstr\n",
        "\n",
        "class Molecule2(list):\n",
        "    def __init__(self, file):\n",
        "        for line in file:\n",
        "            if \"HETATM\" in line and not \"HOH\" in line and not \"CL\" in line and not \"NA\"in line and not \"K\" in line:\n",
        "                self.append(Atom(line))\n",
        "\n",
        "    def __str__(self):\n",
        "        outstr = \"\"\n",
        "        for at in self:\n",
        "            outstr += str(at)\n",
        "\n",
        "        return outstr\n",
        "\n",
        "if number_frames_analysis > 10:\n",
        "  stride_animation = number_frames_analysis/10\n",
        "else:\n",
        "  stride_animation = 1\n",
        "\n",
        "if Remove_waters == \"yes\":\n",
        "  traj = nw_dcd\n",
        "  pdb_ref = nw_pdb\n",
        "else:\n",
        "  traj = whole_dcd\n",
        "  pdb_ref = whole_pdb\n",
        "\n",
        "u = mda.Universe(pdb_ref, traj)\n",
        "\n",
        "# Write out frames for animation\n",
        "protein = u.select_atoms('not (resname HOH)')\n",
        "i = 0\n",
        "for ts in u.trajectory[0:len(u.trajectory):int(stride_animation)]:\n",
        "    if i > -1:\n",
        "        with mda.Writer('' + str(i) + '.pdb', protein.n_atoms) as W:\n",
        "            W.write(protein)\n",
        "    i = i + 1\n",
        "# Load frames as molecules\n",
        "molecules = []\n",
        "for i in range(int(len(u.trajectory)/int(stride_animation))):\n",
        "    with open('' + str(i) + '.pdb') as ifile:\n",
        "        molecules.append(Molecule(ifile))\n",
        "\n",
        "models = \"\"\n",
        "for i in range(len(molecules)):\n",
        "  models += \"MODEL \" + str(i) + \"\\n\"\n",
        "  for j,mol in enumerate(molecules[i]):\n",
        "    models += str(mol)\n",
        "  models += \"ENDMDL\\n\"\n",
        "#view.addModelsAsFrames(models)\n",
        "\n",
        "molecules2 = []\n",
        "for i in range(int(len(u.trajectory)/int(stride_animation))):\n",
        "    with open('' + str(i) + '.pdb') as ifile:\n",
        "        molecules2.append(Molecule2(ifile))\n",
        "\n",
        "models2 = \"\"\n",
        "for i in range(len(molecules2)):\n",
        "  models2 += \"MODEL \" + str(i) + \"\\n\"\n",
        "  for j,mol in enumerate(molecules2[i]):\n",
        "    models2 += str(mol)\n",
        "  models2 += \"ENDMDL\\n\"\n",
        "\n",
        "# Animation\n",
        "view = py3Dmol.view(width=800, height=600)\n",
        "view.addModelsAsFrames(models)\n",
        "for i, at in enumerate(molecules[0]):\n",
        "    default = {\"cartoon\": {'color': 'spectrum'}}\n",
        "    view.setViewStyle({'style':'outline','color':'black','width':0.1})\n",
        "    view.setStyle({'model': -1, 'serial': i+1}, at.get(\"pymol\", default))\n",
        "    HP = ['UNK']\n",
        "    view.setStyle({\"model\":-1,'and':[{'resn':HP}]},{'stick':{'radius':0.3}})\n",
        "view.zoomTo()\n",
        "view.animate({'loop': \"forward\"})\n",
        "view.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hRxiVlQt3tC5",
        "cellView": "form"
      },
      "source": [
        "#@title **View and check the Ligand Interaction Network (LigPlot) during MD simulations**\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 (output.html).\n",
        "\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = 'Interaction' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown The frequency with which an interaction is seen will control the width of the corresponding edge. You can hide the least frequent interactions by using a threshold, i.e. threshold=0.3 will hide interactions that occur in less than 30% of frames.\n",
        "Threshold = 0.3 #@param {type:\"slider\", min:0, max:1.0, step:0.1}\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",
        "if Remove_waters == \"yes\":\n",
        "  topology_mda = top_nw\n",
        "else:\n",
        "    topology_mda = pprm\n",
        "\n",
        "# load topology\n",
        "u = mda.Universe(topology_mda, traj)\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",
        "if number_frames_analysis > 10:\n",
        "  stride_animation = number_frames_analysis/10\n",
        "else:\n",
        "  stride_animation = 1\n",
        "\n",
        "fp = plf.Fingerprint()\n",
        "fp.run(u.trajectory[::int(stride_animation)], 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\", threshold=float(Threshold),\n",
        "                          rotation=270)\n",
        "net.save(os.path.join(workDir, Output_name + \".html\"))\n",
        "net.display()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Emh0vU5UjgB6"
      },
      "source": [
        "---\n",
        "---\n",
        "# **Analysis**\n",
        "\n",
        "Although visualizing your trajectory can be quite useful, sometimes you also want more quantitative data.\n",
        "\n",
        "Analyses of MD trajectories vary a lot and we do not intend to cover it all here. However, one can make use of MDanalysis or PyTraj to easily analyze simulations.\n",
        "\n",
        "Below, you can find a few examples of code snippets that can help you to shed some light on your simulation behavior."
      ]
    },
    {
      "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 = 1 #@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",
        "output = workDir\n",
        "\n",
        "u1 = mda.Universe(pdb_ref, traj)\n",
        "u2 = mda.Universe(pdb_ref, 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 MD simulations:**\n",
        "\n",
        "Frame = \"20\" #@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 **Compute RMSD of the ligand**\n",
        "#@markdown **Provide output file names below:**\n",
        "import MDAnalysis.analysis.rms\n",
        "Output_name = 'rmsd_lig' #@param {type:\"string\"}\n",
        "\n",
        "pdb_ref, traj\n",
        "u = mda.Universe(pdb_ref, traj)\n",
        "ref = mda.Universe(pdb_ref, pdb_ref)\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",
        "\n",
        "time = len(rmsd[0])*int(Write_the_trajectory)/1000\n",
        "time_array = np.arange(0,time,int(Write_the_trajectory)/1000)*int(stride_traj)\n",
        "\n",
        "ax = plt.plot(time_array, rmsd[2], alpha=0.6, color = 'blue', linewidth = 1.0)\n",
        "plt.xlim(0, time_array[-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": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kK8YRhrjD6xo",
        "cellView": "form"
      },
      "source": [
        "#@title **Compute distance between the ligand and catalytic site residues**\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = 'distance' #@param {type:\"string\"}\n",
        "#@markdown **Cutoff distance to nearest residues (Angstrons):**\n",
        "Distance = '4' #@param {type:\"string\"}\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",
        "selection = '!(:HOH) & !(:NA) & !(:CL) & !(:K) & !(:' + str(res_number[0]) + '-' + str(res_number[-1]) + ')'\n",
        "selection2 = ':' + str(res_number[0]) + '-' + str(res_number[-1])\n",
        "pt_topology = traj_load.top\n",
        "restraint_array = pt.select_atoms(selection, pt_topology)\n",
        "restraint_array2 = pt.select_atoms(selection2, pt_topology)\n",
        "lig_first = restraint_array[0]\n",
        "lig_last = restraint_array[-1]\n",
        "first_atom = restraint_array2[0]\n",
        "last_atom = restraint_array2[-1]\n",
        "\n",
        "ini = 0\n",
        "top = pt_topology\n",
        "for frame in traj_load:\n",
        "  top.set_reference(traj_load[ini])\n",
        "  indices = traj_load.top.select('(@' + str(lig_first+1)  + \"-\" + str(lig_last+1) + '<:' + str(Distance) + ')&!(:HOH|:NA,CL,K|@' + str(lig_first+1)  + \"-\" + str(lig_last+1) + ')')\n",
        "  residues = [res.original_resid for res in top[indices].residues]\n",
        "res_string = ','.join(str(e) for e in residues)\n",
        "print(\"Selected residues = \" + res_string + \"\\n\")\n",
        "\n",
        "mask = '@' + str(lig_first+1)  + \"-\" + str(lig_last+1) + \" :\" + str(res_string)\n",
        "dist = pt.distance(traj_load, mask)\n",
        "\n",
        "dist_mean = mean(dist)\n",
        "dist_stdev = stdev(dist)\n",
        "print(\"Distance Average = \" + str(\"{:.2f}\".format(dist_mean)) + \" \\u00B1 \" + str(\"{:.2f}\".format(dist_stdev)) + \" Å\")\n",
        "\n",
        "time = len(dist)*int(Write_the_trajectory)/1000\n",
        "time_array = np.arange(0,time,int(Write_the_trajectory)/1000)*int(stride_traj)\n",
        "\n",
        "# Plotting:\n",
        "ax = plt.plot(time_array, dist, alpha=1, color = 'springgreen', linewidth = 1.0)\n",
        "plt.xlim(0, simulation_ns)\n",
        "#plt.ylim(2, 6)\n",
        "\n",
        "plt.xlabel(\"Time (ns)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel(\"Distance [$\\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(dist)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "mVOogip3h08Z"
      },
      "source": [
        "#@title **Compute distance between the ligand and specific residues**\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = 'distance_select' #@param {type:\"string\"}\n",
        "#@markdown **Type the number of residues separated by commas and without spaces (1,2,3...):**\n",
        "Residues = '78,84' #@param {type:\"string\"}\n",
        "\n",
        "mask = '@' + str(lig_first+1)  + \"-\" + str(lig_last+1) + \" :\" + str(Residues)\n",
        "dist = pt.distance(traj_load, mask)\n",
        "print(\"Selected residues = \" + Residues + \"\\n\")\n",
        "dist_mean = mean(dist)\n",
        "dist_stdev = stdev(dist)\n",
        "print(\"Distance Average = \" + str(\"{:.2f}\".format(dist_mean)) + \" \\u00B1 \" + str(\"{:.2f}\".format(dist_stdev)) + \" Å\")\n",
        "\n",
        "time = len(dist)*int(Write_the_trajectory)/1000\n",
        "time_array = np.arange(0,time,int(Write_the_trajectory)/1000)*int(stride_traj)\n",
        "\n",
        "# Plotting:\n",
        "ax = plt.plot(time_array, dist, alpha=1, color = 'magenta', linewidth = 1.0)\n",
        "plt.xlim(0, simulation_ns)\n",
        "#plt.ylim(2, 6)\n",
        "\n",
        "plt.xlabel(\"Time (ns)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel(\"Distance [$\\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(dist)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wBrBMF4Puyv6",
        "cellView": "form"
      },
      "source": [
        "#@title **Compute RMSD of protein's CA atoms**\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = 'rmsd_ca' #@param {type:\"string\"}\n",
        "\n",
        "rmsd = pt.rmsd(traj_load, ref = 0, mask = \"@CA\")\n",
        "\n",
        "time = len(rmsd)*int(Write_the_trajectory)/1000\n",
        "time_array = np.arange(0,time,int(Write_the_trajectory)/1000)*int(stride_traj)\n",
        "\n",
        "# Plotting:\n",
        "ax = plt.plot(time_array, rmsd, alpha=0.6, color = 'blue', linewidth = 1.0)\n",
        "plt.xlim(0, simulation_ns)\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)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "MDGDypFxVHSp"
      },
      "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, 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": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CvOFrXGXwXrV",
        "cellView": "form"
      },
      "source": [
        "#@title **Compute radius of gyration of protein's CA atoms**\n",
        "\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = 'radius_gyration' #@param {type:\"string\"}\n",
        "\n",
        "radgyr = pt.radgyr(traj_load, mask = \"@CA\")\n",
        "time = len(rmsd)*int(Write_the_trajectory)/1000\n",
        "time_array = np.arange(0,time,int(Write_the_trajectory)/1000)*int(stride_traj)\n",
        "\n",
        "# Plotting:\n",
        "plt.plot(time_array, radgyr, alpha=0.6, color = 'green', linewidth = 1.0)\n",
        "plt.xlim(0, simulation_ns)\n",
        "#plt.ylim(2, 6)\n",
        "\n",
        "plt.xlabel(\"Time (ns)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel(\"Radius of gyration ($\\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(radgyr)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2Q7FKg8Fuxr9",
        "cellView": "form"
      },
      "source": [
        "#@title **Plot radius of gyration as a ditribution**\n",
        "\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = 'radius_gyration_dist' #@param {type:\"string\"}\n",
        "\n",
        "ax = sb.kdeplot(radgyr, color=\"green\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('Radius of gyration ($\\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": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p2Y0DgwTxLWc",
        "cellView": "form"
      },
      "source": [
        "#@title **Compute RMSF of protein's CA atoms**\n",
        "\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = 'rmsf_ca' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "rmsf = pt.rmsf(traj_load, \"@CA\")\n",
        "bfactor = pt.bfactors(traj_load, byres=True)\n",
        "\n",
        "# Plotting:\n",
        "plt.plot(rmsf[:,1], alpha=1.0, color = 'red', linewidth = 1.0)\n",
        "\n",
        "plt.xlabel(\"Residue\", fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel(\"RMSF ($\\AA$)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.xlim(0, len(rmsf[:-1]))\n",
        "\n",
        "#plt.xticks(np.arange(min(rmsf[:1]), max(rmsf[:1])))\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(rmsf)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JalicqqrTodW",
        "cellView": "form"
      },
      "source": [
        "#@title **2D RMSD**\n",
        "\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = '2D_rmsd' #@param {type:\"string\"}\n",
        "\n",
        "last_frame = len(time_array)\n",
        "\n",
        "stride_ticks_f = (last_frame)/5\n",
        "ticks_frame = np.arange(0,(len(time_array) + float(stride_ticks_f)), float(stride_ticks_f))\n",
        "a = ticks_frame.astype(float)\n",
        "stride_ticks_t = (simulation_ns)/5\n",
        "tick_time = np.arange(0,(float(simulation_ns) + float(stride_ticks_t)), float(stride_ticks_t))\n",
        "b = tick_time.astype(float)\n",
        "\n",
        "mat1 = pt.pairwise_rmsd(traj_load, mask=\"@CA\", frame_indices=range(int(number_frames_analysis)))\n",
        "\n",
        "\n",
        "ax = plt.imshow(mat1, cmap = 'PRGn', origin='lower', interpolation = 'bicubic')\n",
        "plt.title('2D RMSD')\n",
        "plt.xlabel('Time (ns)', fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel('Time (ns)', fontsize = 14, fontweight = 'bold')\n",
        "# plt.xticks(fontsize = 12)\n",
        "# plt.yticks(fontsize = 12)\n",
        "plt.xticks(a, b.round(decimals=3), fontsize = 12)\n",
        "plt.yticks(a, b.round(decimals=3), fontsize = 12)\n",
        "# plt.xlim(0, a[-1])\n",
        "# plt.ylim(0, a[-1])\n",
        "\n",
        "cbar1 = plt.colorbar()\n",
        "cbar1.set_label(\"RMSD ($\\AA$)\", fontsize = 14, fontweight = 'bold')\n",
        "\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "raw_data=pd.DataFrame(mat1)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_mgVSbBshWFV",
        "cellView": "form"
      },
      "source": [
        "#@title **Calculate eigvenctors of Principle Component Analysis (PCA)**\n",
        "data = pt.pca(traj_load, fit=True, ref=0, mask='@CA', n_vecs=2)\n",
        "#print('projection values of each frame to first mode = {} \\n'.format(data[0][0]))\n",
        "#print('projection values of each frame to second mode = {} \\n'.format(data[0][1]))\n",
        "#print('eigvenvalues of first two modes', data[1][0])\n",
        "#print(\"\")\n",
        "#print('eigvenvectors of first two modes: \\n', data[1][1])\n",
        "\n",
        "last_frame = len(time_array)\n",
        "\n",
        "stride_ticks_f = (last_frame)/5\n",
        "ticks_frame = np.arange(0,(len(time_array) + float(stride_ticks_f)), float(stride_ticks_f))\n",
        "a = ticks_frame.astype(float)\n",
        "a2 = a.tolist()\n",
        "stride_ticks_t = (simulation_ns)/5\n",
        "tick_time = np.arange(0,(float(simulation_ns) + float(stride_ticks_t)), float(stride_ticks_t))\n",
        "b = tick_time.astype(float)\n",
        "\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = 'PCA' #@param {type:\"string\"}\n",
        "\n",
        "Output_PC1 = 'PC1' #@param {type:\"string\"}\n",
        "Output_PC2 = 'PC2' #@param {type:\"string\"}\n",
        "\n",
        "%matplotlib inline\n",
        "%config InlineBackend.figure_format = 'retina'  # high resolution\n",
        "projection_data = data[0]\n",
        "plt.title(r'PCA of C-$\\alpha$')\n",
        "PC1 = data[0][0]\n",
        "PC2 = data[0][1]\n",
        "\n",
        "a = plt.scatter(PC1,PC2, c=range(int(number_frames_analysis)), cmap='Greens', marker='o',s=8, alpha=1)\n",
        "plt.clim(0, last_frame)\n",
        "\n",
        "plt.xlabel('PC1', fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel('PC2', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "# N = len(number_frames)\n",
        "# x2 = np.arange(N)\n",
        "\n",
        "cbar1 = plt.colorbar(a, orientation=\"vertical\")\n",
        "cbar1.set_label('Time(ns)', fontsize = 14, fontweight = 'bold')\n",
        "cbar1.set_ticks(a2)\n",
        "cbar1.set_ticklabels(b.round(decimals=3))\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "pc1=pd.DataFrame(PC1)\n",
        "pc1.to_csv(os.path.join(workDir, Output_PC1 + \".csv\"))\n",
        "pc2=pd.DataFrame(PC2)\n",
        "pc2.to_csv(os.path.join(workDir, Output_PC2 + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yce9RfNtpl-J",
        "cellView": "form"
      },
      "source": [
        "#@title **Plot Principal Component 1 (PC1) and Principal Component 2 (PC2) as a ditribution**\n",
        "Output_name = 'PCA_dist' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "fig = plt.figure(figsize=(9,5))\n",
        "\n",
        "plt.subplot(1, 2, 1)\n",
        "ax = sb.kdeplot(PC1, color=\"green\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('PC1', 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",
        "plt.subplot(1, 2, 2)\n",
        "ax2 = sb.kdeplot(PC2, color=\"purple\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('PC2', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks([])\n",
        "plt.ylabel('')\n",
        "ax2.spines['top'].set_visible(False)\n",
        "ax2.spines['right'].set_visible(False)\n",
        "ax2.spines['bottom'].set_visible(True)\n",
        "ax2.spines['left'].set_visible(False)\n",
        "\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pTDb7CEfkLq1",
        "cellView": "form"
      },
      "source": [
        "#@title **Pearson's Cross Correlation (CC)**\n",
        "\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = 'cross_correlation' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "traj_align = pt.align(traj_load, mask='@CA', ref=0)\n",
        "\n",
        "mat_cc = matrix.correl(traj_align, '@CA')\n",
        "\n",
        "ax = plt.imshow(mat_cc, cmap = 'PiYG_r', interpolation = 'bicubic', vmin = -1, vmax = 1, origin='lower')\n",
        "\n",
        "plt.xlabel('Residues', fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel('Residues', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "cbar1 = plt.colorbar()\n",
        "cbar1.set_label('$CC_ij$', fontsize = 14, fontweight = 'bold')\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "raw_data=pd.DataFrame(mat_cc)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}