EquiBind+MD.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "EquiBind+MD.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyONBq/2niW4Saek/znwcjs1",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "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/EquiBind%2BMD.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Hi there!**\n",
        "\n",
        "This is a Jupyter notebook to perform direct-shot prediction of both the receptor binding location (blind docking) and the ligand’s bound pose and orientation using [EquiBind](https://github.com/HannesStark/EquiBind) (geometric deep learning model) and for running Molecular Dynamics (MD) simulations using OpenMM engine.\n",
        "\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",
        " **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",
        "**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 [EquiBind](https://github.com/HannesStark/EquiBind) and [GNINA](https://github.com/gnina/gnina) 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",
        "- 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)"
      ],
      "metadata": {
        "id": "n-Xw3hPplMsi"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "HBqsph65FmvO"
      },
      "outputs": [],
      "source": [
        "#@title **Install dependencies**\n",
        "#@markdown It will take a few minutes (around 12 min), please, drink a coffee and wait. ;-)\n",
        "# install dependencies\n",
        "%%capture\n",
        "import sys\n",
        "!pip -q install py3Dmol\n",
        "!pip install --upgrade MDAnalysis\n",
        "!pip install Cython\n",
        "\n",
        "# install conda\n",
        "!wget -qnc https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \n",
        "!bash Miniconda3-latest-Linux-x86_64.sh -bfp /usr/local 2>&1 1>/dev/null\n",
        "!rm -r Miniconda3-latest-Linux-x86_64.sh\n",
        "!conda install -y -q -c conda-forge openmm=7.6 python=3.7 pdbfixer 2>&1 1>/dev/null\n",
        "!conda install -c ambermd pytraj  --yes 2>&1 1>/dev/null\n",
        "!conda install -c conda-forge parmed  --yes 2>&1 1>/dev/null\n",
        "!conda install --yes -c conda-forge openmmforcefields 2>&1 1>/dev/null\n",
        "!conda install -c conda-forge openff-toolkit --yes 2>&1 1>/dev/null\n",
        "!conda install -c bioconda pybel --yes\n",
        "!conda install -c openbabel openbabel --yes\n",
        "!git clone https://github.com/HannesStark/EquiBind\n",
        "!pip install torch\n",
        "!pip install torchvision\n",
        "!pip install dgl-cu102\n",
        "!pip install bio\n",
        "!pip install biopandas\n",
        "!pip install joblib\n",
        "!pip install tensorboard\n",
        "!pip install POT\n",
        "!pip install rdkit-pypi\n",
        "!pip install pyaml\n",
        "!pip install torchani\n",
        "!pip install ase\n",
        "!pip install pyfiglet\n",
        "!git clone https://github.com/pablo-arantes/ProLIF.git\n",
        "prolif1 = \"cd /content/ProLIF\"\n",
        "prolif2 = \"sed -i 's/mdanalysis.*/mdanalysis==2.0.0/' setup.cfg\"\n",
        "prolif3 = \"pip install .\"\n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "\n",
        "with open('prolif.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(prolif1)\n",
        "    print(prolif2)\n",
        "    print(prolif3)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "\n",
        "!chmod 700 prolif.sh 2>&1 1>/dev/null\n",
        "!bash prolif.sh >/dev/null 2>&1\n",
        "!rm -r /content/ProLIF prolif.sh\n",
        "\n",
        "#load dependencies\n",
        "sys.path.append('/usr/local/lib/python3.7/site-packages/')\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 os\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\n",
        "import pyfiglet"
      ]
    },
    {
      "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.flush_and_unmount()\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",
      "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 ID 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.**"
      ],
      "metadata": {
        "id": "qntfzVRElUbb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from openmm.app.pdbfile import PDBFile\n",
        "#@title **Please, provide the necessary input files below for receptor**: \n",
        "#@markdown **Important:** Run the cell to prepare your receptor.\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",
        "\n",
        "Google_Drive_Path = '/content/MyDrive/' #@param {type:\"string\"}\n",
        "workDir = Google_Drive_Path\n",
        "\n",
        "if os.path.exists(os.path.join(workDir, \"name_residues_receptor.txt\")):\n",
        "  os.remove(os.path.join(workDir,\"name_residues_receptor.txt\"))\n",
        "else:\n",
        "  pass\n",
        "\n",
        "receptor = os.path.join(workDir, \"receptor.pdb\")\n",
        "ligand = os.path.join(workDir, \"ligand.sdf\")\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",
        "#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",
        "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(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_ligands(path)\n",
        "\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",
        "if os.path.exists(receptor):\n",
        "  print(\"Successfully generated the files! :-)\")\n",
        "else:\n",
        "  print(\"ERROR: Check your inputs! \")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "CM2kzZ-3Gfz_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Receptor Visualization**: \n",
        "#@markdown Now the protein has been sanitized, it would be recomended to visualize and check the protein.\n",
        "import py3Dmol\n",
        "\n",
        "view = py3Dmol.view()\n",
        "view.removeAllModels()\n",
        "view.setViewStyle({'style':'outline','color':'black','width':0.1})\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.zoomTo()\n",
        "view.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "YMRen7pTGnFD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import requests\n",
        "import rdkit\n",
        "from rdkit import Chem\n",
        "from rdkit.Chem import AllChem,Draw\n",
        "import IPython\n",
        "from IPython.display import Image\n",
        "import pybel\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.image as mpimg\n",
        "\n",
        "\n",
        "# from rdkit.Chem.Draw import IPythonConsole\n",
        "\n",
        "#@title **Please, provide the necessary input files for the ligand**: \n",
        "\n",
        "#@markdown Type the name of your ligand. If you don't know the exactly name, please, check at https://pubchem.ncbi.nlm.nih.gov/\n",
        "Name = '2-PROPYLPHENOL' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "def get_cids(text):\n",
        "    \"\"\"\n",
        "    Search pubchem and return best matching record\n",
        "    \"\"\"  \n",
        "    url = 'https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/{}/cids/TXT'.format(text)\n",
        "    response = requests.get(url)\n",
        "    cids = response.text.split()\n",
        "    if len(cids) == 0:\n",
        "        return None\n",
        "    else:\n",
        "        return cids\n",
        "def get_record(cid):\n",
        "    \"\"\"\n",
        "    Get pubchem record for a given cid and returns molecule as rdkit\n",
        "    \"\"\"\n",
        "    url = 'https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{}/record/SDF'.format(cid)\n",
        "    response = requests.get(url)\n",
        "    mol = Chem.MolFromMolBlock(response.text)\n",
        "    smi = Draw.MolToFile(mol, size=(600, 600), filename=str(Name) + '.png')\n",
        "    img = mpimg.imread(str(Name) + '.png')\n",
        "    plt.figure(figsize = (8,8))\n",
        "    imgplot = plt.imshow(img)\n",
        "    plt.axis('off')\n",
        "    plt.show()\n",
        "    hmol = Chem.AddHs(mol)\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",
        "    # mol= [m for m in pybel.readfile(filename=os.path.join(workDir, \"ligand.mol\"), format='mol')][0]\n",
        "    # out=pybel.Outputfile(filename=ligand,format='sdf',overwrite=True)\n",
        "    # out.write(mol)\n",
        "    # out.close()\n",
        "  # return mol\n",
        "\n",
        "def get_molecule(text, n_results=1):\n",
        "    \"\"\"\n",
        "    Search pubchem and return best matching record\n",
        "    \"\"\"  \n",
        "    cids = get_cids(text)\n",
        "    if cids is None:\n",
        "        return None\n",
        "    else:\n",
        "        if n_results == 1:\n",
        "            return get_record(cids[0])\n",
        "        else:\n",
        "            return [get_record(cid) for cid in cids[:n_results]]\n",
        "\n",
        "get_molecule(Name, n_results=1)"
      ],
      "metadata": {
        "id": "y-a05s1XGzqI",
        "cellView": "form"
      },
      "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",
        "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='sdf',overwrite=True)\n",
        "out.write(mol)\n",
        "out.close()\n",
        "\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": "bEYsHGvq7gmw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Ligand Visualization**: \n",
        "#@markdown Now the ligand has been sanitized, it would be recomended to visualize and check the ligand.\n",
        "\n",
        "view = py3Dmol.view()\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(ligand,'r').read(),format='mol2')\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": "cnY-Ouuxt6pL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Parameters for the EquiBind calculation:**\n",
        "Output_folder_name = 'final' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Explicit random seed:\n",
        "seed = \"0\" #@param {type:\"string\"}\n",
        "\n",
        "font = pyfiglet.figlet_format(\"EquiBind\")\n",
        "print(font)\n",
        "\n",
        "docking_output = os.path.join(workDir, Output_folder_name)\n",
        "\n",
        "# if os.path.exists(docking_output):\n",
        "#   os.remove(docking_output)\n",
        "# else:\n",
        "#   pass\n",
        "\n",
        "equibind = \"python3 /content/EquiBind/multiligand_inference.py -o \" + str(docking_output) + \" -r \" + str(receptor) + \" -l \" + str(ligand) + \" --seed \" + str(int(seed))\n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "with open('equibind.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(equibind)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "\n",
        "!chmod 700 equibind.sh 2>&1 1>/dev/null\n",
        "!bash equibind.sh\n",
        "\n",
        "import gzip\n",
        "v = py3Dmol.view()\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(os.path.join(workDir, Output_folder_name + \"/output.sdf\")).read())\n",
        "v.setStyle({'model':1},{'stick':{'colorscheme':'greenCarbon'}})\n",
        "v.zoomTo({'model':1})\n",
        "v.rotate(90)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "Xwdv_ZltqfUn"
      },
      "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",
        "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",
        "# Load a molecule from a SDF file\n",
        "ligand = Molecule.from_file(os.path.join(workDir, Output_folder_name + \"/output.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(), 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": {
        "cellView": "form",
        "id": "Kf88JFKUJYc_"
      },
      "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",
        "\n",
        "#@markdown Now the system has been built, it would be recomended to visualize and check the system.\n",
        "\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",
        "\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 \"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",
        "view = py3Dmol.view()\n",
        "view.removeAllModels()\n",
        "view.setViewStyle({'style':'outline','color':'black','width':0.1})\n",
        "\n",
        "view.addModel(open(system_pdb,'r').read(),format='pdb')\n",
        "Prot=view.getModel()\n",
        "Prot.setStyle({'cartoon':{'arrows':True, 'tubes':True, 'style':'oval', 'color':'white'}})\n",
        "view.addSurface(py3Dmol.SAS,{'opacity':0.6,'color':'white'})\n",
        "\n",
        "with open(system_pdb) as ifile:\n",
        "    mol = Molecule(ifile)\n",
        "\n",
        "view.addModelsAsFrames(str(mol))\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": "D99VTrFjvASM"
      },
      "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(pprm) \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[0:1], 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=\"frame\", 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 = \"20000\" #@param [\"1000\", \"5000\", \"10000\", \"20000\", \"50000\", \"100000\"]\n",
        "\n",
        "#@markdown Simulation time (in nanoseconds) and integration time (in femtoseconds): \n",
        "Time = \"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\t\t\t\t\t\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, maxIterations=int(Minimization_steps))\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 = \"5\" #@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\t\t\t\t\t\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/MyDrive/' #@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 = \"5\" #@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",
        "view.addModelsAsFrames(models2)\n",
        "for i, at in enumerate(molecules2[0]):\n",
        "    view.setStyle({\"model\":-1},{'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",
      "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": {
        "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, 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": []
    }
  ]
}