Run a REST simulation

Replica exchange solute tempering (REST) is an enhanced sampling approach that aims to improve sampling for a user-defined subset of atoms (i.e., REST atoms) by increasing its effective temperature.

REST involves running replica exchange to sample a set of thermodynamic states, where one state contains the unmodified Hamiltonian and all other states involve modified Hamiltonians where the REST atom interactions are scaled to increased temperatures (Wang et al. 2011, DOI: 10.1021/jp204407d, available at: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3172817/).

In this tutorial, we show how to 1) build a REST-capable OpenMM system and then 2) run REST (using OpenMMTools (https://openmmtools.readthedocs.io) objects).

Overview of the tutorial:

Build a REST-capable system - We first build a “vanilla” OpenMM System, which represents the unmodified Hamiltonian from which we will construct the modified REST Hamiltonian (i.e., the REST system). - We then define the REST atoms, i.e., the atoms whose interactions will be heated up. - We begin creating the REST system by initializing an empty OpenMM System and copying the particles, barostat, box vectors, and constraints from the vanilla system. - We then create custom forces for each valence force in the REST system: CustomBondForce, CustomAngleForce, CustomTorsionForce, where each force has a custom expression that allow the energy to be scaled according to REST. Each custom expression utilizes a global parameter to control the REST scale factor and per-bond (or per-angle or per-torsion) parameters to keep track of parameters needed to compute the energy for each bond (or angle or torsion). - For the REST system’s nonbonded force, we create a NonbondedForce and use offsets to scale the energies for interactions involving REST atoms. We use the NonbondedForce (and not custom forces) for simplicity. If we used custom forces, we would need to use a CustomNonbondedForce for nonbonded interactions and a CustomBondForce for nonbonded exceptions, and these forces can be tricky to work with. - Finally, we copy the bonds, angles, torsions, particles, and nonbonded exceptions from each force in the vanilla system to each new force in the REST system.

Run REST - We define a class called RESTState that creates a “state” (aka version of the REST system) with a particular set of global parameters. - We choose the temperatures we’d like use for REST scaling and then create thermodynamic states corresponding to each temperature using the previously defined RESTState class. For each thermodynamic state, we minimize the positions with respect to the modified Hamiltonian. - Finally, we create a replica exchange sampler that samples each of the thermodynamic states and use the sampler to run the replica exchange simulation.

Note that this tutorial is inspired by the RESTCapableHybridTopologyFactory in Perses (https://github.com/choderalab/perses), which can be used to run relative alchemical free energy calculations.

01 - Construct a vanilla OpenMM system

We first load a PDB file consisting of villin in water and build a System from it, which we refer to as the “vanilla system”. The villin PDB file is present in the current directory (openmm-cookbook/notebooks/cookbook/), but can be also be downloaded here: https://raw.githubusercontent.com/openmm/openmm-cookbook/main/notebooks/cookbook/villin.pdb

[3]:
import openmm
from openmm import app, unit
[4]:
pdb = app.PDBFile('villin.pdb')
forcefield = app.ForceField('amber14-all.xml', 'amber14/tip3pfb.xml')
system = forcefield.createSystem(pdb.topology,
                                 nonbondedMethod=app.PME,
                                 nonbondedCutoff=1*unit.nanometer,
                                 constraints=app.HBonds,
                                 removeCMMotion=False,
                                 hydrogenMass=3*unit.amu)
barostat = openmm.MonteCarloBarostat(1*unit.atmosphere, 300*unit.kelvin, 50)
system.addForce(barostat)
[4]:
4

We check the forces in the vanilla system to make sure they’re as expected.

[5]:
system.getForces()
[5]:
[<openmm.openmm.HarmonicBondForce; proxy of <Swig Object of type 'OpenMM::HarmonicBondForce *' at 0x11fae7450> >,
 <openmm.openmm.NonbondedForce; proxy of <Swig Object of type 'OpenMM::NonbondedForce *' at 0x11fae7e70> >,
 <openmm.openmm.PeriodicTorsionForce; proxy of <Swig Object of type 'OpenMM::PeriodicTorsionForce *' at 0x11fae7f60> >,
 <openmm.openmm.HarmonicAngleForce; proxy of <Swig Object of type 'OpenMM::HarmonicAngleForce *' at 0x11fae7300> >,
 <openmm.openmm.MonteCarloBarostat; proxy of <Swig Object of type 'OpenMM::MonteCarloBarostat *' at 0x11fae72d0> >]

02 - Construct a REST-capable OpenMM system

[6]:
import copy

We now choose the REST atoms to be all atoms in one residue (residue index 3) of villin.

[7]:
res = list(pdb.topology.residues())[3]
rest_atoms = [atom.index for atom in res.atoms()]

There are 15 REST atoms.

[8]:
len(rest_atoms)
[8]:
15

We create a new System (which we refer to as the “REST system”) and copy over the particles, barostat, box vectors, and constraints from the vanilla system.

[9]:
# Create REST system
rest_system = openmm.System()

# Create dict of vanilla system forces (for easy retrieval of force objects later)
system_forces = {type(force).__name__ : force for force in system.getForces()}

# Add particles
for particle_idx in range(system.getNumParticles()):
    particle_mass = system.getParticleMass(particle_idx)
    rest_system.addParticle(particle_mass)

# Copy barostat
if "MonteCarloBarostat" in system_forces:
    barostat = copy.deepcopy(system_forces["MonteCarloBarostat"])
    rest_system.addForce(barostat)

# Copy box vectors
box_vectors = system.getDefaultPeriodicBoxVectors()
rest_system.setDefaultPeriodicBoxVectors(*box_vectors)

# Copy constraints
for constraint_idx in range(system.getNumConstraints()):
    atom1, atom2, length = system.getConstraintParameters(constraint_idx)
    rest_system.addConstraint(atom1, atom2, length)

We create the REST system’s bond force as a CustomBondForce which uses a custom energy expression that allows for REST scaling. We add a global parameter (lambda_rest_bonds) to act as the REST scale factor for bonds, i.e., to control the extent to which the bonds involving REST atoms are scaled.

We add per-bond parameters that are booleans indicating whether each bond is: - REST (if all atoms are REST atoms, is_rest is 1, otherwise 0) - inter (if at least one, but not all atoms are REST atoms, is_inter is 1, otherwise 0), or - non-REST (if no atoms are REST atoms, is_nonrest is 1, otherwise 0).

We also add per-bond parameters that keep track of the length and spring constant for each bond.

[10]:
# Define the custom expression
bond_expression = "rest_scale * (K / 2) * (r - length)^2;"
bond_expression += "rest_scale = is_rest * lambda_rest_bonds * lambda_rest_bonds " \
                   "+ is_inter * lambda_rest_bonds " \
                   "+ is_nonrest;"

# Create custom force
rest_bond_force = openmm.CustomBondForce(bond_expression)
rest_system.addForce(rest_bond_force)

# Add global parameters
rest_bond_force.addGlobalParameter("lambda_rest_bonds", 1.0)

# Add per-bond parameters for rest scaling
rest_bond_force.addPerBondParameter("is_rest")
rest_bond_force.addPerBondParameter("is_inter")
rest_bond_force.addPerBondParameter("is_nonrest")

# Add per-bond parameters for defining bond energy
rest_bond_force.addPerBondParameter('length')  # equilibrium bond length
rest_bond_force.addPerBondParameter('K')  # force constant
[10]:
4

We next want to add bonds to the REST system’s bond force. To do so, we need to determine whether each bond should be considered REST, inter, or non-REST, so that we can set each of the following per-bond parameters: is_rest, is_inter, is_nonrest. We define a function get_rest_identifier to do this.

[11]:
def get_rest_identifier(atoms, rest_atoms):
    """
    For a given atom or set of atoms, get the rest_id which is a list of binary ints that defines which
    (mutually exclusive) set the atom(s) belong to.

    If there is a single atom, the sets are: is_rest, is_nonrest
    If there is a set of atoms, the sets are: is_rest, is_inter, is_nonrest

    Example: if there is a single atom that is in the nonrest set, the rest_id is [0, 1]

    Arguments
    ---------
    atoms : set or int
        a set of hybrid atom indices or single atom
    rest_atoms : set or list
        a list (or list-like) of atoms whose interactions will be scaled by REST
    Returns
    -------
    rest_id : list
        list of binaries indicating which set the atom(s) belong to
    """

    if isinstance(atoms, int):
        rest_id = [0, 1] # Set the default rest_id to non-REST
        if atoms in rest_atoms:
            rest_id = [1, 0]
        return rest_id

    elif isinstance(atoms, set):
        rest_id = [0, 0, 1] # Set the default rest_id to non-REST
        if atoms.intersection(rest_atoms) != set(): # At least one of the atoms is REST
            if atoms.issubset(rest_atoms): # All atoms are REST
                rest_id = [1, 0, 0]
            else: # At least one (but not all) of the atoms is are REST
                rest_id = [0, 1, 0]
        return rest_id

    else:
        raise Exception(f"atoms is of type {type(atoms)}, but only `int` and `set` are allowable")

Now we can add the bonds to the REST system’s bond force. We copy the periodic boundary conditions and bonds from the vanilla system’s bond force to the REST system’s bond force. Each bond is defined by the indices of the two atoms that form the bond (p1, p1), the bond length (r0), and the force constant (k). Each bond also has a rest_id (determined using get_rest_identifier(), defined above), which indicates whether each bond should be considered REST, inter, or non-REST.

[12]:
# Get vanilla system bond force
bond_force = system_forces['HarmonicBondForce']

# Set periodicity
if bond_force.usesPeriodicBoundaryConditions():
    rest_bond_force.setUsesPeriodicBoundaryConditions(True)

# Add bonds to rest_system
for term_idx in range(bond_force.getNumBonds()):
    # Get the bond parameters and rest id
    p1, p2, r0, k = bond_force.getBondParameters(term_idx)
    idx_set = set([p1, p2])
    rest_id = get_rest_identifier(idx_set, rest_atoms)

    # Add the bond
    bond_term = (p1, p2, rest_id + [r0, k])
    rest_bond_force.addBond(*bond_term)

We then repeat the same process to 1) construct the REST system’s angle force (i.e., a CustomAngleForce) and 2) copy the angles from the vanilla system’s angle force to the REST system’s angle force.

[13]:
# Define the custom expression
angle_expression = "rest_scale * (K / 2) * (theta - theta0)^2;"
angle_expression += "rest_scale = is_rest * lambda_rest_angles * lambda_rest_angles " \
                    "+ is_inter * lambda_rest_angles " \
                    "+ is_nonrest;"

# Create custom force
rest_angle_force = openmm.CustomAngleForce(angle_expression)
rest_system.addForce(rest_angle_force)

# Add global parameters
rest_angle_force.addGlobalParameter("lambda_rest_angles", 1.0)

# Add per-angle parameters for rest scaling
rest_angle_force.addPerAngleParameter("is_rest")
rest_angle_force.addPerAngleParameter("is_inter")
rest_angle_force.addPerAngleParameter("is_nonrest")

# Add per-angle parameters for defining angle energy
rest_angle_force.addPerAngleParameter('theta0')  # equilibrium angle
rest_angle_force.addPerAngleParameter('K')  # force constant

# Get vanilla system angle force
angle_force = system_forces['HarmonicAngleForce']

# Set periodicity
if angle_force.usesPeriodicBoundaryConditions():
    rest_angle_force.setUsesPeriodicBoundaryConditions(True)

# Add angles to rest_system
for term_idx in range(angle_force.getNumAngles()):
    # Get the angle parameters and rest id
    p1, p2, p3, theta0, k = angle_force.getAngleParameters(term_idx)
    idx_set = set([p1, p2, p3])
    rest_id = get_rest_identifier(idx_set, rest_atoms)

    # Add the angle
    angle_term = (p1, p2, p3, rest_id + [theta0, k])
    rest_angle_force.addAngle(*angle_term)

We next repeat the same process to 1) construct the REST system’s torsion force (i.e., a PeriodicTorsionForce) and 2) copy the torsions from the vanilla system’s torsion force to the REST system’s torsion force.

