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,
'update_alchemical_charges': [True, True, 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=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 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 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_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 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 run_replica_exchange(topology,system,positions,temperature_list=None,simulation_time_step=None,total_simulation_time=1.0 * unit.picosecond,output_data='output.nc',print_frequency=100,verbose_simulation=False,exchange_attempts=None,test_time_step=False,output_directory=None):
"""
Run a Yank 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 temperature_list: List of temperatures for which to perform replica exchange simulations, default = None
:type temperature: List( float * simtk.unit.temperature )
: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 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 output_data: Name of NETCDF file where we will write simulation data
:type output_data: string
:param print_frequency: Number of simulation steps to skip when writing to output, Default = 100
:type print_frequence: int
:param verbose_simulation: Determines how much output is printed during a simulation run. Default = False
:type verbose_simulation: Logical
:param exchange_attempts: Number of exchange attempts to make during a replica exchange simulation run, Default = None
:type exchange_attempts: int
:param test_time_step: Logical variable determining if a test of the time step will be performed, Default = False
:type test_time_step: Logical
:param output_directory: Path to which we will write the output from simulation runs.
:type output_directory: str
: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)
"""
if simulation_time_step == None:
simulation_time_step,force_threshold = get_simulation_time_step(topology,system,positions,temperature_list[-1],total_simulation_time)
simulation_steps = int(round(total_simulation_time.__div__(simulation_time_step)))
if exchange_attempts == None:
if simulation_steps > 10000:
exchange_attempts = round(simulation_steps/1000)
else:
exchange_attempts = 10
if temperature_list == None:
temperature_list = [(300.0 * unit.kelvin).__add__(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 = mmtools.states.ThermodynamicState(system=system, temperature=temperature)
thermodynamic_states.append(thermodynamic_state)
sampler_states.append(mmtools.states.SamplerState(positions,box_vectors=box_vectors))
# Create and configure simulation object.
move = mmtools.mcmc.LangevinDynamicsMove(timestep=simulation_time_step,collision_rate=5.0/unit.picosecond,n_steps=exchange_attempts, reassign_velocities=True)
simulation = ReplicaExchangeSampler(mcmc_moves=move, number_of_iterations=exchange_attempts)
if os.path.exists(output_data): os.remove(output_data)
reporter = MultiStateReporter(output_data, checkpoint_interval=1)
simulation.create(thermodynamic_states, sampler_states, reporter)
config_root_logger(verbose_simulation)
if not test_time_step:
num_attempts = 0
while num_attempts < 5:
try:
simulation.run()
#print("Replica exchange simulations succeeded with a time step of: "+str(simulation_time_step))
break
except:
num_attempts = num_attempts + 1
if num_attempts >= 5:
print("Replica exchange simulation attempts failed, try verifying your model/simulation settings.")
exit()
else:
simulation_time_step,force_threshold = get_simulation_time_step(topology,system,positions,temperature_list[-1],total_simulation_time)
print("The suggested time step for a simulation with this model is: "+str(simulation_time_step))
while simulation_time_step.__div__(2.0) > 0.001 * unit.femtosecond:
try:
print("Running replica exchange simulations with Yank...")
print("Using a time step of "+str(simulation_time_step))
print("Running each trial simulation for 1000 steps, with 10 exchange attempts.")
move = mmtools.mcmc.LangevinDynamicsMove(timestep=simulation_time_step,collision_rate=20.0/unit.picosecond,n_steps=10, reassign_velocities=True)
simulation = ReplicaExchangeSampler(replica_mixing_scheme='swap-neighbors',mcmc_moves=move,number_of_iterations=10)
reporter = MultiStateReporter(output_data, checkpoint_interval=1)
simulation.create(thermodynamic_states, sampler_states, reporter)
simulation.run()
print("Replica exchange simulations succeeded with a time step of: "+str(simulation_time_step))
break
except:
del simulation
os.remove(output_data)
print("Simulation attempt failed with a time step of: "+str(simulation_time_step))
if simulation_time_step.__div__(2.0) > 0.001 * unit.femtosecond:
simulation_time_step = simulation_time_step.__div__(2.0)
else:
print("Error: replica exchange simulation attempt failed with a time step of: "+str(simulation_time_step))
print("Please check the model and simulations settings, and try again.")
exit()
replica_energies,replica_positions,replica_state_indices = read_replica_exchange_data(system=system,topology=topology,temperature_list=temperature_list,output_data=output_data,print_frequency=print_frequency)
steps_per_stage = round(simulation_steps/exchange_attempts)
if output_directory != None:
plot_replica_exchange_energies(replica_energies,temperature_list,simulation_time_step,steps_per_stage=steps_per_stage,output_directory=output_directory)
plot_replica_exchange_summary(replica_state_indices,temperature_list,simulation_time_step,steps_per_stage=steps_per_stage,output_directory=output_directory)
else:
plot_replica_exchange_energies(replica_energies,temperature_list,simulation_time_step,steps_per_stage=steps_per_stage)
plot_replica_exchange_summary(replica_state_indices,temperature_list,simulation_time_step,steps_per_stage=steps_per_stage)
return(replica_energies,replica_positions,replica_state_indices)
def main():
parser = argparse.ArgumentParser(
description=
'Compute a potential of mean force (PMF) for porin permeation.')
parser.add_argument('--index',
dest='index',
action='store',
type=int,
help='Index of ')
parser.add_argument('--output',
dest='output_filename',
action='store',
default='output.nc',
help='output netcdf filename (default: output.nc)')
args = parser.parse_args()
index = args.index
output_filename = args.output_filename
logger = logging.getLogger(__name__)
logging.root.setLevel(logging.DEBUG)
logging.basicConfig(level=logging.DEBUG)
yank.utils.config_root_logger(verbose=True, log_file_path=None)
# Configure ContextCache, platform and precision
from yank.experiment import ExperimentBuilder
platform = ExperimentBuilder._configure_platform('CUDA', 'mixed')
try:
openmmtools.cache.global_context_cache.platform = platform
except RuntimeError:
# The cache has been already used. Empty it before switching platform.
openmmtools.cache.global_context_cache.empty()
openmmtools.cache.global_context_cache.platform = platform
# Topology
pdbx = app.PDBxFile('mem_prot_md_system.pdbx')
# This system contains the CVforce with parameters different than zero
with open('openmm_system.xml', 'r') as infile:
openmm_system = XmlSerializer.deserialize(infile.read())
####### Indexes of configurations in trajectory ############################
configs = [
39, 141, 276, 406, 562, 668, 833, 1109, 1272, 1417, 1456, 1471, 1537,
1645, 1777, 1882
]
####### Indexes of states for series of replica exchange simulations #######
limits = [(0, 9), (10, 19), (20, 29), (30, 39), (40, 49), (50, 59),
(60, 69), (70, 79), (80, 89), (90, 99), (100, 109), (110, 119),
(120, 129), (130, 139), (140, 149), (150, 159)]
####### Reading positions from mdtraj trajectory ###########################
topology = md.Topology.from_openmm(pdbx.topology)
t = md.load('../../steered_md/comp7/forward/seed_0/steered_forward.nc',
top=topology)
positions = t.openmm_positions(configs[index])
thermodynamic_state_deserialized = states.ThermodynamicState(
system=openmm_system,
temperature=310 * unit.kelvin,
pressure=1.0 * unit.atmospheres)
sampler_state = states.SamplerState(
positions=positions,
box_vectors=t.unitcell_vectors[configs[index], :, :] * unit.nanometer)
logger.debug(type(sampler_state))
move = mcmc.LangevinDynamicsMove(timestep=2 * unit.femtosecond,
collision_rate=1.0 / unit.picoseconds,
n_steps=500,
reassign_velocities=False)
simulation = ReplicaExchangeSampler(mcmc_moves=move,
number_of_iterations=1)
analysis_particle_indices = topology.select(
'(protein and mass > 3.0) or (resname MER and mass > 3.0)')
reporter = MultiStateReporter(
output_filename,
checkpoint_interval=2000,
analysis_particle_indices=analysis_particle_indices)
first, last = limits[index]
# Initialize compound thermodynamic states
protocol = {
'lambda_restraints': [i / 159 for i in range(first, last + 1)],
'K_parallel': [
1250 * unit.kilojoules_per_mole / unit.nanometer**2
for i in range(first, last + 1)
],
'Kmax': [
500 * unit.kilojoules_per_mole / unit.nanometer**2
for i in range(first, last + 1)
],
'Kmin': [
500 * unit.kilojoules_per_mole / unit.nanometer**2
for i in range(first, last + 1)
]
}
my_composable_state = MyComposableState.from_system(openmm_system)
compound_states = states.create_thermodynamic_state_protocol(
thermodynamic_state_deserialized,
protocol=protocol,
composable_states=[my_composable_state])
simulation.create(thermodynamic_states=compound_states,
sampler_states=[sampler_state],
storage=reporter)
simulation.equilibrate(50, mcmc_moves=move)
simulation.run()
ts = simulation._thermodynamic_states[0]
context, _ = openmmtools.cache.global_context_cache.get_context(ts)
files_names = [
'state_{}_{}.log'.format(index, i) for i in range(first, last + 1)
]
files = []
for i, file in enumerate(files_names):
files.append(open(file, 'w'))
mpi.distribute(write_cv,
range(simulation.n_replicas),
context,
simulation,
files,
send_results_to=None)
for i in range(10000):
simulation.extend(n_iterations=2)
mpi.distribute(write_cv,
range(simulation.n_replicas),
context,
simulation,
files,
send_results_to=None)