choderalab / integrator-benchmark

Code for enumerating and evaluating numerical methods for Langevin dynamics using near-equilibrium estimates of the KL-divergence. Accompanies https://doi.org/10.3390/e20050318

Not all Langevin integrators are equal

Enumerating and evaluating numerical schemes for Langevin dynamics

Molecular dynamics requires methods for simulating continuous equations in discrete steps

A widely-used approach to investigating equilibrium properties is to simulate Langevin dynamics. Langevin dynamics is a system of stochastic differential equations defined on a state space of configurations $\mathbf{x}$ and velocities $\mathbf{v}$.

To simulate those equations on a computer, we need to provide explicit instructions for advancing the state of the system $(\mathbf{x},\mathbf{v})$ by very small time increments.

Here, we will consider the family of methods that can be derived by splitting the Langevin system into a sum of three simpler systems, labeled O, R, and V. We then define a numerical method by approximately propagating each of those simpler systems for small increments of time in a specified order.

(TODO: Add LaTeX-rendered Langevin system with underbraces around O, R, V components, using readme2tex)

We will refer to a numerical scheme by its encoding string. For example, OVRVO means: simulate the O component for a time increment of $\Delta t/2$, then the V component for $\Delta t/2$, then the R component for $\Delta t$, then the V component for $\Delta t/2$, and finally the O component for $\Delta t/2$. This approximately propagates the entire system for a total time increment of $\Delta t$.

This introduces error that can be sensitive to details

Using subtly different numerical schemes for the same continuous equations can lead to drastically different behavior at finite timesteps. As a prototypical example, consider the difference between the behavior of schemes VRORV and OVRVO on a 1D quartic system:

($\rho$ is the density sampled by the numerical scheme, and $\pi$ is the exact target density. Each column illustrates an increasing finite timestep $\Delta t$, below the "stability threshold." Rows 2 and 4 illustrate error in the sampled joint distribution, and rows 1 and 3 illustrate error in the sampled $\mathbf{x}$-marginal distribution.)

The two schemes have the same computational cost per iteration, and introduce nearly identical levels of error into the sampled joint $(\mathbf{x}, \mathbf{v})$ distribution -- but one of these methods introduces about 100x more error in the $\mathbf{x}$ marginal than the other at large timesteps!

Toy implementation

To illustrate the scheme, here is a toy Python implementation for each of the explicit updates:

# Defined elsewhere: friction coefficient `gamma`, mass `m`

def propagate_R(x, v, h): 
    """Linear "drift" -- deterministic position update
    using current velocities"""
    return (x + (h * v), v)

def propagate_V(x, v, h):
    """Linear "kick" -- deterministic velocity update
    using current forces"""
    return (x, v + (h * force(x) / m))

def propagate_O(x, v, h):
    """Ornstein-Uhlenbeck -- stochastic velocity update
    using a ficticious "heat-bath""""
    a, b = exp(-gamma * h), sqrt(1 - exp(-2 * gamma * h))
    return (x, (a * v) + b * draw_maxwell_boltzmann_velocities())

propagate = {"O": propagate_O, "R": propagate_R, "V": propagate_V}

where draw_maxwell_boltzmann_velocities() draws an independent sample from the velocity distribution given by the masses and temperature.

Using the functions we just defined, here's how to implement the inner-loop of the scheme denoted OVRVO:

x, v = propagate_O(x, v, dt / 2)
x, v = propagate_V(x, v, dt / 2)
x, v = propagate_R(x, v, dt)
x, v = propagate_V(x, v, dt / 2)
x, v = propagate_O(x, v, dt / 2)

And here's the VRORV inner-loop:

x, v = propagate_V(x, v, dt / 2)
x, v = propagate_R(x, v, dt / 2)
x, v = propagate_O(x, v, dt)
x, v = propagate_R(x, v, dt / 2)
x, v = propagate_V(x, v, dt / 2)

As suggested from these examples, the generic recipe for turning a splitting string into a Langevin integrator is:

def get_n(substep="O", splitting="OVRVO"):
    return sum([s == substep for s in splitting])

def simulate_timestep(x, v, dt, splitting="OVRVO"):
    for substep in splitting:
        x, v = propagate[substep](x, v, dt / get_n(substep, splitting))
    return x, v

Systematically enumerating numerical schemes and measuring their error

In this repository, we enumerate numerical schemes for Langevin dynamics by associating strings over the alphabet {O, R, V} with explicit numerical methods using OpenMM CustomIntegrators. We provide schemes for approximating the error introduced by that method in the sampled distribution over $(\mathbf{x},\mathbf{v}) jointly or $\mathbf{x}$ alone using nonequilibrium work theorems. We further investigate the effects of modifying the mass matrix (aka "hydrogen mass repartitioning") and/or evaluating subsets of the forces per substep (aka "multi-timestep" methods, obtained by expanding the alphabet to {O, R, V0, V1, ..., V32}).

