Getting started tutorial

This tutorial will give you an overview of what you can find in OpenMMTools and how you can use the library.

Test systems, integrators, and forces

In its basic usage, OpenMMTools extends OpenMM by providing pre-packaged systems, integrators of force objects that are not natively implemented in OpenMM.

Test systems

The testsystems module comes with many simulation-ready molecular systems (from analytically-solvable systems to a kinase in explicit solvent) that can be useful for prototyping, validation, testing, and benchmarking. The code below creates a TIP3P water cubic box of 2nm side using PME.

from simtk import openmm, unit
from openmmtools import testsystems

water_box = testsystems.WaterBox(box_edge=2.0*unit.nanometer)
system = water_box.system  # An OpenMM System object.
positions = water_box.positions  # Initial coordinates for the system with associated units.

You can use select a subset of the system atoms using the atom selection domain-specific language (DSL) implemented in MDTraj. For example, the following snippet create a T4-Lysozyme system in implicit OBC GBSA solvent bound to a p-xylene molecule, and finds the atom indices corresponding to the heavy atoms of p-xylene and few residues surrounding the binding site of T4-Lysozyme.

lysozyme_pxylene = testsystems.LysozymeImplicit()
t4_system = lysozyme_pxylene.system
pxylene_dsl = '(resname TMP) and (mass > 1.5)'  # Select heavy atoms of p-xylene.
binding_site_dsl = ('(resi 77 or resi 86 or resi 101 or resi 110 or '
                    ' resi 117 or resi 120) and (mass > 1.5)')
pxylene_atom_indices = lysozyme_pxylene.mdtraj_topology.select(pxylene_dsl).tolist()
binding_site_atom_indices = lysozyme_pxylene.mdtraj_topology.select(binding_site_dsl).tolist()

Integrators

The systems created by testsystems can then be propagated in the usual way with OpenMM. The integrators module provide several high-quality integrators for equilibrium and non-equilibrium simulations in OpenMM.

from openmmtools import integrators

integrator = integrators.LangevinIntegrator(temperature=298.0*unit.kelvin,
                                            collision_rate=1.0/unit.picoseconds,
                                            timestep=1.0*unit.femtoseconds)
context = openmm.Context(t4_system, integrator)
context.setPositions(lysozyme_pxylene.positions)
integrator.step(n_steps)

Our LangevinIntegrator allows you to specify the splitting used to carry out the numerical integration. By default, OpenMMTools will construct a BAOAB integrator (i.e. with V R O R V splitting), which was shown empirically to add a very small integration error in configurational space, but other solutions are possible.

integrator = integrators.LangevinIntegrator(splitting="V0 V1 R R O R R V1 R R O R R V1 V0",
                                            measure_shadow_work=True, measure_heat=True)
context = openmm.Context(t4_system, integrator)
context.setPositions(lysozyme_pxylene.positions)
integrator.step(n_steps)

# Obtain the dissipated heat accumulated during Langevin dynamics in molar energy units.
heat = integrator.get_heat()

The integrator above, for example, implements the geodesic-BAOAB Langevin integrator with solute-solvent splitting, and it collects statistics on the dissipated heat and the shadow work during the propagation (at the cost of a computational overhead).

Forces

The forces module is still under construction, but it already provides a few convenient utility functions and force objects. Let’s create a T4-Lysozyme system in implicit OBC GBSA solvent bound to a p-xylene and add a harmonic restraint between the two molecules.

from openmmtools import forces

harmonic_restraint = forces.HarmonicRestraintForce(spring_constant=0.2*unit.kilocalories_per_mole/unit.angstrom**2,
                                                   restrained_atom_indices1=binding_site_atom_indices,
                                                   restrained_atom_indices2=pxylene_atom_indices)
t4_system.addForce(harmonic_restraint)

The restraint force above will place a single harmonic potential between the centers of mass of the heavy atoms of the p-xylene molecule and the binding site of T4-Lysozyme.

The function forces.find_forces() provides a convenient way to search for particular force objects in the OpenMM System.

# Retrieve our harmonic restraint force.
forces.find_forces(t4_system, force_type=forces.HarmonicRestraintForce)

