MD_Small_Molecules.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "MD_Small_Molecules.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "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/MD_Small_Molecules.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",
        "This is a Jupyter notebook for running Molecular Dynamics (MD) simulations using OpenMM engine and GAFF force field for small molecules. Here, partial charges will be calculated using Psi4. Psi4 is an open-source suite of ab initio quantum chemistry programs designed for efficient, high-accuracy simulations of molecular properties. 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",
        "\n",
        "\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 [Psi4](https://psicode.org/) team for developing an excellent and open source suite of ab initio quantum chemistry.\n",
        "\n",
        "- Thank you to Doctor [Souvik Sinha](https://twitter.com/sou_svk) for his help in the code process. \n",
        "\n",
        "- Thank you to [Jonas Boström](https://twitter.com/DrBostrom) for his **Psi4** code. If you use the notebook, please, [cite](https://pubs.acs.org/doi/10.1021/acs.jcim.1c01535).\n",
        "\n",
        "- 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",
      "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 SMILES notation of your molecule.\n",
        "\n",
        "\n",
        "In this notebook, we will calculate the partial charges of Ibuprofen using [Psi4](https://psicode.org/). It takes a SMILES notation and outputs a mol2 file with RESP derived partial charges. Options for setting method (e.g. B3LYP), basis set (e.g. 6-31G**) and singlepoint or geometry optimization are available. The mol2 file of ibuprofen with RESP partial charges will be used to build our 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. We will use LEaP program (https://ambermd.org/tutorials/pengfei/index.php) to build our simulation box. 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",
        "The input SMILES example for Ibuprofen: CC(C)CC1=CC=C(C=C1)C(C)C(=O)O \n",
        "\n",
        "## ---\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "Hlnvt2ocgzT9"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "---\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.    Numpy (https://numpy.org/)\n",
        "6.    Matplotlib (https://matplotlib.org/)\n",
        "7.    AmberTools (https://ambermd.org/AmberTools.php)\n",
        "8.    Psi4 (https://psicode.org/)\n",
        "9.    RDKit (https://www.rdkit.org/)"
      ],
      "metadata": {
        "id": "4qWao_WFi7Y_"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wH1oMiVUlxO_",
        "cellView": "form"
      },
      "source": [
        "#@title **Install dependencies**\n",
        "#@markdown It will take a few minutes, please, drink a coffee and wait. ;-)\n",
        "# install dependencies\n",
        "%%capture\n",
        "import sys\n",
        "!pip -q install py3Dmol 2>&1 1>/dev/null\n",
        "!pip install --upgrade MDAnalysis 2>&1 1>/dev/null\n",
        "!pip install git+https://github.com/pablo-arantes/biopandas 2>&1 1>/dev/null\n",
        "!pip install rdkit-pypi\n",
        "!pip install Cython\n",
        "\n",
        "\n",
        "# install conda\n",
        "!wget -qnc https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \n",
        "!bash Miniconda3-latest-Linux-x86_64.sh -bfp /usr/local 2>&1 1>/dev/null\n",
        "!rm -r Miniconda3-latest-Linux-x86_64.sh\n",
        "!conda install -y -q -c conda-forge openmm=7.6 python=3.7 pdbfixer 2>&1 1>/dev/null\n",
        "!conda install -c conda-forge ambertools --yes 2>&1 1>/dev/null\n",
        "!conda install -c ambermd pytraj  --yes 2>&1 1>/dev/null\n",
        "!conda install -c conda-forge parmed  --yes 2>&1 1>/dev/null\n",
        "!conda install -c conda-forge openff-toolkit --yes 2>&1 1>/dev/null\n",
        "!conda install -c bioconda pybel --yes\n",
        "!conda install -c openbabel openbabel --yes\n",
        "\n",
        "#load dependencies\n",
        "sys.path.append('/usr/local/lib/python3.7/site-packages/')\n",
        "from openmm import app, unit\n",
        "from openmm.app import HBonds, NoCutoff, PDBFile\n",
        "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 openbabel\n",
        "from __future__ import print_function\n",
        "import pytraj as pt\n",
        "import platform\n",
        "import scipy.cluster.hierarchy\n",
        "from scipy.spatial.distance import squareform\n",
        "import scipy.stats as stats\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "from scipy.interpolate import griddata\n",
        "import seaborn as sb\n",
        "from statistics import mean, stdev\n",
        "from pytraj import matrix\n",
        "from matplotlib import colors\n",
        "from IPython.display import set_matplotlib_formats"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "%%bash\n",
        "#@title **Install dependencies for Psi4**\n",
        "conda create -n p4env python=3.9 psi4 resp -c psi4/label/dev --yes\n",
        "source activate p4env\n",
        "pip install rdkit-pypi\n",
        "pip install Cython\n",
        "conda install -c conda-forge parmed  --yes 2>&1 1>/dev/null\n",
        "conda install -c bioconda pybel --yes\n",
        "conda install -c conda-forge openbabel --yes"
      ],
      "metadata": {
        "cellView": "form",
        "id": "i3LGhcnyWG9F"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fDQnAKJLFxtt"
      },
      "source": [
        "## Using Google Drive to store topology 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 topology 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",
        "# **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",
        "Below, you should provide your smiles and the pathway of your Google Drive folder where the notebook will create the folder and save your files.\n",
        "\n",
        "**Please, don't use spaces in smiles and folders names, use _ instead i.e., MyDrive/small_molecule and so on.**"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Please, provide the necessary input files below**: \n",
        "\n",
        "#@title Upload Image\n",
        "from google.colab import files\n",
        "from PIL import Image, ImageOps\n",
        "from rdkit import Chem\n",
        "from rdkit.Chem import AllChem,Draw\n",
        "from rdkit.Chem.Draw import IPythonConsole\n",
        "import IPython.display \n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.image as mpimg\n",
        "IPythonConsole.ipython_useSVG=True \n",
        "#@markdown Type the smiles of your molecule. If you don't know your smiles, please, check at https://pubchem.ncbi.nlm.nih.gov/\n",
        "Smiles = 'CC(C)CC1=CC=C(C=C1)C(C)C(=O)O' #@param {type:\"string\"}\n",
        "\n",
        "Google_Drive_Path = '/content/drive/MyDrive/small_molecule' #@param {type:\"string\"}\n",
        "workDir = Google_Drive_Path\n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "with open('mkdir_path.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(\"mkdir \" + workDir )\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "    \n",
        "!chmod 700 mkdir_path.sh 2>&1 1>/dev/null\n",
        "!bash mkdir_path.sh >/dev/null 2>&1\n",
        "!rm mkdir_path.sh\n",
        "\n",
        "smiles_fig = Chem.MolFromSmiles(Smiles)\n",
        "\n",
        "# Draw.MolToMPL(smiles_fig, size=(200, 200))\n",
        "smi_end = Draw.MolToFile(smiles_fig, size=(600, 600), filename='smiles.png')\n",
        "\n",
        "#print(\"SMILES:\", Smiles)\n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "final_smiles = '/content/smiles.smi'\n",
        "with open(final_smiles, 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(Smiles, ' ')\n",
        "    sys.stdout = original_stdout\n",
        "\n",
        "img = mpimg.imread('smiles.png')\n",
        "plt.figure(figsize = (8,8))\n",
        "imgplot = plt.imshow(img)\n",
        "plt.axis('off')\n",
        "plt.show()\n",
        "\n",
        "# Image('smiles.png')\n",
        "a_file = open('/content/smiles.smi')\n",
        "file_contents = a_file.read()\n",
        "print(\"SMILES:\", file_contents)"
      ],
      "metadata": {
        "id": "9_xezAAaLpOc",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%bash\n",
        "source activate p4env\n",
        "\n",
        "python\n",
        "import os\n",
        "import psi4\n",
        "import resp\n",
        "from openbabel import openbabel as ob\n",
        "from rdkit import Chem\n",
        "from rdkit.Chem import AllChem\n",
        "from rdkit.Chem import AllChem,Draw\n",
        "\n",
        "def neutralize_atoms(mol):\n",
        "    pattern = Chem.MolFromSmarts(\"[+1!h0!$([*]~[-1,-2,-3,-4]),-1!$([*]~[+1,+2,+3,+4])]\")\n",
        "    at_matches = mol.GetSubstructMatches(pattern)\n",
        "    at_matches_list = [y[0] for y in at_matches]\n",
        "    if len(at_matches_list) > 0:\n",
        "        for at_idx in at_matches_list:\n",
        "            atom = mol.GetAtomWithIdx(at_idx)\n",
        "            chg = atom.GetFormalCharge()\n",
        "            hcount = atom.GetTotalNumHs()\n",
        "            atom.SetFormalCharge(0)\n",
        "            atom.SetNumExplicitHs(hcount - chg)\n",
        "            atom.UpdatePropertyCache()\n",
        "    return mol\n",
        "\n",
        "def cleanUp(psi4out_xyz):\n",
        "    deleteTheseFiles = ['1_default_grid.dat','1_default_grid_esp.dat','grid.dat','timer.dat']\n",
        "    deleteTheseFiles.append(psi4out_xyz)\n",
        "    for fileName in deleteTheseFiles:\n",
        "        if os.path.exists(fileName):\n",
        "            os.remove(fileName)\n",
        "\n",
        "def get_xyz_coords(mol):\n",
        "    if not mol is None:\n",
        "        num_atoms = mol.GetNumAtoms()\n",
        "        xyz_string=\"\"\n",
        "        for counter in range(num_atoms):\n",
        "            pos=mol.GetConformer().GetAtomPosition(counter)\n",
        "            xyz_string = xyz_string + (\"%s %12.6f %12.6f %12.6f\\n\" % (mol.GetAtomWithIdx(counter).GetSymbol(), pos.x, pos.y, pos.z) )\n",
        "    return xyz_string\n",
        "\n",
        "\n",
        "def calcRESPCharges(mol, basisSet, method, gridPsi4 = 1):\n",
        "    options = {'BASIS_ESP': basisSet,\n",
        "               'METHOD_ESP': method,\n",
        "               'RESP_A': 0.0005,\n",
        "               'RESP_B': 0.1,\n",
        "               'VDW_SCALE_FACTORS':[1.4, 1.6, 1.8, 2.0],\n",
        "               'VDW_POINT_DENSITY':int(gridPsi4)\n",
        "    }\n",
        "\n",
        "    # resp_charges = resp.resp([mol], [options])[0][1]\n",
        "    resp_charges = resp.resp([mol], options)\n",
        "    \n",
        "    return resp_charges\n",
        "\n",
        "#@title **Parameters to generate partial charges:**\n",
        "\n",
        "\n",
        "method = \"B3LYP\" #@param [\"B3LYP\", \"MP2\", \"HF\"]\n",
        "\n",
        "basisSet = \"6-31G**\" #@param [\"STO-3g\", \"3-21G\", \"6-31G*\", \"6-31G**\", \"6-31+G(d,p)\", \"6-311++G(2d,p)\"]\n",
        "\n",
        "Neutralize = \"Yes\" #@param [\"Yes\", \"No\"]\n",
        "\n",
        "if Neutralize == \"Yes\":\n",
        "  neutralize = True\n",
        "else:\n",
        "  neutralize = False\n",
        "\n",
        "Calculation_Type = \"Single Point\" #@param [\"Single Point\", \"Geometry Optimization\"]\n",
        "\n",
        "if Calculation_Type == \"Single Point\":\n",
        "  singlePoint = True\n",
        "else:\n",
        "  singlePoint = False  \n",
        "\n",
        "\n",
        "neutralize = True\n",
        "psi4.set_memory('8 GB')\n",
        "obConversion = ob.OBConversion()\n",
        "obConversion.SetInAndOutFormats(\"xyz\", \"mol2\")\n",
        "# singlePoint = True\n",
        "path = '/content'\n",
        "\n",
        "# inputFile = '/content/smiles.smi'\n",
        "# molList = Chem.SDMolSupplier(inputFile, removeHs=False)\n",
        "\n",
        "SMILESasInput = True\n",
        "\n",
        "if SMILESasInput:\n",
        "    molList = []\n",
        "    inputFile = '/content/smiles.smi'\n",
        "    suppl = Chem.SmilesMolSupplier(inputFile, titleLine = False)\n",
        "\n",
        "    for mol in suppl:\n",
        "        mol = Chem.AddHs(mol)\n",
        "        AllChem.EmbedMolecule(mol)\n",
        "        try: \n",
        "            AllChem.MMFFOptimizeMolecule(mol)\n",
        "        except: \n",
        "            AllChem.UFFOptimizeMolecule(mol)            \n",
        "        molList.append(mol)\n",
        "\n",
        "for mol in molList:\n",
        "    print(mol)\n",
        "    if not mol is None:\n",
        "\n",
        "        molId = mol.GetProp(\"_Name\")\n",
        "        print('Trying:', molId)\n",
        "        \n",
        "        if neutralize:\n",
        "            mol = neutralize_atoms(mol)\n",
        "            mol = Chem.AddHs(mol)\n",
        "\n",
        "        xyz_string = get_xyz_coords(mol) \n",
        "        psi_mol = psi4.geometry(xyz_string)\n",
        "\n",
        "        ### single point calculation\n",
        "        outfile_mol2 = inputFile[:-4]+\".mol2\"\n",
        "\n",
        "        if singlePoint:\n",
        "            print('Running singlepoint...')\n",
        "            resp_charges = calcRESPCharges(psi_mol, basisSet, method, gridPsi4 = 1)\n",
        "\n",
        "        else:\n",
        "            print('Running geometry optimization...')            \n",
        "            methodNbasisSet = method+\"/\"+basisSet\n",
        "            psi4.optimize(methodNbasisSet, molecule=psi_mol)\n",
        "            resp_charges = calcRESPCharges(psi_mol, basisSet, method, gridPsi4 = 1)            \n",
        "\n",
        "        ### save coords to xyz file \n",
        "        psi4out_xyz = molId + '.xyz'\n",
        "        psi_mol.save_xyz_file(psi4out_xyz,1)\n",
        "\n",
        "\n",
        "        ### read xyz file and write as mol2 \n",
        "        ob_mol = ob.OBMol()\n",
        "        obConversion.ReadFile(ob_mol, psi4out_xyz)\n",
        "\n",
        "         ### write as mol2 \n",
        "        outfile_mol2 = path+\"/\"+molId+\"_partialChgs.mol2\"\n",
        "        obConversion.WriteFile(ob_mol, outfile_mol2)\n",
        "\n",
        "        ### set new partial charges\n",
        "        count = 0\n",
        "        newChg_temp = resp_charges[1]\n",
        "        print(\"RESP Charges: \", newChg_temp)\n",
        "        for atom in ob.OBMolAtomIter(ob_mol):\n",
        "            newChg = newChg_temp[count]\n",
        "            atom.SetPartialCharge(newChg)\n",
        "            count += 1\n",
        "\n",
        "        ### write as mol2 \n",
        "        outfile_mol2 = path+\"/\"+molId+\"_partialChgs.mol2\"\n",
        "        outfile_pdb = path+\"/\"+molId+\".pdb\"\n",
        "        print(\"Finished. Saved compound with partial charges as mol2 file: %s\" % outfile_mol2)\n",
        "        obConversion.WriteFile(ob_mol, outfile_mol2)\n",
        "        ## clean up\n",
        "        cleanUp(psi4out_xyz)\n",
        "\n",
        "#draw_with_charges\n",
        "inputFile2 = '/content/smiles.smi'\n",
        "suppl2 = Chem.SmilesMolSupplier(inputFile2, titleLine = False)\n",
        "for mol_2 in suppl2:\n",
        "    m2 = mol_2\n",
        "\n",
        "for at, i in zip(m2.GetAtoms(), newChg_temp):\n",
        "    lbl = '%.2f'%(i)\n",
        "    at.SetProp('atomNote',lbl)\n",
        "    Draw.MolToFile(m2, size=(600, 600), filename='smiles_charges.png')"
      ],
      "metadata": {
        "id": "m9hziHETL1ZS",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import sys\n",
        "import os\n",
        "import IPython\n",
        "from IPython.display import Image\n",
        "# from rdkit.Chem.Draw import IPythonConsole\n",
        "\n",
        "#@title **Check your charges**: \n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.image as mpimg\n",
        "img = mpimg.imread('smiles_charges.png')\n",
        "plt.figure(figsize = (8,8))\n",
        "imgplot = plt.imshow(img)\n",
        "plt.axis('off')\n",
        "plt.show()\n",
        "\n",
        "# Image('smiles.png')\n",
        "a_file = open('/content/smiles.smi')\n",
        "file_contents = a_file.read()\n",
        "print(\"SMILES:\", file_contents)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "mdiWpBBzMwR6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "od6pJ_TiLLUo",
        "cellView": "form"
      },
      "source": [
        "from rdkit import Chem\n",
        "from rdkit.Chem import PandasTools\n",
        "import openbabel\n",
        "\n",
        "# from openff.toolkit.typing.engines.smirnoff import ForceField\n",
        "# import parmed\n",
        "\n",
        "#@title **Parameters to generate the molecule topology on GAFF2 force field:**\n",
        "\n",
        "\n",
        "Force_field = \"ff19SB\" \n",
        "ff = \"leaprc.protein.ff19SB\"\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",
        "# elif Water_type == \"OPC\":\n",
        "  # water = \"leaprc.water.tip3p\"\n",
        "  # water_box =\"set SYS box {50,50,50}\"\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",
        "\n",
        "Ligand_Force_field = \"GAFF2\" \n",
        "\n",
        "Ligand_net_charges = \"0\" #@param {type:\"string\", min:-10, max:10, step:1}\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 = \"./_partialChgs.mol2\"\n",
        "ligand_pdb = os.path.join(workDir, \"_partialChgs.pdb\")\n",
        "ligand_new = \"./ligand.mol2\"\n",
        "ligand_mol2 = \"./new.mol2\"\n",
        "ligand_frcmod = os.path.join(workDir, \"ligand.frcmod\")\n",
        "lig_new = os.path.join(workDir, \"ligand_gaff.pdb\")\n",
        "lib = os.path.join(workDir, \"lig.lib\")\n",
        "\n",
        "obConversion = openbabel.OBConversion()\n",
        "obConversion.SetInAndOutFormats(\"mol2\", \"pdb\")\n",
        "\n",
        "mol = openbabel.OBMol()\n",
        "obConversion.ReadFile(mol, \"./_partialChgs.mol2\")   # Open Babel will uncompress automatically\n",
        "\n",
        "# mol.AddHydrogens()\n",
        "# print (mol.NumAtoms())\n",
        "# print (mol.NumBonds())\n",
        "# print (mol.NumResidues())\n",
        "obConversion.WriteFile(mol, ligand_pdb)\n",
        "\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_pdb) + \" -fi pdb -o \" + str(ligand_new) + \" -fo mol2 -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",
        "\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_command3)\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",
        "!cut -b -72 ligand.mol2 > test_ligand.mol2\n",
        "!awk '{print $9}' _partialChgs.mol2 | sed '1i\\\\'> 9th.txt\n",
        "!paste -d \" \" test_ligand.mol2 9th.txt > new.mol2\n",
        "!rm test_ligand.mol2 9th.txt\n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "\n",
        "with open('gaff2.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(gaff_command4)\n",
        "    sys.stdout = original_stdout\n",
        "\n",
        "!chmod 700 gaff2.sh 2>&1 1>/dev/null\n",
        "!bash gaff2.sh >/dev/null 2>&1\n",
        "\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 = loadmol2 \"\"\" + str(ligand_mol2) + \"\\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 = loadmol2 \"\"\" + str(ligand_mol2) + \"\\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",
        "pdb_amber = os.path.exists(pdb)\n",
        "top_amber = os.path.exists(top)\n",
        "crd_amber = os.path.exists(crd)\n",
        "\n",
        "!zip -FSr \"partial_charges.zip\" *\".smi\" *\".mol2\" *\".out\" *\".png\" >/dev/null 2>&1\n",
        "\n",
        "cp_sys = \"cp \" \"partial_charges.zip \" + workDir \n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "\n",
        "with open('cp_sys.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(cp_sys)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "\n",
        "!chmod 700 cp_sys.sh 2>&1 1>/dev/null\n",
        "!bash cp_sys.sh 2> /dev/null\n",
        "\n",
        "zip_end = os.path.join(workDir, \"partial_charges.zip\")\n",
        "\n",
        "zip_true = os.path.exists(zip_end)\n",
        "\n",
        "if pdb_amber == True and top_amber == True and crd_amber == True and zip_true == 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 ipywidgets\n",
        "from ipywidgets import interact, fixed\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "\n",
        "def show_pdb(show_box=True,\n",
        "             show_molecule=True,\n",
        "             color=\"None\"):\n",
        "\n",
        "  def ligand(p, model=0):\n",
        "    HP = ['LIG']\n",
        "    p.addStyle({\"model\":model,'and':[{'resn':HP}]},\n",
        "                       {'stick':{'colorscheme':'greenCarbon','radius':0.3}})\n",
        "    p.setViewStyle({'style':'outline','color':'black','width':0.1})\n",
        "\n",
        "\n",
        "  def box(p, model=0):\n",
        "      p.addModelsAsFrames(pdb)\n",
        "      p.addSurface(py3Dmol.SAS, {'opacity': 0.6, 'color':'white'}) #comment this line if you dont want to see the water box\n",
        "      p.setViewStyle({'style':'outline','color':'black','width':0.1})\n",
        "\n",
        "\n",
        "  p = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js')\n",
        "  p.addModel(open(pdb,'r').read(),'pdb')\n",
        "  if color == \"rainbow\":\n",
        "    p.setStyle({'cartoon': {'color':'spectrum'}})\n",
        "  else:\n",
        "    p.setStyle({'cartoon':{}})\n",
        "\n",
        "  if show_molecule: ligand(p)\n",
        "  if show_box: box(p)\n",
        "  p.zoomTo()\n",
        "  return p.show()\n",
        "\n",
        "interact(show_pdb,\n",
        "         show_box=ipywidgets.Checkbox(value=True),\n",
        "         show_molecule=ipywidgets.Checkbox(value=True),\n",
        "         color=ipywidgets.Dropdown(options=['None', 'rainbow'], value='None'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "---\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 molecule 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."
      ],
      "metadata": {
        "id": "u5hxFojPOHsG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ### **Parameters for MD Equilibration protocol:**\n",
        "\n",
        "# remove whitespaces\n",
        "Jobname = 'eq' #@param {type:\"string\"}\n",
        "\n",
        "Ligand_Force_field = \"GAFF2\"\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 = 800 #@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"
      ],
      "metadata": {
        "cellView": "form",
        "id": "7QWp-XWaOUbE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Runs an Equilibration MD simulation (NPT ensemble)**\n",
        "#@markdown Now, let's equilibrate our system!\n",
        "\n",
        "###########################################\n",
        "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\t\t\t\t\t\n",
        "temperature = float(temperature_eq)*kelvin\n",
        "savcrd_freq = int(write_the_trajectory_eq)*picosecond\n",
        "print_freq  = int(write_the_log_eq)*picosecond\n",
        "\n",
        "pressure\t= float(pressure_eq)*bar\n",
        "\n",
        "restraint_fc = int(Force_constant) # kJ/mol\n",
        "\n",
        "nsteps  = int(simulation_time.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nprint  = int(print_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nsavcrd = int(savcrd_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "\n",
        "#############################################\n",
        "# Defining functions to use below:\n",
        "def backup_old_log(pattern, string):\n",
        "\tresult = []\n",
        "\tfor root, dirs, files in os.walk(\"./\"):\n",
        "\t\tfor name in files:\n",
        "\t\t\tif fnmatch.fnmatch(name, pattern):\n",
        "\n",
        "\t\t\t\ttry:\n",
        "\t\t\t\t\tnumber = int(name[-2])\n",
        "\t\t\t\t\tavail = isinstance(number, int)\n",
        "\t\t\t\t\t#print(name,avail)\n",
        "\t\t\t\t\tif avail == True:\n",
        "\t\t\t\t\t\tresult.append(number)\n",
        "\t\t\t\texcept:\n",
        "\t\t\t\t\tpass\n",
        "\n",
        "\tif len(result) > 0:\n",
        "\t\tmaxnumber = max(result)\n",
        "\telse:\n",
        "\t\tmaxnumber = 0\n",
        "\n",
        "\tbackup_file = \"\\#\" + string + \".\" + str(maxnumber + 1) + \"#\"\n",
        "\tos.system(\"mv \" + string + \" \" + backup_file)\n",
        "\treturn backup_file\n",
        "\n",
        "def restraints(system, crd, fc, restraint_array):\n",
        "\n",
        "\tboxlx = system.getDefaultPeriodicBoxVectors()[0][0].value_in_unit(nanometers)\n",
        "\tboxly = system.getDefaultPeriodicBoxVectors()[1][1].value_in_unit(nanometers)\n",
        "\tboxlz = system.getDefaultPeriodicBoxVectors()[2][2].value_in_unit(nanometers)\n",
        "\n",
        "\tif fc > 0:\n",
        "\t\t# positional restraints for all heavy-atoms\n",
        "\t\tposresPROT = CustomExternalForce('k*periodicdistance(x, y, z, x0, y0, z0)^2;')\n",
        "\t\tposresPROT.addPerParticleParameter('k')\n",
        "\t\tposresPROT.addPerParticleParameter('x0')\n",
        "\t\tposresPROT.addPerParticleParameter('y0')\n",
        "\t\tposresPROT.addPerParticleParameter('z0')\n",
        "  \n",
        "\t\tfor atom1 in restraint_array:\n",
        "\t\t\tatom1 = int(atom1)\n",
        "               \n",
        "\t\t\txpos  = crd.positions[atom1].value_in_unit(nanometers)[0]\n",
        "\t\t\typos  = crd.positions[atom1].value_in_unit(nanometers)[1]\n",
        "\t\t\tzpos  = crd.positions[atom1].value_in_unit(nanometers)[2]\n",
        "\n",
        "\t\t\tposresPROT.addParticle(atom1, [fc, xpos, ypos, zpos])\n",
        "    \n",
        "\t\tsystem.addForce(posresPROT)\n",
        "\n",
        "\treturn system\n",
        "##############################################\n",
        "\n",
        "#############################################\n",
        "print(\"\\n> Simulation details:\\n\")\n",
        "print(\"\\tJob name = \" + jobname)\n",
        "print(\"\\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",
        "\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\")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "UDlkyvJPOnp4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "---\n",
        "# **Running a Production MD simulation**\n",
        "\n",
        "Finally, we will proceed with the Production simulation itself using the equilibrated system coordinates as input structure.\n",
        "\n",
        "Note that we will use here a *.rst state file* , which contains atomic velocities and positions from the last frame of the equilibration simulation, guaranteeing that our production simulation begins from a thermodynamically equilibrated system.\n",
        "\n",
        "Another important information here is the **Number_of_strides** and the **Stride_Time**. In this notebook, we simulate a defined number of *strides*, so the **simulation time = Number_of_strides*Stride_Time**. For example, we can simulate 100ns by setting *Number_of_strides=10* and *Stride_Time=10 ns*.\n",
        "\n",
        "**Important: at the end of the Production simulation, we concatenate all strides to create a complete trajectory file which can be visualized and analyzed**\n",
        "\n",
        "The idea behind this approach is to make use of the intermitent 12h/24h period that Google Colab allows us to use its GPUs. "
      ],
      "metadata": {
        "id": "UANxGObdPDBh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ### **Provide input file names below:** \n",
        "\n",
        "Equilibrated_PDB = 'eq.pdb' #@param {type:\"string\"}\n",
        "State_file = 'eq.rst' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown ### **Parameters for MD Production protocol:**\n",
        "\n",
        "\n",
        "# remove whitespaces\n",
        "Jobname = 'prod' #@param {type:\"string\"}\n",
        "\n",
        "Ligand_Force_field = \"GAFF2\"\n",
        "\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 ---"
      ],
      "metadata": {
        "cellView": "form",
        "id": "NW4TAla8PKGl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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\t\t\t\t\t\n",
        "temperature = float(temperature_prod)*kelvin\n",
        "savcrd_freq = int(write_the_trajectory_prod)*picosecond\n",
        "print_freq  = int(write_the_log_prod)*picosecond\n",
        "\n",
        "pressure\t= float(pressure_prod)*bar\n",
        "\n",
        "simulation_time = stride_time*nstride\n",
        "nsteps  = int(stride_time.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nprint  = int(print_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nsavcrd = int(savcrd_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "firststride = 1 # must be integer\n",
        "#############################################\n",
        "# Defining functions to use below:\n",
        "def backup_old_log(pattern, string):\n",
        "\tresult = []\n",
        "\tfor root, dirs, files in os.walk(\"./\"):\n",
        "\t\tfor name in files:\n",
        "\t\t\tif fnmatch.fnmatch(name, pattern):\n",
        "\n",
        "\t\t\t\ttry:\n",
        "\t\t\t\t\tnumber = int(name[-2])\n",
        "\t\t\t\t\tavail = isinstance(number, int)\n",
        "\t\t\t\t\t#print(name,avail)\n",
        "\t\t\t\t\tif avail == True:\n",
        "\t\t\t\t\t\tresult.append(number)\n",
        "\t\t\t\texcept:\n",
        "\t\t\t\t\tpass\n",
        "\n",
        "\tif len(result) > 0:\n",
        "\t\tmaxnumber = max(result)\n",
        "\telse:\n",
        "\t\tmaxnumber = 0\n",
        "\n",
        "\tbackup_file = \"\\#\" + string + \".\" + str(maxnumber + 1) + \"#\"\n",
        "\tos.system(\"mv \" + string + \" \" + backup_file)\n",
        "\treturn backup_file\n",
        "##############################################\n",
        "\n",
        "#############################################\n",
        "print(\"\\n> Simulation details:\\n\")\n",
        "print(\"\\tJob name = \" + jobname)\n",
        "print(\"\\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",
        "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\")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "KsBkjxcOPbeB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Concatenate and align the trajectory**\n",
        "Skip = \"1\" #@param [\"1\", \"2\", \"5\", \"10\", \"20\", \"50\"]\n",
        "stride_traj = Skip\n",
        "Output_format = \"dcd\" #@param [\"dcd\", \"pdb\", \"trr\", \"xtc\"]\n",
        "\n",
        "#@markdown **Attention:** A high number of frames can explode the memory on Colab. You should be fine with 5000 frames or less.\n",
        "\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(Write_the_trajectory)\n",
        "number_frames_analysis = number_frames/int(stride_traj)\n",
        "\n",
        "\n",
        "traj_end = os.path.join(workDir, str(Jobname) + \"_all.dcd\")\n",
        "traj_end2 = os.path.join(workDir, str(Jobname) + \"_all.\" + str(Output_format))\n",
        "template =  os.path.join(workDir, str(Jobname) + '_%s.dcd')\n",
        "\n",
        "flist = [template % str(i) for i in range(1, nstride + 1)]\n",
        "#print(flist)\n",
        "\n",
        "trajlist = pt.load(flist, pdb, stride=stride_traj)\n",
        "traj_image = trajlist.iterframe(autoimage=True, rmsfit=0)\n",
        "traj_write = pt.write_traj(traj_end, traj_image, overwrite=True)\n",
        "traj_load = pt.load(traj_end, pdb)\n",
        "traj_align = pt.align(traj_load, mask=\"@LIG\", ref=0) \n",
        "traj_write = pt.write_traj(traj_end, traj_align, overwrite=True, options='dcd')\n",
        "traj_write = pt.write_traj(traj_end2, traj_align, overwrite=True, options=Output_format)\n",
        "traj_load = pt.load(traj_end, os.path.join(workDir, \"SYS_gaff2.prmtop\"))\n",
        "print(traj_load)\n",
        "\n",
        "pt_topology = traj_load.top\n",
        "restraint_array = pt.select_atoms('!(:WAT) & !(:Na+) & !(:Cl-) & !(:Mg+) & !(:K+)', pt_topology)\n",
        "first_atom = restraint_array[0]\n",
        "last_atom = restraint_array[-1]\n",
        "mask_end = \"@\" + str(first_atom+1) + \"-\" + str(last_atom+1)\n",
        "\n",
        "traj_end_check = os.path.exists(traj_end2)\n",
        "\n",
        "if traj_end_check == True:\n",
        "  print(\"Trajectory concatenated successfully! :-)\")\n",
        "else:\n",
        "  print(\"ERROR: Check your inputs! \")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "-gTGkJG8P52I"
      },
      "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",
        "!rm *.pdb 2> /dev/null\n",
        "\n",
        "#py3dmol functions\n",
        "class Atom(dict):\n",
        "  def __init__(self, line):\n",
        "    self[\"type\"] = line[0:6].strip()\n",
        "    self[\"idx\"] = line[6:11].strip()\n",
        "    self[\"name\"] = line[12:16].strip()\n",
        "    self[\"resname\"] = line[17:20].strip()\n",
        "    self[\"resid\"] = int(int(line[22:26]))\n",
        "    self[\"x\"] = float(line[30:38])\n",
        "    self[\"y\"] = float(line[38:46])\n",
        "    self[\"z\"] = float(line[46:54])\n",
        "    self[\"sym\"] = line[76:78].strip()\n",
        "\n",
        "  def __str__(self):\n",
        "    line = list(\" \" * 80)\n",
        "    line[0:6] = self[\"type\"].ljust(6)\n",
        "    line[6:11] = self[\"idx\"].ljust(5)\n",
        "    line[12:16] = self[\"name\"].ljust(4)\n",
        "    line[17:20] = self[\"resname\"].ljust(3)\n",
        "    line[22:26] = str(self[\"resid\"]).ljust(4)\n",
        "    line[30:38] = str(self[\"x\"]).rjust(8)\n",
        "    line[38:46] = str(self[\"y\"]).rjust(8)\n",
        "    line[46:54] = str(self[\"z\"]).rjust(8)\n",
        "    line[76:78] = self[\"sym\"].rjust(2)\n",
        "    return \"\".join(line) + \"\\n\"\n",
        "        \n",
        "class Molecule(list):\n",
        "  def __init__(self, file):\n",
        "    for line in file:\n",
        "      if \"ATOM\" in line or \"HETATM\" in line:\n",
        "        self.append(Atom(line))\n",
        "            \n",
        "    def __str__(self):\n",
        "      outstr = \"\"\n",
        "      for at in self:\n",
        "        outstr += str(at)\n",
        "      return outstr\n",
        "\n",
        "if number_frames_analysis > 10:\n",
        "  stride_animation = number_frames_analysis/10\n",
        "else:\n",
        "  stride_animation = 1\n",
        "\n",
        "u = mda.Universe(pdb, traj_end) \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": "A8wxMnRIQEIS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Emh0vU5UjgB6"
      },
      "source": [
        "---\n",
        "---\n",
        "# **Analysis**\n",
        "\n",
        "Although visualizing your trajectory can be quite useful, sometimes you also want more quantitative data.\n",
        "\n",
        "Analyses of MD trajectories vary a lot and we do not intend to cover it all here. However, one can make use of MDanalysis or PyTraj to easily analyze simulations. \n",
        "\n",
        "Below, you can find a few examples of code snippets that can help you to shed some light on your simulation behavior."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from rdkit import Chem\n",
        "from rdkit.Chem.Draw import IPythonConsole\n",
        "from rdkit.Chem import Draw\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.image as mpimg\n",
        "IPythonConsole.ipython_useSVG=True \n",
        "\n",
        "#@title **Please, run the present cell to check your atom numbers, you will need it for further analysis**: \n",
        "\n",
        "def mol_with_atom_index(mol):\n",
        "    for atom in mol.GetAtoms():\n",
        "        atom.SetAtomMapNum(atom.GetIdx()+1)\n",
        "    return mol\n",
        "\n",
        "pdb_nw = os.path.join(workDir, \"SYS_nw.pdb\")\n",
        "mol = Chem.MolFromPDBFile(pdb_nw)\n",
        "\n",
        "# With atom index\n",
        "mol2 = mol_with_atom_index(mol)\n",
        "Draw.MolToFile(mol2, size=(600, 600), filename='atom_numbers.png')\n",
        "\n",
        "img = mpimg.imread('atom_numbers.png')\n",
        "plt.figure(figsize = (8,8))\n",
        "imgplot = plt.imshow(img)\n",
        "plt.axis('off')\n",
        "plt.show()\n",
        "\n"
      ],
      "metadata": {
        "id": "O9AlazCBXIY1",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "mVOogip3h08Z"
      },
      "source": [
        "#@title **Compute distance between the atoms**\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'distance_atoms' #@param {type:\"string\"}\n",
        "#@markdown **Type the number of atoms separated by commas and without spaces (1,2,3...):**\n",
        "Selection_1 = '3' #@param {type:\"string\"}\n",
        "Selection_2 = '12' #@param {type:\"string\"}\n",
        "\n",
        "mask = \"@\" + str(Selection_1) + \" @\" + str(Selection_2)\n",
        "dist = pt.distance(traj_load, mask)\n",
        "print(\"Selected atoms = \" + Selection_1 + \", \"  + Selection_2 + \"\\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",
        "df = pd.Series(dist)\n",
        "running_aver = df.rolling(window =10).mean()\n",
        "\n",
        "# Plotting:\n",
        "ax = plt.plot(time_array, dist, alpha=0.2, color = 'magenta', linewidth = 1.0)\n",
        "ax = plt.plot(time_array, running_aver, alpha=1, color = 'magenta', linewidth = 1.0)\n",
        "\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",
      "source": [
        "#@title **Plot distance as a distribution**\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'distance_dist' #@param {type:\"string\"}\n",
        "\n",
        "ax = sb.kdeplot(dist, color=\"magenta\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('Distance [$\\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')"
      ],
      "metadata": {
        "cellView": "form",
        "id": "msdL9GHxUY27"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Compute specific dihedral angles**\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'dihedral_angle' #@param {type:\"string\"}\n",
        "#@markdown **Type the number of atoms:**\n",
        "Atom_1 = '1' #@param {type:\"string\"}\n",
        "Atom_2 = '2' #@param {type:\"string\"}\n",
        "Atom_3 = '4' #@param {type:\"string\"}\n",
        "Atom_4 = '5' #@param {type:\"string\"}\n",
        "\n",
        "mask = \"@\" + str(Atom_1) + \" @\" + str(Atom_2) + \" @\" + str(Atom_3) + \" @\" + str(Atom_4)\n",
        "dih = pt.dihedral(traj_load, mask)\n",
        "print(\"Selected atoms = \" + Atom_1 + \", \" + Atom_2 + \", \" + Atom_3 + \", \" + Atom_4 + \"\\n\")\n",
        "dih_mean = mean(dih)\n",
        "dih_stdev = stdev(dih)\n",
        "print(\"Dihedral Angle Average = \" + str(\"{:.2f}\".format(dih_mean)) + \" \\u00B1 \" + str(\"{:.2f}\".format(dih_stdev)) + \"°\")\n",
        "\n",
        "time = len(dih)*int(Write_the_trajectory)/1000\n",
        "time_array = np.arange(0,time,int(Write_the_trajectory)/1000)*int(stride_traj)\n",
        "\n",
        "df = pd.Series(dih)\n",
        "running_aver = df.rolling(window =10).mean()\n",
        "\n",
        "# Plotting:\n",
        "ax = plt.plot(time_array, dih, alpha=0.2, color = 'orange', linewidth = 1.0)\n",
        "ax = plt.plot(time_array, running_aver, alpha=1, color = 'orange', linewidth = 1.0)\n",
        "\n",
        "plt.xlim(0, simulation_ns)\n",
        "#plt.ylim(2, 6)\n",
        "\n",
        "plt.xlabel(\"Time (ns)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel(\"Dihedral Angle ($^\\circ$)\", 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(dih)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "metadata": {
        "cellView": "form",
        "id": "tGaOG8j4R3iB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Plot dihedral angle as a distribution**\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'dihedral_dist' #@param {type:\"string\"}\n",
        "\n",
        "ax = sb.kdeplot(dih, color=\"orange\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('Dihedral Angle ($^\\circ$)', 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')"
      ],
      "metadata": {
        "cellView": "form",
        "id": "oTc58aulUuzI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Compute intramolecular hydrogen bonds**\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'hbond_intra' #@param {type:\"string\"}\n",
        "mask_hbond = 'acceptormask ' + mask_end + ' donormask ' + mask_end\n",
        "hb = pt.hbond(traj_load, options=mask_hbond, dtype='dict', distance =3.5)\n",
        "\n",
        "hb_end = hb['total_solute_hbonds']\n",
        "time = len(hb_end)*int(Write_the_trajectory)/1000\n",
        "time_array = np.arange(0,time,int(Write_the_trajectory)/1000)*int(stride_traj)\n",
        "hb.keys()\n",
        "# print(hb_end)\n",
        "# hb_end_mean = mean(hb_end)\n",
        "# hb_end_stdev = stdev(hb_end)\n",
        "# print(\"Hydrogen Bonds Average = \" + str(\"{:.2f}\".format(hb_end_mean)) + \" \\u00B1 \" + str(\"{:.2f}\".format(hb_end_stdev)))\n",
        "\n",
        "df = pd.Series(hb_end)\n",
        "running_aver = df.rolling(window =10).mean()\n",
        "\n",
        "# Plotting:\n",
        "ax = plt.plot(time_array, hb_end, alpha=0.2, color = 'black', linewidth = 1.0)\n",
        "ax = plt.plot(time_array, running_aver, alpha=1, color = 'black', linewidth = 1.0)\n",
        "\n",
        "plt.xlim(0, simulation_ns)\n",
        "\n",
        "plt.xlabel(\"Time (ns)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel(\"Hydrogen Bonds\", 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(hb_end)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "metadata": {
        "cellView": "form",
        "id": "BQ8jzMGaWjb_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Compute hydrogen bonds between your molecule and waters**\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'hbond_solv' #@param {type:\"string\"}\n",
        "mask_hbond1 = 'acceptormask ' + mask_end + ' donormask :WAT'\n",
        "mask_hbond2 = 'acceptormask :WAT' + ' donormask ' + mask_end\n",
        "\n",
        "hb1 = pt.hbond(traj_load, options=mask_hbond1, dtype='dict', distance =3.5)\n",
        "hb2 = pt.hbond(traj_load, options=mask_hbond2, dtype='dict', distance =3.5)\n",
        "\n",
        "hb_end1 = hb1['total_solute_hbonds']\n",
        "hb_end2 = hb2['total_solute_hbonds']\n",
        "hb_total = hb_end1 + hb_end2\n",
        "time = len(hb_total)*int(Write_the_trajectory)/1000\n",
        "time_array = np.arange(0,time,int(Write_the_trajectory)/1000)*int(stride_traj)\n",
        "# hb.keys()\n",
        "# print(hb_total)\n",
        "# hb_end_mean = mean(hb_total)\n",
        "# hb_end_stdev = stdev(hb_total)\n",
        "# print(\"Hydrogen Bonds Average = \" + str(\"{:.2f}\".format(hb_end_mean)) + \" \\u00B1 \" + str(\"{:.2f}\".format(hb_end_stdev)))\n",
        "\n",
        "# running_averages_hb = moving_average(hb_total,4)\n",
        "df = pd.Series(hb_total)\n",
        "running_aver = df.rolling(window =10).mean()\n",
        "# Plotting:\n",
        "ax = plt.plot(time_array, df, alpha=0.2, color = 'black', linewidth = 1.0)\n",
        "ax = plt.plot(time_array, running_aver, alpha=1, color = 'black', linewidth = 1.0)\n",
        "\n",
        "plt.xlim(0, simulation_ns)\n",
        "\n",
        "plt.xlabel(\"Time (ns)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel(\"Hydrogen Bonds\", 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(hb_total)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "metadata": {
        "id": "-nwOAerkYgvP",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Compute solvent-accessible surface area (SASA) of your molecule**\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'sasa' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "sasa = pt.molsurf(traj_load, mask_end)\n",
        "\n",
        "time = len(sasa)*int(Write_the_trajectory)/1000\n",
        "time_array = np.arange(0,time,int(Write_the_trajectory)/1000)*int(stride_traj)\n",
        "\n",
        "df = pd.Series(sasa)\n",
        "running_aver = df.rolling(window =10).mean()\n",
        "\n",
        "# Plotting:\n",
        "ax = plt.plot(time_array, sasa, alpha=0.2, color = 'darkcyan', linewidth = 1.0)\n",
        "ax = plt.plot(time_array, running_aver, alpha=1, color = 'darkcyan', linewidth = 1.0)\n",
        "\n",
        "plt.xlim(0, simulation_ns)\n",
        "#plt.ylim(2, 6)\n",
        "\n",
        "plt.xlabel(\"Time (ns)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel(\"SASA ($\\AA^{2}$)\", 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(sasa)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "metadata": {
        "id": "wAMV7IEoxJMQ",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Plot solvent-accessible surface area (SASA) as a distribution**\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'sasa_dist' #@param {type:\"string\"}\n",
        "\n",
        "ax = sb.kdeplot(sasa, color=\"darkcyan\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('SASA ($\\AA^{2}$)', 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')"
      ],
      "metadata": {
        "cellView": "form",
        "id": "vyTBTatEpn3Y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wBrBMF4Puyv6",
        "cellView": "form"
      },
      "source": [
        "#@title **Compute RMSD of your molecule**\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'rmsd' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "rmsd = pt.rmsd(traj_load, ref = 0, mask = mask_end)\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",
        "df = pd.Series(rmsd)\n",
        "running_aver = df.rolling(window =10).mean()\n",
        "\n",
        "# Plotting:\n",
        "ax = plt.plot(time_array, rmsd, alpha=0.2, color = 'blue', linewidth = 1.0)\n",
        "ax = plt.plot(time_array, running_aver, alpha=1, color = 'blue', linewidth = 1.0)\n",
        "\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 distribution**\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 your molecule**\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'radius_gyration' #@param {type:\"string\"}\n",
        "\n",
        "radgyr = pt.radgyr(traj_load, mask = mask_end)\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",
        "df = pd.Series(radgyr)\n",
        "running_aver = df.rolling(window =10).mean()\n",
        "\n",
        "# Plotting:\n",
        "plt.plot(time_array, radgyr, alpha=0.2, color = 'green', linewidth = 1.0)\n",
        "plt.plot(time_array, running_aver, alpha=1, color = 'green', linewidth = 1.0)\n",
        "\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 distribution**\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**\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'rmsf' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "rmsf = pt.rmsf(traj_load, mask_end)\n",
        "bfactor = pt.bfactors(traj_load, byres=False)\n",
        "\n",
        "# Plotting:\n",
        "plt.plot(rmsf[:,1], alpha=1.0, color = 'red', linewidth = 1.0)\n",
        "\n",
        "plt.xlabel(\"Atom\", 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": "_mgVSbBshWFV",
        "cellView": "form"
      },
      "source": [
        "#@title **Calculate eigenvectors of Principle Component Analysis (PCA)**\n",
        "data = pt.pca(traj_load, fit=True, ref=0, mask=mask_end, 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')\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 distribution**\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": []
    }
  ]
}