[14]:
# Define the custom expression
torsion_expression = "rest_scale * U;"
torsion_expression += "rest_scale = is_rest * lambda_rest_torsions * lambda_rest_torsions " \
                      "+ is_inter * lambda_rest_torsions " \
                      "+ is_nonrest;"
torsion_expression += "U = (K * (1 + cos(periodicity * theta - phase)));"

# Create custom force
rest_torsion_force = openmm.CustomTorsionForce(torsion_expression)
rest_system.addForce(rest_torsion_force)

# Add global parameters
rest_torsion_force.addGlobalParameter("lambda_rest_torsions", 1.0)

# Add per-torsion parameters for rest scaling
rest_torsion_force.addPerTorsionParameter("is_rest")
rest_torsion_force.addPerTorsionParameter("is_inter")
rest_torsion_force.addPerTorsionParameter("is_nonrest")

# Add per-torsion parameters for defining torsion energy
rest_torsion_force.addPerTorsionParameter('periodicity')
rest_torsion_force.addPerTorsionParameter('phase') # phase offset
rest_torsion_force.addPerTorsionParameter('K') # force constant

# Get vanilla system torsion force
torsion_force = system_forces['PeriodicTorsionForce']

# Set periodicity
if torsion_force.usesPeriodicBoundaryConditions():
    rest_torsion_force.setUsesPeriodicBoundaryConditions(True)