# Find all forces that inherit from an OpenMM CustomBondForce object.
forces.find_forces(t4_system, force_type=openmm.CustomBondForce, include_subclasses=True)

# Search for force names using regular expressions.
# Return all openmm.HarmonicBondForce, openmm.HarmonicAngleForce,
# and forces.HarmonicRestraintForce force objects.
forces.find_forces(t4_system, '.*Harmonic.*')

Alchemical transformations

The alchemy module provides helper classes to perform alchemical transformations with OpenMM.

AbsoluteAlchemicalFactory and AlchemicalState

The AbsoluteAlchemicalFactory class prepare OpenMM System objects for alchemical manipulation. Let’s create an alchemical system that we can use to alchemically decouple p-xylene from T4-lysozyme’s binding pocket.

>>> from openmmtools import alchemy

>>> # Create the reference OpenMM System that will be alchemically modified.
>>> lysozyme_pxylene = testsystems.LysozymeImplicit()
>>> t4_system = lysozyme_pxylene.system

>>> # Define the region of the System to be alchemically modified.
>>> pxylene_atoms = lysozyme_pxylene.mdtraj_topology.select('resname TMP')
>>> alchemical_region = alchemy.AlchemicalRegion(alchemical_atoms=pxylene_atoms)

>>> factory = alchemy.AbsoluteAlchemicalFactory()
>>> alchemical_system = factory.create_alchemical_system(t4_system, alchemical_region)

At this point, the p-xylene in alchemical System is in its interacting state and it can be then simulated normally

>>> integrator = integrators.LangevinIntegrator()
>>> context = openmm.Context(alchemical_system, integrator)
>>> context.setPositions(lysozyme_pxylene.positions)
>>> integrator.step(n_steps)

The alchemical degrees of freedom of the Hamiltonian can be controlled during the simulation through the AlchemicalState class.

>>> alchemical_state = alchemy.AlchemicalState.from_system(alchemical_system)
>>> alchemical_state.lambda_electrostatics = 0.0
>>> alchemical_state.lambda_sterics = 0.5
>>> alchemical_state.apply_to_context(context)

The snippet above modifies the simulated System to completely turn off the electrostatics interaction and halve the Lennard-Jones potential between p-xylene and its environment.

Note

In OpenMMTools, the convention is to have the interacting state at lambda=1 and the non-interacting state at lambda=0. Some packages adopt the opposite convention.

Note

The AbsoluteAlchemicalFactory class is currently specialized for absolute calculations in the sense that it cannot prepare an OpenMM System to have an atom changing its element or turn on part of a molecule while decoupling another set of atoms. We’re planning to provide these capabilities in the near future.

Decoupling vs annihilating and softcore nonbonded interactions

By default, the alchemical System is prepared to annihilate electrostatics (i.e. turn off the alchemical atoms’ charges) and decouple the sterics (i.e. preserve the intra-molecular Lennard-Jones interactions), but you can maintain the intra-molecular charges, for example, by configuring the alchemical region.

alchemical_region = alchemy.AlchemicalRegion(alchemical_atoms=pxylene_atoms,
                                             annihilate_electrostatics=True)
alchemical_system = factory.create_alchemical_system(t4_system, alchemical_region)

Similarly, you can set specific softcore parameters for the sterics and electrostatics interactions (see the API documentation for a detailed explanation of the parameters).

alchemical_region = alchemy.AlchemicalRegion(alchemical_atoms=pxylene_atoms,
                                             softcore_alpha=0.5, softcore_c=6)

Softening torsions, angles, and bonds

Beside nonbonded interactions, it is possible to modify other terms of the potentials. The following alchemical region is configured to modify the OpenMM System to enable torsion softening of all the p-xylene dihedrals. The Hamiltonian parameter controlling the torsion, angles, and bond potential terms can be controlled with AlchemicalState in the same way as with nonbonded interactions.

alchemical_region = alchemy.AlchemicalRegion(alchemical_atoms=pxylene_atoms,
                                             alchemical_torsions=True)
