def trial(self, state: interfaces.IState, runner: interfaces.IRunner) -> Tuple[interfaces.IState, bool]: """ Perform a Metropolis trial Args: starting_state: initial state of system runner: runner to evaluate energies Returns: the system state after Monte Carlo trials """ starting_positions = state.positions.copy() starting_energy = state.energy trial_positions = starting_positions.copy() random_vector = np.random.normal(loc=0.0, scale=self.move_size, size=3) trial_positions[self.atom_indices, :] += random_vector state.positions = trial_positions trial_energy = runner.get_energy(state) accepted = _metropolis(starting_energy, trial_energy, bias=0.0) if accepted: state.energy = trial_energy state.positions = trial_positions else: state.energy = starting_energy state.positions = starting_positions return state, accepted
def trial(self, state: interfaces.IState, runner: interfaces.IRunner) -> Tuple[interfaces.IState, bool]: """ Perform a Metropolis trial Args: starting_state: initial state of system runner: runner to evaluate energies Returns: the system state after Monte Carlo trials """ starting_positions = state.positions.copy() starting_energy = state.energy angle = _generate_uniform_angle() trial_positions = starting_positions.copy() trial_positions[self.atom_indices, :] = _rotate_around_vector( starting_positions[self.index1, :], starting_positions[self.index2, :], angle, starting_positions[self.atom_indices, :], ) state.positions = trial_positions trial_energy = runner.get_energy(state) accepted = _metropolis(starting_energy, trial_energy, 0.0) if accepted: state.energy = trial_energy state.positions = trial_positions else: state.energy = starting_energy state.positions = starting_positions return state, accepted
def _permute_states( permutation_matrix: List[int], states: Sequence[interfaces.IState], system_runner: interfaces.IRunner, ) -> Sequence[interfaces.IState]: old_coords = [s.positions for s in states] old_velocities = [s.velocities for s in states] old_box_vectors = [s.box_vector for s in states] old_energy = [s.energy for s in states] old_params = [s.parameters for s in states] old_mappings = [s.mappings for s in states] assert system_runner.temperature_scaler is not None temperatures = [ system_runner.temperature_scaler(s.alpha) for s in states ] for i, index in enumerate(permutation_matrix): states[i].positions = old_coords[index] states[i].velocities = ( math.sqrt(temperatures[i] / temperatures[index]) * old_velocities[index]) states[i].box_vector = old_box_vectors[index] states[i].energy = old_energy[index] states[i].parameters = old_params[index] states[i].mappings = old_mappings[index] return states
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)
def _compute_energies(states: Sequence[interfaces.IState], system_runner: interfaces.IRunner) -> List[float]: my_energies = [] for state in states: my_energies.append(system_runner.get_energy(state)) return my_energies