# Add torsions to rest_system
for torsion_idx in range(torsion_force.getNumTorsions()):
    # Get the torsion parameters and rest id
    p1, p2, p3, p4, periodicity, phase, K = torsion_force.getTorsionParameters(torsion_idx)
    idx_set = set([p1, p2, p3, p4])
    rest_id = get_rest_identifier(idx_set, rest_atoms)

    # Add torsion
    torsion_term = (p1, p2, p3, p4, rest_id + [periodicity, phase, K])
    rest_torsion_force.addTorsion(*torsion_term)

Finally, we construct the REST system’s nonbonded force (i.e., a NonbondedForce) by copying the nonbonded method and related parameters from the vanilla system’s nonbonded force. We add two global parameters (lambda_rest_electrostatics and lambda_rest_sterics) to act as scale factors for nonbonded interactions involving REST atoms.

[15]:
# Create nonbonded force
rest_nonbonded_force = openmm.NonbondedForce()
rest_system.addForce(rest_nonbonded_force)

# Get vanilla system nonbonded force
nonbonded_force = system_forces['NonbondedForce']

# Set the nonbonded method and related parameters
nonbonded_method = nonbonded_force.getNonbondedMethod()
rest_nonbonded_force.setNonbondedMethod(nonbonded_method)
if nonbonded_method != openmm.NonbondedForce.NoCutoff:
    epsilon_solvent = nonbonded_force.getReactionFieldDielectric()
    cutoff = nonbonded_force.getCutoffDistance()
    rest_nonbonded_force.setReactionFieldDielectric(epsilon_solvent)
    rest_nonbonded_force.setCutoffDistance(cutoff)
