Advanced features and developer’s tutorial¶
This tutorial describes more advanced features that can be useful for people developing their software using or extending OpenMMTools.
Using and implementing integrator and force objects¶
Copy and serialization utilities¶
The integrators and forces in OpenMMTools are usually implemented by extending
custom classes in OpenMM. For example, the declaration of our LangevinIntegrator
and the HarmonicRestraintForce
goes boils down to
class LangevinIntegrator(mmtools.utils.RestorableOpenMMObject, openmm.CustomIntegrator):
pass
class HarmonicRestraintForce(mmtools.utils.RestorableOpenMMObject, openmm.CustomCentroidBondForce):
pass
The purpose of inheriting from openmmtools.utils.RestorableOpenMMObject class has to do with copies and
serialization. Without RestorableOpenMMObject
, these objects can still be copied and go through the standard
serialization and deserialization in OpenMM without errors
from simtk import openmm, unit
class VelocityVerlet(openmm.CustomIntegrator):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.addComputePerDof("x", "x+dt*v")
self.addComputePerDof("v", "v+0.5*dt*f/m")
def my_method(self):
return 0.0
integrator = VelocityVerlet(1*unit.femtosecond)
copied_integrator = copy.deepcopy(integrator)
integrator_serialization = openmm.XmlSerializer.serialize(integrator)
deserialized_integrator = openmm.XmlSerializer.deserialize(integrator_serialization)
However, copies and serializations in OpenMM are performed at the C++ level, and thus they don’t keep track of the Python class and methods.
>>> print(type(copied_integrator))
<class 'simtk.openmm.openmm.CustomIntegrator'>
>>> deserialized_integrator.my_method()
Traceback (most recent call last):
...
AttributeError: type object 'object' has no attribute '__getattr__'
Inheriting from openmmtools.utils.RestorableOpenMMObject, allows you to easily recover the original interface
after copying or deserializing. This happens automatically for copies, but you’ll have to use RestorableOpenMMObject.restore_interface()
after deserialization.
from openmmtools import utils
class VelocityVerlet(utils.RestorableOpenMMObject, openmm.CustomIntegrator):
def __init(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.addComputePerDof("x", "x+dt*v")
self.addComputePerDof("v", "v+0.5*dt*f/m")
def my_method(self):
return 0.0
integrator = VelocityVerlet(1*unit.femtosecond)
>>> copied_integrator = copy.deepcopy(integrator)
>>> isinstance(copied_integrator, VelocityVerlet)
True
>>> integrator_serialization = openmm.XmlSerializer.serialize(integrator)
>>> deserialized_integrator = openmm.XmlSerializer.deserialize(integrator_serialization)
>>> utils.RestorableOpenMMObject.restore_interface(deserialized_integrator)
True
>>> deserialized_integrator.my_method()
0.0
For forces, the function openmmtools.forces.find_forces(system) automatically calls
RestorableOpenMMObject.restore_interface()
on all system
forces so there’s usually no need to perform that
call after deserialization.
Integrators coupled to a heat bath¶
If you implement an integrator coupled to a heat bath, you have to expose getTemperature and setTemperature methods
or ThermodynamicState
won’t have any way to recognize it, and it will add an AndersenThermostat
force when
initializing the OpenMM Context
object.
The base class openmmtools.integrators.ThermostatedIntegrator is a convenience class implemented for
this purpose. Inheriting from ThermostatedIntegrator
will implicitly add the RestorableOpenMMObject
functionalities
as well.
>>> from openmmtools import integrators
>>> class MyIntegrator(integrators.ThermostatedIntegrator):
... def __init__(self, temperature=298.0*unit.kelvin, timestep=1.0*unit.femtoseconds):
... super().__init__(temperature, timestep)
...
>>> integrator = MyIntegrator(temperature=350*unit.kelvin)
>>> integrator.getTemperature()
Quantity(value=350.0, unit=kelvin)
>>> integrator.setTemperature(380.0*unit.kelvin)
Using standard Python attribute in custom integrators and forces¶
You should avoid having pure Python attributes when inheriting from custom OpenMM integrators and forces and instead favor using properties that read that attribute from the underlying OpenMM object as, for example, a global variable.
For example, an integrator exposing the temperature should not hold a simple temperature
Python attribute
internally such as
class INCORRECTIntegrator(openmm.CustomIntegrator):
def __init__(self, *args, temperature=298.15*unit.kelvin, **kwargs):
super().__init__(*args, **kwargs)
self.temperature = temperature
but it expose it as a getter or a property similarly to the follow.
class CorrectIntegrator(openmm.CustomIntegrator):
def __init__(self, *args, temperature=298.15*unit.kelvin, **kwargs):
super().__init__(*args, **kwargs)
self.addGlobalVariable('temperature', temperature)
@property
def temperature(self):
return self.getGlobalVariableByName('temperature') * unit.kelvin
This is because:
If the parameter doesn’t affect serialization
ContextCache
won’t be able to distinguish between two integrators that differ by that parameter, and it may return an incorrect integrator.Python attribute cannot be restored by
RestorableOpenMMObject
since there’s no information about them in the XML string, and thus they will be lost with serialization.
Handling (compound) thermodynamic states¶
In the examples that follow, we’ll use a simple ThermodynamicState
, but everything applies to CompoundThermodynamicState
as well as CompoundThermodynamicState
is a subclass of ThermodynamicState
.
Modifying a System object buried in a ThermodynamicState¶
Setting a thermodynamic parameter in ThermodynamicState
is practically instantaneous, but modifying anything else
involves the copy of the internal System
object so it can be very slow.
from openmmtools import states
from openmmtools import testsystems
system = testsystems.TolueneVacuum().system
thermo_state = states.ThermodynamicState(system, temperature=300*unit.kelvin)
# This is very fast.
thermo_state.temperature = 400.0*unit.kelvin
system = thermo_state.system # This is a copy! Changes to this System won't affect thermo_state.
# Make your changes to system.
thermo_state.system = system # This involves another System copy.
The copies are there to ensure the consistency of ThermodynamicState
internal state. If you need to consistently
modifying part of the systems during the simulation consider implementing a composable state that handle those degrees
of freedom (see section Implementing a new ComposableState).
Another thing to keep in mind is that by default the property ThermodynamicState.system
will return a System
containing an AndersenThermostat
force. If you only use ThermodynaicState.create_context()
or the ContextCache
class to create OpenMM Context
objects, this shouldn’t cause issues, but if for any reason you don’t want that
thermostat you can use the getter instead of the property.
system = thermo_state.get_system(remove_thermostat=True)
Implementation details of compatibility checks¶
Internally, ThermodynamicState
associates a unique hash to a System
in a particular ensemble, and it compares
this hash to check for compatibility. The function that performs this task looks like this:
@classmethod
def _standardize_and_hash(cls, system):
"""Standardize the system and return its hash."""
cls._standardize_system(system)
system_serialization = openmm.XmlSerializer.serialize(system)
return system_serialization.__hash__()
The _standardize_system()
functions sets the thermodynamic parameters controlled by the ThermodynamicState
to a
standard value so that System
that differ by only those parameters will have identical XML serialized strings, and
thus identical hashes.
The section Implementing a new ComposableState has information on how the composable states expand the concept of compatibility to thermodynamic parameters other than temperature and pressure.
Note
As a consequence of how the compatibility hash is computed, two ThermodynamicStates to be compatible must have Systems with the same particles and forces in the same order, or the XML serialization will be different.
Copying and initialization of multiple thermodynamic states¶
Because of some memory optimizations, copying a ThermodynamicState
or a CompoundThermodynamicState
does not copy
the internal System
so it is practically instantaneous. On the other hand, initializing a new ThermodynamicState
or a CompoundThermodynamicState
object does involve a System
copy.
thermo_state1 = states.ThermodynamicState(system, temperature=300*unit.kelvin)
# Very fast.
thermo_state2 = copy.deepcopy(thermo_state)
thermo_state2.temperature = 350*unit.kelvin
# Slow.
thermo_state2 = states.ThermodynamicState(system, temperature=350*unit.kelvin)
The function openmmtools.states.create_thermodynamic_state_protocol takes advantage of this to make it easy
to instantiate a list of ThermodynamicState
or CompoundThermodynamicState
objects that differ only by the controlled
parameters.
Implementing a new ComposableState¶
The IComposableInterface¶
Composable states allow to control thermodynamic parameters of the simulation while masking their implementation details. There are no restrictions on the implementation details, but the class must implement the openmmtools.states.IComposableState interface. You can see the API docs for contract details, but here is a list of the methods.
class IComposableState:
def apply_to_system(self, system):
"""Modify an OpenMM System to be in this thermodynamic state."""
def check_system_consistency(self, system):
"""Raise AlchemicalStateError if system has different parameters."""
def apply_to_context(self, context):
"""Modify an OpenMM Context to be in this thermodynamic state."""
def _standardize_system(cls, system):
"""Modify the System to be in the standard thermodynamic state."""
def _on_setattr(self, standard_system, attribute_name, old_attribute_value):
"""Callback that checks if standard system needs to be updated after a state attribute is set."""
def _find_force_groups_to_update(self, context, current_context_state, memo)
"""Find the force groups whose energy must be recomputed after apply_to_context."""
# Optional. This is used only for optimizations.
The _standardize_system
method effectively determines which other states will be compatible (see also section
Implementation details of compatibility checks). The purpose of _standardize_system
is to set the parameters of
the System
that can be manipulated in the Context
to the same value so that their XML serialization string and
their hash will be identical. Systems that after standardization are identical are assigned to the same Context
by
ContextCache.get_context()
.
Relatedly, the callback _on_setattr()
is called by CompoundThermodynamicState
after a thermodynamic parameter
has been set. The method must return True
if the change in the thermodynamic parameter has caused the standard system
to have a different hash. For example, in the basic ThermodynamicState
class this happens when the pressure
parameter goes from None
to any valid value because states in NVT and NPT are not compatible.
The method _find_force_groups_to_update
is optional and related to the optimization described in
Computing the reduced potential of one configuration at multiple thermodynamic states.
ComposableStates that control a Force global parameter¶
Often, a thermodynamic parameter can be implemented with OpenMM as a global parameter added to a custom force. For
example, to alchemically soften torsions, alchemy.AbsoluteAlchemicalFactory
substitute some of the torsion potential
terms using a openmm.CustomTorsionForce
whose energy is multiplied by a global parameter called lambda_torsions
.
energy_function = "lambda_torsions * k*(1+cos(periodicity*theta-phase))"
custom_force = openmm.CustomTorsionForce(energy_function)
custom_force.addGlobalParameter('lambda_torsions', 1.0)
# Other force configurations.
system.addForce(custom_force)
When this is the case, the base class openmmtools.states.GlobalParameterState
can be used to create a composable state
very quickly.
from openmmtools.states import GlobalParameterState
class MyComposableState(GlobalParameterState):
lambda_torsions = GlobalParameterState.GlobalParameter('lambda_torsions', standard_value=1.0)
It is possible to perform checks on the assigned value by adding a validator.
class MyComposableState(GlobalParameterState):
lambda_torsions = GlobalParameterState.GlobalParameter('lambda_torsions', standard_value=1.0)
@lambda_torsions.validator
def lambda_torsions(self, instance, new_value):
if new_value is not None and not (0.0 <= new_value <= 1.0):
raise ValueError('lambda_torsions must be between 0.0 and 1.0')
return new_value
The example above allows only values between 0.0 and 1.0 for lambda_torsions
.
Computing the reduced potential of one configuration at multiple thermodynamic states¶
When computing the potential energy of a single configuration at multiple thermodynamic states, it is often unnecessary to compute the whole Hamiltonian multiple times but just the terms of the Hamiltonian that change from one state to another. OpenMM makes this possible to compute only the energy of a subset of forces through the force groups mechanism.
force = openmm.CustomBondForce('(K/2)*(r-r0)^2;')
force.setForceGroup(5)
The utility function openmmtools.states.reduced_potential_at_states()
takes advantage of forces separated in different
groups to efficiently compute the reduced potentials at the thermodynamic states.
from openmmtools import alchemy
from openmmtools import cache
alanine = testsystems.AlchemicalAlanineDipeptide()
protocol = {'lambda_sterics': [1.0, 0.5, 0.0],
'lambda_electrostatics': [1.0, 0.5, 0.0]}
constants = {'temperature': 300*unit.kelvin}
composable_states = [alchemy.AlchemicalState.from_system(alanine.system)]
compound_states = states.create_thermodynamic_state_protocol(alanine.system, protocol,
constants, composable_states)
sampler_state = states.SamplerState(positions=alanine.positions)
reduced_potentials = states.reduced_potential_at_states(sampler_state, compound_states,
cache.global_context_cache)
In order for the optimization to take effect, the composable states must implement the method
_find_force_groups_to_update(self, context, current_context_state, memo)
. This method inspects the System
associated to the context
and return the force groups that will have an updated energy after the state will be changed
from current_context_state
to self
. The memo
dictionary can be use to store the force groups to inspect in
subsequent calls of the method within a reduced_potential_at_states
execution so that the System
must be parsed
only the first time.
Implementing a new MCMCMove¶
The MCMCMove interface¶
An MCMCMove
requires exclusively the implementation of an apply
method with the following signature (see the
API documentation for more details.
class MCMCMove(SubhookedABCMeta):
def apply(self, thermodynamic_state, sampler_state):
pass
Anything can happen inside apply
as long as thermodynamic_state
and sampler_state
are updated correctly.
It is usually a good idea to include in the constructor a context_cache
argument to let the user specify how the
Context
should be created and on which platform.
OpenMM integrators that modify the thermodnamic state¶
Custom OpenMM integrators can modify global variables that effectively change the thermodynamic state of the Context
.
Important
Remember to update the thermodynamic_state
object correctly at the end of apply
if the integrator changes the thermodynamic state of the simulation.
When this is the case, it’s not possible to cast your integrator into an MCMCMove
with IntegratorMove
.
Nevertheless, it’s still possible to take advantage of the extra features already offered by IntegratorMove
by
subclassing the openmmtools.mcmc.BaseIntegratorMove <mcmc> class. IntegratorMove
inherits from this base class. An
implementation would look more or less like this (see the API documentation for the details).
class MyMove(BaseIntegratorMove):
def __init__(self, timestep, n_steps, **kwargs):
super(MyMove, self).__init__(n_steps, **kwargs)
self.timestep = timestep
def _get_integrator(self, thermodynamic_state):
return MyIntegrator(self.timestep, thermodynamic_state.temperature)
def _before_integration(self, context, thermodynamic_state):
# Optional: Any operation performed after the context
# was created but before integration.
def _after_integration(self, context, thermodynamic_state):
# Update thermodynamic_state from context parameters.
# Optional: Read statistics from context.getIntegrator() parameters.
Metropolized MCMCMoves¶
The mcmc module contains a base class for Metropolized moves as well. The following class implement an example that simply adds the unit vector to the initial coordinates.
from openmmtools import mcmc
class AddOneVector(mcmc.MetropolizedMove):
def _propose_positions(self, initial_positions):
print('Propose new positions')
displacement = numpy.array([1.0, 1.0, 1.0]) * unit.angstrom
return initial_positions + displacement
The parent class will take care of implementing the Metropolis acceptance criteria, collecting acceptance statistics,
and updating the SamplerState
correctly. The constructor accepts an optional atom_subset
to limit the move to
certain atoms. In this case, the initial_positions
will be the positions of the atom subset only.
>>> alanine = testsystems.AlanineDipeptideVacuum()
>>> sampler_state = states.SamplerState(alanine.positions)
>>> thermodynamic_state = states.ThermodynamicState(alanine.system, 300*unit.kelvin)
>>> move = AddOneVector(atom_subset=list(range(sampler_state.n_particles)))
>>> move.apply(thermodynamic_state, sampler_state)
Propose new positions
>>> move.n_accepted
1
>>> move.n_proposed
1
Real time analysis output¶
The MultiStateReporter
outputs a file with analysis statistics every few interactions dictated by the
online_analysis_interval option from the MultiStateSampler
object. The format of this file consists of several
entries as the following:
- iteration: 4
percent_complete: 25.0
mbar_analysis:
free_energy_in_kT: 20.747354877367165
standard_error_in_kT: 0.7055846807301683
number_of_uncorrelated_samples: 5.0
timing_data:
iteration_seconds: 0.04654264450073242
average_seconds_per_iteration: 0.05740888913472494
estimated_time_remaining: '0:00:00.746316'
estimated_localtime_finish_date: 2022-Mar-24-20:07:21
estimated_total_time: '0:00:00.918542'
ns_per_day: 150.49934130799824
This is intended to be used for monitoring performance of currently running simulations.