AlphaFold2+MD.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "AlphaFold2+MD.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "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/AlphaFold2%2BMD.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G4yBrceuFbf3"
      },
      "source": [
        "# **Hello there!**\n",
        "\n",
        "This is a Jupyter notebook for running Molecular Dynamics (MD) simulations using OpenMM engine and AMBER force field for **PROTEIN** models from AlphaFold2 pipeline. 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",
        "Easy to use version of AlphaFold 2 [(Jumper et al. 2021, Nature)](https://www.nature.com/articles/s41586-021-03819-2) a protein structure prediction pipeline, with an API hosted at the Södinglab based on the MMseqs2 server [(Mirdita et al. 2019, Bioinformatics)](https://academic.oup.com/bioinformatics/article/35/16/2856/5280135) for the multiple sequence alignment creation. \n",
        "\n",
        "**WARNING**: this notebook does NOT use the AlphaFold2 pipeline for MSA/template generation. It may give better or worse results depending on number of sequences that can be found. Check out the [full AlphaFold2 pipeline](https://github.com/deepmind/alphafold) or Deepmind's official [google-colab notebook](https://colab.research.google.com/github/deepmind/alphafold/blob/main/notebooks/AlphaFold.ipynb).\n",
        "\n",
        "\n",
        "The main goal of this notebook is to demonstrate how to harness the power of cloud-computing to run microsecond-long MD simulations in a cheap and yet feasible fashion.\n",
        "\n",
        "---\n",
        "\n",
        " **This notebook is NOT a standard protocol for MD simulations!** It is just simple MD pipeline illustrating each step of a simulation protocol.\n",
        "\n",
        "--- \n",
        "**Bugs**\n",
        "- If you encounter any bugs, please report the issue to https://github.com/pablo-arantes/making-it-rain/issues\n",
        "\n",
        "**Acknowledgments**\n",
        "- We would like to thank the OpenMM team for developing an excellent and open source engine. \n",
        "\n",
        "- We would like to thank the AlphaFold team for developing an excellent model and open sourcing the software. \n",
        "\n",
        "- Credit to Sergey Ovchinnikov ([@sokrypton](https://twitter.com/sokrypton)), Milot Mirdita ([@milot_mirdita](https://twitter.com/milot_mirdita)) and Martin Steinegger ([@thesteinegger](https://twitter.com/thesteinegger)) for their fantastic [ColabFold](https://github.com/sokrypton/ColabFold)\n",
        "\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)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fNMH3_3DD5GH"
      },
      "source": [
        "# Introduction\n",
        "\n",
        "In general, MD simulations rely on 1) a set of atomic coordinates of all atoms on a simulation box and 2) a set of force field parameters that describes the interaction energies between atoms.\n",
        "\n",
        "In terms of inputs, we wil need:\n",
        "*  An amino acid sequence of your protein.\n",
        "\n",
        "In this notebook, we will simulate a hen egg-white lysozyme. To build our simulation box, we will use LEaP program (https://ambermd.org/tutorials/pengfei/index.php). The LEaP program is a portal between many chemical structure file types (.pdb and .mol2, primarily), and the Amber model parameter file types such as .lib, .prepi, parm.dat, and .frcmod. Each of the parameter files contains pieces of information needed for constructing a simulation, whether for energy minimization or molecular dynamics. LEaP functions within a larger workflow described in Section 1.1 of the [Amber Manual](https://ambermd.org/doc12/Amber20.pdf);.\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-MY1WXueB7Pd"
      },
      "source": [
        "## Using Google Drive to store simulation data\n",
        "\n",
        "Google Colab does not allow users to keep data on their computing nodes. However, we can use Google Drive to read, write, and store our simulations files. Therefore, we suggest to you to:\n",
        "\n",
        "1.   Create a folder in your own Google Drive and copy the necessary input files there.\n",
        "2.   Copy the path of your created directory. We will use it below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "jeobKtn1BunO"
      },
      "source": [
        "#@title ### Import Google Drive\n",
        "#@markdown Click in the \"Run\" buttom to make your Google Drive accessible.\n",
        "from google.colab import drive\n",
        "\n",
        "drive.flush_and_unmount()\n",
        "drive.mount('/content/drive', force_remount=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "xCU4NQ5boyMf"
      },
      "source": [
        "#@title Check if you correctly allocated GPU nodes\n",
        "\n",
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Select the Runtime > \"Change runtime type\" menu to enable a GPU accelerator, ')\n",
        "  print('and then re-execute this cell.')\n",
        "else:\n",
        "  print(gpu_info)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7WTscoLfElwX"
      },
      "source": [
        "---\n",
        "---\n",
        "## Setting the environment for our calculations\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.    AlphaFold v2.0 (https://github.com/deepmind/alphafold)\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iccGdbe_Pmt9",
        "cellView": "form"
      },
      "source": [
        "#@title Install dependencies\n",
        "#@markdown It will take a few minutes, please, drink a coffee and wait. ;-)\n",
        "%%bash -s $use_amber $use_msa $use_templates\n",
        "pip install --upgrade MDAnalysis 2>&1 1>/dev/null\n",
        "\n",
        "USE_AMBER=$1\n",
        "USE_MSA=$2\n",
        "USE_TEMPLATES=$3\n",
        "\n",
        "# install dependencies\n",
        "pip -q install biopython dm-haiku ml-collections py3Dmol\n",
        "wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/colabfold.py\n",
        "\n",
        "# download model\n",
        "git clone https://github.com/deepmind/alphafold.git --quiet\n",
        "(cd alphafold; git checkout 1d43aaff941c84dc56311076b58795797e49107b --quiet)\n",
        "mv alphafold alphafold_\n",
        "mv alphafold_/alphafold .\n",
        "# remove \"END\" from PDBs, otherwise biopython complains\n",
        "sed -i \"s/pdb_lines.append('END')//\" /content/alphafold/common/protein.py\n",
        "sed -i \"s/pdb_lines.append('ENDMDL')//\" /content/alphafold/common/protein.py\n",
        "\n",
        "# download model params (~1 min)\n",
        "mkdir params\n",
        "curl -fsSL https://storage.googleapis.com/alphafold/alphafold_params_2021-07-14.tar \\\n",
        "| tar x -C params\n",
        "touch AF2_READY\n",
        "\n",
        "# download libraries for interfacing with MMseqs2 API\n",
        "apt-get -qq -y update 2>&1 1>/dev/null\n",
        "apt-get -qq -y install jq curl zlib1g gawk 2>&1 1>/dev/null\n",
        "touch MMSEQ2_READY\n",
        "\n",
        "# setup conda\n",
        "wget -qnc https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh\n",
        "bash Miniconda3-latest-Linux-x86_64.sh -bfp /usr/local 2>&1 1>/dev/null\n",
        "rm Miniconda3-latest-Linux-x86_64.sh\n",
        "touch CONDA_READY\n",
        "# setup template search\n",
        "conda install -y -q -c conda-forge -c bioconda kalign3=3.2.2 hhsuite=3.3.0 python=3.7 2>&1 1>/dev/null\n",
        "touch HH_READY\n",
        "\n",
        "# setup openmm for amber refinement\n",
        "conda install -y -q -c conda-forge openmm=7.5.1 python=3.7 pdbfixer 2>&1 1>/dev/null\n",
        "(cd /usr/local/lib/python3.7/site-packages; patch -s -p0 < /content/alphafold_/docker/openmm.patch)\n",
        "wget -qnc https://git.scicore.unibas.ch/schwede/openstructure/-/raw/7102c63615b64735c4941278d92b554ec94415f8/modules/mol/alg/src/stereo_chemical_props.txt\n",
        "mv stereo_chemical_props.txt alphafold/common/\n",
        "touch AMBER_READY\n",
        "\n",
        "# install dependencies\n",
        "# conda install -c conda-forge mdanalysis --yes 2>&1 1>/dev/null\n",
        "pip install git+https://github.com/pablo-arantes/biopandas 2>&1 1>/dev/null\n",
        "# install conda\n",
        "conda install -c conda-forge ambertools --yes 2>&1 1>/dev/null\n",
        "conda install -c ambermd pytraj  --yes 2>&1 1>/dev/null\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IdMdkmCxCRFO"
      },
      "source": [
        "---\n",
        "---\n",
        "## AlphaFold2 Instructions\n",
        "\n",
        "AlphaFold2 pipeline on Colab by Sergey Ovchinnikov ([@sokrypton](https://twitter.com/sokrypton)), Milot Mirdita ([@milot_mirdita](https://twitter.com/milot_mirdita)) and Martin Steinegger ([@thesteinegger](https://twitter.com/thesteinegger)).\n",
        "\n",
        "**Quick start**\n",
        "1. Paste your protein sequence in the input field below.\n",
        "4. The pipeline consists of a several steps. The currently running steps is indicated by a circle with a stop sign next to it.\n",
        "\n",
        "**Result contents**\n",
        "\n",
        "1. PDB formatted structures sorted by avg. pIDDT. (relaxed, unrelaxed).\n",
        "2. Plots of the model quality.\n",
        "3. Plots of the MSA coverage.\n",
        "4. Parameter log file.\n",
        "5. A3M formatted input MSA.\n",
        "\n",
        "At the end of the job all the files will be uploaded to your Google Drive.\n",
        "\n",
        "**Using a custom MSA as input**\n",
        "\n",
        "To predict the structure with a custom MSA (A3M formatted): (1) Change the msa_mode: to \"custom\", (2) Wait for an upload box to appear at the end of the \"Input Protein ...\" box. Upload your A3M. The first fasta entry of the A3M must be the query sequence without gaps.\n",
        "\n",
        "To generate good input MSAs the HHblits server can be used here: https://toolkit.tuebingen.mpg.de/tools/hhblits\n",
        "\n",
        "After submitting your query, click \"Query Template MSA\" -> \"Download Full A3M\". Download the a3m file and upload it to the notebook.\n",
        "\n",
        "**Troubleshooting**\n",
        "* Try to restart the session \"Runtime\" -> \"Factory reset runtime\".\n",
        "* Check your input sequence.\n",
        "\n",
        "**Known issues**\n",
        "* Colab assigns different types of GPUs with varying amount of memory. Some might have not enough memory to predict the structure.\n",
        "\n",
        "**Limitations**\n",
        "* MSAs: MMseqs2 is very precise and sensitive but might find less hits compared to HHblits/HMMer searched against BFD or Mgnify.\n",
        "* Computing resources: Our MMseqs2 API can probably handle ~20k requests per day.\n",
        "* For best results, we recommend using the full pipeline: https://github.com/deepmind/alphafold\n",
        "\n",
        "**Description of the plots**\n",
        "*   **Number of sequences per position** - We want to see at least 30 sequences per position, for best performance, ideally 100 sequences.\n",
        "*   **Predicted lDDT per position** - model confidence (out of 100) at each position. Higher the better.\n",
        "*   **Predicted Alignment Error** - For homooligomers, this could be a useful metric to assess how confident the model is about the interface. Lower the better.\n",
        "\n",
        "**Please, don't use spaces in the jobname and folders names, i.e., AlphaFold2, content/drive/MyDrive and so on.**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kOblAo-xetgx",
        "cellView": "form"
      },
      "source": [
        "#@title Input protein sequence, then hit `Run` \n",
        "from google.colab import files\n",
        "import os\n",
        "import os.path\n",
        "import re\n",
        "import hashlib\n",
        "\n",
        "def add_hash(x,y):\n",
        "  return x+\"_\"+hashlib.sha1(y.encode()).hexdigest()[:5]\n",
        "\n",
        "query_sequence = 'KVFGRCELAAAMKRHGLDNYRGYSLGNWVCAAKFESNFNTQATNRNTDGSTDYGILQINSRWWCNDGRTPGSRNLCNIPCSALLSSDITASVNCAKKIVSDGNGMNAWVAWRNRCKGTDVQAWIRGCRL' #@param {type:\"string\"}\n",
        "# remove whitespaces\n",
        "query_sequence = \"\".join(query_sequence.split())\n",
        "query_sequence = re.sub(r'[^a-zA-Z]','', query_sequence).upper()\n",
        "\n",
        "jobname = '1AKI' #@param {type:\"string\"}\n",
        "# remove whitespaces\n",
        "jobname = \"\".join(jobname.split())\n",
        "jobname = re.sub(r'\\W+', '', jobname)\n",
        "jobname = add_hash(jobname, query_sequence)\n",
        "\n",
        "\n",
        "with open(f\"{jobname}.fasta\", \"w\") as text_file:\n",
        "    text_file.write(\">1\\n%s\" % query_sequence)\n",
        "\n",
        "# number of models to use\n",
        "#@markdown ---\n",
        "#@markdown ### Advanced settings\n",
        "msa_mode = \"MMseqs2 (UniRef+Environmental)\" #@param [\"MMseqs2 (UniRef+Environmental)\", \"MMseqs2 (UniRef only)\",\"single_sequence\",\"custom\"]\n",
        "num_models = 1 #@param [1,2,3,4,5] {type:\"raw\"}\n",
        "use_msa = True if msa_mode.startswith(\"MMseqs2\") else False\n",
        "use_env = True if msa_mode == \"MMseqs2 (UniRef+Environmental)\" else False\n",
        "use_custom_msa = True if msa_mode == \"custom\" else False\n",
        "use_amber = True #@param {type:\"boolean\"}\n",
        "use_templates = False #@param {type:\"boolean\"}\n",
        "#@markdown ---\n",
        "#@markdown ### Experimental options\n",
        "homooligomer = 1 #@param [1,2,3,4,5,6,7,8] {type:\"raw\"}\n",
        "\n",
        "#@markdown ### Save to Google Drive\n",
        "\n",
        "Google_Drive_Path = '/content/drive/MyDrive/' #@param {type:\"string\"}\n",
        "workDir = Google_Drive_Path \n",
        "\n",
        "#@title Import libraries\n",
        "# setup the model\n",
        "if \"model\" not in dir():\n",
        "\n",
        "  # hiding warning messages\n",
        "  import warnings\n",
        "  from absl import logging\n",
        "  import os\n",
        "  import tensorflow as tf\n",
        "  warnings.filterwarnings('ignore')\n",
        "  logging.set_verbosity(\"error\")\n",
        "  os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'\n",
        "  tf.get_logger().setLevel('ERROR')\n",
        "\n",
        "  import sys\n",
        "  import numpy as np\n",
        "  import pickle\n",
        "  from alphafold.common import protein\n",
        "  from alphafold.data import pipeline\n",
        "  from alphafold.data import templates\n",
        "  from alphafold.model import data\n",
        "  from alphafold.model import config\n",
        "  from alphafold.model import model\n",
        "  from alphafold.data.tools import hhsearch\n",
        "  import colabfold as cf\n",
        "\n",
        "  # plotting libraries\n",
        "  import py3Dmol\n",
        "  import matplotlib.pyplot as plt\n",
        "  import ipywidgets\n",
        "  from ipywidgets import interact, fixed, GridspecLayout, Output\n",
        "\n",
        "\n",
        "if use_amber and \"relax\" not in dir():\n",
        "  sys.path.insert(0, '/usr/local/lib/python3.7/site-packages/')\n",
        "  from alphafold.relax import relax\n",
        "\n",
        "def mk_mock_template(query_sequence):\n",
        "  # since alphafold's model requires a template input\n",
        "  # we create a blank example w/ zero input, confidence -1\n",
        "  ln = len(query_sequence)\n",
        "  output_templates_sequence = \"-\"*ln\n",
        "  output_confidence_scores = np.full(ln,-1)\n",
        "  templates_all_atom_positions = np.zeros((ln, templates.residue_constants.atom_type_num, 3))\n",
        "  templates_all_atom_masks = np.zeros((ln, templates.residue_constants.atom_type_num))\n",
        "  templates_aatype = templates.residue_constants.sequence_to_onehot(output_templates_sequence,\n",
        "                                                                    templates.residue_constants.HHBLITS_AA_TO_ID)\n",
        "  template_features = {'template_all_atom_positions': templates_all_atom_positions[None],\n",
        "                       'template_all_atom_masks': templates_all_atom_masks[None],\n",
        "                       'template_sequence': [f'none'.encode()],\n",
        "                       'template_aatype': np.array(templates_aatype)[None],\n",
        "                       'template_confidence_scores': output_confidence_scores[None],\n",
        "                       'template_domain_names': [f'none'.encode()],\n",
        "                       'template_release_date': [f'none'.encode()]}\n",
        "  return template_features\n",
        "\n",
        "def mk_template(a3m_lines, template_paths):\n",
        "  template_featurizer = templates.TemplateHitFeaturizer(\n",
        "      mmcif_dir=template_paths,\n",
        "      max_template_date=\"2100-01-01\",\n",
        "      max_hits=20,\n",
        "      kalign_binary_path=\"kalign\",\n",
        "      release_dates_path=None,\n",
        "      obsolete_pdbs_path=None)\n",
        "\n",
        "  hhsearch_pdb70_runner = hhsearch.HHSearch(binary_path=\"hhsearch\", databases=[f\"{template_paths}/pdb70\"])\n",
        "\n",
        "  hhsearch_result = hhsearch_pdb70_runner.query(a3m_lines)\n",
        "  hhsearch_hits = pipeline.parsers.parse_hhr(hhsearch_result)\n",
        "  templates_result = template_featurizer.get_templates(query_sequence=query_sequence,\n",
        "                                                       query_pdb_code=None,\n",
        "                                                       query_release_date=None,\n",
        "                                                       hits=hhsearch_hits)\n",
        "  return templates_result.features\n",
        "\n",
        "def set_bfactor(pdb_filename, bfac, idx_res, chains):\n",
        "  I = open(pdb_filename,\"r\").readlines()\n",
        "  O = open(pdb_filename,\"w\")\n",
        "  for line in I:\n",
        "    if line[0:6] == \"ATOM  \":\n",
        "      seq_id = int(line[22:26].strip()) - 1\n",
        "      seq_id = np.where(idx_res == seq_id)[0][0]\n",
        "      O.write(f\"{line[:21]}{chains[seq_id]}{line[22:60]}{bfac[seq_id]:6.2f}{line[66:]}\")\n",
        "  O.close()\n",
        "\n",
        "def predict_structure(prefix, feature_dict, Ls, model_params, use_model, do_relax=False, random_seed=0):  \n",
        "  \"\"\"Predicts structure using AlphaFold for the given sequence.\"\"\"\n",
        "\n",
        "  # Minkyung's code\n",
        "  # add big enough number to residue index to indicate chain breaks\n",
        "  idx_res = feature_dict['residue_index']\n",
        "  L_prev = 0\n",
        "  # Ls: number of residues in each chain\n",
        "  for L_i in Ls[:-1]:\n",
        "      idx_res[L_prev+L_i:] += 200\n",
        "      L_prev += L_i  \n",
        "  chains = list(\"\".join([ascii_uppercase[n]*L for n,L in enumerate(Ls)]))\n",
        "  feature_dict['residue_index'] = idx_res\n",
        "\n",
        "  # Run the models.\n",
        "  plddts,paes = [],[]\n",
        "  unrelaxed_pdb_lines = []\n",
        "  relaxed_pdb_lines = []\n",
        "\n",
        "  for model_name, params in model_params.items():\n",
        "    if model_name in use_model:\n",
        "      print(f\"running {model_name}\")\n",
        "      # swap params to avoid recompiling\n",
        "      # note: models 1,2 have diff number of params compared to models 3,4,5\n",
        "      if any(str(m) in model_name for m in [1,2]): model_runner = model_runner_1\n",
        "      if any(str(m) in model_name for m in [3,4,5]): model_runner = model_runner_3\n",
        "      model_runner.params = params\n",
        "      \n",
        "      processed_feature_dict = model_runner.process_features(feature_dict, random_seed=random_seed)\n",
        "      prediction_result = model_runner.predict(processed_feature_dict)\n",
        "      unrelaxed_protein = protein.from_prediction(processed_feature_dict,prediction_result)\n",
        "      unrelaxed_pdb_lines.append(protein.to_pdb(unrelaxed_protein))\n",
        "      plddts.append(prediction_result['plddt'])\n",
        "      paes.append(prediction_result['predicted_aligned_error'])\n",
        "\n",
        "      if do_relax:\n",
        "        # Relax the prediction.\n",
        "        amber_relaxer = relax.AmberRelaxation(max_iterations=0,tolerance=2.39,\n",
        "                                              stiffness=10.0,exclude_residues=[],\n",
        "                                              max_outer_iterations=20)      \n",
        "        relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)\n",
        "        relaxed_pdb_lines.append(relaxed_pdb_str)\n",
        "\n",
        "  # rerank models based on predicted lddt\n",
        "  lddt_rank = np.mean(plddts,-1).argsort()[::-1]\n",
        "  out = {}\n",
        "  print(\"reranking models based on avg. predicted lDDT\")\n",
        "  for n,r in enumerate(lddt_rank):\n",
        "    print(f\"model_{n+1} {np.mean(plddts[r])}\")\n",
        "\n",
        "    unrelaxed_pdb_path = f'{prefix}_unrelaxed_model_{n+1}.pdb'    \n",
        "    with open(unrelaxed_pdb_path, 'w') as f: f.write(unrelaxed_pdb_lines[r])\n",
        "    set_bfactor(unrelaxed_pdb_path, plddts[r], idx_res, chains)\n",
        "\n",
        "    if do_relax:\n",
        "      relaxed_pdb_path = f'{prefix}_relaxed_model_{n+1}.pdb'\n",
        "      with open(relaxed_pdb_path, 'w') as f: f.write(relaxed_pdb_lines[r])\n",
        "      set_bfactor(relaxed_pdb_path, plddts[r], idx_res, chains)\n",
        "\n",
        "    out[f\"model_{n+1}\"] = {\"plddt\":plddts[r], \"pae\":paes[r]}\n",
        "  return out\n",
        "\n",
        " #load MD dependencies\n",
        "import sys\n",
        "sys.path.append('/usr/local/lib/python3.7/site-packages/')\n",
        "from biopandas.pdb import PandasPdb\n",
        "import os\n",
        "import urllib.request  \n",
        "import MDAnalysis as mda\n",
        "from __future__ import print_function\n",
        "import pytraj as pt\n",
        "import platform\n",
        "import scipy.cluster.hierarchy\n",
        "from scipy.spatial.distance import squareform\n",
        "import scipy.stats as stats\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "from scipy.interpolate import griddata\n",
        "import seaborn as sb\n",
        "from statistics import mean, stdev\n",
        "from pytraj import matrix\n",
        "from matplotlib import colors\n",
        "from IPython.display import set_matplotlib_formats\n",
        "#%matplotlib inline\n",
        "#set_matplotlib_formats('png')\n",
        "#plt.figure(figsize=(5,7)) \n",
        "\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "if homooligomer > 1:\n",
        "  if use_amber:\n",
        "    print(\"amber disabled: amber is not currently supported for homooligomers\")\n",
        "    use_amber = False\n",
        "  if use_templates:\n",
        "    print(\"templates disabled: templates are not currently supported for homooligomers\")\n",
        "    use_templates = False\n",
        "\n",
        "with open(f\"{jobname}.log\", \"w\") as text_file:\n",
        "    text_file.write(\"num_models=%s\\n\" % num_models)\n",
        "    text_file.write(\"use_amber=%s\\n\" % use_amber)\n",
        "    text_file.write(\"use_msa=%s\\n\" % use_msa)\n",
        "    text_file.write(\"msa_mode=%s\\n\" % msa_mode)\n",
        "    text_file.write(\"use_templates=%s\\n\" % use_templates)\n",
        "    text_file.write(\"homooligomer=%s\\n\" % homooligomer)\n",
        "\n",
        "# decide which a3m to use\n",
        "if use_msa:\n",
        "  a3m_file = f\"{jobname}.a3m\"\n",
        "elif use_custom_msa:\n",
        "  a3m_file = f\"{jobname}.custom.a3m\"\n",
        "  if not os.path.isfile(a3m_file):\n",
        "    custom_msa_dict = files.upload()\n",
        "    custom_msa = list(custom_msa_dict.keys())[0]\n",
        "    header = 0\n",
        "    import fileinput\n",
        "    for line in fileinput.FileInput(custom_msa,inplace=1):\n",
        "      if line.startswith(\">\"):\n",
        "         header = header + 1 \n",
        "      if line.startswith(\"#\"):\n",
        "        continue\n",
        "      if line.rstrip() == False:\n",
        "        continue\n",
        "      if line.startswith(\">\") == False and header == 1:\n",
        "         query_sequence = line.rstrip() \n",
        "      print(line, end='')\n",
        "\n",
        "    os.rename(custom_msa, a3m_file)\n",
        "    print(f\"moving {custom_msa} to {a3m_file}\")\n",
        "else:\n",
        "  a3m_file = f\"{jobname}.single_sequence.a3m\"\n",
        "  with open(a3m_file, \"w\") as text_file:\n",
        "    text_file.write(\">1\\n%s\" % query_sequence)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A9tUpDaikPC8",
        "cellView": "form"
      },
      "source": [
        "#@title Call MMseqs2 to get MSA/templates\n",
        "if use_templates:\n",
        "  a3m_lines, template_paths = cf.run_mmseqs2(query_sequence, jobname, use_env, use_templates=True)\n",
        "  if template_paths is None:\n",
        "    template_features = mk_mock_template(query_sequence * homooligomer)\n",
        "  else:\n",
        "    template_features = mk_template(a3m_lines, template_paths)\n",
        "elif use_msa:\n",
        "  a3m_lines = cf.run_mmseqs2(query_sequence, jobname, use_env)\n",
        "  template_features = mk_mock_template(query_sequence * homooligomer)\n",
        "else:\n",
        "  template_features = mk_mock_template(query_sequence * homooligomer)\n",
        "\n",
        "if use_msa:\n",
        "  with open(a3m_file, \"w\") as text_file:\n",
        "    text_file.write(a3m_lines)\n",
        "else:\n",
        "  a3m_lines = \"\".join(open(a3m_file,\"r\").read())\n",
        "\n",
        "# parse MSA\n",
        "msa, deletion_matrix = pipeline.parsers.parse_a3m(a3m_lines)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hUYApPElB30u",
        "cellView": "form"
      },
      "source": [
        "#@title Gather input features, predict structure\n",
        "from string import ascii_uppercase\n",
        "\n",
        "# collect model weights\n",
        "use_model = {}\n",
        "if \"model_params\" not in dir(): model_params = {}\n",
        "for model_name in [\"model_1\",\"model_2\",\"model_3\",\"model_4\",\"model_5\"][:num_models]:\n",
        "  use_model[model_name] = True\n",
        "  if model_name not in model_params:\n",
        "    model_params[model_name] = data.get_model_haiku_params(model_name=model_name+\"_ptm\", data_dir=\".\")\n",
        "    if model_name == \"model_1\":\n",
        "      model_config = config.model_config(model_name+\"_ptm\")\n",
        "      model_config.data.eval.num_ensemble = 1\n",
        "      model_runner_1 = model.RunModel(model_config, model_params[model_name])\n",
        "    if model_name == \"model_3\":\n",
        "      model_config = config.model_config(model_name+\"_ptm\")\n",
        "      model_config.data.eval.num_ensemble = 1\n",
        "      model_runner_3 = model.RunModel(model_config, model_params[model_name])\n",
        "\n",
        "if homooligomer == 1:\n",
        "  msas = [msa]\n",
        "  deletion_matrices = [deletion_matrix]\n",
        "else:\n",
        "  # make multiple copies of msa for each copy\n",
        "  # AAA------\n",
        "  # ---AAA---\n",
        "  # ------AAA\n",
        "  #\n",
        "  # note: if you concat the sequences (as below), it does NOT work\n",
        "  # AAAAAAAAA\n",
        "  msas = []\n",
        "  deletion_matrices = []\n",
        "  Ln = len(query_sequence)\n",
        "  for o in range(homooligomer):\n",
        "    L = Ln * o\n",
        "    R = Ln * (homooligomer-(o+1))\n",
        "    msas.append([\"-\"*L+seq+\"-\"*R for seq in msa])\n",
        "    deletion_matrices.append([[0]*L+mtx+[0]*R for mtx in deletion_matrix])\n",
        "\n",
        "# gather features\n",
        "feature_dict = {\n",
        "    **pipeline.make_sequence_features(sequence=query_sequence*homooligomer,\n",
        "                                      description=\"none\",\n",
        "                                      num_res=len(query_sequence)*homooligomer),\n",
        "    **pipeline.make_msa_features(msas=msas,deletion_matrices=deletion_matrices),\n",
        "    **template_features\n",
        "}\n",
        "outs = predict_structure(jobname, feature_dict,\n",
        "                         Ls=[len(query_sequence)]*homooligomer,\n",
        "                         model_params=model_params, use_model=use_model,\n",
        "                         do_relax=use_amber)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6xbvRNrwnJqj",
        "cellView": "form"
      },
      "source": [
        "#@title Make plots\n",
        "\n",
        "# gather MSA info\n",
        "deduped_full_msa = list(dict.fromkeys(msa))\n",
        "msa_arr = np.array([list(seq) for seq in deduped_full_msa])\n",
        "seqid = (np.array(list(query_sequence)) == msa_arr).mean(-1)\n",
        "seqid_sort = seqid.argsort() #[::-1]\n",
        "non_gaps = (msa_arr != \"-\").astype(float)\n",
        "non_gaps[non_gaps == 0] = np.nan\n",
        "\n",
        "##################################################################\n",
        "plt.figure(figsize=(14,4),dpi=100)\n",
        "##################################################################\n",
        "plt.subplot(1,2,1); plt.title(\"Sequence coverage\")\n",
        "plt.imshow(non_gaps[seqid_sort]*seqid[seqid_sort,None],\n",
        "           interpolation='nearest', aspect='auto',\n",
        "           cmap=\"rainbow_r\", vmin=0, vmax=1, origin='lower')\n",
        "plt.plot((msa_arr != \"-\").sum(0), color='black')\n",
        "plt.xlim(-0.5,msa_arr.shape[1]-0.5)\n",
        "plt.ylim(-0.5,msa_arr.shape[0]-0.5)\n",
        "plt.colorbar(label=\"Sequence identity to query\",)\n",
        "plt.xlabel(\"Positions\")\n",
        "plt.ylabel(\"Sequences\")\n",
        "\n",
        "##################################################################\n",
        "plt.subplot(1,2,2); plt.title(\"Predicted lDDT per position\")\n",
        "for model_name,value in outs.items():\n",
        "  plt.plot(value[\"plddt\"],label=model_name)\n",
        "if homooligomer > 0:\n",
        "  for n in range(homooligomer+1):\n",
        "    x = n*(len(query_sequence)-1)\n",
        "    plt.plot([x,x],[0,100],color=\"black\")\n",
        "plt.legend()\n",
        "plt.ylim(0,100)\n",
        "plt.ylabel(\"Predicted lDDT\")\n",
        "plt.xlabel(\"Positions\")\n",
        "plt.savefig(jobname+\"_coverage_lDDT.png\")\n",
        "##################################################################\n",
        "plt.show()\n",
        "\n",
        "print(\"Predicted Alignment Error\")\n",
        "##################################################################\n",
        "plt.figure(figsize=(3*num_models,2), dpi=100)\n",
        "for n,(model_name,value) in enumerate(outs.items()):\n",
        "  plt.subplot(1,num_models,n+1)\n",
        "  plt.title(model_name)\n",
        "  plt.imshow(value[\"pae\"],label=model_name,cmap=\"bwr\",vmin=0,vmax=30)\n",
        "  plt.colorbar()\n",
        "plt.savefig(jobname+\"_PAE.png\")\n",
        "plt.show()\n",
        "##################################################################"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "KK7X9T44pWb7"
      },
      "source": [
        "#@title Display 3D structure {run: \"auto\"}\n",
        "model_num = 1 #@param [\"1\", \"2\", \"3\", \"4\", \"5\"] {type:\"raw\"}\n",
        "color = \"lDDT\" #@param [\"chain\", \"lDDT\", \"rainbow\"]\n",
        "show_sidechains = False #@param {type:\"boolean\"}\n",
        "show_mainchains = False #@param {type:\"boolean\"}\n",
        "\n",
        "def plot_plddt_legend():\n",
        "  thresh = ['plDDT:','Very low (<50)','Low (60)','OK (70)','Confident (80)','Very high (>90)']\n",
        "  plt.figure(figsize=(1,0.1),dpi=100)\n",
        "  ########################################\n",
        "  for c in [\"#FFFFFF\",\"#FF0000\",\"#FFFF00\",\"#00FF00\",\"#00FFFF\",\"#0000FF\"]:\n",
        "    plt.bar(0, 0, color=c)\n",
        "  plt.legend(thresh, frameon=False,\n",
        "             loc='center', ncol=6,\n",
        "             handletextpad=1,\n",
        "             columnspacing=1,\n",
        "             markerscale=0.5,)\n",
        "  plt.axis(False)\n",
        "  return plt\n",
        "\n",
        "def plot_confidence(model_num=1):\n",
        "  model_name = f\"model_{model_num}\"\n",
        "  plt.figure(figsize=(10,3),dpi=100)\n",
        "  \"\"\"Plots the legend for plDDT.\"\"\"\n",
        "  #########################################\n",
        "  plt.subplot(1,2,1); plt.title('Predicted lDDT')\n",
        "  plt.plot(outs[model_name][\"plddt\"])\n",
        "  for n in range(homooligomer+1):\n",
        "    x = n*(len(query_sequence))\n",
        "    plt.plot([x,x],[0,100],color=\"black\")\n",
        "  plt.ylabel('plDDT')\n",
        "  plt.xlabel('position')\n",
        "  #########################################\n",
        "  plt.subplot(1,2,2);plt.title('Predicted Aligned Error')\n",
        "  plt.imshow(outs[model_name][\"pae\"], cmap=\"bwr\",vmin=0,vmax=30)\n",
        "  plt.colorbar()\n",
        "  plt.xlabel('Scored residue')\n",
        "  plt.ylabel('Aligned residue')\n",
        "  #########################################\n",
        "  return plt\n",
        "\n",
        "def show_pdb(model_num=1, show_sidechains=False, show_mainchains=False, color=\"lDDT\"):\n",
        "  model_name = f\"model_{model_num}\"\n",
        "  if use_amber:\n",
        "    pdb_filename = f\"{jobname}_relaxed_{model_name}.pdb\"\n",
        "  else:\n",
        "    pdb_filename = f\"{jobname}_unrelaxed_{model_name}.pdb\"\n",
        "\n",
        "  view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',)\n",
        "  view.addModel(open(pdb_filename,'r').read(),'pdb')\n",
        "\n",
        "  if color == \"lDDT\":\n",
        "    view.setStyle({'cartoon': {'colorscheme': {'prop':'b','gradient': 'roygb','min':50,'max':90}}})\n",
        "  elif color == \"rainbow\":\n",
        "    view.setStyle({'cartoon': {'color':'spectrum'}})\n",
        "  elif color == \"chain\":\n",
        "    for n,chain,color in zip(range(homooligomer),list(\"ABCDEFGH\"),\n",
        "                     [\"lime\",\"cyan\",\"magenta\",\"yellow\",\"salmon\",\"white\",\"blue\",\"orange\"]):\n",
        "       view.setStyle({'chain':chain},{'cartoon': {'color':color}})\n",
        "  if show_sidechains:\n",
        "    BB = ['C','O','N']\n",
        "    view.addStyle({'and':[{'resn':[\"GLY\",\"PRO\"],'invert':True},{'atom':BB,'invert':True}]},\n",
        "                        {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "    view.addStyle({'and':[{'resn':\"GLY\"},{'atom':'CA'}]},\n",
        "                        {'sphere':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "    view.addStyle({'and':[{'resn':\"PRO\"},{'atom':['C','O'],'invert':True}]},\n",
        "                        {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})  \n",
        "  if show_mainchains:\n",
        "    BB = ['C','O','N','CA']\n",
        "    view.addStyle({'atom':BB},{'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "\n",
        "  view.zoomTo()\n",
        "  return view\n",
        "\n",
        "show_pdb(model_num,show_sidechains, show_mainchains, color).show()\n",
        "if color == \"lDDT\": plot_plddt_legend().show()  \n",
        "plot_confidence(model_num).show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "64i6tEImyoyM",
        "cellView": "form"
      },
      "source": [
        "#@title Ramachandran plot:\n",
        "!npx degit https://github.com/pablo-arantes/Making-it-rain/ temp 2> /dev/null\n",
        "cp_command = \"cp -r temp/rama-500 .\"\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "with open('cp.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(cp_command)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "!chmod 700 cp.sh 2>&1 1>/dev/null\n",
        "!bash cp.sh 2>&1 1>/dev/null\n",
        "!rm -r temp cp.sh\n",
        "\n",
        "model_num = 1 #@param [\"1\", \"2\", \"3\", \"4\", \"5\"] {type:\"raw\"}\n",
        "model_name = f\"model_{model_num}\"\n",
        "if use_amber == True:\n",
        "  pdb_file = f\"{jobname}_relaxed_{model_name}.pdb\"\n",
        "else:\n",
        "  pdb_file = f\"{jobname}_unrelaxed_{model_name}.pdb\"\n",
        "\n",
        "import math\n",
        "import sys\n",
        "import os\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "from Bio import PDB\n",
        "from matplotlib import colors\n",
        "\n",
        "def plot_ramachandran(file):\n",
        "    __file__=file\n",
        "\n",
        "    \"\"\"\n",
        "    The preferences were calculated from the following artice:\n",
        "    Lovell et al. Structure validation by Calpha geometry: phi,psi and Cbeta deviation. 2003\n",
        "    DOI: 10.1002/prot.10286\n",
        "    \"\"\"\n",
        "\n",
        "    # General variable for the background preferences\n",
        "    rama_preferences = {\n",
        "        \"General\": {\n",
        "            \"file\": \"rama500-general.data\",\n",
        "            \"cmap\": colors.ListedColormap(['#FFFFFF', '#B3E8FF', '#7FD9FF']),\n",
        "            \"bounds\": [0, 0.0005, 0.02, 1],\n",
        "        },\n",
        "        \"GLY\": {\n",
        "            \"file\": \"rama500-gly-sym.data\",\n",
        "            \"cmap\": colors.ListedColormap(['#FFFFFF', '#FFE8C5', '#FFCC7F']),\n",
        "            \"bounds\": [0, 0.002, 0.02, 1],\n",
        "        },\n",
        "        \"PRO\": {\n",
        "            \"file\": \"rama500-pro.data\",\n",
        "            \"cmap\": colors.ListedColormap(['#FFFFFF', '#D0FFC5', '#7FFF8C']),\n",
        "            \"bounds\": [0, 0.002, 0.02, 1],\n",
        "        },\n",
        "        \"PRE-PRO\": {\n",
        "            \"file\": \"rama500-prepro.data\",\n",
        "            \"cmap\": colors.ListedColormap(['#FFFFFF', '#B3E8FF', '#7FD9FF']),\n",
        "            \"bounds\": [0, 0.002, 0.02, 1],\n",
        "        }\n",
        "    }\n",
        "\n",
        "    # Read in the expected torsion angles\n",
        "    __location__ = '/content/rama-500/' #You must set the ptah of the .data files here\n",
        "    rama_pref_values = {}\n",
        "    for key, val in rama_preferences.items():\n",
        "        rama_pref_values[key] = np.full((360, 360), 0, dtype=np.float64)\n",
        "        with open(os.path.join(__location__, val[\"file\"])) as fn:\n",
        "            for line in fn:\n",
        "                if not line.startswith(\"#\"):\n",
        "                    # Preference file has values for every second position only\n",
        "                    rama_pref_values[key][int(float(line.split()[1])) + 180][int(float(line.split()[0])) + 180] = float(\n",
        "                        line.split()[2])\n",
        "                    rama_pref_values[key][int(float(line.split()[1])) + 179][int(float(line.split()[0])) + 179] = float(\n",
        "                        line.split()[2])\n",
        "                    rama_pref_values[key][int(float(line.split()[1])) + 179][int(float(line.split()[0])) + 180] = float(\n",
        "                        line.split()[2])\n",
        "                    rama_pref_values[key][int(float(line.split()[1])) + 180][int(float(line.split()[0])) + 179] = float(\n",
        "                        line.split()[2])\n",
        "\n",
        "    normals = {}\n",
        "    outliers = {}\n",
        "    for key, val in rama_preferences.items():\n",
        "        normals[key] = {\"x\": [], \"y\": [],'Res':[]}\n",
        "        outliers[key] = {\"x\": [], \"y\": []}\n",
        "\n",
        "    # Calculate the torsion angle of the inputs\n",
        "    for inp in sys.argv[1:]:\n",
        "        if not os.path.isfile(inp):\n",
        "            # print(\"{} not found!\".format(inp))\n",
        "            continue\n",
        "    structure = PDB.PDBParser().get_structure('input_structure', __file__)\n",
        "    for model in structure:\n",
        "        for chain in model:\n",
        "            polypeptides = PDB.PPBuilder().build_peptides(chain)\n",
        "            for poly_index, poly in enumerate(polypeptides):\n",
        "                phi_psi = poly.get_phi_psi_list()\n",
        "                for res_index, residue in enumerate(poly):\n",
        "                    res_name = \"{}\".format(residue.resname)\n",
        "                    res_num = residue.id[1]\n",
        "                    phi, psi = phi_psi[res_index]\n",
        "                    if phi and psi:\n",
        "                        aa_type = \"\"\n",
        "                        if str(poly[res_index + 1].resname) == \"PRO\":\n",
        "                            aa_type = \"PRE-PRO\"\n",
        "                        elif res_name == \"PRO\":\n",
        "                            aa_type = \"PRO\"\n",
        "                        elif res_name == \"GLY\":\n",
        "                            aa_type = \"GLY\"\n",
        "                        else:\n",
        "                            aa_type = \"General\"\n",
        "                        if rama_pref_values[aa_type][int(math.degrees(psi)) + 180][int(math.degrees(phi)) + 180] < \\\n",
        "                                rama_preferences[aa_type][\"bounds\"][1]:\n",
        "                            # print(\"{} {} {} {}{} is an outlier\".format(inp, model, chain, res_name, res_num))\n",
        "                            outliers[aa_type][\"x\"].append(math.degrees(phi))\n",
        "                            outliers[aa_type][\"y\"].append(math.degrees(psi))\n",
        "                        else:\n",
        "                            normals[aa_type][\"x\"].append(math.degrees(phi))\n",
        "                            normals[aa_type][\"y\"].append(math.degrees(psi))\n",
        "                            normals[aa_type]['Res'].append(res_name+'_'+str(res_num))\n",
        "\n",
        "    # Generate the plots\n",
        "    plt.figure(figsize=(10,10))\n",
        "    for idx, (key, val) in enumerate(sorted(rama_preferences.items(), key=lambda x: x[0].lower())):\n",
        "        plt.subplot(2, 2, idx + 1)\n",
        "        plt.title(key,fontsize=20)\n",
        "        plt.imshow(rama_pref_values[key], cmap=rama_preferences[key][\"cmap\"],\n",
        "                   norm=colors.BoundaryNorm(rama_preferences[key][\"bounds\"], rama_preferences[key][\"cmap\"].N),\n",
        "                   extent=(-180, 180, 180, -180),alpha=0.7)\n",
        "\n",
        "        plt.scatter(normals[key][\"x\"], normals[key][\"y\"],s=[15],marker='.')\n",
        "\n",
        "        #for key in normals:\n",
        "            #for i, name in enumerate (normals[key]['Res']):\n",
        "                #plt.annotate(name, (normals[key][\"x\"][i], normals[key][\"y\"][i]))\n",
        "\n",
        "        plt.scatter(outliers[key][\"x\"], outliers[key][\"y\"],color=\"red\",s=[15],marker='.')\n",
        "        plt.xlim([-180, 180])\n",
        "        plt.ylim([-180, 180])\n",
        "        plt.plot([-180, 180], [0, 0],color=\"k\",alpha=0.7)\n",
        "        plt.plot([0, 0], [-180, 180],color=\"k\",alpha=0.7)\n",
        "        plt.xlabel(r'$\\phi$',fontsize=12)\n",
        "        plt.ylabel(r'$\\psi$',fontsize=12)\n",
        "        plt.grid(linestyle='dotted')\n",
        "\n",
        "    plt.tight_layout()\n",
        "    plt.savefig(jobname+\"_ramachandran.png\", dpi=600) #Uncommet this line of you want so save the plot in a specific location\n",
        "    plt.show()\n",
        "plot_ramachandran(pdb_file)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "33g5IIegij5R",
        "cellView": "form"
      },
      "source": [
        "#@title Package and upload the AlphaFold2 results on Drive\n",
        "#@markdown Upload all the AlphaFold results in zip format on Google Drive\n",
        "\n",
        "!zip -FSr $jobname\".alphafold2_result.zip\" $jobname\".log\" $a3m_file $jobname\"_\"*\"relaxed_model_\"*\".pdb\" $jobname\"_coverage_lDDT.png\" $jobname\"_PAE.png\" $jobname\"_ramachandran.png\" 2>&1 1>/dev/null\n",
        "\n",
        "cp_sys = \"cp \" + jobname + \".alphafold2_result.zip \" + workDir \n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "\n",
        "with open('cp_sys.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(cp_sys)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "\n",
        "!chmod 700 cp_sys.sh 2>&1 1>/dev/null\n",
        "!bash cp_sys.sh 2> /dev/null\n",
        "\n",
        "zip_end = os.path.join(workDir, jobname + \".alphafold2_result.zip\")\n",
        "\n",
        "zip_true = os.path.exists(zip_end)\n",
        "\n",
        "if zip_true == True:\n",
        "  print(\" Zip file loaded successfully on Google Drive! :-)\")\n",
        "else:\n",
        "  print(\"ERROR: Check your input file! \")\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2gOXPG2dbno4",
        "cellView": "form"
      },
      "source": [
        "#@title Parameters to generate the Amber topology:\n",
        "model_num = 1 #@param [\"1\", \"2\", \"3\", \"4\", \"5\"] {type:\"raw\"}\n",
        "model_name = f\"model_{model_num}\"\n",
        "if use_amber:\n",
        "  pdb_filename = f\"{jobname}_relaxed_{model_name}.pdb\"\n",
        "else:\n",
        "  pdb_filename = f\"{jobname}_unrelaxed_{model_name}.pdb\"\n",
        "\n",
        "starting = os.path.join(workDir, \"starting.pdb\")\n",
        "starting_end = os.path.join(workDir, \"starting2.pdb\")\n",
        "tleap = os.path.join(workDir, \"tleap.in\")\n",
        "top_nw = os.path.join(workDir, \"SYS_nw.prmtop\")\n",
        "crd_nw = os.path.join(workDir, \"SYS_nw.crd\")\n",
        "pdb_nw = os.path.join(workDir, \"SYS_nw.pdb\")\n",
        "top = os.path.join(workDir, \"SYS.prmtop\")\n",
        "crd = os.path.join(workDir, \"SYS.crd\")\n",
        "pdb = os.path.join(workDir, \"SYS.pdb\")\n",
        "\n",
        "pdbfn = pdb_filename\n",
        "ppdb = PandasPdb().read_pdb(pdbfn)\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",
        "pdb4amber_cmd = \"pdb4amber -i \" + str(starting) + \" -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",
        "!chmod 700 pdb4amber.sh 2>&1 1>/dev/null\n",
        "!bash pdb4amber.sh 2> /dev/null\n",
        "\n",
        "Force_field = \"ff19SB\" #@param [\"ff19SB\", \"ff14SB\"]\n",
        "if Force_field == \"ff19SB\":\n",
        "  ff = \"leaprc.protein.ff19SB\"\n",
        "else:\n",
        "  ff = \"leaprc.protein.ff14SB\"\n",
        "\n",
        "Water_type = \"OPC\" #@param [\"TIP3P\", \"OPC\"]\n",
        "if Water_type == \"TIP3P\":\n",
        "  water = \"leaprc.water.tip3p\"\n",
        "  water_box = \"TIP3PBOX\"\n",
        "else:\n",
        "  water = \"leaprc.water.opc\"\n",
        "  water_box = \"OPCBOX\"\n",
        "\n",
        "#@markdown Size Box (Angstrons):\n",
        "\n",
        "Size_box = 12 #@param {type:\"slider\", min:10, max:20, step:1}\n",
        "size_box = Size_box\n",
        "\n",
        "#@markdown **ATTENTION**: Give the concentration in Molar units, AMBER tleap will neutralize your system automatically:\n",
        "\n",
        "Ions = \"NaCl\" #@param [\"NaCl\", \"KCl\" ]\n",
        "\n",
        "Concentration = \"0.15\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "f = open(tleap, \"w\")\n",
        "f.write(\"\"\"source \"\"\" + str(ff) + \"\\n\"\n",
        "\"\"\"source leaprc.DNA.OL15\n",
        "source leaprc.RNA.OL3\n",
        "source leaprc.GLYCAM_06j-1 \n",
        "source leaprc.gaff2\n",
        "source \"\"\"  + str(water) + \"\\n\"\n",
        "\"\"\"SYS = loadpdb \"\"\"  + str(starting_end) + \"\\n\"\n",
        "\"\"\"alignaxes SYS\n",
        "savepdb SYS \"\"\" + str(pdb_nw) + \"\\n\"\n",
        "\"\"\"saveamberparm SYS \"\"\" + str(top_nw) + \" \" + str(crd_nw) + \"\\n\"\n",
        "\"\"\"solvatebox SYS \"\"\" + str(water_box) + \" \" + str(size_box) +  \"\"\" 0.7\n",
        "saveamberparm SYS \"\"\" + str(top) + \" \" + str(crd) + \"\\n\"\n",
        "\"\"\"savepdb SYS \"\"\" + str(pdb) + \"\\n\"\n",
        "\"\"\"quit\"\"\")\n",
        "f.close()\n",
        "\n",
        "tleap_command = \"tleap -f \" + str(tleap)\n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "\n",
        "with open('run_tleap.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(tleap_command)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "\n",
        "SYS = os.path.join(workDir, \"SYS*\")\n",
        "rm_sys = \"rm \" + SYS\n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "\n",
        "with open('rm_sys.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(rm_sys)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "\n",
        "!chmod 700 rm_sys.sh 2>&1 1>/dev/null\n",
        "!bash rm_sys.sh 2> /dev/null\n",
        "\n",
        "!chmod 700 run_tleap.sh 2>&1 1>/dev/null\n",
        "!bash run_tleap.sh 2>&1 1>/dev/null\n",
        "\n",
        "!grep \"Volume:\" leap.log > temp.txt\n",
        "with open(\"temp.txt\", 'r') as f:\n",
        "  for line in f:\n",
        "        vol = float(line.split()[1])\n",
        "\n",
        "vol_lit  = vol * pow(10, -27)\n",
        "atom_lit = 9.03 * pow(10, 22)\n",
        "conc = float(Concentration)\n",
        "num_ion = int(vol_lit * (conc/0.15) * atom_lit)\n",
        "\n",
        "if Ions == \"NaCl\":\n",
        "  pos_neut = \"Na+ 0\"\n",
        "  pos_num = \"Na+ \" + str(num_ion)\n",
        "  Cl_num = num_ion\n",
        "else:\n",
        "  pos_neut = \"K+ 0\"\n",
        "  pos_num = \"K+ \" + str(num_ion)\n",
        "  Cl_num = num_ion\n",
        "\n",
        "f = open(tleap, \"w\")\n",
        "f.write(\"\"\"source \"\"\" + str(ff) + \"\\n\"\n",
        "\"\"\"source leaprc.DNA.OL15\n",
        "source leaprc.RNA.OL3\n",
        "source leaprc.GLYCAM_06j-1 \n",
        "source leaprc.gaff2\n",
        "source \"\"\"  + str(water) + \"\\n\"\n",
        "\"\"\"SYS = loadpdb \"\"\"  + str(starting_end) + \"\\n\"\n",
        "\"\"\"alignaxes SYS\n",
        "check SYS \n",
        "charge SYS\n",
        "addions SYS \"\"\" + str(pos_neut) + \"\\n\"\n",
        "\"\"\"addions SYS Cl- 0\n",
        "check SYS\n",
        "charge SYS\n",
        "savepdb SYS \"\"\" + str(pdb_nw) + \"\\n\"\n",
        "\"\"\"saveamberparm SYS \"\"\" + str(top_nw) + \" \" + str(crd_nw) + \"\\n\"\n",
        "\"\"\"solvatebox SYS \"\"\" + str(water_box) + \" \" + str(size_box) +  \"\"\" 0.7 \"\"\" + \"\\n\"\n",
        "\"\"\"addIonsRand SYS \"\"\" + str(pos_num) + \"\"\" Cl- \"\"\" + str(Cl_num) + \"\\n\"\n",
        "\"\"\"saveamberparm SYS \"\"\" + str(top) + \" \" + str(crd) + \"\\n\"\n",
        "\"\"\"savepdb SYS \"\"\" + str(pdb) + \"\\n\"\n",
        "\"\"\"quit\"\"\")\n",
        "f.close()\n",
        "\n",
        "!chmod 700 run_tleap.sh 2>&1 1>/dev/null\n",
        "!bash run_tleap.sh 2>&1 1>/dev/null\n",
        "\n",
        "pdb_amber = os.path.exists(pdb)\n",
        "top_amber = os.path.exists(top)\n",
        "crd_amber = os.path.exists(crd)\n",
        "\n",
        "if pdb_amber == True and top_amber == True and crd_amber == True:\n",
        "  print(\"Successfully generated topology! :-)\")\n",
        "else:\n",
        "  print(\"ERROR: Check your input file! \")\n",
        "!rm *.sh temp.txt >/dev/null 2>&1  "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "75YWjs-SgA39"
      },
      "source": [
        "## Let's take a look on our simulation box:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Zp0zOaZggFJE",
        "cellView": "form"
      },
      "source": [
        "#@title Show 3D structure\n",
        "import ipywidgets\n",
        "from ipywidgets import interact, fixed\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "\n",
        "def show_pdb(show_box=True,\n",
        "             show_sidechains=False,\n",
        "             show_mainchain=False,\n",
        "             color=\"None\"):\n",
        "\n",
        "  def mainchain(p, color=\"white\", model=0):\n",
        "    BB = ['C','O','N','CA']\n",
        "    p.addStyle({\"model\":model,'atom':BB},\n",
        "                       {'stick':{'colorscheme':f\"{color}Carbon\",'radius':0.4}})\n",
        "\n",
        "  def box(p, model=0):\n",
        "      p.addModelsAsFrames(pdb)\n",
        "      p.addSurface(py3Dmol.SAS, {'opacity': 0.6, 'color':'white'}) #comment this line if you dont want to see the water box\n",
        " \n",
        "\n",
        "  def sidechain(p, model=0):\n",
        "    HP = [\"ALA\",\"GLY\",\"VAL\",\"ILE\",\"LEU\",\"PHE\",\"MET\",\"PRO\",\"TRP\",\"CYS\",\"TYR\"]\n",
        "    BB = ['C','O','N']\n",
        "    p.addStyle({\"model\":model,'and':[{'resn':HP},{'atom':BB,'invert':True}]},\n",
        "              {'stick':{'colorscheme':\"whiteCarbon\",'radius':0.4}})\n",
        "    p.addStyle({\"model\":model,'and':[{'resn':\"GLY\"},{'atom':'CA'}]},\n",
        "              {'sphere':{'colorscheme':\"whiteCarbon\",'radius':0.4}})\n",
        "    p.addStyle({\"model\":model,'and':[{'resn':\"PRO\"},{'atom':['C','O'],'invert':True}]},\n",
        "              {'stick':{'colorscheme':\"whiteCarbon\",'radius':0.4}})  \n",
        "    p.addStyle({\"model\":model,'and':[{'resn':HP,'invert':True},{'atom':BB,'invert':True}]},\n",
        "              {'stick':{'colorscheme':\"whiteCarbon\",'radius':0.4}})\n",
        "\n",
        "  p = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js')\n",
        "  p.addModel(open(pdb,'r').read(),'pdb')\n",
        "  if color == \"rainbow\":\n",
        "    p.setStyle({'cartoon': {'color':'spectrum'}})\n",
        "  else:\n",
        "    p.setStyle({'cartoon':{}})\n",
        "\n",
        "  if show_sidechains: sidechain(p)\n",
        "  if show_mainchain: mainchain(p)\n",
        "  if show_box: box(p)\n",
        "  p.zoomTo()\n",
        "  return p.show()\n",
        "\n",
        "interact(show_pdb,\n",
        "         show_box=ipywidgets.Checkbox(value=True),\n",
        "         show_sidechains=ipywidgets.Checkbox(value=False),\n",
        "         show_mainchain=ipywidgets.Checkbox(value=False),\n",
        "         color=ipywidgets.Dropdown(options=['None', 'rainbow'], value='rainbow'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q2Z_A7c9gO0x"
      },
      "source": [
        "---\n",
        "---\n",
        "# Equilibrating the simulation box\n",
        "\n",
        "Proper MD equilibration protocol is designed to equilibrate both temperature and pressure throughout the simulation box while preserving the protein experimental conformation. In addition, we also allow the solvent to accomodate around the protein, creating proper solvation layers.\n",
        "\n",
        "Below, we will set up the MD equilibration parameters, such as temperature, pressure and the desired simulation time. We will define the force constant used to restraint protein heavy-atoms in place and the frequency at which we want to save atomic coordinates in a trajectory file (.dcd).\n",
        "\n",
        "After you are done, you can run the next 2 cells to equilibrate your system."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rqrxJ-BXgTVT",
        "cellView": "form"
      },
      "source": [
        "#@title ### Parameters for MD Equilibration protocol:\n",
        "\n",
        "# remove whitespaces\n",
        "Jobname = '1aki_equil' #@param {type:\"string\"}\n",
        "\n",
        "Minimization_steps = \"1000\" #@param [\"1000\", \"5000\", \"10000\", \"20000\", \"50000\", \"100000\"]\n",
        "\n",
        "#@markdown Simulation time (in nanoseconds) and integration time (in femtoseconds): \n",
        "Time = \"0.2\" #@param {type:\"string\"}\n",
        "stride_time_eq = Time\n",
        "Integration_timestep = \"2\" #@param [\"0.5\", \"1\", \"2\", \"3\", \"4\"]\n",
        "dt_eq = Integration_timestep\n",
        "\n",
        "#@markdown Temperature (in Kelvin) and Pressure (in bar)\n",
        "Temperature = 298 #@param {type:\"string\"}\n",
        "temperature_eq = Temperature\n",
        "Pressure = 1 #@param {type:\"string\"}\n",
        "pressure_eq = Pressure\n",
        "\n",
        "#@markdown Position restraints force constant (in kJ/mol): \n",
        "Force_constant = 500 #@param {type:\"slider\", min:0, max:2000, step:100}\n",
        "\n",
        "#@markdown Frequency to write the trajectory file (in picoseconds): \n",
        "\n",
        "Write_the_trajectory = \"10\" #@param [\"10\", \"100\", \"200\", \"500\", \"1000\"]\n",
        "write_the_trajectory_eq = Write_the_trajectory\n",
        "#@markdown Frequency to write the log file (in picoseconds): \n",
        "\n",
        "Write_the_log = \"10\" #@param [\"10\", \"100\", \"200\", \"500\", \"1000\"]\n",
        "write_the_log_eq = Write_the_log\n",
        "\n",
        "\n",
        "#@markdown ---\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zoamR9iynphz",
        "cellView": "form"
      },
      "source": [
        "#@title Runs an Equilibration MD simulation (NPT ensemble)\n",
        "#@markdown Now, let's equilibrate our system!\n",
        "\n",
        "###########################################\n",
        "import simtk.openmm as mm\n",
        "from simtk.openmm import *\n",
        "from simtk.openmm.app import *\n",
        "from simtk.unit import *\n",
        "import pytraj as pt\n",
        "\n",
        "from sys import stdout, exit, stderr\n",
        "import os, math, fnmatch\n",
        "\n",
        "#############################################\n",
        "# Defining MD simulation parameters\n",
        "\n",
        "jobname = os.path.join(workDir, Jobname)\n",
        "coordinatefile = os.path.join(workDir, \"SYS.crd\")\n",
        "pdbfile = os.path.join(workDir, \"SYS.pdb\")\n",
        "topologyfile = os.path.join(workDir, \"SYS.prmtop\")\n",
        "\n",
        "time_ps = float(Time)*1000\n",
        "simulation_time = float(time_ps)*picosecond\t\t# in ps\n",
        "dt = int(dt_eq)*femtosecond\t\t\t\t\t\n",
        "temperature = float(temperature_eq)*kelvin\n",
        "savcrd_freq = int(write_the_trajectory_eq)*picosecond\n",
        "print_freq  = int(write_the_log_eq)*picosecond\n",
        "\n",
        "pressure\t= float(pressure_eq)*bar\n",
        "\n",
        "restraint_fc = int(Force_constant) # kJ/mol\n",
        "\n",
        "nsteps  = int(simulation_time.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nprint  = int(print_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nsavcrd = int(savcrd_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "\n",
        "#############################################\n",
        "# Defining functions to use below:\n",
        "def backup_old_log(pattern, string):\n",
        "\tresult = []\n",
        "\tfor root, dirs, files in os.walk(\"./\"):\n",
        "\t\tfor name in files:\n",
        "\t\t\tif fnmatch.fnmatch(name, pattern):\n",
        "\n",
        "\t\t\t\ttry:\n",
        "\t\t\t\t\tnumber = int(name[-2])\n",
        "\t\t\t\t\tavail = isinstance(number, int)\n",
        "\t\t\t\t\t#print(name,avail)\n",
        "\t\t\t\t\tif avail == True:\n",
        "\t\t\t\t\t\tresult.append(number)\n",
        "\t\t\t\texcept:\n",
        "\t\t\t\t\tpass\n",
        "\n",
        "\tif len(result) > 0:\n",
        "\t\tmaxnumber = max(result)\n",
        "\telse:\n",
        "\t\tmaxnumber = 0\n",
        "\n",
        "\tbackup_file = \"\\#\" + string + \".\" + str(maxnumber + 1) + \"#\"\n",
        "\tos.system(\"mv \" + string + \" \" + backup_file)\n",
        "\treturn backup_file\n",
        "\n",
        "def restraints(system, crd, fc, restraint_array):\n",
        "\n",
        "\tboxlx = system.getDefaultPeriodicBoxVectors()[0][0].value_in_unit(nanometers)\n",
        "\tboxly = system.getDefaultPeriodicBoxVectors()[1][1].value_in_unit(nanometers)\n",
        "\tboxlz = system.getDefaultPeriodicBoxVectors()[2][2].value_in_unit(nanometers)\n",
        "\n",
        "\tif fc > 0:\n",
        "\t\t# positional restraints for all heavy-atoms\n",
        "\t\tposresPROT = CustomExternalForce('k*periodicdistance(x, y, z, x0, y0, z0)^2;')\n",
        "\t\tposresPROT.addPerParticleParameter('k')\n",
        "\t\tposresPROT.addPerParticleParameter('x0')\n",
        "\t\tposresPROT.addPerParticleParameter('y0')\n",
        "\t\tposresPROT.addPerParticleParameter('z0')\n",
        "  \n",
        "\t\tfor atom1 in restraint_array:\n",
        "\t\t\tatom1 = int(atom1)\n",
        "               \n",
        "\t\t\txpos  = crd.positions[atom1].value_in_unit(nanometers)[0]\n",
        "\t\t\typos  = crd.positions[atom1].value_in_unit(nanometers)[1]\n",
        "\t\t\tzpos  = crd.positions[atom1].value_in_unit(nanometers)[2]\n",
        "\n",
        "\t\t\tposresPROT.addParticle(atom1, [fc, xpos, ypos, zpos])\n",
        "    \n",
        "\t\tsystem.addForce(posresPROT)\n",
        "\n",
        "\treturn system\n",
        "##############################################\n",
        "\n",
        "#############################################\n",
        "print(\"\\n> Simulation details:\\n\")\n",
        "print(\"\\tJob name = \" + jobname)\n",
        "print(\"\\tCoordinate file = \" + str(coordinatefile))\n",
        "print(\"\\tPDB file = \" + str(pdbfile))\n",
        "print(\"\\tTopology file = \" + str(topologyfile))\n",
        "\n",
        "print(\"\\n\\tSimulation_time = \" + str(simulation_time))\n",
        "print(\"\\tIntegration timestep = \" + str(dt))\n",
        "print(\"\\tTotal number of steps = \" +  str(nsteps))\n",
        "\n",
        "print(\"\\n\\tSave coordinates each \" + str(savcrd_freq))\n",
        "print(\"\\tPrint in log file each \" + str(print_freq))\n",
        "\n",
        "print(\"\\n\\tTemperature = \" + str(temperature))\n",
        "print(\"\\tPressure = \" + str(pressure))\n",
        "#############################################\n",
        "\n",
        "print(\"\\n> Setting the system:\\n\")\n",
        "\n",
        "print(\"\\t- Reading topology and structure file...\")\n",
        "prmtop = AmberPrmtopFile(topologyfile)\n",
        "inpcrd = AmberInpcrdFile(coordinatefile)\n",
        "\n",
        "print(\"\\t- Creating system and setting parameters...\")\n",
        "nonbondedMethod = PME\n",
        "nonbondedCutoff = 1.0*nanometers\n",
        "ewaldErrorTolerance = 0.0005\n",
        "constraints = HBonds\n",
        "rigidWater = True\n",
        "constraintTolerance = 0.000001\n",
        "friction = 1.0\n",
        "system = prmtop.createSystem(nonbondedMethod=nonbondedMethod, nonbondedCutoff=nonbondedCutoff,\n",
        "                           constraints=constraints, rigidWater=rigidWater, ewaldErrorTolerance=ewaldErrorTolerance)\n",
        "\n",
        "print(\"\\t- Applying restraints. Force Constant = \" + str(Force_constant) + \"kJ/mol\")\n",
        "pt_system = pt.iterload(coordinatefile, topologyfile)\n",
        "pt_topology = pt_system.top\n",
        "restraint_array = pt.select_atoms('!(:H*) & !(:WAT) & !(:Na+) & !(:Cl-) & !(:Mg+) & !(:K+)', pt_topology)\n",
        "\n",
        "system = restraints(system, inpcrd, restraint_fc, restraint_array)\n",
        "\n",
        "print(\"\\t- Setting barostat...\")\n",
        "system.addForce(MonteCarloBarostat(pressure, temperature))\n",
        "\n",
        "print(\"\\t- Setting integrator...\")\n",
        "integrator = LangevinIntegrator(temperature, friction, dt)\n",
        "integrator.setConstraintTolerance(constraintTolerance)\n",
        "simulation = Simulation(prmtop.topology, system, integrator)\n",
        "simulation.context.setPositions(inpcrd.positions)\n",
        "if inpcrd.boxVectors is not None:\n",
        "    simulation.context.setPeriodicBoxVectors(*inpcrd.boxVectors)\n",
        "\n",
        "print(\"\\t- Energy minimization: \" + str(Minimization_steps) + \" steps\")\n",
        "simulation.minimizeEnergy(tolerance=10*kilojoule/mole, maxIterations=int(Minimization_steps))\n",
        "print(\"\\t-> Potential Energy = \" + str(simulation.context.getState(getEnergy=True).getPotentialEnergy()))\n",
        "\n",
        "print(\"\\t- Setting initial velocities...\")\n",
        "simulation.context.setVelocitiesToTemperature(temperature)\n",
        "\n",
        "#############################################\n",
        "# Running Equilibration on NPT ensemble\n",
        "\n",
        "dcd_file = jobname + \".dcd\"\n",
        "log_file = jobname + \".log\"\n",
        "rst_file = jobname + \".rst\"\n",
        "prv_rst_file = jobname + \".rst\"\n",
        "pdb_file = jobname + \".pdb\"\n",
        "\n",
        "# Creating a trajectory file and reporters\n",
        "dcd = DCDReporter(dcd_file, nsavcrd)\n",
        "firstdcdstep = (nsteps) + nsavcrd\n",
        "dcd._dcd = DCDFile(dcd._out, simulation.topology, simulation.integrator.getStepSize(), firstdcdstep, nsavcrd) # charmm doesn't like first step to be 0\n",
        "\n",
        "simulation.reporters.append(dcd)\n",
        "simulation.reporters.append(StateDataReporter(stdout, nprint, step=True, speed=True, progress=True, totalSteps=nsteps, remainingTime=True, separator='\\t\\t'))\n",
        "simulation.reporters.append(StateDataReporter(log_file, nprint, step=True, kineticEnergy=True, potentialEnergy=True, totalEnergy=True, temperature=True, volume=True, speed=True))\n",
        "\n",
        "print(\"\\n> Simulating \" + str(nsteps) + \" steps...\")\n",
        "simulation.step(nsteps)\n",
        "\n",
        "simulation.reporters.clear() # remove all reporters so the next iteration don't trigger them.\n",
        "\n",
        "\n",
        "##################################\n",
        "# Writing last frame information of stride\n",
        "print(\"\\n> Writing state file (\" + str(rst_file) + \")...\")\n",
        "state = simulation.context.getState( getPositions=True, getVelocities=True )\n",
        "with open(rst_file, 'w') as f:\n",
        "\tf.write(XmlSerializer.serialize(state))\n",
        "\n",
        "last_frame = int(nsteps/nsavcrd)\n",
        "print(\"> Writing coordinate file (\" + str(pdb_file) + \", frame = \" + str(last_frame) + \")...\")\n",
        "positions = simulation.context.getState(getPositions=True).getPositions()\n",
        "PDBFile.writeFile(simulation.topology, positions, open(pdb_file, 'w'))\n",
        "\n",
        "print(\"\\n> Finished!\\n\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4hcKPNoskvaX"
      },
      "source": [
        "---\n",
        "---\n",
        "# Running a Production MD simulation\n",
        "\n",
        "Finally, we will proceed with the Production simulation itself using the equilibrated system coordinates as input structure.\n",
        "\n",
        "Note that we will use here a *.rst state file* , which contains atomic velocities and positions from the last frame of the equilibration simulation, guaranteeing that our production simulation begins from a thermodynamically equilibrated system.\n",
        "\n",
        "Another important information here is the **Number_of_strides** and the **Stride_Time**. In this notebook, we simulate a defined number of *strides*, so the **simulation time = Number_of_strides*Stride_Time**. For example, we can simulate 100ns by setting *Number_of_strides=10* and *Stride_Time=10 ns*.\n",
        "\n",
        "**Important: at the end of the Production simulation, we concatenate all strides to create a complete trajectory file which can be visualized and analyzed**\n",
        "\n",
        "The idea behind this approach is to make use of the intermitent 12h/24h period that Google Colab allows us to use its GPUs. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "jPEWjcoGky6R"
      },
      "source": [
        "#@markdown ### Provide input file names below:\n",
        "\n",
        "Equilibrated_PDB = '1aki_equil.pdb' #@param {type:\"string\"}\n",
        "State_file = '1aki_equil.rst' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown ### Parameters for MD Production protocol:\n",
        "\n",
        "\n",
        "# remove whitespaces\n",
        "Jobname = '1aki_prod' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Simulation time (in nanoseconds), number of strides (integers) and integration timestep (in femtoseconds): \n",
        "Stride_Time = \"0.2\" #@param {type:\"string\"}\n",
        "stride_time_prod = Stride_Time\n",
        "Number_of_strides = \"1\" #@param {type:\"string\"}\n",
        "nstride = Number_of_strides\n",
        "Integration_timestep = \"2\" #@param [\"0.5\", \"1\", \"2\", \"3\", \"4\"]\n",
        "dt_prod = Integration_timestep\n",
        "\n",
        "#@markdown Temperature (in Kelvin) and Pressure (in bar)\n",
        "Temperature = 298 #@param {type:\"string\"}\n",
        "temperature_prod = Temperature\n",
        "Pressure = 1 #@param {type:\"string\"}\n",
        "pressure_prod = Pressure\n",
        "\n",
        "#@markdown Frequency to write the trajectory file (in picoseconds): \n",
        "Write_the_trajectory = \"10\" #@param [\"10\", \"100\", \"200\", \"500\", \"1000\"]\n",
        "write_the_trajectory_prod = Write_the_trajectory\n",
        "\n",
        "#@markdown Frequency to write the log file (in picoseconds): \n",
        "Write_the_log = \"10\" #@param [\"10\", \"100\", \"200\", \"500\", \"1000\"]\n",
        "write_the_log_prod = Write_the_log\n",
        "\n",
        "#@markdown ---"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0QcjKSoqHHhi",
        "cellView": "form"
      },
      "source": [
        "#@title Runs a Production MD simulation (NPT ensemble) after equilibration\n",
        "#\n",
        "###########################################\n",
        "import simtk.openmm as mm\n",
        "from simtk.openmm import *\n",
        "from simtk.openmm.app import *\n",
        "from simtk.unit import *\n",
        "\n",
        "from sys import stdout, exit, stderr\n",
        "import os, math, fnmatch\n",
        "\n",
        "#############################################\n",
        "# Defining MD simulation parameters\n",
        "\n",
        "jobname = os.path.join(workDir, str(Jobname))\n",
        "coordinatefile = os.path.join(workDir, \"SYS.crd\")\n",
        "pdbfile = os.path.join(workDir, Equilibrated_PDB)\n",
        "topologyfile = os.path.join(workDir, \"SYS.prmtop\")\n",
        "equil_rst_file = os.path.join(workDir, State_file)\n",
        "\n",
        "\n",
        "stride_time_ps = float(stride_time_prod)*1000\n",
        "stride_time = float(stride_time_ps)*picosecond        \n",
        "nstride = int(Number_of_strides)\n",
        "dt = int(dt_prod)*femtosecond\t\t\t\t\t\n",
        "temperature = float(temperature_prod)*kelvin\n",
        "savcrd_freq = int(write_the_trajectory_prod)*picosecond\n",
        "print_freq  = int(write_the_log_prod)*picosecond\n",
        "\n",
        "pressure\t= float(pressure_prod)*bar\n",
        "\n",
        "simulation_time = stride_time*nstride\n",
        "nsteps  = int(stride_time.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nprint  = int(print_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "nsavcrd = int(savcrd_freq.value_in_unit(picosecond)/dt.value_in_unit(picosecond))\n",
        "firststride = 1 # must be integer\n",
        "#############################################\n",
        "# Defining functions to use below:\n",
        "def backup_old_log(pattern, string):\n",
        "\tresult = []\n",
        "\tfor root, dirs, files in os.walk(\"./\"):\n",
        "\t\tfor name in files:\n",
        "\t\t\tif fnmatch.fnmatch(name, pattern):\n",
        "\n",
        "\t\t\t\ttry:\n",
        "\t\t\t\t\tnumber = int(name[-2])\n",
        "\t\t\t\t\tavail = isinstance(number, int)\n",
        "\t\t\t\t\t#print(name,avail)\n",
        "\t\t\t\t\tif avail == True:\n",
        "\t\t\t\t\t\tresult.append(number)\n",
        "\t\t\t\texcept:\n",
        "\t\t\t\t\tpass\n",
        "\n",
        "\tif len(result) > 0:\n",
        "\t\tmaxnumber = max(result)\n",
        "\telse:\n",
        "\t\tmaxnumber = 0\n",
        "\n",
        "\tbackup_file = \"\\#\" + string + \".\" + str(maxnumber + 1) + \"#\"\n",
        "\tos.system(\"mv \" + string + \" \" + backup_file)\n",
        "\treturn backup_file\n",
        "##############################################\n",
        "\n",
        "#############################################\n",
        "print(\"\\n> Simulation details:\\n\")\n",
        "print(\"\\tJob name = \" + jobname)\n",
        "print(\"\\tCoordinate file = \" + str(coordinatefile))\n",
        "print(\"\\tPDB file = \" + str(pdbfile))\n",
        "print(\"\\tTopology file = \" + str(topologyfile))\n",
        "\n",
        "print(\"\\n\\tSimulation_time = \" + str(stride_time*nstride))\n",
        "print(\"\\tIntegration timestep = \" + str(dt))\n",
        "print(\"\\tTotal number of steps = \" +  str(nsteps*nstride))\n",
        "print(\"\\tNumber of strides = \" + str(nstride) + \" (\" + str(stride_time) + \" in each stride)\")\n",
        "\n",
        "print(\"\\n\\tSave coordinates each \" + str(savcrd_freq))\n",
        "print(\"\\tPrint in log file each \" + str(print_freq))\n",
        "\n",
        "print(\"\\n\\tTemperature = \" + str(temperature))\n",
        "print(\"\\tPressure = \" + str(pressure))\n",
        "#############################################\n",
        "\n",
        "print(\"\\n> Setting the system:\\n\")\n",
        "\n",
        "print(\"\\t- Reading topology and structure file...\")\n",
        "prmtop = AmberPrmtopFile(topologyfile)\n",
        "inpcrd = AmberInpcrdFile(coordinatefile)\n",
        "\n",
        "\n",
        "print(\"\\t- Creating system and setting parameters...\")\n",
        "nonbondedMethod = PME\n",
        "nonbondedCutoff = 1.0*nanometers\n",
        "ewaldErrorTolerance = 0.0005\n",
        "constraints = HBonds\n",
        "rigidWater = True\n",
        "constraintTolerance = 0.000001\n",
        "friction = 1.0\n",
        "system = prmtop.createSystem(nonbondedMethod=nonbondedMethod, nonbondedCutoff=nonbondedCutoff,\n",
        "                           constraints=constraints, rigidWater=rigidWater, ewaldErrorTolerance=ewaldErrorTolerance)\n",
        "\n",
        "print(\"\\t- Setting barostat...\")\n",
        "system.addForce(MonteCarloBarostat(pressure, temperature))\n",
        "\n",
        "print(\"\\t- Setting integrator...\")\n",
        "integrator = LangevinIntegrator(temperature, friction, dt)\n",
        "integrator.setConstraintTolerance(constraintTolerance)\n",
        "simulation = Simulation(prmtop.topology, system, integrator)\n",
        "simulation.context.setPositions(inpcrd.positions)\n",
        "if inpcrd.boxVectors is not None:\n",
        "\tsimulation.context.setPeriodicBoxVectors(*inpcrd.boxVectors)\n",
        "\n",
        "#############################################\n",
        "# Opening a loop of extension NSTRIDE to simulate the entire STRIDE_TIME*NSTRIDE\n",
        "for n in range(1, nstride + 1):\n",
        "\n",
        "\tprint(\"\\n\\n>>> Simulating Stride #\" + str(n) + \" <<<\")\n",
        "\n",
        "\tdcd_file = jobname + \"_\" + str(n) + \".dcd\"\n",
        "\tlog_file = jobname + \"_\" + str(n) + \".log\"\n",
        "\trst_file = jobname + \"_\" + str(n) + \".rst\"\n",
        "\tprv_rst_file = jobname + \"_\" + str(n-1) + \".rst\"\n",
        "\tpdb_file = jobname + \"_\" + str(n) + \".pdb\"\n",
        "\n",
        "\tif os.path.exists(rst_file):\n",
        "\t\tprint(\"> Stride #\" + str(n) + \" finished (\" + rst_file + \" present). Moving to next stride... <\")\n",
        "\t\tcontinue\n",
        "\n",
        "\tif n == 1:\n",
        "\t\tprint(\"\\n> Loading previous state from equilibration > \" + equil_rst_file + \" <\")\n",
        "\t\twith open(equil_rst_file, 'r') as f:\n",
        "\t\t\tsimulation.context.setState(XmlSerializer.deserialize(f.read()))\n",
        "\t\t\tcurrstep = int((n-1)*nsteps)\n",
        "\t\t\tcurrtime = currstep*dt.in_units_of(picosecond)\n",
        "\t\t\tsimulation.currentStep = currstep\n",
        "\t\t\tsimulation.context.setTime(currtime)\n",
        "\t\t\tprint(\"> Current time: \" + str(currtime) + \" (Step = \" + str(currstep) + \")\")\n",
        "\n",
        "\telse:\n",
        "\t\tprint(\"> Loading previous state from > \" + prv_rst_file + \" <\")\n",
        "\t\twith open(prv_rst_file, 'r') as f:\n",
        "\t\t\tsimulation.context.setState(XmlSerializer.deserialize(f.read()))\n",
        "\t\t\tcurrstep = int((n-1)*nsteps)\n",
        "\t\t\tcurrtime = currstep*dt.in_units_of(picosecond)\n",
        "\t\t\tsimulation.currentStep = currstep\n",
        "\t\t\tsimulation.context.setTime(currtime)\n",
        "\t\t\tprint(\"> Current time: \" + str(currtime) + \" (Step = \" + str(currstep) + \")\")\n",
        "\n",
        "\n",
        "\tdcd = DCDReporter(dcd_file, nsavcrd)\n",
        "\tfirstdcdstep = (currstep) + nsavcrd\n",
        "\tdcd._dcd = DCDFile(dcd._out, simulation.topology, simulation.integrator.getStepSize(), firstdcdstep, nsavcrd) # first step should not be 0\n",
        "\n",
        "\tsimulation.reporters.append(dcd)\n",
        "\tsimulation.reporters.append(StateDataReporter(stdout, nprint, step=True, speed=True, progress=True, totalSteps=(nsteps*nstride), remainingTime=True, separator='\\t\\t'))\n",
        "\tsimulation.reporters.append(StateDataReporter(log_file, nprint, step=True, kineticEnergy=True, potentialEnergy=True, totalEnergy=True, temperature=True, volume=True, speed=True))\n",
        "\n",
        "\tprint(\"\\n> Simulating \" + str(nsteps) + \" steps... (Stride #\" + str(n) + \")\")\n",
        "\tsimulation.step(nsteps)\n",
        "\n",
        "\tsimulation.reporters.clear() # remove all reporters so the next iteration don't trigger them.\n",
        "\n",
        "\n",
        "\t##################################\n",
        "\t# Writing last frame information of stride\n",
        "\tprint(\"\\n> Writing state file (\" + str(rst_file) + \")...\")\n",
        "\tstate = simulation.context.getState( getPositions=True, getVelocities=True )\n",
        "\twith open(rst_file, 'w') as f:\n",
        "\t\tf.write(XmlSerializer.serialize(state))\n",
        "\n",
        "\tlast_frame = int(nsteps/nsavcrd)\n",
        "\tprint(\"> Writing coordinate file (\" + str(pdb_file) + \", frame = \" + str(last_frame) + \")...\")\n",
        "\tpositions = simulation.context.getState(getPositions=True).getPositions()\n",
        "\tPDBFile.writeFile(simulation.topology, positions, open(pdb_file, 'w'))\n",
        "\n",
        "print(\"\\n> Finished!\\n\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ULa1WR_3k9Jd",
        "cellView": "form"
      },
      "source": [
        "#@title **Concatenate and align the trajectory**\n",
        "Skip = \"1\" #@param [\"1\", \"2\", \"5\", \"10\", \"20\", \"50\"]\n",
        "stride_traj = Skip\n",
        "Output_format = \"dcd\" #@param [\"dcd\", \"pdb\", \"trr\", \"xtc\"]\n",
        "\n",
        "#@markdown **Attention:** A high number of frames can explode the memory on Colab. You should be fine with 5000 frames or less.\n",
        "\n",
        "simulation_time_analysis = stride_time_ps*nstride\n",
        "simulation_ns = float(Stride_Time)*int(Number_of_strides)\n",
        "number_frames = int(simulation_time_analysis)/int(Write_the_trajectory)\n",
        "number_frames_analysis = number_frames/int(stride_traj)\n",
        "\n",
        "\n",
        "traj_end = os.path.join(workDir, str(Jobname) + \"_all.dcd\")\n",
        "traj_end2 = os.path.join(workDir, str(Jobname) + \"_all.\" + str(Output_format))\n",
        "template =  os.path.join(workDir, str(Jobname) + '_%s.dcd')\n",
        "\n",
        "flist = [template % str(i) for i in range(1, nstride + 1)]\n",
        "#print(flist)\n",
        "\n",
        "trajlist = pt.load(flist, pdb, stride=stride_traj)\n",
        "traj_image = trajlist.iterframe(autoimage=True, rmsfit=0)\n",
        "traj_write = pt.write_traj(traj_end, traj_image, overwrite=True)\n",
        "traj_load = pt.load(traj_end, pdb)\n",
        "traj_align = pt.align(traj_load, mask=\"@CA\", ref=0) \n",
        "traj_write = pt.write_traj(traj_end, traj_align, overwrite=True, options='dcd')\n",
        "traj_write = pt.write_traj(traj_end2, traj_align, overwrite=True, options=Output_format)\n",
        "traj_load = pt.load(traj_end, pdb)\n",
        "print(traj_load)\n",
        "\n",
        "traj_end_check = os.path.exists(traj_end2)\n",
        "\n",
        "if traj_end_check == True:\n",
        "  print(\"Trajectory concatenated successfully! :-)\")\n",
        "else:\n",
        "  print(\"ERROR: Check your inputs! \")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rcTxS3T5lBNl",
        "cellView": "form"
      },
      "source": [
        "#@title Load, view and check the trajectory\n",
        "#@markdown This will take a few minutes. Another coffee would be great. :-)\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "!rm *.pdb 2> /dev/null\n",
        "\n",
        "#py3dmol functions\n",
        "class Atom(dict):\n",
        "  def __init__(self, line):\n",
        "    self[\"type\"] = line[0:6].strip()\n",
        "    self[\"idx\"] = line[6:11].strip()\n",
        "    self[\"name\"] = line[12:16].strip()\n",
        "    self[\"resname\"] = line[17:20].strip()\n",
        "    self[\"resid\"] = int(int(line[22:26]))\n",
        "    self[\"x\"] = float(line[30:38])\n",
        "    self[\"y\"] = float(line[38:46])\n",
        "    self[\"z\"] = float(line[46:54])\n",
        "    self[\"sym\"] = line[76:78].strip()\n",
        "\n",
        "  def __str__(self):\n",
        "    line = list(\" \" * 80)\n",
        "    line[0:6] = self[\"type\"].ljust(6)\n",
        "    line[6:11] = self[\"idx\"].ljust(5)\n",
        "    line[12:16] = self[\"name\"].ljust(4)\n",
        "    line[17:20] = self[\"resname\"].ljust(3)\n",
        "    line[22:26] = str(self[\"resid\"]).ljust(4)\n",
        "    line[30:38] = str(self[\"x\"]).rjust(8)\n",
        "    line[38:46] = str(self[\"y\"]).rjust(8)\n",
        "    line[46:54] = str(self[\"z\"]).rjust(8)\n",
        "    line[76:78] = self[\"sym\"].rjust(2)\n",
        "    return \"\".join(line) + \"\\n\"\n",
        "        \n",
        "class Molecule(list):\n",
        "  def __init__(self, file):\n",
        "    for line in file:\n",
        "      if \"ATOM\" in line or \"HETATM\" in line:\n",
        "        self.append(Atom(line))\n",
        "            \n",
        "    def __str__(self):\n",
        "      outstr = \"\"\n",
        "      for at in self:\n",
        "        outstr += str(at)\n",
        "      return outstr\n",
        "\n",
        "if number_frames_analysis > 10:\n",
        "  stride_animation = number_frames_analysis/10\n",
        "else:\n",
        "  stride_animation = 1\n",
        "\n",
        "u = mda.Universe(top, traj_end) \n",
        "\n",
        "# Write out frames for animation\n",
        "protein = u.select_atoms('not (resname WAT)')\n",
        "i = 0\n",
        "for ts in u.trajectory[0:len(u.trajectory):int(stride_animation)]: \n",
        "    if i > -1:\n",
        "        with mda.Writer('' + str(i) + '.pdb', protein.n_atoms) as W:\n",
        "            W.write(protein)\n",
        "    i = i + 1\n",
        "# Load frames as molecules\n",
        "molecules = []\n",
        "for i in range(int(len(u.trajectory)/int(stride_animation))):\n",
        "    with open('' + str(i) + '.pdb') as ifile:\n",
        "        molecules.append(Molecule(ifile))\n",
        "\n",
        "models = \"\"\n",
        "for i in range(len(molecules)):\n",
        "  models += \"MODEL \" + str(i) + \"\\n\"\n",
        "  for j,mol in enumerate(molecules[i]):\n",
        "    models += str(mol)\n",
        "  models += \"ENDMDL\\n\"\n",
        "#view.addModelsAsFrames(models)\n",
        "\n",
        "# Animation\n",
        "view = py3Dmol.view(width=800, height=600)\n",
        "view.addModelsAsFrames(models)\n",
        "for i, at in enumerate(molecules[0]):\n",
        "    default = {\"cartoon\": {'color': 'spectrum'}}\n",
        "    view.setStyle({'model': -1, 'serial': i+1}, at.get(\"pymol\", default))\n",
        "\n",
        "view.zoomTo()\n",
        "view.animate({'loop': \"forward\"})\n",
        "view.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Emh0vU5UjgB6"
      },
      "source": [
        "---\n",
        "---\n",
        "# Analysis\n",
        "\n",
        "Although visualizing your trajectory can be quite useful, sometimes you also want more quantitative data.\n",
        "\n",
        "Analyses of MD trajectories vary a lot and we do not intend to cover it all here. However, one can make use of MDanalysis or PyTraj to easily analyze simulations. \n",
        "\n",
        "Below, you can find a few examples of code snippets that can help you to shed some light on your simulation behavior."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wBrBMF4Puyv6",
        "cellView": "form"
      },
      "source": [
        "#@title Compute RMSD of protein's CA atoms\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'rmsd_ca' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "rmsd = pt.rmsd(traj_load, ref = 0, mask = \"@CA\")\n",
        "\n",
        "time = len(rmsd)*int(Write_the_trajectory)/1000\n",
        "time_array = np.arange(0,time,int(Write_the_trajectory)/1000)*int(stride_traj)\n",
        "\n",
        "# Plotting:\n",
        "ax = plt.plot(time_array, rmsd, alpha=0.6, color = 'blue', linewidth = 1.0)\n",
        "plt.xlim(0, simulation_ns)\n",
        "#plt.ylim(2, 6)\n",
        "\n",
        "plt.xlabel(\"Time (ns)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel(\"RMSD [$\\AA$]\", fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "raw_data=pd.DataFrame(rmsd)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZHyMpikjuaLT",
        "cellView": "form"
      },
      "source": [
        "#@title Plot RMSD as a ditribution\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'rmsd_dist' #@param {type:\"string\"}\n",
        "\n",
        "ax = sb.kdeplot(rmsd, color=\"blue\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('RMSD [$\\AA$]', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks([])\n",
        "plt.ylabel('')\n",
        "ax.spines['top'].set_visible(False)\n",
        "ax.spines['right'].set_visible(False)\n",
        "ax.spines['bottom'].set_visible(True)\n",
        "ax.spines['left'].set_visible(False)\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CvOFrXGXwXrV",
        "cellView": "form"
      },
      "source": [
        "#@title Compute radius of gyration of protein's CA atoms\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'radius_gyration' #@param {type:\"string\"}\n",
        "\n",
        "radgyr = pt.radgyr(traj_load, mask = \"@CA\")\n",
        "time = len(rmsd)*int(Write_the_trajectory)/1000\n",
        "time_array = np.arange(0,time,int(Write_the_trajectory)/1000)*int(stride_traj)\n",
        "\n",
        "# Plotting:\n",
        "plt.plot(time_array, radgyr, alpha=0.6, color = 'green', linewidth = 1.0)\n",
        "plt.xlim(0, simulation_ns)\n",
        "#plt.ylim(2, 6)\n",
        "\n",
        "plt.xlabel(\"Time (ns)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel(\"Radius of gyration ($\\AA$)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "raw_data=pd.DataFrame(radgyr)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2Q7FKg8Fuxr9",
        "cellView": "form"
      },
      "source": [
        "#@title Plot radius of gyration as a ditribution\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'radius_gyration_dist' #@param {type:\"string\"}\n",
        "\n",
        "ax = sb.kdeplot(radgyr, color=\"green\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('Radius of gyration ($\\AA$)', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks([])\n",
        "plt.ylabel('')\n",
        "ax.spines['top'].set_visible(False)\n",
        "ax.spines['right'].set_visible(False)\n",
        "ax.spines['bottom'].set_visible(True)\n",
        "ax.spines['left'].set_visible(False)\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p2Y0DgwTxLWc",
        "cellView": "form"
      },
      "source": [
        "#@title Compute RMSF of protein's CA atoms\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'rmsf_ca' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "rmsf = pt.rmsf(traj_load, \"@CA\")\n",
        "bfactor = pt.bfactors(traj_load, byres=True)\n",
        "\n",
        "# Plotting:\n",
        "plt.plot(rmsf[:,1], alpha=1.0, color = 'red', linewidth = 1.0)\n",
        "\n",
        "plt.xlabel(\"Residue\", fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel(\"RMSF ($\\AA$)\", fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.xlim(0, len(rmsf[:-1]))\n",
        "\n",
        "#plt.xticks(np.arange(min(rmsf[:1]), max(rmsf[:1])))\n",
        "plt.yticks(fontsize = 12)\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "raw_data=pd.DataFrame(rmsf)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JalicqqrTodW",
        "cellView": "form"
      },
      "source": [
        "#@title 2D RMSD\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = '2D_rmsd' #@param {type:\"string\"}\n",
        "\n",
        "last_frame = len(time_array)\n",
        "\n",
        "stride_ticks_f = (last_frame)/5\n",
        "ticks_frame = np.arange(0,(len(time_array) + float(stride_ticks_f)), float(stride_ticks_f))\n",
        "a = ticks_frame.astype(float)\n",
        "stride_ticks_t = (simulation_ns)/5\n",
        "tick_time = np.arange(0,(float(simulation_ns) + float(stride_ticks_t)), float(stride_ticks_t))\n",
        "b = tick_time.astype(float)\n",
        "\n",
        "mat1 = pt.pairwise_rmsd(traj_load, mask=\"@CA\", frame_indices=range(int(number_frames_analysis)))\n",
        "\n",
        "\n",
        "ax = plt.imshow(mat1, cmap = 'PRGn', origin='lower', interpolation = 'bicubic')\n",
        "plt.title('2D RMSD')\n",
        "plt.xlabel('Time (ns)', fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel('Time (ns)', fontsize = 14, fontweight = 'bold')\n",
        "# plt.xticks(fontsize = 12)\n",
        "# plt.yticks(fontsize = 12)\n",
        "plt.xticks(a, b.round(decimals=3), fontsize = 12)\n",
        "plt.yticks(a, b.round(decimals=3), fontsize = 12)\n",
        "# plt.xlim(0, a[-1])\n",
        "# plt.ylim(0, a[-1])\n",
        "\n",
        "cbar1 = plt.colorbar()\n",
        "cbar1.set_label(\"RMSD ($\\AA$)\", fontsize = 14, fontweight = 'bold')\n",
        "\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "raw_data=pd.DataFrame(mat1)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_mgVSbBshWFV",
        "cellView": "form"
      },
      "source": [
        "#@title Calculate eigvenctors of Principle Component Analysis (PCA)\n",
        "data = pt.pca(traj_load, fit=True, ref=0, mask='@CA', n_vecs=2)\n",
        "#print('projection values of each frame to first mode = {} \\n'.format(data[0][0]))\n",
        "#print('projection values of each frame to second mode = {} \\n'.format(data[0][1]))\n",
        "#print('eigvenvalues of first two modes', data[1][0])\n",
        "#print(\"\")\n",
        "#print('eigvenvectors of first two modes: \\n', data[1][1])\n",
        "\n",
        "last_frame = len(time_array)\n",
        "\n",
        "stride_ticks_f = (last_frame)/5\n",
        "ticks_frame = np.arange(0,(len(time_array) + float(stride_ticks_f)), float(stride_ticks_f))\n",
        "a = ticks_frame.astype(float)\n",
        "a2 = a.tolist()\n",
        "stride_ticks_t = (simulation_ns)/5\n",
        "tick_time = np.arange(0,(float(simulation_ns) + float(stride_ticks_t)), float(stride_ticks_t))\n",
        "b = tick_time.astype(float)\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'PCA' #@param {type:\"string\"}\n",
        "\n",
        "Output_PC1 = 'PC1' #@param {type:\"string\"}\n",
        "Output_PC2 = 'PC2' #@param {type:\"string\"}\n",
        "\n",
        "%matplotlib inline\n",
        "%config InlineBackend.figure_format = 'retina'  # high resolution\n",
        "projection_data = data[0]\n",
        "plt.title(r'PCA of C-$\\alpha$')\n",
        "PC1 = data[0][0] \n",
        "PC2 = data[0][1]\n",
        "\n",
        "a = plt.scatter(PC1,PC2, c=range(int(number_frames_analysis)), cmap='Greens', marker='o',s=8, alpha=1)\n",
        "plt.clim(0, last_frame)\n",
        "\n",
        "plt.xlabel('PC1', fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel('PC2', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "# N = len(number_frames)\n",
        "# x2 = np.arange(N)\n",
        "\n",
        "cbar1 = plt.colorbar(a, orientation=\"vertical\")\n",
        "cbar1.set_label('Time(ns)', fontsize = 14, fontweight = 'bold')\n",
        "cbar1.set_ticks(a2)\n",
        "cbar1.set_ticklabels(b.round(decimals=3))\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "pc1=pd.DataFrame(PC1)\n",
        "pc1.to_csv(os.path.join(workDir, Output_PC1 + \".csv\"))\n",
        "pc2=pd.DataFrame(PC2)\n",
        "pc2.to_csv(os.path.join(workDir, Output_PC2 + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yce9RfNtpl-J",
        "cellView": "form"
      },
      "source": [
        "#@title Plot Principal Component 1 (PC1) and Principal Component 2 (PC2) as a ditribution\n",
        "Output_name = 'PCA_dist' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "fig = plt.figure(figsize=(9,5))\n",
        "\n",
        "plt.subplot(1, 2, 1)\n",
        "ax = sb.kdeplot(PC1, color=\"green\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('PC1', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks([])\n",
        "plt.ylabel('')\n",
        "ax.spines['top'].set_visible(False)\n",
        "ax.spines['right'].set_visible(False)\n",
        "ax.spines['bottom'].set_visible(True)\n",
        "ax.spines['left'].set_visible(False)\n",
        "plt.subplot(1, 2, 2)\n",
        "ax2 = sb.kdeplot(PC2, color=\"purple\", shade=True, alpha=0.2, linewidth=0.5)\n",
        "plt.xlabel('PC2', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks([])\n",
        "plt.ylabel('')\n",
        "ax2.spines['top'].set_visible(False)\n",
        "ax2.spines['right'].set_visible(False)\n",
        "ax2.spines['bottom'].set_visible(True)\n",
        "ax2.spines['left'].set_visible(False)\n",
        "\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pTDb7CEfkLq1",
        "cellView": "form"
      },
      "source": [
        "#@title Pearson's Cross Correlation (CC)\n",
        "\n",
        "#@markdown **Provide output file names below:** \n",
        "Output_name = 'cross_correlation' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "traj_align = pt.align(traj_load, mask='@CA', ref=0)\n",
        "\n",
        "mat_cc = matrix.correl(traj_align, '@CA')\n",
        "\n",
        "ax = plt.imshow(mat_cc, cmap = 'PiYG_r', interpolation = 'bicubic', vmin = -1, vmax = 1, origin='lower')\n",
        "\n",
        "plt.xlabel('Residues', fontsize = 14, fontweight = 'bold')\n",
        "plt.ylabel('Residues', fontsize = 14, fontweight = 'bold')\n",
        "plt.xticks(fontsize = 12)\n",
        "plt.yticks(fontsize = 12)\n",
        "cbar1 = plt.colorbar()\n",
        "cbar1.set_label('$CC_ij$', fontsize = 14, fontweight = 'bold')\n",
        "\n",
        "plt.savefig(os.path.join(workDir, Output_name + \".png\"), dpi=600, bbox_inches='tight')\n",
        "\n",
        "raw_data=pd.DataFrame(mat_cc)\n",
        "raw_data.to_csv(os.path.join(workDir, Output_name + \".csv\"))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}