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: [![Open On Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/sef43/openmm_workshop/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 is taking a protein structure in PDB file format, solvating it in water, equilibrating, and running NVT or NPT simulations. Additionally it will demonstrate how to setup simulation checkpoints so you can run long simulations in HPC clusters.\n",
    "\n",
    "## Table of Contents\n",
    "- [Setup conda environment](#setup)\n",
    "- [Download the protein file](#download)\n",
    "- [Load a PDB file into OpenMM](#load)\n",
    "- [Choose the force-field](#ff)\n",
    "- [Solvate the protein with water and ions](#solvate)\n",
    "- [Setup system and integrator](#system)\n",
    "- [Run local minimization](#minim)\n",
    "- [Setup reporting](#reporting)\n",
    "- [Run NVT equilibration](#nvt)\n",
    "- [Run NPT production molecular dynamics](#npt)\n",
    "- [Basic analysis](#analysis)\n",
    "- [How to use checkpoints](#checkpoints)\n",
    "- [Visualization](#viz)\n",
    "\n",
    "\n",
    "The exercises are announced by the keyword **Exercise** and followed by an incomplete cell. Missing parts are indicated by:\n",
    "```python\n",
    "FIXME\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",
    "If you are using Colab you can run the cell below to install mamba in the Colab environment.\n",
    "\n",
    "**First try and change runtime type to GPU!**  \n",
    "Click \"runtime\">\"change runtime type\" and select \"GPU\" from the \"Hardware accelerator\" dropdown menu.\n",
    "\n",
    "It will work on CPU but will be slower.\n",
    "\n",
    "If you want to run on your own machine take a look at the [setup instructions](https://github.com/sef43/openmm_workshop/tree/main/setup). \n",
    "Remember you can replace `mamba` with `conda` if you have not installed `mamba`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mD8aoBG3heyP"
   },
   "outputs": [],
   "source": [
    "# Execute this cell to install mamba in the Colab environment\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": {
    "id": "bGVSPx2uheyQ"
   },
   "source": [
    "Now we can install openmm from conda-forge"
   ]
  },
  {
   "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/sef43/openmm_workshop/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.\n",
    "\n",
    "![villin](./villin.png)\n",
    "**Figure**. 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 in 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 forcefield we want to use. We will use the Amber14 forcefield and the TIP3P-FB water model. You can find out about 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 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 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",
    "\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 1nm away from the solute and fills it with water molecules. Additionally, it adds in the required number of CL- and Na+ ions to make the system charge neutral. 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 an example of the powerful units tracking and automatic conversion facility built into the OpenMM Python API that makes specifying unit-bearing quantities convenient and less error-prone. We could have equivalently specified `10*angstrom` instead of `1*nanometer` and achieved 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 combine our molecular topology and the forcefield to create a complete description of the system. This is done 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. We then create the integrator, and combine the integrator and system to create the Simulation object. Finally we set the initial atomic positions."
   ]
  },
  {
   "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)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The [System](http://docs.openmm.org/latest/api-python/generated/openmm.openmm.System.html#system) is an object than contains the complete mathematical description of the system we want to simulate. 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` advances 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 process 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 local energy minimization at the start of a simulation because the coordinates in starting configuration file might produce very large forces. \n",
    "\n",
    "Due to how the minimizer is implemented it will not print out information during the run. You will need to be patient and wait for it to complete. This minimization step should take ~1 minute on CPU and a few seconds on 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 get output from a simulation you need to add \"reporters\". We use [`DCDReporter`](http://docs.openmm.org/latest/api-python/generated/openmm.app.dcdreporter.DCDReporter.html) to write the coordinates every 1000 timesteps to 'traj.dcd' and we use [`StateDataReporter`](http://docs.openmm.org/development/api-python/generated/openmm.app.statedatareporter.StateDataReporter.html) to print the timestep, potential energy, temperature, and volume to the screen; and the same to a file called 'md_log.txt'. The Simulation object contains a list of reporters in `simulation.reporters` and we use the append method to add the reporters to it."
   ]
  },
  {
   "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 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 in a barostat to control the pressure. 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). The barostat assumes the simulation is being run at constant temperature, but it does not itself do anything to regulate the temperature. It is therefore critical that you always use it along with a Langevin integrator or Andersen thermostat, and that you specify the same temperature for both the barostat and the integrator or thermostat. Otherwise, you will get incorrect results.\n",
    "\n",
    "**Exercise 1**. You will need to replace the `FIXME` in the cell below with code to run the simulation for 10000 steps."
   ]
  },
  {
   "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)\n",
    "\n",
    "print('Running NPT')\n",
    "\n",
    "# run for 10000 steps\n",
    "FIXME\n"
   ]
  },
  {
   "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": [
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "data = np.loadtxt('md_log.txt', delimiter=',')\n",
    "\n",
    "step = data[:,0]\n",
    "potential_energy = data[:,1]\n",
    "temperature = data[:,2]\n",
    "volume = data[:,3]\n",
    "\n",
    "plt.plot(step, potential_energy)\n",
    "plt.xlabel(\"Step\")\n",
    "plt.ylabel(\"Potential energy (kJ/mol)\")\n",
    "plt.show()\n",
    "plt.plot(step, temperature)\n",
    "plt.xlabel(\"Step\")\n",
    "plt.ylabel(\"Temperature (K)\")\n",
    "plt.show()\n",
    "plt.plot(step, volume)\n",
    "plt.xlabel(\"Step\")\n",
    "plt.ylabel(\"Volume (nm^3)\")\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 be able to save checkpoints. This means you can restart them in the case of a crash. Or it means you can resume them 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 three files saved to disk that we can load in:\n",
    "1. The topology - this will be a PDB file of our solvated system.\n",
    "2. A serialized `System` -  this is an xml file that contains the forcefield 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 random number generators.\n",
    "\n",
    "The first two 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 a simulation from the timestep when the checkpoint file was last saved.\n",
    "\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. It will save\n",
    "# a checkpoint every 1000 timesteps.\n",
    "simulation.reporters.append(CheckpointReporter('checkpoint.chk', 1000))\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`CheckpointReporter` saves periodic checkpoints of a simulation. The checkpoints will overwrite one another - 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 for a set amount of wall-clock time using the [`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 in 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",
    "**Exercise 2.** You will need to add a line of code to make the simulation run for 30 seconds of wall time."
   ]
  },
  {
   "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",
    "\n",
    "# Setup a checkpoint reporter. This stores the positions, velocities, and box vectors.\n",
    "simulation.reporters.append(CheckpointReporter('checkpoint.chk', 1000))\n",
    "\n",
    "\n",
    "# write the code to run for 30 seconds of wall clock time\n",
    "FIXME\n",
    "\n"
   ]
  },
  {
   "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",
    "**Exercise 3.** You will need to add the code to create the `Simulation` object."
   ]
  },
  {
   "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",
    "\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",
    "\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": [
    "data = np.loadtxt('md_log.txt', delimiter=',')\n",
    "\n",
    "step = data[:,0]\n",
    "potential_energy = data[:,1]\n",
    "temperature = data[:,2]\n",
    "volume = data[:,3]\n",
    "\n",
    "plt.plot(step, potential_energy)\n",
    "plt.xlabel(\"Step\")\n",
    "plt.ylabel(\"Potential energy (kJ/mol)\")\n",
    "plt.show()\n",
    "plt.plot(step, temperature)\n",
    "plt.xlabel(\"Step\")\n",
    "plt.ylabel(\"Temperature (K)\")\n",
    "plt.show()\n",
    "plt.plot(step, volume)\n",
    "plt.xlabel(\"Step\")\n",
    "plt.ylabel(\"Volume (nm^3)\")\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "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 a variety of programs available (https://en.wikipedia.org/wiki/List_of_molecular_graphics_systems). A couple of the most popular ones are [VMD](https://www.ks.uiuc.edu/Research/vmd/) and [PyMol](https://pymol.org/).\n",
    "\n",
    "**Note this does not currently work in Colab**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!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": [
    "## 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
}