def _broadcast_database(self):
"""Load the positions, replica_states, u_kl, proposed, and accepted from root node database."""
if self.mpicomm.rank == 0:
positions = self.database.last_positions
replica_states = self.database.last_replica_states
u_kl = self.database.last_u_kl
Nij_proposed = self.database.last_proposed
Nij_accepted = self.database.last_accepted
iteration = self.database.last_iteration
iteration += 1 # Want to begin with the NEXT step of repex
parameters = self.database.parameters
else:
positions, replica_states, u_kl, Nij_proposed, Nij_accepted, parameters, iteration = None, None, None, None, None, None, None
positions = self.mpicomm.bcast(positions, root=0)
self.replica_states = self.mpicomm.bcast(replica_states, root=0)
self.u_kl = self.mpicomm.bcast(u_kl, root=0)
self.iteration = self.mpicomm.bcast(iteration, root=0)
self.Nij_proposed = self.mpicomm.bcast(Nij_proposed, root=0)
self.Nij_accepted = self.mpicomm.bcast(Nij_accepted, root=0)
self.parameters = self.mpicomm.bcast(parameters, root=0) # Send out as dictionary
self.parameters = dict_to_named_tuple(self.parameters) # Convert to named_tuple for const-ness
self._check_run_parameter_consistency()
self.sampler_states = [SamplerState(self.thermodynamic_states[k].system, positions[k], self.platform) for k in range(len(self.thermodynamic_states))]
def extend(self, n_iter):
"""Extend an existing repex run and modify its database.
Parameters
----------
n_iter : int
How many repex iterations to append.
Notes
-----
This function is MPI aware and only makes database changes on the root
node.
"""
value = self.parameters.number_of_iterations + n_iter
if self.mpicomm.rank == 0:
self.database._store_parameter("number_of_iterations", value)
self.database.sync()
self.parameters = self.parameters._asdict() # Make dict for mutability
self.parameters["number_of_iterations"] = value # extend
self.parameters = dict_to_named_tuple(self.parameters) # Convert to namedtuple for const-ness
def __init__(self, thermodynamic_states, sampler_states=None, database=None, mpicomm=None, platform=None, parameters={}):
"""Create a ReplicaExchange simulation object.
Parameters
----------
thermodynamic_states : list of ThermodynamicState objects
List of thermodynamic states to simulate.
sampler_states : list of SamplerState objects, optional?
List of MCMC sampler states to initialize replica exchange simulations with.
database : Database object, optional
Database to use to write/append simulation data to.
mpicomm : mpicomm implementation, optional?
Communicator (real or dummy) implementation to use for parallelization.
platform : simtk.openmm.Platform, optional, default None
Platform to use for execution.
Notes
-----
Use ReplicaExchange.create() to create a new repex simulation.
To resume an existing ReplicaExchange (or subclass) simulation,
use the `resume()` function. In general, you will not need to
directory call the ReplicaExchange() constructor.
"""
if mpicomm is None:
self.mpicomm = dummympi.COMM_WORLD
else:
self.mpicomm = mpicomm
self.platform = platform
self.database = database
self.thermodynamic_states = thermodynamic_states
self.n_states = len(self.thermodynamic_states)
self.n_atoms = self.thermodynamic_states[0].system.getNumParticles()
if sampler_states is not None: # New Repex job
self.sampler_states = sampler_states
self.parameters = self.process_parameters(parameters) # Fill in missing parameters with defaults
if self.mpicomm.rank == 0: # Store the filled-in parameter namedtuple
self.database.store_parameters(self.parameters)
self.parameters = dict_to_named_tuple(self.parameters) # Convert to namedtuple for const-ness
self._check_run_parameter_consistency()
self._allocate_arrays()
else: # Resume repex job
self._broadcast_database()
self.current_timestep = self.parameters.timestep
self.n_replicas = len(self.thermodynamic_states) # Determine number of replicas from the number of specified thermodynamic states.
# Check to make sure all states have the same number of atoms and are in the same thermodynamic ensemble.
for state in self.thermodynamic_states:
if not state.is_compatible_with(self.thermodynamic_states[0]):
raise ValueError("Provided ThermodynamicState states must all be from the same thermodynamic ensemble.")
if self.database is not None:
self.database.ncfile.repex_classname = self.__class__.__name__
# Eventually, we might want to wrap a setter around the ncfile
logger.debug("Initialized node %d / %d" % (self.mpicomm.rank, self.mpicomm.size))
citations.display_citations(self.parameters.replica_mixing_scheme, self.parameters.online_analysis)
