1_tps_sampling_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TPS Sampling Tutorial\n",
    "\n",
    "In this tutorial, you will set up and run a transition path sampling simulation for the $C_{7eq}$ to $\\alpha_R$ transition in alanine dipeptide. This is similar to the work in [Bolhuis, Dellago, and Chandler. PNAS **97**, 5877 (2000)](http://dx.doi.org/10.1073/pnas.100127697).\n",
    "\n",
    "You will learn aspects of three different tools:\n",
    "\n",
    "* [`OpenMM`](http://openmm.org), which you will use to perform the actual MD\n",
    "* [`MDTraj`](http://mdtraj.org), which you will use to select atoms from the topology (similar to making Gromacs `idx` files) and to calculate collective variables\n",
    "* [`OpenPathSampling`](http://openpathsampling.org), which will manage the path sampling simulation\n",
    "\n",
    "You will also use a few other useful libraries:\n",
    "\n",
    "* [`openmmtools`](https://github.com/choderalab/openmmtools), a library of additions that are not included in the OpenMM core\n",
    "* [`matplotlib`](http://matplotlib.org/), a library for creating plots\n",
    "\n",
    "We start off by importing the libraries we'll need. It's common to keep all the imports at the top of the file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from openmm import app\n",
    "import openmm as mm\n",
    "import openmm.unit as unit\n",
    "import openmmtools\n",
    "\n",
    "import openpathsampling as paths\n",
    "import openpathsampling.engines.openmm as ops_openmm\n",
    "\n",
    "import mdtraj as md"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the engine\n",
    "\n",
    "Path sampling methods wrap around molecular dynamics, so one of the first things you need to do when setting up a path sampling simulation is to set up a molecular dynamics engine. In this example, we'll use OpenMM as the underlying engine."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next cell shows you how to set up several things specific to OpenMM. We'll be running in the $NVT$ ensemble, with $T=300 K$. We're using the [Velocity Verlet with Velocity Randomization (VVVR) integrator](http://arxiv.org/abs/1301.3800), which simulates Langevin dynamics. Note that the integrator itself comes from [`openmmtools`](https://github.com/choderalab/openmmtools), a library that extends OpenMM. You should always use a reversible integrator when performing path sampling simulations. The default integrators in OpenMM are leapfrog-based, and therefore not reversible.\n",
    "\n",
    "You can learn a lot more about setting up OpenMM simulations from the [OpenMM documentation](http://docs.openmm.org/). However, it is often even easier to use [OpenMM-Setup](https://github.com/openmm/openmm-setup/) to learn how to set up the simulation the way you'd like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this cell is all OpenMM specific\n",
    "forcefield = app.ForceField('amber96.xml', 'tip3p.xml')\n",
    "system = forcefield.createSystem(\n",
    "    topology=app.PDBFile(\"AD_initial_frame.pdb\").topology, \n",
    "    nonbondedMethod=app.PME, \n",
    "    nonbondedCutoff=1.0*unit.nanometers,\n",
    "    constraints=app.HBonds, \n",
    "    rigidWater=True,\n",
    "    ewaldErrorTolerance=0.0005\n",
    ")\n",
    "\n",
    "integrator = openmmtools.integrators.VVVRIntegrator(\n",
    "    300*unit.kelvin,\n",
    "    1.0/unit.picoseconds,\n",
    "    2.0*unit.femtoseconds\n",
    ")\n",
    "integrator.setConstraintTolerance(0.00001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we deal with the interface between OpenMM and OPS. We create an OPS `Engine` object that wraps around the OpenMM system and integrator you just created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "openmm_properties = {}\n",
    "engine_options = {\n",
    "    'n_frames_max': 2000,\n",
    "    'n_steps_per_frame': 10\n",
    "}\n",
    "ops_topology = ops_openmm.tools.topology_from_pdb(\"AD_initial_frame.pdb\")\n",
    "engine = ops_openmm.Engine(\n",
    "    topology=ops_topology, \n",
    "    system=system, \n",
    "    integrator=integrator, \n",
    "    openmm_properties=openmm_properties,\n",
    "    options=engine_options\n",
    ")\n",
    "engine.name = 'TPS MD Engine'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining collective variables\n",
    "\n",
    "We'll define our states based on the [Ramachandran angles](https://en.wikipedia.org/wiki/Ramachandran_plot) $\\phi$ and $\\psi$. These are dihedrals along the protein backbone, defined for residue $i$ as involving the following atoms:\n",
    "\n",
    "* $\\phi_i$: C$^{(i-1)}$, N$^{(i)}$, C$_\\alpha^{(i)}$, C$^{(i)}$\n",
    "* $\\psi_i$: N$^{(i)}$, C$_\\alpha^{(i)}$, C$^{(i)}$, N$^{(i+1)}$\n",
    "\n",
    "where the parenthetical superscript indicates residue number. Note that the order of the atoms is important for the computation!\n",
    "\n",
    "We'll need the atom numbers for these (just as Gromacs put those atom numbers into an `idx` file). To find them, we'll use MDTraj's `Topology` object. First, we convert the OPS `Topology` object to an MDTraj `Topology` object.\n",
    "\n",
    "> NB: There are several types of `Topology` objects. See, for example, what you get from `type(engine.topology)` vs. `type(engine.simulation.topology)` and then `type(md_topology)` after you make it. The atom selection language discussed below is only for MDTraj topology objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "md_topology = md.Topology.from_openmm(engine.simulation.topology)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MDTraj has an advanced [atom selection language](http://mdtraj.org/1.8.0/atom_selection.html). Here we'll print out the backbone atoms and their atom numbers, since we know the atoms for our collective variables are from the backbone."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for atom_number in md_topology.select(\"backbone\"):\n",
    "    print(atom_number, md_topology.atom(atom_number))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There's only one pair of Ramachandran angles of interest here: the ones for residue 1, alanine. To figure out which atoms are needed, we use the definition of $\\phi_i$, which tells us we need C$^{(0)}$, N$^{(1)}$, C$_\\alpha^{(1)}$, and C$^{(1)}$; i.e., `ACE1-C`, `ALA2-N`, `ALA2-CA`, and `ALA2-C`. We can find the correct atom indices from the result of the previous cell. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ACE1-C, ALA2-N, ALA2-CA, ALA2-C: 4,6,8,14\n",
    "phi = paths.MDTrajFunctionCV(\"phi\", md.compute_dihedrals, engine.topology, indices=[[4,6,8,14]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# YOUR TURN: Define `psi`, similarly to `phi` above.\n",
    "# Figure out which atoms make the dihedral, find their atom indices, then create the CV."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining the stable states\n",
    "\n",
    "Once we've created our collective variables, we use them to define our stable states. For this system, we'll define the stable states as follows:\n",
    "\n",
    "* $C_{7\\text{eq}}$: $-180 \\le \\phi \\le 0$ and $100 \\le \\psi \\le 200$\n",
    "* $\\alpha_R$: $-180 \\le \\phi \\le 0$ and $-100 \\le \\psi \\le 0$\n",
    "\n",
    "with the angles measured in degrees.\n",
    "\n",
    "Note that OPS allows you to intersect two volumes with `&`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deg = 180.0/np.pi  # for conversion between radians and degrees\n",
    "\n",
    "C_7eq = (\n",
    "    paths.PeriodicCVDefinedVolume(phi, lambda_min=-180/deg, lambda_max=0/deg, \n",
    "                                  period_min=-np.pi, period_max=np.pi)\n",
    "    & paths.PeriodicCVDefinedVolume(psi, lambda_min=100/deg, lambda_max=200/deg,\n",
    "                                    period_min=-np.pi, period_max=np.pi)\n",
    ").named(\"C_7eq\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# YOUR TURN: define the `alpha_R` state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may have noticed that, even though the angles are defined from -180 degrees to 180 degrees, the C$_{7\\text{eq}}$ state is defined $\\psi$ up to 200 degrees. `PeriodicCVDefinedVolume` knows how to correctly handle volumes that wrap around a periodic boundary."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the reaction network and move scheme\n",
    "\n",
    "In OPS, path sampling simulations sample networks of reactions (even if the network only involves a simple $A\\to B$ transition, as is the case here). For (flexible length) transition path sampling, use the `TPSNetwork` object.\n",
    "\n",
    "OPS also allows you to set up many different kinds of move schemes. In flexible length transition path sampling, there's usually only one mover in the move scheme: a shooting mover. The `OneWayShootingMoveScheme` sets that up quickly. We will use uniform shooting point selection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "network = paths.TPSNetwork(initial_states=C_7eq, final_states=alpha_R)\n",
    "scheme = paths.OneWayShootingMoveScheme(network=network, \n",
    "                                        selector=paths.UniformSelector(),\n",
    "                                        engine=engine).named(\"tps_scheme\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting an initial trajectory\n",
    "\n",
    "In practice, getting the initial trajectory can be one of the hardest parts of path sampling. For this example, you could easily do it by running high temperature MD to get an unphysical path, and then equilibrate it for a while. However, for the purposes of this tutorial, I already did that for you. So we'll just load the equilibrated trajectory from a file.\n",
    "\n",
    "Note that this trajectory comes from an OPS NetCDF file. This should not be confused with the Amber NetCDF trajectory format. NetCDF is a storage backend that can handle arbitrary data. The way that data is organized for OPS is different from the way it is organized in Amber, so the file formats are not compatible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_traj_storage = paths.Storage(\"initial_trajectory.nc\", 'r')\n",
    "init_traj = init_traj_storage.trajectories[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The CVs you created act like functions when applied to a snapshot or trajectory. This allows us to easily plot the trajectory in $\\phi$ and $\\psi$, using `matplotlib`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(phi(init_traj), psi(init_traj))\n",
    "plt.xlabel(\"$\\phi$\")\n",
    "plt.ylabel(\"$\\psi$\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step associates the initial trajectory with the ensemble we'll be sampling. This becomes important when there are multiple ensembles, as with transition interface sampling, and you must ensure that the initial trajectories are correct for each ensemble."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_conditions = scheme.initial_conditions_from_trajectories(init_traj)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This should tell you that there are no missing ensembles and no extra ensembles. If not, then there's a problem with your definition of the $\\alpha_R$ state or with your definition of the $\\psi$ collective variable. Check those very carefully."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the simulation\n",
    "\n",
    "Now all the parts are in place, and we can run the simulation! First, we open a file to store everything in, and then we build the simulation and run it.\n",
    "\n",
    "In practice, for a long-running simulation you might save all the relevant information to a \"setup\" file, and then use the OPS Command Line Interface (CLI) to run it (perhaps on a remote computer). The advantage of that approach is that you can re-use the same exact objects, guaranteeing that they will behave the same way and making analysis easier. That approach is described in notebook 5 of this tutorial. In this notebook, we'll directly run the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "storage = paths.Storage(\"tps_simulation_results.nc\", \"w\", template=init_traj[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you get a `PermissionError` when doing that, it probably means that you have already opened this file for writing. Use `storage.close()` to close the old file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = paths.PathSampling(storage=storage,\n",
    "                             move_scheme=scheme,\n",
    "                             sample_set=initial_conditions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler.run(40)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "storage.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You're done running the (short) simulation. In the next notebook, we'll analyze these results, and then do a more thorough analysis of a longer simulation."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}