5_custom_shooting_setup.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Customizing shooting moves: Setup\n",
    "\n",
    "The tutorial focuses on how to customize shooting moves in the move scheme. We discuss three major reasons to do that:\n",
    "\n",
    "1. To use a different kind of shooting move, such as two-way shooting, when needed.\n",
    "2. To use a different shooting point selector, such as a Gaussian biased selector, to get better efficiency.\n",
    "3. To create move schemes that only sample part of the network, for example, to perform TIS with each ensemble sampled in parallel.\n",
    "\n",
    "In addition, you'll learn about:\n",
    "\n",
    "* Creating a \"setup\" file with engine and other simulation information.\n",
    "* Using the OpenPathSampling command line interface (CLI) to run simulations from the setup file.\n",
    "\n",
    "Note that you'll need to install the CLI, which can be done with either `pip install openpathsampling-cli` or `conda install -c conda-forge openpathsampling-cli`. (Already installed in Binder.)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import openpathsampling as paths\n",
    "from openpathsampling import strategies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A simple toy model example system\n",
    "\n",
    "In this notebook, we'll use a simple toy system. For the most part, the ideas here directly generalize to other engines, although we'll make a few comments where the units associated with OpenMM require special care.\n",
    "\n",
    "We'll start by loading a number things from the storage file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "if sys.version_info < (3, 8):\n",
    "    f_version = \"\"\n",
    "else:\n",
    "    f_version = \"_38\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "storage = paths.Storage(f\"./inputs/2_state_toy{f_version}.nc\", mode='r')\n",
    "state_A = storage.volumes['A']\n",
    "state_B = storage.volumes['B']\n",
    "cv = storage.cvs['x']\n",
    "engine = storage.engines['toy_engine']\n",
    "initial_conditions = storage.tags['initial_conditions']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The toy model here uses a simple 2D potential energy surface, described by:\n",
    "\n",
    "$$V(x, y) = x^6 + y^6 - e^{-12(x+0.6)^2 - 5 y^2} - e^{-12(x-0.6)^2 - 5 y^2}$$\n",
    "\n",
    "We can visualize the toy engine using the `toy_plot_helpers` included in this repository:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run toy_plot_helpers.py\n",
    "\n",
    "pes = engine.topology.pes\n",
    "\n",
    "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": [
    "We run this simulation at a temperature of $T=0.1$, so the barrier between those wells is about $10\\ k_\\text{B} T$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Changing the shooting move type (e.g., two-way shooting)\n",
    "\n",
    "TPS is usually introduced with two-way shooting. In that algorithm, you select a frame of the trajectory, make some modification (typically changing the velocities in some way consistent with the thermodynamic ensemble), and then integrate the equations of motion forward and backward.\n",
    "\n",
    "However, this is problematic if the velocity memory is short compared to your TPS trajectory length, as is often the case in condensed phase systems. In this case, each shot might as well be a committor shot. The ideal place to shoot from is the where the committor is 1/2, where each shot has only a 50% chance of landing in the correct state. But this gives you a maximum of a 25% acceptance rate.\n",
    "\n",
    "To get around this, we typically use the one-way shooting algorithm when working with condensed matter systems. One-way shooting works under the assumption that you are using a stochastic integrator. If that's the case, then instead of changing the velocities at the shooting point, you can use the fact that the integrator will generate a new sequence of random numbers. Therefore, you create a new trajectory by only running in one direction, and that trajectory is still a valid, physical trajectory.\n",
    "\n",
    "In general, OpenPathSampling defaults to using one-way shooting. However, there are circumstances where you might want to use the older two-way shooting algorithm. For example, if the velocity memory is shorter than the period between saved frames, it might be better to use two-way shooting. If you intend to use deterministic dynamics, then one-way shooting is not possible.\n",
    "\n",
    "This part of the tutorial will show you how to use two-way shooting instead of one-way shooting in OPS. There are two ways of doing this: either replace the existing one-way shooting strategy, or create a new move scheme from scratch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tps_network = paths.TPSNetwork(initial_states=state_A, final_states=state_B)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Replacing parts of an existing move scheme\n",
    "\n",
    "Frequently, the easiest way to modify a move scheme (especially a complex move scheme) is to replace the parts that you want to change. In this case, we can start with the move scheme that only includes 1-way shooting, and work from there."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modified_scheme = paths.OneWayShootingMoveScheme(\n",
    "    network=tps_network,\n",
    "    engine=engine\n",
    ").named(\"2_way_from_1_way\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For our two-way shooting, we'll need to create a `TwoWayShootingStrategy`. This will require a modifier; to keep things simple, we'll completely randomize velocities (consistent with a given temperature) using `paths.RandomVelocities`. In order to be consistent between engines, `RandomVelocities` takes its input as the inverse temperature, $\\beta = 1/(k_\\text{B}T)$.\n",
    "\n",
    "For the toy engine, we can obtain the temperature from `engine.integ.temperature`, and we work in units where $k_\\text{B}=1$, so it is easy to calculate $\\beta$. For other engines, you'll need to use the correct value of $k_\\text{B}$. For OpenMM, you also need to worry about units: use `openmm.unit.BOLTZMANN_CONSTANT_kB`.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(engine.integ.temperature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# YOUR TURN: Set beta correctly (yes, this is as easy as you think).\n",
    "# beta = ... # fill in the ellipsis and uncomment this line"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modifier = paths.RandomVelocities(beta=beta, engine=engine)\n",
    "shooting_strategy = strategies.TwoWayShootingStrategy(modifier=modifier, engine=engine)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To modify a move scheme, just `append` new strategies. In this case, the new strategy for the group of moves called `'shooting'` will be overwritten. If you gave the `TwoWayShootingStrategy` a different string for its `group` argument, then you would create a second group of movers, and each move would have a 50/50 chance of using one-way shooting or of using two-way shooting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modified_scheme.append(shooting_strategy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building a move scheme from scratch\n",
    "\n",
    "In this particular case, our overall move scheme will not be too complicated. In fact, let's start by looking at the at the code for "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import inspect\n",
    "from IPython.display import Code\n",
    "\n",
    "Code(inspect.getsource(paths.OneWayShootingMoveScheme))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output above is the full source code for the `OneWayShootingMoveScheme`. You can see that all it does is pass the network up to the superclass's initialization, and then `append` two move strategies to itself. So we can recreate this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "two_way_scheme = paths.MoveScheme(network=tps_network).named(\"2_way\")\n",
    "global_strategy = strategies.OrganizeByMoveGroupStrategy()\n",
    "two_way_scheme.append([shooting_strategy, global_strategy])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Changing the shooting point selection\n",
    "\n",
    "A common problem in one-way shooting is that paths do not decorrelate quickly enough. This typically happens when there is a significant barrier withing the transition region, such that shooting points before the barrier always go back to the initial state (so only backward shots are accepted) and shooting points after the barrier always to do the final state (so only forward shots are accepted).\n",
    "\n",
    "The same issue manifests in two-way shooting as a significant decrease in the acceptance rate for the shooting move. One side of the barrier always creates $A\\to A$ trials, while the other side always creates $B\\to B$ trials. With fewer $A\\to B$ trials, the acceptance rate drops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "xvals = np.arange(-1.0, 1.0, 0.01)\n",
    "\n",
    "def plot_gaussian_bias(selector):\n",
    "    alpha = selector.alpha\n",
    "    x_0 = selector.l_0\n",
    "    gaussians = np.exp(-alpha*(xvals - x_0)**2)\n",
    "    plt.plot(xvals, gaussians)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "biased_shooting_scheme = paths.MoveScheme(network=tps_network).named(\"biased_shooting\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gaussian_sel = paths.GaussianBiasSelector(cv, alpha=100.0, l_0=0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_gaussian_bias(gaussian_sel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "paths.GaussianBiasSelector?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* The Gaussian selector has two parameters, $\\alpha$ and $l_0$. What do each of these parameters change?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# YOUR TURN: Create a move scheme using a two-way shooting strategy with this selector.\n",
    "# 1. Create a two-way shooting strategy that uses this Gaussian selector. (Use the \n",
    "#    selector keyword argument of TwoWayShootingStrategy.)\n",
    "# 2. Append that strategy and the global_strategy to the biased_shooting_scheme"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparing uniform vs Gaussian bias shooting point selection\n",
    "\n",
    "Now we have two different move schemes that sample the same TPS network. Both use two-way shooting with complete velocity randomization; the only difference is that one uses uniform shooting point selection whereas the other uses a Gaussian biased selection.\n",
    "\n",
    "Let's see how these two approaches compare with each other -- we'll run them for the same number of steps, and in the next notebook, we'll compare the acceptance rates.\n",
    "\n",
    "First, let's verify that the initial conditions will work for the move schemes we've created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = two_way_scheme.initial_conditions_from_trajectories(initial_conditions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = biased_shooting_scheme.initial_conditions_from_trajectories(initial_conditions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating a setup file for the OpenPathSampling CLI\n",
    "\n",
    "The files that OPS stores to always include details of how the simulation is run, as well as storing the newly generated data. However, we can also store only the information needed to run the simulation. That's the idea behind a \"setup\" file for an OPS simulation.\n",
    "\n",
    "This is the same storage you've used before, but we're using it in a different way. One of the main reasons to use \"setup\" files is that path sampling is rarely just a single computational run. Often you run several copies of your simulation simultaneously. Or perhaps after running TPS you'll want to do a committor analysis with the same state definitions. Using a setup file ensures that the objects you use are identical (no accidental copy-paste errors). Because the OPS analysis tools can also verify that these objects are the same, this also makes analysis easier.\n",
    "\n",
    "In our case, we'll use a single setup file that stores our initial conditions and both of our move schemes (and therefore, the TPS network, state definitions, etc.) Then we'll use that one setup file to run both of our simulations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shooting_setup = paths.Storage(\"shooting_setup.nc\", mode='w')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# saving everything will take a few minutes\n",
    "shooting_setup.tags['initial_conditions'] = initial_conditions\n",
    "shooting_setup.save(two_way_scheme)\n",
    "shooting_setup.save(biased_shooting_scheme)\n",
    "shooting_setup.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running simulations with the CLI\n",
    "\n",
    "If you've installed the OpenPathSampling CLI, you will have the command `openpathsampling` available in your shell. The `openpathsampling` command has several subcommands. You can see some of the common subcommands with `openpathsampling --help`, and you can also use the `--help` flag to get more information about a subcommand. For example, you can get more information about the `pathsampling` subcommand with `openpathsampling pathsampling --help`.\n",
    "\n",
    "* In a terminal, use the `openpathsampling contents` on the file `shooting_setup.nc` to list the contents of the file. How many snapshots are stored in that file? How many move schemes?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, use the `pathsampling` subcommand to run simulations with your two move schemes. Run at least 500 steps (`-n 500`), or run more if you'd like:\n",
    "\n",
    "```\n",
    "$ openpathsampling pathsampling shooting_setup.nc -o 2_way.nc --scheme 2_way -n 500\n",
    "\n",
    "$ openpathsampling pathsampling shooting_setup.nc -o biased.nc --scheme biased_shooting -n 500\n",
    "```\n",
    "\n",
    "Those simulations will take a few minutes, so this is a good time to take a quick break. In the next notebook, we'll analyze the results."
   ]
  }
 ],
 "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": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "278.594px"
   },
   "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
}