alchemical_system = factory.create_alchemical_system(t4_system, alchemical_region)
context = openmm.Context(alchemical_system, integrators.LangevinIntegrator())

alchemical_state = alchemy.AlchemicalState.from_system(alchemical_system)
alchemical_state.lambda_torsions = 0.8
alchemical_state.apply_to_context(context)

Alchemical functions

Finally you can enslave the degrees of freedom of the Hamiltonian to a variable through a custom function. The code below configure the AlchemicalState to turn off first electrostatic and the steric interactions one after the other as a generic variable called lambda goes from 1.0 to 0.0.

# Enslave lambda_sterics and lambda_electrostatics to a generic lambda variable.
alchemical_state.set_function_variable('lambda', 1.0)

# The functions here turn off first electrostatic and the steric interactions
# in sequence as lambda goes from 1.0 to 0.0.
f_electrostatics = '2*(lambda-0.5)*step(lambda-0.5)'
f_sterics = '2*lambda*step_hm(0.5-lambda) + step_hm(lambda-0.5)'
alchemical_state.lambda_electrostatics = alchemy.AlchemicalFunction(f_electrostatics)
alchemical_state.lambda_sterics = alchemy.AlchemicalFunction(f_sterics)

alchemical_state.set_function_variable('lambda', 0.75)
assert alchemical_state.lambda_electrostatics == 0.5
assert alchemical_state.lambda_sterics == 1.0

alchemical_state.set_function_variable('lambda', 0.25)
assert alchemical_state.lambda_electrostatics == 0.0
assert alchemical_state.lambda_sterics == 0.5

# Set the alchemical state of the simulated system.
alchemical_state.apply_to_context(context)

In the example above, step_hm is the Heaviside step function with half-maximum convention (i.e. step_hm(0.0) == 0.5), while step(0.0) == 0.0. All the functions in the Python standard module math can be specified in the string.


Manipulating the thermodynamic state of your simulation

The classes in the states module provide a framework to decouple the degrees of freedom (or parameters) of the simulated thermodynamic state from their implementation details in OpenMM.

Defining temperature and pressure

The fundamental class in the states module is ThermodynamicState. This class hold a System object and controls the ensemble parameters of temperature and pressure. For example, the code below creates a water box in NVT ensemble at 298 K.

>>> from openmmtools import states

>>> waterbox = testsystems.WaterBox(box_edge=2*unit.nanometers)
>>> thermo_state = states.ThermodynamicState(system=waterbox.system,
...                                          temperature=298.0*unit.kelvin)
>>> thermo_state.volume.format('%.1f')
'8.0 nm**3'
>>> assert thermo_state.pressure is None

The volume is computed from the box vectors associated to the System object. To convert the system to an NPT state at 298 K and 1 atm pressure, you can set the pressure attribute.

thermo_state.pressure = 1.0*unit.atmosphere
assert thermo_state.volume is None

Note that the operation of specifying a constant pressure result in a null volume, as the volume will fluctuate during the simulation. You can then create an OpenMM Context object that is guaranteed to be in the specified thermodynamic state.

>>> integrator = integrators.LangevinIntegrator(temperature=298.0*unit.kelvin)
>>> context = thermo_state.create_context(integrator)
>>> context.setPositions(waterbox.positions)
>>> integrator.step(n_steps)

>>> # ThermodynamicState takes care of adding and configuring a MonteCarloBarostatForce
>>> # to keep the pressure at 1atm.
>>> force_index, barostat = forces.find_forces(context.getSystem(),
...                                            openmm.MonteCarloBarostat,
...                                            only_one=True)
>>> barostat.getDefaultTemperature().format('%.1f')
'298.0 K'
>>> print(barostat.getDefaultPressure())
1.01325 bar

Consistency checks for free

Using the ThermodynamicState class means to take advantage of several consistency checks that can avoid bugs in your application that can be very hard to detect in the first place and then to track down (we speak from personal experience).

For example, trying to create a Context using Langevin integrator set to the incorrect temperature or trying to add a barostat to a system in vacuum raises an error.

