section_1/protein_in_water.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "dnVTT6WIheyN"
   },
   "source": [
    "# Protein in Water\n",
    "\n",
    "You can run this notebook in your browser: \n",
    "\n",
    "[![Open On Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/openmm/openmm_workshops/blob/main/section_1/protein_in_water.ipynb)\n",
    "\n",
    "\n",
    "This notebook is a good place to start if you have not used OpenMM before. It covers one of the most basic procedures in molecular simulation, which includes:\n",
    "\n",
    "1. Loading a protein structure from a PDB file\n",
    "1. Solvating the protein in water\n",
    "1. Equilibrating the system\n",
    "1. Running NVT or NPT production simulations\n",
    "\n",
    "Additionally, this notebook demonstrates how to setup simulation checkpoints, enabling you to restart simulations and perform long production runs on HPC clusters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Table of Contents\n",
    "\n",
    "- Setup conda environment\n",
    "- Download the protein file\n",
    "- Load a PDB file into OpenMM\n",
    "- Define the force field\n",
    "- Solvate the protein with water and ions\n",
    "- Setup system and integrator\n",
    "- Run local minimization\n",
    "- Setup reporting\n",
    "- Run NVT equilibration\n",
    "- Run NPT production molecular dynamics\n",
    "- Basic analysis\n",
    "- How to use checkpoints\n",
    "- Visualization\n",
    "\n",
    "\n",
    "The exercises are announced by the block\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "    ℹ️ <b>Exercise:</b> Description of the exercise.\n",
    "</div>\n",
    "\n",
    "and followed by an incomplete cell. Missing parts are indicated by:\n",
    "\n",
    "```python\n",
    "FIXME\n",
    "```\n",
    "\n",
    "which will throw an error when the cell is run."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "TUXynx95heyP"
   },
   "source": [
    "## Setup\n",
    "<a id=\"setup\"></a>\n",
    "\n",
    "### OpenMM on a local machine\n",
    "\n",
    "- If you want to run OpenMM on your own machine, please take a look at the [setup instructions](../setup/conda_setup.md). \n",
    "\n",
    "\n",
    "### OpenMM on Colab\n",
    "- If you are using Colab, you can run the cell below to install `mamba` in the Colab environment.\n",
    "\n",
    "    <div class=\"alert alert-block alert-info\">\n",
    "    ⚠️ <b>First try and change runtime type to GPU</b>\n",
    "\n",
    "    You can change to a GPU instance on Colab by clicking `runtime`→`change runtime type` and selecting `T4 GPU` from the `Hardware accelerator` dropdown menu. OpenMM runs on CPUs, but will be slower.\n",
    "    </div>\n",
    "\n",
    "- Remember that you can replace `mamba` with `conda` if you have not installed `mamba`.\n",
    "\n",
    "- The first time you run a notebook in Colab, you will get a warning like this:\n",
    "\n",
    "    ```\n",
    "    Warning: This notebook was not authored by Google\n",
    "    This notebook is being loaded from GitHub.\n",
    "    ```\n",
    "\n",
    "    You can ignore it and click the `run anyway` button."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mD8aoBG3heyP"
   },
   "outputs": [],
   "source": [
    "# Execute this cell to install mamba in the Colab environment and then install openmm\n",
    "\n",
    "if 'google.colab' in str(get_ipython()):\n",
    "    print('Running on colab')\n",
    "    !pip install -q condacolab\n",
    "    import condacolab\n",
    "    condacolab.install_mambaforge()\n",
    "else:\n",
    "    print('Not running on colab.')\n",
    "    print('Make sure you create and activate a new conda environment!')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Notes:** \n",
    "- During this step on Colab the kernel will be restarted. This will produce the error message:\n",
    "\"Your session crashed for an unknown reason. \" This is normal and can be safely ignored.\n",
    "- Installing the necessary packages may take several minutes."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "bGVSPx2uheyQ"
   },
   "source": [
    "### Install OpenMM\n",
    "\n",
    "Now we can install `openmm` from the `conda-forge` repository:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hlC5gumMheyQ"
   },
   "outputs": [],
   "source": [
    "!mamba install -y -c conda-forge openmm"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "LM6ipuwPheyQ"
   },
   "source": [
    "Test the installation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZgorYvvKheyQ"
   },
   "outputs": [],
   "source": [
    "!python -m openmm.testInstallation"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "mOLccFBeheyQ"
   },
   "source": [
    "## Download the protein structure file\n",
    "<a id=\"download\"></a>\n",
    "\n",
    "We will download the file from the workshop github repo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8uTn83fsheyR"
   },
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/openmm/openmm_workshops/main/section_1/villin.pdb"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The protein is the villin headpiece. This is small fast folding protein commonly used as a toy system. Note that this PDB file has been cleaned up and is ready for use in OpenMM. If you try and use a PDB file directly from the protein data bank you may encounter errors. Please look at the [OpenMM FAQs](https://github.com/openmm/openmm/wiki/Frequently-Asked-Questions) and [PDBfixer](https://github.com/openmm/pdbfixer).\n",
    "\n",
    "![villin](./images/villin.png)\n",
    "**Figure 1:** Villin headpiece protein."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "fmRKRiDLheyR"
   },
   "source": [
    "## Load the PDB File into OpenMM\n",
    "<a id=\"load\"></a>\n",
    "\n",
    "First we need to import OpenMM.\n",
    "We then then load in the PDB file using the [PDBFile](http://docs.openmm.org/latest/api-python/generated/openmm.app.pdbfile.PDBFile.html#openmm.app.pdbfile.PDBFile) class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Qz92cI6dheyR"
   },
   "outputs": [],
   "source": [
    "from openmm.app import *\n",
    "from openmm import *\n",
    "from openmm.unit import *\n",
    "from sys import stdout\n",
    "\n",
    "# Load the pdb file\n",
    "pdb = PDBFile('villin.pdb')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`PDBFile('file_name.pdb')` loads the PDB file from disk and puts the information into a `PDBFile` object which we have assign to the variable `pdb`. The object contains the molecular topology (atom names, residue types, bonds etc) and the atomic positions. These can be accessed as `pdb.topology` and `pdb.positions`. Take a look at the [API documentation](http://docs.openmm.org/latest/api-python/generated/openmm.app.pdbfile.PDBFile.html#openmm.app.pdbfile.PDBFile). All OpenMM classes have documentation available on the Python API reference: http://docs.openmm.org/latest/api-python/."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "Ww5gjnUiheyR"
   },
   "source": [
    "## Define the Force Field\n",
    "<a id=\"ff\"></a>\n",
    "\n",
    "We need to define the force field we want to use. In this example, we will use the Amber14 forcefield and the TIP3P-FB water model. You can explore all the forcefields available by default in OpenMM in the [documentation](http://docs.openmm.org/latest/userguide/application/02_running_sims.html?highlight=forcefield#force-fields)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Gbv1xLdpheyR"
   },
   "outputs": [],
   "source": [
    "# Specify the forcefield\n",
    "forcefield = ForceField('amber14-all.xml', 'amber14/tip3pfb.xml')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Force fields in OpenMM are defined by XML files. The line above loads in specified files. You can look at them in the OpenMM source code (e.g. [`amber14/tip3pfb.xml`](https://github.com/openmm/openmm/blob/master/wrappers/python/openmm/app/data/amber14/tip3pfb.xml)). It is also possible to create your own XML force field file. You can find details in the [user guide](http://docs.openmm.org/latest/userguide/application/05_creating_ffs.html#creating-force-fields)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "9ZNOQY8jheyR"
   },
   "source": [
    "## Solvate\n",
    "<a id=\"solvate\"></a>\n",
    "\n",
    "We can use the [`Modeller`](http://docs.openmm.org/latest/userguide/application/03_model_building_editing.html#model-building-and-editing) class to solvate the protein in a waterbox. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CVxXXCzoheyR"
   },
   "outputs": [],
   "source": [
    "# Create a Modeller object\n",
    "modeller = Modeller(pdb.topology, pdb.positions)\n",
    "\n",
    "# Solvate the protein in a box of water\n",
    "modeller.addSolvent(forcefield, padding=1.0*nanometer)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "3pJIJqvsheyR"
   },
   "source": [
    "This command creates a box that has edges at least 1 nm away from the solute and fills it with water molecules. Additionally, it adds the required number of Cl- and Na+ ions to neutralize the system's charge. Optionally, you can specify the ion concentration as an argument to [`addSolvent`](http://docs.openmm.org/latest/api-python/generated/openmm.app.modeller.Modeller.html#openmm.app.modeller.Modeller.addSolvent). \n",
    "\n",
    "Note that the `nanometer` variable is a unit definition that was imported from `openmm.unit`. This is part of the powerful units tracking and automatic conversion system built into the OpenMM Python API, which makes working with unit-bearing quantities both convenient and less error-prone. For example, we could have equivalently specified `10*angstrom` instead of `1*nanometer` to achieve the same result. You can read more about the units library [here](http://docs.openmm.org/latest/userguide/library/05_languages_not_cpp.html#units-and-dimensional-analysis)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Setup System and Integrator\n",
    "<a id='system'></a>\n",
    "\n",
    "We now need to perform the following steps to create a simulation:\n",
    "\n",
    "1. Combine the molecular topology and force field to create a complete system description using the [`ForceField`](http://docs.openmm.org/latest/api-python/generated/openmm.app.forcefield.ForceField.html#forcefield) object's [`createSystem()`](http://docs.openmm.org/latest/api-python/generated/openmm.app.forcefield.ForceField.html#openmm.app.forcefield.ForceField.createSystem) method.\n",
    "1. Create an integrator to control the simulation dynamics.\n",
    "1. Combine the integrator and system to create the `Simulation` object.\n",
    "1. Set the initial atomic positions for the system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PlyqnotRheyR"
   },
   "outputs": [],
   "source": [
    "# Create a system. Here we define some forcefield settings such as the nonbonded method\n",
    "system = forcefield.createSystem(modeller.topology, nonbondedMethod=PME, nonbondedCutoff=1.0*nanometer, constraints=HBonds)\n",
    "\n",
    "# Define the integrator. The Langevin integrator is also a thermostat\n",
    "integrator = LangevinMiddleIntegrator(300*kelvin, 1/picosecond, 0.004*picoseconds)\n",
    "\n",
    "# Create the Simulation\n",
    "simulation = Simulation(modeller.topology, system, integrator)\n",
    "simulation.context.setPositions(modeller.positions)\n",
    "simulation.context.setVelocitiesToTemperature(300*kelvin)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The [System](http://docs.openmm.org/latest/api-python/generated/openmm.openmm.System.html#system) is an important object in OpenMM than contains the complete mathematical description of the system we want to simulate. Specifically, it contains four key bits of information:\n",
    " - The set of particles in the simulation\n",
    " - The forces acting on them\n",
    " - Details of any constraints\n",
    " - The dimensions of the periodic box\n",
    "\n",
    "The `integrator` propagates the equations of motion. There are a variety of [integrators available in OpenMM](http://docs.openmm.org/latest/api-python/library.html#integrators). We are using the [LangevinMiddleIntegrator](http://docs.openmm.org/latest/api-python/generated/openmm.openmm.LangevinMiddleIntegrator.html#openmm.openmm.LangevinMiddleIntegrator), which performs Langevin dynamics.\n",
    "\n",
    "The [`Simulation`](http://docs.openmm.org/latest/api-python/generated/openmm.app.simulation.Simulation.html#openmm.app.simulation.Simulation) object manages all the processes involved in running a simulation, such as advancing time and writing output. \n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "QFmwHMqhheyR"
   },
   "source": [
    "## Local Energy Minimization\n",
    "<a id=\"minim\"></a>\n",
    "\n",
    "It is a good idea to run a local energy minimization at the start of a simulation. The initial coordinates of the system might be far from an energetically stable state, leading to very large forces that could cause the simulation to crash.\n",
    "\n",
    "Because of how the minimizer is implemented, no progress information will be printed during the run. You’ll need to be patient and wait for it to finish. This minimization step typically takes around 1 minute on a CPU and only a few seconds on a GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8XngRzXUheyS"
   },
   "outputs": [],
   "source": [
    "print(\"Minimizing energy\")\n",
    "simulation.minimizeEnergy()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "_oEETniPheyS"
   },
   "source": [
    "## Setup Reporting\n",
    "<a id=\"reporting\"></a>\n",
    "\n",
    "To obtain output from a simulation, you need to add \"reporters\" to your `Simulation` object. In OpenMM, we commonly use two types of reporters:\n",
    "\n",
    "- **`DCDReporter`**: This reporter writes the coordinates of the simulation system to a file at specified intervals. For instance, you can configure `DCDReporter` to save the coordinates every 1000 timesteps to a file named `traj.dcd`. You can find more information about it in the [DCDReporter documentation](http://docs.openmm.org/latest/api-python/generated/openmm.app.dcdreporter.DCDReporter.html).\n",
    "\n",
    "- **`StateDataReporter`**: This reporter outputs important simulation data, such as timestep, potential energy, temperature, and volume. It prints this data to the screen and also writes it to a file called `md_log.txt`. More details are available in the [StateDataReporter documentation](http://docs.openmm.org/development/api-python/generated/openmm.app.statedatareporter.StateDataReporter.html).\n",
    "\n",
    "To add these reporters to your simulation, you need to modify the `simulation.reporters` list by appending the desired reporter objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mbG1iFPEheyS"
   },
   "outputs": [],
   "source": [
    "# Write trajectory to a file called traj.dcd every 1000 steps\n",
    "simulation.reporters.append(DCDReporter('traj.dcd', 1000))\n",
    "\n",
    "# Print state information to the screen every 1000 steps\n",
    "simulation.reporters.append(StateDataReporter(stdout, 1000, step=True,\n",
    "        potentialEnergy=True, temperature=True, volume=True))\n",
    "\n",
    "# Print the same info to a log file every 100 steps\n",
    "simulation.reporters.append(StateDataReporter('md_log.txt', 100, step=True,\n",
    "        potentialEnergy=True, temperature=True, volume=True))\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "dL4ykymVheyS"
   },
   "source": [
    "## NVT Equilibration\n",
    "<a id=nvt></a>\n",
    "\n",
    "We are using a Langevin integrator, which means we are simulating in the NVT ensemble. To equilibrate the temperature, we just need to run the simulation for a given number of timesteps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "d0LsqNQUheyS"
   },
   "outputs": [],
   "source": [
    "print('Running NVT')\n",
    "simulation.step(10000)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "8QLi5oWCheyS"
   },
   "source": [
    "## NPT Production MD\n",
    "<a id=npt></a>\n",
    "\n",
    "To run our simulation in the NPT ensemble, we need to add a barostat to the system to control the pressure. For this, we can use [`MonteCarloBarostat`](http://docs.openmm.org/latest/api-python/generated/openmm.openmm.MonteCarloBarostat.html#openmm.openmm.MonteCarloBarostat). The parameters are the pressure (1 bar) and the temperature (300 K). \n",
    "\n",
    "The barostat assumes that the simulation is being run at constant temperature, but it does not regulate the temperature itself. Therefore, it is critical to always use it along with a Langevin integrator or Andersen thermostat, ensuring you specify the same temperature for both the barostat and the integrator or thermostat. Failing to do so will result in incorrect results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A66nwWBTheyS"
   },
   "outputs": [],
   "source": [
    "system.addForce(MonteCarloBarostat(1*bar, 300*kelvin))\n",
    "\n",
    "# It is important to call the reinitialize method on the simulation\n",
    "# otherwise the modifications will not be applied.\n",
    "simulation.context.reinitialize(preserveState=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then run the simulation for 10000 steps.\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "ℹ️ <b>Exercise 1</b>\n",
    "\n",
    "Replace the `FIXME` in the cell below with code to run the simulation for 10000 steps.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Running NPT')\n",
    "\n",
    "# run for 10000 steps\n",
    "FIXME"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "nrzBoLSnheyS"
   },
   "source": [
    "## Analysis\n",
    "<a id=analysis></a>\n",
    "\n",
    "We can now do some basic analysis using Python. We will plot the time evolution of the potential energy, temperature, and box volume. Remember that OpenMM itself is primarily an MD engine. For in-depth analysis of your simulations you can use other Python packages such as [MDtraj](https://www.mdtraj.org/) or [MDAnalysis](https://www.mdanalysis.org/).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Co3JDaJDheyS"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Load the data and extract \n",
    "data = np.loadtxt('md_log.txt', delimiter=',')\n",
    "step = data[:,0]\n",
    "potential_energy = data[:,1]\n",
    "temperature = data[:,2]\n",
    "volume = data[:,3]\n",
    "\n",
    "# Potential Energy\n",
    "plt.figure(figsize=(10, 10)) \n",
    "plt.subplot(3, 1, 1) \n",
    "plt.plot(step, potential_energy, color='b', linewidth=1.5)\n",
    "plt.xlabel(\"Step\")\n",
    "plt.ylabel(\"Potential Energy (kJ/mol)\")\n",
    "\n",
    "# Temperature\n",
    "plt.subplot(3, 1, 2) \n",
    "plt.plot(step, temperature, color='r', linewidth=1.5)\n",
    "plt.xlabel(\"Step\")\n",
    "plt.ylabel(\"Temperature (K)\")\n",
    "\n",
    "# Volume\n",
    "plt.subplot(3, 1, 3) \n",
    "plt.plot(step, volume, color='g', linewidth=1.5)\n",
    "plt.xlabel(\"Step\")\n",
    "plt.ylabel(\"Volume (nm³)\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "My7Z10-QheyS"
   },
   "source": [
    "## Checkpointing\n",
    "<a id=checkpoints></a>\n",
    "\n",
    "When you run long simulations, it is useful to save checkpoints. This means you can restart the simulation in the case of a crash, and also resume it if you need to fit within the time constraints of a HPC job scheduler.\n",
    "\n",
    "To run a resume a simulation we need to have 3 files saved to disk that we can load:\n",
    "\n",
    "1. The topology file. This will be a PDB file of our solvated system.\n",
    "2. A serialized `System`. This is an XML file that contains the force field settings.\n",
    "3. A checkpoint file. This is a binary file that contains the positions, velocities, box vectors, and other internal data such as the states of the random number generators used in the simulation.\n",
    "\n",
    "The first 2 files only need to be saved once, because they are constant throughout the simulation. The checkpoint file needs to be saved frequently. You can then resume the simulation from the timestep at which the checkpoint file was last saved.\n",
    "\n",
    "### Setup the Checkpoint\n",
    "\n",
    "We will create the topology file using `PDBFile` to write a PDB file of the system. We will use `XmlSerializer` of save the serialized system to an xml file. And we will use `CheckpointReporter` to regularly create checkpoint files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "imxKHC32heyS"
   },
   "outputs": [],
   "source": [
    "# Save the toplogy as a PDB file.\n",
    "with open('topology.pdb', 'w') as output:\n",
    "    PDBFile.writeFile(simulation.topology, simulation.context.getState(getPositions=True).getPositions(),output)\n",
    "\n",
    "# save a serialized version of the system. This stores the forcefield parameters.\n",
    "with open('system.xml', 'w') as output:\n",
    "    output.write(XmlSerializer.serialize(system))\n",
    "\n",
    "# Setup a checkpoint reporter. This stores the positions, velocities, and box vectors. \n",
    "# It will save a checkpoint every 1000 timesteps.\n",
    "simulation.reporters.append(CheckpointReporter('checkpoint.chk', 1000))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`CheckpointReporter` saves periodic checkpoints of a simulation. The checkpoints will overwrite one another, i.e., only the last checkpoint will be saved in the file. Loading a checkpoint will restore a simulation to a reasonably close, but usually not identical, state to when it was written. The checkpoint contains data that is highly specific to the `System`, `Platform`, and the hardware and software of the computer it was created on. If you try and load it on a computer with different hardware, it is likely to fail. Checkpoints created with different versions of OpenMM are often incompatible. \n",
    "\n",
    "For a more portable way of saving the state of a simulation, you can save the checkpoint as an XML state file. Read the [API docs](http://docs.openmm.org/development/api-python/generated/openmm.app.checkpointreporter.CheckpointReporter.html) for more information."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "LbSN1vjdheyS"
   },
   "source": [
    "### Running for a Set Time Limit\n",
    "\n",
    "We can run a simulation for a set amount of wall clock time using the `Simulation`'s [`runForClockTime`](http://docs.openmm.org/latest/api-python/generated/openmm.app.simulation.Simulation.html#openmm.app.simulation.Simulation.runForClockTime) method. By [wall clock](https://en.wikipedia.org/wiki/Elapsed_real_time) time, we mean the actual time a program runs for, as measured by looking at a clock on a wall (or a watch, or a timer, etc), as opposed to the simulated time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BpQYtum5heyS"
   },
   "outputs": [],
   "source": [
    "# run for 30 seconds\n",
    "simulation.runForClockTime(30.0*seconds)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "sVyOa336heyS"
   },
   "source": [
    "### Resume From a Checkpoint\n",
    "\n",
    "We now have the required files `'topology.pdb'`, `'system.xml'`, and `'checkpoint.chk'`. We will need to load them so we can resume the simulation from the last checkpoint. Note that we have to define the integrator again, as well as the simulation reporters. Furthermore, we have set the `append=True` flag to the DCD and StateData reporters.\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "ℹ️ <b>Exercise 2</b>\n",
    "\n",
    "Add a line of code to make the simulation run for 30 seconds of wall time.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DgAu-mflheyS"
   },
   "outputs": [],
   "source": [
    "pdb = PDBFile('topology.pdb')\n",
    "\n",
    "with open('system.xml') as input:\n",
    "    system = XmlSerializer.deserialize(input.read())\n",
    "\n",
    "# Define the integrator.\n",
    "integrator = LangevinMiddleIntegrator(300*kelvin, 1/picosecond, 0.004*picoseconds)\n",
    "\n",
    "# Create the Simulation\n",
    "simulation = Simulation(pdb.topology, system, integrator)\n",
    "\n",
    "# set the positions, velocities, and box vectors from the checkpoint file\n",
    "simulation.loadCheckpoint('checkpoint.chk')\n",
    "\n",
    "# We still need to define the reporters again\n",
    "\n",
    "# Write trajectory to a file called traj.dcd every 1000 steps\n",
    "simulation.reporters.append(DCDReporter('traj.dcd', 1000, append=True))\n",
    "\n",
    "# Print state information to the screen every 1000 steps\n",
    "simulation.reporters.append(StateDataReporter(stdout, 1000, step=True,\n",
    "        potentialEnergy=True, temperature=True, volume=True))\n",
    "\n",
    "# Print the same info to a log file every 100 steps\n",
    "simulation.reporters.append(StateDataReporter('md_log.txt', 100, step=True,\n",
    "        potentialEnergy=True, temperature=True, volume=True, append=True))\n",
    "\n",
    "# Setup a checkpoint reporter. This stores the positions, velocities, and box vectors.\n",
    "simulation.reporters.append(CheckpointReporter('checkpoint.chk', 1000))\n",
    "\n",
    "# write the code to run for 30 seconds of wall clock time\n",
    "FIXME"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "H1WMB3XLheyS"
   },
   "source": [
    "### Resume Multiple Times\n",
    "\n",
    "We can practice resuming multiple times. This is something you might have to do to fit a long simulation within the limits of a HPC job scheduler.\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "ℹ️ <b>Exercise 2</b>\n",
    "\n",
    "Add the code required to create the `Simulation` object.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RW2SPj2XheyS"
   },
   "outputs": [],
   "source": [
    "for i in range(3):\n",
    "    print(\"Resuming from checkpoint iteration = \", i)\n",
    "\n",
    "    pdb = PDBFile('topology.pdb')\n",
    "\n",
    "    with open('system.xml') as input:\n",
    "        system = XmlSerializer.deserialize(input.read())\n",
    "\n",
    "    # Define the integrator.\n",
    "    integrator = LangevinMiddleIntegrator(300*kelvin, 1/picosecond, 0.004*picoseconds)\n",
    "\n",
    "    # Create the Simulation\n",
    "    # write the code to create the simulation object\n",
    "    simulation = FIXME\n",
    "\n",
    "    # set the positions, velocities, and box vectors from the checkpoint file\n",
    "    simulation.loadCheckpoint('checkpoint.chk')\n",
    "\n",
    "    # We still need to define the reporters again\n",
    "\n",
    "    # Write trajectory to a file called traj.dcd every 1000 steps\n",
    "    simulation.reporters.append(DCDReporter('traj.dcd', 1000, append=True))\n",
    "\n",
    "    # Print state information to the screen every 1000 steps\n",
    "    simulation.reporters.append(StateDataReporter(stdout, 1000, step=True,\n",
    "            potentialEnergy=True, temperature=True, volume=True))\n",
    "\n",
    "    # Print the same info to a log file every 100 steps\n",
    "    simulation.reporters.append(StateDataReporter('md_log.txt', 100, step=True,\n",
    "        potentialEnergy=True, temperature=True, volume=True, append=True))\n",
    "\n",
    "    # Setup a checkpoint reporter. This stores the positions, velocities, and box vectors.\n",
    "    simulation.reporters.append(CheckpointReporter('checkpoint.chk', 1000))\n",
    "\n",
    "    # run for 30 seconds\n",
    "    simulation.runForClockTime(30.0*seconds)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "X6UXRboOheyT"
   },
   "source": [
    "### Analysis\n",
    "\n",
    "We can redo the analysis on the longer trajectory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "e2ft2bO2heyT"
   },
   "outputs": [],
   "source": [
    "# Load the data and extract \n",
    "data = np.loadtxt('md_log.txt', delimiter=',')\n",
    "step = data[:,0]\n",
    "potential_energy = data[:,1]\n",
    "temperature = data[:,2]\n",
    "volume = data[:,3]\n",
    "\n",
    "# Potential Energy\n",
    "plt.figure(figsize=(10, 10)) \n",
    "plt.subplot(3, 1, 1) \n",
    "plt.plot(step, potential_energy, color='b', linewidth=1.5)\n",
    "plt.xlabel(\"Step\")\n",
    "plt.ylabel(\"Potential Energy (kJ/mol)\")\n",
    "\n",
    "# Temperature\n",
    "plt.subplot(3, 1, 2) \n",
    "plt.plot(step, temperature, color='r', linewidth=1.5)\n",
    "plt.xlabel(\"Step\")\n",
    "plt.ylabel(\"Temperature (K)\")\n",
    "\n",
    "# Volume\n",
    "plt.subplot(3, 1, 3) \n",
    "plt.plot(step, volume, color='g', linewidth=1.5)\n",
    "plt.xlabel(\"Step\")\n",
    "plt.ylabel(\"Volume (nm³)\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "KH4DdmhAheyT"
   },
   "source": [
    "## Visualization\n",
    "<a id=\"viz\"></a>\n",
    "\n",
    "We can use the `nglview` package to view the simulation structures and trajectories in the Juyter notebook.\n",
    "\n",
    "For more serious visualization and rendering, there is a variety of programs available (https://en.wikipedia.org/wiki/List_of_molecular_graphics_systems). A couple of the most popular ones are:\n",
    "- [VMD](https://www.ks.uiuc.edu/Research/vmd/)\n",
    "- [PyMol](https://pymol.org/)\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "⚠️ <b>Note this part does not currently work in Colab</b>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if 'google.colab' in str(get_ipython()):\n",
    "    # https://github.com/googlecolab/colabtools/issues/3409\n",
    "    import locale\n",
    "    locale.getpreferredencoding = lambda: \"UTF-8\"\n",
    "\n",
    "!mamba install -y -c conda-forge nglview mdtraj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mdtraj\n",
    "import nglview\n",
    "\n",
    "traj = mdtraj.load(\"traj.dcd\", top=\"topology.pdb\")\n",
    "view = nglview.show_mdtraj(traj)\n",
    "view.add_representation('licorice',selection=\"water\")\n",
    "view"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "ℹ️ <b>Exercise 4</b>\n",
    "\n",
    "Download the files `topology.pdb` and `traj.dcd` from Colab. You can find the files in the left side pane:\n",
    "\n",
    "<img src=\"images/screenshot1.png\" alt=\"screenshot of Colab interface\">\n",
    "</div>\n",
    "\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "ℹ️ <b>Exercise 5 (optional)</b>\n",
    "    \n",
    "Open `topology.pdb` with VMD or another program of your choice. In VMD, you can then go `file`→`load data into molecule` and select `traj.dcd`. You should now be able to view the trajectory.\n",
    "\n",
    "When you first open `topology.pdb` with VMD, it will look like this:\n",
    "\n",
    "<img src=\"images/screenshot2.png\" alt=\"screenshot of Colab interface\">\n",
    "\n",
    "The blue region of atoms in the middle is the protein and the red atoms are the water box.\n",
    "</div>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next\n",
    "Now go to the protein-ligand complex [notebook](./protein_ligand_complex.ipynb)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solutions\n",
    "<a id=\"solutions\"></a>\n",
    "\n",
    "*Exercise 1.* Run for 1000 steps:\n",
    "```python\n",
    "simulation.step(10000)\n",
    "```\n",
    "\n",
    "*Exercise 2.* Run for 30 seconds of wall time:\n",
    "```python\n",
    "simulation.runForClockTime(30.0*seconds)\n",
    "```\n",
    "\n",
    "*Exercise 3.* Create the simulation:\n",
    "```python\n",
    "simulation = Simulation(pdb.topology, system, integrator)\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}