CHARMM_GUI.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "CHARMM-GUI.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "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/CHARMM_GUI.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5CS1elvgo1Qr"
      },
      "source": [
        "# **Hello there!**\n",
        "\n",
        "This is a Jupyter notebook for running Molecular Dynamics (MD) simulations using OpenMM engine and CHARMM force field files. This notebook is a supplementary material of the paper \"***Making it rain: Cloud-based molecular simulations for everyone***\" ([link here](https://doi.org/10.1021/acs.jcim.1c00998)) and we encourage you to read it before using this pipeline.\n",
        "\n",
        "The main goal of this notebook is to demonstrate how to harness the power of cloud-computing to run microsecond-long MD simulations in a cheap and yet feasible fashion.\n",
        "\n",
        "---\n",
        "\n",
        " **This notebook is NOT a standard protocol for MD simulations!** It is just simple MD pipeline illustrating each step of a simulation protocol.\n",
        "\n",
        "--- \n",
        "**Bugs**\n",
        "- If you encounter any bugs, please report the issue to https://github.com/pablo-arantes/making-it-rain/issues\n",
        "\n",
        "**Acknowledgments**\n",
        "- We would like to thank the OpenMM team for developing an excellent and open source engine. \n",
        "\n",
        "- A Making-it-rain by **Pablo R. Arantes** ([@pablitoarantes](https://twitter.com/pablitoarantes)), **Marcelo D. Polêto** ([@mdpoleto](https://twitter.com/mdpoleto)), **Conrado Pedebos** ([@ConradoPedebos](https://twitter.com/ConradoPedebos)) and **Rodrigo Ligabue-Braun** ([@ligabue_braun](https://twitter.com/ligabue_braun)).\n",
        "\n",
        "\n",
        "- Also, credit to [David Koes](https://github.com/dkoes) for his awesome [py3Dmol](https://3dmol.csb.pitt.edu/) plugin.\n",
        "\n",
        "- For related notebooks see: [Making-it-rain](https://github.com/pablo-arantes/making-it-rain)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gJBeLHQ0AQ3y"
      },
      "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 CHARMM inputs, we wil need:\n",
        "\n",
        "*   A .crd and a .pdb file containing a set of atomic coordinates.\n",
        "*   The respective .psf file containing the system topology.\n",
        "*   A *toppar* folder containing force field parameter files.\n",
        "*   A *toppar.str* file pointing to the force field files that we will use in our simulation.\n",
        "\n",
        "In this notebook, we will simulate PDB 1AKI, a hen egg-white lysozyme. To build our simulation box, we will use CHARMM-GUI (https://www.charmm-gui.org/) solution builder feature. (Important: to use CHARMM-GUI, you must register for a free account). You can download the necessary input files from [here](https://github.com/pablo-arantes/making-it-rain/tree/main/Colab_OpenMM_CHARMM);\n",
        "\n",
        "---\n",
        "\n",
        "For more information on how to create a solvated simulation box, please refer to CHARMM-GUI YouTube channel (https://www.youtube.com/channel/UCtHN7aNAjDet_JKWPxSioLQ)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z6Lfp_NZPaC-"
      },
      "source": [
        "---\n",
        "---\n",
        "# **Setting the environment for MD calculation**\n",
        "\n",
        "Firstly, we need to install all necessary libraries and packages for our simulation. The main packages we will be installing are:\n",
        "\n",
        "1.    Anaconda (https://docs.conda.io/en/latest/miniconda.html)\n",
        "2.    OpenMM (https://openmm.org/)\n",
        "3.    PyTraj (https://amber-md.github.io/pytraj/latest/index.html)\n",
        "4.    py3Dmol (https://pypi.org/project/py3Dmol/)\n",
        "5.    Numpy (https://numpy.org/)\n",
        "6.    Matplotlib (https://matplotlib.org/)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wH1oMiVUlxO_",
        "cellView": "form"
      },
      "source": [
        "#@title **Install dependencies**\n",
        "#@markdown This will take a few minutes. Please, have a coffee while you wait ;-)\n",
        "# install dependencies\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",
        "# 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 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",
        "\n",
        "#load dependencies\n",
        "import sys\n",
        "sys.path.append('/usr/local/lib/python3.7/site-packages/')\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 numpy as np\n",
        "import MDAnalysis as mda\n",
        "import py3Dmol\n",
        "import pytraj as pt\n",
        "import platform\n",
        "import scipy.cluster.hierarchy\n",
        "from scipy.spatial.distance import squareform\n",
        "import scipy.stats as stats\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "from scipy.interpolate import griddata\n",
        "import seaborn as sb\n",
        "from statistics import mean, stdev\n",
        "from pytraj import matrix\n",
        "from matplotlib import colors\n",
        "from IPython.display import set_matplotlib_formats"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lbg5J03ImeTX"
      },
      "source": [
        "## Using Google Drive to store simulation data\n",
        "\n",
        "Google Colab does not allow users to keep data on their computing nodes. However, we can use Google Drive to read, write, and store our simulations files. Therefore, we suggest to you to:\n",
        "\n",
        "1.   Create a folder in your own Google Drive and copy the necessary input files there.\n",
        "2.   Copy the path of your created directory. We will use it below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "Lm7Akepv_vl-"
      },
      "source": [
        "#@title ### **Import Google Drive**\n",
        "#@markdown Click in the \"Run\" buttom to make your Google Drive accessible.\n",
        "from google.colab import drive\n",
        "\n",
        "drive.flush_and_unmount()\n",
        "drive.mount('/content/drive', force_remount=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "lOKg9eH_ueRn"
      },
      "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": "IcUZIEX1A3g4"
      },
      "source": [
        "---\n",
        "---\n",
        "# **Loading the necessary input files**\n",
        "\n",
        "At this point, we should have all libraries and dependencies installed and all necessary input files already at your Google Drive folder.\n",
        "\n",
        "**IMPORTANT**: make sure the *toppar.str* file points to the correct pathway of all force field files. CHARMM-GUI assumes a certain directory structure that is not applicable here. Incorrect pathways will cause the job to crash. Correct the pathway and re-upload the files if necessary.\n",
        "\n",
        "Below, you should provide the names of all input files and the pathway of your Google Drive folder containing them. Please, don't use spaces in the files and folders names, i.e., CHARMM_GUI, step3_input.pdb and so on.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wZfALbPc6t5H",
        "cellView": "form"
      },
      "source": [
        "#@title **Please, provide the necessary input files below:** \n",
        "PSF_filename = 'step3_input.psf' #@param {type:\"string\"}\n",
        "CRD_filename = 'step3_input.crd' #@param {type:\"string\"}\n",
        "PDB_filename = 'step3_input.pdb' #@param {type:\"string\"}\n",
        "toppar_filename = 'toppar.str' #@param {type:\"string\"}\n",
        "\n",
        "Google_Drive_Path = '/content/drive/MyDrive/CHARMM_GUI/' #@param {type:\"string\"}\n",
        "workDir = Google_Drive_Path\n",
        "\n",
        "top = os.path.join(workDir, str(PSF_filename))\n",
        "toppar = os.path.join(workDir, str(toppar_filename))\n",
        "crd = os.path.join(workDir, str(CRD_filename))\n",
        "pdb = os.path.join(workDir, str(PDB_filename))\n",
        "\n",
        "pdb_check = os.path.exists(pdb)\n",
        "crd_check = os.path.exists(crd)\n",
        "top_check = os.path.exists(top)\n",
        "toppar_check = os.path.exists(toppar)\n",
        "\n",
        "if pdb_check == True and crd_check == True and top_check == True and toppar_check == True:\n",
        "    print(\"Files loaded succesfully! ;-)\")\n",
        "else:\n",
        "    print(\"ERROR! Check your input names and Google Drive path\")\n",
        "\n",
        "#@markdown ---"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yr0canfGqXNy"
      },
      "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_sidechains=False,\n",
        "             show_mainchain=False,\n",
        "             color=\"None\"):\n",
        "\n",
        "  def mainchain(p, color=\"white\", model=0):\n",
        "    BB = ['C','O','N','CA']\n",
        "    p.addStyle({\"model\":model,'atom':BB},\n",
        "                       {'stick':{'colorscheme':f\"{color}Carbon\",'radius':0.4}})\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",
        " \n",
        "\n",
        "  def sidechain(p, model=0):\n",
        "    HP = [\"ALA\",\"GLY\",\"VAL\",\"ILE\",\"LEU\",\"PHE\",\"MET\",\"PRO\",\"TRP\",\"CYS\",\"TYR\"]\n",
        "    BB = ['C','O','N']\n",
        "    p.addStyle({\"model\":model,'and':[{'resn':HP},{'atom':BB,'invert':True}]},\n",
        "              {'stick':{'colorscheme':\"whiteCarbon\",'radius':0.4}})\n",
        "    p.addStyle({\"model\":model,'and':[{'resn':\"GLY\"},{'atom':'CA'}]},\n",
        "              {'sphere':{'colorscheme':\"whiteCarbon\",'radius':0.4}})\n",
        "    p.addStyle({\"model\":model,'and':[{'resn':\"PRO\"},{'atom':['C','O'],'invert':True}]},\n",
        "              {'stick':{'colorscheme':\"whiteCarbon\",'radius':0.4}})  \n",
        "    p.addStyle({\"model\":model,'and':[{'resn':HP,'invert':True},{'atom':BB,'invert':True}]},\n",
        "              {'stick':{'colorscheme':\"whiteCarbon\",'radius':0.4}})\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_sidechains: sidechain(p)\n",
        "  if show_mainchain: mainchain(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_sidechains=ipywidgets.Checkbox(value=False),\n",
        "         show_mainchain=ipywidgets.Checkbox(value=False),\n",
        "         color=ipywidgets.Dropdown(options=['None', 'rainbow'], value='rainbow'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g6RWtNE-qnzQ"
      },
      "source": [
        "---\n",
        "---\n",
        "# **Equilibrating the simulation box**\n",
        "\n",
        "Proper MD equilibration protocol is designed to equilibrate both temperature and pressure throughout the simulation box while preserving the protein experimental conformation. In addition, we also allow the solvent to accomodate around the protein, creating proper solvation layers.\n",
        "\n",
        "Below, we will set up the MD equilibration parameters, such as temperature, pressure and the desired simulation time. We will define the force constant used to restraint protein heavy-atoms in place and the frequency at which we want to save atomic coordinates in a trajectory file (.dcd).\n",
        "\n",
        "After you are done, you can run the next 2 cells to equilibrate your system."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "8x9Qp_dbr9HP"
      },
      "source": [
        "#@title ### **Parameters for MD Equilibration protocol:**\n",
        "\n",
        "# remove whitespaces\n",
        "Jobname = '1aki_equil' #@param {type:\"string\"}\n",
        "jobname = os.path.join(workDir,Jobname)\n",
        "\n",
        "Minimization_steps = \"1000\" #@param [\"1000\", \"5000\", \"10000\", \"20000\", \"50000\", \"100000\"]\n",
        "\n",
        "#@markdown Simulation time (in nanoseconds) and integration time (in femtoseconds): \n",
        "Time = \"0.1\" #@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",
        "\n",
        "#@markdown Position restraints force constant (in kJ/mol): \n",
        "Force_constant = 500 #@param {type:\"slider\", min:0, max:2000, step:100}\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_eq = 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_eq = Write_the_log\n",
        "\n",
        "\n",
        "#@markdown ---"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zoamR9iynphz",
        "cellView": "form"
      },
      "source": [
        "#@title **Runs an Equilibration MD simulation (NPT ensemble)**\n",
        "\n",
        "#@markdown Now, let's equilibrate our system!\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",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\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 gen_box(psf, crd):\n",
        "\tcoords = crd.positions\n",
        "\n",
        "\tmin_crds = [coords[0][0], coords[0][1], coords[0][2]]\n",
        "\tmax_crds = [coords[0][0], coords[0][1], coords[0][2]]\n",
        "\n",
        "\tfor coord in coords:\n",
        "\t\tmin_crds[0] = min(min_crds[0], coord[0])\n",
        "\t\tmin_crds[1] = min(min_crds[1], coord[1])\n",
        "\t\tmin_crds[2] = min(min_crds[2], coord[2])\n",
        "\t\tmax_crds[0] = max(max_crds[0], coord[0])\n",
        "\t\tmax_crds[1] = max(max_crds[1], coord[1])\n",
        "\t\tmax_crds[2] = max(max_crds[2], coord[2])\n",
        "\n",
        "\tboxlx = max_crds[0]-min_crds[0]\n",
        "\tboxly = max_crds[1]-min_crds[1]\n",
        "\tboxlz = max_crds[2]-min_crds[2]\n",
        "\n",
        "\tpsf.setBox(boxlx, boxly, boxlz)\n",
        "\treturn psf\n",
        "\n",
        "def read_toppar(filename,path):\n",
        "\textlist = ['rtf', 'prm', 'str']\n",
        "\n",
        "\tparFiles = ()\n",
        "\tfor line in open(filename, 'r'):\n",
        "\t\tif '!' in line: line = line.split('!')[0]\n",
        "\t\tparfile = line.strip()\n",
        "\t\tif len(parfile) != 0:\n",
        "\t\t\text = parfile.lower().split('.')[-1]\n",
        "\t\t\tif not ext in extlist: continue\n",
        "\t\t\tparFiles += ( path + parfile, )\n",
        "\n",
        "\tparams = CharmmParameterSet( *parFiles )\n",
        "\treturn params, parFiles\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",
        "#############################################\n",
        "# Defining MD simulation parameters\n",
        "\n",
        "top = os.path.join(workDir, str(PSF_filename))\n",
        "toppar = os.path.join(workDir, str(toppar_filename))\n",
        "crd = os.path.join(workDir, str(CRD_filename))\n",
        "pdb = os.path.join(workDir, str(PDB_filename))\n",
        "\n",
        "toppar = toppar\n",
        "coordinatefile = crd\n",
        "pdbfile = pdb\n",
        "topologyfile = top\n",
        "\n",
        "time_ps = float(Time)*1000 #convert to picosecond\n",
        "simulation_time = float(time_ps)*picosecond\t\t# in ps\n",
        "dt = int(dt_eq)*femtosecond\n",
        "temperature = float(temperature_eq)*kelvin\n",
        "savcrd_freq = int(write_the_trajectory_eq)*picosecond\n",
        "print_freq  = int(write_the_log_eq)*picosecond\n",
        "\n",
        "pressure\t= float(pressure_eq)*bar\n",
        "\n",
        "restraint_fc = int(Force_constant) # kJ/mol\n",
        "\n",
        "nsteps  = int(simulation_time.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nprint  = int(print_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nsavcrd = int(savcrd_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "\n",
        "charmm_params, toppar_files = read_toppar(toppar,workDir)\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",
        "print(\"\\tForce Field files = \" + str(toppar_files))\n",
        "\n",
        "print(\"\\n\\tSimulation_time = \" + str(simulation_time))\n",
        "print(\"\\tIntegration timestep = \" + str(dt))\n",
        "print(\"\\tTotal number of steps = \" +  str(nsteps))\n",
        "\n",
        "print(\"\\n\\tSave coordinates each \" + str(savcrd_freq))\n",
        "print(\"\\tPrint in log file each \" + str(print_freq))\n",
        "\n",
        "print(\"\\n\\tTemperature = \" + str(temperature))\n",
        "print(\"\\tPressure = \" + str(pressure))\n",
        "#############################################\n",
        "\n",
        "print(\"\\n> Setting the system:\\n\")\n",
        "\n",
        "print(\"\\t- Reading topology and structure file...\")\n",
        "psf = CharmmPsfFile(topologyfile)\n",
        "crd = CharmmCrdFile(coordinatefile)\n",
        "\n",
        "print(\"\\t- Setting box (using user's information)...\")\n",
        "psf = gen_box(psf, crd)\n",
        "\n",
        "print(\"\\t- Creating system and setting parameters...\")\n",
        "system = psf.createSystem(charmm_params, nonbondedMethod=PME, nonbondedCutoff=1.2*nanometers, switchDistance=1.0*nanometer, ewaldErrorTolerance = 0.0005, constraints=HBonds)\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",
        "system = restraints(system, crd, 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, 1/picosecond, dt)\n",
        "simulation = Simulation(psf.topology, system, integrator)\n",
        "simulation.context.setPositions(crd.positions)\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",
        "if os.path.exists(rst_file):\n",
        "    print(\"\\n\\n> Equilibration finished (\" + rst_file + \" present). Nothing to do here... <\")\n",
        "    exit\n",
        "    sys.exit()\n",
        "\n",
        "else:\n",
        "\tdcd = DCDReporter(dcd_file, nsavcrd)\n",
        "\tfirstdcdstep = (nsteps) + nsavcrd\n",
        "\tdcd._dcd = DCDFile(dcd._out, simulation.topology, simulation.integrator.getStepSize(), firstdcdstep, nsavcrd) # charmm doesn't like first step to be 0\n",
        "\n",
        "\tsimulation.reporters.append(dcd)\n",
        "\tsimulation.reporters.append(StateDataReporter(stdout, nprint, step=True, speed=True, progress=True, totalSteps=nsteps, 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...\")\n",
        "\tsimulation.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\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yNsUJYs7A-G5"
      },
      "source": [
        "---\n",
        "---\n",
        "# **Running a Production MD simulation**\n",
        "\n",
        "Finally, we will proceed with the Production simulation itself using the equilibrated system coordinates as input structure. We will use the same .psf topology file and the same force field files. \n",
        "\n",
        "Note that we will use here a *.rst state file* , which contains atomic velocities and positions from the last frame of the equilibration simulation, guaranteeing that our production simulation begins from a thermodynamically equilibrated system.\n",
        "\n",
        "Another important information here is the **Number_of_strides** and the **Stride_Time**. In this notebook, we simulate a defined number of *strides*, so the **simulation time = Number_of_strides*Stride_Time**. For example, we can simulate 100ns by setting *Number_of_strides=10* and *Stride_Time=10 ns*.\n",
        "\n",
        "**Important: at the end of the Production simulation, we concatenate all strides to create a complete trajectory file which can be visualized and analyzed**\n",
        "\n",
        "The idea behind this approach is to make use of the intermitent 12h/24h period that Google Colab allows us to use its GPUs. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z0JV6Zid50_o",
        "cellView": "form"
      },
      "source": [
        "#@markdown ---\n",
        "#@markdown ### **Provide input file names below:** \n",
        "\n",
        "PSF_file = 'step3_input.psf' #@param {type:\"string\"}\n",
        "Equilibrated_PDB = '1aki_equil.pdb' #@param {type:\"string\"}\n",
        "Toppar_file = 'toppar.str' #@param {type:\"string\"}\n",
        "State_file = '1aki_equil.rst' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown ### **Parameters for MD Production protocol:**\n",
        "\n",
        "\n",
        "# remove whitespaces\n",
        "Jobname = 'test_prod' #@param {type:\"string\"}\n",
        "jobname = os.path.join(workDir, Jobname)\n",
        "\n",
        "#@markdown Simulation time (in nanoseconds), number of strides (integers) and integration timestep (in femtoseconds): \n",
        "Stride_Time = \"0.2\" #@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_eq = 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_eq = Write_the_log\n",
        "\n",
        "#@markdown ---"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0QcjKSoqHHhi",
        "cellView": "form"
      },
      "source": [
        "#@title **Runs an Production MD simulation (NPT ensemble) after equilibration**\n",
        "#\n",
        "#@markdown\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 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 gen_box(psf, crd):\n",
        "\tcoords = crd.positions\n",
        "\n",
        "\tmin_crds = [coords[0][0], coords[0][1], coords[0][2]]\n",
        "\tmax_crds = [coords[0][0], coords[0][1], coords[0][2]]\n",
        "\n",
        "\tfor coord in coords:\n",
        "\t\tmin_crds[0] = min(min_crds[0], coord[0])\n",
        "\t\tmin_crds[1] = min(min_crds[1], coord[1])\n",
        "\t\tmin_crds[2] = min(min_crds[2], coord[2])\n",
        "\t\tmax_crds[0] = max(max_crds[0], coord[0])\n",
        "\t\tmax_crds[1] = max(max_crds[1], coord[1])\n",
        "\t\tmax_crds[2] = max(max_crds[2], coord[2])\n",
        "\n",
        "\tboxlx = max_crds[0]-min_crds[0]\n",
        "\tboxly = max_crds[1]-min_crds[1]\n",
        "\tboxlz = max_crds[2]-min_crds[2]\n",
        "\n",
        "\tpsf.setBox(boxlx, boxly, boxlz)\n",
        "\treturn psf\n",
        "\n",
        "def read_toppar(filename,path):\n",
        "\textlist = ['rtf', 'prm', 'str']\n",
        "\n",
        "\tparFiles = ()\n",
        "\tfor line in open(filename, 'r'):\n",
        "\t\tif '!' in line: line = line.split('!')[0]\n",
        "\t\tparfile = line.strip()\n",
        "\t\tif len(parfile) != 0:\n",
        "\t\t\text = parfile.lower().split('.')[-1]\n",
        "\t\t\tif not ext in extlist: continue\n",
        "\t\t\tparFiles += ( path + parfile, )\n",
        "\n",
        "\tparams = CharmmParameterSet( *parFiles )\n",
        "\treturn params, parFiles\n",
        "##############################################\n",
        "\n",
        "\n",
        "#############################################\n",
        "# Defining MD simulation parameters\n",
        "\n",
        "topologyfile = os.path.join(workDir, PSF_file)\n",
        "pdbfile = os.path.join(workDir, Equilibrated_PDB)\n",
        "toppar = os.path.join(workDir,Toppar_file)\n",
        "equil_rst_file = os.path.join(workDir,State_file)\n",
        "\n",
        "\n",
        "stride_time_ps = float(stride_time_prod)*1000 #convert to picosecond\n",
        "stride_time = float(stride_time_ps)*picosecond\t\t# in ps\n",
        "nstride = int(nstride)\n",
        "dt = int(dt_prod)*femtosecond\n",
        "simulation_time = stride_time*nstride\n",
        "\n",
        "temperature = float(temperature_prod)*kelvin\n",
        "pressure\t= float(pressure_prod)*bar\n",
        "\n",
        "savcrd_freq = int(write_the_trajectory_eq)*picosecond\n",
        "print_freq  = int(write_the_log_eq)*picosecond\n",
        "\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",
        "\n",
        "charmm_params, toppar_files = read_toppar(toppar,workDir)\n",
        "\n",
        "#############################################\n",
        "print(\"\\n> Simulation details:\\n\")\n",
        "print(\"\\tJob name = \" + jobname)\n",
        "print(\"\\tPDB file = \" + str(pdbfile))\n",
        "print(\"\\tTopology file = \" + str(topologyfile))\n",
        "print(\"\\tForce Field files = \" + str(toppar_files))\n",
        "\n",
        "print(\"\\n\\tSimulation_time = \" + str(simulation_time))\n",
        "print(\"\\tIntegration timestep = \" + str(dt))\n",
        "print(\"\\tTotal number of steps = \" +  str(nsteps*nstride))\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",
        "print(\"\\t- Reading topology and structure file...\")\n",
        "psf = CharmmPsfFile(topologyfile)\n",
        "crd = PDBFile(pdbfile)\n",
        "\n",
        "print(\"\\t- Setting box (using user's information)...\")\n",
        "psf = gen_box(psf, crd)\n",
        "\n",
        "print(\"\\t- Creating system and setting parameters...\")\n",
        "system = psf.createSystem(charmm_params, nonbondedMethod=PME, nonbondedCutoff=1.2*nanometers, switchDistance=1.0*nanometer, ewaldErrorTolerance = 0.0005, constraints=HBonds)\n",
        "\n",
        "print(\"\\t- Setting barostat...\")\n",
        "system.addForce(MonteCarloBarostat(pressure, temperature))\n",
        "\n",
        "print(\"\\t- Setting integrator...\")\n",
        "integrator = LangevinIntegrator(temperature, 1/picosecond, dt)\n",
        "simulation = Simulation(psf.topology, system, integrator)\n",
        "simulation.context.setPositions(crd.positions)\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\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DaLoQlJGf57o",
        "cellView": "form"
      },
      "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=\"@CA\", 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, pdb)\n",
        "print(traj_load)\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! \")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x_awbwdZnp6Q",
        "cellView": "form"
      },
      "source": [
        "#@title **Load, view and check the trajectory**\n",
        "#@markdown This will take a few minutes. Another coffee would be great. :-)\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "!rm *.pdb 2> /dev/null\n",
        "\n",
        "#py3dmol functions\n",
        "class Atom(dict):\n",
        "  def __init__(self, line):\n",
        "    self[\"type\"] = line[0:6].strip()\n",
        "    self[\"idx\"] = line[6:11].strip()\n",
        "    self[\"name\"] = line[12:16].strip()\n",
        "    self[\"resname\"] = line[17:20].strip()\n",
        "    self[\"resid\"] = int(int(line[22:26]))\n",
        "    self[\"x\"] = float(line[30:38])\n",
        "    self[\"y\"] = float(line[38:46])\n",
        "    self[\"z\"] = float(line[46:54])\n",
        "    self[\"sym\"] = line[76:78].strip()\n",
        "\n",
        "  def __str__(self):\n",
        "    line = list(\" \" * 80)\n",
        "    line[0:6] = self[\"type\"].ljust(6)\n",
        "    line[6:11] = self[\"idx\"].ljust(5)\n",
        "    line[12:16] = self[\"name\"].ljust(4)\n",
        "    line[17:20] = self[\"resname\"].ljust(3)\n",
        "    line[22:26] = str(self[\"resid\"]).ljust(4)\n",
        "    line[30:38] = str(self[\"x\"]).rjust(8)\n",
        "    line[38:46] = str(self[\"y\"]).rjust(8)\n",
        "    line[46:54] = str(self[\"z\"]).rjust(8)\n",
        "    line[76:78] = self[\"sym\"].rjust(2)\n",
        "    return \"\".join(line) + \"\\n\"\n",
        "        \n",
        "class Molecule(list):\n",
        "  def __init__(self, file):\n",
        "    for line in file:\n",
        "      if \"ATOM\" in line or \"HETATM\" in line:\n",
        "        self.append(Atom(line))\n",
        "            \n",
        "    def __str__(self):\n",
        "      outstr = \"\"\n",
        "      for at in self:\n",
        "        outstr += str(at)\n",
        "      return outstr\n",
        "\n",
        "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 TIP3)')\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.setStyle({'model': -1, 'serial': i+1}, at.get(\"pymol\", default))\n",
        "\n",
        "view.zoomTo()\n",
        "view.animate({'loop': \"forward\"})\n",
        "view.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Emh0vU5UjgB6"
      },
      "source": [
        "---\n",
        "---\n",
        "# **Analysis**\n",
        "\n",
        "Although visualizing your trajectory can be quite useful, sometimes you also want more quantitative data.\n",
        "\n",
        "Analyses of MD trajectories vary a lot and we do not intend to cover it all here. However, one can make use of MDanalysis or PyTraj to easily analyze simulations. \n",
        "\n",
        "Below, you can find a few examples of code snippets that can help you to shed some light on your simulation behavior."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wBrBMF4Puyv6",
        "cellView": "form"
      },
      "source": [
        "#@title **Compute RMSD of protein's CA atoms**\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'rmsd_ca' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "rmsd = pt.rmsd(traj_load, ref = 0, mask = \"@CA\")\n",
        "\n",
        "time = len(rmsd)*int(Write_the_trajectory)/1000\n",
        "time_array = np.arange(0,time,int(Write_the_trajectory)/1000)*int(stride_traj)\n",
        "\n",
        "# Plotting:\n",
        "ax = plt.plot(time_array, rmsd, alpha=0.6, color = 'blue', linewidth = 1.0)\n",
        "plt.xlim(0, simulation_ns)\n",
        "#plt.ylim(2, 6)\n",
        "\n",
        "plt.xlabel(\"Time (ns)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel(\"RMSD [$\\AA$]\", fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "raw_data=pd.DataFrame(rmsd)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZHyMpikjuaLT",
        "cellView": "form"
      },
      "source": [
        "#@title **Plot RMSD as a ditribution**\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'rmsd_dist' #@param {type:\"string\"}\n",
        "\n",
        "ax = sb.kdeplot(rmsd, color=\"blue\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('RMSD [$\\AA$]', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks([])\n",
        "plt.ylabel('')\n",
        "ax.spines['top'].set_visible(False)\n",
        "ax.spines['right'].set_visible(False)\n",
        "ax.spines['bottom'].set_visible(True)\n",
        "ax.spines['left'].set_visible(False)\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CvOFrXGXwXrV",
        "cellView": "form"
      },
      "source": [
        "#@title **Compute radius of gyration of protein's CA atoms**\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'radius_gyration' #@param {type:\"string\"}\n",
        "\n",
        "radgyr = pt.radgyr(traj_load, mask = \"@CA\")\n",
        "time = len(rmsd)*int(Write_the_trajectory)/1000\n",
        "time_array = np.arange(0,time,int(Write_the_trajectory)/1000)*int(stride_traj)\n",
        "\n",
        "# Plotting:\n",
        "plt.plot(time_array, radgyr, alpha=0.6, color = 'green', linewidth = 1.0)\n",
        "plt.xlim(0, simulation_ns)\n",
        "#plt.ylim(2, 6)\n",
        "\n",
        "plt.xlabel(\"Time (ns)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel(\"Radius of gyration ($\\AA$)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "raw_data=pd.DataFrame(radgyr)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2Q7FKg8Fuxr9",
        "cellView": "form"
      },
      "source": [
        "#@title **Plot radius of gyration as a ditribution**\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'radius_gyration_dist' #@param {type:\"string\"}\n",
        "\n",
        "ax = sb.kdeplot(radgyr, color=\"green\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('Radius of gyration ($\\AA$)', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks([])\n",
        "plt.ylabel('')\n",
        "ax.spines['top'].set_visible(False)\n",
        "ax.spines['right'].set_visible(False)\n",
        "ax.spines['bottom'].set_visible(True)\n",
        "ax.spines['left'].set_visible(False)\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p2Y0DgwTxLWc",
        "cellView": "form"
      },
      "source": [
        "#@title **Compute RMSF of protein's CA atoms**\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'rmsf_ca' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "rmsf = pt.rmsf(traj_load, \"@CA\")\n",
        "bfactor = pt.bfactors(traj_load, byres=True)\n",
        "\n",
        "# Plotting:\n",
        "plt.plot(rmsf[:,1], alpha=1.0, color = 'red', linewidth = 1.0)\n",
        "\n",
        "plt.xlabel(\"Residue\", fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel(\"RMSF ($\\AA$)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.xlim(0, len(rmsf[:-1]))\n",
        "\n",
        "#plt.xticks(np.arange(min(rmsf[:1]), max(rmsf[:1])))\n",
        "plt.yticks(fontsize = 12)\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "raw_data=pd.DataFrame(rmsf)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JalicqqrTodW",
        "cellView": "form"
      },
      "source": [
        "#@title **2D RMSD**\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = '2D_rmsd' #@param {type:\"string\"}\n",
        "\n",
        "last_frame = len(time_array)\n",
        "\n",
        "stride_ticks_f = (last_frame)/5\n",
        "ticks_frame = np.arange(0,(len(time_array) + float(stride_ticks_f)), float(stride_ticks_f))\n",
        "a = ticks_frame.astype(float)\n",
        "stride_ticks_t = (simulation_ns)/5\n",
        "tick_time = np.arange(0,(float(simulation_ns) + float(stride_ticks_t)), float(stride_ticks_t))\n",
        "b = tick_time.astype(float)\n",
        "\n",
        "mat1 = pt.pairwise_rmsd(traj_load, mask=\"@CA\", frame_indices=range(int(number_frames_analysis)))\n",
        "\n",
        "\n",
        "ax = plt.imshow(mat1, cmap = 'PRGn', origin='lower', interpolation = 'bicubic')\n",
        "plt.title('2D RMSD')\n",
        "plt.xlabel('Time (ns)', fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel('Time (ns)', fontsize = 14, fontweight = 'bold')\n",
        "# plt.xticks(fontsize = 12)\n",
        "# plt.yticks(fontsize = 12)\n",
        "plt.xticks(a, b.round(decimals=3), fontsize = 12)\n",
        "plt.yticks(a, b.round(decimals=3), fontsize = 12)\n",
        "# plt.xlim(0, a[-1])\n",
        "# plt.ylim(0, a[-1])\n",
        "\n",
        "cbar1 = plt.colorbar()\n",
        "cbar1.set_label(\"RMSD ($\\AA$)\", fontsize = 14, fontweight = 'bold')\n",
        "\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "raw_data=pd.DataFrame(mat1)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_mgVSbBshWFV",
        "cellView": "form"
      },
      "source": [
        "#@title **Calculate eigvenctors of Principle Component Analysis (PCA)**\n",
        "data = pt.pca(traj_load, fit=True, ref=0, mask='@CA', n_vecs=2)\n",
        "#print('projection values of each frame to first mode = {} \\n'.format(data[0][0]))\n",
        "#print('projection values of each frame to second mode = {} \\n'.format(data[0][1]))\n",
        "#print('eigvenvalues of first two modes', data[1][0])\n",
        "#print(\"\")\n",
        "#print('eigvenvectors of first two modes: \\n', data[1][1])\n",
        "\n",
        "last_frame = len(time_array)\n",
        "\n",
        "stride_ticks_f = (last_frame)/5\n",
        "ticks_frame = np.arange(0,(len(time_array) + float(stride_ticks_f)), float(stride_ticks_f))\n",
        "a = ticks_frame.astype(float)\n",
        "a2 = a.tolist()\n",
        "stride_ticks_t = (simulation_ns)/5\n",
        "tick_time = np.arange(0,(float(simulation_ns) + float(stride_ticks_t)), float(stride_ticks_t))\n",
        "b = tick_time.astype(float)\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'PCA' #@param {type:\"string\"}\n",
        "\n",
        "Output_PC1 = 'PC1' #@param {type:\"string\"}\n",
        "Output_PC2 = 'PC2' #@param {type:\"string\"}\n",
        "\n",
        "%matplotlib inline\n",
        "%config InlineBackend.figure_format = 'retina'  # high resolution\n",
        "projection_data = data[0]\n",
        "plt.title(r'PCA of C-$\\alpha$')\n",
        "PC1 = data[0][0] \n",
        "PC2 = data[0][1]\n",
        "\n",
        "a = plt.scatter(PC1,PC2, c=range(int(number_frames_analysis)), cmap='Greens', marker='o',s=8, alpha=1)\n",
        "plt.clim(0, last_frame)\n",
        "\n",
        "plt.xlabel('PC1', fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel('PC2', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "# N = len(number_frames)\n",
        "# x2 = np.arange(N)\n",
        "\n",
        "cbar1 = plt.colorbar(a, orientation=\"vertical\")\n",
        "cbar1.set_label('Time(ns)', fontsize = 14, fontweight = 'bold')\n",
        "cbar1.set_ticks(a2)\n",
        "cbar1.set_ticklabels(b.round(decimals=3))\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "pc1=pd.DataFrame(PC1)\n",
        "pc1.to_csv(os.path.join(workDir, Output_PC1 + \".csv\"))\n",
        "pc2=pd.DataFrame(PC2)\n",
        "pc2.to_csv(os.path.join(workDir, Output_PC2 + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yce9RfNtpl-J",
        "cellView": "form"
      },
      "source": [
        "#@title **Plot Principal Component 1 (PC1) and Principal Component 2 (PC2) as a ditribution**\n",
        "Output_name = 'PCA_dist' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "fig = plt.figure(figsize=(9,5))\n",
        "\n",
        "plt.subplot(1, 2, 1)\n",
        "ax = sb.kdeplot(PC1, color=\"green\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('PC1', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks([])\n",
        "plt.ylabel('')\n",
        "ax.spines['top'].set_visible(False)\n",
        "ax.spines['right'].set_visible(False)\n",
        "ax.spines['bottom'].set_visible(True)\n",
        "ax.spines['left'].set_visible(False)\n",
        "plt.subplot(1, 2, 2)\n",
        "ax2 = sb.kdeplot(PC2, color=\"purple\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('PC2', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks([])\n",
        "plt.ylabel('')\n",
        "ax2.spines['top'].set_visible(False)\n",
        "ax2.spines['right'].set_visible(False)\n",
        "ax2.spines['bottom'].set_visible(True)\n",
        "ax2.spines['left'].set_visible(False)\n",
        "\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pTDb7CEfkLq1",
        "cellView": "form"
      },
      "source": [
        "#@title **Pearson's Cross Correlation (CC)**\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'cross_correlation' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "traj_align = pt.align(traj_load, mask='@CA', ref=0)\n",
        "\n",
        "mat_cc = matrix.correl(traj_align, '@CA')\n",
        "\n",
        "ax = plt.imshow(mat_cc, cmap = 'PiYG_r', interpolation = 'bicubic', vmin = -1, vmax = 1, origin='lower')\n",
        "\n",
        "plt.xlabel('Residues', fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel('Residues', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "cbar1 = plt.colorbar()\n",
        "cbar1.set_label('$CC_ij$', fontsize = 14, fontweight = 'bold')\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "raw_data=pd.DataFrame(mat_cc)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}