>>> thermo_state.temperature = 298.0*unit.kelvin
>>> integrator = integrators.LangevinIntegrator(temperature=310.0*unit.kelvin)
>>> thermo_state.create_context(integrator)
Traceback (most recent call last):
...
ThermodynamicsError: Integrator is coupled to a heat bath at a different temperature.
>>> vacuum_system = testsystems.TolueneVacuum().system
>>> thermo_state = states.ThermodynamicState(system=vacuum_system,
...                                          temperature=298.15*unit.kelvin,
...                                          pressure=1.0*unit.atmosphere)
Traceback (most recent call last):
...
ThermodynamicsError: Non-periodic systems cannot have a barostat.

While, if you create a Context with an integrator that is not coupled to a heat bath, ThermodynamicState will take care of adding an AndersenThermostat.

>>> # Use a non-thermostated integrator.
>>> thermo_state_nvt = states.ThermodynamicState(system=vacuum_system,
...                                              temperature=298.15*unit.kelvin)
>>> integrator = openmm.VerletIntegrator(2.0*unit.femtoseconds)
>>> context_nvt = thermo_state_nvt.create_context(integrator)
>>> len(forces.find_forces(context_nvt.getSystem(), openmm.AndersenThermostat))
1

Manipulating the thermodynamic state: Compatible thermodynamic states

Once a Context has been created, is is possible to change the simulation thermodynamic state through the method ThermodynamicState.apply_to_context(). The method will mask the implementation details and take care of modifying all the OpenMM forces and integrators that depend on the temperature and pressure parameters. In this sense, the ThermodynamicState class decouples the representation of the thermodynamic parameters from their implementation details.

>>> # Modify temperature and pressure of a system employing a Langevin
>>> # thermostat and a Monte Carlo barostat.
>>> thermo_state.temperature = 400.0*unit.kelvin
>>> thermo_state.pressure = 1.2*unit.bar
>>> thermo_state.apply_to_context(context)
>>> context.getIntegrator().getTemperature().format('%.1f')
'400.0 K'
>>> context.getParameter(openmm.MonteCarloBarostat.Pressure())
1.2
>>> # The MonteCarloBarostat requires also a temperature parameter for the acceptance probability.
>>> context.getParameter(openmm.MonteCarloBarostat.Temperature())
400.0
>>> # Modify the temperature of a system using an Andersen thermostat.
>>> thermo_state_nvt.temperature = 400.0*unit.kelvin
>>> thermo_state_nvt.apply_to_context(context_nvt)
>>> context_nvt.getParameter(openmm.AndersenThermostat.Temperature())
400.0

A ThermodynamicState can be applied to any Context that was created from a compatible thermodynamic state.

Important

Two ThermodynamicState objects x, y are compatible if a context created by x can be modified to be in the y thermodynamic state through y.apply_to_context(context) and viceversa.

This is not always possible in OpenMM because of some implementation details related to optimizations. In short, two ThermodynamicState``s are compatible if they have the same ``System and they are in the same ensemble (i.e. NVT and NPT thermodynamic states are incompatible).

>>> alanine = testsystems.AlanineDipeptideExplicit()
>>> state1 = states.ThermodynamicState(alanine.system, 273*unit.kelvin)
>>> state2 = states.ThermodynamicState(alanine.system, 310*unit.kelvin)
>>> state1.is_state_compatible(state2)
True

# Switch state1 from NVT to NPT ensemble.
>>> state1.pressure = 1.0*unit.atmosphere
>>> state1.is_state_compatible(state2)
False

Luckily, the class openmmtools.cache.ContextCache takes care of checking for compatibility and decide whether it’s possible to modifying a previously created Context object or if it is necessary to create a separate one.

Using the ContextCache

Important

Using ContextCache is the recommended way of creating Context objects within the OpenMMTools framework.

The openmmtools.cache.ContextCache class has the role of maintaining the minimum number of compatible Contexts allocated on the GPU, allowing virtually an infinite number of thermodynamic states to be simulated on finite-memory hardware, and minimizing the number of expensive Context creation/destruction.

To obtain a Context simply use the ContextCache.get_context() method.

from openmmtools import cache