if nonbonded_method in [openmm.NonbondedForce.PME, openmm.NonbondedForce.Ewald]:
    [alpha_ewald, nx, ny, nz] = nonbonded_force.getPMEParameters()
    delta = nonbonded_force.getEwaldErrorTolerance()
    rest_nonbonded_force.setPMEParameters(alpha_ewald, nx, ny, nz)
    rest_nonbonded_force.setEwaldErrorTolerance(delta)

# Copy switching function from vanilla system
switch_bool = nonbonded_force.getUseSwitchingFunction()
rest_nonbonded_force.setUseSwitchingFunction(switch_bool)
if switch_bool:
    switching_distance = nonbonded_force.getSwitchingDistance()
    rest_nonbonded_force.setSwitchingDistance(switching_distance)

# Copy dispersion correction
dispersion_bool = nonbonded_force.getUseDispersionCorrection()
rest_nonbonded_force.setUseDispersionCorrection(dispersion_bool)

# Add global parameters
rest_nonbonded_force.addGlobalParameter('lambda_rest_electrostatics', 0.)
rest_nonbonded_force.addGlobalParameter('lambda_rest_sterics', 0.)


[15]:
1

We then copy the nonbonded particles and exceptions from the vanilla system’s nonbonded force to the rest system’s nonbonded force. We add offsets to the particles and exceptions in order to apply REST scaling. The offsets for electrostatics and sterics depend on lambda_rest_electrostatics and lambda_rest_sterics, respectively. The offsets for REST and inter exceptions conveniently depend on lambda_rest_sterics and lambda_rest_electrostatics, respectively (derivation available in the Appendix).

[16]:
# Add nonbondeds to rest_system
for particle_idx in range(nonbonded_force.getNumParticles()):
    # Get the nonbonded parameters and rest id
    q, sigma, epsilon = nonbonded_force.getParticleParameters(particle_idx)
    rest_id = get_rest_identifier(particle_idx, rest_atoms)

    # Add particles and offsets
    if rest_id == [0, 1]: # nonrest
        rest_nonbonded_force.addParticle(q, sigma, epsilon)

    else: # rest
        rest_nonbonded_force.addParticle(q, sigma, epsilon)
        rest_nonbonded_force.addParticleParameterOffset('lambda_rest_electrostatics', particle_idx, q, 0.0*sigma, epsilon*0.0)
        rest_nonbonded_force.addParticleParameterOffset('lambda_rest_sterics', particle_idx, q*0.0, 0.0*sigma, epsilon)

# Handle exceptions
for exception_idx in range(nonbonded_force.getNumExceptions()):
    # Get exception parameters and rest id
    p1, p2, chargeProd, sigma, epsilon = nonbonded_force.getExceptionParameters(exception_idx)
    idx_set = set([p1, p2])
    rest_id = get_rest_identifier(idx_set, rest_atoms)

    # Add exceptions and offsets
    exc_idx = rest_nonbonded_force.addException(p1, p2, chargeProd, sigma, epsilon)
    if rest_id == [0, 0, 1]: # nonrest
        pass

    elif rest_id == [1, 0, 0]: # rest
        rest_nonbonded_force.addExceptionParameterOffset('lambda_rest_sterics', exc_idx, chargeProd, 0.0*sigma, epsilon)

    elif rest_id == [0, 1, 0]: # inter
        rest_nonbonded_force.addExceptionParameterOffset('lambda_rest_electrostatics', exc_idx, chargeProd, 0.0*sigma, epsilon)

We’ve now finished building a REST-capable OpenMM system for villin, rest_system!

03 - Run replica exchange

This section requires openmmtools

[17]:
!mamba install -qy -c conda-forge openmmtools
[18]:
import math
import logging
import numpy as np
from openmmtools.constants import kB
from openmmtools import cache, mcmc, multistate
from openmmtools.multistate import ReplicaExchangeSampler
from openmmtools.states import GlobalParameterState, SamplerState, ThermodynamicState, CompoundThermodynamicState

Now, we would like to use rest_system to run REST, which consists of using replica exchange to sample a number of thermodynamic states, each of which has a different REST-modified Hamiltonian. To construct each of the thermodynamic states, we define a class called RESTState that creates a “state” (aka version of the system) with a particular set of global parameters (lambda_rest_bonds, lambda_rest_angles, lambda_rest_torsions, lambda_rest_electrostatics, lambda_rest_sterics). To set the global parameters for each thermodynamic state, we just need to supply beta_m (proportional to the temperature m we wish to scale to) and beta_0 (proportional to the temperature of distribution of interest), where beta = \(\frac{1}{k_b T}\). See Wang et al. 2011 (DOI: 10.1021/jp204407d) for more details.

