Created ExpandedEnsembleSampler (#5265)

* Created ExpandedEnsembleSampler

* Attempt at fixing test errors on Windows

* Another attempt at fixing test errors on Windows

* More output options

* Minor fixes

* Still trying to fix Windows errors

* Debugging

* Just skip the test on Windows

* Fix error on older Python

wrappers/python/openmm/app/__init__.py

@@ -39,6 +39,7 @@ from .simulatedtempering import SimulatedTempering
 from .metadynamics import Metadynamics, BiasVariable
 from .replicaexchangesampler import ReplicaExchangeSampler
 from .replicaexchangereporter import ReplicaExchangeReporter
+from .expandedensemblesampler import ExpandedEnsembleSampler
 
 # Enumerated values
 

wrappers/python/openmm/app/expandedensemblesampler.py

+from __future__ import print_function
+
+"""
+expandedensemblesampler.py: Performs multistate sampling with the expanded ensemble method
+
+This is part of the OpenMM molecular simulation toolkit.
+See https://openmm.org/development.
+
+Portions copyright (c) 2015-2026 Stanford University and the Authors.
+Authors: Peter Eastman
+Contributors:
+
+Permission is hereby granted, free of charge, to any person obtaining a 
+copy of this software and associated documentation files (the "Software"),
+to deal in the Software without restriction, including without limitation
+the rights to use, copy, modify, merge, publish, distribute, sublicense,
+and/or sell copies of the Software, and to permit persons to whom the
+Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
+USE OR OTHER DEALINGS IN THE SOFTWARE.
+"""
+__author__ = "Peter Eastman"
+__version__ = "1.0"
+
+import openmm
+import openmm.unit as unit
+from openmm.app.internal import safesave
+from openmm.app.internal.multistatesampler import MultistateSampler
+import math
+import pickle
+import random
+
+class ExpandedEnsembleSampler(object):
+    """
+    ExpandedEnsembleSampler uses the expanded ensemble method to simulate a system in a collection of thermodynamic
+    states.  It supports both the temperature version, in which the states correspond to different temperatures, and the
+    Hamiltonian version, in which they correspond to different potential functions.  It also can combine them to vary
+    temperature and potential function at the same time.
+
+    You provide a Simulation describing the system to simulate and a list of states.  It simulates the system while
+    periodically moving between states in a way that ensures correct sampling.  The method is based on Gibbs sampling,
+    which alternates sampling of conformations x and thermodynamic states s.  It first performs molecular dynamics to
+    sample P(x|s), the distribution of conformations for a fixed state.  It then chooses a new state from the
+    distribution P(s|x) for the current conformation.  For more information on this method, see
+    https://doi.org/10.33011/livecoms.4.1.1583.
+
+    Expanded ensemble sampling is used for a variety of purposes.  Here are some of the more common ones.
+
+    - The temperature version is most often used to accelerate sampling of rare events.  A particular transition might
+      only happen rarely at physiological temperature.  You therefore simulate both the temperature you care about and
+      also higher temperatures at which the transition happens more easily.  Allowing it to move between temperatures
+      produces accurate sampling at low temperature of the entire phase space.
+
+    - The Hamiltonian version may be used to sample an alchemical transition between two endpoints.  It produces
+      efficient sampling at every point along the transition pathway from which a free energy difference can be computed.
+
+    Each thermodynamic state (not to be confused with a State object) is represented by a dict containing property
+    values.  All states must specify values for the same properties.  They describe the ways in which the states differ
+    from each other.  The following properties are supported.
+
+    - 'temperature': the simulation temperature
+    - 'groups': a set containing the force groups to include when computing the energy, for example {0, 2}.  It also may
+      be a weighted sum specified as a dict.  For example, {0:1.0, 2:0.5} means to compute the energy of group 0 plus
+      half the energy of group 2.
+    - Context parameters, specified by the parameter name
+    - Global variables defined by a CustomIntegrator, specified by the variable name
+
+    For example, a typical use of temperature sampling might define the states as
+
+    >>> states = [{'temperature':t*kelvin} for t in np.geomspace(300.0, 500.0, 10)]
+
+    A typical use of Hamiltonian sampling might include forces that depend on a parameter `lambda`.  The states could
+    be defined as
+
+    >>> states = [{'lambda':x} for x in np.linspace(0.0, 1.0, 11)]
+
+    States may also specify multiple properties.  This example includes all combinations of two different parameters.
+
+    >>> states = [{'lambda1':x, 'lambda2':y} for x in np.linspace(0.0, 1.0, 6) for y in np.linspace(0.0, 1.0, 6)]
+
+    To run an expanded ensemble simulation, first create a Simulation object in the normal way.  Then create a
+    ExpandedEnsembleSampler:
+
+    >>> sampler = ExpandedEnsembleSampler(states, simulation, stepsPerIteration)
+
+    The third argument is the number of time steps to integrate between attempted state changes.  Then perform a
+    simulation by calling simulation.step() in the usual way.  The ExpandedEnsembleSampler adds a reporter to the
+    Simulation that manages state changes as it runs.
+
+    The probability distribution P(s|x) is often very nonuniform, which can lead to the simulation spending much more
+    time in some states than others.  This is addressed by including a weight factor for each state when computing the
+    probability.  Ideally the weights should be chosen so it spends equal time in every state.  You can specify the
+    weights yourself with a constructor argument, but usually it is easier to let the ExpandedEnsembleSampler choose
+    them automatically.  This is done with a combination of the Wang-Landau (https://doi.org/10.1103/PhysRevLett.86.2050)
+    and Self-Adjusted Mixture Sampling (http://dx.doi.org/10.1080/10618600.2015.1113975) algorithms.  You should monitor
+    the weights and discard the initial part of the simulation where they are changing rapidly.  Until the weights have
+    converged, the simulation will not sample the correct distribution.
+
+    An ExpandedEnsembleSampler can act as a reporter, generating output to files at regular intervals.  The reporting
+    interval should generally be the same as the one used for writing a trajectory so they will be synchronized with
+    each other.  The following files can optionally be written.
+
+    - A log file that records the current thermodynamic state and weights.  To analyze the results of a simulation, it
+    is essential to know what state it was in at every point in time.
+    - A file recording the current reduced energy (potential energy divided by kT) of every thermodynamic state.  This
+    information is useful for calculating free energies.
+    - A checkpoint file containing all information necessary to resume the simulation.  This includes internal fields
+    of the ExpandedEnsembleSampler itself, as well as a State object recording the current positions, velocities,
+    parameters, etc.
+
+    To resume a simulation from the saved checkpoint, pass `resume=True` to the constructor.  It will load all necessary
+    information and configure the ExpandedEnsembleSampler correctly.
+
+    Attributes
+    ----------
+    states: list[dict]
+        The states to sample
+    simulation: openmm.app.Simulation
+        The Simulation defining the System, Context, and Integrator to use
+    stepsPerIteration: int
+        The number of time steps to integrate between attempted state changes
+    reinitializeVelocities: bool
+        If true, the simulation's velocities are reinitialized from a Boltzmann distribution every time its state
+        changes.  This may sometimes improve stability, but also decreases efficiency.
+    currentIteration: int
+        The number of iterations that have been completed, each consisting of stepsPerIteration time steps followed by
+        an attempted state change
+    currentStateIndex: int
+        The current state of the simulation, specified as an index into states.
+    """
+
+    def __init__(self, states: list[dict], simulation: "openmm.app.Simulation", stepsPerIteration: int,
+                 reinitializeVelocities: bool = False, weights: list[float] | None = None, reportInterval: int = 1000,
+                 logFile: str | object | None = None, energyFile: str | object | None = None,
+                 checkpointFile: str | None = None, resume: bool = False):
+        """Create a new ExpandedEnsembleSampler.
+
+        Parameters
+        ----------
+        states: list[dict]
+            The states to sample
+        simulation: openmm.app.Simulation
+            The Simulation defining the System, Context, and Integrator to use
+        stepsPerIteration: int
+            The number of time steps to integrate between attempted state changes
+        reinitializeVelocities: bool
+            If true, the simulation's velocities are reinitialized from a Boltzmann distribution every time its state
+            changes.  This may sometimes improve stability, but also decreases efficiency.
+        weights: list[float] | None
+            The weights to use for each state.  If None, weights are chosen automatically as the simulation runs.
+        reportInterval: int
+            The frequency at which to write output, measured in time steps.  This must be a multiple of
+            stepsPerIteration.
+        logFile: str | object | None
+            An optional file to write a log to.  This may be either a file-like object or a string containing the path
+            to the file.
+        energyFile: str | object | None
+            An optional file to write reduced energies to.  This may be either a file-like object or a string containing
+            the path to the file.
+        checkpointFile: str | None
+            The path to an optional file for saving checkpointing information
+        resume: bool
+            Specifies whether to resume an earlier simulation.  If True, the checkpoint will be loaded and future output
+            will be appended to the existing files.
+        """
+        self.states = states
+        self.simulation = simulation
+        self.stepsPerIteration = stepsPerIteration
+        self.reinitializeVelocities = reinitializeVelocities
+        self.currentIteration = 0
+        self._stage = 0
+        self.currentStateIndex = 0
+        self._sampler = MultistateSampler(states, simulation.context)
+        if 'temperature' in states[0]:
+            temperature = [s['temperature'] for s in states]
+            for i, t in enumerate(temperature):
+                if not unit.is_quantity(t):
+                    temperature[i] = t*unit.kelvin
+            self._kT = [unit.MOLAR_GAS_CONSTANT_R*t for t in temperature]
+        else:
+            self._kT = None
+            if not hasattr(simulation.integrator, 'getTemperature'):
+                raise ValueError('Cannot determine temperature because the integrator does not have a getTemperature() method. '
+                                 'Specify the temperature in each state dict.')
+        self.reportInterval = reportInterval
+
+        # Initialize the weights.
+
+        if weights is None:
+            self._weights = [0.0]*len(states)
+            self._updateWeights = True
+            self._weightUpdateFactor = 1.0
+            self._histogram = [0]*len(states)
+            self._hasMadeTransition = False
+        else:
+            self._weights = weights
+            self._updateWeights = False
+
+        # Restore from the checkpoint file.
+
+        if resume:
+            if checkpointFile is None:
+                raise ValueError('resume=True, but no checkpoint file was specified to restore from.')
+            with open(checkpointFile, 'rb') as input:
+                checkpoint = pickle.load(input)
+                self._applyCheckpoint(checkpoint)
+
+        # Apply the current state to the simulation.
+
+        self._sampler.applyState(self.currentStateIndex)
+
+        # Open output files.
+
+        if (logFile is not None or energyFile is not None) and reportInterval%stepsPerIteration != 0:
+            raise ValueError('The reporting interval must be a multiple iteration length.')
+        mode = 'a' if resume else 'w'
+        self._openedLogFile = isinstance(logFile, str)
+        if self._openedLogFile:
+            self._log = open(logFile, mode, 1)
+        else:
+            self._log = logFile
+        self._openedEnergyFile = isinstance(energyFile, str)
+        if self._openedEnergyFile:
+            self._energy = open(energyFile, mode, 1)
+        else:
+            self._energy = energyFile
+        self._checkpointFile = checkpointFile
+
+        # Add a reporter to the simulation which will handle the updates and reports.
+
+        class EEReporter(object):
+            def __init__(self, sampler):
+                self.sampler = sampler
+
+            def describeNextReport(self, simulation):
+                sampler = self.sampler
+                steps1 = sampler.stepsPerIteration - simulation.currentStep%sampler.stepsPerIteration
+                steps2 = sampler.reportInterval - simulation.currentStep%sampler.reportInterval
+                steps = min(steps1, steps2)
+                isUpdateAttempt = (steps1 == steps)
+                if isUpdateAttempt and not sampler.reinitializeVelocities:
+                    return {'steps': steps, 'periodic':None, 'include':['velocities']}
+                else:
+                    return {'steps': steps, 'periodic':None, 'include':[]}
+
+            def report(self, simulation, state):
+                sampler = self.sampler
+                if simulation.currentStep%sampler.stepsPerIteration == 0:
+                    reducedEnergy = sampler.attemptStateChange(state)
+                    if simulation.currentStep%sampler.reportInterval == 0:
+                        sampler._writeReport(reducedEnergy)
+
+        simulation.reporters.append(EEReporter(self))
+
+        # Write header lines to the output files.
+
+        if not resume:
+            if self._log is not None:
+                headers = ['"Steps"', '"Iteration"', '"State"'] + [f'"Weight {i}"' for i in range(len(self.states))]
+                print('%s' % (',').join(headers), file=self._log)
+            if self._energy is not None:
+                headers = ['"Steps"'] + [f'"u{i}"' for i in range(len(self.states))]
+                print('%s' % (',').join(headers), file=self._energy)
+
+    def __del__(self):
+        if self._openedLogFile:
+            self._log.close()
+        if self._openedEnergyFile:
+            self._energy.close()
+
+    @property
+    def weights(self):
+        return [x-self._weights[0] for x in self._weights]
+
+    def step(self, steps: int):
+        """Advance the simulation by integrating a specified number of time steps."""
+        self.simulation.step(steps)
+
+    def attemptStateChange(self, state: openmm.State) -> list[float]:
+        """Attempt to move to a different state.  This is called automatically during the simulation, and there is not
+        normally a reason to call it directly.
+
+        You can create subclasses that override this method to select states in different ways."""
+
+        # Compute the probability for each state.  This is done in log space to avoid overflow.
+
+        self.currentIteration += 1
+        energies = self._sampler.computeAllEnergies()
+        if self._kT is None:
+            kT = [unit.MOLAR_GAS_CONSTANT_R*self.simulation.integrator.getTemperature()]*len(self.states)
+        else:
+            kT = self._kT
+        reducedEnergy = [energies[i]/kT[i] for i in range(len(self._weights))]
+        logProbability = [(self._weights[i]-reducedEnergy[i]) 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]
+
+        # Select a new state.
+
+        prevState = self.currentStateIndex
+        r = random.random()
+        for j in range(len(probability)):
+            if r < probability[j]:
+                if j != self.currentStateIndex:
+                    # Select this state.
+
+                    self._hasMadeTransition = True
+                    self.currentStateIndex = j
+                if self._updateWeights:
+                    # Update the weights.
+
+                    updateFactor = 1/self.currentIteration
+                    if self._stage == 0:
+                        if self._weightUpdateFactor >= updateFactor:
+                            updateFactor = self._weightUpdateFactor
+                        else:
+                            self._stage = 1
+                    for k in range(len(probability)):
+                        self._weights[k] -= updateFactor*probability[k]
+                    if self._stage == 0:
+                        # Early in the simulation we reduce the update factor when the histogram is sufficiently flat.
+                        # Later we will switch over to just using 1/iteration.
+
+                        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.states)
+                            self._weights = [x-self._weights[0] for x in self._weights]
+                        elif not self._hasMadeTransition and probability[self.currentStateIndex] > 0.99 and sum(self._histogram) >= 10:
+                            # 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.states)
+                break
+            r -= probability[j]
+
+        # Apply the state.
+
+        self._sampler.applyState(self.currentStateIndex)
+
+        # Reinitialize or rescale the velocities.
+
+        if self.reinitializeVelocities:
+            if 'temperature' in self.states[self.currentStateIndex]:
+                temperature = self.states[self.currentStateIndex]['temperature']
+            else:
+                temperature = self.simulation.integrator.getTemperature()
+            self.simulation.context.setVelocitiesToTemperature(temperature)
+        elif kT[prevState] != kT[self.currentStateIndex]:
+            scale = math.sqrt(kT[self.currentStateIndex]/kT[prevState])
+            self.simulation.context.setVelocities(scale*state.getVelocities(asNumpy=True))
+        return reducedEnergy
+
+    def _writeReport(self, reducedEnergy: list[float]):
+        """Write output to the files."""
+        if self._log is not None:
+            print(f'{self.simulation.currentStep},{self.currentIteration},{self.currentStateIndex},' + ','.join('%g' % v for v in self.weights), file=self._log)
+        if self._energy is not None:
+            print(f'{self.simulation.currentStep},' + ','.join('%g' % v for v in reducedEnergy), file=self._energy)
+        if self._checkpointFile is not None:
+            checkpoint = self._createCheckpoint()
+            safesave.save(pickle.dumps(checkpoint), self._checkpointFile)
+
+    def _createCheckpoint(self) -> dict:
+        """Create a dict containing the information to save to a checkpoint file."""
+        checkpoint = {}
+        checkpoint['currentIteration'] = self.currentIteration
+        checkpoint['stage'] = self._stage
+        checkpoint['currentStateIndex'] = self.currentStateIndex
+        checkpoint['weights'] = self._weights
+        checkpoint['weightUpdateFactor'] = self._weightUpdateFactor
+        checkpoint['histogram'] = self._histogram
+        checkpoint['hasMadeTransition'] = self._hasMadeTransition
+        checkpoint['state'] = self.simulation.context.getState(positions=True, velocities=True, parameters=True, integratorParameters=True)
+        return checkpoint
+
+    def _applyCheckpoint(self, checkpoint: dict):
+        """Record the information that was loaded from a checkpoint file."""
+        self.currentIteration = checkpoint['currentIteration']
+        self._stage = checkpoint['stage']
+        self.currentStateIndex = checkpoint['currentStateIndex']
+        self._weights = checkpoint['weights']
+        self._weightUpdateFactor = checkpoint['weightUpdateFactor']
+        self._histogram = checkpoint['histogram']
+        self._hasMadeTransition = checkpoint['hasMadeTransition']
+        self.simulation.context.setState(checkpoint['state'])

