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)