4_mstis_sampling_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Running an MSTIS simulation\n",
    "\n",
    "This notebook gives a run-through of how to set up multiple state transition interface sampling. Here, we will use the toy model engine that is built into OPS. However, you'll see that the process of setting up a path sampling simulation is the same: set up engine, then CVs, then states (and interfaces for TIS), then the sampling network, then the move scheme. After that, you're ready to run your simulation!\n",
    "\n",
    "This notebook does not include any \"YOUR TURN\" steps; it can just be run directly.\n",
    "\n",
    "Tasks covered in this notebook:\n",
    "* Loading OPS objects from storage\n",
    "* Ways of assigning initial trajectories to initial samples\n",
    "* Setting up a path sampling simulation with various move schemes\n",
    "* Visualizing trajectories while the path sampling is running"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "%matplotlib inline\n",
    "import sys\n",
    "import openpathsampling as paths\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "# toy_plot_helpers.py has conveniences for plotting the 2D contour plots, etc.\n",
    "%run ./toy_plot_helpers.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the engine\n",
    "\n",
    "First we set up our system: for the toy dynamics, this involves defining a potential energy surface (PES), setting up an integrator, and giving the simulation an initial configuration. In real MD systems, the PES is handled by the combination of a topology file and a force field definition, and the initial configuration would come from a file instead of being described by hand.\n",
    "\n",
    "First we need to describe the system we'll be simulating. With biomolecular systems, this is often done with an initial PDB structure and a choice of force field. For the toy model, we need to give a snapshot as a template, as well as a potential energy surface. The template snapshot also includes a pointer to the topology information (which is relatively simple for the toy systems.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convenience for the toy dynamics\n",
    "import openpathsampling.engines.toy as toys\n",
    "\n",
    "# Toy_PES supports adding/subtracting various PESs. \n",
    "# The OuterWalls PES type gives an x^6+y^6 boundary to the system.\n",
    "pes = (\n",
    "    toys.OuterWalls(sigma=[1.0, 1.0], x0=[0.0, 0.0])\n",
    "    + toys.Gaussian(A=-0.7, alpha=[12.0, 12.0], x0=[0.0, 0.4])\n",
    "    + toys.Gaussian(A=-0.7, alpha=[12.0, 12.0], x0=[-0.5, -0.5]) \n",
    "    + toys.Gaussian(A=-0.7, alpha=[12.0, 12.0], x0=[0.5, -0.5])\n",
    ")\n",
    "\n",
    "topology=toys.Topology(\n",
    "    n_spatial=2,\n",
    "    masses=[1.0, 1.0],\n",
    "    pes=pes\n",
    ")\n",
    "\n",
    "integ = toys.LangevinBAOABIntegrator(dt=0.02, temperature=0.1, gamma=2.5)\n",
    "\n",
    "options={\n",
    "    'integ' : integ,\n",
    "    'n_frames_max' : 5000,\n",
    "    'n_steps_per_frame' : 1\n",
    "}\n",
    "\n",
    "engine = toys.Engine(\n",
    "    options=options,\n",
    "    topology=topology\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's look at the potential energy surface we've created:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot = ToyPlot()\n",
    "plot.contour_range = np.arange(-1.5, 1.0, 0.1)\n",
    "plot.add_pes(pes)\n",
    "fig = plot.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining collective variables\n",
    "\n",
    "TIS methods usually require that you define states and interfaces before starting the simulation. State and interfaces are both defined in terms of `Volume` objects. The most common type of `Volume` is one based on some set of collective variables, so the first thing we have to do is to define the collective variable.\n",
    "\n",
    "For this system, we'll define the collective variables as circles centered on the middle of the state. OPS allows us to define one function for the circle, which is parameterized by different centers. Note that each collective variable is in fact a separate function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def circle(snapshot, center):\n",
    "    import math\n",
    "    return math.sqrt((snapshot.xyz[0][0]-center[0])**2 + (snapshot.xyz[0][1]-center[1])**2)\n",
    "    \n",
    "opA = paths.CoordinateFunctionCV(name=\"opA\", f=circle, center=[-0.5, -0.5])\n",
    "opB = paths.CoordinateFunctionCV(name=\"opB\", f=circle, center=[0.5, -0.5])\n",
    "opC = paths.CoordinateFunctionCV(name=\"opC\", f=circle, center=[0.0, 0.4])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining states and interfaces\n",
    "\n",
    "Next we'll use those collective variables to define both states and interfaces. In this example, the innermost interface is the same as the state definition; this does not have to be the case, but when it is not, you should make sure that all frames in the state definition are also inside the innermost interface. The `VolumeInterfaceSet` gives a shortcut to create the full set of volume objects using the same collective variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stateA = paths.CVDefinedVolume(collectivevariable=opA, lambda_min=0.0, lambda_max=0.2)\n",
    "stateB = paths.CVDefinedVolume(collectivevariable=opB, lambda_min=0.0, lambda_max=0.2)\n",
    "stateC = paths.CVDefinedVolume(collectivevariable=opC, lambda_min=0.0, lambda_max=0.2)\n",
    "\n",
    "interfacesA = paths.VolumeInterfaceSet(cv=opA, minvals=0.0, maxvals=[0.2, 0.3, 0.4])\n",
    "interfacesB = paths.VolumeInterfaceSet(cv=opB, minvals=0.0, maxvals=[0.2, 0.3, 0.4])\n",
    "interfacesC = paths.VolumeInterfaceSet(cv=opC, minvals=0.0, maxvals=[0.2, 0.3, 0.4])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the MSTIS transition network\n",
    "\n",
    "Once we have the collective variables, states, and interfaces defined, we can create the entire transition network. In this one small piece of code, we create all the path ensembles needed for the simulation, organized into structures to assist with later analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ms_outers = paths.MSOuterTISInterface.from_lambdas(\n",
    "    {ifaces: 0.5\n",
    "     for ifaces in [interfacesA, interfacesB, interfacesC]}\n",
    ")\n",
    "mstis = paths.MSTISNetwork(\n",
    "    [(stateA, interfacesA),\n",
    "     (stateB, interfacesB),\n",
    "     (stateC, interfacesC)],\n",
    "    ms_outers=ms_outers\n",
    ").named('mstis')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Equilibration: Setting up the move scheme\n",
    "\n",
    "In the following, we will first do a (very) brief equilibration, and then "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "equil_scheme = paths.OneWayShootingMoveScheme(mstis, engine=engine)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Equilibration: loading the samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading from storage is very easy. Each store is a list. We take the 0th snapshot as a template (it doesn't actually matter which one) for the next storage we'll create. There's only one engine stored, so we take the only one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# some aspects of storage depend on Python version\n",
    "if sys.version_info < (3,):\n",
    "    filename = \"./inputs/mstis_bootstrap_py2.nc\"\n",
    "elif (3, 6) <= sys.version_info < (3, 8):\n",
    "    filename = \"./inputs/mstis_bootstrap_py3.nc\"\n",
    "elif (3, 8) <= sys.version_info < (3, 11):\n",
    "    filename = \"./inputs/mstis_bootstrap_py38.nc\"\n",
    "else:\n",
    "    raise RuntimeError(\n",
    "        \"Uh oh! Looks like we don't have an input file for your Python version: \"\n",
    "        + \".\".join(str(x) for x in sys.version_info[:3])\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "old_store = paths.AnalysisStorage(filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "template = old_store.snapshots[0]\n",
    "old_sampleset = old_store.samplesets[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sset = equil_scheme.initial_conditions_from_trajectories(\n",
    "    trajectories=[s.trajectory for s in old_sampleset.samples]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(sset))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point, we've loaded 9 samples, which is one for each ensemble that plays a role in the TIS sampling, which are the ensembles sampled by the one-way shooting scheme we use for equilibration."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Equilibration\n",
    "\n",
    "In molecular dynamics, you need to equilibrate if you don't start with an equilibrium frame (e.g., if you start with solvent molecules on a grid, your system should equilibrate before you start taking statistics). Similarly, if you start with a set of paths which are far from the path ensemble equilibrium, you need to equilibrate. This could either be because your trajectories are not from the real dynamics (generated with metadynamics, high temperature, etc.) or because your trajectories are not representative of the path ensemble (e.g., if you put transition trajectories into all interfaces).\n",
    "\n",
    "As with MD, running equilibration can be the same process as running the total simulation. However, in path sampling, it doesn't have to be: we can equilibrate without replica exchange moves or path reversal moves, for example. In the example below, we create a `MoveScheme` that only includes shooting movers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "equilibration = paths.PathSampling(\n",
    "    storage=None,\n",
    "    sample_set=sset,\n",
    "    move_scheme=equil_scheme\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "equilibration.run(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To continue with the equilibrated results (without having used a storage to save them), we extract the `sample_set` at the end of the calculation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sset = equilibration.sample_set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main simulation\n",
    "\n",
    "Again we set up a scheme. This time we use `DefaultScheme`, which includes shooting moves, replica exchange, path reversals, minus moves, and shooting in the multiple state outer interface (if one exists). Note that we use the same network, just a different move scheme."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scheme = paths.DefaultScheme(mstis, engine=engine).named(\"full TIS\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting up additional intial conditions\n",
    "\n",
    "Unlike the equilibration scheme we used, `DefaultScheme` involves the multiple state outer interface and the minus interfaces. Our current sample set doesn't have them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(scheme.initial_conditions_report(sset))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Minus interface ensemble\n",
    "\n",
    "The minus interface ensembles do not yet have a trajectory. We will generate them by starting with same-state trajectories (A-to-A, B-to-B, C-to-C) in each interface, and extending into the minus ensemble.\n",
    "\n",
    "* check whether the traj is A-to-A\n",
    "* extend\n",
    "\n",
    "First we need to make sure that the trajectory in the innermost ensemble of each state also ends in that state. This is necessary so that when we extend the trajectory, it can extends into the minus ensemble.\n",
    "\n",
    "If the trajectory isn't right, we run a shooting move on it until it is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this first part is only really important when not working interactively\n",
    "# interactively, you can probably find an appropriate trajectory on your own\n",
    "for transition in mstis.sampling_transitions:\n",
    "    innermost_ensemble = transition.ensembles[0]\n",
    "    shooter = None\n",
    "    if not transition.stateA(sset[innermost_ensemble].trajectory[-1]):\n",
    "        shooter = paths.OneWayShootingMover(ensemble=innermost_ensemble,\n",
    "                                            selector=paths.UniformSelector(),\n",
    "                                            engine=engine)\n",
    "        pseudoscheme = paths.LockedMoveScheme(root_mover=shooter)\n",
    "        pseudosim = paths.PathSampling(storage=None, \n",
    "                                       move_scheme=pseudoscheme, \n",
    "                                       sample_set=sset,\n",
    "                                      )\n",
    "    while not transition.stateA(sset[innermost_ensemble].trajectory[-1]):\n",
    "        pseudosim.run(1)\n",
    "        sset = pseudosim.sample_set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that all the innermost ensembles are safe to use for extending into a minus interface, we extend them into a minus interface:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "minus_samples = []\n",
    "for transition in mstis.sampling_transitions:\n",
    "    minus_samples.append(transition.minus_ensemble.extend_sample_from_trajectories(\n",
    "        sset[transition.ensembles[0]].trajectory,\n",
    "        replica=-len(minus_samples)-1,\n",
    "        engine=engine\n",
    "    ))\n",
    "sset = sset.apply_samples(minus_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(scheme.initial_conditions_report(sset))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Mutliple state outer ensemble (`UnionEnsemble`)\n",
    "\n",
    "The missing ensemble is the multiple state outer ensemble. As it happens, there's actually a trajectory from the initial file that will satisfy it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sset = scheme.initial_conditions_from_trajectories(\n",
    "    trajectories=[s.trajectory for s in old_sampleset.samples],\n",
    "    sample_set=sset\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the differences between this and when we filled the equilibration list: first, the ensembles that get filled by `MoveScheme.initial_conditions_from_trajectories` depend on which ensembles are used by the move scheme. The network may define ensembles that are unused. Second, here we want to keep most of the sample set unchanged, so we give an initial sample set. In this case, the method only appends new samples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running RETIS\n",
    "\n",
    "Now we have our initial conditions, so we'll make the storage and run the full calculation. \n",
    "\n",
    "Up to here, we haven't been storing any of our results. This time, we'll start a storage object, and we'll save the network we've created. Then we'll run a new `PathSampling` calculation object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "storage = paths.storage.Storage(\"mstis.nc\", \"w\", template=template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mstis_calc = paths.PathSampling(\n",
    "    storage=storage,\n",
    "    sample_set=sset,\n",
    "    move_scheme=scheme\n",
    ")\n",
    "mstis_calc.save_frequency = 50"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next block sets up a live visualization. This is optional, and only recommended if you're using OPS interactively (which would only be for very small systems). Some of the same tools can be used to play back the behavior after the fact if you want to see the behavior for more complicated systems. You can create a background (here we use the PES contours), and the visualization will plot the trajectories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xval = paths.FunctionCV(\"xval\", lambda snap : snap.xyz[0][0])\n",
    "yval = paths.FunctionCV(\"yval\", lambda snap : snap.xyz[0][1])\n",
    "mstis_calc.live_visualizer = paths.StepVisualizer2D(\n",
    "    network=mstis, \n",
    "    cv_x=xval, \n",
    "    cv_y=yval,\n",
    "    xlim=[-1.0, 1.0],\n",
    "    ylim=[-1.0, 1.0]\n",
    ")\n",
    "background = ToyPlot()\n",
    "background.contour_range = np.arange(-1.5, 1.0, 0.1)\n",
    "background.add_pes(engine.pes)\n",
    "mstis_calc.live_visualizer.background = background.plot()\n",
    "mstis_calc.status_update_frequency = 1 # increasing this number speeds things up, but isn't as pretty"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next question is, how many steps do we want to run? You can just choose an arbitrary number, but often we want to think in terms of how many shooting moves per ensemble. Assuming each ensemble has the same probability of doing a shooting move (as is the case in the `DefaultScheme`) we can select an arbitrary shooting mover as `scheme.movers['shooting'][0]`. Let's say we want to run enough moves that, on average, each shooting mover would do 10 steps. (In practice, this number would be more like 1000)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_steps = scheme.n_steps_for_trials(scheme.movers['shooting'][0], 10)\n",
    "print(n_steps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another important consideration is how many times other movers will be called. For example, we wouldn't want to have any mover completely left out because the probability of running it is too low. In the `DefaultScheme`, the minus movers are the least likely to run. So we check how many trials we expect for an arbitrary minus mover (all have same probability):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(scheme.n_trials_for_steps(mover=scheme.movers['minus'][0], n_steps=n_steps))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now everything is ready: let's run the simulation!\n",
    "\n",
    "Note that the `n_steps` defined above is a `float`, so we need to turn it into an `int` first:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mstis_calc.run(int(n_steps))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "storage.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.9.13"
  },
  "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": true
  },
  "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
}