Given beta_m and beta_0, the RESTState takes care of actually computing the value for each global parameter (based on lambda functions). For lambda_rest_bonds, lambda_rest_angles, lambda_rest_torsions, the value will be computed according to \(\sqrt{\frac{\beta_m}{\beta_0}}\). For lambda_rest_electrostatics, the value will be computed as \(\sqrt{\frac{\beta_m}{\beta_0}} - 1\). For lambda_rest_sterics, the value will be computed as \(\frac{\beta_m}{\beta_0} - 1\). See Appendix for derivation of the functions for lambda_rest_electrostatics and lambda_rest_sterics.

[19]:
class RESTState(GlobalParameterState):
    lambda_rest_bonds = GlobalParameterState.GlobalParameter('lambda_rest_bonds', standard_value=1.0)
    lambda_rest_angles = GlobalParameterState.GlobalParameter('lambda_rest_angles', standard_value=1.0)
    lambda_rest_torsions = GlobalParameterState.GlobalParameter('lambda_rest_torsions', standard_value=1.0)
    lambda_rest_electrostatics = GlobalParameterState.GlobalParameter('lambda_rest_electrostatics', standard_value=0.0)
    lambda_rest_sterics = GlobalParameterState.GlobalParameter('lambda_rest_sterics', standard_value=0.0)

    def set_rest_parameters(self, beta_m, beta_0):
        """Set all defined lambda parameters to the given value.

        The undefined parameters (i.e. those being set to None) remain undefined.

        Parameters
        ----------
        new_value : float
            The new value for all defined parameters.
        """
        lambda_functions = {'lambda_rest_bonds': lambda beta_m, beta_0 : np.sqrt(beta_m / beta_0),
                 'lambda_rest_angles' : lambda beta_m, beta_0 : np.sqrt(beta_m / beta_0),
                 'lambda_rest_torsions' : lambda beta_m, beta_0 : np.sqrt(beta_m / beta_0),
                 'lambda_rest_electrostatics' : lambda beta_m, beta_0 : np.sqrt(beta_m / beta_0) - 1,
                 'lambda_rest_sterics' : lambda beta_m, beta_0 : beta_m / beta_0 - 1
                 }

        for parameter_name in self._parameters:
            if self._parameters[parameter_name] is not None:
                new_value = lambda_functions[parameter_name](beta_m, beta_0)
                setattr(self, parameter_name, new_value)

We next determine the number of replicas and states (both are 12) and choose the T_min (the minimum temperature, i.e., the temperature of the desired distribution) and T_max (the maximum temperature to be scaled to during REST) to be 300 and 600 Kelvin, respectively. We set the temperatures of each thermodynamic state to be exponentially spaced between 300 and 600 Kelvin.

[20]:
# Set temperatures for each thermodynamic state
n_replicas = 12  # Number of temperature replicas
T_min = 300 * unit.kelvin  # Minimum temperature (i.e., temperature of desired distribution)
T_max = 600 * unit.kelvin  # Maximum temperature
temperatures = [T_min + (T_max - T_min) * (math.exp(float(i) / float(n_replicas-1)) - 1.0) / (math.e - 1.0)
                for i in range(n_replicas)]


We now create a reference thermodynamic state (i.e., RESTState) from the rest_system and use that state to create OpenMMTools ThermodynamicState and CompoundThermodynamicState objects which will be used with other OpenMMTools objects below to run replica exchange.

[21]:
# Create reference thermodynamic state
rest_state = RESTState.from_system(rest_system)
thermostate = ThermodynamicState(rest_system, temperature=T_min)
compound_thermodynamic_state = CompoundThermodynamicState(thermostate, composable_states=[rest_state])

Using the reference thermodynamic state, we create a list of 12 thermodynamic states, one for each temperature in temperatures. Each of these thermodynamic states also has an associated OpenMMTools SamplerState object which contains the corresponding box vectors and minimized positions.

[22]:
# Create thermodynamic states
sampler_state =  SamplerState(pdb.positions, box_vectors=rest_system.getDefaultPeriodicBoxVectors())
beta_0 = 1/(kB*T_min)
thermodynamic_state_list = []
sampler_state_list = []
for temperature in temperatures:
    # Create a thermodynamic state with REST interactions scaled to the given temperature
    beta_m = 1/(kB*temperature)
    compound_thermodynamic_state_copy = copy.deepcopy(compound_thermodynamic_state)
    compound_thermodynamic_state_copy.set_rest_parameters(beta_m, beta_0)
    thermodynamic_state_list.append(compound_thermodynamic_state_copy)

    # Generate a sampler_state with minimized positions
    context, integrator = cache.global_context_cache.get_context(compound_thermodynamic_state_copy)
    sampler_state.apply_to_context(context, ignore_velocities=True)
    openmm.LocalEnergyMinimizer.minimize(context)
    sampler_state.update_from_context(context)
    sampler_state_list.append(copy.deepcopy(sampler_state))

We now create OpenMMTools objects to run replica exchange. We instantiate a LangevinDynamicsMove which tells the ReplicaExchangeSampler to run 1 ps of molecular dynamics (here, Langevin dynamics) between swap attempts. We instantiate the ReplicaExchangeSampler with 100 iterations (i.e., swap attempts). We then initialize a MultiStateReporter to save critical information from the replica exchange simulation, e.g., reduced potentials, positions, box vectors, etc. And finally, we run the REST simulation!