wrappers/python/openmm/app/internal/multistatesampler.py

@@ -209,7 +209,7 @@ class MultistateSampler(object):
         -------
         an array containing the potential energies of all states in the order they appear in self.states
         """
-        energies = [0*kilojoules_per_mole for _ in self.states]
+        energies = [0 for _ in self.states]*kilojoules_per_mole
         for i, subset in enumerate(self.subsets):
             if self.groups is None:
                 # States don't depend on force groups, so we can just evaluate the energy of one state.
@@ -241,6 +241,6 @@ class MultistateSampler(object):
         an array containing the shifted potential energies of all states in the order they appear in self.states
         """
         if len(self.subsets) == 1 and self.groups is None:
-            return [0*kilojoules_per_mole]*len(self.states)
+            return [0]*len(self.states)*kilojoules_per_mole
         energies = self.computeAllEnergies()
         return [e-energies[0] for e in energies]
\ No newline at end of file

wrappers/python/openmm/app/internal/safesave.py

@@ -4,9 +4,9 @@ safesave.py: Helper module to ensure atomic overwrite/backup of existing files.
 This is part of the OpenMM molecular simulation toolkit.
 See https://openmm.org/development.
 
-Portions copyright (c) 2025 Stanford University and the Authors.
+Portions copyright (c) 2025-2026 Stanford University and the Authors.
 Authors: Evan Pretti