alanine = testsystems.AlanineDipeptideExplicit()
thermo_state = states.ThermodynamicState(alanine.system, 310*unit.kelvin)
integrator = integrators.LangevinIntegrator(temperature=310*unit.kelvin)

context_cache = cache.ContextCache()
context, context_integrator = context_cache.get_context(thermo_state,
                                                        integrator)
context.setPositions(alanine.positions)
context_integrator.step(n_steps)

Note that get_context() returns also an Integrator that may be a different instance of the integrator passed as a parameter. This is because an OpenMM Context can be associated with a single integrator instance, thus reusing a previously instantiated Context requires using the previously instantiated integrator as well. Nevertheless, context_integrator is guaranteed to be identical to integrator.

Requesting a context in a compatible ThermodynamicState returns the same Context object correctly configured to simulate the requested thermodynamic state.

>>> compatible_state = states.ThermodynamicState(alanine.system, 400*unit.kelvin)
>>> compatible_integrator = integrators.LangevinIntegrator(temperature=400*unit.kelvin)
>>> compatible_context, compatible_integrator = context_cache.get_context(compatible_state,
...                                                                       compatible_integrator)
>>> id(context) == id(compatible_context)
True
>>> len(context_cache)  # The number of Contexts maintained in memory.
1
>>> compatible_integrator.getTemperature().format('%.1f')
'400.0 K'

Requesting a context in a different ensemble causes the creation of another Context.

>>> thermo_state_npt = copy.deepcopy(thermo_state)
>>> thermo_state_npt.pressure = 1.0*unit.atmosphere
>>> integrator = integrators.LangevinIntegrator(temperature=thermo_state_npt.temperature)
>>> context_npt, integrator_npt = context_cache.get_context(thermo_state_npt, integrator)
>>> id(context) == id(context_npt)
False
>>> len(context_cache)
2

You can set a capacity and a time to live for contexts. The time to live is currently measured in number of accesses to the ContextCache.

>>> context_cache = cache.ContextCache(capacity=1, time_to_live=5)
>>> integrator = openmm.VerletIntegrator(1.0*unit.femtosecond)
>>> context1, integrator1 = context_cache.get_context(thermo_state,
...                                                   copy.deepcopy(integrator))
>>> context2, integrator2 = context_cache.get_context(thermo_state_npt,
...                                                   copy.deepcopy(integrator))
>>> len(context_cache)
1

In the example above, the maximum capacity of the cache is 1, so the first context is deallocated to make space for the second Context created with the incompatible thermodynamic state.

Finally, you can force the ContextCache to create contexts on a specific platform.

platform = openmm.Platform.getPlatformByName('Reference')
context_cache = cache.ContextCache(platform=platform)

The global ContextCache

The openmmtools.cache module exposes a global variable that provides a shared ContextCache for all the classes in the framework.

cache.global_context_cache.platform = openmm.Platform.getPlatformByName('CPU')
cache.global_context_cache.capacity = 2
cache.global_context_cache.time_to_live = 10
verlet_integrator = openmm.VerletIntegrator(1.0*unit.femtosecond)
context, integrator = cache.global_context_cache.get_context(thermo_state,
                                                             verlet_integrator)

Usually, you’ll want to create a Context using the global_context_cache to minimize the number of created contexts overall. This is, for example, the context cache used by default by all the MCMCMove objects internally, which we’ll touch shortly.

Extending ThermodynamicState to control arbitrary parameters

It is possible to extend the ThermodynamicState to manipulate other thermodynamic parameters of the System through the states.CompoundThermodynamicState class and one or more composable states. An example may clarify this. Remember the alchemy.AlchemicalState class we discussed above? AlchemicalState is a composable state.

# Prepare T4-Lysozyme + p-xylene system for alchemical perturbation.
factory = alchemy.AbsoluteAlchemicalFactory()
alchemical_region = alchemy.AlchemicalRegion(alchemical_atoms=pxylene_atoms)
alchemical_system = factory.create_alchemical_system(t4_system, alchemical_region)

# Define the basic thermodynamic state of the system.
thermo_state = states.ThermodynamicState(alchemical_system, temperature=298*unit.kelvin)

