def run_replica_exchange( topology, system, positions, total_simulation_time=1.0 * unit.picosecond, simulation_time_step=None, temperature_list=None, friction=1.0 / unit.picosecond, minimize=True, exchange_frequency=1000, output_data="output/output.nc", ): """ Run a OpenMMTools replica exchange simulation using an OpenMM coarse grained model. :param topology: OpenMM Topology :type topology: `Topology() <https://simtk.org/api_docs/openmm/api4_1/python/classsimtk_1_1openmm_1_1app_1_1topology_1_1Topology.html>`_ :param system: OpenMM System() :type system: `System() <https://simtk.org/api_docs/openmm/api4_1/python/classsimtk_1_1openmm_1_1openmm_1_1System.html>`_ :param positions: Positions array for the model we would like to test :type positions: `Quantity() <http://docs.openmm.org/development/api-python/generated/simtk.unit.quantity.Quantity.html>`_ ( np.array( [cgmodel.num_beads,3] ), simtk.unit ) :param total_simulation_time: Total run time for individual simulations :type total_simulation_time: `SIMTK <https://simtk.org/>`_ `Unit() <http://docs.openmm.org/7.1.0/api-python/generated/simtk.unit.unit.Unit.html>`_ :param simulation_time_step: Simulation integration time step :type simulation_time_step: `SIMTK <https://simtk.org/>`_ `Unit() <http://docs.openmm.org/7.1.0/api-python/generated/simtk.unit.unit.Unit.html>`_ :param temperature_list: List of temperatures for which to perform replica exchange simulations, default = None :type temperature: List( float * simtk.unit.temperature ) :param friction: Langevin thermostat friction coefficient, default = 1 / ps :type friction: `SIMTK <https://simtk.org/>`_ `Unit() <http://docs.openmm.org/7.1.0/api-python/generated/simtk.unit.unit.Unit.html>`_ :param minimize: Whether minimization is done before running the simulation :type minimize: bool :param output_data: Name of NETCDF file where we will write simulation data :type output_data: string :param exchange_frequency: Number of time steps between replica exchange attempts, Default = None :type exchange_frequency: int :param output_data: file to put the output .nc :type output_data: netCDF4 file as generated by OpenMM :returns: - replica_energies ( `Quantity() <http://docs.openmm.org/development/api-python/generated/simtk.unit.quantity.Quantity.html>`_ ( np.float( [number_replicas,number_simulation_steps] ), simtk.unit ) ) - The potential energies for all replicas at all (printed) time steps - replica_positions ( `Quantity() <http://docs.openmm.org/development/api-python/generated/simtk.unit.quantity.Quantity.html>`_ ( np.float( [number_replicas,number_simulation_steps,cgmodel.num_beads,3] ), simtk.unit ) ) - The positions for all replicas at all (printed) time steps - replica_state_indices ( np.int64( [number_replicas,number_simulation_steps] ), simtk.unit ) - The thermodynamic state assignments for all replicas at all (printed) time steps :Example: >>> from foldamers.cg_model.cgmodel import CGModel >>> from cg_openmm.simulation.rep_exch import * >>> cgmodel = CGModel() >>> replica_energies,replica_positions,replica_state_indices = run_replica_exchange(cgmodel.topology,cgmodel.system,cgmodel.positions) """ simulation_steps = int( np.floor(total_simulation_time / simulation_time_step)) exchange_attempts = int(np.floor(simulation_steps / exchange_frequency)) if temperature_list is None: temperature_list = [((300.0 + i) * unit.kelvin) for i in range(-50, 50, 10)] num_replicas = len(temperature_list) sampler_states = list() thermodynamic_states = list() # Define thermodynamic states. # box_vectors = system.getDefaultPeriodicBoxVectors() for temperature in temperature_list: thermodynamic_state = openmmtools.states.ThermodynamicState( system=system, temperature=temperature) thermodynamic_states.append(thermodynamic_state) sampler_states.append(openmmtools.states.SamplerState( positions)) # no box vectors, non-periodic system. # Create and configure simulation object. move = openmmtools.mcmc.LangevinDynamicsMove( timestep=simulation_time_step, collision_rate=friction, n_steps=exchange_frequency, reassign_velocities=False, ) simulation = ReplicaExchangeSampler( mcmc_moves=move, number_of_iterations=exchange_attempts, replica_mixing_scheme='swap-neighbors', ) if os.path.exists(output_data): os.remove(output_data) reporter = MultiStateReporter(output_data, checkpoint_interval=1) simulation.create(thermodynamic_states, sampler_states, reporter) if minimize: simulation.minimize() print("Running replica exchange simulations with OpenMM...") print(f"Using a time step of {simulation_time_step}") try: simulation.run() except BaseException: print( "Replica exchange simulation failed, try verifying your model/simulation settings." ) exit()
class MCMCReplicaExchangeOpenMMSimulationWrapper: class Config(Config): def __init__(self, args): self.t_max_k = None self.t_min_k = None self.n_replicas = None self.temps_in_k = None self.hybrid = None self.ligand_pertubation_samples = None self.displacement_sigma = None self.verbose = None self.n_steps = None self.parameters = mmWrapperUtils.SystemParams(args['params']) self.warmupparameters = None if "warmupparams" in args: self.warmupparameters = mmWrapperUtils.SystemParams( args['warmupparams']) self.systemloader = None if args is not None: self.__dict__.update(args) if self.temps_in_k is None: self.T_min = self.t_min_k * unit.kelvin # Minimum temperature. self.T_max = self.t_max_k * unit.kelvin # Maximum temperature. self.temps_in_k = [ self.T_min + (self.T_max - self.T_min) * (math.exp(float(i) / float(self.n_replicas - 1)) - 1.0) / (math.e - 1.0) for i in range(self.n_replicas) ] print("MCMCReplicaExchange Temps", self.temps_in_k) elif None in [self.T_min, self.T_max, self.n_replicas]: self.temps_in_k = self.temps_in_k * unit.kelvin self.T_min = min(self.temps_in_k) self.T_max = max(self.temps_in_k) self.n_replicas = len(self.temps_in_k) else: assert (False) def get_obj(self, system_loader, *args, **kwargs): self.systemloader = system_loader return MCMCReplicaExchangeOpenMMSimulationWrapper( self, *args, **kwargs) def __init__(self, config_: Config, old_sampler_state=None): """ :param systemLoader: :param config: """ self._times = None self.config = config_ self.logger = make_message_writer(self.config.verbose, self.__class__.__name__) with self.logger("__init__") as logger: self.explicit = self.config.systemloader.explicit self.amber = bool( self.config.systemloader.config.method == 'amber') self._trajs = np.zeros((1, 1)) self._id_number = int(self.config.systemloader.params_written) if self.config.systemloader.system is None: self.system = self.config.systemloader.get_system( self.config.parameters.createSystem) cache.global_context_cache.set_platform( self.config.parameters.platform, self.config.parameters.platform_config) cache.global_context_cache.time_to_live = 10 else: self.system = self.config.systemloader.system self.topology = self.config.systemloader.topology positions, velocities = self.config.systemloader.get_positions( ), None sequence_move = mmWrapperUtils.prepare_mcmc( self.topology, self.config) thermo_states = [ openmmtools.states.ThermodynamicState( copy.deepcopy(self.system), temperature=t, pressure=1.0 * unit.atmosphere if self.config.systemloader.explicit else None) for t in self.config.temps_in_k ] if self.explicit: sampler_states = [ openmmtools.states.SamplerState( positions=positions, velocities=velocities, box_vectors=self.config.systemloader.boxvec) for _ in thermo_states ] else: sampler_states = [ openmmtools.states.SamplerState(positions=positions, velocities=velocities) for _ in thermo_states ] self.simulation = ReplicaExchangeSampler(mcmc_moves=sequence_move, number_of_iterations=500) self.storage_path = tempfile.NamedTemporaryFile( delete=True).name + '.nc' self.reporter = openmmtools.multistate.MultiStateReporter( self.storage_path, checkpoint_interval=10) self.simulation.create(thermodynamic_states=thermo_states, sampler_states=sampler_states, storage=self.reporter) logger.log("Minimizing...", self.simulation.Status) self.simulation.minimize( max_iterations=self.config.parameters.minMaxIters) self.simulation.equilibrate(1) logger.log("Done, minimizing...", self.simulation.Status) def run(self, iters, steps_per_iter, idx=0): """ :param steps: """ with self.logger("run") as logger: if 'cur_sim_steps' not in self.__dict__: self.cur_sim_steps = 0.0 * unit.picosecond pbar = tqdm( range(1), desc="running {} steps per sample".format(steps_per_iter)) # self._trajs = np.zeros((iters, self.system.getNumParticles(), 3)) # self._times = np.zeros((iters)) # dcdreporter = DCDReporter(f"{self.config.tempdir()}/traj.dcd", 1, append=False) for i in pbar: self.simulation.run(1) # self.cur_sim_steps += (steps_per_iter * self.get_sim_time()) # # positions = self.simulation.sampler_states[idx].positions # boxvectors = self.simulation.sampler_states[idx].box_vectors # # dcdreporter.report_ns(self.topology, positions, boxvectors, (i + 1), 0.5 * unit.femtosecond) # # log trajectory # self._trajs[i] = np.array(positions.value_in_unit(unit.angstrom)).reshape( # (self.system.getNumParticles(), 3)) # self._times[i] = self.cur_sim_steps.value_in_unit(unit.picosecond) pbar.close() exit() def writetraj(self, idx=0): if self.explicit: lengths, angles = mmWrapperUtils.get_mdtraj_box( boxvec=self.simulation.sampler_states[idx].box_vectors, iterset=self._trajs.shape[0]) traj = md.Trajectory(self._trajs, md.Topology.from_openmm(self.topology), unitcell_lengths=lengths, unitcell_angles=angles, time=self._times) traj.image_molecules(inplace=True) else: traj = md.Trajectory(self._trajs, md.Topology.from_openmm(self.topology), time=self._times) traj.save_hdf5(f'{self.config.tempdir()}/mdtraj_traj.h5') def run_amber_mmgbsa(self, run_decomp=False): mmWrapperUtils.run_amber_mmgbsa(self.logger, self.explicit, self.config.tempdir(), run_decomp=run_decomp) def get_sim_time(self): return self.config.n_steps * self.config.parameters.integrator_params[ 'timestep'] def get_velocities(self, idx=0): return self.simulation.sampler_states[idx].velocities def get_coordinates(self, idx=0): return mmWrapperUtils.get_coordinates_samplers( self.topology, self.simulation.sampler_states[idx], self.explicit) def get_pdb(self, file_name=None, idx=0): mmWrapperUtils.get_pdb(self.topology, self.get_coordinates(idx), file_name=file_name)