Harmonic_Bias.md

jupyter
jupytext kernelspec formats text_representation ipynb,md extension format_name format_version jupytext_version .md markdown 1.3 1.16.4 display_name language name Python 3 python python3

Setting up the environment

First, we are setting up our environment. We use an already compiled and packaged installation of OpenMM and the DLExt plugin. We copy it from Google Drive and install PySAGES for it. We also have a Google Colab that performs this installation for reference, but that requires permissions that we do not want on our Google Drive.

BASE_URL="https://drive.google.com/u/0/uc?id=1hsKkKtdxZTVfHKgqVF6qV2e-4SShmhr7&export=download"
wget -q --load-cookies /tmp/cookies.txt "$BASE_URL&confirm=$(wget -q --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate $BASE_URL -O- | sed -rn 's/.*confirm=(\w+).*/\1\n/p')" -O pysages-env.zip
rm -rf /tmp/cookies.txt

%env PYSAGES_ENV=/env/pysages

mkdir -p $PYSAGES_ENV .
unzip -qquo pysages-env.zip -d $PYSAGES_ENV

import os
import sys

ver = sys.version_info

sys.path.append(os.environ["PYSAGES_ENV"] + "/lib/python" + str(ver.major) + "." + str(ver.minor) + "/site-packages/")

PySAGES

The next step is to install PySAGES. First, we install the jaxlib version that matches the CUDA installation of this Colab setup. See the JAX documentation here for more details.

We test the jax installation and check the versions.

import jax
import jaxlib
print(jax.__version__)
print(jaxlib.__version__)

Now we can finally install PySAGES. We clone the newest version from here and build the remaining pure python dependencies and PySAGES itself.

rm -rf PySAGES
git clone https://github.com/SSAGESLabs/PySAGES.git &> /dev/null
cd PySAGES
pip install -q . &> /dev/null

Harmonic Bias simulations

mkdir /content/harmonic-bias
cd /content/harmonic-bias

A harmonic bias simulation constraints a collective variable with a harmonic potential. This is useful for a variety of advanced sampling methods, in particular, umbrella sampling.

For this Colab, we are using alanine dipeptide as the example molecule, a system widely-used for benchmarking enhanced sampling methods. So first, we fetch the molecule from the examples of PySAGES.

cp /content/PySAGES/examples/inputs/alanine-dipeptide/adp-explicit.pdb ./

Next we load the PySAGES/OpenMM environment.

from pysages.colvars import DihedralAngle
from pysages.methods import HarmonicBias, HistogramLogger
import numpy as np
from pysages.utils import try_import

import pysages

openmm = try_import("openmm", "simtk.openmm")
unit = try_import("openmm.unit", "simtk.unit")
app = try_import("openmm.app", "simtk.openmm.app")

Next, we write a function that can generate an execution context for OpenMM. This is everything you would normally write in an OpenMM script, just wrapped as a function that returns the context of the simulation.

"""
Generates a simulation context, we pass this function to the attribute `run` of our sampling method.
"""
def generate_simulation(**kwargs):
    pdb_filename = "adp-explicit.pdb"
    T = 298.15 * unit.kelvin
    dt = 2.0 * unit.femtoseconds
    pdb = app.PDBFile(pdb_filename)

    ff = app.ForceField("amber99sb.xml", "tip3p.xml")
    cutoff_distance = 1.0 * unit.nanometer
    topology = pdb.topology
    system = ff.createSystem(
        topology, constraints = app.HBonds, nonbondedMethod = app.NoCutoff,
        nonbondedCutoff = cutoff_distance
    )

    positions = pdb.getPositions(asNumpy = True)

    integrator = openmm.LangevinIntegrator(T, 1 / unit.picosecond, dt)

    simulation = app.Simulation(topology, system, integrator)
    simulation.context.setPositions(positions)
    simulation.minimizeEnergy()

    return simulation

The next step is to define the collective variable (CV). In this case, we choose the two dihedral angles on the molecule as defined by the atom positions. We also choose an equilibrium value to constrain the dihedrals and the corresponding spring constant. The HarmonicBias class is responsible for introducing the bias into the simulation run.

cvs = (DihedralAngle((4, 6, 8, 14)), DihedralAngle((6, 8, 14, 16)))
center =[ -0.33*np.pi, -0.4*np.pi]
k = 100
method = HarmonicBias(cvs, k, center)

We now define a Histogram callback to log the measured values of the CVs and run the simulation for $10^4$ time steps. The HistogramLogger collects the state of the collective variable during the run. Make sure to run with GPU support. Using the CPU platform with OpenMM is possible and supported, but can take a very long time.

callback = HistogramLogger(50)
pysages.run(method, generate_simulation, int(1e4), callback)

Next, we want to plot the histogram as recorded from the simulations.

bins = 25
lim = (-np.pi/2, -np.pi/4)
lims = [lim for i in range(2)]
hist, edges = callback.get_histograms(bins=bins, range=lims)
hist_list = [np.sum(hist, axis=(0)), np.sum(hist, axis=(1))]

import matplotlib.pyplot as plt
fig, ax = plt.subplots()

ax.set_xlabel(r"CV $\xi_i$")
ax.set_ylabel(r"$p(\xi_i)$")

x = np.linspace(lim[0], lim[1], hist_list[0].shape[0])

for i in range(len(hist_list)):
    (line,) = ax.plot(x, hist_list[i], label="i= {0}".format(i))

ax.legend(loc="best")

We see how the dihedral angles are distributed. The histograms are not perfect in this example because we ran the simulation only for a few time steps and hence sampling is quite limited.