docs/BLUES_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\"nbsphinx-toctree\": {\n",
    "  \"maxdepth\": 2\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to BLUES\n",
    "\n",
    "In this Jupyter Notebook, we will cover the following topics:\n",
    "\n",
    "- YAML configuration\n",
    "- Setting up a system for BLUES\n",
    "- Advanced options (HMR, restraints, freezing)\n",
    "- Configuring reporters\n",
    "- Running a BLUES simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Background\n",
    "\n",
    "## Coupling MD simulations with random NCMC moves for enhanced sampling of ligand binding modes via BLUES\n",
    "  [**BLUES**](https://pubs.acs.org/doi/abs/10.1021/acs.jpcb.7b11820) is an approach that combines molecular dynamics ([MD](http://www.ch.embnet.org/MD_tutorial/pages/MD.Part1.html)) simulations and the Non-equilibrium Candidate Monte Carlo ([NCMC](http://www.pnas.org/content/108/45/E1009)) framework to enhance ligand binding mode sampling [(Github)](https://github.com/MobleyLab/blues) \n",
    "  \n",
    "During a MD simulation, BLUES will perform a random rotation of the bound ligand and then allow the system to relax through [_alchemically_](http://www.alchemistry.org/wiki/Free_Energy_Fundamentals#Why_the_name_.22Alchemical.22.3F) scaling off/on the ligand-receptor interactions. BLUES enables us to sample alternative ligand binding modes, that would normally take _very_ long simulations to capture in a traditional MD simulation because of the large gap in [timescales](https://www.ncbi.nlm.nih.gov/pubmed/20934381) between atomistic motions and biological motions.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# YAML Configuration\n",
    "\n",
    "BLUES can be configured in either pure python through dictionaries of the appropriate parameters or using a YAML file (which is converted to a dict under the hood). Below we will walk through the keywords for configuring BLUES. An example YAML configuration file can be found in `rotmove_cuda.yaml`. \n",
    "\n",
    "Note: Code blocks in this notebook denoted with ------ indicate a section from a YAML configuration file."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Input/Output\n",
    "\n",
    "```\n",
    "------\n",
    "output_dir: .\n",
    "outfname: t4-toluene\n",
    "logger_level: info #critical, error, warning, info, debug\n",
    "-----\n",
    "```\n",
    "\n",
    "### Output files\n",
    "Specify the directory you want all the simulation output files to be saved to with **`output_dir`**. By default, BLUES will save them in the current directory that you're running BLUES in. The parameter **`outfname`** will be used for the filename prefix for all output files (e.g. `t4-toluene.nc`, `t4-toleune.log`). The level of verbosity can be controlled by **`logger_level`**. The default **`logger_level`** is set to `info` and valid choices are `critical`, `error`, `warning`, `info` or `debug`.\n",
    "\n",
    "### Input files to generate the structure of your system\n",
    "\n",
    "- Input a Parameter/topology file and a Coordinate file, which will be used to generate the ParmEd Structure. \n",
    "\n",
    "- The ParmEd Structure is a chemical structure composed of atoms, bonds, angles, torsions, and other topological features. \n",
    "\n",
    "To see a full list of supported file formats: https://parmed.github.io/ParmEd/html/readwrite.html\n",
    "\n",
    "```\n",
    "----------------\n",
    "structure:\n",
    "  filename: tests/data/eqToluene.prmtop\n",
    "  xyz: tests/data/eqToluene.inpcrd\n",
    "  restart: t4-toluene_2.rst7\n",
    "----------------\n",
    "```\n",
    "\n",
    "BLUES simulations all begin from a `parmed.Structure`. Keywords nested under **`structure`** are for generating the `parmed.Structure` by calling [parmed.load_file()](https://parmed.github.io/ParmEd/html/readwrite.html#reading-files-with-load-file) under the hood. The **`filename`** keyword should point to the file containing the parameters for your system. For example, if coming from AMBER you would specify the `.prmtop` file. The **`xyz`** keyword is intended for specifying the coordinates of your system (e.g. `.inpcrd` or `.pdb`).\n",
    "\n",
    "### Restart files\n",
    "BLUES supports \"soft\" restarts of simulations from AMBER restart files **`.rst7`** via [parmed.amber.Rst7](https://parmed.github.io/ParmEd/html/amberobj/parmed.amber.Rst7.html?highlight=rst7#parmed-amber-rst7). \"Soft\" restart implies the simulation will begin from the saved positions, velocities, and box vectors but does not store any internal data such as the states of random number generators. It should be noted that velocities are re-initialized at every BLUES iteration, so storing these is not so important."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## System configuration \n",
    "\n",
    "From the YAML file, there is a section dedicated for generating an `openmm.System` from a `parmed.Structure`. For definitions of **`system`** keywords and valid options see [parmed.Structure.createSystem()]( https://parmed.github.io/ParmEd/html/structobj/parmed.structure.Structure.html#parmed.structure.Structure.createSystem)\n",
    "\n",
    "Below we provide an example for generating a system in a cubic box with explicit solvent.\n",
    "\n",
    "```\n",
    "--------\n",
    "system:\n",
    "  nonbondedMethod: PME\n",
    "  nonbondedCutoff: 10 * angstroms\n",
    "  constraints: HBonds\n",
    "  rigidWater: True\n",
    "  removeCMMotion: True\n",
    "  ewaldErrorTolerance: 0.005\n",
    "  flexibleConstraints: True\n",
    "  splitDihedrals: False\n",
    "---------\n",
    "```\n",
    "\n",
    "### (Optional) Hydrogen mass repartitioning\n",
    "BLUES has the option to use the hydrogen mass repartitioning scheme [HMR](https://pubs.acs.org/doi/abs/10.1021/ct5010406) to allow use of longer time steps in the simulation. Simply provide the keyword **`hydrogenMass`** in the YAML file like below:\n",
    "\n",
    "```\n",
    "--------\n",
    "system:\n",
    "  nonbondedMethod: PME\n",
    "  nonbondedCutoff: 10 * angstroms\n",
    "  constraints: HBonds\n",
    "  rigidWater: True\n",
    "  removeCMMotion: True\n",
    "  ewaldErrorTolerance: 0.005\n",
    "  flexibleConstraints: True\n",
    "  splitDihedrals: False\n",
    "  hydrogenMass: 3.024 * daltons\n",
    "---------\n",
    "```\n",
    "\n",
    "If using HMR, you can set the timestep (**`dt`**) for the simulation to 4fs and **`constraints`** should be set to either **`HBonds`** or **`AllBonds`**. \n",
    "\n",
    "\n",
    "### (Optional) Alchemical system configuration \n",
    "\n",
    "Nested under the system parameters, you can modify parameters for the alchemical system. Below are the default settings and are not required to be specified in the YAML configuration. **Modifications to these parameters are for advanced users.** \n",
    "\n",
    "```\n",
    "-------\n",
    "system:\n",
    "  nonbondedMethod: PME\n",
    "  nonbondedCutoff: 10 * angstroms\n",
    "  constraints: HBonds\n",
    "  rigidWater: True\n",
    "  removeCMMotion: True\n",
    "  ewaldErrorTolerance: 0.005\n",
    "  flexibleConstraints: True\n",
    "  splitDihedrals: False\n",
    "  alchemical:\n",
    "    # Sterics\n",
    "    softcore_alpha: 0.5\n",
    "    softcore_a : 1\n",
    "    softcore_b : 1\n",
    "    softcore_c : 6\n",
    "    \n",
    "    # Electrostatics\n",
    "    softcore_beta : 0.0\n",
    "    softcore_d : 1\n",
    "    softcore_e : 1\n",
    "    softcore_f : 2\n",
    "    \n",
    "    annihilate_electrostatics : True \n",
    "    annihilate_sterics : False\n",
    "-------\n",
    "```\n",
    "\n",
    "For further details on alchemical parameters see: http://getyank.org/0.16.2/yamlpages/options.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulation Configuration\n",
    "The keywords for configuring the Simulations for BLUES are explained below:\n",
    "\n",
    "- **`dt`**: timestep\n",
    "- **`nIter`**: number of iterations or proposed moves\n",
    "- **`nstepsMD`**: number of MD steps\n",
    "- **`nstepsNC`**: number NCMC steps\n",
    "\n",
    "The configuration below will run BLUES in NVT with 2fs timesteps for 10 iterations. The MD and NCMC simulation will run 10,000 steps per iteration. \n",
    "```\n",
    "----------\n",
    "simulation:\n",
    "  platform: CUDA \n",
    "  dt: 0.002 * picoseconds\n",
    "  friction: 1 * 1/picoseconds\n",
    "  temperature: 300 * kelvin\n",
    "  nIter: 10\n",
    "  nstepsMD: 10000\n",
    "  nstepsNC: 10000\n",
    "----------\n",
    "```\n",
    "\n",
    "### NPT Simulation\n",
    "\n",
    "To run BLUES in NPT, simply specify a **`pressure`**:\n",
    "```\n",
    "----------\n",
    "simulation:\n",
    "  platform: CUDA \n",
    "  dt: 0.002 * picoseconds\n",
    "  friction: 1 * 1/picoseconds\n",
    "  temperature: 300 * kelvin\n",
    "  nIter: 10\n",
    "  nstepsMD: 10000\n",
    "  nstepsNC: 10000\n",
    "  pressure: 1 * atmospheres\n",
    "----------\n",
    "```\n",
    "\n",
    "### (Optional) Additional relaxation steps in the NCMC simulation\n",
    "Keywords **`nprop`** and **`propLambda`** allow you to add additional relxation steps between a set range in the lambda schedule for the alchemical process in the NCMC simulation. \n",
    "Setting **`propLambda`** to 0.3 will select a lambda range of -/+ 0.3 from the midpoint (0.5), giving [0.2, 0.8]. During the alchemical process, when lambda is between 0.2 to 0.8, **`nprop`** controls the number of additional relaxation steps to add at each lambda step (change in lambda). Additional relaxation steps has been show to increase acceptance proposed NCMC moves.\n",
    "```\n",
    "----------\n",
    "simulation:\n",
    "  platform: CUDA \n",
    "  dt: 0.002 * picoseconds\n",
    "  friction: 1 * 1/picoseconds\n",
    "  temperature: 300 * kelvin\n",
    "  nIter: 10\n",
    "  nstepsMD: 10000\n",
    "  nstepsNC: 10000\n",
    "  pressure: 1 * atmospheres\n",
    "  nprop: 3\n",
    "  propLambda: 0.3\n",
    "----------\n",
    "```\n",
    "\n",
    "\n",
    "### (Optional) Platform Properties\n",
    "If you need to modify platform properties for the simulation, you can set the keyword **`properties`** like below:\n",
    "\n",
    "#### Example: OpenCL in single precision on GPU device 2 \n",
    "*Note: works for running on the GPU on MacBook Pro 2017*\n",
    "\n",
    "```\n",
    "----------\n",
    "simulation:\n",
    "  platform: OpenCL\n",
    "  properties:\n",
    "    OpenCLPrecision: single\n",
    "    OpenCLDeviceIndex: 2\n",
    "  dt: 0.002 * picoseconds\n",
    "  friction: 1 * 1/picoseconds\n",
    "  temperature: 300 * kelvin\n",
    "  nIter: 10\n",
    "  nstepsMD: 10000\n",
    "  nstepsNC: 10000\n",
    "  pressure: 1 * atmospheres\n",
    "----------\n",
    "```\n",
    "\n",
    "#### Example: CUDA in double precision on GPU device 0\n",
    "\n",
    "```\n",
    "----------\n",
    "simulation:\n",
    "  platform: CUDA\n",
    "  properties:\n",
    "    CudaPrecision: double\n",
    "    CudaDeviceIndex: 0\n",
    "  dt: 0.002 * picoseconds\n",
    "  friction: 1 * 1/picoseconds\n",
    "  temperature: 300 * kelvin\n",
    "  nIter: 10\n",
    "  nstepsMD: 10000\n",
    "  nstepsNC: 10000\n",
    "  pressure: 1 * atmospheres\n",
    "----------\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reporter Configuration\n",
    "\n",
    "We provide functionality to configure a recommended set of reporters from the YAML file. These are used to record information for either the MD or NCMC simulation. Below are the keywords for each reporter. Each reporter will require the **`reportInterval`** keyword to specify the frequency to store the simulation data:\n",
    "- **`state`** : State data reporter for OpenMM simulations, but it is a little more generalized. Writes to a `.ene` file.\n",
    "  - For full list of parameters see  [parmed.openmm.reporters.StateDataReporter](https://parmed.github.io/ParmEd/html/api/parmed/parmed.openmm.reporters.html#parmed.openmm.reporters.StateDataReporter)\n",
    "- **`traj_netcdf`** : Customized AMBER NetCDF (`.nc`) format reporter \n",
    "- **`restart`** : Restart AMBER NetCDF (`.rst7`) format reporter\n",
    "- **`progress`** : Write to a file (`.prog`), the progress report of how many steps has been done, how fast the simulation is running, and how much time is left (similar to the mdinfo file in Amber). File is overwritten at each reportInterval.\n",
    "  - For full list of parameters see  [parmed.openmm.reporters.ProgressReporter](https://parmed.github.io/ParmEd/html/api/parmed/parmed.openmm.reporters.html#parmed.openmm.reporters.ProgressReporter)\n",
    "- **`stream`** : Customized version of openmm.app.StateDataReporter. This will instead stream/print the information to the terminal as opposed to writing to a file. \n",
    "  - takes the same parameters as the [openmm.app.StateDataReporter](http://docs.openmm.org/development/api-python/generated/simtk.openmm.app.statedatareporter.StateDataReporter.html#simtk.openmm.app.statedatareporter.StateDataReporter)\n",
    "\n",
    "To attach them to the MD simulation. You nest the reporter keywords under the keyword **`md_reporters`** like below. To attach the reporters to NCMC simulation, use the **`ncmc_reporters`** keyword instead.\n",
    "\n",
    "```\n",
    "------\n",
    "md_reporters:\n",
    "  state:\n",
    "    reportInterval: 250\n",
    "  traj_netcdf:\n",
    "    reportInterval: 250\n",
    "  restart:\n",
    "    reportInterval: 1000\n",
    "  progress:\n",
    "    totalSteps: 10000\n",
    "    reportInterval: 10\n",
    "  stream:\n",
    "    title: md\n",
    "    reportInterval: 250\n",
    "    totalSteps: 10000\n",
    "    step: True\n",
    "    speed: True\n",
    "    progress: True\n",
    "    remainingTime: True\n",
    "----\n",
    "```\n",
    "\n",
    "In the above example, we are using the **`stream`** reporter to print the speeds on the intergrator at regular intervals. This may be a bit redudant with the **`progress`** reporter if are running the job remotely and don't need the information streamed to terminal.\n",
    "\n",
    "### (Optional) Advanced options to the **`traj_netcdf`** reporter\n",
    "The **`traj_netcdf`** reporter can store additional information that may be useful for the NCMC simulation or record at specific frames. In the example below, we will store the first, midpoint (when the move is applied), and last frame of each NCMC iteration, along with the **`alchemicalLambda`** step and the **`protocolWork`**. \n",
    "\n",
    "```\n",
    "----\n",
    "ncmc_reporters:\n",
    "  traj_netcdf:\n",
    "    frame_indices: [1, 0.5, -1]\n",
    "    alchemicalLambda: True\n",
    "    protocolWork: True\n",
    "-----\n",
    "```\n",
    "\n",
    "To access the numerical data stored in the NetCDF file:\n",
    "```\n",
    "from netCDF4 import Dataset\n",
    "\n",
    "f = Dataset(\"t4-toluene-ncmc.nc\")\n",
    "print(f.variables['alchemicalLambda'][:])\n",
    "print(f.variables['protocolWork'][:])\n",
    "\n",
    ">>> [ 0.001  0.5    1.   ]\n",
    ">>> [  0.03706791  30.72696877  25.708498  ]\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Running a BLUES simulation\n",
    "\n",
    "Below we will provide an example for running an NPT BLUES simulation which applies random rotational moves to the toluene ligand in T4-lysozyme from a YAML configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "yaml_cfg = \"\"\"\n",
    "output_dir: .\n",
    "outfname: t4-toluene\n",
    "logger_level: info \n",
    "\n",
    "structure:\n",
    "  filename: ../blues/tests/data/eqToluene.prmtop\n",
    "  xyz: ../blues/tests/data/eqToluene.inpcrd\n",
    "\n",
    "system:\n",
    "  nonbondedMethod: PME\n",
    "  nonbondedCutoff: 10 * angstroms\n",
    "  constraints: HBonds\n",
    "  rigidWater: True\n",
    "  removeCMMotion: True\n",
    "  hydrogenMass: 3.024 * daltons\n",
    "  ewaldErrorTolerance: 0.005\n",
    "  flexibleConstraints: True\n",
    "  splitDihedrals: False\n",
    "\n",
    "freeze:\n",
    "  freeze_center: ':LIG'\n",
    "  freeze_solvent: ':WAT,Cl-'\n",
    "  freeze_distance: 5 * angstroms\n",
    "\n",
    "simulation:\n",
    "  platform: CUDA\n",
    "  properties:\n",
    "    CudaPrecision: single\n",
    "    CudaDeviceIndex: 0\n",
    "  dt: 0.004 * picoseconds\n",
    "  friction: 1 * 1/picoseconds\n",
    "  pressure: 1 * atmospheres\n",
    "  temperature: 300 * kelvin\n",
    "  nIter: 5\n",
    "  nstepsMD: 1000\n",
    "  nstepsNC: 1000\n",
    "\n",
    "md_reporters:\n",
    "  state:\n",
    "    reportInterval: 250\n",
    "  traj_netcdf:\n",
    "    reportInterval: 250\n",
    "  restart:\n",
    "    reportInterval: 1000\n",
    "  stream:\n",
    "    title: md\n",
    "    reportInterval: 250\n",
    "    totalSteps: 5000 # nIter * nstepsMD\n",
    "    step: True\n",
    "    speed: True\n",
    "    progress: True\n",
    "    remainingTime: True\n",
    "    \n",
    "ncmc_reporters:\n",
    "  stream:\n",
    "    title: ncmc\n",
    "    reportInterval: 250\n",
    "    totalSteps: 1000 # Use nstepsNC\n",
    "    step: True\n",
    "    speed: True\n",
    "    progress: True\n",
    "    remainingTime: True\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import the following BLUES modules required for the following steps.\n",
    "\n",
    "- Make sure to specify the type of move that you want to import from __blues.moves__\n",
    "- Available moves can be veiwed in __moves.py__, supported/tested moves include:\n",
    "    - RandomLigandRotationMove\n",
    "    - SideChainMove"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from blues.moves import RandomLigandRotationMove, MoveEngine\n",
    "from blues.simulation import *\n",
    "from blues.settings import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./t4-toluene\n"
     ]
    }
   ],
   "source": [
    "#Read in the YAML file\n",
    "cfg = Settings(yaml_cfg).asDict()\n",
    "#Shortcut to access `parmed.Structure` from dict\n",
    "structure = cfg['Structure']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below is what the resulting configuration dictionary looks like (formatted into JSON for readability)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  \"Logger\": \"<logging.RootLogger object at 0x7f96b1bff940>\",\n",
      "  \"Structure\": \"../blues/tests/data/eqToluene.prmtop\",\n",
      "  \"freeze\": {\n",
      "    \"freeze_center\": \":LIG\",\n",
      "    \"freeze_distance\": \"5.0 A\",\n",
      "    \"freeze_solvent\": \":WAT,Cl-\"\n",
      "  },\n",
      "  \"logger_level\": \"info\",\n",
      "  \"md_reporters\": [\n",
      "    \"<parmed.openmm.reporters.StateDataReporter object at 0x7f966832c128>\",\n",
      "    \"<blues.reporters.NetCDF4Reporter object at 0x7f966574e898>\",\n",
      "    \"<parmed.openmm.reporters.RestartReporter object at 0x7f966574e860>\",\n",
      "    \"<blues.reporters.BLUESStateDataReporter object at 0x7f966574e828>\"\n",
      "  ],\n",
      "  \"ncmc_reporters\": [\n",
      "    \"<blues.reporters.BLUESStateDataReporter object at 0x7f966574e908>\"\n",
      "  ],\n",
      "  \"outfname\": \"./t4-toluene\",\n",
      "  \"output_dir\": \".\",\n",
      "  \"simulation\": {\n",
      "    \"dt\": \"0.004 ps\",\n",
      "    \"friction\": \"1.0 /ps\",\n",
      "    \"md_trajectory_interval\": 250,\n",
      "    \"moveStep\": 500,\n",
      "    \"nIter\": 5,\n",
      "    \"nprop\": 1,\n",
      "    \"nstepsMD\": 1000,\n",
      "    \"nstepsNC\": 1000,\n",
      "    \"outfname\": \"./t4-toluene\",\n",
      "    \"platform\": \"CUDA\",\n",
      "    \"pressure\": \"1.0 atm\",\n",
      "    \"propLambda\": 0.3,\n",
      "    \"propSteps\": 1000,\n",
      "    \"properties\": {\n",
      "      \"CudaDeviceIndex\": 0,\n",
      "      \"CudaPrecision\": \"single\"\n",
      "    },\n",
      "    \"temperature\": \"300.0 K\",\n",
      "    \"verbose\": false\n",
      "  },\n",
      "  \"structure\": {\n",
      "    \"filename\": \"../blues/tests/data/eqToluene.prmtop\",\n",
      "    \"xyz\": \"../blues/tests/data/eqToluene.inpcrd\"\n",
      "  },\n",
      "  \"system\": {\n",
      "    \"constraints\": \"HBonds\",\n",
      "    \"ewaldErrorTolerance\": 0.005,\n",
      "    \"flexibleConstraints\": true,\n",
      "    \"hydrogenMass\": \"3.024 Da\",\n",
      "    \"nonbondedCutoff\": \"10.0 A\",\n",
      "    \"nonbondedMethod\": \"PME\",\n",
      "    \"removeCMMotion\": true,\n",
      "    \"rigidWater\": true,\n",
      "    \"splitDihedrals\": false,\n",
      "    \"verbose\": false\n",
      "  },\n",
      "  \"verbose\": false\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "print(json.dumps(cfg, sort_keys=True, indent=2, skipkeys=True, default=str))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Selecting a move and initialize the move engine\n",
    "\n",
    "Here, we initialize the **`RandomLigandRotationMove`** from the **`blues.moves`** module which proposes random rotations on the toluene ligand. We select the toluene ligand by providing the residue name `LIG` and the `parmed.Structure` to select the atoms from. If we begin BLUES from our YAML configuration, the `parmed.Structure` for our system is generated from the call to `startup()`. We can access it at the top level with `cfg['Structure']`\n",
    "\n",
    "After initialization of the selected move, we pass the move object to the **`MoveEngine`** from the **`blues.engine`** module. The **`MoveEngine`** controls what types of moves will be performed during the NCMC protocol and with a given probability. This will be more useful when we use multiple move types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Initialize the move class and pass it to the engine\n",
    "ligand = RandomLigandRotationMove(structure, 'LIG')\n",
    "ligand_mover = MoveEngine(ligand)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating the Systems for openMM: **`SystemFactory`**\n",
    "Next, we must generate the `openmm.System` from the `parmed.Structure` by calling the **`SystemFactory`** class from the `blues.simulation` module. The class must be initialized by providing 3 required arguments:\n",
    "```\n",
    "structure : parmed.Structure\n",
    "    A chemical structure composed of atoms, bonds, angles, torsions, and\n",
    "    other topological features.\n",
    "atom_indices : list of int\n",
    "    Atom indicies of the move or designated for which the nonbonded forces\n",
    "    (both sterics and electrostatics components) have to be alchemically\n",
    "    modified.\n",
    "config : dict, parameters for generating the `openmm.System for the MD\n",
    "    and NCMC simulation. For complete parameters, see docs for `generateSystem`\n",
    "    and `generateAlchSystem`\n",
    "```\n",
    "\n",
    "Upon initialization, this class will create the system for the MD simulation and the NCMC simulation. They can be accessed through the attributes **`systems.md`** or **`systems.alch`**. Any modifications to either of these Systems should be done within the context of this object. Once the systems are passed into an `openmm.Simulation`, you will not be able to modify the system easily."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO: Adding bonds...\n",
      "INFO: Adding angles...\n",
      "INFO: Adding dihedrals...\n",
      "INFO: Adding Ryckaert-Bellemans torsions...\n",
      "INFO: Adding Urey-Bradleys...\n",
      "INFO: Adding improper torsions...\n",
      "INFO: Adding CMAP torsions...\n",
      "INFO: Adding trigonal angle terms...\n",
      "INFO: Adding out-of-plane bends...\n",
      "INFO: Adding pi-torsions...\n",
      "INFO: Adding stretch-bends...\n",
      "INFO: Adding torsion-torsions...\n",
      "INFO: Adding Nonbonded force...\n"
     ]
    }
   ],
   "source": [
    "systems = SystemFactory(structure, ligand.atom_indices, cfg['system'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (Optional) Applying restraints or freezing atoms\n",
    "The **`SystemFactory`** class also provides functionality for restraining or freezing the atoms. **Use extreme caution when freezing/restraining atoms. You should consider if freezing/restraining should be applied to BOTH the MD and alchemical system.**\n",
    "\n",
    "\n",
    "Selections for either restraining or freezing atoms in your system use the [Amber mask syntax](http://parmed.github.io/ParmEd/html/amber.html#amber-mask-syntax).\n",
    "\n",
    "#### Positional restraints: **`SystemFactory.restrain_positions()`**\n",
    "To apply positional restraints, you can call **`SystemFactory.restrain_positions()`**.\n",
    "You can specify the parameters/selection for applying positional restraints in the YAML file.\n",
    "\n",
    "```\n",
    "------\n",
    "restraints:\n",
    "    selection: '@CA,C,N' \n",
    "    weight: 5 \n",
    "------\n",
    "```\n",
    "\n",
    "restrain keywords:\n",
    "- **`selection`**: Specify what to apply positional restraints to using Amber mask syntax. Default = '@CA,C,N'\n",
    "- **`weight`**: Restraint weight for xyz atom restraints in kcal/(mol A^2). Default = 5\n",
    "\n",
    "\n",
    "From the YAML example above, we would be applying positional restraints to the backbone atoms of the protein. If applying restraints, you most likely will want to apply it to BOTH the MD and alchemical systems like below:\n",
    "\n",
    "```\n",
    "systems.md = SystemFactory.restrain_positions(structure, systems.md, **cfg['restraints'])\n",
    "systems.alch = SystemFactory.restrain_positions(structure, systems.alch, **cfg['restraints'])\n",
    "```\n",
    "\n",
    "#### Freezing selected atoms: **`SystemFactory.freeze_atoms()`**\n",
    "To freeze a selection, call **`SystemFactory.freeze_atoms()`**. Atoms that have a mass of zero will be ignored by the integrator and will not change positions during the simulation, effectively they are frozen.\n",
    "\n",
    "To freeze atoms using a given selection string. Use the keyword **`freeze_selection`**.\n",
    "\n",
    "```\n",
    "-----\n",
    "freeze:\n",
    "  freeze_selection: ':LIG'\n",
    "------\n",
    "```\n",
    "\n",
    "From the YAML example above, we would be freezing only the atoms belonging to the residue `LIG`. Although freezing the ligand in this example wouldn't be very useful. It would be applied like\n",
    "\n",
    "```\n",
    "systems.md = SystemFactory.freeze_atoms(structure, systems.md, **cfg['freeze'])\n",
    "```\n",
    "\n",
    "\n",
    "#### Freezing atoms around a selection: **`SystemFactory.freeze_radius()`**\n",
    "Alternatively, you can choose to freeze atoms around a given selection. To do so, call **`SystemFactory.freeze_radius()`**. For example, you may want to freeze atoms that are 5 angstroms away from the ligand and include the solvent.\n",
    "\n",
    "```\n",
    "-----\n",
    "freeze:\n",
    "  freeze_center: ':LIG' \n",
    "  freeze_solvent: ':WAT,Cl-' \n",
    "  freeze_distance: 5 * angstroms\n",
    "------\n",
    "```\n",
    "\n",
    "- **`freeze_center`**: Specifies the center of the object for freezing, masses will be zeroed. Default = ':LIG'\n",
    "- **`freeze_solvent`**: select which solvent atoms should have their masses zeroed. Default = ':HOH,NA,CL' \n",
    "- **`freeze_distance`**: Distance ( in angstroms) to select atoms for retaining their masses. Atoms outside the set distance will have their masses set to 0.0. Default = 5.0\n",
    "\n",
    "We often utilize this type of freezing to speed up the alchemical process during the NCMC simulation while leaving them completely free in the MD simulation for proper relaxation. \n",
    "```\n",
    "systems.alch = SystemFactory.freeze_radius(structure, systems.alch, **cfg['freeze'])\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook example, our YAML config indicates we will be freezing around the ligand (keyword: **`freeze_center`**). So we will call the **`freeze_radius`** function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO: Freezing 22065 atoms 5.0 Angstroms from ':LIG' on <simtk.openmm.openmm.System; proxy of <Swig Object of type 'OpenMM::System *' at 0x7f96aa74fab0> >\n"
     ]
    }
   ],
   "source": [
    "systems.alch = SystemFactory.freeze_radius(structure, systems.alch, **cfg['freeze'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating the OpenMM Simulations: **`SimulationFactory`**\n",
    "\n",
    "Now that we have generated our `openmm.System` for the MDand frozen the solvent around the ligand. We are now ready to create the set of simulations for running BLUES. We do this by calling the **`SimulationFactory`** class . The expected parameters are:\n",
    "\n",
    "```\n",
    "systems : blues.simulation.SystemFactory object\n",
    "    The object containing the MD and alchemical openmm.Systems\n",
    "move_engine : blues.engine.MoveEngine object\n",
    "    MoveProposal object which contains the dict of moves performed\n",
    "    in the NCMC simulation.\n",
    "config : dict of parameters for the simulation (i.e timestep, temperature, etc.)\n",
    "md_reporters : list of Reporter objects for the MD openmm.Simulation\n",
    "ncmc_reporters : list of Reporter objects for the NCMC openmm.Simulation\n",
    "```\n",
    "\n",
    "If you wish to use your own openmm reporters, simply pass them into the arguments as a *list* of Reporter objects. Since we have configured our reporters from the YAML file, we can pass them into the arguments **`md_reporters`** and **`ncmc_reporters`**. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<parmed.openmm.reporters.StateDataReporter at 0x7f966832c128>,\n",
       " <blues.reporters.NetCDF4Reporter at 0x7f966574e898>,\n",
       " <parmed.openmm.reporters.RestartReporter at 0x7f966574e860>,\n",
       " <blues.reporters.BLUESStateDataReporter at 0x7f966574e828>]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# List of MD reporters\n",
    "cfg['md_reporters']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<blues.reporters.BLUESStateDataReporter at 0x7f966574e908>]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# List of NCMC reporters\n",
    "cfg['ncmc_reporters']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO: Adding MonteCarloBarostat with 1.0 atm. MD simulation will be 300.0 K NPT.\n",
      "WARNING: NCMC simulation will NOT have pressure control. NCMC will use pressure from last MD state.\n",
      "INFO: OpenMM(7.1.1.dev-c1a64aa) simulation generated for CUDA platform\n",
      "system = Linux \n",
      "node = titanpascal \n",
      "release = 4.13.0-41-generic \n",
      "version = #46~16.04.1-Ubuntu SMP Thu May 3 10:06:43 UTC 2018 \n",
      "machine = x86_64 \n",
      "processor = x86_64 \n",
      "DeviceIndex = 0 \n",
      "DeviceName = TITAN Xp \n",
      "UseBlockingSync = true \n",
      "Precision = single \n",
      "UseCpuPme = false \n",
      "CudaCompiler = /usr/local/cuda-8.0/bin/nvcc \n",
      "TempDirectory = /tmp \n",
      "CudaHostCompiler =  \n",
      "DisablePmeStream = false \n",
      "DeterministicForces = false \n",
      "\n"
     ]
    }
   ],
   "source": [
    "simulations = SimulationFactory(systems, ligand_mover, cfg['simulation'], \n",
    "                                cfg['md_reporters'], cfg['ncmc_reporters'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Accessing the MD or NCMC simulation\n",
    "If you would like to access the MD or NCMC simulation. You can access them as attributes to the **`SimulationFactory`** class with **`simulations.md`** or **`simulations.ncmc`**. This will allow you to do things like \n",
    "energy minimize the system or run a few steps of regular dynamics before running the hybrid (MD+NCMC) BLUES approach. The NCMC simulation will automatically be synced to the state of the MD simulation when running the BLUES simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pre-Minimized energy = -69057.34671058532 kcal/mol\n",
      "Minimized energy = -87007.38938198877 kcal/mol\n"
     ]
    }
   ],
   "source": [
    "# Energy minimization\n",
    "state = simulations.md.context.getState(getPositions=True, getEnergy=True)\n",
    "print('Pre-Minimized energy = {}'.format(state.getPotentialEnergy().in_units_of(unit.kilocalorie_per_mole)))\n",
    "\n",
    "simulations.md.minimizeEnergy(maxIterations=0)\n",
    "state = simulations.md.context.getState(getPositions=True, getEnergy=True)\n",
    "print('Minimized energy = {}'.format(state.getPotentialEnergy().in_units_of(unit.kilocalorie_per_mole)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#\"Progress (%)\"\t\"Step\"\t\"Speed (ns/day)\"\t\"Time Remaining\"\n",
      "md: 5.0%\t250\t0\t--\n",
      "md: 10.0%\t500\t271\t0:05\n"
     ]
    }
   ],
   "source": [
    "# Running only the MD simulation\n",
    "simulations.md.step(500)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the BLUES Simulation\n",
    "To run the full BLUES simulation, where we apply NCMC moves and follow-up with the MD simulation, we simply pass the **`SimulationFactory`** object to the **`Simulation`** class and call the **`run()`** function which takes **`nIter`, `nstepsNC`, `nstepsMD` ** as arguments (or we can pass it the simulation configuration from the YAML on initialization of the class)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO: Total BLUES Simulation Time = 40.0 ps (8.0 ps/Iter)\n",
      "Total Force Evaluations = 10000 \n",
      "Total NCMC time = 20.0 ps (4.0 ps/iter)\n",
      "Total MD time = 20.0 ps (4.0 ps/iter)\n",
      "Trajectory Interval = 2.0 ps/frame (4.0 frames/iter)\n",
      "INFO: Running 5 BLUES iterations...\n",
      "INFO: BLUES Iteration: 0\n",
      "INFO: Advancing 1000 NCMC switching steps...\n",
      "#\"Progress (%)\"\t\"Step\"\t\"Speed (ns/day)\"\t\"Time Remaining\"\n",
      "ncmc: 25.0%\t250\t0\t--\n",
      "ncmc: 50.0%\t500\t110\t0:01\n",
      "Performing RandomLigandRotationMove...\n",
      "ncmc: 75.0%\t750\t93.3\t0:00\n",
      "ncmc: 100.0%\t1000\t93\t0:00\n",
      "NCMC MOVE REJECTED: work_ncmc -17.61007128922719 < -3.476065340494845\n",
      "Advancing 1000 MD steps...\n",
      "md: 15.0%\t750\t31.9\t0:46\n",
      "md: 20.0%\t1000\t43.1\t0:32\n",
      "md: 25.0%\t1250\t54.8\t0:23\n",
      "md: 30.0%\t1500\t65.5\t0:18\n",
      "BLUES Iteration: 1\n",
      "Advancing 1000 NCMC switching steps...\n",
      "ncmc: 25.0%\t250\t57.2\t--\n",
      "ncmc: 50.0%\t500\t62.6\t0:13\n",
      "Performing RandomLigandRotationMove...\n",
      "ncmc: 75.0%\t750\t65.4\t0:03\n",
      "ncmc: 100.0%\t1000\t69.1\t0:00\n",
      "NCMC MOVE REJECTED: work_ncmc -28.930984747001787 < -2.602259467723195\n",
      "Advancing 1000 MD steps...\n",
      "md: 35.0%\t1750\t45.5\t0:24\n",
      "md: 40.0%\t2000\t50.5\t0:20\n",
      "md: 45.0%\t2250\t56.3\t0:16\n",
      "md: 50.0%\t2500\t61.9\t0:13\n",
      "BLUES Iteration: 2\n",
      "Advancing 1000 NCMC switching steps...\n",
      "ncmc: 25.0%\t250\t57.9\t--\n",
      "ncmc: 50.0%\t500\t60.1\t0:25\n",
      "Performing RandomLigandRotationMove...\n",
      "ncmc: 75.0%\t750\t62.4\t0:06\n",
      "ncmc: 100.0%\t1000\t64.6\t0:00\n",
      "NCMC MOVE REJECTED: work_ncmc -9.178847171172448 < -0.3081492900307722\n",
      "Advancing 1000 MD steps...\n",
      "md: 55.0%\t2750\t49.7\t0:15\n",
      "md: 60.0%\t3000\t53.1\t0:13\n",
      "md: 65.0%\t3250\t56.7\t0:10\n",
      "md: 70.0%\t3500\t60.4\t0:08\n",
      "BLUES Iteration: 3\n",
      "Advancing 1000 NCMC switching steps...\n",
      "ncmc: 25.0%\t250\t58.4\t--\n",
      "ncmc: 50.0%\t500\t60.4\t0:37\n",
      "Performing RandomLigandRotationMove...\n",
      "ncmc: 75.0%\t750\t61.6\t0:09\n",
      "ncmc: 100.0%\t1000\t63.6\t0:00\n",
      "NCMC MOVE REJECTED: work_ncmc -6.022089748510081 < -0.28447339331708726\n",
      "Advancing 1000 MD steps...\n",
      "md: 75.0%\t3750\t52.6\t0:08\n",
      "md: 80.0%\t4000\t55\t0:06\n",
      "md: 85.0%\t4250\t58\t0:04\n",
      "md: 90.0%\t4500\t60.8\t0:02\n",
      "BLUES Iteration: 4\n",
      "Advancing 1000 NCMC switching steps...\n",
      "ncmc: 25.0%\t250\t58.6\t--\n",
      "ncmc: 50.0%\t500\t59.8\t0:49\n",
      "Performing RandomLigandRotationMove...\n",
      "ncmc: 75.0%\t750\t60.7\t0:12\n",
      "ncmc: 100.0%\t1000\t61.9\t0:00\n",
      "NCMC MOVE REJECTED: work_ncmc -63.43688730563919 < -0.5417610940298393\n",
      "Advancing 1000 MD steps...\n",
      "md: 95.0%\t4750\t53.3\t0:01\n",
      "md: 100.0%\t5000\t55.2\t0:00\n",
      "md: 105.0%\t5250\t57.5\t23:59:59\n",
      "md: 110.0%\t5500\t59.8\t23:59:58\n",
      "Acceptance Ratio: 0.0\n",
      "nIter: 5 \n"
     ]
    }
   ],
   "source": [
    "blues = BLUESSimulation(simulations, cfg['simulation'])\n",
    "blues.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}