[23]:
# Set up sampler
_logger = logging.getLogger()
_logger.setLevel(logging.DEBUG)
move = mcmc.LangevinDynamicsMove(timestep=4*unit.femtoseconds, n_steps=250) # each move is 1 ps
simulation = ReplicaExchangeSampler(mcmc_moves=move, number_of_iterations=100)

# Run repex
reporter_file = "test.nc"
reporter = multistate.MultiStateReporter(reporter_file, checkpoint_interval=100)
simulation.create(thermodynamic_states=thermodynamic_state_list,
                  sampler_states=sampler_state_list,
                  storage=reporter)
simulation.run()
Warning: The openmmtools.multistate API is experimental and may change in future releases
/bin/sh: nvidia-smi: command not found
Warning: The openmmtools.multistate API is experimental and may change in future releases
Please cite the following:

        Friedrichs MS, Eastman P, Vaidyanathan V, Houston M, LeGrand S, Beberg AL, Ensign DL, Bruns CM, and Pande VS. Accelerating molecular dynamic simulations on graphics processing unit. J. Comput. Chem. 30:864, 2009. DOI: 10.1002/jcc.21209
        Eastman P and Pande VS. OpenMM: A hardware-independent framework for molecular simulations. Comput. Sci. Eng. 12:34, 2010. DOI: 10.1109/MCSE.2010.27
        Eastman P and Pande VS. Efficient nonbonded interactions for molecular dynamics on a graphics processing unit. J. Comput. Chem. 31:1268, 2010. DOI: 10.1002/jcc.21413
        Eastman P and Pande VS. Constant constraint matrix approximation: A robust, parallelizable constraint method for molecular simulations. J. Chem. Theor. Comput. 6:434, 2010. DOI: 10.1021/ct900463w
        Chodera JD and Shirts MR. Replica exchange and expanded ensemble simulations as Gibbs multistate: Simple improvements for enhanced mixing. J. Chem. Phys., 135:194110, 2011. DOI:10.1063/1.3660669

Appendix

REST potential

According to Wang et al. 2011 (DOI: 10.1021/jp204407d, available at: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3172817/), the modified REST potential is defined as:

\(U_\text{total}(x) = \frac{\beta_m}{\beta_0} U_\text{REST}(x) + \sqrt{\frac{\beta_m}{\beta_0}} U_\text{inter}(x) + U_\text{non-REST}(x)\)

where \(U_\text{total}(x)\) represents the total potential energy of configuration \(x\), \(U_\text{REST}(x)\) represents the potential energy of interactions between REST atoms, \(U_\text{inter}(x)\) represents the potential energy of interactions between REST and non-REST atoms (i.e., “inter”), and \(U_\text{non-REST}(x)\) represents interactions between non-REST atoms.

Electrostatic interactions

For the NonbondedForce, we use offsets to scale the energies according to REST. We first explore how offsets are applied for electrostatic interactions.

For electrostatics, the potential energy is defined by: \(U_\text{elec} = \frac{q_1 q_2}{4\pi\epsilon_0}\frac{1}{r}\) where \(q\) is the charge of an atom, \(\epsilon_0\) is the vacuum permittivity, and \(r\) is the distance between atoms.

We apply offsets to the charge (\(q\)) according to:

charge = baseCharge + lambda_rest_electrostatics*chargeScale

where baseCharge is the current charge value for the atom in the NonbondedForce, lambda_rest_electrostatics is a global parameter and must be between 0 and 1, and chargeScale (when multiplied with lambda_rest_electrostatics) controls the value to be added to baseCharge.

We set chargeScale = baseCharge and derive the functional form for lambda_rest_electrostatics to enable REST scaling for each of the following cases:

Case 1: Both atoms are non-REST

  • The interaction should not be scaled, meaning no offset is needed.

Case 2: Atom 1 is non-REST, Atom 2 is REST (or vice-versa)

  • We can express the electrostatics energy as

    \(U_\text{elec} = \sqrt{\frac{\beta_m}{\beta_0}} \frac{q_1 q_2}{4\pi\epsilon_0}\frac{1}{r} = \frac{q_1 (q_2 + \lambda q_2)}{4\pi\epsilon_0}\frac{1}{r}\)

  • which is equivalent to

    \(U_\text{elec} = \sqrt{\frac{\beta_m}{\beta_0}} = 1 + \lambda\)

  • meaning

    \(\lambda = \sqrt{\frac{\beta_m}{\beta_0}} - 1\)

