Protein_ligand.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/pablo-arantes/making-it-rain/blob/main/Protein_ligand.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pj2BWZxUDbEE"
      },
      "source": [
        "# **Hello there!**\n",
        "\n",
        "This is a Jupyter notebook for running Molecular Dynamics (MD) simulations using OpenMM engine and AMBER force field for **Protein and Ligand** systems. This notebook is a supplementary material of the paper \"***Making it rain: Cloud-based molecular simulations for everyone***\" ([link here](https://doi.org/10.1021/acs.jcim.1c00998)) and we encourage you to read it before using this pipeline.\n",
        "\n",
        "The main goal of this notebook is to demonstrate how to harness the power of cloud-computing to run microsecond-long MD simulations in a cheap and yet feasible fashion.\n",
        "\n",
        "---\n",
        "\n",
        " **This notebook is NOT a standard protocol for MD simulations!** It is just simple MD pipeline illustrating each step of a simulation protocol.\n",
        "\n",
        "---\n",
        "**Bugs**\n",
        "- If you encounter any bugs, please report the issue to https://github.com/pablo-arantes/making-it-rain/issues\n",
        "\n",
        "**Acknowledgments**\n",
        "- We would like to thank the OpenMM team for developing an excellent and open source engine.\n",
        "\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",
        "\n",
        "- A Making-it-rain by **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)).\n",
        "\n",
        "\n",
        "- Also, credit to [David Koes](https://github.com/dkoes) for his awesome [py3Dmol](https://3dmol.csb.pitt.edu/) plugin.\n",
        "\n",
        "- For related notebooks see: [Making-it-rain](https://github.com/pablo-arantes/making-it-rain)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hoyY6XonD1UX"
      },
      "source": [
        "# **Introduction**\n",
        "\n",
        "In general, MD simulations rely on 1) a set of atomic coordinates of all atoms on a simulation box and 2) a set of force field parameters that describes the interaction energies between atoms.\n",
        "\n",
        "In terms of inputs, we wil need:\n",
        "*  A .pdb file of the protein and a .pdb file of the ligand containing a set of atomic coordinates.\n",
        "\n",
        "In this notebook, we will simulate PDB 3HTB. To build our simulation box, we will use LEaP program (https://ambermd.org/tutorials/pengfei/index.php). The LEaP program is a portal between many chemical structure file types (.pdb and .mol2, primarily), and the Amber model parameter file types such as .lib, .prepi, parm.dat, and .frcmod. Each of the parameter files contains pieces of information needed for constructing a simulation, whether for energy minimization or molecular dynamics. LEaP functions within a larger workflow described in Section 1.1 of the [Amber Manual](https://ambermd.org/doc12/Amber20.pdf).\n",
        "\n",
        "To build ligand topology we will use general AMBER force field (GAFF - http://ambermd.org/antechamber/gaff.html). GAFF is compatible with the AMBER force field and it has parameters for almost all the organic molecules made of C, N, O, H, S, P, F, Cl, Br and I. As a complete force field, GAFF is suitable for study of a great number of molecules in an automatic fashion.\n",
        "\n",
        "You can download the input files examples from [here](https://github.com/pablo-arantes/making-it-rain/tree/main/PROTEIN_LIGAND);\n",
        "## ---\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lh96y6mGFY1D"
      },
      "source": [
        "---\n",
        "---\n",
        "# **Setting the environment for MD calculation**\n",
        "\n",
        "Firstly, we need to install all necessary libraries and packages for our simulation. The main packages we will be installing are:\n",
        "\n",
        "1.    Anaconda (https://docs.conda.io/en/latest/miniconda.html)\n",
        "2.    OpenMM (https://openmm.org/)\n",
        "3.    PyTraj (https://amber-md.github.io/pytraj/latest/index.html)\n",
        "4.    py3Dmol (https://pypi.org/project/py3Dmol/)\n",
        "5.    ProLIF (https://github.com/chemosim-lab/ProLIF)\n",
        "6.    Numpy (https://numpy.org/)\n",
        "7.    Matplotlib (https://matplotlib.org/)\n",
        "8.    AmberTools (https://ambermd.org/AmberTools.php)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Install Conda Colab**\n",
        "#@markdown It will restart the kernel (session), don't worry.\n",
        "!pip install -q condacolab\n",
        "import condacolab\n",
        "condacolab.install()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "oZReYQjmBLbG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Install dependencies**\n",
        "#@markdown It will take a few minutes, please, drink a coffee and wait. ;-)\n",
        "# install dependencies\n",
        "import subprocess\n",
        "import sys\n",
        "subprocess.run(\"rm -rf /usr/local/conda-meta/pinned\", shell=True)\n",
        "subprocess.run(\"mamba install -c conda-forge ambertools -y\", shell=True)\n",
        "import pytraj as pt\n",
        "subprocess.run(\"pip -q install py3Dmol\", shell=True)\n",
        "subprocess.run(\"pip install git+https://github.com/pablo-arantes/biopandas\", shell=True)\n",
        "subprocess.run(\"mamba install openmm=7.7.0 pdbfixer -y\", shell=True)\n",
        "subprocess.run(\"mamba install openmmforcefields -c conda-forge -y\", shell=True)\n",
        "subprocess.run(\"pip install prolif==1.1.0\", shell=True)\n",
        "subprocess.run(\"mamba install -c conda-forge openbabel -y\", shell=True)\n",
        "subprocess.run(\"mamba install -c conda-forge parmed -y\", shell=True)\n",
        "subprocess.run(\"pip install --upgrade MDAnalysis==2.4.2\", shell=True)\n",
        "subprocess.run(\"pip install numpy==1.23.5\", shell=True)\n",
        "\n",
        "#load dependencies\n",
        "from openmm import app, unit\n",
        "from openmm.app import HBonds, NoCutoff, PDBFile\n",
        "from openff.toolkit.topology import Molecule, Topology\n",
        "from openff.toolkit.typing.engines.smirnoff import ForceField\n",
        "from openff.toolkit.utils import get_data_file_path\n",
        "import parmed as pmd\n",
        "from biopandas.pdb import PandasPdb\n",
        "import openmm as mm\n",
        "from openmm import *\n",
        "from openmm.app import *\n",
        "from openmm.unit import *\n",
        "import 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",
        "# !wget  https://raw.githubusercontent.com/openforcefield/openff-forcefields/master/openforcefields/offxml/openff_unconstrained-2.0.0.offxml"
      ],
      "metadata": {
        "id": "H-oBjHKEBPJY",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fDQnAKJLFxtt"
      },
      "source": [
        "## Using Google Drive to store simulation data\n",
        "\n",
        "Google Colab does not allow users to keep data on their computing nodes. However, we can use Google Drive to read, write, and store our simulations files. Therefore, we suggest to you to:\n",
        "\n",
        "1.   Create a folder in your own Google Drive and copy the necessary input files there.\n",
        "2.   Copy the path of your created directory. We will use it below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Lm7Akepv_vl-",
        "cellView": "form"
      },
      "source": [
        "#@title ### **Import Google Drive**\n",
        "#@markdown Click in the \"Run\" buttom to make your Google Drive accessible.\n",
        "from google.colab import drive\n",
        "\n",
        "drive.flush_and_unmount()\n",
        "drive.mount('/content/drive', force_remount=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lOKg9eH_ueRn",
        "cellView": "form"
      },
      "source": [
        "#@title **Check if you correctly allocated GPU nodes**\n",
        "\n",
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Select the Runtime > \"Change runtime type\" menu to enable a GPU accelerator, ')\n",
        "  print('and then re-execute this cell.')\n",
        "else:\n",
        "  print(gpu_info)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BZwl66HTGI7v"
      },
      "source": [
        "---\n",
        "---\n",
        "# **Loading the necessary input files**\n",
        "\n",
        "At this point, we should have all libraries and dependencies installed and all necessary input files already at your Google Drive folder.\n",
        "\n",
        "**Important**: Make sure the PDB file points to the correct pathway. If necessary, correct the pathway and re-upload the files. We will merge the receptor and ligand structure objects to form the complex. Note that the coordinates of protein and ligand are determined by the PDB file, and they should be consistent with the ligand being positioned in the binding pocket.\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., protein_inpupt.pdb, MyDrive/protein_ligand and so on.**"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Please, provide the necessary input files below**:\n",
        "#@markdown **Important:** The protonation of your ligand is crucial for the correct parameterization of the molecule.\n",
        "# %%capture\n",
        "import rdkit\n",
        "import mdtraj as md\n",
        "from rdkit import Chem\n",
        "from rdkit.Chem import AllChem\n",
        "from rdkit.Chem import Draw\n",
        "from rdkit.Chem import rdMolTransforms\n",
        "from rdkit.Chem.Draw import rdMolDraw2D\n",
        "from rdkit.Chem import rdDepictor\n",
        "from rdkit.Chem import rdForceFieldHelpers\n",
        "from IPython.display import SVG\n",
        "import ipywidgets as widgets\n",
        "import rdkit\n",
        "from rdkit.Chem.Draw import IPythonConsole\n",
        "AllChem.SetPreferCoordGen(True)\n",
        "from IPython.display import Image\n",
        "from pdbfixer import PDBFixer\n",
        "from openbabel import pybel\n",
        "import os\n",
        "import subprocess\n",
        "import warnings\n",
        "\n",
        "# Suppress UserWarnings\n",
        "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
        "\n",
        "\n",
        "Protein_PDB_file_name = 'protein.pdb' #@param {type:\"string\"}\n",
        "remove_waters = \"yes\" #@param [\"yes\", \"no\" ]\n",
        "if remove_waters == \"yes\":\n",
        "  no_waters = \"nowat\"\n",
        "else:\n",
        "  no_waters = ''\n",
        "\n",
        "Ligand_PDB_file_name = 'ligand.pdb'  #@param {type:\"string\"}\n",
        "\n",
        "Add_ligand_hydrogens = \"Yes\" #@param [\"Yes\", \"No\"]\n",
        "Charge = 0 #@param {type:\"slider\", min:-10, max:10, step:1}\n",
        "\n",
        "\n",
        "ligand_name = Ligand_PDB_file_name\n",
        "Google_Drive_Path = '/content/drive/MyDrive/' #@param {type:\"string\"}\n",
        "workDir = Google_Drive_Path\n",
        "initial_pdb = os.path.join(workDir, str(Protein_PDB_file_name))\n",
        "prepareforleap = os.path.join(workDir, \"prepareforleap.in\")\n",
        "ligand_pdb = os.path.join(workDir, str(ligand_name))\n",
        "ligand_pdb2 = os.path.join(workDir, \"ligand_H.pdb\")\n",
        "starting = os.path.join(workDir, \"starting1.pdb\")\n",
        "starting2 = os.path.join(workDir, \"starting2.pdb\")\n",
        "starting_end = os.path.join(workDir, \"starting_end.pdb\")\n",
        "\n",
        "\n",
        "def mol_with_atom_index(mol):\n",
        "    for atom in mol.GetAtoms():\n",
        "        atom.SetAtomMapNum(atom.GetIdx())\n",
        "    return mol\n",
        "\n",
        "def remove_lines(filename):\n",
        "    with open(filename, 'r') as file:\n",
        "        ter_count = 0\n",
        "        for line in file:\n",
        "            if line.startswith('TER'):\n",
        "                ter_count += 1\n",
        "                if ter_count >= 1:\n",
        "                    yield line\n",
        "                    for i in range(3):\n",
        "                        line = next(file, None)\n",
        "                        if line is not None and line.startswith('ATOM') and line.split()[2] in ['P', 'OP1', 'OP2']:\n",
        "                            continue\n",
        "                        else:\n",
        "                            yield line\n",
        "                else:\n",
        "                    yield line\n",
        "            else:\n",
        "                yield line\n",
        "\n",
        "# if Add_ligand_hydrogens == \"Yes\":\n",
        "#   mol= [m for m in pybel.readfile(filename=ligand_pdb, format='pdb')][0]\n",
        "#   out=pybel.Outputfile(filename=\"temp.mol\",format='mol',overwrite=True)\n",
        "#   out.write(mol)\n",
        "#   out.close()\n",
        "\n",
        "#   mol = Chem.MolFromMolFile('temp.mol', removeHs=True)\n",
        "#   hmol = Chem.AddHs(mol)\n",
        "#   mp = AllChem.MMFFGetMoleculeProperties(hmol)\n",
        "#   ff = AllChem.MMFFGetMoleculeForceField(hmol, mp)\n",
        "#   for a in hmol.GetAtoms():\n",
        "#     if (a.GetAtomicNum() > 1):\n",
        "#       ff.MMFFAddPositionConstraint(a.GetIdx(), 0, 1.e4)\n",
        "#   ff.Minimize(maxIts=1000)\n",
        "#   charge_mol = Chem.rdPartialCharges.ComputeGasteigerCharges(hmol)\n",
        "#   charge = Chem.GetFormalCharge(hmol)\n",
        "#   print(\"Charge = \" + str(charge))\n",
        "#   # AllChem.MolToMolFile(hmol, (os.path.join(workDir, f\"start_min.mol\")))\n",
        "#   AllChem.MolToPDBFile(hmol, ligand_pdb2)\n",
        "#   mol_end = mol_with_atom_index(hmol)\n",
        "#   IPythonConsole.drawMol3D(hmol)\n",
        "# else:\n",
        "#   mol= [m for m in pybel.readfile(filename=ligand_pdb, format='pdb')][0]\n",
        "#   out=pybel.Outputfile(filename=\"temp.mol\",format='mol',overwrite=True)\n",
        "#   out.write(mol)\n",
        "#   out.close()\n",
        "\n",
        "#   hmol = Chem.MolFromMolFile('temp.mol', removeHs=False)\n",
        "#   mp = AllChem.MMFFGetMoleculeProperties(hmol)\n",
        "#   ff = AllChem.MMFFGetMoleculeForceField(hmol, mp)\n",
        "#   for a in hmol.GetAtoms():\n",
        "#     if (a.GetAtomicNum() > 1):\n",
        "#       ff.MMFFAddPositionConstraint(a.GetIdx(), 0, 1.e4)\n",
        "#   ff.Minimize(maxIts=1000)\n",
        "#   charge_mol = Chem.rdPartialCharges.ComputeGasteigerCharges(hmol)\n",
        "#   charge = Chem.GetFormalCharge(hmol)\n",
        "#   print(\"Charge = \" + str(charge))\n",
        "#   # AllChem.MolToMolFile(hmol, (os.path.join(workDir, f\"start_min.mol\")))\n",
        "#   AllChem.MolToPDBFile(hmol, ligand_pdb2)\n",
        "#   mol_end = mol_with_atom_index(hmol)\n",
        "#   IPythonConsole.drawMol3D(hmol)\n",
        "\n",
        "\n",
        "\n",
        "#Add hydrogens in the ligand\n",
        "if Add_ligand_hydrogens == \"Yes\":\n",
        "  fixer = PDBFixer(filename=ligand_pdb)\n",
        "  PDBFile.writeFile(fixer.topology, fixer.positions, open(\"temp.pdb\", 'w'))\n",
        "\n",
        "  ppdb = PandasPdb().read_pdb(\"temp.pdb\")\n",
        "  ppdb.df['ATOM'] = ppdb.df['ATOM']\n",
        "  ppdb.df['HETATM']= ppdb.df['HETATM'][ppdb.df['HETATM']['element_symbol'] != 'H']\n",
        "  ppdb.to_pdb(path=\"temp.pdb\", records=['ATOM', 'HETATM'], gz=False, append_newline=True)\n",
        "\n",
        "  mol= [m for m in pybel.readfile(filename=\"temp.pdb\", format='pdb')][0]\n",
        "  mol.calccharges\n",
        "  mol.addh()\n",
        "  out=pybel.Outputfile(filename=\"temp2.pdb\",format='pdb',overwrite=True)\n",
        "  out.write(mol)\n",
        "  out.close()\n",
        "\n",
        "  md.load(\"temp2.pdb\").save(\"temp2.pdb\")\n",
        "\n",
        "  halogens = ['Cl', 'F', 'Br', 'I']\n",
        "  atom_id = []\n",
        "  H_id = []\n",
        "  with open(\"temp2.pdb\") as f:\n",
        "      for line in f:\n",
        "        data = line.split()\n",
        "        if data[0] == \"ATOM\":\n",
        "          if data[2] in halogens:\n",
        "            atom_id.append(data[1])\n",
        "        if data[0] == \"CONECT\":\n",
        "          if data[1] in atom_id:\n",
        "            if len(data) > 3:\n",
        "              H_id.append(data[3])\n",
        "              H_id.append(data[4])\n",
        "              H_id.append(data[5])\n",
        "\n",
        "  # with open(ligand_pdb2, 'w') as h:\n",
        "  #   with open(\"temp2.pdb\") as f:\n",
        "  #     for line in f:\n",
        "  #       data = line.split()\n",
        "  #       if data[0] == \"ATOM\":\n",
        "  #         if data[1] not in H_id:\n",
        "  #           print(line, file=h)\n",
        "  #       elif data[0] == \"CONECT\":\n",
        "  #           if data[1] not in atom_id:\n",
        "  #             print(line, file=h)\n",
        "  #       else:\n",
        "  #         print(line, file=h)\n",
        "\n",
        "  with open(ligand_pdb2, 'w') as h:\n",
        "    with open(\"temp2.pdb\") as f:\n",
        "        for line in f:\n",
        "            if line.strip():  # Check if line is not empty or just whitespace\n",
        "                data = line.split()\n",
        "                if len(data) > 0 and data[0] not in [\"TER\", \"ENDMDL\"]:  # Exclude lines starting with TER or ENDMDL\n",
        "                    if data[0] == \"ATOM\":\n",
        "                        if data[1] not in H_id:\n",
        "                            print(line, end='', file=h)  # Avoid adding extra newline\n",
        "                    elif data[0] == \"CONECT\":\n",
        "                        if data[1] not in atom_id:\n",
        "                            print(line, end='', file=h)\n",
        "                    else:\n",
        "                        print(line, end='', file=h)\n",
        "\n",
        "  mol= [m for m in pybel.readfile(filename=ligand_pdb2, format='pdb')][0]\n",
        "  out=pybel.Outputfile(filename=\"temp.mol\",format='mol',overwrite=True)\n",
        "  out.write(mol)\n",
        "  out.close()\n",
        "  hmol = Chem.MolFromMolFile('temp.mol', removeHs=False)\n",
        "  # charge_mol = Chem.rdPartialCharges.ComputeGasteigerCharges(hmol)\n",
        "  charge = Charge\n",
        "  print(\"Charge = \" + str(charge))\n",
        "  mol_end = mol_with_atom_index(hmol)\n",
        "  IPythonConsole.drawMol3D(hmol)\n",
        "\n",
        "else:\n",
        "  ppdb = PandasPdb().read_pdb(ligand_pdb)\n",
        "  ppdb.df['ATOM'] = ppdb.df['ATOM']\n",
        "  ppdb.to_pdb(path=\"temp.pdb\", records=['ATOM', 'HETATM'], gz=False, append_newline=True)\n",
        "  mol= [m for m in pybel.readfile(filename=\"temp.pdb\", format='pdb')][0]\n",
        "  mol.calccharges\n",
        "  out=pybel.Outputfile(filename=\"temp2.pdb\",format='pdb',overwrite=True)\n",
        "  out.write(mol)\n",
        "  out.close()\n",
        "\n",
        "  md.load(\"temp2.pdb\").save(\"temp2.pdb\")\n",
        "\n",
        "  with open(ligand_pdb2, 'w') as h:\n",
        "    with open(\"temp2.pdb\") as f:\n",
        "        for line in f:\n",
        "            if line.strip() and not line.startswith((\"TER\", \"ENDMDL\")):\n",
        "                print(line, end='', file=h)\n",
        "\n",
        "  mol= [m for m in pybel.readfile(filename=ligand_pdb2, format='pdb')][0]\n",
        "  out=pybel.Outputfile(filename=\"temp.mol\",format='mol',overwrite=True)\n",
        "  out.write(mol)\n",
        "  out.close()\n",
        "  hmol = Chem.MolFromMolFile('temp.mol', removeHs=False)\n",
        "  # charge_mol = Chem.rdPartialCharges.ComputeGasteigerCharges(hmol)\n",
        "  charge = Charge\n",
        "  print(\"Charge = \" + str(charge))\n",
        "  mol_end = mol_with_atom_index(hmol)\n",
        "  IPythonConsole.drawMol3D(hmol)\n",
        "\n",
        "\n",
        "#Fix protein\n",
        "f = open(prepareforleap, \"w\")\n",
        "f.write(\"\"\"parm \"\"\" + str(initial_pdb) + \"\\n\"\n",
        "\"\"\"loadcrd \"\"\" + str(initial_pdb) + \"\"\" name edited\"\"\" + \"\\n\"\n",
        "\"\"\"prepareforleap crdset edited name from-prepareforleap \\ \"\"\" + \"\\n\"\n",
        "\"\"\"pdbout \"\"\" + str(starting) + \" \" + str(no_waters) + \"\"\" noh\"\"\" + \"\\n\"\n",
        "\"\"\"go \"\"\")\n",
        "f.close()\n",
        "\n",
        "prepareforleap_command = \"cpptraj -i \" + str(prepareforleap)\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "with open('prepareforleap.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(prepareforleap_command)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "\n",
        "subprocess.run([\"chmod 700 prepareforleap.sh\"], shell=True)\n",
        "subprocess.run([\"./prepareforleap.sh\"], shell=True,)\n",
        "\n",
        "\n",
        "pdb4amber_cmd = \"pdb4amber -i \" + str(starting) + \" -o \" + str(starting_end) + \" -a\"\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "\n",
        "with open('pdb4amber.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(pdb4amber_cmd)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "\n",
        "subprocess.run([\"chmod 700 pdb4amber.sh\"], shell=True)\n",
        "subprocess.run([\"./pdb4amber.sh\"], shell=True,)\n",
        "\n",
        "# with open(starting_end, 'w') as out_file:\n",
        "#     for line in remove_lines(starting2):\n",
        "#         out_file.write(line)\n",
        "\n",
        "protein_check = os.path.exists(starting_end)\n",
        "ligand_check = os.path.exists(ligand_pdb2)\n",
        "\n",
        "if protein_check == True and ligand_check == True:\n",
        "  print(\"Successfully generated protein and ligand files! :-)\")\n",
        "else:\n",
        "  print(\"ERROR: Check your inputs! \")\n",
        "\n",
        "#@markdown ---"
      ],
      "metadata": {
        "cellView": "form",
        "id": "lAXZm7bIZiVf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "od6pJ_TiLLUo",
        "cellView": "form"
      },
      "source": [
        "from rdkit import Chem\n",
        "from openff.toolkit.typing.engines.smirnoff import ForceField\n",
        "import parmed\n",
        "\n",
        "#@title **Parameters to generate the topology:**\n",
        "\n",
        "#@markdown **Parameters to generate the protein topology:**\n",
        "\n",
        "Force_field = \"ff19SB\" #@param [\"ff19SB\", \"ff14SB\"]\n",
        "if Force_field == \"ff19SB\":\n",
        "  ff = \"leaprc.protein.ff19SB\"\n",
        "else:\n",
        "  ff = \"leaprc.protein.ff14SB\"\n",
        "\n",
        "Water_type = \"TIP3P\" #@param [\"TIP3P\", \"OPC\"]\n",
        "if Water_type == \"TIP3P\":\n",
        "  water = \"leaprc.water.tip3p\"\n",
        "  water_box = \"TIP3PBOX\"\n",
        "else:\n",
        "  water = \"leaprc.water.opc\"\n",
        "  water_box = \"OPCBOX\"\n",
        "\n",
        "#@markdown Size Box (Angstrons):\n",
        "\n",
        "Size_box = 12 #@param {type:\"slider\", min:10, max:20, step:1}\n",
        "size_box = Size_box\n",
        "\n",
        "#@markdown **ATTENTION**: Give the concentration in Molar units, AMBER tleap will neutralize your system automatically:\n",
        "\n",
        "Ions = \"NaCl\" #@param [\"NaCl\", \"KCl\" ]\n",
        "\n",
        "Concentration = \"0.15\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown **Parameters to generate the ligand topology:**\n",
        "\n",
        "Ligand_Force_field = \"GAFF2\" #@param [\"GAFF2\"]\n",
        "\n",
        "# Ligand_isomer = \"1\" #@param {type:\"string\", min:1, max:10, step:100}\n",
        "# if chiralCentres == []:\n",
        "#   isomer_end = isomer1\n",
        "# else:\n",
        "#   isomer_end = globals()[f\"isomer{Ligand_isomer}\"]\n",
        "\n",
        "Ligand_net_charges = charge\n",
        "\n",
        "#@markdown ---\n",
        "tleap = os.path.join(workDir, \"tleap.in\")\n",
        "top_nw = os.path.join(workDir, \"SYS_nw.prmtop\")\n",
        "crd_nw = os.path.join(workDir, \"SYS_nw.crd\")\n",
        "pdb_nw = os.path.join(workDir, \"SYS_nw.pdb\")\n",
        "top = os.path.join(workDir, \"SYS_gaff2.prmtop\")\n",
        "crd = os.path.join(workDir, \"SYS_gaff2.crd\")\n",
        "pdb = os.path.join(workDir, \"SYS.pdb\")\n",
        "ligand_noh = os.path.join(workDir, \"ligand_noh.pdb\")\n",
        "ligand_h = os.path.join(workDir, \"ligand_h.pdb\")\n",
        "ligand_mol2 = os.path.join(workDir, \"ligand.mol2\")\n",
        "ligand_frcmod = os.path.join(workDir, \"ligand.frcmod\")\n",
        "lig_new = os.path.join(workDir, \"ligand_gaff.pdb\")\n",
        "protein_ligand = os.path.join(workDir, \"protein_ligand.pdb\")\n",
        "lib = os.path.join(workDir, \"lig.lib\")\n",
        "\n",
        "#gaff_command1 = \"pdb4amber -i \" + str(ligand_pdb2) + \" -o \" + str(ligand_h)\n",
        "gaff_command1 = \"pdb4amber -i \" + str(ligand_pdb2) + \" -o \" + str(ligand_h)\n",
        "gaff_command3 = \"antechamber -i \" + str(ligand_h) + \" -fi pdb -o \" + str(ligand_mol2) + \" -fo mol2 -c bcc -nc \" + str(Ligand_net_charges) + \" -rn LIG -at gaff2\"\n",
        "gaff_command4 = \"parmchk2 -i \" + str(ligand_mol2) + \" -f mol2 -o \" + str(ligand_frcmod) + \" -s gaff2\"\n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "\n",
        "with open('gaff.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(gaff_command1)\n",
        "    print(gaff_command3)\n",
        "    print(gaff_command4)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "\n",
        "!chmod 700 gaff.sh 2>&1 1>/dev/null\n",
        "!bash gaff.sh >/dev/null 2>&1\n",
        "\n",
        "f = open(tleap, \"w\")\n",
        "f.write(\"\"\"source \"\"\" + str(ff) + \"\\n\"\n",
        "\"\"\"source leaprc.gaff2\n",
        "LIG = loadmol2 \"\"\" + str(ligand_mol2) + \"\\n\"\n",
        "\"\"\"loadamberparams \"\"\" + str(ligand_frcmod) + \"\\n\"\n",
        "\"\"\"saveoff LIG \"\"\" + str(lib) + \"\\n\"\n",
        "\"\"\"savepdb LIG \"\"\" + str(lig_new) + \"\\n\"\n",
        "\"\"\"quit\"\"\")\n",
        "f.close()\n",
        "\n",
        "tleap_command = \"tleap -f \" + str(tleap)\n",
        "cat_command = \"cat \" + str(starting_end) + \" \" + str(lig_new) + str(\" > \") + str(protein_ligand)\n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "\n",
        "with open('run_tleap.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(tleap_command)\n",
        "    print(cat_command)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "\n",
        "!chmod 700 run_tleap.sh 2>&1 1>/dev/null\n",
        "!bash run_tleap.sh 2>&1 1>/dev/null\n",
        "\n",
        "ppdb = PandasPdb().read_pdb(protein_ligand)\n",
        "ppdb.df['ATOM'] = ppdb.df['ATOM']\n",
        "ppdb.df['OTHERS'] = [ppdb.df['OTHERS'] != 'OTHERS']\n",
        "ppdb.to_pdb(path=protein_ligand, records=['ATOM', 'HETATM'], gz=False, append_newline=True)\n",
        "\n",
        "f = open(tleap, \"w\")\n",
        "f.write(\"\"\"source \"\"\" + str(ff) + \"\\n\"\n",
        "\"\"\"source leaprc.DNA.OL15\n",
        "source leaprc.RNA.OL3\n",
        "source leaprc.GLYCAM_06j-1\n",
        "source leaprc.gaff2\n",
        "source \"\"\"  + str(water) + \"\\n\"\n",
        "\"\"\"loadamberparams \"\"\" + str(ligand_frcmod) + \"\\n\"\n",
        "\"\"\"loadoff \"\"\" + str(lib) + \"\\n\"\n",
        "\"\"\"SYS = loadpdb \"\"\" + str(protein_ligand) + \"\\n\"\n",
        "\"\"\"alignaxes SYS\n",
        "savepdb SYS \"\"\" + str(pdb_nw) + \"\\n\"\n",
        "\"\"\"saveamberparm SYS \"\"\" + str(top_nw) + \" \" + str(crd_nw) + \"\\n\"\n",
        "\"\"\"solvatebox SYS \"\"\" + str(water_box) + \" \" + str(size_box) +  \"\"\" 0.7\n",
        "saveamberparm SYS \"\"\" + str(top) + \" \" + str(crd) + \"\\n\"\n",
        "\"\"\"savepdb SYS \"\"\" + str(pdb) + \"\\n\"\n",
        "\"\"\"quit\"\"\")\n",
        "f.close()\n",
        "\n",
        "tleap_command = \"tleap -f \" + str(tleap)\n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "\n",
        "with open('run_tleap.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(tleap_command)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "\n",
        "SYS = os.path.join(workDir, \"SYS*\")\n",
        "rm_sys = \"rm \" + SYS\n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "\n",
        "with open('rm_sys.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(rm_sys)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "\n",
        "!chmod 700 rm_sys.sh 2>&1 1>/dev/null\n",
        "!bash rm_sys.sh 2> /dev/null\n",
        "\n",
        "!chmod 700 run_tleap.sh 2>&1 1>/dev/null\n",
        "!bash run_tleap.sh 2>&1 1>/dev/null\n",
        "\n",
        "\n",
        "!grep \"Volume:\" leap.log > temp.txt\n",
        "with open(\"temp.txt\", 'r') as f:\n",
        "  for line in f:\n",
        "        vol = float(line.split()[1])\n",
        "\n",
        "vol_lit  = vol * pow(10, -27)\n",
        "atom_lit = 9.03 * pow(10, 22)\n",
        "conc = float(Concentration)\n",
        "num_ion = int(vol_lit * (conc/0.15) * atom_lit)\n",
        "\n",
        "if Ions == \"NaCl\":\n",
        "  pos_neut = \"Na+ 0\"\n",
        "  pos_num = \"Na+ \" + str(num_ion)\n",
        "  Cl_num = num_ion\n",
        "else:\n",
        "  pos_neut = \"K+ 0\"\n",
        "  pos_num = \"K+ \" + str(num_ion)\n",
        "  Cl_num = num_ion\n",
        "\n",
        "f = open(tleap, \"w\")\n",
        "f.write(\"\"\"source \"\"\" + str(ff) + \"\\n\"\n",
        "\"\"\"source leaprc.DNA.OL15\n",
        "source leaprc.RNA.OL3\n",
        "source leaprc.GLYCAM_06j-1\n",
        "source leaprc.gaff2\n",
        "source \"\"\"  + str(water) + \"\\n\"\n",
        "\"\"\"loadamberparams \"\"\" + str(ligand_frcmod) + \"\\n\"\n",
        "\"\"\"loadoff \"\"\" + str(lib) + \"\\n\"\n",
        "\"\"\"SYS = loadpdb \"\"\" + str(protein_ligand) + \"\\n\"\n",
        "\"\"\"alignaxes SYS\n",
        "check SYS\n",
        "charge SYS\n",
        "addions SYS \"\"\" + str(pos_neut) + \"\\n\"\n",
        "\"\"\"addions SYS Cl- 0\n",
        "check SYS\n",
        "charge SYS\n",
        "savepdb SYS \"\"\" + str(pdb_nw) + \"\\n\"\n",
        "\"\"\"saveamberparm SYS \"\"\" + str(top_nw) + \" \" + str(crd_nw) + \"\\n\"\n",
        "\"\"\"solvatebox SYS \"\"\" + str(water_box) + \" \" + str(size_box) +  \"\"\" 0.7 \"\"\" + \"\\n\"\n",
        "\"\"\"addIonsRand SYS \"\"\" + str(pos_num) + \"\"\" Cl- \"\"\" + str(Cl_num) + \"\\n\"\n",
        "\"\"\"saveamberparm SYS \"\"\" + str(top) + \" \" + str(crd) + \"\\n\"\n",
        "\"\"\"savepdb SYS \"\"\" + str(pdb) + \"\\n\"\n",
        "\"\"\"quit\"\"\")\n",
        "f.close()\n",
        "\n",
        "\n",
        "!chmod 700 run_tleap.sh 2>&1 1>/dev/null\n",
        "!bash run_tleap.sh 2>&1 1>/dev/null\n",
        "\n",
        "# if Ligand_Force_field == \"OpenFF 2.0.0 (Sage)\":\n",
        "#   mol = Chem.MolFromPDBFile(lig_new, removeHs=False)\n",
        "#   Chem.MolToPDBFile(mol, os.path.join(workDir, \"ligand_openFF.pdb\"))\n",
        "\n",
        "#   in_prmtop = top\n",
        "#   in_crd = crd\n",
        "#   orig_structure = parmed.amber.AmberParm(in_prmtop, in_crd)\n",
        "#   pieces = orig_structure.split()\n",
        "#   for piece in pieces:\n",
        "#       print(f\"There are {len(piece[1])} instance(s) of {piece[0]}\")\n",
        "\n",
        "#   from openmm.app import PDBFile\n",
        "#   from openff.toolkit.topology import Molecule, Topology\n",
        "#   from openff.toolkit.tests.utils import get_data_file_path\n",
        "\n",
        "#   # rdmol = Chem.MolFromMolFile(os.path.join(workDir, \"ligand_openFF.sdf\"))\n",
        "#   # ligand_off_molecule = Molecule.from_rdkit(rdmol, hydrogens_are_explicit=True)\n",
        "\n",
        "#   ligand_off_molecule = Molecule.from_smiles(isomer_end)\n",
        "#   ligand_pdbfile = PDBFile(os.path.join(workDir, \"ligand_openFF.pdb\"))\n",
        "#   ligand_off_topology = Topology.from_openmm(\n",
        "#       ligand_pdbfile.topology,\n",
        "#       unique_molecules=[ligand_off_molecule],)\n",
        "\n",
        "#   force_field = ForceField(\"openff_unconstrained-2.0.0.offxml\")\n",
        "\n",
        "#   ligand_system = force_field.create_openmm_system(ligand_off_topology)\n",
        "#   new_ligand_structure = parmed.openmm.load_topology(\n",
        "#       ligand_off_topology.to_openmm(),\n",
        "#       ligand_system,\n",
        "#       xyz=pieces[1][0].positions,)\n",
        "\n",
        "#   new_ligand_structure.save(os.path.join(workDir, \"ligand.prmtop\"), overwrite=True)\n",
        "#   new_ligand_structure.save(os.path.join(workDir, \"ligand.inpcrd\"), overwrite=True)\n",
        "\n",
        "#   # Check how many atoms and which order elements are in the new ligand\n",
        "#   n_atoms_new = len(new_ligand_structure.atoms)\n",
        "#   elements_new = [atom.element for atom in new_ligand_structure.atoms]\n",
        "\n",
        "#   # Check how many atoms and which order elements are in the old ligand\n",
        "#   old_ligand_structure, n_copies = pieces[1]\n",
        "#   n_atoms_old = len(old_ligand_structure.atoms)\n",
        "#   elements_old = [atom.element for atom in old_ligand_structure.atoms]\n",
        "\n",
        "#   print(\n",
        "#       f\"There are {n_atoms_old} in the old ligand structure and {n_atoms_new} atoms \"\n",
        "#       f\"in the new ligand structure\")\n",
        "\n",
        "#   # Print out error message if number of atoms doesn't match\n",
        "#   if n_atoms_new != n_atoms_old:\n",
        "#       print(\n",
        "#           \"Error: Number of atoms in input ligand doesn't match number extracted \"\n",
        "#           \"from prmtop file.\")\n",
        "\n",
        "#   if elements_new != elements_old:\n",
        "#       print(\n",
        "#           \"Error: Elements in input ligand don't match elements in the ligand \"\n",
        "#           \"from the prmtop file.\")\n",
        "#       print(f\"Old elements: {elements_old}\")\n",
        "#       print(f\"New elements: {elements_new}\")\n",
        "\n",
        "#   # Create a new, empty system\n",
        "#   complex_structure = parmed.Structure()\n",
        "\n",
        "#   # Add the protein. Convert explicitly to an AmberParm object to ensure that 1-4 scaling factors are preserved.\n",
        "#   complex_structure += parmed.amber.AmberParm.from_structure(pieces[0][0])\n",
        "\n",
        "#   # Add the ligand\n",
        "#   complex_structure += parmed.amber.AmberParm.from_structure(new_ligand_structure)\n",
        "\n",
        "#   # Add ions and Waters\n",
        "#   ppdb = PandasPdb().read_pdb(pdb)\n",
        "#   Cl = [ppdb.df['ATOM']['atom_name'] == 'Cl-']\n",
        "#   Na = [ppdb.df['ATOM']['atom_name'] == 'Na+']\n",
        "#   K = [ppdb.df['ATOM']['atom_name'] == 'K+']\n",
        "#   Cl = np.array(Cl)\n",
        "#   Na = np.array(Na)\n",
        "#   K = np.array(K)\n",
        "\n",
        "#   if True in Cl and True in Na:\n",
        "#     just_ion1_structure = parmed.Structure()\n",
        "#     just_ion1_structure += pieces[2][0]\n",
        "#     just_ion1_structure *= len(pieces[2][1])\n",
        "\n",
        "#     just_ion2_structure = parmed.Structure()\n",
        "#     just_ion2_structure += pieces[3][0]\n",
        "#     just_ion2_structure *= len(pieces[3][1])\n",
        "\n",
        "#     complex_structure += parmed.amber.AmberParm.from_structure(just_ion1_structure)\n",
        "#     complex_structure += parmed.amber.AmberParm.from_structure(just_ion2_structure)\n",
        "\n",
        "#     just_water_structure = parmed.Structure()\n",
        "#     just_water_structure += pieces[4][0]\n",
        "#     just_water_structure *= len(pieces[4][1])\n",
        "\n",
        "#     complex_structure += parmed.amber.AmberParm.from_structure(just_water_structure)\n",
        "#   elif True in Cl and True in K:\n",
        "#     just_ion1_structure = parmed.Structure()\n",
        "#     just_ion1_structure += pieces[2][0]\n",
        "#     just_ion1_structure *= len(pieces[2][1])\n",
        "\n",
        "#     just_ion2_structure = parmed.Structure()\n",
        "#     just_ion2_structure += pieces[3][0]\n",
        "#     just_ion2_structure *= len(pieces[3][1])\n",
        "\n",
        "#     complex_structure += parmed.amber.AmberParm.from_structure(just_ion1_structure)\n",
        "#     complex_structure += parmed.amber.AmberParm.from_structure(just_ion2_structure)\n",
        "\n",
        "#     just_water_structure = parmed.Structure()\n",
        "#     just_water_structure += pieces[4][0]\n",
        "#     just_water_structure *= len(pieces[4][1])\n",
        "\n",
        "#     complex_structure += parmed.amber.AmberParm.from_structure(just_water_structure)\n",
        "#   elif True in Cl:\n",
        "#     just_ion1_structure = parmed.Structure()\n",
        "#     just_ion1_structure += pieces[2][0]\n",
        "#     just_ion1_structure *= len(pieces[2][1])\n",
        "\n",
        "#     complex_structure += parmed.amber.AmberParm.from_structure(just_ion1_structure)\n",
        "\n",
        "#     just_water_structure = parmed.Structure()\n",
        "#     just_water_structure += pieces[3][0]\n",
        "#     just_water_structure *= len(pieces[3][1])\n",
        "\n",
        "#     complex_structure += parmed.amber.AmberParm.from_structure(just_water_structure)\n",
        "#   elif True in Na:\n",
        "#     just_ion1_structure = parmed.Structure()\n",
        "#     just_ion1_structure += pieces[2][0]\n",
        "#     just_ion1_structure *= len(pieces[2][1])\n",
        "\n",
        "#     complex_structure += parmed.amber.AmberParm.from_structure(just_ion1_structure)\n",
        "\n",
        "#     just_water_structure = parmed.Structure()\n",
        "#     just_water_structure += pieces[3][0]\n",
        "#     just_water_structure *= len(pieces[3][1])\n",
        "\n",
        "#     complex_structure += parmed.amber.AmberParm.from_structure(just_water_structure)\n",
        "#   elif True in K:\n",
        "#     just_ion1_structure = parmed.Structure()\n",
        "#     just_ion1_structure += pieces[2][0]\n",
        "#     just_ion1_structure *= len(pieces[2][1])\n",
        "\n",
        "#     complex_structure += parmed.amber.AmberParm.from_structure(just_ion1_structure)\n",
        "\n",
        "#     just_water_structure = parmed.Structure()\n",
        "#     just_water_structure += pieces[3][0]\n",
        "#     just_water_structure *= len(pieces[3][1])\n",
        "\n",
        "#     complex_structure += parmed.amber.AmberParm.from_structure(just_water_structure)\n",
        "#   else:\n",
        "#     just_water_structure = parmed.Structure()\n",
        "#     just_water_structure += pieces[2][0]\n",
        "#     just_water_structure *= len(pieces[2][1])\n",
        "\n",
        "#     complex_structure += parmed.amber.AmberParm.from_structure(just_water_structure)\n",
        "\n",
        "#   # Copy over the original coordinates and box vectors\n",
        "#   complex_structure.coordinates = orig_structure.coordinates\n",
        "#   complex_structure.box_vectors = orig_structure.box_vectors\n",
        "\n",
        "#   # Export the Structure to AMBER files\n",
        "#   top = os.path.join(workDir, \"SYS_openff.prmtop\")\n",
        "#   crd = os.path.join(workDir, \"SYS_openff.inpcrd\")\n",
        "#   complex_structure.save(top, overwrite=True)\n",
        "#   complex_structure.save(crd, overwrite=True)\n",
        "#   top_openff = os.path.exists(top)\n",
        "#   crd_openff = os.path.exists(crd)\n",
        "#   if top_openff == True and crd_openff == True:\n",
        "#     print(\"Successfully generated topology! :-)\")\n",
        "#   else:\n",
        "#     print(\"ERROR: Check your inputs! \")\n",
        "\n",
        "pdb_amber = os.path.exists(pdb)\n",
        "top_amber = os.path.exists(top)\n",
        "crd_amber = os.path.exists(crd)\n",
        "\n",
        "if pdb_amber == True and top_amber == True and crd_amber == True:\n",
        "  print(\"Successfully generated topology! :-)\")\n",
        "else:\n",
        "  print(\"ERROR: Check your inputs! \")\n",
        "!rm *.sh  ANTECHAMBER* ATOMTYPE* temp.txt >/dev/null 2>&1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C8kKR7bpI86W"
      },
      "source": [
        "## Let's take a look on our simulation box:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vmQ27nZLssjv",
        "cellView": "form"
      },
      "source": [
        "#@title **Show 3D structure**\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "import py3Dmol\n",
        "\n",
        "color = \"gray\" #@param [\"gray\", \"rainbow\"]\n",
        "show_sidechains = False #@param {type:\"boolean\"}\n",
        "show_mainchains = False #@param {type:\"boolean\"}\n",
        "show_ligand = True #@param {type:\"boolean\"}\n",
        "show_box = True #@param {type:\"boolean\"}\n",
        "box_opacity = 0.6 #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
        "\n",
        "\n",
        "def show_pdb(show_sidechains=False, show_mainchains=False, show_ligand = False, show_box = False, color=\"rainbow\"):\n",
        "  view = py3Dmol.view(width=800, height=600)\n",
        "  view.addModel(open(pdb,'r').read(),'pdb')\n",
        "\n",
        "  if color == \"gray\":\n",
        "    view.setStyle({'cartoon':{}})\n",
        "  elif color == \"rainbow\":\n",
        "    view.setStyle({'cartoon': {'color':'spectrum'}})\n",
        "\n",
        "  if show_sidechains:\n",
        "    BB = ['C','O','N']\n",
        "    view.addStyle({'and':[{'resn':[\"GLY\",\"PRO\"],'invert':True},{'atom':BB,'invert':True}]},\n",
        "                        {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "    view.addStyle({'and':[{'resn':\"GLY\"},{'atom':'CA'}]},\n",
        "                        {'sphere':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "    view.addStyle({'and':[{'resn':\"PRO\"},{'atom':['C','O'],'invert':True}]},\n",
        "                        {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "  if show_mainchains:\n",
        "    BB = ['C','O','N','CA']\n",
        "    view.addStyle({'atom':BB},{'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "\n",
        "  if show_box:\n",
        "    view.addSurface(py3Dmol.SAS, {'opacity': box_opacity, 'color':'white'})\n",
        "\n",
        "  if show_ligand:\n",
        "    HP = ['LIG']\n",
        "    view.addStyle({'and':[{'resn':HP}]},\n",
        "                       {'stick':{'colorscheme':'greenCarbon','radius':0.3}})\n",
        "    view.setViewStyle({'style':'outline','color':'black','width':0.1})\n",
        "\n",
        "  view.zoomTo()\n",
        "  return view\n",
        "\n",
        "\n",
        "show_pdb(show_sidechains, show_mainchains, show_ligand, show_box, color).show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NwczIRoxxnVz",
        "cellView": "form"
      },
      "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(os.path.join(workDir, \"SYS_gaff2.prmtop\"), pdb)\n",
        "lig = u.select_atoms(\"resname LIG\")\n",
        "prot = u.select_atoms(\"protein\")\n",
        "\n",
        "# create RDKit-like molecules for visualisation\n",
        "lmol = plf.Molecule.from_mda(lig)\n",
        "pmol = plf.Molecule.from_mda(prot)\n",
        "\n",
        "fp = plf.Fingerprint()\n",
        "fp.run(u.trajectory[::10], lig, prot)\n",
        "df = fp.to_dataframe(return_atoms=True)\n",
        "\n",
        "net = LigNetwork.from_ifp(df, lmol,\n",
        "                          # replace with `kind=\"frame\", frame=0` for the other depiction\n",
        "                          kind=\"frame\", frame=0,\n",
        "                          rotation=270)\n",
        "net.save(os.path.join(workDir, \"initial.html\"))\n",
        "net.display()"
      ],
      "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",
        "# Ligand_Force_field = \"GAFF2\" #@param [\"GAFF2\", \"OpenFF 2.0.0 (Sage)\"]\n",
        "\n",
        "# if Ligand_Force_field == \"OpenFF 2.0.0 (Sage)\":\n",
        "#   top = os.path.join(workDir, \"SYS_openff.prmtop\")\n",
        "#   crd = os.path.join(workDir, \"SYS_openff.inpcrd\")\n",
        "#   pdb = os.path.join(workDir, \"SYS.pdb\")\n",
        "# else:\n",
        "top = os.path.join(workDir, \"SYS_gaff2.prmtop\")\n",
        "crd = os.path.join(workDir, \"SYS_gaff2.crd\")\n",
        "pdb = os.path.join(workDir, \"SYS.pdb\")\n",
        "\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",
      "metadata": {
        "id": "zoamR9iynphz",
        "cellView": "form"
      },
      "source": [
        "#@title **Runs an Equilibration MD simulation (NPT ensemble)**\n",
        "#@markdown Now, let's equilibrate our system!\n",
        "\n",
        "###########################################\n",
        "import openmm as mm\n",
        "from openmm import *\n",
        "from openmm.app import *\n",
        "from openmm.unit import *\n",
        "import pytraj as pt\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",
        "coordinatefile = crd\n",
        "pdbfile = pdb\n",
        "topologyfile = top\n",
        "\n",
        "time_ps = float(Time)*1000\n",
        "simulation_time = float(time_ps)*picosecond\t\t# in ps\n",
        "dt = int(dt_eq)*femtosecond\n",
        "temperature = float(temperature_eq)*kelvin\n",
        "savcrd_freq = int(write_the_trajectory_eq)*picosecond\n",
        "print_freq  = int(write_the_log_eq)*picosecond\n",
        "\n",
        "pressure\t= float(pressure_eq)*bar\n",
        "\n",
        "restraint_fc = int(Force_constant) # kJ/mol\n",
        "\n",
        "nsteps  = int(simulation_time.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nprint  = int(print_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nsavcrd = int(savcrd_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "\n",
        "#############################################\n",
        "# Defining functions to use below:\n",
        "def backup_old_log(pattern, string):\n",
        "\tresult = []\n",
        "\tfor root, dirs, files in os.walk(\"./\"):\n",
        "\t\tfor name in files:\n",
        "\t\t\tif fnmatch.fnmatch(name, pattern):\n",
        "\n",
        "\t\t\t\ttry:\n",
        "\t\t\t\t\tnumber = int(name[-2])\n",
        "\t\t\t\t\tavail = isinstance(number, int)\n",
        "\t\t\t\t\t#print(name,avail)\n",
        "\t\t\t\t\tif avail == True:\n",
        "\t\t\t\t\t\tresult.append(number)\n",
        "\t\t\t\texcept:\n",
        "\t\t\t\t\tpass\n",
        "\n",
        "\tif len(result) > 0:\n",
        "\t\tmaxnumber = max(result)\n",
        "\telse:\n",
        "\t\tmaxnumber = 0\n",
        "\n",
        "\tbackup_file = \"\\#\" + string + \".\" + str(maxnumber + 1) + \"#\"\n",
        "\tos.system(\"mv \" + string + \" \" + backup_file)\n",
        "\treturn backup_file\n",
        "\n",
        "def restraints(system, crd, fc, restraint_array):\n",
        "\n",
        "\tboxlx = system.getDefaultPeriodicBoxVectors()[0][0].value_in_unit(nanometers)\n",
        "\tboxly = system.getDefaultPeriodicBoxVectors()[1][1].value_in_unit(nanometers)\n",
        "\tboxlz = system.getDefaultPeriodicBoxVectors()[2][2].value_in_unit(nanometers)\n",
        "\n",
        "\tif fc > 0:\n",
        "\t\t# positional restraints for all heavy-atoms\n",
        "\t\tposresPROT = CustomExternalForce('k*periodicdistance(x, y, z, x0, y0, z0)^2;')\n",
        "\t\tposresPROT.addPerParticleParameter('k')\n",
        "\t\tposresPROT.addPerParticleParameter('x0')\n",
        "\t\tposresPROT.addPerParticleParameter('y0')\n",
        "\t\tposresPROT.addPerParticleParameter('z0')\n",
        "\n",
        "\t\tfor atom1 in restraint_array:\n",
        "\t\t\tatom1 = int(atom1)\n",
        "\n",
        "\t\t\txpos  = crd.positions[atom1].value_in_unit(nanometers)[0]\n",
        "\t\t\typos  = crd.positions[atom1].value_in_unit(nanometers)[1]\n",
        "\t\t\tzpos  = crd.positions[atom1].value_in_unit(nanometers)[2]\n",
        "\n",
        "\t\t\tposresPROT.addParticle(atom1, [fc, xpos, ypos, zpos])\n",
        "\n",
        "\t\tsystem.addForce(posresPROT)\n",
        "\n",
        "\treturn system\n",
        "##############################################\n",
        "\n",
        "#############################################\n",
        "print(\"\\n> Simulation details:\\n\")\n",
        "print(\"\\tJob name = \" + jobname)\n",
        "print(\"\\tCoordinate file = \" + str(coordinatefile))\n",
        "print(\"\\tPDB file = \" + str(pdbfile))\n",
        "print(\"\\tTopology file = \" + str(topologyfile))\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",
        "if Ligand_Force_field == \"OpenFF 2.0.0 (Sage)\":\n",
        "  print(\"\\t- Reading topology and structure file...\")\n",
        "  prmtop = pmd.load_file(topologyfile)\n",
        "  inpcrd = AmberInpcrdFile(coordinatefile)\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 = complex_structure.createSystem(nonbondedMethod=nonbondedMethod, nonbondedCutoff=nonbondedCutoff,\n",
        "                                          constraints=constraints, rigidWater=rigidWater, ewaldErrorTolerance=ewaldErrorTolerance)\n",
        "else:\n",
        "  print(\"\\t- Reading topology and structure file...\")\n",
        "  prmtop = AmberPrmtopFile(topologyfile)\n",
        "  inpcrd = AmberInpcrdFile(coordinatefile)\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 = prmtop.createSystem(nonbondedMethod=nonbondedMethod, nonbondedCutoff=nonbondedCutoff,\n",
        "                                          constraints=constraints, rigidWater=rigidWater, ewaldErrorTolerance=ewaldErrorTolerance)\n",
        "\n",
        "\n",
        "print(\"\\t- Applying restraints. Force Constant = \" + str(Force_constant) + \"kJ/mol\")\n",
        "pt_system = pt.iterload(coordinatefile, topologyfile)\n",
        "pt_topology = pt_system.top\n",
        "restraint_array = pt.select_atoms('!(:H*) & !(:WAT) & !(:Na+) & !(:Cl-) & !(:Mg+) & !(: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(prmtop.topology, system, integrator)\n",
        "simulation.context.setPositions(inpcrd.positions)\n",
        "if inpcrd.boxVectors is not None:\n",
        "    simulation.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, simulation.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\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LXyL26HCO8Bu"
      },
      "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."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z0JV6Zid50_o",
        "cellView": "form"
      },
      "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",
        "Ligand_Force_field = \"GAFF2\"\n",
        "#@markdown ---\n",
        "#@markdown ### **Parameters for MD Production protocol:**\n",
        "\n",
        "\n",
        "# remove whitespaces\n",
        "Jobname = 'prot_lig_prod' #@param {type:\"string\"}\n",
        "\n",
        "# Ligand_Force_field = \"GAFF2\" #@param [\"GAFF2\", \"OpenFF 2.0.0 (Sage)\"]\n",
        "\n",
        "# if Ligand_Force_field == \"OpenFF 2.0.0 (Sage)\":\n",
        "#   top = os.path.join(workDir, \"SYS_openff.prmtop\")\n",
        "#   crd = os.path.join(workDir, \"SYS_openff.inpcrd\")\n",
        "#   pdb = os.path.join(workDir, \"SYS.pdb\")\n",
        "\n",
        "# else:\n",
        "top = os.path.join(workDir, \"SYS_gaff2.prmtop\")\n",
        "crd = os.path.join(workDir, \"SYS_gaff2.crd\")\n",
        "pdb = os.path.join(workDir, \"SYS.pdb\")\n",
        "\n",
        "\n",
        "#@markdown Simulation time (in nanoseconds), number of strides (integers) and integration timestep (in femtoseconds):\n",
        "Stride_Time = \"10\" #@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 ---"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0QcjKSoqHHhi",
        "cellView": "form"
      },
      "source": [
        "#@title **Runs a Production MD simulation (NPT ensemble) after equilibration**\n",
        "#\n",
        "###########################################\n",
        "import openmm as mm\n",
        "from openmm import *\n",
        "from openmm.app import *\n",
        "from openmm.unit import *\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",
        "coordinatefile = crd\n",
        "pdbfile = os.path.join(workDir, Equilibrated_PDB)\n",
        "topologyfile = top\n",
        "equil_rst_file = os.path.join(workDir, State_file)\n",
        "\n",
        "\n",
        "stride_time_ps = float(stride_time_prod)*1000\n",
        "stride_time = float(stride_time_ps)*picosecond\n",
        "nstride = int(Number_of_strides)\n",
        "dt = int(dt_prod)*femtosecond\n",
        "temperature = float(temperature_prod)*kelvin\n",
        "savcrd_freq = int(write_the_trajectory_prod)*picosecond\n",
        "print_freq  = int(write_the_log_prod)*picosecond\n",
        "\n",
        "pressure\t= float(pressure_prod)*bar\n",
        "\n",
        "simulation_time = stride_time*nstride\n",
        "nsteps  = int(stride_time.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nprint  = int(print_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nsavcrd = int(savcrd_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "firststride = 1 # must be integer\n",
        "#############################################\n",
        "# Defining functions to use below:\n",
        "def backup_old_log(pattern, string):\n",
        "\tresult = []\n",
        "\tfor root, dirs, files in os.walk(\"./\"):\n",
        "\t\tfor name in files:\n",
        "\t\t\tif fnmatch.fnmatch(name, pattern):\n",
        "\n",
        "\t\t\t\ttry:\n",
        "\t\t\t\t\tnumber = int(name[-2])\n",
        "\t\t\t\t\tavail = isinstance(number, int)\n",
        "\t\t\t\t\t#print(name,avail)\n",
        "\t\t\t\t\tif avail == True:\n",
        "\t\t\t\t\t\tresult.append(number)\n",
        "\t\t\t\texcept:\n",
        "\t\t\t\t\tpass\n",
        "\n",
        "\tif len(result) > 0:\n",
        "\t\tmaxnumber = max(result)\n",
        "\telse:\n",
        "\t\tmaxnumber = 0\n",
        "\n",
        "\tbackup_file = \"\\#\" + string + \".\" + str(maxnumber + 1) + \"#\"\n",
        "\tos.system(\"mv \" + string + \" \" + backup_file)\n",
        "\treturn backup_file\n",
        "##############################################\n",
        "\n",
        "#############################################\n",
        "print(\"\\n> Simulation details:\\n\")\n",
        "print(\"\\tJob name = \" + jobname)\n",
        "print(\"\\tCoordinate file = \" + str(coordinatefile))\n",
        "print(\"\\tPDB file = \" + str(pdbfile))\n",
        "print(\"\\tTopology file = \" + str(topologyfile))\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",
        "if Ligand_Force_field == \"OpenFF 2.0.0 (Sage)\":\n",
        "  print(\"\\t- Reading topology and structure file...\")\n",
        "  prmtop = pmd.load_file(topologyfile)\n",
        "  inpcrd = AmberInpcrdFile(coordinatefile)\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 = complex_structure.createSystem(nonbondedMethod=nonbondedMethod, nonbondedCutoff=nonbondedCutoff,\n",
        "                                          constraints=constraints, rigidWater=rigidWater, ewaldErrorTolerance=ewaldErrorTolerance)\n",
        "else:\n",
        "  print(\"\\t- Reading topology and structure file...\")\n",
        "  prmtop = AmberPrmtopFile(topologyfile)\n",
        "  inpcrd = AmberInpcrdFile(coordinatefile)\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 = prmtop.createSystem(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(prmtop.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, simulation.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\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Concatenate and align the trajectory**\n",
        "#@markdown **Important**: The **Google Drive Path**, **Jobname**, **Number of strides**, **stride time** and **trajectory saved frequency** should be the same you have been used to run your simulation in the previous steps.\n",
        "\n",
        "import MDAnalysis as mda\n",
        "from MDAnalysis.analysis import align, rms\n",
        "\n",
        "Google_Drive_Path = '/content/drive/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 = \"10\" #@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(Skip)\n",
        "\n",
        "\n",
        "nw_dcd = os.path.join(workDir, str(Jobname) + output_prefix + \"_nw.\" + 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",
        "\n",
        "if Remove_waters == \"yes\":\n",
        "  #Save topology without waters\n",
        "  gaff_top = pt.load_topology(os.path.join(workDir, \"SYS_gaff2.prmtop\"))\n",
        "  gaff_nw = gaff_top['!:WAT']\n",
        "  gaff_nw.save(os.path.join(workDir, \"SYS_gaff2_nw.prmtop\"))\n",
        "  # Save trajectory without waters\n",
        "  trajlist = pt.load(flist, os.path.join(workDir, \"SYS_gaff2.prmtop\"), stride=Skip)\n",
        "  t0 = trajlist.strip(':WAT')\n",
        "  traj_image = t0.iterframe(autoimage=True, rmsfit=0)\n",
        "  traj_nw = pt.write_traj(nw_dcd, traj_image, overwrite=True, options=Output_format)\n",
        "  traj_dcd_check = os.path.exists(nw_dcd)\n",
        "  traj = nw_dcd\n",
        "  pdb_ref = os.path.join(workDir, \"SYS_gaff2_nw.prmtop\")\n",
        "else:\n",
        "  trajlist = pt.load(flist, os.path.join(workDir, \"SYS_gaff2.prmtop\"), stride=Skip)\n",
        "  traj_image = trajlist.iterframe(autoimage=True, rmsfit=0)\n",
        "  traj = pt.write_traj(whole_dcd, traj_image, overwrite=True, options=Output_format)\n",
        "  traj_dcd_check = os.path.exists(whole_dcd)\n",
        "  traj = whole_dcd\n",
        "  pdb_ref = os.path.join(workDir, \"SYS_gaff2.prmtop\")\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": {
        "cellView": "form",
        "id": "nr0D7rwoc-LZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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",
        "\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",
        "if number_frames_analysis > 10:\n",
        "  stride_animation = number_frames_analysis/10\n",
        "else:\n",
        "  stride_animation = 1\n",
        "\n",
        "u = mda.Universe(pdb_ref, traj)\n",
        "\n",
        "# Write out frames for animation\n",
        "protein = u.select_atoms('not (resname WAT)')\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",
        "# Animation\n",
        "view = py3Dmol.view(width=800, height=600)\n",
        "view.addModelsAsFrames(models)\n",
        "for i, at in enumerate(molecules[0]):\n",
        "    default = {\"cartoon\": {'color': 'spectrum'}}\n",
        "    view.setViewStyle({'style':'outline','color':'black','width':0.1})\n",
        "    view.setStyle({'model': -1, 'serial': i+1}, at.get(\"pymol\", default))\n",
        "    HP = ['LIG']\n",
        "    view.setStyle({\"model\":-1,'and':[{'resn':HP}]},{'stick':{'radius':0.3}})\n",
        "view.zoomTo()\n",
        "view.animate({'loop': \"forward\"})\n",
        "view.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "dkZnDwundH78"
      },
      "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",
        "# load topology\n",
        "u = mda.Universe(pdb_ref, traj)\n",
        "lig = u.select_atoms(\"resname LIG\")\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": "XvaTEfsPR8hZ",
        "cellView": "form"
      },
      "source": [
        "#@title **MM-PBSA method to calculate the binding free energy**\n",
        "#@markdown **Important:** We will now calculate the interaction energy and solvation free energy for the complex, receptor and ligand and average the results to obtain an estimate of the binding free energy. Please note that we will not perform a calculation of the entropy contribution to binding and so strictly speaking our result will not be a true free energy but could be used to compare against similar systems. We will carry out the binding energy calculation using both the MM-GBSA method and the MM-PBSA method for comparison.\n",
        "\n",
        "#@markdown Select the GB/SA input parameters,  the \"OBC\" models (igb=2 and 5) are newer, but appear to give significant improvements and are recommended for most projects (For more information check the Section 4.1 of the [Amber Manual](https://ambermd.org/doc12/Amber20.pdf)):\n",
        "igb = \"2\" #@param [\"1\", \"2\", \"5\", \"7\", \"8\"]\n",
        "\n",
        "import os\n",
        "import locale\n",
        "def getpreferredencoding(do_setlocale = True):\n",
        "    return \"UTF-8\"\n",
        "locale.getpreferredencoding = getpreferredencoding\n",
        "\n",
        "if igb == \"1\":\n",
        "  mbondi = 'mbondi'\n",
        "elif igb == \"2\" or igb == \"5\":\n",
        "  mbondi = 'mbondi2'\n",
        "elif igb == \"7\":\n",
        "  mbondi = 'bondi'\n",
        "elif igb == \"8\":\n",
        "  mbondi = 'mbondi3'\n",
        "else:\n",
        "  pass\n",
        "\n",
        "Salt_concentration = '0.15' #@param {type:\"string\"}\n",
        "fold_MMPBSA = \"MMPBSA_igb_\" + igb\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = 'FINAL_RESULTS_MMPBSA' #@param {type:\"string\"}\n",
        "\n",
        "final_mmpbsa = os.path.join(workDir, Output_name)\n",
        "\n",
        "if number_frames_analysis > 10:\n",
        "  stride = number_frames_analysis/10\n",
        "else:\n",
        "  stride = 1\n",
        "\n",
        "f = open(\"mmpbsa.in\", \"w\")\n",
        "f.write(\"\"\"&general \"\"\"  + \"\\n\"\n",
        "\"\"\"  endframe=\"\"\" + str(int(number_frames_analysis)) + \"\"\",  interval=\"\"\" + str(int(stride)) + \"\"\", strip_mask=:WAT:Na+:Cl-:Mg+:K+, \"\"\" + \"\\n\"\n",
        "\"\"\"/ \"\"\" + \"\\n\"\n",
        "\"\"\"&gb \"\"\" + \"\\n\"\n",
        "\"\"\" igb=\"\"\" + str(igb) +  \"\"\", saltcon=\"\"\" + str(Salt_concentration) +  \"\"\", \"\"\" + \"\\n\"\n",
        "\"\"\"/ \"\"\" + \"\\n\"\n",
        "\"\"\"&pb \"\"\" + \"\\n\"\n",
        "\"\"\" istrng=\"\"\" + str(Salt_concentration) +  \"\"\", inp=2, radiopt=0, prbrad=1.4, \"\"\" + \"\\n\"\n",
        "\"\"\"/\"\"\")\n",
        "f.close()\n",
        "\n",
        "amberhome = \"source /usr/local/amber.sh\"\n",
        "ante_MMPBSA = \"ante-MMPBSA.py  -p \" + str(pdb_ref) + \" -c com.prmtop -r rec.prmtop -l ligand.prmtop -s :WAT:Na+:Cl-:Mg+:K+ -n :LIG --radii \" + str(mbondi)\n",
        "MMPBSA = \"MMPBSA.py -O -i mmpbsa.in -o \" + str(final_mmpbsa) +  \".dat -sp \" + str(pdb_ref) + \" -cp com.prmtop -rp rec.prmtop -lp ligand.prmtop -y \"  + str(traj)\n",
        "mkdir = \"mkdir \" + os.path.join(workDir, fold_MMPBSA)\n",
        "mv = \"mv _MMPBSA* com.prmtop rec.prmtop ligand.prmtop reference.frc mmpbsa.in \" + os.path.join(workDir, fold_MMPBSA)\n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "\n",
        "with open('run_MMPBSA.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(amberhome)\n",
        "    print(ante_MMPBSA)\n",
        "    print(MMPBSA)\n",
        "    print(mkdir)\n",
        "    print(mv)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "\n",
        "!chmod 700 run_MMPBSA.sh 2>&1 1>/dev/null\n",
        "!bash run_MMPBSA.sh 2>&1 1>/dev/null\n",
        "\n",
        "f_mmpbsa = open(final_mmpbsa + '.dat', 'r')\n",
        "file_contents = f_mmpbsa.read()\n",
        "print(file_contents)\n",
        "f_mmpbsa.close()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c_cCDsfF9WJF",
        "cellView": "form"
      },
      "source": [
        "#@title **Interaction Energy**\n",
        "#@markdown **Important:** To quantify the strength of the interaction between the ligand and the protein, we will compute the nonbonded interaction energy between these two species.  It is important to note that this quantity is NOT a free energy or a binding energy.\n",
        "\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = 'Interaction_energy' #@param {type:\"string\"}\n",
        "\n",
        "pt_topology = traj_load.top\n",
        "restraint_array = pt.select_atoms('!(:WAT) & !(:Na+) & !(:Cl-) & !(:Mg+) & !(:K+) & !(:LIG)', pt_topology)\n",
        "first_atom = restraint_array[0]\n",
        "last_atom = restraint_array[-1]\n",
        "mask = \"LIE :LIG @\" + str(first_atom+1) + \"-\" + str(last_atom+1)\n",
        "\n",
        "lie = pt.analysis.energy_analysis.lie(traj_load, mask=mask, options='cutvdw 12.0 cutelec 12.0 diel 2.0', dtype='dict')\n",
        "\n",
        "lie_elec = lie['LIE[EELEC]']\n",
        "lie_vdw = lie['LIE[EVDW]']\n",
        "lie_total = lie_elec + lie_vdw\n",
        "Write_the_trajectory = traj_save_freq\n",
        "time = len(lie_total)*int(Write_the_trajectory)/1000\n",
        "time_array = np.arange(0,time,int(Write_the_trajectory)/1000)*int(stride_traj)\n",
        "\n",
        "def filter_outliers(data):\n",
        "    \"\"\"Return a mask of booleans to filter out outliers.\"\"\"\n",
        "    Q1 = np.percentile(data, 25)\n",
        "    Q3 = np.percentile(data, 75)\n",
        "    IQR = Q3 - Q1\n",
        "    lower_bound = Q1 - 1.5 * IQR\n",
        "    upper_bound = Q3 + 1.5 * IQR\n",
        "    return (data >= lower_bound) & (data <= upper_bound)\n",
        "\n",
        "# Create masks for non-outliers\n",
        "mask_total = filter_outliers(lie_total)\n",
        "mask_elec = filter_outliers(lie_elec)\n",
        "mask_vdw = filter_outliers(lie_vdw)\n",
        "\n",
        "# Filter values based on the mask\n",
        "filtered_time_total = time_array[mask_total]\n",
        "filtered_lie_total = lie_total[mask_total]\n",
        "\n",
        "filtered_time_elec = time_array[mask_elec]\n",
        "filtered_lie_elec = lie_elec[mask_elec]\n",
        "\n",
        "filtered_time_vdw = time_array[mask_vdw]\n",
        "filtered_lie_vdw = lie_vdw[mask_vdw]\n",
        "\n",
        "lie_total_mean = mean(filtered_lie_total)\n",
        "lie_total_stdev = stdev(filtered_lie_total)\n",
        "print(\"Interaction Energy Average = \" + str(\"{:.2f}\".format(lie_total_mean)) + \" \\u00B1 \" + str(\"{:.2f}\".format(lie_total_stdev)) + \" kcal/mol\")\n",
        "\n",
        "# Plotting:\n",
        "plt.plot(filtered_time_total, filtered_lie_total, alpha=0.6, color='blue', linewidth=1.5, label=\"Total Energy\")\n",
        "plt.plot(filtered_time_elec, filtered_lie_elec, alpha=0.6, color='green', linewidth=1.5, label=\"Electrostatic Energy\")\n",
        "plt.plot(filtered_time_vdw, filtered_lie_vdw, alpha=0.6, color='red', linewidth=1.5, label=\"van der Waals Energy\")\n",
        "\n",
        "plt.xlim(0, simulation_ns)\n",
        "# plt.ylim(-50, 0)\n",
        "\n",
        "plt.xlabel(\"Time (ns)\", fontsize=14, fontweight='bold')\n",
        "plt.ylabel('Interaction Energy \\n (kcal/mol)', fontsize=14, fontweight='bold')\n",
        "plt.xticks(fontsize=12)\n",
        "plt.yticks(fontsize=12)\n",
        "plt.legend(frameon=False, loc='center left', bbox_to_anchor=(1, 0.5))\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "lie_eelec = pd.DataFrame(lie['LIE[EELEC]'])\n",
        "lie_eelec.to_csv(os.path.join(workDir, Output_name + \"_eelec.csv\"))\n",
        "lie_evdw = pd.DataFrame(lie['LIE[EVDW]'])\n",
        "lie_evdw.to_csv(os.path.join(workDir, Output_name + \"_evdw.csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kK8YRhrjD6xo",
        "cellView": "form"
      },
      "source": [
        "#@title **Compute distance between the ligand and catalytic site residues**\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = 'distance' #@param {type:\"string\"}\n",
        "#@markdown **Cutoff distance to nearest residues (Angstrons):**\n",
        "Distance = '5' #@param {type:\"string\"}\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('(:LIG<:' + str(Distance) + ')&!(:WAT|:Na+,Cl-,LIG)')\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 = \":LIG :\" + 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,85' #@param {type:\"string\"}\n",
        "\n",
        "mask = \":LIG :\" + 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",
        "\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": []
    }
  ]
}