Martini+cg2all.ipynb

{
  "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/Martini%2Bcg2all.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pj2BWZxUDbEE"
      },
      "source": [
        "# **Hello there!**\n",
        "\n",
        "This is a Jupyter notebook for running Molecular Dynamics (MD) simulations of protein systems using OpenMM engine and Martini force field. 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 coarse-grained MD simulations in a cheap and yet feasible fashion.\n",
        "\n",
        "---\n",
        "\n",
        " **This notebook is NOT a standard protocol for MD simulations!** It is just simple MD pipeline illustrating each step of a simulation protocol.\n",
        "\n",
        "---\n",
        "**Bugs**\n",
        "- If you encounter any bugs, please report the issue to https://github.com/pablo-arantes/making-it-rain/issues\n",
        "\n",
        "**Acknowledgments**\n",
        "- We would like to thank the OpenMM team for developing an excellent and open source engine.\n",
        "\n",
        "- We would like to thank the [Martini](http://cgmartini.nl/index.php) team for their amazing  work in the coarse-grained force fields.\n",
        "\n",
        "- We would like to thank [Paulo C. T. Souza ](https://twitter.com/SouzaPauloCT),  [Justin L. MacCallum](https://twitter.com/jlmaccal), [Valentina Corradi](https://www.ucalgary.ca/biocomputing/vcorradi) and [D. Peter Tieleman](https://live-ucalgary.ucalgary.ca/biocomputing/tieleman) for their work in the Martini3 force field and for the implementation of Martini in OpenMM.\n",
        "\n",
        "-  We would like to thank [Feig lab Team](https://feig.bch.msu.edu/) for his incredible [cg2all](https://www.biorxiv.org/content/10.1101/2023.05.22.541652v1).\n",
        "\n",
        "- A Making-it-rain by **Pablo R. Arantes** ([@pablitoarantes](https://twitter.com/pablitoarantes)), **Marcelo D. Polêto** ([@mdpoleto](https://twitter.com/mdpoleto)), **Conrado Pedebos** ([@ConradoPedebos](https://twitter.com/ConradoPedebos)) and **Rodrigo Ligabue-Braun** ([@ligabue_braun](https://twitter.com/ligabue_braun)).\n",
        "\n",
        "\n",
        "- Also, credit to [David Koes](https://github.com/dkoes) for his awesome [py3Dmol](https://3dmol.csb.pitt.edu/) plugin.\n",
        "\n",
        "- For related notebooks see: [Making-it-rain](https://github.com/pablo-arantes/making-it-rain)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hoyY6XonD1UX"
      },
      "source": [
        "# **Introduction**\n",
        "\n",
        "In general, MD simulations rely on 1) a set of atomic coordinates of all atoms on a simulation box and 2) a set of force field parameters that describes the interaction energies between atoms.\n",
        "\n",
        "In terms of MARTINI inputs, we wil need:\n",
        "*  A PDB ID or a .pdb file containing a set of atomic coordinates of your molecule.\n",
        "\n",
        "In this notebook, we will simulate PDB 1UBQ, a structure of ubiquitin. To build our simulation box, we will use martinize2, a tools for setting up starting structures for molecular dynamics (MD) simulations starting from atomistic coordinates, with a special focus on polymeric systems (including proteins and DNA). Martinize2 is a rewrite of Martinize. It is aimed at producing coarse-grained structures and topologies from an atomistic structure. Here, we will use [Vermouth](https://github.com/marrink-lab/vermouth-martinize) (VERsatile, MOdular, and Universal Tranformation Helper), a python library that powers Martinize2. It allows to describe and apply transformation on molecular structures and topologies using graph algorithms.\n",
        "## ---\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lh96y6mGFY1D"
      },
      "source": [
        "---\n",
        "---\n",
        "# **Setting the environment for MD calculation**\n",
        "\n",
        "Firstly, we need to install all necessary libraries and packages for our simulation. The main packages we will be installing are:\n",
        "\n",
        "1.    Anaconda (https://docs.conda.io/en/latest/miniconda.html)\n",
        "2.    OpenMM (https://openmm.org)\n",
        "3.    PyTraj (https://amber-md.github.io/pytraj/latest/index.html)\n",
        "4.    py3Dmol (https://pypi.org/project/py3Dmol)\n",
        "5.    Numpy (https://numpy.org)\n",
        "6.    Matplotlib (https://matplotlib.org)\n",
        "7.    AmberTools (https://ambermd.org/AmberTools.php)\n",
        "8.    MDAnalysis (https://www.mdanalysis.org)\n",
        "9.    Vermouth (https://github.com/marrink-lab/vermouth-martinize)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "mBFASF7ICxWi"
      },
      "outputs": [],
      "source": [
        "#@title **Install Conda Colab**\n",
        "#@markdown It will restart the kernel (session), don't worry.\n",
        "!pip install -q condacolab\n",
        "import condacolab\n",
        "condacolab.install()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n0XmYfH1C2we",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title **Install dependencies**\n",
        "#@markdown It will take a few minutes, please, drink a coffee and wait. ;-)\n",
        "# install dependencies\n",
        "%%capture\n",
        "\n",
        "import sys\n",
        "import tarfile\n",
        "import os\n",
        "import subprocess\n",
        "import sys\n",
        "\n",
        "commands = [\n",
        "    \"mamba install -c conda-forge ambertools -y\",\n",
        "    \"pip install git+https://github.com/pablo-arantes/biopandas\",\n",
        "    \"mamba install openmm=7.7.0\",\n",
        "    \"pip install --upgrade MDAnalysis==2.4.2\",\n",
        "    \"pip install vermouth\",\n",
        "    \"pip install git+https://github.com/pablo-arantes/martini_openmm\",\n",
        "    \"pip install git+https://github.com/Tsjerk/simopt\",\n",
        "    \"pip install git+https://github.com/pablo-arantes/Insane\",\n",
        "    \"pip install py3Dmol\",\n",
        "    \"mamba install -c conda-forge mdtraj\"\n",
        "]\n",
        "\n",
        "for cmd in commands:\n",
        "    subprocess.run(cmd, shell=True)\n",
        "\n",
        "os.system(\"npx degit https://github.com/pablo-arantes/Making-it-rain/ temp\")\n",
        "os.system(\"cp -r temp/martini .\")\n",
        "os.system(\"rm -r temp\")\n",
        "\n",
        "\n",
        "#load dependencies\n",
        "import sys\n",
        "from biopandas.pdb import PandasPdb\n",
        "import openmm as mm\n",
        "from openmm import *\n",
        "from openmm.app import *\n",
        "from openmm.unit import *\n",
        "import os\n",
        "import urllib.request\n",
        "import numpy as np\n",
        "import MDAnalysis as mda\n",
        "import py3Dmol\n",
        "import pytraj as pt\n",
        "import platform\n",
        "import scipy.cluster.hierarchy\n",
        "from scipy.spatial.distance import squareform\n",
        "import scipy.stats as stats\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "from scipy.interpolate import griddata\n",
        "import seaborn as sb\n",
        "from statistics import mean, stdev\n",
        "from pytraj import matrix\n",
        "from matplotlib import colors\n",
        "from IPython.display import set_matplotlib_formats\n",
        "import subprocess\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Install dependencies for cg2all**\n",
        "#@markdown **cg2all** is a novel deep-learning model capable of efficiently predicting all-atom protein structures from coarse-grained (CG) representations.\n",
        "\n",
        "#@markdown It will take around 4-5 minutes, please, continue your coffee and wait. ;-)\n",
        "# install dependencies\n",
        "\n",
        "%%capture\n",
        "%%bash\n",
        "mamba create --name cg2all pip cudatoolkit=11.3 dgl=1.0 -c dglteam/label/cu113\n",
        "source activate cg2all\n",
        "pip install git+http://github.com/huhlim/cg2all"
      ],
      "metadata": {
        "cellView": "form",
        "id": "hJ12UubTiTay"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fDQnAKJLFxtt"
      },
      "source": [
        "## Using Google Drive to store simulation data\n",
        "\n",
        "Google Colab does not allow users to keep data on their computing nodes. However, we can use Google Drive to read, write, and store our simulations files. Therefore, we suggest to you to:\n",
        "\n",
        "1.   Create a folder in your own Google Drive and copy the necessary input files there.\n",
        "2.   Copy the path of your created directory. We will use it below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Lm7Akepv_vl-"
      },
      "outputs": [],
      "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)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "lOKg9eH_ueRn"
      },
      "outputs": [],
      "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)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BZwl66HTGI7v"
      },
      "source": [
        "---\n",
        "---\n",
        "# **Loading the necessary input files**\n",
        "\n",
        "At this point, we should have all libraries and dependencies installed and all necessary input files already at your Google Drive folder.\n",
        "\n",
        "**Important**: You can submit your own PDB file, make sure the PDB file points to the correct structure. Please type the entire file name, i.e. file.pdb\n",
        "\n",
        "Below, you should provide the names of all input files and the pathway of your Google Drive folder containing them.\n",
        "\n",
        "**Please, don't use spaces in the files and folders names, i.e. MyDrive/protein_martini and so on.**\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-7em-dUO51Nw",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title **Please, provide the necessary input below:**\n",
        "%%capture\n",
        "Google_Drive_Path = '/content/drive/MyDrive/' #@param {type:\"string\"}\n",
        "workDir = Google_Drive_Path\n",
        "\n",
        "if os.path.exists(os.path.join(workDir)):\n",
        "  pass\n",
        "else:\n",
        "  os.system(\"mkdir \" + str(workDir))\n",
        "\n",
        "foldername = 'martini'\n",
        "top_folder = os.path.join(workDir, str(foldername))\n",
        "top_folder_check = os.path.exists(top_folder)\n",
        "if top_folder_check == False:\n",
        "  os.system(\"cp -r ./martini \" + str(workDir))\n",
        "else:\n",
        "  pass\n",
        "\n",
        "Own_PDB = \"no\" #@param [\"yes\", \"no\"]\n",
        "\n",
        "PDB_ID_or_filename = '1UBQ' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "if Own_PDB == \"no\":\n",
        "  pdbfn = PDB_ID_or_filename + \".pdb\"\n",
        "  pdbfn = PDB_ID_or_filename + \".pdb\"\n",
        "  url = 'https://files.rcsb.org/download/' + pdbfn\n",
        "  outfnm = os.path.join(workDir, pdbfn)\n",
        "  urllib.request.urlretrieve(url, outfnm)\n",
        "else:\n",
        "  outfnm = os.path.join(workDir, PDB_ID_or_filename)\n",
        "\n",
        "import numpy as np\n",
        "import mdtraj as md\n",
        "from mdtraj.utils.six import PY2\n",
        "from mdtraj.utils import ensure_type\n",
        "from mdtraj.geometry.hbond import _prep_kabsch_sander_arrays\n",
        "from mdtraj.geometry import _geometry\n",
        "if PY2:\n",
        "    from string import maketrans\n",
        "else:\n",
        "    maketrans = str.maketrans\n",
        "\n",
        "SIMPLIFIED_CODE_TRANSLATION = maketrans('HGIEBTS ', 'HHHEECCC')\n",
        "__all__ = ['compute_dssp']\n",
        "\n",
        "def compute_dssp(traj, simplified=True):\n",
        "    \"\"\"Compute Dictionary of protein secondary structure (DSSP) secondary structure assignments\n",
        "    Parameters\n",
        "    ----------\n",
        "    traj : md.Trajectory\n",
        "        A trajectory\n",
        "    simplified : bool, default=True\n",
        "        Use the simplified 3-category assignment scheme. Otherwise the original\n",
        "        8-category scheme is used.\n",
        "    Returns\n",
        "    -------\n",
        "    assignments : np.ndarray, shape=(n_frames, n_residues), dtype=S1\n",
        "        The assignments is a 2D array of character codes (see below), giving\n",
        "        the secondary structure of each residue in each frame.\n",
        "    Notes\n",
        "    -----\n",
        "    The DSSP assignment codes are:\n",
        "       - 'H' : Alpha helix\n",
        "       - 'B' : Residue in isolated beta-bridge\n",
        "       - 'E' : Extended strand, participates in beta ladder\n",
        "       - 'G' : 3-helix (3/10 helix)\n",
        "       - 'I' : 5 helix (pi helix)\n",
        "       - 'T' : hydrogen bonded turn\n",
        "       - 'S' : bend\n",
        "       - ' ' : Loops and irregular elements\n",
        "    The simplified DSSP codes are:\n",
        "       - 'H' : Helix. Either of the 'H', 'G', or 'I' codes.\n",
        "       - 'E' : Strand. Either of the 'E', or 'B' codes.\n",
        "       - 'C' : Coil. Either of the 'T', 'S' or ' ' codes.\n",
        "    A special 'NA' code will be assigned to each 'residue' in the topology which\n",
        "    isn't actually a protein residue (does not contain atoms with the names\n",
        "    'CA', 'N', 'C', 'O'), such as water molecules that are listed as 'residue's\n",
        "    in the topology.\n",
        "    Our implementation is based on DSSP-2.2.0, written by Maarten L. Hekkelman\n",
        "    and distributed under the Boost Software license.\n",
        "    References\n",
        "    ----------\n",
        "    .. [1] Kabsch W, Sander C (1983). \"Dictionary of protein secondary\n",
        "       structure: pattern recognition of hydrogen-bonded and geometrical\n",
        "       features\". Biopolymers 22 (12): 2577-637. doi:10.1002/bip.360221211\n",
        "    \"\"\"\n",
        "    if traj.topology is None:\n",
        "        raise ValueError('kabsch_sander requires topology')\n",
        "\n",
        "    xyz, nco_indices, ca_indices, proline_indices, protein_indices \\\n",
        "        = _prep_kabsch_sander_arrays(traj)\n",
        "    chain_ids = np.array([r.chain.index for r in traj.top.residues], dtype=np.int32)\n",
        "\n",
        "    value = _geometry._dssp(xyz, nco_indices, ca_indices, proline_indices, chain_ids)\n",
        "    if simplified:\n",
        "        value = value.translate(SIMPLIFIED_CODE_TRANSLATION)\n",
        "\n",
        "    n_frames = xyz.shape[0]\n",
        "    n_residues = nco_indices.shape[0]\n",
        "    if PY2:\n",
        "        array = np.fromiter(value, dtype=np.dtype('S2'))\n",
        "    else:\n",
        "        array = np.fromiter(value, dtype=np.dtype('U2'))\n",
        "\n",
        "    array = array.reshape(n_frames, n_residues)\n",
        "    array[:, np.logical_not(protein_indices)] = 'NA'\n",
        "    return array\n",
        "\n",
        "# remove_waters = \"yes\" #@param [\"yes\", \"no\" ]\n",
        "# if remove_waters == \"yes\":\n",
        "#   no_waters = \"nowat\"\n",
        "# else:\n",
        "#   no_waters = ''\n",
        "\n",
        "starting = os.path.join(workDir, \"starting.pdb\")\n",
        "starting2 = os.path.join(workDir, \"starting2.pdb\")\n",
        "starting_end = os.path.join(workDir, \"starting_end.pdb\")\n",
        "ppdb = PandasPdb().read_pdb(outfnm)\n",
        "ppdb.df['ATOM'] = ppdb.df['ATOM']\n",
        "ppdb.df['HETATM'] = ppdb.df['HETATM'][ppdb.df['HETATM']['residue_name'] == 'HOH']\n",
        "ppdb.df['ATOM'] = ppdb.df['ATOM'][ppdb.df['ATOM']['atom_name'] != 'OXT']\n",
        "ppdb.df['ATOM']= ppdb.df['ATOM'][ppdb.df['ATOM']['element_symbol'] != 'H']\n",
        "ppdb.to_pdb(path=starting, records=['ATOM', 'HETATM'], gz=False, append_newline=True)\n",
        "\n",
        "from Bio.PDB import is_aa\n",
        "from Bio.PDB import PDBParser, PDBIO, Select\n",
        "\n",
        "\n",
        "class ProtSelect(Select):\n",
        "    def accept_residue(self, residue):\n",
        "        print(f\"{residue} -> {is_aa(residue)}\")\n",
        "        return is_aa(residue, standard=True)\n",
        "\n",
        "\n",
        "from Bio import PDB\n",
        "\n",
        "pdb_ini = PDBParser().get_structure(\"pdb\", starting)\n",
        "io = PDBIO()\n",
        "io.set_structure(pdb_ini)\n",
        "io.save(starting2, ProtSelect());\n",
        "\n",
        "\n",
        "\n",
        "pdb4amber_cmd = \"pdb4amber -i \" + str(starting2) + \" -o \" + str(starting_end) + \" -p\"\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "\n",
        "with open('pdb4amber.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(pdb4amber_cmd)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "\n",
        "subprocess.run([\"chmod 700 pdb4amber.sh\"], shell=True)\n",
        "subprocess.run([\"./pdb4amber.sh\"], shell=True,)\n",
        "\n",
        "\n",
        "pdb_ss = md.load_pdb(starting2)\n",
        "ss = compute_dssp(pdb_ss, simplified=True,)\n",
        "\n",
        "my_string = ''.join([str(item) for sublist in ss for item in sublist])\n",
        "# print(\"Secondary Structure = \" + str(my_string))\n",
        "\n",
        "#@markdown ---"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "v7CLq3rVWUcn"
      },
      "outputs": [],
      "source": [
        "#@title **Show 3D structure**\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "import py3Dmol\n",
        "\n",
        "color = \"rainbow\" #@param [\"gray\", \"rainbow\"]\n",
        "show_sidechains = False #@param {type:\"boolean\"}\n",
        "show_mainchains = False #@param {type:\"boolean\"}\n",
        "show_box = True #@param {type:\"boolean\"}\n",
        "opacity = 0.3 #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
        "\n",
        "def show_pdb(show_sidechains=False, show_mainchains=False, show_box = False, color=\"rainbow\"):\n",
        "  view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',)\n",
        "  view.addModel(open(starting2,'r').read(),'pdb')\n",
        "\n",
        "  if color == \"gray\":\n",
        "    view.setStyle({'cartoon':{}})\n",
        "  elif color == \"rainbow\":\n",
        "    view.setStyle({'cartoon': {'color':'spectrum'}})\n",
        "\n",
        "  if show_sidechains:\n",
        "    BB = ['C','O','N']\n",
        "    view.addStyle({'and':[{'resn':[\"GLY\",\"PRO\"],'invert':True},{'atom':BB,'invert':True}]},\n",
        "                        {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "    view.addStyle({'and':[{'resn':\"PW\"}]},\n",
        "                        {'sphere':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "    view.addStyle({'and':[{'resn':\"PRO\"},{'atom':['C','O'],'invert':True}]},\n",
        "                        {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "  if show_mainchains:\n",
        "    BB = ['C','O','N','CA']\n",
        "    view.addStyle({'atom':BB},{'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "\n",
        "  if show_box:\n",
        "    view.addSurface(py3Dmol.SAS, {'opacity': opacity, 'color':'white'})\n",
        "  view.zoomTo()\n",
        "  return view\n",
        "\n",
        "print(\"Secondary Structure = \" + str(my_string))\n",
        "show_pdb(show_sidechains, show_mainchains, show_box, color).show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "LHAQVag1aSX4"
      },
      "outputs": [],
      "source": [
        "#@title **Parameters to generate the Martini topology:**\n",
        "import glob\n",
        "\n",
        "top_martini = os.path.join(workDir, \"martini.top\")\n",
        "top_temp = os.path.join(workDir, \"martini_temp.top\")\n",
        "gro_martini = os.path.join(workDir, \"system.gro\")\n",
        "pdb_nw_martini = os.path.join(workDir, \"martini.pdb\")\n",
        "pdb_martini = os.path.join(workDir, \"ions.pdb\")\n",
        "\n",
        "force_field = \"martini3001\" #@param [\"martini3001\", \"martini22\"]\n",
        "#@markdown Position Restraints:\n",
        "position_restraints = \"backbone\" #@param [\"none\",\"all\",\"backbone\"]\n",
        "#@markdown Position restraints force constant in kJ/mol (default: 1000):\n",
        "posres_force = 1000 #@param {type:\"slider\", min:0, max:2000, step:100}\n",
        "secondary_structure_handling = \"yes\" #@param [\"yes\",\"no\"]\n",
        "if secondary_structure_handling == \"yes\":\n",
        "  ss_h = my_string\n",
        "else:\n",
        "  ss_h = ''\n",
        "\n",
        "#@markdown Protein elastic network:\n",
        "write_elastic_bonds = \"False\" #@param [\"True\",\"False\"]\n",
        "#@markdown Elastic bond force constant Fc in kJ/mol (default: 500):\n",
        "elastic_force = 0 #@param {type:\"slider\", min:0, max:2000, step:100}\n",
        "#@markdown Elastic bond lower cutoff (default: 0):\n",
        "elastic_lower = 0 #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
        "#@markdown Elastic bond upper cutoff (default: 0.9):\n",
        "elastic_upper = 0.9 #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
        "if write_elastic_bonds == \"True\":\n",
        "  elastic = \" -elastic -ef \" + str(elastic_force) + \" -el \" + str(elastic_lower) + \" -eu \" + str(elastic_upper)\n",
        "else:\n",
        "  elastic = ''\n",
        "\n",
        "side_chain_corrections = \"yes\" #@param [\"yes\",\"no\"]\n",
        "if side_chain_corrections == \"yes\":\n",
        "  scfix = \" -scfix\"\n",
        "else:\n",
        "  scfix = ''\n",
        "cystein_bonds = \"yes\" #@param [\"yes\",\"no\"]\n",
        "if cystein_bonds == \"yes\":\n",
        "  cys = \"auto\"\n",
        "else:\n",
        "  cys = \"none\"\n",
        "\n",
        "if os.path.exists(top_martini):\n",
        "  os.remove(top_martini)\n",
        "  os.remove(gro_martini)\n",
        "  os.remove(pdb_nw_martini)\n",
        "else:\n",
        "  pass\n",
        "\n",
        "folder_workdir = \"cd \" + str(workDir)\n",
        "martinize2 = \"martinize2 -f \" + str(starting2) + \" -x \" + str(pdb_nw_martini)  + \" -ss \" + str(ss_h) + \" -ff \" + str(force_field) +  str(elastic) + \" -p \" + str(position_restraints) +  \" -pf \" + str(posres_force) + \" -cys \" + str(cys) + str(scfix) +  \" -o \" + str(top_martini) + \" -maxwarn 10 -ignh\"\n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "with open('martinize.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(folder_workdir)\n",
        "    print(martinize2)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "\n",
        "os.system(\"chmod 700 martinize.sh\")\n",
        "os.system(\"bash martinize.sh\")\n",
        "\n",
        "#@markdown **Parameters to solvate the system:**\n",
        "\n",
        "box_type = \"cubic\" #@param [\"square\", \"cubic\", \"optimal\"]\n",
        "\n",
        "#@markdown Minimum distance from the protein to box edge (nm):\n",
        "\n",
        "size_box = 1.5 #@param {type:\"slider\", min:0, max:10, step:0.1}\n",
        "\n",
        "#@markdown **ATTENTION**: Give the concentration in Molar units (NaCl):\n",
        "\n",
        "Concentration = 0.15 #@param {type:\"slider\", min:0, max:2, step:0.01}\n",
        "\n",
        "os.system(\"insane -o \" + str(gro_martini) + \" -p \" + str(top_temp) + \" -f \"  + str(pdb_nw_martini) + \" -d \" + str(size_box) + \" -sol W -salt \" + str(Concentration) + \" -charge auto -pbc \" + str(box_type))\n",
        "\n",
        "with open(top_martini, 'r') as f:\n",
        "    lines = f.readlines()\n",
        "# Remove the first line\n",
        "lines.pop(0)\n",
        "lines.pop(0)\n",
        "if force_field == \"martini3001\":\n",
        "  new_lines = ['#include \"martini/martini_v3.0.0.itp\"\\n', '#include \"martini/martini_v3.0.0_ions.itp\"\\n', '#include \"martini/martini_v3.0.0_solvents.itp\"\\n']\n",
        "elif force_field == \"martini22\":\n",
        "  new_lines = ['#include \"martini/martini_v2.2.itp\"\\n', '#include \"martini/martini_v2.2_ions.itp\"\\n', '#include \"martini/martini_v2.2_aminoacids.itp\"\\n']\n",
        "else:\n",
        "  pass\n",
        "lines = new_lines + lines\n",
        "# Write the updated content back to the file\n",
        "with open(top_martini, 'w') as f:\n",
        "    f.writelines(lines)\n",
        "\n",
        "\n",
        "if Concentration == 0:\n",
        "  with open(top_temp, \"r\") as file:\n",
        "    lines = file.readlines()\n",
        "  last_line = lines[-1].split()\n",
        "  if last_line[0] == \"W\":\n",
        "      with open(top_temp) as f_in, open(top_martini, \"a\") as f_out :\n",
        "        for row in f_in.readlines()[-1:] :\n",
        "          f_out.write(row)\n",
        "  else:\n",
        "      with open(top_temp) as f_in, open(top_martini, \"a\") as f_out :\n",
        "        for row in f_in.readlines()[-2:] :\n",
        "          f_out.write(row)\n",
        "else:\n",
        "  with open(top_temp) as f_in, open(top_martini, \"a\") as f_out :\n",
        "      for row in f_in.readlines()[-3:] :\n",
        "          f_out.write(row)\n",
        "\n",
        "universe = mda.Universe(gro_martini)\n",
        "with mda.Writer(pdb_martini) as pdb:\n",
        "  pdb.write(universe)\n",
        "\n",
        "os.system(\"rm \" + str(workDir) + \"/*#\")\n",
        "os.remove(top_temp)\n",
        "\n",
        "gro_check = os.path.exists(gro_martini)\n",
        "top_check = os.path.exists(top_martini)\n",
        "\n",
        "if gro_check == True and top_check == True:\n",
        "  print(\"Successfully generated topology! :-)\")\n",
        "else:\n",
        "  print(\"ERROR: Check your input file! \")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C8kKR7bpI86W"
      },
      "source": [
        "## Let's take a look on our simulation box:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "oDYDCmc1Vn9e"
      },
      "outputs": [],
      "source": [
        "#@title **Check your conversion to Martini**\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "import py3Dmol\n",
        "show_waters = True #@param {type:\"boolean\"}\n",
        "opacity = 0.4 #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
        "\n",
        "if show_waters == False:\n",
        "  input_file = pdb_nw_martini\n",
        "else:\n",
        "  input_file = pdb_martini\n",
        "\n",
        "def show_pdb(color=\"rainbow\"):\n",
        "  view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',)\n",
        "  view.addModel(open(input_file,'r').read(),'pdb')\n",
        "  view.setStyle({'sphere': {'color':'spectrum'}})\n",
        "  view.addSurface(py3Dmol.SAS, {'opacity': opacity, 'color':'white'})\n",
        "  view.zoomTo()\n",
        "  return view\n",
        "\n",
        "show_pdb(color).show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n85MrAO7M7uQ"
      },
      "source": [
        "---\n",
        "---\n",
        "# **Equilibrating the simulation box**\n",
        "\n",
        "Proper MD equilibration protocol is designed to equilibrate both temperature and pressure throughout the simulation box while preserving the protein experimental conformation. In addition, we also allow the solvent to accomodate around the protein, creating proper solvation layers.\n",
        "\n",
        "Below, we will set up the MD equilibration parameters, such as temperature, pressure and the desired simulation time. We will define the force constant used to restraint protein heavy-atoms in place and the frequency at which we want to save atomic coordinates in a trajectory file (.dcd).\n",
        "\n",
        "After you are done, you can run the next 2 cells to equilibrate your system on NVT (first) and NPT ensembles."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "8x9Qp_dbr9HP"
      },
      "outputs": [],
      "source": [
        "#@title ### **Parameters for MD Equilibration (NVT and NPT):**\n",
        "\n",
        "# remove whitespaces\n",
        "Jobname = '1UBQ' #@param {type:\"string\"}\n",
        "\n",
        "Minimization_steps = \"5000\" #@param [\"1000\", \"5000\", \"10000\", \"20000\", \"50000\", \"100000\"]\n",
        "\n",
        " #@markdown Integration time (in femtoseconds):\n",
        "Integration_timestep = \"20\" #@param [\"5\", \"10\", \"20\", \"30\", \"40\"]\n",
        "dt_eq = Integration_timestep\n",
        "\n",
        "#@markdown Simulation time (in nanoseconds) for NVT equilibration:\n",
        "Time = \"5\" #@param {type:\"string\"}\n",
        "stride_time_eq_nvt = Time\n",
        "#@markdown Simulation time (in nanoseconds) for NPT equilibration:\n",
        "Time = \"5\" #@param {type:\"string\"}\n",
        "stride_time_eq_npt = Time\n",
        "\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 = 500 #@param {type:\"slider\", min:0, max:2000, step:100}\n",
        "\n",
        "#@markdown Frequency to write the trajectories files (in picoseconds):\n",
        "\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",
        "\n",
        "Write_the_log = \"10\" #@param [\"10\", \"100\", \"200\", \"500\", \"1000\"]\n",
        "write_the_log_eq = Write_the_log\n",
        "\n",
        "#@markdown ---\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "OiUDr7xkAkoU"
      },
      "outputs": [],
      "source": [
        "#@title **Runs an Equilibration coarse-grained MD simulation (NVT and NPT ensembles)**\n",
        "#@markdown Now, let's run 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 martini_openmm as martini\n",
        "from mdtraj.reporters import XTCReporter\n",
        "import pytraj as pt\n",
        "import parmed\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 = os.path.join(workDir, \"system.gro\")\n",
        "pdbfile = os.path.join(workDir, \"ions.pdb\")\n",
        "topologyfile = os.path.join(workDir, \"martini.top\")\n",
        "\n",
        "time_ps = float(stride_time_eq_nvt)*1000\n",
        "simulation_time = float(time_ps)*picosecond\t\t# in ps\n",
        "dt = int(dt_eq)*femtosecond\n",
        "temperature = float(temperature_eq)*kelvin\n",
        "savcrd_freq = int(write_the_trajectory_eq)*picosecond\n",
        "print_freq  = int(write_the_log_eq)*picosecond\n",
        "\n",
        "pressure\t= float(pressure_eq)*bar\n",
        "\n",
        "# restraint_fc = int(Force_constant) # kJ/mol\n",
        "\n",
        "nsteps  = int(simulation_time.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nprint  = int(print_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nsavcrd = int(savcrd_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "\n",
        "#############################################\n",
        "# Defining functions to use below:\n",
        "def backup_old_log(pattern, string):\n",
        "\tresult = []\n",
        "\tfor root, dirs, files in os.walk(\"./\"):\n",
        "\t\tfor name in files:\n",
        "\t\t\tif fnmatch.fnmatch(name, pattern):\n",
        "\n",
        "\t\t\t\ttry:\n",
        "\t\t\t\t\tnumber = int(name[-2])\n",
        "\t\t\t\t\tavail = isinstance(number, int)\n",
        "\t\t\t\t\t#print(name,avail)\n",
        "\t\t\t\t\tif avail == True:\n",
        "\t\t\t\t\t\tresult.append(number)\n",
        "\t\t\t\texcept:\n",
        "\t\t\t\t\tpass\n",
        "\n",
        "\tif len(result) > 0:\n",
        "\t\tmaxnumber = max(result)\n",
        "\telse:\n",
        "\t\tmaxnumber = 0\n",
        "\n",
        "\tbackup_file = \"\\#\" + string + \".\" + str(maxnumber + 1) + \"#\"\n",
        "\tos.system(\"mv \" + string + \" \" + backup_file)\n",
        "\treturn backup_file\n",
        "\n",
        "# def restraints(system, crd, fc, restraint_array):\n",
        "\n",
        "# \tboxlx = system.getDefaultPeriodicBoxVectors()[0][0].value_in_unit(nanometers)\n",
        "# \tboxly = system.getDefaultPeriodicBoxVectors()[1][1].value_in_unit(nanometers)\n",
        "# \tboxlz = system.getDefaultPeriodicBoxVectors()[2][2].value_in_unit(nanometers)\n",
        "\n",
        "# \tif fc > 0:\n",
        "# \t\t# positional restraints for all heavy-atoms\n",
        "# \t\tposresPROT = CustomExternalForce('k*periodicdistance(x, y, z, x0, y0, z0)^2;')\n",
        "# \t\tposresPROT.addPerParticleParameter('k')\n",
        "# \t\tposresPROT.addPerParticleParameter('x0')\n",
        "# \t\tposresPROT.addPerParticleParameter('y0')\n",
        "# \t\tposresPROT.addPerParticleParameter('z0')\n",
        "\n",
        "# \t\tfor atom1 in restraint_array:\n",
        "# \t\t\tatom1 = int(atom1)\n",
        "\n",
        "# \t\t\txpos  = crd.positions[atom1].value_in_unit(nanometers)[0]\n",
        "# \t\t\typos  = crd.positions[atom1].value_in_unit(nanometers)[1]\n",
        "# \t\t\tzpos  = crd.positions[atom1].value_in_unit(nanometers)[2]\n",
        "\n",
        "# \t\t\tposresPROT.addParticle(atom1, [fc, xpos, ypos, zpos])\n",
        "\n",
        "# \t\tsystem.addForce(posresPROT)\n",
        "\n",
        "# \treturn system\n",
        "##############################################\n",
        "\n",
        "#############################################\n",
        "print(\"\\n> Simulation details:\\n\")\n",
        "print(\"\\tJob name = \" + jobname)\n",
        "print(\"\\tCoordinate file = \" + str(coordinatefile))\n",
        "print(\"\\tPDB file = \" + str(pdbfile))\n",
        "print(\"\\tTopology file = \" + str(topologyfile))\n",
        "\n",
        "print(\"\\n\\tSimulation_time = \" + str(simulation_time))\n",
        "print(\"\\tIntegration timestep = \" + str(dt))\n",
        "print(\"\\tTotal number of steps = \" +  str(nsteps))\n",
        "\n",
        "print(\"\\n\\tSave coordinates each \" + str(savcrd_freq))\n",
        "print(\"\\tPrint in log file each \" + str(print_freq))\n",
        "\n",
        "print(\"\\n\\tTemperature = \" + str(temperature))\n",
        "#############################################\n",
        "\n",
        "# get any defines\n",
        "defines = {}\n",
        "try:\n",
        "\twith open(\"defines.txt\") as def_file:\n",
        "\t\tfor line in def_file:\n",
        "\t\t\tline = line.strip()\n",
        "\t\t\tdefines[line] = True\n",
        "except FileNotFoundError:\n",
        "\tpass\n",
        "\n",
        "print(\"\\n> Setting the system:\\n\")\n",
        "\n",
        "print(\"\\t- Reading topology and structure file...\")\n",
        "gro = GromacsGroFile(coordinatefile)\n",
        "top = martini.MartiniTopFile(topologyfile, periodicBoxVectors=gro.getPeriodicBoxVectors(), defines=defines)\n",
        "\n",
        "\n",
        "print(\"\\t- Creating system and setting parameters...\")\n",
        "friction = 10.0 / picosecond\n",
        "system = top.create_system(nonbonded_cutoff=1.1*nanometer)\n",
        "\n",
        "# print(\"\\t- Applying restraints. Force Constant = \" + str(Force_constant) + \"kJ/mol\")\n",
        "# pt_system = pt.iterload(coordinatefile, pdb)\n",
        "# pt_topology = pt_system.top\n",
        "# restraint_array = pt.select_atoms('!(:SOL) & !(:NA) & !(:CL) & !(:MG) & !(:K)', pt_topology)\n",
        "\n",
        "# system = restraints(system, gro, 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.setRandomNumberSeed(0)\n",
        "# integrator.setConstraintTolerance(constraintTolerance)\n",
        "simulation = Simulation(top.topology, system, integrator)\n",
        "simulation.context.setPositions(gro.positions)\n",
        "\n",
        "parmed_structure = parmed.openmm.load_topology(top.topology, system, gro.positions)\n",
        "\n",
        "parmed_structure.save(os.path.join(workDir, 'system_openMM.pdb'), overwrite=True)\n",
        "with open(os.path.join(workDir, 'system_openMM.xml'), 'w') as output:\n",
        "  output.write(XmlSerializer.serialize(system))\n",
        "\n",
        "print(\"\\t- Energy minimization: \" + str(Minimization_steps) + \" steps\")\n",
        "simulation.minimizeEnergy(tolerance=1.0, 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 NVT ensemble\n",
        "\n",
        "dcd_file = jobname + \"_eq_nvt.dcd\"\n",
        "log_file = jobname + \"_eq_nvt.log\"\n",
        "rst_file_nvt = jobname + \"_eq_nvt.rst\"\n",
        "prv_rst_file = jobname + \"_eq_nvt.rst\"\n",
        "pdb_file = jobname + \"_eq_nvt.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_nvt) + \")...\")\n",
        "state = simulation.context.getState( getPositions=True, getVelocities=True )\n",
        "with open(rst_file_nvt, '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",
        "\n",
        "#############################################\n",
        "# Running Equilibration on NPT ensemble\n",
        "\n",
        "time_ps = float(stride_time_eq_npt)*1000\n",
        "simulation_time = float(time_ps)*picosecond\t\t# in ps\n",
        "dt = int(dt_eq)*femtosecond\n",
        "temperature = float(temperature_eq)*kelvin\n",
        "savcrd_freq = int(write_the_trajectory_eq)*picosecond\n",
        "print_freq  = int(write_the_log_eq)*picosecond\n",
        "pressure\t= float(pressure_eq)*bar\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",
        "print(\"\\tPressure = \" + str(pressure))\n",
        "print(\"\\t- Setting barostat...\")\n",
        "system.addForce(MonteCarloBarostat(pressure, temperature))\n",
        "\n",
        "print(\"\\n> Loading previous state from equilibration > \" + rst_file_nvt + \" <\")\n",
        "with open(rst_file_nvt, 'r') as f:\n",
        "\tsimulation.context.setState(XmlSerializer.deserialize(f.read()))\n",
        "\tcurrstep = int((1-1)*nsteps)\n",
        "\tcurrtime = currstep*dt.in_units_of(picosecond)\n",
        "\tsimulation.currentStep = currstep\n",
        "\tsimulation.context.setTime(currtime)\n",
        "\tprint(\"> Current time: \" + str(currtime) + \" (Step = \" + str(currstep) + \")\")\n",
        "\n",
        "dcd_file = jobname + \"_eq_npt.dcd\"\n",
        "log_file = jobname + \"_eq_npt.log\"\n",
        "rst_file_npt = jobname + \"_eq_npt.rst\"\n",
        "prv_rst_file = jobname + \"_eq_npt.rst\"\n",
        "pdb_file = jobname + \"_eq_npt.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",
        "# Writing last frame information of stride\n",
        "print(\"\\n> Writing state file (\" + str(rst_file_npt) + \")...\")\n",
        "state = simulation.context.getState( getPositions=True, getVelocities=True )\n",
        "with open(rst_file_npt, '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\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LXyL26HCO8Bu"
      },
      "source": [
        "---\n",
        "---\n",
        "# **Running a Production MD simulation**\n",
        "\n",
        "Finally, we will proceed with the Production simulation itself using the equilibrated system coordinates as input structure.\n",
        "\n",
        "Note that we will use here a *.rst state file* , which contains atomic velocities and positions from the last frame of the equilibration simulation, guaranteeing that our production simulation begins from a thermodynamically equilibrated system.\n",
        "\n",
        "Another important information here is the **Number_of_strides** and the **Stride_Time**. In this notebook, we simulate a defined number of *strides*, so the **simulation time = Number_of_strides*Stride_Time**. For example, we can simulate 100ns by setting *Number_of_strides=10* and *Stride_Time=10 ns*.\n",
        "\n",
        "**Important: at the end of the Production simulation, we concatenate all strides to create a complete trajectory file which can be visualized and analyzed**\n",
        "\n",
        "The idea behind this approach is to make use of the intermitent 12h/24h period that Google Colab allows us to use its GPUs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "MXPXFKQIQwmk"
      },
      "outputs": [],
      "source": [
        "# Equilibrated_PDB = '1aki_equil.pdb' #@param {type:\"string\"}\n",
        "# State_file = '1aki_equil.rst' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown ### **Parameters for MD Production protocol:**\n",
        "#@markdown **Important**: The **Jobname** should be the same you have been used to run your simulation in the previous step.\n",
        "Jobname = '1UBQ' #@param {type:\"string\"}\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 = \"10\" #@param {type:\"string\"}\n",
        "nstride = Number_of_strides\n",
        "Integration_timestep = \"20\" #@param [\"5\", \"10\", \"20\", \"30\", \"40\"]\n",
        "\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 = \"200\" #@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",
        "continue_simulation = \"no\" #@param [\"yes\", \"no\"]\n",
        "\n",
        "#@markdown **Important:** If you are going to continue your simulation, all the parameters should be the same, including the **Jobname**. The only parameter you should change is the **Number of strides**, which you will increase.\n",
        "\n",
        "\n",
        "#@markdown ---"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "efprpueWaLrO"
      },
      "outputs": [],
      "source": [
        "#@title **Runs a Production coarse-grained 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",
        "import martini_openmm as martini\n",
        "from mdtraj.reporters import XTCReporter\n",
        "import pytraj as pt\n",
        "\n",
        "from sys import stdout, exit, stderr\n",
        "import os, math, fnmatch\n",
        "#############################################\n",
        "\n",
        "jobname = os.path.join(workDir, str(Jobname))\n",
        "coordinatefile = os.path.join(workDir, \"system.gro\")\n",
        "pdbfile = os.path.join(workDir, \"ions.pdb\")\n",
        "topologyfile = os.path.join(workDir, \"martini.top\")\n",
        "\n",
        "stride_time_ps = float(stride_time_prod)*1000\n",
        "stride_time = float(stride_time_ps)*picosecond\n",
        "nstride = int(Number_of_strides)\n",
        "dt = int(dt_prod)*femtosecond\n",
        "temperature = float(temperature_prod)*kelvin\n",
        "savcrd_freq = int(write_the_trajectory_prod)*picosecond\n",
        "print_freq  = int(write_the_log_prod)*picosecond\n",
        "\n",
        "pressure\t= float(pressure_prod)*bar\n",
        "\n",
        "simulation_time = stride_time*nstride\n",
        "nsteps  = int(stride_time.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nprint  = int(print_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nsavcrd = int(savcrd_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "firststride = 1 # must be integer\n",
        "\n",
        "# get any defines\n",
        "defines = {}\n",
        "try:\n",
        "\twith open(\"defines.txt\") as def_file:\n",
        "\t\tfor line in def_file:\n",
        "\t\t\tline = line.strip()\n",
        "\t\t\tdefines[line] = True\n",
        "except FileNotFoundError:\n",
        "\tpass\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",
        "\n",
        "#############################################\n",
        "print(\"\\n> Simulation details:\\n\")\n",
        "print(\"\\tJob name = \" + jobname)\n",
        "print(\"\\tCoordinate file = \" + str(coordinatefile))\n",
        "print(\"\\tPDB file = \" + str(pdbfile))\n",
        "print(\"\\tTopology file = \" + str(topologyfile))\n",
        "\n",
        "print(\"\\n\\tSimulation_time = \" + str(stride_time*nstride))\n",
        "print(\"\\tIntegration timestep = \" + str(dt))\n",
        "print(\"\\tTotal number of steps = \" +  str(nsteps*nstride))\n",
        "print(\"\\tNumber of strides = \" + str(nstride) + \" (\" + str(stride_time) + \" in each stride)\")\n",
        "\n",
        "print(\"\\n\\tSave coordinates each \" + str(savcrd_freq))\n",
        "print(\"\\tSave checkpoint each \" + str(savcrd_freq))\n",
        "print(\"\\tPrint in log file each \" + str(print_freq))\n",
        "\n",
        "print(\"\\n\\tTemperature = \" + str(temperature))\n",
        "print(\"\\tPressure = \" + str(pressure))\n",
        "#############################################\n",
        "\n",
        "print(\"\\n> Setting the system:\\n\")\n",
        "\n",
        "if continue_simulation == \"yes\":\n",
        "\tfor i in range(1,int(nstride)+1):\n",
        "\t\tname = str(jobname) + \"_\" + str(i) + \"_prod_npt.pdb\"\n",
        "\t\tname2 = str(jobname) + \"_\" + str(i) + \"_prod_npt.rst\"\n",
        "\t\tif os.path.exists(name) and os.path.exists(name2):\n",
        "\t\t\tlast_stride_pdb = name\n",
        "\t\t\tlast_stride_rst = name2\n",
        "\t\telse:\n",
        "\t\t\tpass\n",
        "\tgro = PDBFile(last_stride_pdb)\n",
        "\twith open(os.path.join(workDir, 'system_openMM.xml')) as input:\n",
        "  \t\tsystem = XmlSerializer.deserialize(input.read())\n",
        "\ttop = gro.getTopology()\n",
        "else:\n",
        "\ttop = top.topology\n",
        "\n",
        "print(\"\\t- Setting integrator...\")\n",
        "friction = 10.0 / picosecond\n",
        "integrator = LangevinIntegrator(temperature, friction, dt)\n",
        "integrator.setRandomNumberSeed(0)\n",
        "simulation = Simulation(top, system, integrator)\n",
        "simulation.context.setPositions(gro.positions)\n",
        "\n",
        "print(\"\\tPressure = \" + str(pressure))\n",
        "print(\"\\t- Setting barostat...\")\n",
        "system.addForce(MonteCarloBarostat(pressure, temperature))\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",
        "\trst_file_npt = jobname + \"_eq_npt.rst\"\n",
        "\tdcd_file = jobname + \"_\" + str(n) + \"_prod_npt.dcd\"\n",
        "\tlog_file = jobname + \"_\" + str(n) + \"_prod_npt.log\"\n",
        "\trst_file = jobname + \"_\" + str(n) + \"_prod_npt.rst\"\n",
        "\tprv_rst_file = jobname + \"_\" + str(n-1) + \"_prod_npt.rst\"\n",
        "\tpdb_file = jobname + \"_\" + str(n) + \"_prod_npt.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 > \" + rst_file_npt + \" <\")\n",
        "\t\twith open(rst_file_npt, '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",
        "\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\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "kw9yWoU14506"
      },
      "outputs": [],
      "source": [
        "#@title **Concatenate and align the trajectory**\n",
        "#@markdown **Important**: The **Google Drive Path**, **Jobname**, **Number of strides**, **stride time** and **trajectory saved frequency** should be the same you have been used to run your simulation in the previous steps.\n",
        "\n",
        "#@markdown The code will save two trajectories: One with waters and other without waters.\n",
        "import MDAnalysis as mda\n",
        "from MDAnalysis.analysis import align, rms\n",
        "\n",
        "Google_Drive_Path = '/content/drive/MyDrive/' #@param {type:\"string\"}\n",
        "workDir = Google_Drive_Path\n",
        "Jobname = \"1UBQ\" #@param {type: \"string\"}\n",
        "# force_field = \"martini3001\" #@param [\"martini22p\", \"martini3001\", \"martini22\"]\n",
        "Equilibrated_PDB = Jobname + \"_eq_npt.pdb\"\n",
        "Skip = \"1\" #@param [\"1\", \"2\", \"5\", \"10\", \"20\", \"50\"]\n",
        "stride_traj = Skip\n",
        "# Output_format = \"dcd\" #[\"dcd\", \"pdb\", \"trr\", \"xtc\"]\n",
        "first_stride = \"1\" #@param {type:\"string\"}\n",
        "Number_of_strides = \"10\" #@param {type:\"string\"}\n",
        "nstride = int(Number_of_strides)\n",
        "stride_time = \"10\" #@param {type:\"string\"}\n",
        "trajectory_saved_frequency = \"200\" #@param [\"10\", \"100\", \"200\", \"500\", \"1000\"]\n",
        "Write_trajectory = trajectory_saved_frequency\n",
        "traj_save_freq = trajectory_saved_frequency\n",
        "# Remove_waters = \"no\" #@param [\"yes\", \"no\"]\n",
        "# stride_id_as_ref_for_alignment = \"1\" #@param {type: \"string\"}\n",
        "# output_prefix = first_stride+\"-\"+str(int(first_stride)+nstride-1)\n",
        "\n",
        "stride_time_ps = float(stride_time)*1000\n",
        "simulation_time_analysis = stride_time_ps*nstride\n",
        "simulation_ns = float(stride_time)*int(Number_of_strides)\n",
        "number_frames = int(simulation_time_analysis)/int(traj_save_freq)\n",
        "number_frames_analysis = number_frames/int(stride_traj)\n",
        "\n",
        "nw_dcd = os.path.join(workDir, str(Jobname) + \"_cg_nw.dcd\")\n",
        "nw_pdb = os.path.join(workDir, str(Jobname) +  \"_nw.pdb\")\n",
        "backmap_dcd = os.path.join(workDir, str(Jobname) + \"_100_frames_nw.dcd\")\n",
        "whole_pdb = os.path.join(workDir, str(Jobname) +  \"_whole.pdb\")\n",
        "whole_dcd = os.path.join(workDir, str(Jobname) +  \"_cg_whole.dcd\")\n",
        "template =  os.path.join(workDir, str(Jobname) + '_%s_prod_npt.dcd')\n",
        "pdb = os.path.join(workDir, Equilibrated_PDB)\n",
        "\n",
        "# if force_field == \"martini3001\":\n",
        "#   mask_mda = \"not (resname W or name NA+ or name CL-)\"\n",
        "# elif force_field == \"martini22\":\n",
        "#   mask_mda = \"not (resname W or name NA+ or name CL-)\"\n",
        "# else:\n",
        "#   pass\n",
        "\n",
        "mask_mda = \"not (resname W or name NA+ or name CL-)\"\n",
        "\n",
        "flist = [template % str(i) for i in range(int(first_stride), int(first_stride) + nstride)]\n",
        "ref = [template % int(1)]\n",
        "\n",
        "u1 = mda.Universe(pdb, flist)\n",
        "u2 = mda.Universe(pdb, ref)\n",
        "\n",
        "num_frames = len(u1.trajectory)\n",
        "if num_frames >= 10:\n",
        "  stride_backmapping = num_frames/10\n",
        "else:\n",
        "  \"ATTENTION: You need at least 100 frames to proceed, please increase your MD simulation time.\"\n",
        "\n",
        "u2.trajectory[0] # set u2 to first frame\n",
        "\n",
        "# print(rms.rmsd(u1.select_atoms('name CA').positions, u2.select_atoms('name CA').positions, superposition=False))\n",
        "\n",
        "align.AlignTraj(u1, u2, select=mask_mda, in_memory=True).run()\n",
        "\n",
        "nw = u1.select_atoms(mask_mda)\n",
        "\n",
        "#Remove the Waters\n",
        "with mda.Writer(nw_dcd, nw.n_atoms) as W:\n",
        "  for ts in u1.trajectory[::int(Skip)]:\n",
        "      W.write(nw, )\n",
        "not_waters = u2.select_atoms(mask_mda)\n",
        "not_waters.write(nw_pdb)\n",
        "traj_dcd_check = os.path.exists(nw_dcd)\n",
        "traj = nw_dcd\n",
        "pdb_ref = nw_pdb\n",
        "#For Backmapping\n",
        "# with mda.Writer(backmap_dcd, nw.n_atoms) as W:\n",
        "#   for ts in u1.trajectory[::int(stride_backmapping)]:\n",
        "#   # for ts in u1.trajectory[0:100:int(1)]:\n",
        "#       W.write(nw, )\n",
        "#keep the waters\n",
        "with mda.Writer(whole_dcd, u1.select_atoms(\"all\").n_atoms) as W:\n",
        "  for ts in u1.trajectory[::int(Skip)]:\n",
        "      W.write(u1.select_atoms(\"all\"))\n",
        "whole = u2.select_atoms(\"all\")\n",
        "whole.write(whole_pdb)\n",
        "traj_dcd_check = os.path.exists(whole_dcd)\n",
        "  # traj = whole_dcd\n",
        "  # pdb_ref = whole_pdb\n",
        "\n",
        "traj_load_cg = pt.load(traj, pdb_ref)\n",
        "print(traj_load_cg)\n",
        "\n",
        "if traj_dcd_check == True:\n",
        "  print(\"Trajectory concatenated successfully! :-)\")\n",
        "else:\n",
        "  print(\"ERROR: Check your inputs! \")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "DD9LqaCx5E6L"
      },
      "outputs": [],
      "source": [
        "#@title **Load, view and check the trajectory**\n",
        "#@markdown This will take a few minutes. Another coffee would be great. :-)\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "\n",
        "#py3dmol functions\n",
        "class Atom(dict):\n",
        "  def __init__(self, line):\n",
        "    self[\"type\"] = line[0:6].strip()\n",
        "    self[\"idx\"] = line[6:11].strip()\n",
        "    self[\"name\"] = line[12:16].strip()\n",
        "    self[\"resname\"] = line[17:20].strip()\n",
        "    self[\"resid\"] = int(int(line[22:26]))\n",
        "    self[\"x\"] = float(line[30:38])\n",
        "    self[\"y\"] = float(line[38:46])\n",
        "    self[\"z\"] = float(line[46:54])\n",
        "    self[\"sym\"] = line[76:78].strip()\n",
        "\n",
        "  def __str__(self):\n",
        "    line = list(\" \" * 80)\n",
        "    line[0:6] = self[\"type\"].ljust(6)\n",
        "    line[6:11] = self[\"idx\"].ljust(5)\n",
        "    line[12:16] = self[\"name\"].ljust(4)\n",
        "    line[17:20] = self[\"resname\"].ljust(3)\n",
        "    line[22:26] = str(self[\"resid\"]).ljust(4)\n",
        "    line[30:38] = str(self[\"x\"]).rjust(8)\n",
        "    line[38:46] = str(self[\"y\"]).rjust(8)\n",
        "    line[46:54] = str(self[\"z\"]).rjust(8)\n",
        "    line[76:78] = self[\"sym\"].rjust(2)\n",
        "    return \"\".join(line) + \"\\n\"\n",
        "\n",
        "class Molecule(list):\n",
        "  def __init__(self, file):\n",
        "    for line in file:\n",
        "      if \"ATOM\" in line or \"HETATM\" in line:\n",
        "        self.append(Atom(line))\n",
        "\n",
        "    def __str__(self):\n",
        "      outstr = \"\"\n",
        "      for at in self:\n",
        "        outstr += str(at)\n",
        "      return outstr\n",
        "\n",
        "if number_frames_analysis > 10:\n",
        "  stride_animation = number_frames_analysis/10\n",
        "else:\n",
        "  stride_animation = 1\n",
        "\n",
        "u = mda.Universe(pdb_ref, traj)\n",
        "\n",
        "# Write out frames for animation\n",
        "protein = u.select_atoms('all')\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 = {\"sphere\": {'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()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5mXX1yidlhV"
      },
      "source": [
        "---\n",
        "---\n",
        "# **Back-mapping**\n",
        "\n",
        "Here we are going to use a general approach to transforming molecular models between different levels of resolution, based on machine learning methods. To perform this transformation we will use the [cg2all](https://www.biorxiv.org/content/10.1101/2023.05.22.541652v1) code."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L_ScY4GqDir3",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Conversion of a CG simulation trajectory to an atomistic simulation trajectory using [__convert_cg2all__](https://github.com/huhlim/cg2all/tree/main)\n",
        "#@markdown Back-map a coarse-grained model (trajectory) to a finer level of representation (atomistic).\n",
        "\n",
        "#@markdown This step will take a few minutes. it will be proportional to the number of frames you have in your final coarse-grained trajectory.\n",
        "\n",
        "\n",
        "%%capture\n",
        "%%bash\n",
        "\n",
        "source activate cg2all\n",
        "python\n",
        "import cg2all.lib.libmodel\n",
        "from cg2all.lib.libconfig import MODEL_HOME\n",
        "for model_type in [\"Martini\", \"Martini3\",]:\n",
        "    ckpt_fn = MODEL_HOME / f\"{model_type}.ckpt\"\n",
        "    if not ckpt_fn.exists():\n",
        "        cg2all.lib.libmodel.download_ckpt_file(model_type, ckpt_fn)\n",
        "\n",
        "import os\n",
        "import subprocess\n",
        "import sys\n",
        "\n",
        "Google_Drive_Path = '/content/drive/MyDrive/' #@param {type:\"string\"}\n",
        "workDir = Google_Drive_Path\n",
        "Jobname = \"1UBQ\" #@param {type: \"string\"}\n",
        "\n",
        "Coarse_grained_model = \"Martini3\" #@param [\"Martini\", \"Martini3\"]\n",
        "\n",
        "batch_size = 1 #@param {type: \"number\"}\n",
        "#@markdown - Batch size should be a divisor of the total number of frames.\n",
        "# device = \"cuda\" #@param [\"cpu\", \"cuda\"]\n",
        "\n",
        "traj_input = os.path.join(workDir, str(Jobname) + \"_cg_nw.dcd\")\n",
        "pdb_input = os.path.join(workDir, str(Jobname) +  \"_nw.pdb\")\n",
        "pdb_ini_output = os.path.join(workDir, str(Jobname) +  \"_frame0_backmapped.pdb\")\n",
        "traj_output = os.path.join(workDir, str(Jobname) +  \"_backmapped.dcd\")\n",
        "pdb_output = os.path.join(workDir, str(Jobname) +  \"_backmapped.pdb\")\n",
        "backmap_align_dcd = os.path.join(workDir, str(Jobname) +  \"_backmapped_align.dcd\")\n",
        "backmap_align_pdb = os.path.join(workDir, str(Jobname) +  \"_backmapped_align.pdb\")\n",
        "\n",
        "subprocess.run(\"convert_cg2all -p \" + str(pdb_input) + \" -o \" + str(pdb_ini_output) + \" --cg \" + str(Coarse_grained_model), shell=True)\n",
        "subprocess.run(\"convert_cg2all -p \" + str(pdb_input) + \" --dcd \" + str(traj_input) + \" -o \" + str(traj_output) + \" -opdb \" + str(pdb_output) + \" --cg \" + str(Coarse_grained_model) + \" --batch \" + str(batch_size), shell=True)\n",
        "# print(f\"Converted {traj_input} in Martini to {traj_output}\")\n",
        "\n",
        "\n",
        "# traj_dcd_check = os.path.exists(traj_output)\n",
        "\n",
        "# if traj_dcd_check == True:\n",
        "#   print(\"Trajectory back-mapped successfully! :-)\")\n",
        "#   print(\"Number of frames:\", num_frames)\n",
        "# else:\n",
        "#   print(\"ERROR: Check your inputs! \")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Align the backmapped trajectory**\n",
        "\n",
        "Google_Drive_Path = '/content/drive/MyDrive/' #@param {type:\"string\"}\n",
        "workDir = Google_Drive_Path\n",
        "Jobname = \"1UBQ\" #@param {type: \"string\"}\n",
        "\n",
        "pdb_ini_output = os.path.join(workDir, str(Jobname) +  \"_frame0_backmapped.pdb\")\n",
        "traj_output = os.path.join(workDir, str(Jobname) +  \"_backmapped.dcd\")\n",
        "pdb_output = os.path.join(workDir, str(Jobname) +  \"_backmapped.pdb\")\n",
        "backmap_align_dcd = os.path.join(workDir, str(Jobname) +  \"_backmapped_align.dcd\")\n",
        "backmap_align_pdb = os.path.join(workDir, str(Jobname) +  \"_backmapped_align.pdb\")\n",
        "\n",
        "\n",
        "\n",
        "u1 = mda.Universe(pdb_ini_output, traj_output)\n",
        "u2 = mda.Universe(pdb_ini_output, pdb_ini_output)\n",
        "\n",
        "u2.trajectory[0] # set u2 to first frame\n",
        "\n",
        "align.AlignTraj(u1, u2, select='name CA', in_memory=True).run()\n",
        "\n",
        "sele = u1.select_atoms(\"all\")\n",
        "\n",
        "with mda.Writer(backmap_align_dcd, sele.n_atoms) as W:\n",
        "  for ts in u1.trajectory[::]:\n",
        "      W.write(sele, )\n",
        "sele2 = u2.select_atoms(\"all\")\n",
        "sele2.write(backmap_align_pdb)\n",
        "traj_dcd_check = os.path.exists(backmap_align_dcd)\n",
        "traj_atomistic = backmap_align_dcd\n",
        "pdb_ref_atomistic = backmap_align_pdb\n",
        "\n",
        "# Calculate the length of the trajectory\n",
        "u = mda.Universe(pdb_ref_atomistic,traj_atomistic)\n",
        "num_frames = len(u.trajectory)\n",
        "\n",
        "traj_load_atomistic = pt.load(traj_atomistic, pdb_ref_atomistic)\n",
        "\n",
        "if traj_dcd_check == True:\n",
        "  print(\"Trajectory aligned successfully! :-)\")\n",
        "  print(\"Number of frames:\", num_frames)\n",
        "else:\n",
        "  print(\"ERROR: Check your inputs! \")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "l6st3SFfY4IU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "p3Ys_MW5qOxE"
      },
      "outputs": [],
      "source": [
        "#@title **Load, view and check the back-mapped trajectory**\n",
        "#@markdown This will take a few minutes. Another coffee would be great. :-)\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "\n",
        "#py3dmol functions\n",
        "class Atom(dict):\n",
        "  def __init__(self, line):\n",
        "    self[\"type\"] = line[0:6].strip()\n",
        "    self[\"idx\"] = line[6:11].strip()\n",
        "    self[\"name\"] = line[12:16].strip()\n",
        "    self[\"resname\"] = line[17:20].strip()\n",
        "    self[\"resid\"] = int(int(line[22:26]))\n",
        "    self[\"x\"] = float(line[30:38])\n",
        "    self[\"y\"] = float(line[38:46])\n",
        "    self[\"z\"] = float(line[46:54])\n",
        "    self[\"sym\"] = line[76:78].strip()\n",
        "\n",
        "  def __str__(self):\n",
        "    line = list(\" \" * 80)\n",
        "    line[0:6] = self[\"type\"].ljust(6)\n",
        "    line[6:11] = self[\"idx\"].ljust(5)\n",
        "    line[12:16] = self[\"name\"].ljust(4)\n",
        "    line[17:20] = self[\"resname\"].ljust(3)\n",
        "    line[22:26] = str(self[\"resid\"]).ljust(4)\n",
        "    line[30:38] = str(self[\"x\"]).rjust(8)\n",
        "    line[38:46] = str(self[\"y\"]).rjust(8)\n",
        "    line[46:54] = str(self[\"z\"]).rjust(8)\n",
        "    line[76:78] = self[\"sym\"].rjust(2)\n",
        "    return \"\".join(line) + \"\\n\"\n",
        "\n",
        "class Molecule(list):\n",
        "  def __init__(self, file):\n",
        "    for line in file:\n",
        "      if \"ATOM\" in line or \"HETATM\" in line:\n",
        "        self.append(Atom(line))\n",
        "\n",
        "    def __str__(self):\n",
        "      outstr = \"\"\n",
        "      for at in self:\n",
        "        outstr += str(at)\n",
        "      return outstr\n",
        "\n",
        "backmap_align_dcd = os.path.join(workDir, str(Jobname) +  \"_backmapped_align.dcd\")\n",
        "backmap_align_pdb = os.path.join(workDir, str(Jobname) +  \"_backmapped_align.pdb\")\n",
        "\n",
        "\n",
        "number_frames_analysis = num_frames\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(backmap_align_pdb,backmap_align_dcd)\n",
        "\n",
        "# Write out frames for animation\n",
        "protein = u.select_atoms('all')\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()"
      ]
    },
    {
      "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",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "wBrBMF4Puyv6"
      },
      "outputs": [],
      "source": [
        "#@title **Compute protein RMSD**\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = 'rmsd' #@param {type:\"string\"}\n",
        "\n",
        "nw_dcd = os.path.join(workDir, str(Jobname) + \"_cg_nw.dcd\")\n",
        "nw_pdb = os.path.join(workDir, str(Jobname) +  \"_nw.pdb\")\n",
        "traj_load_cg = pt.load(nw_dcd, nw_pdb)\n",
        "traj_atomistic = os.path.join(workDir, str(Jobname) +  \"_backmapped_align.dcd\")\n",
        "pdb_ref_atomistic = os.path.join(workDir, str(Jobname) +  \"_backmapped_align.pdb\")\n",
        "traj_load_atomistic = pt.load(traj_atomistic, pdb_ref_atomistic)\n",
        "\n",
        "ref_top = pt.load(pdb_ref_atomistic, pdb_ref_atomistic)\n",
        "rmsd_cg = pt.rmsd(traj_load_cg, ref = 0)\n",
        "df = pd.Series(rmsd_cg)\n",
        "running_aver_cg = df.rolling(window =10).mean()\n",
        "rmsd_atomistic = pt.rmsd(traj_load_atomistic, ref = 0, mask = \"@CA,C,O,N,H\")\n",
        "df = pd.Series(rmsd_atomistic)\n",
        "running_aver_atom = df.rolling(window =10).mean()\n",
        "\n",
        "\n",
        "time = len(rmsd_cg)*int(Write_trajectory)/1000\n",
        "time_array = np.arange(0,time,int(Write_trajectory)/1000)*int(stride_traj)\n",
        "\n",
        "# Plotting:\n",
        "fig, axs = plt.subplots(2, 2, figsize=(16, 5))\n",
        "plt.subplot(1, 2, 1)\n",
        "ax = plt.plot(time_array, rmsd_cg, alpha=0.2, color = 'blue', linewidth = 1.0)\n",
        "ax = plt.plot(time_array, running_aver_cg, alpha=0.8, color = 'blue', linewidth = 1.0)\n",
        "plt.xlim(0, simulation_ns)\n",
        "plt.title(\"Coarse-grained\")\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",
        "\n",
        "plt.subplot(1, 2, 2)\n",
        "ax = plt.plot(time_array, rmsd_atomistic, alpha=0.2, color = 'blue', linewidth = 1.0)\n",
        "ax = plt.plot(time_array, running_aver_atom, alpha=0.8, color = 'blue', linewidth = 1.0)\n",
        "plt.xlim(0, simulation_ns)\n",
        "plt.title(\"Back-mapped\")\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",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "raw_data=pd.DataFrame(rmsd_cg)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \"_cg.csv\"))\n",
        "raw_data=pd.DataFrame(rmsd_atomistic)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \"_atomistic.csv\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ZHyMpikjuaLT"
      },
      "outputs": [],
      "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",
        "fig, axs = plt.subplots(2, 2, figsize=(16, 5))\n",
        "plt.subplot(1, 2, 1)\n",
        "ax = sb.kdeplot(rmsd_cg, color=\"blue\", fill=True, alpha=0.2, linewidth=0.5)\n",
        "plt.title(\"Coarse-grained\")\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.subplot(1, 2, 2)\n",
        "ax = sb.kdeplot(rmsd_atomistic, color=\"blue\", fill=True, alpha=0.2, linewidth=0.5)\n",
        "plt.title(\"Back-mapped\")\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')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "CvOFrXGXwXrV"
      },
      "outputs": [],
      "source": [
        "#@title **Compute protein radius of gyration**\n",
        "\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = 'radius_gyration' #@param {type:\"string\"}\n",
        "\n",
        "radgyr_cg = pt.radgyr(traj_load_cg)\n",
        "df = pd.Series(radgyr_cg)\n",
        "running_aver_cg = df.rolling(window =10).mean()\n",
        "radgyr_atom = pt.radgyr(traj_load_atomistic)\n",
        "df = pd.Series(radgyr_atom)\n",
        "running_aver_atom = df.rolling(window =10).mean()\n",
        "\n",
        "time = len(running_aver_cg)*int(Write_trajectory)/1000\n",
        "time_array = np.arange(0,time,int(Write_trajectory)/1000)*int(stride_traj)\n",
        "\n",
        "# Plotting:\n",
        "fig, axs = plt.subplots(2, 2, figsize=(16, 5))\n",
        "plt.subplot(1, 2, 1)\n",
        "plt.plot(time_array, radgyr_cg, alpha=0.2, color = 'green', linewidth = 1.0)\n",
        "plt.plot(time_array, running_aver_cg, alpha=0.8, color = 'green', linewidth = 1.0)\n",
        "plt.xlim(0, simulation_ns)\n",
        "plt.title(\"Coarse-grained\")\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",
        "\n",
        "plt.subplot(1, 2, 2)\n",
        "plt.plot(time_array, radgyr_atom, alpha=0.2, color = 'green', linewidth = 1.0)\n",
        "plt.plot(time_array, running_aver_atom, alpha=0.8, color = 'green', linewidth = 1.0)\n",
        "plt.xlim(0, simulation_ns)\n",
        "plt.title(\"Back-mapped\")\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",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "raw_data=pd.DataFrame(radgyr_cg)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \"_cg.csv\"))\n",
        "raw_data=pd.DataFrame(radgyr_atom)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \"_atomistic.csv\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2Q7FKg8Fuxr9"
      },
      "outputs": [],
      "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",
        "\n",
        "fig, axs = plt.subplots(2, 2, figsize=(16, 5))\n",
        "plt.subplot(1, 2, 1)\n",
        "ax = sb.kdeplot(radgyr_cg, color=\"green\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('Radius of gyration ($\\AA$)', fontsize = 14, fontweight = 'bold')\n",
        "plt.title(\"Coarse-grained\")\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.subplot(1, 2, 2)\n",
        "ax = sb.kdeplot(radgyr_atom, color=\"green\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('Radius of gyration ($\\AA$)', fontsize = 14, fontweight = 'bold')\n",
        "plt.title(\"Back-mapped\")\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')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "p2Y0DgwTxLWc"
      },
      "outputs": [],
      "source": [
        "#@title **Compute RMSF of protein**\n",
        "\n",
        "#@markdown **Provide output file names below:**\n",
        "Output_name = 'rmsf' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "rmsf_cg = pt.rmsf(traj_load_cg, \"byres\")\n",
        "# df = pd.Series(rmsf_cg[:,1])\n",
        "# running_aver_cg = df.rolling(window =5).mean()\n",
        "bfactor_cg = pt.bfactors(traj_load_cg, byres=False)\n",
        "rmsf_atom = pt.rmsf(traj_load_atomistic,\"byres\")\n",
        "bfactor_atom = pt.bfactors(traj_load_atomistic, byres=True)\n",
        "\n",
        "\n",
        "\n",
        "# Plotting:\n",
        "fig, axs = plt.subplots(2, 2, figsize=(16, 5))\n",
        "plt.subplot(1, 2, 1)\n",
        "plt.plot(rmsf_cg[:,1], alpha=1.0, color = 'red', linewidth = 1.0)\n",
        "# plt.plot(running_aver_cg, alpha=1.0, color = 'red', linewidth = 1.0)\n",
        "plt.title(\"Coarse-grained\")\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_cg[:-1]))\n",
        "\n",
        "plt.subplot(1, 2, 2)\n",
        "plt.plot(rmsf_atom[:,1], alpha=1.0, color = 'red', linewidth = 1.0)\n",
        "plt.title(\"Back-mapped\")\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_atom[:-1]))\n",
        "\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_cg)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \"_cg.csv\"))\n",
        "raw_data=pd.DataFrame(rmsf_atom)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \"_atomistic.csv\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "JalicqqrTodW"
      },
      "outputs": [],
      "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",
        "mat_cg = pt.pairwise_rmsd(traj_load_cg, frame_indices=range(int(number_frames_analysis)), metric=\"rms\")\n",
        "mat_atom = pt.pairwise_rmsd(traj_load_atomistic, mask = \"@CA,C,O,N,H\", frame_indices=range(int(number_frames_analysis)), metric=\"rms\")\n",
        "\n",
        "fig, axs = plt.subplots(2, 2, figsize=(16, 5))\n",
        "plt.subplot(1, 2, 1)\n",
        "ax = plt.imshow(mat_cg, cmap = 'PRGn', origin='lower', interpolation = 'bicubic')\n",
        "plt.title(\"Coarse-grained\")\n",
        "plt.xlabel('Time (ns)', fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel('Time (ns)', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(a, b.round(decimals=3), fontsize = 12)\n",
        "plt.yticks(a, b.round(decimals=3), fontsize = 12)\n",
        "# plt.clim(0,3)\n",
        "cbar1 = plt.colorbar()\n",
        "cbar1.set_label(\"RMSD ($\\AA$)\", fontsize = 14, fontweight = 'bold')\n",
        "\n",
        "\n",
        "plt.subplot(1, 2, 2)\n",
        "ax = plt.imshow(mat_atom, cmap = 'PRGn', origin='lower', interpolation = 'bicubic')\n",
        "plt.title(\"Back-mapped\")\n",
        "plt.xlabel('Time (ns)', fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel('Time (ns)', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(a, b.round(decimals=3), fontsize = 12)\n",
        "plt.yticks(a, b.round(decimals=3), fontsize = 12)\n",
        "# plt.clim(0,3)\n",
        "cbar1 = plt.colorbar()\n",
        "cbar1.set_label(\"RMSD ($\\AA$)\", 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_cg)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \"_cg.csv\"))\n",
        "raw_data=pd.DataFrame(mat_atom)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \"_atomistic.csv\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "_mgVSbBshWFV"
      },
      "outputs": [],
      "source": [
        "#@title **Calculate eigvenctors of Principle Component Analysis (PCA)**\n",
        "data_cg = pt.pca(traj_load_cg, fit=True, ref=0, mask='', n_vecs=2)\n",
        "data_atom = pt.pca(traj_load_atomistic, fit=True, ref=0, mask = \"@CA,C,O,N,H\", n_vecs=2)\n",
        "\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",
        "fig, axs = plt.subplots(2, 2, figsize=(16, 5))\n",
        "plt.subplot(1, 2, 1)\n",
        "projection_data = data_cg[0]\n",
        "plt.title(\"Coarse-grained\")\n",
        "PC1_cg = data_cg[0][0]\n",
        "PC2_cg = data_cg[0][1]\n",
        "a = plt.scatter(PC1_cg,PC2_cg, c=range(int(number_frames_analysis)), cmap='Greens', marker='o',s=8, alpha=1)\n",
        "plt.clim(0, last_frame)\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",
        "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.subplot(1, 2, 2)\n",
        "projection_data = data_atom[0]\n",
        "plt.title(\"Back-mapped\")\n",
        "PC1_atom = data_atom[0][0]\n",
        "PC2_atom = data_atom[0][1]\n",
        "a = plt.scatter(PC1_atom,PC2_atom, c=range(int(number_frames_analysis)), cmap='Greens', marker='o',s=8, alpha=1)\n",
        "plt.clim(0, last_frame)\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",
        "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_cg)\n",
        "pc1.to_csv(os.path.join(workDir, Output_PC1 + \"_cg.csv\"))\n",
        "pc2=pd.DataFrame(PC2_cg)\n",
        "pc2.to_csv(os.path.join(workDir, Output_PC2 + \"_cg.csv\"))\n",
        "pc1=pd.DataFrame(PC1_atom)\n",
        "pc1.to_csv(os.path.join(workDir, Output_PC1 + \"_atomistic.csv\"))\n",
        "pc2=pd.DataFrame(PC2_atom)\n",
        "pc2.to_csv(os.path.join(workDir, Output_PC2 + \"__atomistic.csv\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "pTDb7CEfkLq1"
      },
      "outputs": [],
      "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",
        "cc_cg = matrix.correl(traj_load_cg)\n",
        "cc_atom = matrix.correl(traj_load_atomistic, '@CA')\n",
        "\n",
        "fig, axs = plt.subplots(2, 2, figsize=(16, 5))\n",
        "plt.subplot(1, 2, 1)\n",
        "ax = plt.imshow(cc_cg, cmap = 'PiYG_r', interpolation = 'bicubic', vmin = -1, vmax = 1, origin='lower')\n",
        "plt.title(\"Coarse-grained\")\n",
        "plt.xlabel('Beads', fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel('Beads', 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.subplot(1, 2, 2)\n",
        "ax = plt.imshow(cc_atom, cmap = 'PiYG_r', interpolation = 'bicubic', vmin = -1, vmax = 1, origin='lower')\n",
        "plt.title(\"Back-mapped\")\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(cc_cg)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \"_cg.csv\"))\n",
        "raw_data=pd.DataFrame(cc_atom)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \"_atomistic.csv\"))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}