simstore_and_cli/1_simstore_setup.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Using SimStore and the OpenPathSampling CLI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "This tutorial covers both the new storage subsystem, SimStore, and the OpenPathSampling Command Line Interface (CLI). It also shows how to use test systems from OpenMMTools as toy examples. This can be particularly useful for method development, as the OpenMMTools test systems add a step up in complexity from the OPS internal toy engine.\n",
    "\n",
    "The OpenPathSampling Command Line Interface (CLI) makes it easier to run your OPS simulations, especially in cluster environments. The basic approach is to first create the simulation objects, including initial conditions, and to save those to a file. Think of this file as a database of simulation setup information that you can later load with the CLI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "import openpathsampling as paths\n",
    "import openmmtools\n",
    "from openmm import unit\n",
    "import numpy as np\n",
    "import mdtraj as md"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# to use SimStore, we need to monkey patch and import specific CVs, Storage\n",
    "from openpathsampling.experimental.storage import monkey_patch_all, Storage\n",
    "from openpathsampling.experimental.storage.collective_variables import MDTrajFunctionCV\n",
    "\n",
    "paths = monkey_patch_all(paths)\n",
    "paths.InterfaceSet.simstore = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## TIS Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Here we'll set up a double-well dimer in a bath of WCA particles. This is a very common toy system for rare events, and is included in OpenMMTools.\n",
    "\n",
    "All particles in this system have a WCA nonbonded interaction, but two of them have a quartic double well \"bond\", which allows them to either be in a \"condensed\" (short) state or an \"extended\" (long) state. The quartic potential is defined by:\n",
    "\n",
    "$$\n",
    "V_{dw}(r) = h \\left(1 - \\left(\\frac{r - r_0 - w}{w}\\right)^2\\right)^2\n",
    "$$\n",
    "\n",
    "where $r$ is the distance between the particles, and the parameters $r_0$, $w$, and $h$ define the minima and maximum of the quartic potential such that the two wells have minima at $(r_0, 0)$ and $(r_0 + 2w, 0)$, and the barrier has a maximum at $(r_0 + w, h)$. The potential (with the interface locations shown in grey) looks like this:\n",
    "\n",
    "![Double well PES](double_well.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Create \"dimensionless\" units"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "# energy in OpenMM is energy/mol, so k in kT is k_B*N_A\n",
    "kB = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# argon-like WCA fluid parameters\n",
    "# these are units of mass, energy, and distance\n",
    "mass = 39.9 * unit.dalton\n",
    "epsilon = 120. * unit.kelvin * kB\n",
    "sigma = 3.4 * unit.angstroms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# tau is the \"unitless\" time unit\n",
    "tau = np.sqrt(sigma**2 * mass / epsilon)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# this is so we can use MDTraj more easily\n",
    "sigma_nm = sigma.value_in_unit(unit.nanometer)\n",
    "print(sigma_nm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Set some parameters we'll use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "temperature = 0.824 / kB * epsilon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "# double-well parameters based on van Erp et al. JCP 2003\n",
    "h = 6.0 * kB * temperature\n",
    "r0 = 2.**(1. / 6.) * sigma\n",
    "w = 0.25 * sigma"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Engine Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "# 1. Set up the OpenMM side\n",
    "\n",
    "# these are all the default parameters, but we'll make it explicit\n",
    "testsystem = openmmtools.testsystems.DoubleWellDimer_WCAFluid(\n",
    "    ndimers=1, \n",
    "    nparticles=216,\n",
    "    density=0.96,\n",
    "    mass=mass,\n",
    "    epsilon=epsilon,\n",
    "    sigma=sigma,\n",
    "    h=h,\n",
    "    r0=r0,\n",
    "    w=w\n",
    ")\n",
    "\n",
    "integrator = openmmtools.integrators.VVVRIntegrator(\n",
    "    timestep=0.001 * tau,\n",
    "    collision_rate=2.0 / tau,\n",
    "    temperature=temperature\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "# 2. Use the OPS wrappers for OpenMM\n",
    "\n",
    "topology = paths.engines.MDTrajTopology(testsystem.mdtraj_topology)\n",
    "engine = paths.engines.openmm.Engine(\n",
    "    topology=topology,\n",
    "    system=testsystem.system,\n",
    "    integrator=integrator,\n",
    "    options={'n_frames_max': 2000,\n",
    "             'n_steps_per_frame': 10}\n",
    ").named('engine')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "# 3. Get an initial snapshot\n",
    "\n",
    "# use OpenMM simulation object to set information\n",
    "engine.simulation.context.setPositions(testsystem.positions)\n",
    "engine.simulation.minimizeEnergy()\n",
    "snap = engine.current_snapshot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Define CVs, stable states, and interfaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "# OLD:\n",
    "#cv = paths.MDTrajFunctionCV(\"r\", md.compute_distances, topology=topology,\n",
    "#                            atom_pairs=[[0,1]])\n",
    "cv = MDTrajFunctionCV(\n",
    "    md.compute_distances,\n",
    "    topology=topology,\n",
    "    atom_pairs=[[0,1]]\n",
    ").named(\"r\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "# stable states\n",
    "condensed = paths.CVDefinedVolume(\n",
    "    cv,\n",
    "    lambda_min=float(\"-inf\"),\n",
    "    lambda_max=1.20 * sigma_nm\n",
    ").named(\"condensed\")\n",
    "\n",
    "extended = paths.CVDefinedVolume(\n",
    "    cv,\n",
    "    lambda_min=1.54 * sigma_nm,\n",
    "    lambda_max=float(\"inf\")\n",
    ").named(\"extended\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "# TIS interfaces\n",
    "interfaces = paths.VolumeInterfaceSet(\n",
    "    cv,\n",
    "    minvals=float(\"-inf\"),\n",
    "    maxvals=np.array([1.20, 1.26, 1.32]) * sigma_nm\n",
    ").named(\"interfaces\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Create network and move scheme"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "network = paths.MISTISNetwork([(condensed, interfaces, extended)]).named(\"tis\")\n",
    "scheme = paths.DefaultScheme(network, engine).named(\"retis_scheme\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Save everything"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "# OLD\n",
    "#storage = paths.Storage(\"setup.nc\", mode='w')\n",
    "#storage.save(paths.Trajectory([snap]))  # save as a trajectory for templating\n",
    "\n",
    "storage = Storage(\"setup.db\", mode='w')\n",
    "storage.save(snap)  # now because we need initial conditions, not a template!\n",
    "storage.save(scheme)  # scheme contains all the simulation info!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# we may find tau and sigma_nm useful in analysis, so we store them, too\n",
    "storage.tags['tau'] = tau\n",
    "storage.tags['sigma_nm'] = sigma_nm\n",
    "\n",
    "storage.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## High-Temperature MD Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "To obtain an initial trajectory (which we will then need to equilibrate), we will use high temperature dynamics for this example. High temperature dynamics is a reasonable way to get an initial trajectory, but depending on your system and what information you have about it, there may be better ways.\n",
    "\n",
    "We'll create the high-temperature engine and save it by appending to the `setup.db` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "hi_temp = openmmtools.integrators.VVVRIntegrator(\n",
    "    timestep=0.001 * tau,\n",
    "    collision_rate=2.0 / tau,\n",
    "    temperature=2 * temperature\n",
    ")\n",
    "\n",
    "hi_temp_engine = paths.engines.openmm.Engine(\n",
    "    topology=topology,\n",
    "    system=testsystem.system,\n",
    "    integrator=hi_temp,\n",
    "    options={'n_frames_max': 10000,\n",
    "             'n_steps_per_frame': 10}\n",
    ").named('hi_temp')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "storage = Storage(\"setup.db\", mode='a')\n",
    "storage.save(hi_temp_engine)\n",
    "storage.close()"
   ]
  }
 ],
 "metadata": {
  "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.9"
  },
  "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
}