section_2/custom_forces.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom Forces - Umbrella Sampling\n",
    "\n",
    "You can run this notebook in your browser: \n",
    "\n",
    "[![Open On Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/openmm/openmm_workshops/blob/main/section_2/custom_forces.ipynb)\n",
    "\n",
    "\n",
    "This notebook introduces the custom forces feature available in OpenMM. First, we apply them to a Lennard-Jones toy system. Then, we cover an in-depth example of implementing Umbrella Sampling in OpenMM using custom forces.\n",
    "\n",
    "## Table of Contents\n",
    "- Setting up environment\n",
    "- Custom Forces\n",
    "    - Creating a toy system\n",
    "    - Custom non-bonded force\n",
    "    - Custom external force\n",
    "- Umbrella Sampling\n",
    "    - System\n",
    "    - Setting up the windows with SMD\n",
    "    - Running the windows\n",
    "    - Analysis\n",
    "- Extra exercises\n",
    "- References\n",
    "- Solutions"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the Environment\n",
    "<a id=\"env\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute this cell to install mamba in the Colab environment and then install openmm\n",
    "\n",
    "if 'google.colab' in str(get_ipython()):\n",
    "    print('Running on colab')\n",
    "    !pip install -q condacolab\n",
    "    import condacolab\n",
    "    condacolab.install_mambaforge()\n",
    "else:\n",
    "    print('Not running on colab.')\n",
    "    print('Make sure you create and activate a new conda environment!')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** During this step on Colab the kernel will be restarted. This will produce the error message:\n",
    "\"Your session crashed for an unknown reason. \" This is normal and you can safely ignore it."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** Installing the packages will take several minutes!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install OpenMM\n",
    "!mamba install -y -c conda-forge openmm  "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Custom Forces\n",
    "<a id=\"customforces\"></a>\n",
    "\n",
    "OpenMM provides several [Custom Force](http://docs.openmm.org/latest/userguide/theory/03_custom_forces.html) classes. These classes provide great flexibility by allowing the user to specify arbitrary algebraic expressions. In this example, we will apply a custom non-bonded force and a custom external force to a toy system consisting of Lennard-Jones particles"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Toy System\n",
    "<a id=\"toysystem\"></a>\n",
    "\n",
    "We will use the OpenMM API to create a LJ system. Note that there are a variety of different test systems available in the OpenMM tools package: https://github.com/choderalab/openmmtools. For our test system we will generate of 3D grid of 1000 Lennard Jones particles using parameters for Argon."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmm as mm\n",
    "import openmm.app as app\n",
    "import openmm.unit as unit\n",
    "import numpy as np\n",
    "import itertools\n",
    "from sys import stdout\n",
    "\n",
    "# Number of atoms\n",
    "N=1000\n",
    "\n",
    "# Box length\n",
    "L=5.0*unit.nanometer\n",
    "\n",
    "# Lennard-Jones parameters\n",
    "# Use values for Argon from openmm-tools test systems.\n",
    "sigma = 3.4 * unit.angstrom\n",
    "epsilon = 0.238 * unit.kilocalories_per_mole\n",
    "\n",
    "# Make an empty topology\n",
    "topology = app.Topology()\n",
    "\n",
    "# Add a single chain\n",
    "chain = topology.addChain()\n",
    "\n",
    "# Add atoms to the topology\n",
    "for i in range(N):\n",
    "    residue = topology.addResidue(name='Ar', chain=chain)\n",
    "    atom = topology.addAtom(name='Ar', element=app.element.get_by_symbol('Ar'), residue=residue)\n",
    "\n",
    "topology.setPeriodicBoxVectors(np.eye(3)*L)\n",
    "\n",
    "# Check the topology\n",
    "print(topology)\n",
    "\n",
    "# Generate positions on a grid\n",
    "M=int(np.cbrt(N))\n",
    "dx = L._value/M\n",
    "print(dx)\n",
    "x = np.linspace(0,L._value-dx,int(np.cbrt(N)))\n",
    "positions = np.array(list(itertools.product(x,x,x)))*unit.nanometers\n",
    "\n",
    "# Create the system and add the particles to it\n",
    "system = mm.System()\n",
    "system.setDefaultPeriodicBoxVectors(*topology.getPeriodicBoxVectors())\n",
    "for atom in topology.atoms():\n",
    "    system.addParticle(atom.element.mass)\n",
    "\n",
    "# Output the topology and positions to a PDB file\n",
    "with open('topology_lj.pdb', 'w') as f:\n",
    "    app.PDBFile.writeFile(topology, positions, f)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The initial topology looks like this:\n",
    "\n",
    "![topology_lj](images/topology_lj.png)\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "⚠️ <b>Note that this image was generated using Ovito, https://www.ovito.org/, an alternative molecular visulization software.</b>\n",
    "</div>\n",
    "\n",
    "We now have a `Topology` and a `System` but no forces defined yet."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom Non-Bonded Force\n",
    "<a id=\"customnonbonded\"></a>\n",
    "\n",
    "We will create a [`CustomNonBondedForce`](http://docs.openmm.org/latest/api-python/generated/openmm.openmm.CustomNonbondedForce.html) that, for simplicity, uses a Lennard-Jones interaction, i.e.:\n",
    "\n",
    "$V(r) = 4 \\epsilon \\left[ \\left( \\frac{\\sigma}{r} \\right)^{12} - \\left( \\frac{\\sigma}{r} \\right)^{6} \\right]$\n",
    "\n",
    "\n",
    "with the Lorentz-Berthelot combination rules:\n",
    "\n",
    "$\\epsilon_{12} = \\sqrt{\\epsilon_1 \\epsilon_2}$\n",
    "\n",
    "$\\sigma_{12} = \\frac{\\sigma_1 + \\sigma_2}{2}$\n",
    "\n",
    "When you create a custom force, you must write the potential energy expression as a string. The `CustomForce` class then computes the gradient of this energy expression to determine the forces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_lj = '4*epsilon*((sigma/r)^12-(sigma/r)^6);'\n",
    "custom_lj += 'sigma=0.5*(sigma1+sigma2);'\n",
    "custom_lj += 'epsilon=sqrt(epsilon1*epsilon2)'\n",
    "custom_nb_force = mm.CustomNonbondedForce(custom_lj)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The potential energy depends on two parameters, `sigma` and `epsilon`, and we have defined the combination rules to calculate them. \n",
    "We need to tell the custom force that these are `perParticleParameters`, meaning that each particle will have its own values for these parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_nb_force.addPerParticleParameter('sigma')\n",
    "custom_nb_force.addPerParticleParameter('epsilon')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then need to set the values of `sigma` and `epsilon` for each particle. This is done using the `addParticle` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(N):\n",
    "    custom_nb_force.addParticle([sigma, epsilon])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We set the cutoff method settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_nb_force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)\n",
    "custom_nb_force.setCutoffDistance(3.0*sigma)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now add the force to the system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system.addForce(custom_nb_force)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The system is now ready for simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "integrator = mm.LangevinMiddleIntegrator(300*unit.kelvin, 1/unit.picosecond, 0.004*unit.picoseconds)\n",
    "simulation = app.Simulation(topology, system, integrator)\n",
    "simulation.context.setPositions(positions)\n",
    "simulation.context.setVelocitiesToTemperature(300*unit.kelvin)\n",
    "\n",
    "# Add reporters to output data\n",
    "simulation.reporters.append(app.PDBReporter(\"lj_traj.pdb\", 100))\n",
    "simulation.reporters.append(app.StateDataReporter(stdout, 100, step=True,\n",
    "        potentialEnergy=True, temperature=True, speed=True))\n",
    "\n",
    "simulation.step(1000)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compare the speed to that of a standard nonbonded force."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a standard system with a standard nonbonded force\n",
    "standard_system = mm.System()\n",
    "standard_system.setDefaultPeriodicBoxVectors(*topology.getPeriodicBoxVectors())\n",
    "for atom in topology.atoms():\n",
    "    standard_system.addParticle(atom.element.mass)\n",
    "\n",
    "nb_force = mm.NonbondedForce()\n",
    "nb_force.setNonbondedMethod(mm.NonbondedForce.CutoffPeriodic)\n",
    "nb_force.setCutoffDistance(3.0*sigma)\n",
    "charge = 0.0\n",
    "\n",
    "# Add particles to the standard system\n",
    "for i in range(N):\n",
    "    nb_force.addParticle(charge, sigma, epsilon)\n",
    "\n",
    "# Add the force to the system\n",
    "standard_system.addForce(nb_force)\n",
    "\n",
    "# Create a simulation with the standard system\n",
    "integrator = mm.LangevinMiddleIntegrator(300*unit.kelvin, 1/unit.picosecond, 0.004*unit.picoseconds)\n",
    "standard_simulation = app.Simulation(topology, standard_system, integrator)\n",
    "standard_simulation.context.setPositions(positions)\n",
    "simulation.context.setVelocitiesToTemperature(300*unit.kelvin)\n",
    "\n",
    "# Add reporters to output data\n",
    "standard_simulation.reporters.append(app.PDBReporter(\"lj_traj.pdb\", 100))\n",
    "standard_simulation.reporters.append(app.StateDataReporter(stdout, 100, step=True,\n",
    "        potentialEnergy=True, temperature=True, speed=True))\n",
    "\n",
    "standard_simulation.step(1000)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Depending on the hardware, you will experience different speeds. Generally, the standard `NonBondedForce` is expected to be faster, though it is typically not more than twice as fast as an equivalent `CustomNonbondedForce`.\n",
    "\n",
    "The functional form of the custom nonbonded force can be as complex as you need."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom External Force\n",
    "<a id=\"customexternal\"></a>\n",
    "\n",
    "Another very flexible custom force is [`CustomExternalForce`](http://docs.openmm.org/latest/api-python/generated/openmm.openmm.CustomExternalForce.html#openmm.openmm.CustomExternalForce). This class implements an 'external' force on particles. We can create a `CustomExternalForce` that attracts each particle to a target position $x_0$ in the $x$ direction using a harmonic potential. This force depends on a global parameter `k`, which defines the force constant of the potential, and a per-particle parameter `x0`.\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "ℹ️ <b>Exercise 1</b>\n",
    "\n",
    "You will need to add the per particle parameter `x0` to the `CustomExternalForce`.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "external_force = mm.CustomExternalForce('k*periodicdistance(x,0,0,x0,0,0)^2')\n",
    "external_force.addGlobalParameter('k', 10.0)\n",
    "\n",
    "# Add the per particle parameter 'x0' to the force\n",
    "FIXME\n",
    "\n",
    "# Add each particle to the external_force\n",
    "# For this custom force we give the index of each particle to add,\n",
    "# and then a list of the perParticleParameters\n",
    "for i in range(N):\n",
    "    external_force.addParticle(i, [0.0])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have now defined a force that pulls all particles to the $x$ origin. We can add this force to a system and run a simulation.\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "ℹ️ <b>Exercise 2</b>\n",
    "\n",
    "Add the external custom force to the OpenMM system.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "\n",
    "# Create a new system with the same particles\n",
    "new_system = mm.System()\n",
    "new_system.setDefaultPeriodicBoxVectors(*topology.getPeriodicBoxVectors())\n",
    "for atom in topology.atoms():\n",
    "    new_system.addParticle(atom.element.mass)\n",
    "\n",
    "# Add the standard nonbonded force to the new system\n",
    "new_system.addForce(deepcopy(nb_force))\n",
    "\n",
    "# Add the external_force to the system\n",
    "FIXME\n",
    "\n",
    "# Create and run a simulation with the new system\n",
    "integrator = mm.LangevinMiddleIntegrator(300*unit.kelvin, 1/unit.picosecond, 0.004*unit.picoseconds)\n",
    "new_simulation = app.Simulation(topology, new_system, integrator)\n",
    "new_simulation.context.setPositions(positions)\n",
    "simulation.context.setVelocitiesToTemperature(300*unit.kelvin)\n",
    "\n",
    "# Add reporters to output data\n",
    "new_simulation.reporters.append(app.StateDataReporter(stdout, 100, step=True,\n",
    "        potentialEnergy=True, temperature=True, speed=True))\n",
    "new_simulation.reporters.append(app.PDBReporter(\"lj_traj.pdb\", 100))\n",
    "\n",
    "new_simulation.step(2000)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "ℹ️ <b>Exercise 3</b>\n",
    "\n",
    "Visualize the trajectory file `'lj_traj.pdb'` to check if the external force is working correctly. Then, change the expression to `'k*(x-x0)^2'` and rerun the simulation. What happens to the trajectory? Why does this occur?\n",
    "</div>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Umbrella Sampling\n",
    "<a id=\"umbrellasampling\"></a>\n",
    "\n",
    "We will now use the custom forces to implement umbrella sampling in OpenMM.\n",
    " \n",
    "Umbrella sampling is a technique used in molecular dynamics simulations to enhance sampling in systems containing high barriers in the free energy landscape [[1,2]](#references). In standard MD simulations, the system can remain trapped in free energy minima. Umbrella sampling adds extra biasing potentials to drive the system in the direction of a chosen collective variable (CV), enabling the calculation of free energy profiles much more efficiently than with standard MD alone.\n",
    "\n",
    "\n",
    "Figure 1 shows an overview of the umbrella sampling method. For more detalied explanation of the theory of umbrella sampling, see [umbrella_sampling_theory.md](umbrella_sampling_theory.md).\n",
    "\n",
    "\n",
    "![umbrellasampling](./images/umbrella_sampling.svg)\n",
    "\n",
    "**Figure 1:** Umbrella sampling method to compute a free energy profile. (a) Multiple biasing potentials are placed across the collective variable $x$. The blue curve is the free energy of the system which we are trying to calculate. (b) Simulations are run for each window $w$. The resulting biased probability distributions $P'(x)$ are plotted. (c) The unbiased free energies $F_i$ from each window. They are each offset by a different $C_i$. (d) The Weighted Histogram Analysis Method (WHAM) is used to combine the windows and compute the free energy curve."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Umbrella Sampling Simulations\n",
    "\n",
    "Umbrella sampling simulations generally involve three main steps:\n",
    "\n",
    "1. **Preparing the Windows**: This is usually done using Steered Molecular Dynamics (SMD) to set up the initial configurations.\n",
    "2. **Running the Windows**: These simulations can be performed in parallel to leverage computational resources effectively.\n",
    "3. **Analyzing the Results**: This step typically involves computing the potential of mean force (PMF) using the Weighted Histogram Analysis Method (WHAM).\n",
    "\n",
    "It's common to find in step 3 that the simulations do not produce a well-converged PMF. If this happens, you may need to return to step 1 and adjust some settings. This trial-and-error feedback loop is a normal part of the process. For this tutorial, we will use settings that are already known to work."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### System\n",
    "<a id=\"umbrellasystem\"></a>\n",
    " \n",
    "In this tutorial, we will use umbrella sampling to compute the free energy profile of the end-to-end distance ($r$) of deca-alanine in vacuum. Deca-alanine is commonly used as a toy system [[3](#References)]. Its equilibrium structure is a stable alpha-helix. Starting with the alpha-helix structure, we will first perform a SMD simulation to pull it from a helix into a coil. Next, we will run 24 umbrella sampling windows in the range of 1.3 nm to 3.3 nm. Finally, we will compute the PMF along $r$. Figure 2 shows the initial alpha-helix structure and the final extended coil structure.\n",
    "\n",
    "![deca-alanine](./images/deca-alanine.png)\n",
    "\n",
    "**Figure 2:** Structure of deca-alanine for end-to-end distance $r$=1.3 and $r$=3.3nm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the deca-alanine PDB file\n",
    "!wget https://raw.githubusercontent.com/openmm/openmm_workshops/main/section_2/deca-ala.pdb"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### Step 1 - Setting up the Windows with SMD\n",
    "<a id=\"setupwindows\"></a> \n",
    "\n",
    "The script below performs the following steps:\n",
    "- Loads a PDB file.\n",
    "- Defines a collective variable between the first and last alpha-carbon. This is done by using [`CustomBondForce`](http://docs.openmm.org/latest/api-python/generated/openmm.openmm.CustomBondForce.html#openmm.openmm.CustomBondForce) to record the distance between these two atoms.\n",
    "- Adds a harmonic restraint to the CV using [`CustomCVForce`](http://docs.openmm.org/latest/api-python/generated/openmm.openmm.CustomCVForce.html#openmm.openmm.CustomCVForce).\n",
    "- Runs a simulation where the location of the harmonic restraint is moved with constant velocity from 1.3 nm to 3.3 nm (this is called constant velocity steered MD).\n",
    "- Saves a configuration for each of the 24 equally spaced windows between 1.3 nm and 3.3 nm.\n",
    "\n",
    "Make sure you read through the full script before you run the cell!\n",
    "\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "ℹ️ <b>Exercise 3</b>\n",
    "\n",
    "Add code to the cell below to complete setting up the `CustomCVForce`.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmm as mm\n",
    "import openmm.app as app\n",
    "import openmm.unit as unit\n",
    "from sys import stdout\n",
    "import numpy as np\n",
    "\n",
    "pdb = app.PDBFile('deca-ala.pdb')\n",
    "\n",
    "forcefield = app.ForceField('amber14-all.xml')\n",
    "\n",
    "# We have a single molecule in vacuum so we use no cutoff.\n",
    "system = forcefield.createSystem(pdb.topology, nonbondedMethod=app.NoCutoff, constraints=app.HBonds, hydrogenMass=1.5*unit.amu)\n",
    "integrator = mm.LangevinMiddleIntegrator(300*unit.kelvin, 1/unit.picosecond, 0.004*unit.picoseconds)\n",
    "simulation = app.Simulation(pdb.topology, system, integrator)\n",
    "simulation.context.setPositions(pdb.positions)\n",
    "simulation.context.setVelocitiesToTemperature(300*unit.kelvin)\n",
    "\n",
    "# Add reporters to output data\n",
    "simulation.reporters.append(app.DCDReporter('smd_traj.dcd', 10000))\n",
    "simulation.reporters.append(app.StateDataReporter(stdout, 10000, step=True, time=True, potentialEnergy=True, temperature=True, speed=True))\n",
    "\n",
    "# Equilibrate\n",
    "simulation.step(1000)\n",
    "\n",
    "# Define the CV as the distance between the CAs of the two end residues.\n",
    "# You can open the \"deca-ala.pdb\" file to check these indices.\n",
    "index1 = 8\n",
    "index2 = 98\n",
    "cv = mm.CustomBondForce('r')\n",
    "cv.addBond(index1, index2)\n",
    "\n",
    "# Now setup SMD\n",
    "# Starting value\n",
    "r0 = 1.3*unit.nanometers\n",
    "\n",
    "# Force constant\n",
    "fc_pull = 1000.0*unit.kilojoules_per_mole/unit.nanometers**2 \n",
    "\n",
    "# Pulling speed\n",
    "v_pulling = 0.02*unit.nanometers/unit.picosecond # nm/ps\n",
    "\n",
    "# Simulation time step\n",
    "dt = simulation.integrator.getStepSize()\n",
    "\n",
    "# Total number of steps\n",
    "total_steps = 30000 # 120ps\n",
    "\n",
    "# Number of steps to run between incrementing r0 (1 makes the simulation slow)\n",
    "increment_steps = 10\n",
    "\n",
    "# Define a harmonic restraint on the CV\n",
    "# The location of the restrain will be moved as we run the simulation\n",
    "# This is constant velocity steered MD\n",
    "pullingForce = mm.CustomCVForce('0.5 * fc_pull * (cv-r0)^2')\n",
    "\n",
    "# Add the global parameters fc_pull and r0 to the CustomCVforce\n",
    "FIXME\n",
    "\n",
    "pullingForce.addCollectiveVariable(\"cv\", cv)\n",
    "system.addForce(pullingForce)\n",
    "simulation.context.reinitialize(preserveState=True)\n",
    "\n",
    "# Define the windows\n",
    "# During the pulling loop we will save specific configurations corresponding to the windows\n",
    "windows = np.linspace(1.3, 3.3, 24)\n",
    "window_coords = []\n",
    "window_index = 0\n",
    "\n",
    "# SMD pulling loop\n",
    "for i in range(total_steps//increment_steps):\n",
    "    simulation.step(increment_steps)\n",
    "    current_cv_value = pullingForce.getCollectiveVariableValues(simulation.context)\n",
    "    \n",
    "    if (i * increment_steps) % 5000 == 0:\n",
    "        print(\"r0 = \", r0, \"r = \", current_cv_value)\n",
    "    \n",
    "    # Increment the location of the CV based on the pulling velocity\n",
    "    r0 += v_pulling * dt * increment_steps\n",
    "    simulation.context.setParameter('r0', r0)\n",
    "\n",
    "    # Check if we should save this config as a window starting structure\n",
    "    if (window_index < len(windows) and current_cv_value >= windows[window_index]):\n",
    "        window_coords.append(simulation.context.getState(getPositions=True, enforcePeriodicBox=False).getPositions())\n",
    "        window_index += 1\n",
    "\n",
    "# Save the window structures\n",
    "for i, coords in enumerate(window_coords):\n",
    "    outfile = open(f'window_{i}.pdb', 'w')\n",
    "    app.PDBFile.writeFile(simulation.topology,coords, outfile)\n",
    "    outfile.close()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the script has completed running there will be 24 new pdb files called `'window_n.pdb'` where `n` is an integer ranging from 0 to 23.\n",
    " \n",
    "We now have the initial configurations for the umbrella sampling windows."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2 - Running the Windows\n",
    "<a id=\"runningwindows\"></a>\n",
    " \n",
    "The script to run the windows is very similar to the script in step 1. The key differences are that we load in an initial structure that corresponds to each specific window and that the harmonic restraint on the CV does not move. The script below defines a function to run one window. It re-uses the `Simulation` we created in step 1.\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "ℹ️ <b>Exercise 4</b>\n",
    "\n",
    "Add code to set the value of `r0` for the window.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_window(window_index):\n",
    "    print('running window', window_index)\n",
    "    \n",
    "    # Load the starting configuration for this window\n",
    "    pdb = app.PDBFile(f'window_{window_index}.pdb')\n",
    "\n",
    "    # We can reuse the existing Simulation object\n",
    "    simulation.context.setPositions(pdb.positions)\n",
    "    simulation.context.setVelocitiesToTemperature(300*unit.kelvin)\n",
    "\n",
    "    # Set the fixed location of the harmonic restraint for this window\n",
    "    r0 = windows[window_index]\n",
    "\n",
    "    # Set the value of r0 in the simulation\n",
    "    FIXME\n",
    "\n",
    "    # Run short equilibration with new positions and r0   \n",
    "    simulation.step(1000)\n",
    "\n",
    "    # Run the data collection\n",
    "    # Total number of steps\n",
    "    total_steps = 10000 # 400 ps\n",
    " \n",
    "    # Frequency to record the current CV value\n",
    "    record_steps = 100\n",
    "\n",
    "    # Run the simulation and record the value of the CV.\n",
    "    cv_values=[]\n",
    "    for i in range(total_steps//record_steps):\n",
    "        simulation.step(record_steps)\n",
    "\n",
    "        # get the current value of the cv\n",
    "        current_cv_value = pullingForce.getCollectiveVariableValues(simulation.context)\n",
    "        cv_values.append([i, current_cv_value[0]])\n",
    "\n",
    "    # Save the CV timeseries to a file so we can postprocess \n",
    "    np.savetxt(f'cv_values_window_{window_index}.txt', np.array(cv_values))\n",
    "\n",
    "    print('Completed window', window_index)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then run all 24 windows/."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for n in range(24):\n",
    "    run_window(n)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once all the window simulations have completed you will have the CV timeseries files: `'cv_values_window_n.txt'`."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3 - Analysis - Compute the PMF\n",
    "<a id=\"computepmf\"></a>\n",
    "\n",
    "The first thing to check is that the histograms of the CV timeseries have good overlap. Here is an example:\n",
    "\n",
    "![histogram](./images/hist.png)\n",
    "\n",
    "You can plot yours by running the cell below. This cell also produces the metadata file we will need for the next step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.figure(figsize=(10, 7)) \n",
    "\n",
    "# Plot the histograms\n",
    "metafilelines = []\n",
    "for i in range(len(windows)):\n",
    "    data = np.loadtxt(f'cv_values_window_{i}.txt')\n",
    "    plt.hist(data[:,1], alpha=0.6)\n",
    "    metafileline = f'cv_values_window_{i}.txt {windows[i]} 1000\\n'\n",
    "    metafilelines.append(metafileline)\n",
    "\n",
    "plt.xlabel(\"r (nm)\")\n",
    "plt.ylabel(\"count\")\n",
    "\n",
    "with open(\"metafile.txt\", \"w\") as f:\n",
    "    f.writelines(metafilelines)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To compute the PMF we can use WHAM [[1](#References)]. An easy to use and widely compatible implementation is the WHAM program by Alan Grossfield, which can be downloaded [here](http://membrane.urmc.rochester.edu/?page_id=126). It is a C program, so it will need to be compiled. The command below should work on Linux and Mac.\n",
    "\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "⚠️ <b>If the next two cells don't work for you, then skip them and move onto the plotting phase with the file we have provided for you, `'pmf_backup.txt'`, which you can use instead of `'pmf.txt'`.</b>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget http://membrane.urmc.rochester.edu/sites/default/files/wham/wham-release-2.0.11.tgz\n",
    "!tar xf wham-release-2.0.11.tgz\n",
    "!cd wham/wham && make"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use `wham`, we need a metadata file that lists the names of each CV timeseries file, the location of the harmonic restraints, and the value of the spring constant. We created this in the histogram plotting script earlier.\n",
    " \n",
    "The `wham` program is run using command line arguments. Read the documentation to find out more: http://membrane.urmc.rochester.edu/sites/default/files/wham/doc.pdf\n",
    " \n",
    "The command below will compute the PMF from our data. The command line arguments correspond to a range of 1.3 nm to 3.3 nm, 50 histogram bins, a tolerance of 1e-6, and a temperature of 300K."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!./wham/wham/wham 1.3 3.3 50 1e-6 300 0 metafile.txt pmf.txt > wham_log.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then plot the computed PMF. It should look something like this:\n",
    "\n",
    "![pmf.png](./images/pmf.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the PMF\n",
    "pmf = np.loadtxt(\"pmf.txt\")\n",
    "plt.figure(figsize=(7, 5)) \n",
    "plt.plot(pmf[:,0], pmf[:,1])\n",
    "plt.xlabel(\"r (nm)\")\n",
    "plt.ylabel(\"PMF (kJ/mol)\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extra Exercises\n",
    "<a id=\"extra\"></a>\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "ℹ️ <b>Exercise 5</b>\n",
    "\n",
    "In the [umbrella sampling setup stage](#setupwindows), we used constant velocity SMD to pull the peptide apart. There is also constant force steered MD. This is where a constant force is applied to the system to pull it into the desired configuration. Modify the window setup script to perform constant force steered MD. \n",
    "\n",
    "Tip: Use `CustomExternalForce`. \n",
    "</div>\n",
    "\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "ℹ️ <b>Exercise 6</b>\n",
    "\n",
    "Run the full umbrella sampling procedure for the deca-alanine peptide solvated in a water box.\n",
    "</div>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "<a id=\"References\"></a>\n",
    "\n",
    "[1] Kumar, S., Rosenberg, J.M., Bouzida, D., Swendsen, R.H. and Kollman, P.A. (1992), *The weighted histogram analysis method for free-energy calculations on biomolecules. I. The method.* J. Comput. Chem., 13: 1011-1021. https://doi.org/10.1002/jcc.540130812  \n",
    "[2] Kästner, J. (2011), *Umbrella sampling.* WIREs Comput Mol Sci, 1: 932-942. https://doi.org/10.1002/wcms.66  \n",
    "[3] Park, S., Khalili-Araghi, F., Tajkhorshid, E., and Schulten, K. (2003), *Free energy calculation from steered molecular dynamics simulations using Jarzynski’s equality.* J. Chem. Phys., 119 (6): 3559–3566. https://doi.org/10.1063/1.1590311  "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solutions\n",
    "<a id=\"solutions\"></a>\n",
    "\n",
    "*Exercise 1.* Add the per particle parameter:\n",
    "```python\n",
    "external_force.addPerParticleParameter('x0')\n",
    "```\n",
    "\n",
    "*Exercise 2.* Add the external force:\n",
    "```python\n",
    "new_system.addForce(external_force)\n",
    "```\n",
    "\n",
    "*Exercise 3.* If you visualize the trajectory you will see all the atoms get pulled to $x=0$ as show in the image on the left below. If you use `(x-x0)^2` instead of  `periodicdistance()`, you will see that the custom force does not work as expected due to the PBCs not being taken into account (shown in the right image below). \n",
    "\n",
    "![solution 3](images/solution3.png)\n",
    "\n",
    "*Exercise 4.* Add the global parameters fc_pull and r0 to the CustomCVforce:\n",
    "```python\n",
    "pullingForce.addGlobalParameter('fc_pull', fc_pull)\n",
    "pullingForce.addGlobalParameter('r0', r0)\n",
    "```\n",
    "\n",
    "*Exercise 5.* \n",
    "```python\n",
    "simulation.context.setParameter('r0',r0)\n",
    "```\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cookbook3",
   "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.16"
  },
  "nbsphinx": {
   "execute": "never"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}