Unverified Commit aa8ec1a8 authored by Peter Eastman's avatar Peter Eastman Committed by GitHub
Browse files

SimulatedTempering uses ExpandedEnsembleSampler (#5281)

parent ce9fcace
...@@ -1617,16 +1617,35 @@ These are known as enhanced sampling methods. OpenMM offers several that you ...@@ -1617,16 +1617,35 @@ These are known as enhanced sampling methods. OpenMM offers several that you
can choose from. They are briefly described here. Consult the API documentation can choose from. They are briefly described here. Consult the API documentation
for more detailed descriptions and example code. for more detailed descriptions and example code.
Simulated Tempering Replica Exchange
------------------- ----------------
Replica exchange involves simulating several copies of the system at once, each in a
different thermodynamic state. It periodically exchanges states between the replicas, allowing
each trajectory to explore multiple states.
In one common version, the thermodynamic states correspond to different temperatures. A
single replica might spend time at a high temperature where barriers are easily crossed, then
transition to the lower temperature you are most interested in to sample the distribution of
conformations at that temperature. This is a powerful method to speed up sampling when you
do not know in advance what motions you want to sample.
In another common version, the thermodynamic states correspond to different potential
functions, for example different values of a global parameter that controls the strength of
a force. A single trajectory might sample the distribution of conformations when the force
has full strength, then reduce the strength of the force so it can cross barriers, and then
turn it back on again.
Expanded Ensemble
-----------------
The expanded ensemble method is closely related to replica exchange, but instead of
simulating many replicas and exchanging thermodynamic states between them, it simulates a
single trajectory that jumps between states. When the states correspond to temperatures,
this method is also known as simulated tempering.
Simulated tempering\ :cite:`Marinari1992` involves making frequent changes to the Expanded ensemble sampling can be more efficient than replica exchange, but also requires a
temperature of a simulation. At high temperatures, it can quickly cross energy barriers bit more care to ensure it samples all thermodynamic states evenly.
and explore a wide range of configurations. At lower temperatures, it more thoroughly
explores each local region of configuration space. This is a powerful method to
speed up sampling when you do not know in advance what motions you want to sample.
Simply specify the range of temperatures to simulate and the algorithm handles
everything for you mostly automatically.
Metadynamics Metadynamics
------------ ------------
......
...@@ -6,7 +6,7 @@ simulatedtempering.py: Implements simulated tempering ...@@ -6,7 +6,7 @@ simulatedtempering.py: Implements simulated tempering
This is part of the OpenMM molecular simulation toolkit. This is part of the OpenMM molecular simulation toolkit.
See https://openmm.org/development. See https://openmm.org/development.
Portions copyright (c) 2015 Stanford University and the Authors. Portions copyright (c) 2015-2026 Stanford University and the Authors.
Authors: Peter Eastman Authors: Peter Eastman
Contributors: Contributors:
...@@ -32,28 +32,14 @@ __author__ = "Peter Eastman" ...@@ -32,28 +32,14 @@ __author__ = "Peter Eastman"
__version__ = "1.0" __version__ = "1.0"
import openmm.unit as unit import openmm.unit as unit
import math
import random
from sys import stdout from sys import stdout
try:
import bz2
have_bz2 = True
except: have_bz2 = False
try:
import gzip
have_gzip = True
except: have_gzip = False
try:
import numpy
have_numpy = True
except: have_numpy = False
class SimulatedTempering(object): class SimulatedTempering(object):
"""SimulatedTempering implements the simulated tempering algorithm for accelerated sampling. """SimulatedTempering implements the simulated tempering algorithm for accelerated sampling.
@deprecated This class exists mainly for historical reasons. It now is just a thin wrapper around ExpandedEnsembleSampler. It is still supported for backward compatibility, but using ExpandedEnsembleSampler directly is recommended since it is much more flexible.
It runs a simulation while allowing the temperature to vary. At high temperatures, it can more easily cross It runs a simulation while allowing the temperature to vary. At high temperatures, it can more easily cross
energy barriers to explore a wider area of conformation space. At low temperatures, it can thoroughly energy barriers to explore a wider area of conformation space. At low temperatures, it can thoroughly
explore each local region. For details, see Marinari, E. and Parisi, G., Europhys. Lett. 19(6). pp. 451-458 (1992). explore each local region. For details, see Marinari, E. and Parisi, G., Europhys. Lett. 19(6). pp. 451-458 (1992).
...@@ -106,160 +92,27 @@ class SimulatedTempering(object): ...@@ -106,160 +92,27 @@ class SimulatedTempering(object):
The interval (in time steps) at which to write information to the report file The interval (in time steps) at which to write information to the report file
reportFile: string or file reportFile: string or file
The file to write reporting information to, specified as a file name or file object The file to write reporting information to, specified as a file name or file object
""" """
self.simulation = simulation
if temperatures is None: if temperatures is None:
if unit.is_quantity(minTemperature): if unit.is_quantity(minTemperature):
minTemperature = minTemperature.value_in_unit(unit.kelvin) minTemperature = minTemperature.value_in_unit(unit.kelvin)
if unit.is_quantity(maxTemperature): if unit.is_quantity(maxTemperature):
maxTemperature = maxTemperature.value_in_unit(unit.kelvin) maxTemperature = maxTemperature.value_in_unit(unit.kelvin)
self.temperatures = [minTemperature*((float(maxTemperature)/minTemperature)**(i/float(numTemperatures-1))) for i in range(numTemperatures)]*unit.kelvin temperatures = [minTemperature*((float(maxTemperature)/minTemperature)**(i/float(numTemperatures-1))) for i in range(numTemperatures)]*unit.kelvin
else: self.temperatures = [(t.value_in_unit(unit.kelvin) if unit.is_quantity(t) else t)*unit.kelvin for t in temperatures]
numTemperatures = len(temperatures) states = [{'temperature':t} for t in temperatures]
self.temperatures = [(t.value_in_unit(unit.kelvin) if unit.is_quantity(t) else t)*unit.kelvin for t in temperatures] from . import ExpandedEnsembleSampler
if any(self.temperatures[i] >= self.temperatures[i+1] for i in range(numTemperatures-1)): self._sampler = ExpandedEnsembleSampler(states, simulation, tempChangeInterval, weights=weights,
raise ValueError('The temperatures must be in strictly increasing order') reportInterval=reportInterval, logFile=reportFile)
self.tempChangeInterval = tempChangeInterval
self.reportInterval = reportInterval
self.inverseTemperatures = [1.0/(unit.MOLAR_GAS_CONSTANT_R*t) for t in self.temperatures]
# If necessary, open the file we will write reports to.
self._openedFile = isinstance(reportFile, str)
if self._openedFile:
# Detect the desired compression scheme from the filename extension
# and open all files unbuffered
if reportFile.endswith('.gz'):
if not have_gzip:
raise RuntimeError("Cannot write .gz file because Python could not import gzip library")
self._out = gzip.GzipFile(fileobj=open(reportFile, 'wb', 0))
elif reportFile.endswith('.bz2'):
if not have_bz2:
raise RuntimeError("Cannot write .bz2 file because Python could not import bz2 library")
self._out = bz2.BZ2File(reportFile, 'w', 0)
else:
self._out = open(reportFile, 'w', 1)
else:
self._out = reportFile
# Initialize the weights.
if weights is None:
self._weights = [0.0]*numTemperatures
self._updateWeights = True
self._weightUpdateFactor = 1.0
self._histogram = [0]*numTemperatures
self._hasMadeTransition = False
else:
self._weights = weights
self._updateWeights = False
# Select the initial temperature.
self.currentTemperature = 0
self.simulation.integrator.setTemperature(self.temperatures[self.currentTemperature])
for param in self.simulation.context.getParameters():
if 'MonteCarloTemperature' in param:
self.simulation.context.setParameter(param, self.temperatures[self.currentTemperature])
# Add a reporter to the simulation which will handle the updates and reports.
class STReporter(object):
def __init__(self, st):
self.st = st
def describeNextReport(self, simulation):
st = self.st
steps1 = st.tempChangeInterval - simulation.currentStep%st.tempChangeInterval
steps2 = st.reportInterval - simulation.currentStep%st.reportInterval
steps = min(steps1, steps2)
isUpdateAttempt = (steps1 == steps)
if isUpdateAttempt:
return {'steps': steps, 'periodic':None, 'include':['velocities', 'energy']}
else:
return {'steps': steps, 'periodic':None, 'include':[]}
def report(self, simulation, state):
st = self.st
if simulation.currentStep%st.tempChangeInterval == 0:
st._attemptTemperatureChange(state)
if simulation.currentStep%st.reportInterval == 0:
st._writeReport()
simulation.reporters.append(STReporter(self))
# Write out the header line.
headers = ['Steps', 'Temperature (K)']
for t in self.temperatures:
headers.append('%gK Weight' % t.value_in_unit(unit.kelvin))
print('#"%s"' % ('"\t"').join(headers), file=self._out)
def __del__(self):
if self._openedFile:
self._out.close()
@property @property
def weights(self): def weights(self):
return [x-self._weights[0] for x in self._weights] return self._sampler.weights
@property
def currentTemperature(self):
return self._sampler.currentStateIndex
def step(self, steps): def step(self, steps):
"""Advance the simulation by integrating a specified number of time steps.""" """Advance the simulation by integrating a specified number of time steps."""
self.simulation.step(steps) self._sampler.step(steps)
def _attemptTemperatureChange(self, state):
"""Attempt to move to a different temperature."""
# Compute the probability for each temperature. This is done in log space to avoid overflow.
logProbability = [(self._weights[i]-self.inverseTemperatures[i]*state.getPotentialEnergy()) for i in range(len(self._weights))]
maxLogProb = max(logProbability)
offset = maxLogProb + math.log(sum(math.exp(x-maxLogProb) for x in logProbability))
probability = [math.exp(x-offset) for x in logProbability]
r = random.random()
for j in range(len(probability)):
if r < probability[j]:
if j != self.currentTemperature:
# Rescale the velocities.
scale = math.sqrt(self.temperatures[j]/self.temperatures[self.currentTemperature])
if have_numpy:
velocities = scale*state.getVelocities(asNumpy=True).value_in_unit(unit.nanometers/unit.picoseconds)
else:
velocities = [v*scale for v in state.getVelocities().value_in_unit(unit.nanometers/unit.picoseconds)]
self.simulation.context.setVelocities(velocities)
# Select this temperature.
self._hasMadeTransition = True
self.currentTemperature = j
self.simulation.integrator.setTemperature(self.temperatures[j])
for param in self.simulation.context.getParameters():
if 'MonteCarloTemperature' in param:
self.simulation.context.setParameter(param, self.temperatures[self.currentTemperature])
if self._updateWeights:
# Update the weight factors.
self._weights[j] -= self._weightUpdateFactor
self._histogram[j] += 1
minCounts = min(self._histogram)
if minCounts > 20 and minCounts >= 0.2*sum(self._histogram)/len(self._histogram):
# Reduce the weight update factor and reset the histogram.
self._weightUpdateFactor *= 0.5
self._histogram = [0]*len(self.temperatures)
self._weights = [x-self._weights[0] for x in self._weights]
elif not self._hasMadeTransition and probability[self.currentTemperature] > 0.99:
# Rapidly increase the weight update factor at the start of the simulation to find
# a reasonable starting value.
self._weightUpdateFactor *= 2.0
self._histogram = [0]*len(self.temperatures)
return
r -= probability[j]
def _writeReport(self):
"""Write out a line to the report."""
temperature = self.temperatures[self.currentTemperature].value_in_unit(unit.kelvin)
values = [temperature]+self.weights
print(('%d\t' % self.simulation.currentStep) + '\t'.join('%g' % v for v in values), file=self._out)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment