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)