# Extend the definition of thermodynamic state to consider alchemical parameters as well.
alchemical_state = alchemy.AlchemicalState.from_system(alchemical_system)
compound_state = states.CompoundThermodynamicState(thermodynamic_state=thermo_state,
                                                   composable_states=[alchemical_state])

At this point, compound_state is both a ThermodynamicState and an AlchemicalState in the sense that it exposes the interface to modify the thermodynamic parameters controlled by both objects.

context = compound_state.create_context(integrators.LangevinIntegrator())
compound_state.temperature = 350*unit.kelvin  # Increase temperature of simulation.
compound_state.lambda_sterics = 0.2  # Soften torsions.
compound_state.apply_to_context(context)

Obviously, CompoundThermodynamicState is not compatible exclusively with AlchemicalState but with any object implementing the states.IComposableState interface. A quick way to define your own composable state is described in the developer’s tutorial.

The power of this abstraction will become evident when we’ll implement a simple replica-exchange algorithm at the end of the tutorial.


MCMC framework

The Markov chain Monte Carlo (MCMC) framework implemented in the mcmc module take advantage of the thermodynamic state objects described above to provide an easy way to experiment with different propagation schemes mixing Monte Carlo moves and dynamics.

Basic usage

The basic object in the module is the mcmc.MCMCMove abstract class that provides a common interface for both integrators and Monte Carlo to propagate the state of the system.

from openmmtools import mcmc

# Define the thermodynamic state of the T4-Lysozyme + p-xylene system
thermo_state = states.ThermodynamicState(t4_system, temperature=300*unit.kelvin)

# Create a SamplerState system holding the coordinates of the system.
sampler_state = states.SamplerState(positions=lysozyme_pxylene.positions)

# Propagate the system with a GHMC integrator.
ghmc_move = mcmc.GHMCMove(timestep=1.0*unit.femtosecond, n_steps=n_steps)
ghmc_move.apply(thermo_state, sampler_state)

The SamplerState object in the snippet above holds the configurational degrees of freedom of the System (e.g., positions, velocities, and eventually box vectors). The sampler state is updated by MCMCMove.apply to hold the coordinates and velocities after 1000 steps of GHMC integration. Note however that, in princple, the framework allows an MCMCMove to change also the thermodynamic degrees of freedom in thermo_state.

OpenMM integrators as MCMCMoves

The mcmc module provides a few integrators in the form of an MCMCMove, including openmmtools.integrators.LangevinIntegrator. Casting integrators in the form of an MCMCMove object makes it easy to combine them with Monte Carlo techniques. Moreover, integrator ``MCMCMove``s provide a few extra features such as automatic recovery after a NaN.

langevin_move = mcmc.LangevinSplittingDynamicsMove(splitting='V R O R V',
                                                   n_steps=n_steps,
                                                   n_restart_attempts=5)
langevin_move.apply(thermo_state, sampler_state)

Propagating your system through Langevin dynamics has always a non-zero probability of incurring into a NaN error. When this happens, instead of crashing, the Langevin move above will restore the state of the System before integrating and try again, relying on the stochastic component of the propagation to obtain a different solution. This is repeated to a maximum of 5 times before giving up and throwing an error. The raised exception exposes a method to serialize the simulation objects automatically for further debugging.

try:
    langevin_move.apply(thermo_state, sampler_state)
except mcmc.IntegratorMoveError as e:
    # This saves to disk the OpenMM System, Integrator, and State objects.
    e.serialize_error(path_files_prefix='debug/langevin')

When a NaN occurr, the code above serializes the OpenMM System, Integrator, and State objects on disk at debug/langevin-system.xml, debug/langevin-integrator.xml, and debug/langevin-state.xml respectively.

This feature can easily be extended to other integrators that are not explicitly provided in the mcmc module.

integrator = integrators.HMCIntegrator(timestep=1.0*unit.femtosecond)
HMC_move = mcmc.IntegratorMove(integrator, n_steps=n_steps, n_restart_attempts=4)

Combining Monte Carlo and dynamics

Combining and mixing multiple MCMCMove is usually performed through the mcmc.SequenceMove object

