#!/usr/local/bin/env python
# ==============================================================================
# MODULE DOCSTRING
# ==============================================================================
"""
ParallelTemperingSampler
========================
Derived multi-thermodynamic state multistate class with exchanging configurations between replicas of different
temperatures. This is a special case which accepts a single thermodynamic_state and different temperatures to sample.
If you want different temperatures and Hamiltonians, use ReplicaExchangeSampler with temperatures pre-set.
COPYRIGHT
Current version by Andrea Rizzi <andrea.rizzi@choderalab.org>, Levi N. Naden <levi.naden@choderalab.org> and
John D. Chodera <john.chodera@choderalab.org> while at Memorial Sloan Kettering Cancer Center.
Original version by John D. Chodera <jchodera@gmail.com> while at the University of
California Berkeley.
LICENSE
This code is licensed under the latest available version of the MIT License.
"""
import copy
import math
import logging
import numpy as np
import openmmtools as mmtools
from .replicaexchange import ReplicaExchangeSampler, ReplicaExchangeAnalyzer
from .multistatereporter import MultiStateReporter
logger = logging.getLogger(__name__)
# ==============================================================================
# PARALLEL TEMPERING
# ==============================================================================
[docs]class ParallelTemperingSampler(ReplicaExchangeSampler):
"""Parallel tempering simulation facility.
This class provides a facility for parallel tempering simulations. It
is a subclass of :class:`ReplicaExchange`, but provides efficiency improvements
for parallel tempering simulations, so should be preferred for this type
of simulation. In particular, this makes use of the fact that the reduced
potentials are linear in inverse temperature.
Examples
--------
Create the system.
>>> from simtk import unit
>>> from openmmtools import testsystems, states, mcmc
>>> import tempfile
>>> testsystem = testsystems.AlanineDipeptideImplicit()
Create thermodynamic states for parallel tempering with exponentially-spaced schedule.
>>> n_replicas = 3 # Number of temperature replicas.
>>> T_min = 298.0 * unit.kelvin # Minimum temperature.
>>> T_max = 600.0 * unit.kelvin # Maximum temperature.
>>> reference_state = states.ThermodynamicState(system=testsystem.system, temperature=T_min)
Initialize simulation object with options. Run with a GHMC integrator.
>>> move = mcmc.GHMCMove(timestep=2.0*unit.femtoseconds, n_steps=50)
>>> simulation = ParallelTemperingSampler(mcmc_moves=move, number_of_iterations=2)
Create simulation with its storage file (in a temporary directory) and run.
>>> storage_path = tempfile.NamedTemporaryFile(delete=False).name + '.nc'
>>> reporter = MultiStateReporter(storage_path, checkpoint_interval=10)
>>> simulation.create(reference_state,
... states.SamplerState(testsystem.positions),
... reporter, min_temperature=T_min,
... max_temperature=T_max, n_temperatures=n_replicas)
>>> simulation.run(n_iterations=1)
Clean up.
>>> os.remove(storage_path)
See Also
--------
MultiStateSampler
ReplicaExchangeSampler
"""
_TITLE_TEMPLATE = ('Parallel tempering simulation created using ParallelTempering '
'class of yank.multistate on {}')
[docs] def create(self, thermodynamic_state, sampler_states: list, storage,
min_temperature=None, max_temperature=None, n_temperatures=None,
temperatures=None, **kwargs):
"""Initialize a parallel tempering simulation object.
Parameters
----------
thermodynamic_state : openmmtools.states.ThermodynamicState
Reference thermodynamic state that will be simulated at the given
temperatures.
WARNING: This is a SINGLE state, not a list of states!
sampler_states : openmmtools.states.SamplerState or list
One or more sets of initial sampler states. If a list of SamplerStates,
they will be assigned to replicas in a round-robin fashion.
storage : str or Reporter
If str: path to the storage file, checkpoint options are default
If Reporter: Instanced :class:`Reporter` class, checkpoint information is read from
In the future this will be able to take a Storage class as well.
min_temperature : simtk.unit.Quantity, optional
Minimum temperature (units of temperature, default is None).
max_temperature : simtk.unit.Quantity, optional
Maximum temperature (units of temperature, default is None).
n_temperatures : int, optional
Number of exponentially-spaced temperatures between ``min_temperature``
and ``max_temperature`` (default is None).
temperatures : list of simtk.unit.Quantity, optional
If specified, this list of temperatures will be used instead of
``min_temperature``, ``max_temperature``, and ``n_temperatures`` (units of temperature,
default is None).
metadata : dict, optional
Simulation metadata to be stored in the file.
Notes
-----
Either (``min_temperature``, ``max_temperature``, ``n_temperatures``) must all be
specified or the list of '`temperatures`' must be specified.
"""
# Create thermodynamic states from temperatures.
if not isinstance(thermodynamic_state, mmtools.states.ThermodynamicState):
raise ValueError("ParallelTempering only accepts a single ThermodynamicState!\n"
"If you have already set temperatures in your list of states, please use the "
"standard ReplicaExchange class with your list of states.")
if temperatures is not None:
logger.debug("Using provided temperatures")
elif min_temperature is not None and max_temperature is not None and n_temperatures is not None:
temperatures = [min_temperature + (max_temperature - min_temperature) *
(math.exp(i / n_temperatures-1) - 1.0) / (math.e - 1.0)
for i in range(n_temperatures)] # Python 3 uses true division for /
logger.debug('using temperatures {}'.format(temperatures))
else:
raise ValueError("Either 'temperatures' or ('min_temperature', 'max_temperature', "
"and 'n_temperatures') must be provided.")
thermodynamic_states = [copy.deepcopy(thermodynamic_state) for _ in range(n_temperatures)]
for state, temperature in zip(thermodynamic_states, temperatures):
state.temperature = temperature
# Initialize replica-exchange simulation.
super(ParallelTemperingSampler, self).create(thermodynamic_states, sampler_states, storage=storage, **kwargs)
def _compute_replica_energies(self, replica_id):
"""Compute the energy for the replica at every temperature.
Because only the temperatures differ among replicas, we replace the generic O(N^2)
replica-exchange implementation with an O(N) implementation.
"""
# Initialize replica energies for each thermodynamic state.
energy_thermodynamic_states = np.zeros(self.n_states)
energy_unsampled_states = np.zeros(len(self._unsampled_states))
# Determine neighborhood
state_index = self._replica_thermodynamic_states[replica_id]
neighborhood = self._neighborhood(state_index)
# Only compute energies over neighborhoods
energy_neighborhood_states = energy_thermodynamic_states[neighborhood] # Array, can be indexed like this
neighborhood_thermodynamic_states = [self._thermodynamic_states[n] for n in neighborhood] # List
# Retrieve sampler states associated to this replica.
sampler_state = self._sampler_states[replica_id]
# Thermodynamic state differ only by temperatures.
reference_thermodynamic_state = self._thermodynamic_states[0]
# Get the context, any Integrator works.
context, integrator = mmtools.cache.global_context_cache.get_context(reference_thermodynamic_state)
# Update positions and box vectors.
sampler_state.apply_to_context(context)
# Compute energy.
reference_reduced_potential = reference_thermodynamic_state.reduced_potential(context)
# Strip reference potential of reference state's beta.
reference_beta = 1.0 / (mmtools.constants.kB * reference_thermodynamic_state.temperature)
reference_reduced_potential /= reference_beta
# Update potential energy by temperature.
for thermodynamic_state_id, thermodynamic_state in enumerate(neighborhood_thermodynamic_states):
beta = 1.0 / (mmtools.constants.kB * thermodynamic_state.temperature)
energy_neighborhood_states[thermodynamic_state_id] = beta * reference_reduced_potential
# Since no assumptions can be made about the unsampled thermodynamic states, do it the hard way
for unsampled_id, state in enumerate(self._unsampled_states):
if unsampled_id == 0 or not state.is_state_compatible(context_state):
context_state = state
# Get the context, any Integrator works.
context, integrator = mmtools.cache.global_context_cache.get_context(state)
# Update positions and box vectors. We don't need
# to set Context velocities for the potential.
sampler_state.apply_to_context(context, ignore_velocities=True)
else:
# If this state is compatible with the context, just fix the
# thermodynamic state as positions/box vectors are the same.
state.apply_to_context(context)
# Compute energy.
energy_unsampled_states[unsampled_id] = state.reduced_potential(context)
# Return the new energies.
return energy_neighborhood_states, energy_unsampled_states
[docs]class ParallelTemperingAnalyzer(ReplicaExchangeAnalyzer):
"""
The ParallelTemperingAnalyzer is the analyzer for a simulation generated from a Parallel Tempering sampler
simulation, implemented as an instance of the :class:`ReplicaExchangeAnalyzer` as the sampler is a subclass of
the :class:`yank.multistate.ReplicaExchangeSampler`
See Also
--------
PhaseAnalyzer
ReplicaExchangeAnalyzer
"""
pass
# ==============================================================================
# MAIN AND TESTS
# ==============================================================================
if __name__ == "__main__":
import doctest
doctest.testmod()