7_parallel_tis_setup.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3df2ac2d-30f1-4010-ae44-d024bedb1544",
   "metadata": {},
   "source": [
    "# (Embarrassingly) Parallel TIS vs. RETIS: Setup\n",
    "\n",
    "The OPS `DefaultScheme` is designed to provide reasonable default behaviors for TIS. These include replica exchange moves, path reversal moves, the minus interface move, as well as shooting. In general, replica exchange TIS is a more efficient way to sample than TIS without replica exchange. However, replica exchange TIS is much harder to parallelize, because the duration of each trajectory is not known before running the trajectory.\n",
    "\n",
    "Therefore, in some cases you may want to sample each interface independently. This allows a naïve parallelization, since each interface is its own simulation.\n",
    "\n",
    "In this notebook, we'll make move schemes for a custom version of RETIS (similar to the default scheme, but without the minus move) and for running each interface of a TIS simulation independently (in parallel). We'll use the OPS CLI to run it, and in the next notebook we'll compare the results.\n",
    "\n",
    "We'll use the same simple toy model as we used in notebooks 5 and 6."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c1fb77d-673a-4b59-bfb9-5c048476b414",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openpathsampling as paths\n",
    "from openpathsampling import strategies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1166098e",
   "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,
   "id": "de08a51c-cc74-4c8a-ba97-4049a17e8d65",
   "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": "code",
   "execution_count": null,
   "id": "69347fa2-0532-413d-b2e3-adbfc6d5728c",
   "metadata": {},
   "outputs": [],
   "source": [
    "interfaces = paths.VolumeInterfaceSet(cv=cv, minvals=float(\"-inf\"), \n",
    "                                      maxvals=[-0.6, -0.45, -0.375, -0.3, -0.2])\n",
    "tis_network = paths.MISTISNetwork([(state_A, interfaces, state_B)]).named(\"tis\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70bbba4c-d885-4d4c-a7ce-194f2c864447",
   "metadata": {},
   "source": [
    "All move schemes need a \"global\" level strategy. This is usually the `OrganizeByMoveGroupStrategy`. We can re-use the same strategy in all our move schemes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3c99594-cb39-45a9-a9e8-feda08b49f4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "global_strategy = strategies.OrganizeByMoveGroupStrategy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7399b89-9acd-4003-a961-05ecb0bd8882",
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is basically the DefaultScheme without the minus interface move\n",
    "retis_scheme = paths.MoveScheme(network=tis_network).named(\"retis\")\n",
    "retis_scheme.append([\n",
    "    strategies.OneWayShootingStrategy(engine=engine),\n",
    "    strategies.PathReversalStrategy(),\n",
    "    strategies.NearestNeighborRepExStrategy(),\n",
    "    global_strategy\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0da16b81-da5b-4d89-995d-ec50514730cc",
   "metadata": {},
   "source": [
    "The `OneWayShootingStrategy` includes an `ensembles` option, which selects specific ensembles to use. The (normal TIS) ensembles sampled by the TIS network are in the attribute `sampling_ensembles`. (Aside: other ensembles, such as the minus interface ensembles, are in the `special_ensembles` attribute.) This means that we can create a shooting strategy only samples a single ensemble (instead of the default, which is to sample all the `sampling_ensembles` in the network)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fb5c252-17c7-4501-8832-f1bab84fbaf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "ens_0_strategy = strategies.OneWayShootingStrategy(\n",
    "    ensembles=[tis_network.sampling_ensembles[0]],\n",
    "    engine=engine\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9641b304-0a4a-45b0-8e5a-bb2934169086",
   "metadata": {},
   "source": [
    "From here, we can make a moves scheme for each interface. I'll do that manually, although in practice you might come up with a loop to make this easier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "184183fd-795e-4d4b-af1b-8714ba4700d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "scheme_0 = paths.MoveScheme(tis_network).named(\"scheme_0\")\n",
    "scheme_1 = paths.MoveScheme(tis_network).named(\"scheme_1\")\n",
    "scheme_2 = paths.MoveScheme(tis_network).named(\"scheme_2\")\n",
    "scheme_3 = paths.MoveScheme(tis_network).named(\"scheme_3\")\n",
    "scheme_4 = paths.MoveScheme(tis_network).named(\"scheme_4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c46b83a-a25e-4cb8-8e80-683f33f14276",
   "metadata": {},
   "outputs": [],
   "source": [
    "# YOUR TURN: Make the correct strategies and append things to the scheme\n",
    "# 1. Create a OneWayShootingStrategy for each ensemble\n",
    "# 2. Append the global_strategy and the appropriate shooting strategy to each scheme"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48d17350-73c8-402b-b3ba-156cb7a611cc",
   "metadata": {},
   "source": [
    "## Running TIS\n",
    "\n",
    "Again, we'll use the command line interface to run the TIS. So the first stage is to save the relevant things to a setup file, and then we can use the `--scheme` option in the `pathsampling` command to select which scheme to run.\n",
    "\n",
    "First, we check that all our schemes match our initial conditions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b44299f-23a4-4255-bfc2-12caef6ef913",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = retis_scheme.initial_conditions_from_trajectories(initial_conditions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95ed4cb2-4786-4822-a378-db3137263a74",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = scheme_0.initial_conditions_from_trajectories(initial_conditions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12e8d979-8a26-4b97-9f01-874c71ef58e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = scheme_1.initial_conditions_from_trajectories(initial_conditions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f427e394-faba-45c3-8b54-da62990476f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = scheme_2.initial_conditions_from_trajectories(initial_conditions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48056996-0f5c-4c84-8058-7bd744c63093",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = scheme_3.initial_conditions_from_trajectories(initial_conditions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0cf2b3f-908e-4427-96d5-334cecc46717",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = scheme_4.initial_conditions_from_trajectories(initial_conditions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "500beb8c-645c-4710-b38d-9120b163911e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# saving everything will take a few minutes\n",
    "parallel_setup = paths.Storage(\"parallel_setup.nc\", mode='w')\n",
    "parallel_setup.tags['initial_conditions'] = initial_conditions\n",
    "parallel_setup.save(retis_scheme)\n",
    "parallel_setup.save(scheme_0)\n",
    "parallel_setup.save(scheme_1)\n",
    "parallel_setup.save(scheme_2)\n",
    "parallel_setup.save(scheme_3)\n",
    "parallel_setup.save(scheme_4)\n",
    "parallel_setup.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bc1f3d5-94cd-4bf9-9695-3cccc24a5f7d",
   "metadata": {},
   "source": [
    "In OPS, each individual move, such as an attempt to swap a specific pair of replicas, counts as a Monte Carlo step. So in order to make a fair comparison of the approaches with an without replica exchange, we want to ensure that they both have about the same number of shooting moves for each ensemble.\n",
    "\n",
    "However, the `MoveScheme` can give us an estimate of how many total moves are required to get a certain number of moves of a certain mover. To get 300 trials of the 0th (only!) shooting mover in `scheme_0`, how many total steps do we need?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd514607-c420-47b1-8f69-d85ef9e7f1ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "scheme_0.n_steps_for_trials(scheme_0.movers['shooting'][0], 300)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dcb1f5e-8fde-4763-b4a1-3779658fc8e7",
   "metadata": {},
   "source": [
    "That answer is probably pretty obvious. But what about our RETIS scheme? How many total steps to we need (on average) to get 300 trials of the 0th shooting mover from that move scheme?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8df4dcf3-8712-4108-9865-e50b912d660d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# YOUR TURN: Answer the question above"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b524c507-e7e6-4f03-af94-9f2585d1b8c3",
   "metadata": {},
   "source": [
    "This should be a significantly larger number, and it is due to the many replica exchange and path reversal moves in that move scheme, as well as the fact that there are multiple ensembles to sample."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce816c75-03c9-467c-b78b-5e1649d6e114",
   "metadata": {},
   "source": [
    "Now let's run the simulations. First, we equilibrate the initial condition, since it is a transition trajectory, which is highly unlikely in inner interfaces. You can do this with:\n",
    "\n",
    "```\n",
    "$ openpathsampling equilibrate parallel_setup.nc -o equil_retis.nc --scheme retis --extra-steps 50\n",
    "```\n",
    "\n",
    "That will first run until the first decorrelated path (no frames in common with the initial trajectory), and then run an additional 50 MC steps. The results will be saved in the `equil_retis.nc` file.\n",
    "\n",
    "Then you can run the full simulation with:\n",
    "```\n",
    "$ openpathsampling pathsampling equil_retis.nc -o retis.nc -n $NSTEPS > retis.out &\n",
    "```\n",
    "where you should replace `$NSTEPS` with the number of steps you found for RETIS above. \n",
    "\n",
    "If you're not familiar, the `> retis.out` redirects the output to the file `retis.out` (you can `tail retis.out` to see progress updates), and the `&` at the end of the command forces the command to run in the background, so that you can issue more commands from the same command line (i.e., run multiple things in parallel).\n",
    "\n",
    "* Why don't you need to specify a `--scheme` with the second command? (Hint: use the `openpathsampling contents` command on `retis_equil.nc` and `parallel_setup.nc`. How many move schemes are saved in each?) \n",
    "\n",
    "\n",
    "Do the same for `scheme_0`, `scheme_1`, `scheme_2`, `scheme_3` and `scheme_4`, running the path sampling with 300 steps each. In this way, you will be running all 5 interfaces in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed256067",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}