def test_non_alchemical_protocol(self): """Any property of the ThermodynamicSystem can be specified in the protocol.""" name, thermodynamic_state, sampler_state, topography = self.host_guest_implicit protocol = { 'lambda_sterics': [0.0, 0.5, 1.0], 'temperature': [300, 320, 300] * unit.kelvin } alchemical_phase = AlchemicalPhase(sampler=ReplicaExchangeSampler()) with self.temporary_storage_path() as storage_path: alchemical_phase.create(thermodynamic_state, sampler_state, topography, protocol, storage_path, restraint=yank.restraints.Harmonic()) self.check_protocol(alchemical_phase, protocol) # If temperatures of the end states is different, an error is raised. protocol['temperature'][-1] = 330 * unit.kelvin alchemical_phase = AlchemicalPhase(sampler=ReplicaExchangeSampler()) with nose.tools.assert_raises(ValueError): alchemical_phase.create(thermodynamic_state, sampler_state, topography, protocol, 'not_created.nc', restraint=yank.restraints.Harmonic())
def mix_replicas(n_swaps=100, n_states=16, u_kl=None, nswap_attempts=None): """ Utility function to generate replicas and call the mixing function a certain number of times Arguments --------- n_swaps : int (optional) The number of times to call the mixing code (default 100) n_states : int (optional) The number of replica states to include (default 16) u_kl : n_states x n_states ndarray of float64 (optional) Energies for each state. If None, will be initialized to zeros Returns ------- permutation_list : n_states x n_swaps ndarray of np.int64 Contains the result of each swap """ if u_kl is None: u_kl = np.zeros([n_states, n_states], dtype=np.float64) replica_states = np.array(range(n_states), np.int64) if nswap_attempts is None: nswap_attempts = n_states**4 Nij_proposed = np.zeros([n_states,n_states], dtype=np.int64) Nij_accepted = np.zeros([n_states,n_states], dtype=np.int64) permutation_list = [] from openmmtools.multistate import ReplicaExchangeSampler for i in range(n_swaps): ReplicaExchangeSampler._mix_all_replicas_numba(nswap_attempts, n_states, replica_states, u_kl, Nij_proposed, Nij_accepted) permutation_list.append(copy.deepcopy(replica_states)) permutation_list_np = np.array(permutation_list, dtype=np.int64) return permutation_list_np
def test_randomize_ligand(self): """Test method AlchemicalPhase.randomize_ligand.""" _, thermodynamic_state, sampler_state, topography = self.host_guest_implicit restraint = yank.restraints.Harmonic() ligand_atoms, receptor_atoms = topography.ligand_atoms, topography.receptor_atoms ligand_positions = sampler_state.positions[ligand_atoms] receptor_positions = sampler_state.positions[receptor_atoms] with self.temporary_storage_path() as storage_path: alchemical_phase = AlchemicalPhase(ReplicaExchangeSampler()) alchemical_phase.create(thermodynamic_state, sampler_state, topography, self.protocol, storage_path, restraint=restraint) # Randomize ligand positions. alchemical_phase.randomize_ligand() # The new sampler states have the same receptor positions # but different ligand positions. for sampler_state in alchemical_phase._sampler.sampler_states: assert np.allclose(sampler_state.positions[receptor_atoms], receptor_positions) assert not np.allclose(sampler_state.positions[ligand_atoms], ligand_positions)
def test_create(self): """Alchemical state correctly creates the simulation object.""" available_restraints = [ restraint for restraint in yank.restraints.available_restraint_classes().values() if not (hasattr(restraint, "dev_validate") and not restraint.dev_validate) ] for test_index, test_case in enumerate(self.all_test_cases): test_name, thermodynamic_state, sampler_state, topography = test_case # Add random restraint if this is ligand-receptor system in implicit solvent. if len(topography.ligand_atoms) > 0: restraint_cls = available_restraints[np.random.randint( 0, len(available_restraints))] restraint = restraint_cls() protocol = self.restrained_protocol test_name += ' with restraint {}'.format( restraint_cls.__name__) else: restraint = None protocol = self.protocol # Add either automatic of fixed correction cutoff. if test_index % 2 == 0: correction_cutoff = 12 * unit.angstroms else: correction_cutoff = 'auto' # Replace the reaction field of the reference system to compare # also cutoff and switch width for the electrostatics. reference_system = thermodynamic_state.system mmtools.forcefactories.replace_reaction_field(reference_system, return_copy=False) alchemical_phase = AlchemicalPhase( sampler=ReplicaExchangeSampler()) with self.temporary_storage_path() as storage_path: alchemical_phase.create( thermodynamic_state, sampler_state, topography, protocol, storage_path, restraint=restraint, anisotropic_dispersion_cutoff=correction_cutoff) yield prepare_yield(self.check_protocol, test_name, alchemical_phase, protocol) yield prepare_yield(self.check_standard_state_correction, test_name, alchemical_phase, topography, restraint) yield prepare_yield(self.check_expanded_states, test_name, alchemical_phase, protocol, correction_cutoff, reference_system) # Free memory. del alchemical_phase
def run_replica_exchange_done(job): output_directory = os.path.join(job.workspace(), "output") output_data = os.path.join(output_directory, "output.nc") rep_exch_completed = 0 if os.path.isfile(output_data): rep_exch_status = ReplicaExchangeSampler.read_status(output_data) rep_exch_completed = rep_exch_status.is_completed return rep_exch_completed
def test_illegal_restraint(self): """Raise an error when restraint is handled incorrectly.""" # An error is raised with ligand-receptor systems in implicit without restraint. test_name, thermodynamic_state, sampler_state, topography = self.host_guest_implicit alchemical_phase = AlchemicalPhase(sampler=ReplicaExchangeSampler()) with nose.tools.assert_raises(ValueError): alchemical_phase.create(thermodynamic_state, sampler_state, topography, self.protocol, 'not_created.nc')
def CEI_replica_exchange_done(job): output_directory = os.path.join(job.workspace(),"output_CEI") output_data = os.path.join(output_directory, "output.nc") rep_exch_completed = 0 if os.path.isfile(output_data): # ***Note: this can sometimes fail if the .nc files are currently being written to rep_exch_status = ReplicaExchangeSampler.read_status(output_data) rep_exch_completed = rep_exch_status.is_completed return rep_exch_completed
def test_minimize(self): """Test AlchemicalPhase minimization of positions in reference state.""" # Ligand-receptor in implicit solvent. test_system = testsystems.AlanineDipeptideVacuum() thermodynamic_state = states.ThermodynamicState(test_system.system, temperature=300 * unit.kelvin) topography = Topography(test_system.topology) # We create 3 different sampler states that will be distributed over # replicas in a round-robin fashion. displacement_vector = np.ones(3) * unit.nanometer positions2 = test_system.positions + displacement_vector positions3 = positions2 + displacement_vector box_vectors = test_system.system.getDefaultPeriodicBoxVectors() sampler_state1 = states.SamplerState(positions=test_system.positions, box_vectors=box_vectors) sampler_state2 = states.SamplerState(positions=positions2, box_vectors=box_vectors) sampler_state3 = states.SamplerState(positions=positions3, box_vectors=box_vectors) sampler_states = [sampler_state1, sampler_state2, sampler_state3] with self.temporary_storage_path() as storage_path: # Create alchemical phase. alchemical_phase = AlchemicalPhase(ReplicaExchangeSampler()) alchemical_phase.create(thermodynamic_state, sampler_states, topography, self.protocol, storage_path) # Measure the average distance between positions. This should be # maintained after minimization. sampler_states = alchemical_phase._sampler.sampler_states original_diffs = [ np.average(sampler_states[i].positions - sampler_states[i + 1].positions) for i in range(len(sampler_states) - 1) ] # Minimize. alchemical_phase.minimize() # The minimized positions should be still more or less # one displacement vector from each other. sampler_states = alchemical_phase._sampler.sampler_states new_diffs = [ np.average(sampler_states[i].positions - sampler_states[i + 1].positions) for i in range(len(sampler_states) - 1) ] assert np.allclose( original_diffs, new_diffs, rtol=0.1 ), "original_diffs {} are not close to new_diffs {}".format( original_diffs, new_diffs)
def test_illegal_protocol(self): """An error is raised when the protocol parameters have a different number of states.""" name, thermodynamic_state, sampler_state, topography = self.host_guest_implicit protocol = { 'lambda_sterics': [0.0, 1.0], 'lambda_electrostatics': [1.0] } alchemical_phase = AlchemicalPhase(sampler=ReplicaExchangeSampler()) restraint = yank.restraints.Harmonic() with nose.tools.assert_raises(ValueError): alchemical_phase.create(thermodynamic_state, sampler_state, topography, protocol, 'not_created.nc', restraint=restraint)
def test_from_storage(self): """When resuming, the AlchemicalPhase recover the correct sampler.""" _, thermodynamic_state, sampler_state, topography = self.host_guest_implicit restraint = yank.restraints.Harmonic() with self.temporary_storage_path() as storage_path: alchemical_phase = AlchemicalPhase(ReplicaExchangeSampler()) alchemical_phase.create(thermodynamic_state, sampler_state, topography, self.protocol, storage_path, restraint=restraint) # Delete old alchemical phase to close storage file. del alchemical_phase # Resume, the sampler has the correct class. alchemical_phase = AlchemicalPhase.from_storage(storage_path) assert isinstance(alchemical_phase._sampler, ReplicaExchangeSampler)
def test_illegal_restraint(self): """Raise an error when restraint is handled incorrectly.""" # An error is raised with ligand-receptor systems in implicit without restraint. test_name, thermodynamic_state, sampler_state, topography = self.host_guest_implicit alchemical_phase = AlchemicalPhase(sampler=ReplicaExchangeSampler()) with nose.tools.assert_raises(ValueError): alchemical_phase.create(thermodynamic_state, sampler_state, topography, self.protocol, 'not_created.nc') # An error is raised when trying to apply restraint to non ligand-receptor systems. restraint = yank.restraints.Harmonic() test_name, thermodynamic_state, sampler_state, topography = self.alanine_explicit with nose.tools.assert_raises(RuntimeError): alchemical_phase.create(thermodynamic_state, sampler_state, topography, self.protocol, 'not_created.nc', restraint=restraint)
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()
else: restraint = None protocol = self.protocol # Add either automatic of fixed correction cutoff. if test_index % 2 == 0: correction_cutoff = 12 * unit.angstroms else: correction_cutoff = 'auto' # Replace the reaction field of the reference system to compare # also cutoff and switch width for the electrostatics. reference_system = thermodynamic_state.system mmtools.forcefactories.replace_reaction_field(reference_system, return_copy=False) alchemical_phase = AlchemicalPhase(sampler=ReplicaExchangeSampler()) with self.temporary_storage_path() as storage_path: alchemical_phase.create(thermodynamic_state, sampler_state, topography, protocol, storage_path, restraint=restraint, anisotropic_dispersion_cutoff=correction_cutoff) yield prepare_yield(self.check_protocol, test_name, alchemical_phase, protocol) yield prepare_yield(self.check_standard_state_correction, test_name, alchemical_phase, topography, restraint) yield prepare_yield(self.check_expanded_states, test_name, alchemical_phase, protocol, correction_cutoff, reference_system) # Free memory. del alchemical_phase def test_default_alchemical_region(self):
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)
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)