Case 3: Both atoms are REST

  • We can express the electrostatics energy as

    \(U_\text{elec} = \frac{\beta_m}{\beta_0} \frac{q_1 q_2}{4\pi\epsilon_0}\frac{1}{r} = \frac{(q_1 + q_1 \lambda) (q_2 + \lambda q_2)}{4\pi\epsilon_0}\frac{1}{r}\)

  • which is equivalent to

    \(U_\text{elec} = \frac{\beta_m}{\beta_0} \frac{q_1 q_2}{4\pi\epsilon_0}\frac{1}{r} = \frac{q_1 q_2 (1 + \lambda)^2}{4\pi\epsilon_0}\frac{1}{r}\)

  • meaning

    \(\lambda = \sqrt{\frac{\beta_m}{\beta_0}} - 1\)

In both cases 2 and 3, \(\lambda = \sqrt{\frac{\beta_m}{\beta_0}} - 1\). We use this as our lambda function for lambda_rest_electrostatics.

Steric interactions

We also use offsets to modify steric interactions involving REST atoms.

For sterics, the potential energy is defined by: \(U_\text{sterics} = 4\epsilon\left(\left(\frac{\sigma}{r}\right)^{12} - \left(\frac{\sigma}{r}\right)^6\right)\) where \(\sigma\) is the distance at which the energy equals zero, \(r\) is the distance between atoms, and \(\epsilon\) is the strength of the interaction (computed from the the \(\epsilon\)s of atoms 1 and 2: \(\epsilon = \sqrt{\epsilon_1\epsilon_2}\)).

We apply offsets to \(\epsilon_i\) according to:

epsilon = baseEpsilon + lambda_rest_sterics*epsilonScale$

where baseEpsilon is the current \(\epsilon\) value for the atom in the NonbondedForce, lambda_rest_sterics is a global parameter and must be between 0 and 1, and epsilonScale (when multiplied with lambda_rest_sterics) controls the value to be added to baseEpsilon.

We set epsilonScale = baseEpsilon and derive the functional form for lambda_rest_sterics to enable REST scaling for each of the following cases:

Case 1: Both atoms are non-REST

  • The interaction should not be scaled, meaning no offset is needed.

Case 2: Atom 1 is non-REST, Atom 2 is REST (or vice-versa)

  • We can express the sterics energy as

    \(U_\text{sterics} = \sqrt{\frac{\beta_m}{\beta_0}} \cdot 4\sqrt{\epsilon_1\epsilon_2}\left(\left(\frac{\sigma}{r}\right)^{12} - \left(\frac{\sigma}{r}\right)^6\right) = 4\sqrt{\epsilon_1(\epsilon_2 + \lambda\epsilon_2)}\left(\left(\frac{\sigma}{r}\right)^{12} - \left(\frac{\sigma}{r}\right)^6\right)\)

  • which is equivalent to

    \(U_\text{sterics} = \sqrt{\frac{\beta_m}{\beta_0}} \cdot \sqrt{\epsilon_1\epsilon_2} = \sqrt{\epsilon_1\epsilon_2(1 + \lambda)}\)

  • meaning

    \(\lambda = \frac{\beta_m}{\beta_0} - 1\)

Case 3: Both atoms are REST

  • We can express the sterics energy as

    \(U_\text{sterics} = \frac{\beta_m}{\beta_0} \cdot 4\sqrt{\epsilon_1\epsilon_2}\left(\left(\frac{\sigma}{r}\right)^{12} - \left(\frac{\sigma}{r}\right)^6\right) = 4\sqrt{\epsilon_1\epsilon_2 (1 + \lambda)^2}\left(\left(\frac{\sigma}{r}\right)^{12} - \left(\frac{\sigma}{r}\right)^6\right)\)

  • which is equivalent to

    \(U_\text{sterics} = \frac{\beta_m}{\beta_0} \cdot \sqrt{\epsilon_1\epsilon_2} = \sqrt{\epsilon_1\epsilon_2(1 + \lambda)^2}\)

  • meaning

    \(\lambda = \frac{\beta_m}{\beta_0} - 1\)

In both cases 2 and 3, \(\lambda = \frac{\beta_m}{\beta_0} - 1\). We use this as our lambda function for lambda_rest_sterics.

Electrostatic exceptions

We also use offsets to modify electrostatics exceptions involving REST atoms.

For electrostatic exceptions, the functional form of the potential energy is the same as for electrostatic interactions, except that for exceptions, the functional form accepts the charge product directly, instead of the individual charges for each atom. Therefore, lambda_rest_electrostatics will scale the product, not each individual charge.

We apply offsets to the chargeProd (\(q_1 q_2\)) according to:

chargeProd = baseChargeProd + lambda_rest_electrostatics*chargeProdScale

where baseChargeProd is the current charge product for the two atoms in the NonbondedForce, lambda_rest_electrostatics is a global parameter and must be between 0 and 1, and chargeProdScale (when multiplied with lambda_rest_electrostatics) controls the value to be added to baseChargeProd.

We set chargeProdScale = baseChargeProd and derive the functional form for the rest scale factor (i.e., \(\lambda\)) for each of the following cases:

Case 1: Both atoms are non-REST

  • The interaction should not be scaled, meaning no offset is needed.