(TODO: Details on nonequilibrium error measurement scheme.)

Relation to prior work

We did not introduce the concept of splitting the Langevin system into these three "building blocks" -- this decomposition is developed lucidly in Chapter 7 of [Leimkuhler and Matthews, 2015]. We also did not discover the VRORV integrator -- Leimkuhler and Matthews have studied the favorable properties of particular integrator VRORV ("BAOAB" in their notation) in great detail. Here, we have:

1. provided a method to translate these strings into efficient CustomIntegrators in OpenMM,
2. provided a uniform scheme for measuring the sampling error introduced by any member of this family of methods on any target density (approximating the KL divergence directly, rather than monitoring error in a system-specific choice of low-dimensional observable),
3. considered an expanded alphabet, encompassing many widely-used variants as special cases.

References

(TODO: Add references from paper!)

integrator-benchmark's Issues

Lazy sample generation

Currently, I fetch or generate equilibrium samples upon import. It would be more appropriate to generate samples only when needed, using a @property decorator, as suggested by John (#18 (comment))

Cache equilibrium data

As suggested by John: #5 (comment)

Add citation for Entropy paper to README

We should cite the Entropy paper in the README

Have tests that require OpenCL fall back to CPU (or be marked as `slow`)?

Some of the tests require OpenCL with mixed precision, which is not available on some platforms (travis and my MacBook Pro). It would be great if we could have these fall back to CPU or label them as slow to skip them on travis.

Bookkeepers should handle NPT when generating equilibrium samples for periodic systems

For systems like water boxes and Lennard-Jones fluids, we should be sure to allow the box size to fluctuate (and record the box vectors) when generating initial equilibrium samples. Currently, code like this assumes the volume remains fixed.

We should be careful to remove or disable the barostat during work measurement trajectories, but it is important to have an appropriate equilibrium distribution sampled from NPT for these systems since the initial densities may have unphysical pressures if the volume is constrained.

Note that WaterBox, DHFRExplicit, and SrcExplicit should use 298K and 1atm, but the LennardJonesFluid should use a different set of parameters to keep it in the liquid or gas phases (depending on which one we want to simulate).

Remaining Travis issues: TypeError in Python 3 on the serialize testsystems script ( https://travis-ci.org/choderalab/integrator-benchmark/jobs/226168539#L667 ). Python 3 import error in tests.py

Parallelize

After populating the sample cache, all experiments are independent. I'll parallelize over these conditions to speed up experiments.

Get Travis tests to pass!

Six remaining failures (2 import errors, 2 missing arguments, 1 unknown global, 1 call signature error)

Avoid storing energies or work values as unit-bearing quantities

We should adopt the convention that all dimensionless energy or work-like quantities are always kT, not kJ/mol.

Code like here can often trip us up.

Missing baoab_vs_aboba_analysis.py file

This line in the master branch refers to a baoab_vs_aboba_analysis.py file that doesn't seem to be in the repository.

I'm a bit confused as to which branch is also the current one here...

Integrator precision and PME settings

Do we mean to have CUDA set to double precision in testsystems.configuration? while OpenCL is set to mixed precision?

Handle NaNs correctly

NaNs can be encountered in the nonequilibrium protocols either in the 0 -> M trajectory or in the M -> 2M trajectory, leading to runtime errors. Will want to catch Exception: Particle coordinate is nan.

Also, instead of terminating early when NaNs are encountered, @jchodera suggested to count the number of NaNs encountered for each scheme.

Nightly build failed

Nightly build of master passed yesterday, but nightly build of master failed today during conda install.

Seems to be a conda-build bug introduced a few hours ago, which should hopefully be resolved soon, see conda-build issue #2008 and PR #2009 ...

Update openmmtools.testsystems with better defaults?

benchmark.testsystems contains a bunch of settings that override the defaults to build useful testsystems with appropriate PME tolerances, etc. When appropriate, it would be useful to update the defaults in openmmtools.testsystems and create testsystem variants there that we can use directly, since this will make it easier to just use the "off the shelf" testsystems in other projects (e.g. the shadow work paper).

Potential unit inconsistencies

I find it's much more foolproof to store work in units of kT whenever we need to store them as a unitless array. We appear to use kJ/mol in the code at various places, such as here, which has the potential to cause some confusion at best and errors at worst (if we accidentally feed them to, say, EXP or BAR before converting to kT).

If we do store energies or work with specific units, we should make sure to actually store unit-bearing quantities. Code like this has the potential to cause confusion by returning unitless quantities that must be interpreted in kJ/mol.

The best policy is

• If the quantity is supposed to be unit-bearing, make sure it is a simtk.unit.Quantity with the appropriate units
• If you want it to be dimensionless, convert it to units of kT; it's then safe to be unitless

Replace serialize_integrators.py with PrettyPrintableIntegrator

Replace the pretty printing code here with use of PrettyPrintableIntegrator

Use centrally-stored global parameters for temperature, collision rate, etc

Some parameters, like temperature, are explicitly written in multiple places, such as

• temperature used for generating equilibrium samples
• default parameters for LangevinSplittingIntegrator

It may be useful to use a centrally-located set of parameters (like the system_params in testsystems.py) to make sure there is no chance these are accidentally set to different values such that the equilibrium samples are generated at one temperature and the work measurements are done at a different temperature.

Issues in measuring work with constrained systems

@maxentile noticed that the estimates of $DeltaF_neq$ for constrained alanine dipeptide were large and negative, which shouldn't happen. Examination of the shadow work distributions shows the centers of these distributions are displaced by roughly -10 kJ/mol, suggesting we're seeing a loss of kinetic energy roughly equal to the velocity components along all 12 constrained bonds.

For alanine dipeptide with bonds to hydrogen constrained, the mean is about -10 kJ/mol.
There are 12 bonds to hydrogen constrained, so on average, this is a loss of kinetic energy of roughly 12 * (kT/2) = 12 * (0.6 kcal/mol)/2 * (4.184 kJ/kcal) = 15 kJ/mol.
So we're roughly in the ballpark of each trajectory losing an amount of kinetic energy equal to constraining those bonds to hydrogen.

My current belief is that this is due to a failure to constrain the velocities before measuring the initial kinetic energies, which would affect the shadow work for integrators that begin with either a V or R step (but not O step, where it will only affect the heat).

The way the baoab_vs_vvvr.py experiments work now is to:

• (1) Draw a position sample from equilibrium
• (2) Assign velocities from the Maxwell-Boltzmann distribution without removing the components of velocities along constrained bonds
• (3) The first step in BAOAB is a V step, which measures the kinetic energy without first constraining the velocities. The velocities are then constrained at the end of the V step, and the kinetic energy is measured again. Unsurprisingly, we lose an amount of energy equal to the bond constraint velocity components.

We have to fix either (or both) of (2) and (3).

For constrained systems, we should be removing the component of the velocity along the constrained bonds.

For safety, however, if we want to ensure our integrators always work, we may want to constrain the velocities at the beginning of each step where we measure the kinetic energy change.
We can't guarantee that the user will do this otherwise.

This comes with some overhead, so we could alternatively just constrain the positions and velocities on the first step using an idiom like this:
https://github.com/choderalab/perses/blob/master/perses/annihilation/ncmc_switching.py#L1227-L1242

            self.beginIfBlock('step = 0')
            # Constrain initial positions and velocities
            self.addConstrainPositions()
            self.addConstrainVelocities()
            # End block
            self.endBlock()

Does DeltaF_neq measure KL-divergence in (x,v) or (x,p) space?

Travis failures

@maxentile : It looks like travis is failing for odd reasons:

• The same alanine_constrained system generates samples multiple times
• Downstream processing can't find the .npy files generated

Is this primarily an issue of needing to provide a way to control where these data files are placed via an environment variable?

Automatically determine protocol length

The variance of the DeltaF_neq estimator (and computation time per protocol) increases if the protocol length M is too large. If M is too small (i.e. haven't reached nonequilibrium steady state), our DeltaF_neq estimates are more biased.

Currently, I set M "by eye." I'll want to add an automatic way to select M.

Modify testsystems so contexts aren't created upon import

This was just inefficient for CPU / Reference platforms, but causes runtime errors on CUDA platform...

Constrained velocities sampled from wrong distribution

As John points out here, openmm/openmm#1788 (comment) , I had been sampling velocities incorrectly.

To fix, I'll modify this function to sample v conditioned on x (setPositions, then setVelocitiesToTemperature).

Debugging issues with experiments/baoab_vs_vvvr.py

@maxentile : I'm trying to dig into what might be going on with negative estimates of $\Delta F_{neq}$ from the baoab_vs_vvvr.py example you pointed us to, but after installing the integrator-benchmark package via

python setup.py develop

there seems to be some issue with importing benchmark.testsystems:

lski1962:experiments choderaj$ python baoab_vs_vvvr.py 
Traceback (most recent call last):
  File "baoab_vs_vvvr.py", line 1, in <module>
    import benchmark.testsystems
  File "/Users/choderaj/github/choderalab/integrator-benchmark/integrator-benchmark/benchmark/experiments/benchmark.py", line 5, in <module>
    from code.integrators import LangevinSplittingIntegrator
ImportError: No module named 'code.integrators'; 'code' is not a package

My guess is

• There is some code that hasn't been checked in. master was last updated 24 days ago (!!!)
• You might have some old code installed locally as code.integrators that makes this work locally for you
• You are using a branch other than master. Could this be refactor (updated 19 days ago)?

Any idea what the issue might be?