def run(self, communicator: interfaces.ICommunicator,
system_runner: interfaces.IRunner) -> None:
"""
Continue running worker jobs until done.
Args:
communicator: a communicator object for talking to the leader
system_runner: a system runner object for actually running the simulations
"""
# we always minimize when we first start, either on the first
# stage or the first stage after a restart
minimize = True
while self._step <= self._max_steps:
# update simulation conditions
state = communicator.receive_state_from_leader()
new_alpha = communicator.receive_alpha_from_leader()
state.alpha = new_alpha
system_runner.prepare_for_timestep(state, new_alpha, self._step)
# do one round of simulation
if minimize:
state = system_runner.minimize_then_run(state)
minimize = False # we don't need to minimize again
else:
state = system_runner.run(state)
# compute energies
states = communicator.exchange_states_for_energy_calc(state)
energies = []
for state in states:
energy = system_runner.get_energy(state)
energies.append(energy)
communicator.send_energies_to_leader(energies)
self._step += 1
def run(
self,
communicator: interfaces.ICommunicator,
system_runner: interfaces.IRunner,
store: vault.DataStore,
):
"""
Run replica exchange until finished
Args:
communicator: a communicator object to talk with workers
system_runner: a interfaces.IRunner object to run the simulations
store: a store object to handle storing data to disk
"""
logger.info("Beginning replica exchange")
# check to make sure n_replicas matches
assert self._n_replicas == communicator.n_replicas
assert self._n_replicas == store.n_replicas
# load previous state from the store
states = store.load_states(stage=self.step - 1)
# we always minimize when we first start, either on the first
# stage or the first stage after a restart
minimize = True
while self._step <= self._max_steps:
logger.info("Running replica exchange step %d of %d.", self._step,
self._max_steps)
# communicate state
my_state = communicator.broadcast_states_to_workers(states)
# update alphas
system_runner.prepare_for_timestep(my_state, 0.0, self._step)
self._alphas = self.adaptor.adapt(self._alphas, self._step)
communicator.broadcast_alphas_to_workers(self._alphas)
# do one step
if minimize:
logger.info("First step, minimizing and then running.")
my_state = system_runner.minimize_then_run(my_state)
minimize = False # we don't need to minimize again
else:
logger.info("Running molecular dynamics.")
my_state = system_runner.run(my_state)
# gather all of the states
states = communicator.exchange_states_for_energy_calc(my_state)
# compute our energy for each state
my_energies = self._compute_energies(states, system_runner)
energies = communicator.gather_energies_from_workers(my_energies)
# ask the ladder how to permute things
permutation_vector = self.ladder.compute_exchanges(
energies, self.adaptor)
states = self._permute_states(permutation_vector, states,
system_runner)
# store everything
store.save_states(states, self.step)
store.append_traj(states[0], self.step)
store.save_alphas(np.array(self._alphas), self.step)
store.save_permutation_vector(permutation_vector, self.step)
store.save_energy_matrix(energies, self.step)
store.save_acceptance_probabilities(
self.adaptor.get_acceptance_probabilities(), self.step)
store.save_data_store()
# on to the next step!
self._step += 1
store.save_remd_runner(self)
store.backup(self.step - 1)
logger.info("Finished %d steps of replica exchange successfully.",
self._max_steps)