sequence_move = mcmc.SequenceMove(move_list=[
    mcmc.MCDisplacementMove(atom_subset=pxylene_atoms),
    mcmc.MCRotationMove(atom_subset=pxylene_atoms),
    mcmc.LangevinSplittingDynamicsMove(timestep=2.0*unit.femtoseconds, n_steps=n_steps,
                                       reassign_velocities=True, n_restart_attempts=6)
])

sequence_move.apply(thermo_state, sampler_state)

The MCMCMove above performs in sequence a Metropolized Monte Carlo rigid translation and rotation of the p-xylene molecule followed by Langevin dynamics after randomizing the velocities according to the Boltzmann distribution at the temperature of thermo_state.

ContextCache and Platform with MCMCMoves

All MCMCMove objects implemented in OpenMMTools accept a context_cache in the constructor. This parameter defaults to mmtools.cache.global_context_cache, but you can pass a local cache to trigger other behaviors.

local_cache = cache.ContextCache(platform=openmm.Platform.getPlatformByName('CPU'))
dummy_cache = cache.DummyContextCache()  # Disable caching.
move = mcmc.SequenceMove(move_list=[
    mcmc.MCDisplacementMove(atom_subset=pxylene_atoms, context_cache=local_cache),
    mcmc.MCRotationMove(atom_subset=pxylene_atoms, context_cache=dummy_cache),
    mcmc.LangevinSplittingDynamicsMove(n_steps=n_steps)  # Uses global_context_cache.
])

In the example above, applying the move will perform an MC translation of the ligands atom using a local ContextCache that runs on the CPU, then an MC rotation using the DummyContextCache, which recreates context every time effectively deactivating caching, and finally propagates the system with Langevin dynamics using the global cache on the fastest platform available.


Example: A minimal implementation of a general replica-exchange simulation class

Our most recent enhanced-sampling facilities are currently hosted in YANK, and they are still waiting to be moved to OpenMMTools. However, the following minimal implementation of a replica exchange simulation class should give you an idea of what is possible to do when taking advantage of the full framework.

import math
from random import random, randint

class ReplicaExchange:

    def __init__(self, thermodynamic_states, sampler_states, mcmc_move):
        self._thermodynamic_states = thermodynamic_states
        self._replicas_sampler_states = sampler_states
        self._mcmc_move = mcmc_move

    def run(self, n_iterations=1):
        for iteration in range(n_iterations):
            self._mix_replicas()
            self._propagate_replicas()

    def _propagate_replicas(self):
        # _thermodynamic_state[i] is associated to the replica configuration in _replicas_sampler_states[i].
        for thermo_state, sampler_state in zip(self._thermodynamic_states, self._replicas_sampler_states):
            self._mcmc_move.apply(thermo_state, sampler_state)

    def _mix_replicas(self, n_attempts=1):
        # Attempt to switch two replicas at random. Obviously, this scheme can be improved.
        for attempt in range(n_attempts):
            # Select two replicas at random.
            i = randint(0, len(self._thermodynamic_states)-1)
            j = randint(0, len(self._thermodynamic_states)-1)
            sampler_state_i, sampler_state_j = (self._replicas_sampler_states[k] for k in [i, j])
            thermo_state_i, thermo_state_j = (self._thermodynamic_states[k] for k in [i, j])

            # Compute the energies.
            energy_ii = self._compute_reduced_potential(sampler_state_i, thermo_state_i)
            energy_jj = self._compute_reduced_potential(sampler_state_j, thermo_state_j)
            energy_ij = self._compute_reduced_potential(sampler_state_i, thermo_state_j)
            energy_ji = self._compute_reduced_potential(sampler_state_j, thermo_state_i)

            # Accept or reject the swap.
            log_p_accept = - (energy_ij + energy_ji) + energy_ii + energy_jj
            if log_p_accept >= 0.0 or random() < math.exp(log_p_accept):
                # Swap states in replica slots i and j.
                self._thermodynamic_states[i] = thermo_state_j
                self._thermodynamic_states[j] = thermo_state_i

    def _compute_reduced_potential(self, sampler_state, thermo_state):
        # Obtain a Context to compute the energy with OpenMM. Any integrator will do.
        context, integrator = cache.global_context_cache.get_context(thermo_state)
        # Compute the reduced potential of the sampler_state configuration
        # in the given thermodynamic state.
        sampler_state.apply_to_context(context)
        return thermo_state.reduced_potential(context)