Case 2: Atom 1 is non-REST, Atom 2 is REST (or vice-versa)

  • We can express the electrostatics energy as

    \(U_\text{elec} = \sqrt{\frac{\beta_m}{\beta_0}} \frac{q_1 q_2}{4\pi\epsilon_0}\frac{1}{r} = \frac{q_1 q_2 + \lambda q_1 q_2}{4\pi\epsilon_0}\frac{1}{r}\)

  • which is equivalent to

    \(U_\text{elec} = \sqrt{\frac{\beta_m}{\beta_0}} = 1 + \lambda\)

  • meaning

    \(\lambda = \sqrt{\frac{\beta_m}{\beta_0}} - 1\)

Case 3: Both atoms are REST

  • We can express the electrostatics energy as

    \(U_\text{elec} = \frac{\beta_m}{\beta_0} \frac{q_1 q_2}{4\pi\epsilon_0}\frac{1}{r} = \frac{q_1 q_2 + \lambda q_1 q_2}{4\pi\epsilon_0}\frac{1}{r}\)

  • which is equivalent to

    \(U_\text{elec} = \frac{\beta_m}{\beta_0} = 1 + \lambda\)

  • meaning

    \(\lambda = \frac{\beta_m}{\beta_0} - 1\)

In case 2, \(\lambda = \sqrt{\frac{\beta_m}{\beta_0}} - 1\), which matches the lambda_rest_electrostatics function. In case 3, \(\lambda = \frac{\beta_m}{\beta_0} - 1\), which matches the lambda_rest_sterics function. We therefore use lambda_rest_electrostatics and lambda_rest_sterics to control the offsets for inter and REST exceptions, respectively.

Steric exceptions

We also use offsets to modify steric exceptions involving REST atoms.

For steric exceptions, the functional form of the potential energy is the same as for steric interactions, except that for exceptions, the functional form accepts the combined \(\epsilon\) directly, instead of the individual \(\epsilon\)s for each atom. Therefore, lambda_rest_sterics will scale the product, not each individual \(\epsilon\).

We apply offsets to the combined \(\epsilon\) (computed from \(\sqrt{\epsilon_1 \epsilon_2}\)) according to:

epsilonCombined = baseEpsilonCombined + lambda_rest_sterics*epsilonCombinedScale

where baseEpsilonCombined is the current combined epsilon for the two atoms in the NonbondedForce, lambda_rest_sterics is a global parameter and must be between 0 and 1, and epsilonCombinedScale (when multiplied with lambda_rest_sterics) controls the value to be added to baseEpsilonCombined.

We set epsilonCombinedScale = baseEpsilonCombined and derive the functional form for the rest scale factor (i.e., \(\lambda\)) for each of the following cases:

Case 1: Both atoms are non-REST

  • The interaction should not be scaled, meaning no offset is needed.

Case 2: Atom 1 is non-REST, Atom 2 is REST (or vice-versa)

  • We can express the sterics energy as

    \(U_\text{sterics} = \sqrt{\frac{\beta_m}{\beta_0}} \cdot 4\sqrt{\epsilon_1\epsilon_2}\left(\left(\frac{\sigma}{r}\right)^{12} - \left(\frac{\sigma}{r}\right)^6\right) = 4\left(\sqrt{\epsilon_1\epsilon_2} + \lambda\sqrt{\epsilon_1\epsilon_2}\right)\left(\left(\frac{\sigma}{r}\right)^{12} - \left(\frac{\sigma}{r}\right)^6\right)\)

  • which is equivalent to

    \(U_\text{sterics} = \sqrt{\frac{\beta_m}{\beta_0}} \cdot \sqrt{\epsilon_1\epsilon_2} = \sqrt{\epsilon_1\epsilon_2}(1 + \lambda)\)

  • meaning

    \(\lambda = \sqrt{\frac{\beta_m}{\beta_0}} - 1\)

Case 3: Both atoms are REST

  • We can express the sterics energy as

    \(U_\text{sterics} = \frac{\beta_m}{\beta_0} \cdot 4\sqrt{\epsilon_1\epsilon_2}\left(\left(\frac{\sigma}{r}\right)^{12} - \left(\frac{\sigma}{r}\right)^6\right) = 4\left(\sqrt{\epsilon_1\epsilon_2} + \lambda \sqrt{\epsilon_1\epsilon_2}\right)\left(\left(\frac{\sigma}{r}\right)^{12} - \left(\frac{\sigma}{r}\right)^6\right)\)

  • which is equivalent to

    \(U_\text{sterics} = \frac{\beta_m}{\beta_0} \cdot \sqrt{\epsilon_1\epsilon_2} = \sqrt{\epsilon_1\epsilon_2}(1 + \lambda)\)

  • meaning

    \(\lambda = \frac{\beta_m}{\beta_0} - 1\)

In case 2, \(\lambda = \sqrt{\frac{\beta_m}{\beta_0}} - 1\), which matches the lambda_rest_electrostatics function. In case 3, \(\lambda = \frac{\beta_m}{\beta_0} - 1\), which matches the lambda_rest_sterics function. We therefore use lambda_rest_electrostatics and lambda_rest_sterics to control the offsets for inter and REST exceptions, respectively.