8_parallel_tis_analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d975f240-bdca-458f-acee-d9dd72e04ade",
   "metadata": {},
   "source": [
    "# (Embarrassingly) Parallel TIS vs. RETIS: Analysis\n",
    "\n",
    "Now that you've run your embarrassingly parallel simulation and your RETIS simulation, let's compare the results. We'll check if the simulations seem consistent by looking at the acceptance of the shooting moves in each ensemble (you could also compare path length distributions or even do per-interface path density plots to give a more thorough check.) Then we'll look at the crossing probabilities to see if they seem to be converged (and perhaps whether one seems closer to converged than the other)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4504965b-df8b-4130-a19f-3f9ab1e5e93f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import openpathsampling as paths\n",
    "from openpathsampling.analysis import tis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15588f71-805d-4374-aa9b-580d8f1a4c65",
   "metadata": {},
   "source": [
    "We'll start out by loading the output files we generated, and from each file we load in the move scheme that it used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73d2926f-7f26-44cd-85cf-b445156efeac",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "storage_0 = paths.Storage('scheme_0.nc', mode='r')\n",
    "storage_1 = paths.Storage('scheme_1.nc', mode='r')\n",
    "storage_2 = paths.Storage('scheme_2.nc', mode='r')\n",
    "storage_3 = paths.Storage('scheme_3.nc', mode='r')\n",
    "storage_4 = paths.Storage('scheme_4.nc', mode='r')\n",
    "storage_retis = paths.Storage('retis.nc', mode='r')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32cebb28-471c-467e-bcdf-89b25a9bb19a",
   "metadata": {},
   "outputs": [],
   "source": [
    "scheme_0 = storage_0.schemes['scheme_0']\n",
    "scheme_1 = storage_1.schemes['scheme_1']\n",
    "scheme_2 = storage_2.schemes['scheme_2']\n",
    "scheme_3 = storage_3.schemes['scheme_3']\n",
    "scheme_4 = storage_4.schemes['scheme_4']\n",
    "scheme_retis = storage_retis.schemes['retis']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71e4b708-a37f-4cf6-83d2-d3d5a80323fb",
   "metadata": {},
   "source": [
    "## Comparing the acceptance rates\n",
    "\n",
    "Recall that we can get per-mover acceptance by passing the appropriate string to the `movers` keyword of the `move_summary`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "822feb97-0d21-4bb7-bb85-77dbfd7b86ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "scheme_retis.move_summary(storage_retis.steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "727e1c33-0679-422b-bb3c-5b9435e1fad7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "scheme_retis.move_summary(movers='shooting')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baf70277-3c76-440b-a1d1-1e6d6d3ff4db",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "scheme_retis.move_summary(movers='repex')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb6b9de7-35e7-419e-94ec-2b0928563196",
   "metadata": {},
   "outputs": [],
   "source": [
    "scheme_0.move_summary(storage_0.steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef124f49-ac95-4149-a63f-d0e1309b610b",
   "metadata": {},
   "outputs": [],
   "source": [
    "scheme_1.move_summary(storage_1.steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "704453f9-8d5f-4a7b-8bd4-6d938dfacdea",
   "metadata": {},
   "outputs": [],
   "source": [
    "scheme_2.move_summary(storage_2.steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "135c7d76-2213-4280-8e8b-be842ecfa3e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "scheme_3.move_summary(storage_3.steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c769c6f-b98d-432c-a23d-d49f236589ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "scheme_4.move_summary(storage_4.steps)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36962444-6bd1-40ef-8751-b70dbf181198",
   "metadata": {},
   "source": [
    "* Is there a major difference in the shooting move acceptance for any ensemble? Would you expect there to be one?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "043d97eb-17e6-4520-854d-34cc3d670a9f",
   "metadata": {},
   "source": [
    "## TIS analysis (crossing probabilities, etc.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e58c4ea-0ef5-43ad-8e6a-c96337bb2f71",
   "metadata": {},
   "source": [
    "We don't have the flux here, so we can't calculate the actual rates. However, we can create a fake flux that says that the flux through the out of state $A$ and through the innermost interface is `1.0`. This allows us to use the rest of the `StandardTISAnalysis` object. It just means that the rate that gets reported is actually the total transition probability.\n",
    "\n",
    "You can get the actual flux either from including a minus interface move in your TIS simulation, or from using direct MD. The `paths.TrajectoryTransitionAnalysis` class will analyze existing MD trajectories, or the `paths.DirectSimulation` class can run MD and analyze the flux on the fly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ca161ea-a4a5-4307-84ab-c8c189e120a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# because we used a setup file, the netword/state/interface are the same in both\n",
    "network = storage_retis.networks[0]\n",
    "state_A = storage_retis.volumes['A']\n",
    "state_B = storage_retis.volumes['B']\n",
    "interface_0 = network.sampling_transitions[0].interfaces[0]\n",
    "fake_flux = tis.DictFlux({(state_A, interface_0): 1.0})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ae2ed96-3d8a-439c-98dd-0165d0342ef7",
   "metadata": {},
   "source": [
    "Finally, we assemble the `StandardTISAnalysis` and perform the analysis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff24a22f-c0b3-4135-8e67-c7c7d054879f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "retis_analysis = tis.StandardTISAnalysis(\n",
    "    network=network,\n",
    "    flux_method=fake_flux,\n",
    "    max_lambda_calcs={t: {'bin_width': 0.025, 'bin_range': (-0.6, 0.6)}\n",
    "                      for t in network.sampling_transitions},\n",
    "    combiners={t.interfaces: paths.numerics.WHAM(cutoff=0.01,  # lower cutoff, default is 0.05\n",
    "                                                 interfaces=t.interfaces.lambdas)\n",
    "               for t in network.sampling_transitions},\n",
    "    steps=storage_retis.steps\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8be276a9-82c2-45c0-8205-e35321eb5ae1",
   "metadata": {},
   "source": [
    "Currently, the parallel analysis needs an extra step to run correctly. We need to create the `weighted_trajectories` object from the steps, and then perform the overall analysis using that as the input, instead of the steps themselves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a86866f-59e4-4ba1-8ab7-06643fc8a86d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# currently we need to manually join the weighted trajectories from each storage\n",
    "# Future versions of OPS will simplify this\n",
    "weighted_trajectories = {}\n",
    "storages = [storage_0, storage_1, storage_2, storage_3, storage_4]\n",
    "for storage, ensemble in zip(storages, network.sampling_ensembles):\n",
    "    weighted_trajectories.update(\n",
    "        tis.core.steps_to_weighted_trajectories(storage.steps, [ensemble])\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f974b4d-7f1a-4da4-8274-b1ae4656c54b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#state_A = storage_0.volumes['A']\n",
    "#interface_0 = network.sampling_transitions[0].interfaces[0]\n",
    "#fake_flux = tis.DictFlux({(state_A, interface_0): 1.0})\n",
    "parallel_analysis = tis.StandardTISAnalysis(\n",
    "    network=network,\n",
    "    flux_method=fake_flux,\n",
    "    max_lambda_calcs={t: {'bin_width': 0.025, 'bin_range': (-0.6, 0.6)}\n",
    "                      for t in network.sampling_transitions},\n",
    "    combiners={t.interfaces: paths.numerics.WHAM(cutoff=0.01,  # c cutoff, default is 0.05\n",
    "                                                 interfaces=t.interfaces.lambdas)\n",
    "               for t in network.sampling_transitions}\n",
    ")\n",
    "parallel_analysis.results['flux'] = fake_flux.calculate('foo')\n",
    "parallel_analysis.results = parallel_analysis.from_weighted_trajectories(weighted_trajectories)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "423f6e36-a97f-4fe3-9509-49a215ea5fd4",
   "metadata": {},
   "source": [
    "### Plotting the crossing probabilities\n",
    "\n",
    "One of the spot-checks to see if your simulation is converged is to plot the crossing probabilities functions. For each ensemble, the `StandardTISAnalysis` calculates a crossing probability along the order parameter, defined as the fraction of paths in that ensemble that reach at least the given value on the $x$ axis. As such, the crossing probability is always 1 for values less than the cutoff for the interface. Additionally, two ensemble crossing probabilities should never cross; the one from an outer interface should always be higher at a given value of the order parameter than one from an inner interface.\n",
    "\n",
    "There is also the *total* crossing probability, which is generated by using a histogram combining algorithm (usually WHAM) to combine the individual ensemble crossing probabilities into a good estimate for the true crossing probability (from the innermost interface). Like all crossing probabilities, this should be monotonically decreasing; if it is not, that is a sign of insufficient sampling.\n",
    "\n",
    "Since the y-axis is probability, and we're looking at rare events, we frequently plot crossing probabilities on a semi-log plot."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eecd58a1",
   "metadata": {},
   "source": [
    "Note: In some cases, you may get an error here about a lack of overlap between ensembles. Normally, this can be solved by continuing the simulation longer. For the purposes of this tutorial, you can probably just re-run the sampling stage and then re-run this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "372e1eb0-d3e2-49ed-9ad0-2ed5405a8a23",
   "metadata": {},
   "outputs": [],
   "source": [
    "for ensemble in network.transitions[(state_A, state_B)].ensembles:\n",
    "    crossing = retis_analysis.crossing_probability(ensemble)\n",
    "    label = \"Interface at $x$={:3.2f}\".format(ensemble.lambda_i)\n",
    "    plt.plot(crossing.x, crossing, label=label)\n",
    "\n",
    "tcp_AB = retis_analysis.total_crossing_probability[(state_A, state_B)]\n",
    "plt.plot(tcp_AB.x, tcp_AB, lw=2, color='k', label=\"Total crossing probability\")\n",
    "plt.legend(loc='upper right')\n",
    "plt.yscale('log')\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('Crossing probability')\n",
    "plt.title(\"RETIS\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98fc8567-14d5-4f03-9391-9f598eef0d7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "for ensemble in network.transitions[(state_A, state_B)].ensembles:\n",
    "    crossing = parallel_analysis.crossing_probability(ensemble)\n",
    "    label = \"Interface at x={:3.2f}\".format(ensemble.lambda_i)\n",
    "    plt.plot(crossing.x, crossing, label=label)\n",
    "\n",
    "tcp_AB = parallel_analysis.total_crossing_probability[(state_A, state_B)]\n",
    "plt.plot(tcp_AB.x, tcp_AB, lw=2, color='k', label=\"Total crossing probability\")\n",
    "plt.legend(loc='upper right')\n",
    "plt.yscale('log')\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('Crossing probability')\n",
    "plt.title('Parallel TIS');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4c66d93-fd79-4745-9ffb-45c97a01859a",
   "metadata": {},
   "source": [
    "* Are your total crossing probabilities monotonically decreasing?\n",
    "* Do your individual ensemble crossing probabilities cross each other?\n",
    "* Based on the crossing probability plots, which approach seems more converged?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d233219d-2d88-4f69-b98e-d5a9a3dc4545",
   "metadata": {},
   "source": [
    "OPS keeps track of how long each sampling step took. We can look at that to get a rough comparison of the sampokling times (note that other aspects of the computing environment, such as other processes running at the same time, may have a significant effect on the timings here)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5721a2d6-3e8f-448c-8586-35d430e702e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_total_time(storage):\n",
    "    # step 0 (init conds) doesn't have timing data; all the others do\n",
    "    return sum(step.change.details.timing for step in storage.steps[1:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82efcac2-253b-4e53-be84-e41e803643e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "calculate_total_time(storage_retis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7db6cf9-2a56-49ae-b29d-02ef7762a6e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_times = [calculate_total_time(storage)\n",
    "               for storage in [storage_0, storage_1, storage_2, storage_3, storage_4]]\n",
    "print(sum(total_times), total_times)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d1ea868-295d-410d-bdb4-40ae6f518aa5",
   "metadata": {},
   "source": [
    "* The times reported here do not include the cost of storing results to disk. How do they compare to the time it actually took to run the simulation? Is storing to disk a significant overhead for simulations of toy models?\n",
    "* The RETIS simulation includes replica exchange moves and path reversals, as well as the shooting moves used by both approaches. Do replica exchange and path reversal contribute significantly to the total simulation time?"
   ]
  }
 ],
 "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
}