-Contributors:
+Contributors: Peter Eastman
 
 Permission is hereby granted, free of charge, to any person obtaining a
 copy of this software and associated documentation files (the "Software"),
@@ -51,11 +51,8 @@ def _getTempFilename(prefix):
 
     for index in itertools.count():
         name = f'{prefix}.{index}.tmp'
-        try:
-            with open(name, 'x'):
-                return name
-        except FileExistsError:
-            pass
+        if not os.path.exists(name):
+            return name
 
 def save(data, filename):
     """

wrappers/python/tests/TestExpandedEnsembleSampler.py

+from openmm import *
+from openmm.app import *
+from openmm.unit import *
+import numpy as np
+import os
+import tempfile
+import unittest
+
+class TestExpandedEnsembleSampler(unittest.TestCase):
+    def testTemperature(self):
+        """Test a set of states that differ in temperature."""
+        system = System()
+        system.addParticle(1.0)
+        force = CustomExternalForce('x*x+y*y+z*z')
+        force.addParticle(0)
+        system.addForce(force)
+        states = [{'temperature':t*kelvin} for t in np.geomspace(300.0, 600.0, 5)]
+        for reinitialize in [False, True]:
+            integrator = LangevinIntegrator(300*kelvin, 10/picosecond, 0.01*picosecond)
+            simulation = Simulation(Topology(), system, integrator, Platform.getPlatform('Reference'))
+            simulation.context.setPositions([Vec3(0, 0, 0)])
+            sampler = ExpandedEnsembleSampler(states, simulation, 10, reinitialize)
+
+            # Run for a little while to let the weights stabilize.
+
+            sampler.step(10000)
+
+            # Run for a while and record the states and energies.
+
+            energies = [[] for _ in range(len(states))]
+            iterations = 20000
+            for i in range(iterations):
+                sampler.step(10)
+                energies[sampler.currentStateIndex].append(simulation.context.getState(energy=True).getPotentialEnergy())
+
+            # Check that it spent roughly equal time in each state, and that the energies are correct.
+
+            for energy, state in zip(energies, states):
+                n = len(energy)
+                assert iterations/10 < n < iterations/2
+                average = sum(energy)/n
+                expected = 1.5*(state['temperature']*MOLAR_GAS_CONSTANT_R)
+                self.assertTrue(0.7 < average/expected < 1.3)
+
+    def testParameter(self):
+        """Test a set of states that differ in a force parameter."""
+        system = System()
+        system.addParticle(1.0)
+        force = CustomExternalForce('0.5*k*x*x')
+        force.addGlobalParameter('k', 1.0)
+        force.addParticle(0)
+        system.addForce(force)
+        states = [{'k':k*kilojoules_per_mole/(nanometer**2)} for k in np.geomspace(10.0, 100.0, 5)]
+        for reinitialize in [False, True]:
+            integrator = LangevinIntegrator(300*kelvin, 10/picosecond, 0.01*picosecond)
+            simulation = Simulation(Topology(), system, integrator, Platform.getPlatform('Reference'))
+            simulation.context.setPositions([Vec3(0, 0, 0)])
+            sampler = ExpandedEnsembleSampler(states, simulation, 10, reinitialize)
+
+            # Run for a little while to let the weights stabilize.
+
+            sampler.step(10000)
+
+            # Run for a while and record the states and displacements.
+
+            r2 = [[] for _ in range(len(states))]
+            iterations = 20000
+            for i in range(iterations):
+                sampler.step(10)
+                x = simulation.context.getState(positions=True).getPositions()[0][0]
+                r2[sampler.currentStateIndex].append(x*x)
+
+            # Check that it spent roughly equal time in each state, and that the energies are correct.
+
+            expected = 0.5*integrator.getTemperature()*MOLAR_GAS_CONSTANT_R
+            for i in range(len(r2)):
+                n = len(r2[i])
+                assert iterations/10 < n < iterations/2
+                average = 0.5*states[i]['k']*sum(r2[i])/n
+                self.assertTrue(0.7 < average/expected < 1.3)
+
+    def testReporter(self):
+        """Test reporting output from an expanded ensemble simulation."""
+        system = System()
+        force = CustomExternalForce('0.5*k*(x*x+y*y+z*z)')
+        force.addGlobalParameter('k', 1.0)
+        system.addForce(force)
+        for i in range(3):
+            system.addParticle(1.0)
+            force.addParticle(0)
+        states = [{'k':k} for k in (200.0, 300.0, 400.0)]
+        with tempfile.NamedTemporaryFile(mode='w', delete=False) as logFile:
+            with tempfile.NamedTemporaryFile(mode='w', delete=False) as energyFile:
+                with tempfile.NamedTemporaryFile(mode='w', delete=False) as checkpointFile:
+                    integrator = LangevinIntegrator(300*kelvin, 1/picosecond, 0.001*picosecond)
+                    simulation = Simulation(Topology(), system, integrator, Platform.getPlatform('Reference'))
+                    simulation.context.setPositions([Vec3(0, 0, 0)]*3)
+                    sampler = ExpandedEnsembleSampler(states, simulation, 5, reportInterval=5, logFile=logFile.name,
+                                                      energyFile=energyFile.name, checkpointFile=checkpointFile.name)
+
+                    # Run a simulation.
+
+                    step = []
+                    iteration = []
+                    stateIndex = []
+                    weights = []
+                    energies = []
+
+                    def runIteration():
+                        simulation.step(5)
+                        step.append(simulation.currentStep)
+                        iteration.append(sampler.currentIteration)
+                        stateIndex.append(sampler.currentStateIndex)
+                        weights.append(sampler.weights)
+                        kT = MOLAR_GAS_CONSTANT_R*simulation.integrator.getTemperature()
+                        energies.append(sampler._sampler.computeAllEnergies()/kT)
+                        sampler._sampler.applyState(sampler.currentStateIndex)
+
+                    try:
+                        for _ in range(4):
+                            runIteration()
+                    except PermissionError:
+                        # tempfile is kind of broken on Windows.  Just skip the test.
+                        return
+                    state1 = simulation.context.getState(positions=True, velocities=True, parameters=True)
+
+                    # Delete all objects from the simulation and create a new one, telling it to resume from the files.
+
+                    del sampler
+                    del simulation
+                    del integrator
+                    integrator = LangevinIntegrator(300*kelvin, 1/picosecond, 0.001*picosecond)
+                    simulation = Simulation(Topology(), system, integrator, Platform.getPlatform('Reference'))
+                    sampler = ExpandedEnsembleSampler(states, simulation, 5, reportInterval=5, logFile=logFile.name,
+                                                      energyFile=energyFile.name, checkpointFile=checkpointFile.name,
+                                                      resume=True)
+
+                    # Make sure everything was loaded correctly.
+
+                    state2 = simulation.context.getState(positions=True, velocities=True, parameters=True)
+                    self.assertEqual(XmlSerializer.serialize(state1), XmlSerializer.serialize(state2))
+                    self.assertEqual(step[-1], simulation.currentStep)
+                    self.assertEqual(iteration[-1], sampler.currentIteration)
+                    self.assertEqual(stateIndex[-1], sampler.currentStateIndex)
+                    self.assertEqual(weights[-1], sampler.weights)
+
+                    # Generate some more output.
+
+                    for _ in range(4):
+                        runIteration()
+
+                    # Check the log file.
+
+                    logFile.close()
+                    with open(logFile.name) as input:
+                        lines = input.readlines()[1:]
+                    os.remove(logFile.name)
+                    self.assertEqual(8, len(lines))
+                    for i, line in enumerate(lines):
+                        fields = line.split(',')
+                        self.assertEqual(int(fields[0]), step[i])
+                        self.assertEqual(int(fields[1]), iteration[i])
+                        self.assertEqual(int(fields[2]), stateIndex[i])
+                        self.assertTrue(np.allclose([float(x) for x in fields[3:]], weights[i]))
+
+                    # Check the energy file.
+
+                    energyFile.close()
+                    with open(energyFile.name) as input:
+                        lines = input.readlines()[1:]
+                    os.remove(energyFile.name)
+                    self.assertEqual(8, len(lines))
+                    for i, line in enumerate(lines):
+                        fields = line.split(',')
+                        self.assertEqual(int(fields[0]), step[i])
+                        self.assertTrue(np.allclose([float(x) for x in fields[1:]], energies[i]))

wrappers/python/tests/TestReplicaExchangeSampler.py

@@ -19,6 +19,7 @@ class TestReplicaExchangeSampler(unittest.TestCase):
         for reinitialize in [False, True]:
             integrator = LangevinIntegrator(300*kelvin, 10/picosecond, 0.01*picosecond)
             simulation = Simulation(Topology(), system, integrator, Platform.getPlatform('Reference'))
+            simulation.context.setPositions([Vec3(0, 0, 0)])
             repex = ReplicaExchangeSampler(states, simulation, 20, reinitialize)
             energies = [0.0*kilojoules_per_mole]*len(states)
             exchanged = False
@@ -51,10 +52,11 @@ class TestReplicaExchangeSampler(unittest.TestCase):
         force.addGlobalParameter('k', 1.0)
         force.addParticle(0)
         system.addForce(force)
-        states = [{'k':k*kilojoules_per_mole/(nanometer**2)} for k in np.geomspace(5.0, 100.0, 5)]
+        states = [{'k':k*kilojoules_per_mole/(nanometer**2)} for k in np.geomspace(10.0, 100.0, 5)]
         for reinitialize in [False, True]:
             integrator = LangevinIntegrator(300*kelvin, 10/picosecond, 0.01*picosecond)
             simulation = Simulation(Topology(), system, integrator, Platform.getPlatform('Reference'))
+            simulation.context.setPositions([Vec3(0, 0, 0)])
             repex = ReplicaExchangeSampler(states, simulation, 20, reinitialize)
             r2 = [0.0*nanometer**2]*len(states)
             exchanged = False
@@ -135,7 +137,8 @@ class TestReplicaExchangeSampler(unittest.TestCase):
 
             # Check the log file.
 
-            lines = open(os.path.join(directory, 'log.csv')).readlines()[1:]
+            with open(os.path.join(directory, 'log.csv')) as input:
+                lines = input.readlines()[1:]
             for i, line in enumerate(lines):
                 fields = [int(x) for x in line.split(',')]
                 self.assertEqual(fields[0], 3*(i+1))
@@ -176,4 +179,7 @@ class TestReplicaExchangeSampler(unittest.TestCase):
             # Creating a new reporter for the same directory should fail.
 
             with self.assertRaises(ValueError):
-                ReplicaExchangeReporter(directory, 3, sampler)
\ No newline at end of file
+                ReplicaExchangeReporter(directory, 3, sampler)
+            del sampler
+            del simulation
+            del integrator