The first observation is that the bulk of the code complexity lies in the replica swapping code, while most of the other details are handled by the specialized classes of the framework. From a software engineering perspective, this is a good sign as it is compatible with the single responsibility principle.

Secondly, the class can be used to implement a variety of algorithm. A few examples follow.

Parallel tempering

To run a parallel tempering simulation, we just have initialize the ReplicaExchange object with a list of thermodynamic states that vary in temperature. You can make use of the utility function states.create_thermodynamic_state_protocol to initialize efficiently a list of ThermodynamicState or CompoundThermodynamicState.

>>> # Initialize thermodynamic states at different temperatures.
>>> host_guest = testsystems.HostGuestVacuum()
>>> protocol = {'temperature': [300, 310, 330, 370, 450] * unit.kelvin}
>>> thermo_states = states.create_thermodynamic_state_protocol(host_guest.system, protocol)

>>> # Initialize replica initial configurations.
>>> sampler_states = [states.SamplerState(positions=host_guest.positions)
...                   for _ in thermo_states]

>>> # Propagate the replicas with Langevin dynamics.
>>> langevin_move = mcmc.LangevinSplittingDynamicsMove(timestep=2.0*unit.femtosecond,
...                                                    n_steps=n_steps)

>>> # Run the parallel tempering simulation.
>>> parallel_tempering = ReplicaExchange(thermo_states, sampler_states, langevin_move)
>>> parallel_tempering.run()

This example creates 5 replicas starting from the same configurations but at the temperatures of 300, 310, …, 450 K, and propagates the system with Langevin dynamics.

Hamiltonian replica exchange + parallel tempering

Let’s say we want to implement an enhanced sampling scheme that increases the temperature while alchemically softening part of a system.

>>> # Prepare the host-guest system for alchemical modification.
>>> guest_atoms = host_guest.mdtraj_topology.select('resname B2')
>>> alchemical_region = alchemy.AlchemicalRegion(alchemical_atoms=guest_atoms)
>>> factory = alchemy.AbsoluteAlchemicalFactory()
>>> alchemical_system = factory.create_alchemical_system(host_guest.system, alchemical_region)

>>> # Initialize compound thermodynamic states at different temperatures and alchemical states.
>>> protocol = {'temperature': [300, 310, 330, 370, 450] * unit.kelvin,
...             'lambda_electrostatics': [1.0, 0.5, 0.0, 0.0, 0.0],
...             'lambda_sterics': [1.0, 1.0, 1.0, 0.5, 0.0]}
>>> alchemical_state = alchemy.AlchemicalState.from_system(alchemical_system)
>>> compound_states = states.create_thermodynamic_state_protocol(
...     alchemical_system, protocol=protocol, composable_states=[alchemical_state])

>>> # Run the combined Hamiltonian replica exchange + parallel tempering simulation.
>>> hrex_tempering = ReplicaExchange(compound_states, sampler_states, langevin_move)

Hamiltonian replica exchange + parallel tempering mixing Monte Carlo and dynamics

Finally, let’s mix Monte Carlo and dynamics for propagation.

>>> sequence_move = mcmc.SequenceMove(move_list=[
...     mcmc.MCDisplacementMove(atom_subset=guest_atoms),
...     mcmc.MCRotationMove(atom_subset=guest_atoms),
...     mcmc.LangevinSplittingDynamicsMove(timestep=2.0*unit.femtoseconds, n_steps=n_steps,
...                                        reassign_velocities=True, n_restart_attempts=6)
... ])

>>> # Run the combined Hamiltonian replica exchange + parallel tempering simulation
>>> # using a combination of Monte Carlo moves and Langevin dynamics.
>>> hrex_tempering = ReplicaExchange(compound_states, sampler_states, sequence_move)