def relax_ligand(self, positions, velocities):
with self.logger("relax_ligand") as logger:
cache.global_context_cache.empty()
system = copy.deepcopy(
self.config.systemloader._unconstrained_system)
thermo_state_ = ThermodynamicState(
system=system,
temperature=self.config.parameters.
integrator_params['temperature'],
pressure=1.0 *
unit.atmosphere if self.config.systemloader.explicit else None)
pforce = mmWrapperUtils.get_protein_restraint_force(self.topology,
positions,
self.explicit,
K=5.0)
lforce = mmWrapperUtils.get_ligand_restraint_force(
self.topology,
positions,
self.explicit,
K=thermo_state_.kT / 3.0**2)
del thermo_state_
system.addForce(pforce)
system.addForce(lforce)
thermo_state = ThermodynamicState(
system=system,
temperature=self.config.parameters.
integrator_params['temperature'],
pressure=1.0 *
unit.atmosphere if self.config.systemloader.explicit else None)
sampler_state = SamplerState(
positions=positions,
velocities=velocities,
box_vectors=self.config.systemloader.boxvec)
sampler = MCMCSampler(
thermo_state,
sampler_state,
move=LangevinSplittingDynamicsMove(
timestep=0.1 * unit.femtosecond,
n_steps=1000,
collision_rate=self.config.parameters.
integrator_params['collision_rate'],
reassign_velocities=True,
n_restart_attempts=6,
constraint_tolerance=self.config.parameters.
integrator_setConstraintTolerance))
sampler.minimize(max_iterations=self.config.parameters.minMaxIters)
logger.log("Build unconstrained system. Relaxing for 1 ps")
sampler.run(10)
positions, velocities = sampler.sampler_state.positions, sampler.sampler_state.velocities
cache.global_context_cache.empty()
return positions, velocities
def test_metropolized_moves():
"""Test Displacement and Rotation moves."""
testsystem = testsystems.AlanineDipeptideVacuum()
original_sampler_state = SamplerState(testsystem.positions)
thermodynamic_state = ThermodynamicState(testsystem.system, 300*unit.kelvin)
all_metropolized_moves = MetropolizedMove.__subclasses__()
for move_class in all_metropolized_moves:
move = move_class(atom_subset=range(thermodynamic_state.n_particles))
sampler_state = copy.deepcopy(original_sampler_state)
# Start applying the move and remove one at each iteration tyring
# to generate both an accepted and rejected move.
old_n_accepted, old_n_proposed = 0, 0
while len(move.atom_subset) > 0:
initial_positions = copy.deepcopy(sampler_state.positions)
move.apply(thermodynamic_state, sampler_state)
final_positions = copy.deepcopy(sampler_state.positions)
# If the move was accepted the positions should be different.
if move.n_accepted > old_n_accepted:
assert not np.allclose(initial_positions, final_positions)
# If we have generated a rejection and an acceptance, test next move.
if move.n_accepted > 0 and move.n_accepted != move.n_proposed:
break
# Try with a smaller subset.
move.atom_subset = move.atom_subset[:-1]
old_n_accepted, old_n_proposed = move.n_accepted, move.n_proposed
# Check that we were able to generate both an accepted and a rejected move.
assert len(move.atom_subset) != 0, ('Could not generate an accepted and rejected '
'move for class {}'.format(move_class.__name__))
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)
self.topology = self.config.systemloader.topology
cache.global_context_cache.set_platform(
self.config.parameters.platform,
self.config.parameters.platform_config)
positions, velocities = self.config.systemloader.get_positions(
), None
else:
self.system = self.config.systemloader.system
self.topology = self.config.systemloader.topology
positions, velocities = self.config.systemloader.get_positions(
), None
# if self.config.systemloader.config.relax_ligand:
# positions, velocities = self.relax_ligand(positions, velocities)
thermo_state = ThermodynamicState(
system=self.system,
temperature=self.config.parameters.
integrator_params['temperature'],
pressure=1.0 *
unit.atmosphere if self.config.systemloader.explicit else None)
sampler_state = SamplerState(
positions=positions,
velocities=velocities,
box_vectors=self.config.systemloader.boxvec)
self.sampler = MCMCSampler(thermo_state,
sampler_state,
move=mmWrapperUtils.prepare_mcmc(
self.topology, self.config))
self.sampler.minimize(
max_iterations=self.config.parameters.minMaxIters)
ctx = cache.global_context_cache.get_context(
self.sampler.thermodynamic_state)[0]
ctx.setVelocitiesToTemperature(
self.config.parameters.integrator_params['temperature'])
self.sampler.sampler_state.velocities = ctx.getState(
getVelocities=True).getVelocities()
def get_harmonic_oscillator(cls):
"""Create a harmonic oscillator thermodynamic state to test the trailblaze algorithm."""
from openmmtools.states import (GlobalParameterState,
ThermodynamicState,
CompoundThermodynamicState,
SamplerState)
# Create composable state that control offset of harmonic oscillator.
class X0State(GlobalParameterState):
testsystems_HarmonicOscillator_x0 = GlobalParameterState.GlobalParameter(
cls.PAR_NAME_X0, 1.0)
testsystems_HarmonicOscillator_K = GlobalParameterState.GlobalParameter(
cls.PAR_NAME_K,
(1.0 * unit.kilocalories_per_mole /
unit.nanometer**2).value_in_unit_system(unit.md_unit_system))
# Create a harmonic oscillator thermo state.
k = 1.0 * unit.kilocalories_per_mole / unit.nanometer**2
oscillator = mmtools.testsystems.HarmonicOscillator(K=k)
sampler_state = SamplerState(positions=oscillator.positions)
thermo_state = ThermodynamicState(oscillator.system,
temperature=300 * unit.kelvin)
x0_state = X0State(
testsystems_HarmonicOscillator_x0=0.0,
testsystems_HarmonicOscillator_K=k.value_in_unit_system(
unit.md_unit_system))
compound_state = CompoundThermodynamicState(
thermo_state, composable_states=[x0_state])
return compound_state, sampler_state
def compute_reduced_potential(thermodynamic_state: states.ThermodynamicState,
sampler_state: states.SamplerState) -> float:
"""
Compute the reduced potential of the given SamplerState under the given ThermodynamicState.
Arguments
----------
thermodynamic_state : openmmtools.states.ThermodynamicState
The thermodynamic state under which to compute the reduced potential
sampler_state : openmmtools.states.SamplerState
The sampler state for which to compute the reduced potential
Returns
-------
reduced_potential : float
unitless reduced potential (kT)
"""
if type(cache.global_context_cache) == cache.DummyContextCache:
integrator = openmm.VerletIntegrator(
1.0) #we won't take any steps, so use a simple integrator
context, integrator = cache.global_context_cache.get_context(
thermodynamic_state, integrator)
else:
context, integrator = cache.global_context_cache.get_context(
thermodynamic_state)
sampler_state.apply_to_context(context, ignore_velocities=True)
return thermodynamic_state.reduced_potential(context)
def _topology_proposal_to_thermodynamic_states(self, topology_proposal):
"""
Convert a topology proposal to thermodynamic states for the end systems. This will be used to compute the
"logP_energy" quantity.
Parameters
----------
topology_proposal : perses.rjmc.TopologyProposal
topology proposal for whose endpoint systems we want ThermodynamicStates
Returns
-------
old_thermodynamic_state : openmmtools.states.ThermodynamicState
The old system (nonalchemical) thermodynamic state
new_thermodynamic_state : openmmtools.states.ThermodynamicState
The new system (nonalchemical) thermodynamic state
"""
systems = [topology_proposal.old_system, topology_proposal.new_system]
thermostates = []
for system in systems:
thermodynamic_state = ThermodynamicState(
system, temperature=self._temperature, pressure=self._pressure)
thermostates.append(thermodynamic_state)
return thermostates[0], thermostates[1]
def generate_thermodynamic_states(system: openmm.System,
topology_proposal: TopologyProposal):
"""
Generate endpoint thermodynamic states for the system
Parameters
----------
system : openmm.System
System object corresponding to thermodynamic state
topology_proposal : perses.rjmc.topology_proposal.TopologyProposal
TopologyProposal representing transformation
Returns
-------
nonalchemical_zero_thermodynamic_state : ThermodynamicState
Nonalchemical thermodynamic state for lambda zero endpoint
nonalchemical_one_thermodynamic_state : ThermodynamicState
Nonalchemical thermodynamic state for lambda one endpoint
lambda_zero_thermodynamic_state : ThermodynamicState
Alchemical (hybrid) thermodynamic state for lambda zero
lambda_one_thermodynamic_State : ThermodynamicState
Alchemical (hybrid) thermodynamic state for lambda one
"""
#create the thermodynamic state
lambda_zero_alchemical_state = alchemy.AlchemicalState.from_system(system)
lambda_one_alchemical_state = copy.deepcopy(lambda_zero_alchemical_state)
#ensure their states are set appropriately
lambda_zero_alchemical_state.set_alchemical_parameters(0.0)
lambda_one_alchemical_state.set_alchemical_parameters(0.0)
#create the base thermodynamic state with the hybrid system
thermodynamic_state = ThermodynamicState(system, temperature=temperature)
#Create thermodynamic states for the nonalchemical endpoints
nonalchemical_zero_thermodynamic_state = ThermodynamicState(
topology_proposal.old_system, temperature=temperature)
nonalchemical_one_thermodynamic_state = ThermodynamicState(
topology_proposal.new_system, temperature=temperature)
#Now create the compound states with different alchemical states
lambda_zero_thermodynamic_state = CompoundThermodynamicState(
thermodynamic_state, composable_states=[lambda_zero_alchemical_state])
lambda_one_thermodynamic_state = CompoundThermodynamicState(
thermodynamic_state, composable_states=[lambda_one_alchemical_state])
return nonalchemical_zero_thermodynamic_state, nonalchemical_one_thermodynamic_state, lambda_zero_thermodynamic_state, lambda_one_thermodynamic_state
def test_move_restart():
"""Test optional restart move if NaN is detected."""
n_restart_attempts = 5
# We define a Move that counts the times it is attempted.
class MyMove(BaseIntegratorMove):
def __init__(self, **kwargs):
super(MyMove, self).__init__(n_steps=1,
n_restart_attempts=n_restart_attempts,
**kwargs)
self.attempted_count = 0
def _get_integrator(self, thermodynamic_state):
return integrators.GHMCIntegrator(temperature=300 * unit.kelvin)
def _before_integration(self, context, thermodynamic_state):
self.attempted_count += 1
# Create a system with an extra NaN particle.
testsystem = testsystems.AlanineDipeptideVacuum()
system = testsystem.system
for force in system.getForces():
if isinstance(force, openmm.NonbondedForce):
break
# Add a non-interacting particle to the system at NaN position.
system.addParticle(39.9 * unit.amu)
force.addParticle(0.0, 1.0, 0.0)
particle_position = np.array([np.nan, 0.2, 0.2])
positions = unit.Quantity(np.vstack(
(testsystem.positions, particle_position)),
unit=testsystem.positions.unit)
# Create and run move. An IntegratoMoveError is raised.
sampler_state = SamplerState(positions)
thermodynamic_state = ThermodynamicState(system, 300 * unit.kelvin)
# We use a local context cache with Reference platform since on the
# CPU platform CustomIntegrators raises an error with NaN particles.
reference_platform = openmm.Platform.getPlatformByName('Reference')
move = MyMove(context_cache=cache.ContextCache(
platform=reference_platform))
with nose.tools.assert_raises(IntegratorMoveError) as cm:
move.apply(thermodynamic_state, sampler_state)
# We have counted the correct number of restart attempts.
assert move.attempted_count == n_restart_attempts + 1
# Test serialization of the error.
with utils.temporary_directory() as tmp_dir:
prefix = os.path.join(tmp_dir, 'prefix')
cm.exception.serialize_error(prefix)
assert os.path.exists(prefix + '-move.json')
assert os.path.exists(prefix + '-system.xml')
assert os.path.exists(prefix + '-integrator.xml')
assert os.path.exists(prefix + '-state.xml')
def runEthyleneTest(dir, N):
filename = dir.join('ethylene-test_%s' % N)
print('Running %s...' % filename)
# Set Simulation parameters
temperature = 200 * unit.kelvin
collision_rate = 1 / unit.picoseconds
timestep = 1.0 * unit.femtoseconds
n_steps = 20
nIter = 100
reportInterval = 5
alchemical_atoms = [2, 3, 4, 5, 6, 7]
platform = openmm.Platform.getPlatformByName('CPU')
context_cache = cache.ContextCache(platform)
# Load a Parmed Structure for the Topology and create our openmm.Simulation
structure_pdb = utils.get_data_filename(
'blues', 'tests/data/ethylene_structure.pdb')
structure = parmed.load_file(structure_pdb)
nc_reporter = NetCDF4Storage(filename + '_MD.nc', reportInterval)
# Iniitialize our Move set
rot_move = RandomLigandRotationMove(timestep=timestep,
n_steps=n_steps,
atom_subset=alchemical_atoms,
context_cache=context_cache,
reporters=[nc_reporter])
langevin_move = ReportLangevinDynamicsMove(timestep=timestep,
collision_rate=collision_rate,
n_steps=n_steps,
reassign_velocities=True,
context_cache=context_cache)
# Load our OpenMM System and create Integrator
system_xml = utils.get_data_filename('blues',
'tests/data/ethylene_system.xml')
with open(system_xml, 'r') as infile:
xml = infile.read()
system = openmm.XmlSerializer.deserialize(xml)
thermodynamic_state = ThermodynamicState(system=system,
temperature=temperature)
sampler_state = SamplerState(
positions=structure.positions.in_units_of(unit.nanometers))
sampler = BLUESSampler(thermodynamic_state=thermodynamic_state,
sampler_state=sampler_state,
ncmc_move=rot_move,
dynamics_move=langevin_move,
topology=structure.topology)
sampler.run(nIter)
return filename
def test_langevin_splitting_move():
"""Test that the langevin splitting mcmc move works with different splittings"""
splittings = ["V R O R V", "V R R R O R R R V", "O { V R V } O"]
testsystem = testsystems.AlanineDipeptideVacuum()
sampler_state = SamplerState(testsystem.positions)
thermodynamic_state = ThermodynamicState(testsystem.system, 300*unit.kelvin)
for splitting in splittings:
move = LangevinSplittingDynamicsMove(splitting=splitting)
# Create MCMC sampler
sampler = MCMCSampler(thermodynamic_state, sampler_state, move=move)
sampler.run(1)
def test_context_cache():
"""Test configuration of the context cache."""
testsystem = testsystems.AlanineDipeptideImplicit()
sampler_state = SamplerState(testsystem.positions)
thermodynamic_state = ThermodynamicState(testsystem.system,
300 * unit.kelvin)
# By default the global context cache is used.
cache.global_context_cache.empty() # Clear cache from previous tests.
move = SequenceMove([LangevinDynamicsMove(n_steps=5), GHMCMove(n_steps=5)])
move.apply(thermodynamic_state, sampler_state)
assert len(cache.global_context_cache) == 2
# Configuring the global cache works correctly.
cache.global_context_cache = cache.ContextCache(time_to_live=1)
move.apply(thermodynamic_state, sampler_state)
assert len(cache.global_context_cache) == 1
# The ContextCache creates only one context with compatible moves.
cache.global_context_cache = cache.ContextCache(capacity=10,
time_to_live=None)
move = SequenceMove([
LangevinDynamicsMove(n_steps=1),
LangevinDynamicsMove(n_steps=1),
LangevinDynamicsMove(n_steps=1),
LangevinDynamicsMove(n_steps=1)
])
move.apply(thermodynamic_state, sampler_state)
assert len(cache.global_context_cache) == 1
# We can configure a local context cache instead of global.
local_cache = cache.ContextCache()
move = SequenceMove([LangevinDynamicsMove(n_steps=5),
GHMCMove(n_steps=5)],
context_cache=local_cache)
for m in move:
assert m.context_cache == local_cache
# Running with the local cache doesn't affect the global one.
cache.global_context_cache = cache.ContextCache() # empty global
move.apply(thermodynamic_state, sampler_state)
assert len(cache.global_context_cache) == 0
assert len(local_cache) == 2
# DummyContextCache works for all platforms.
platforms = utils.get_available_platforms()
dummy_cache = cache.DummyContextCache()
for platform in platforms:
dummy_cache.platform = platform
move = LangevinDynamicsMove(n_steps=5, context_cache=dummy_cache)
move.apply(thermodynamic_state, sampler_state)
assert len(cache.global_context_cache) == 0
def get_states():
# Set Simulation parameters
temperature = 200 * unit.kelvin
collision_rate = 1 / unit.picoseconds
timestep = 1.0 * unit.femtoseconds
n_steps = 20
nIter = 100
alchemical_atoms = [2, 3, 4, 5, 6, 7]
platform = openmm.Platform.getPlatformByName('CPU')
context_cache = cache.ContextCache(platform)
# Load a Parmed Structure for the Topology and create our openmm.Simulation
structure_pdb = utils.get_data_filename(
'blues', 'tests/data/ethylene_structure.pdb')
structure = parmed.load_file(structure_pdb)
# Load our OpenMM System and create Integrator
system_xml = utils.get_data_filename('blues',
'tests/data/ethylene_system.xml')
with open(system_xml, 'r') as infile:
xml = infile.read()
system = openmm.XmlSerializer.deserialize(xml)
thermodynamic_state = ThermodynamicState(system=system,
temperature=temperature)
sampler_state = SamplerState(
positions=structure.positions.in_units_of(unit.nanometers))
alch_system = generateAlchSystem(thermodynamic_state.get_system(),
alchemical_atoms)
alch_state = alchemy.AlchemicalState.from_system(alch_system)
alch_thermodynamic_state = ThermodynamicState(
alch_system, thermodynamic_state.temperature)
alch_thermodynamic_state = CompoundThermodynamicState(
alch_thermodynamic_state, composable_states=[alch_state])
return structure, thermodynamic_state, alch_thermodynamic_state
def from_amber(cls, prmtop, inpcrd, temperature=50 * u.kelvin):
prmtop = app.AmberPrmtopFile(prmtop)
inpcrd = app.AmberInpcrdFile(inpcrd)
system = prmtop.createSystem(nonbondedMethod=app.PME,
constraints=app.HBonds,
nonbondedCutoff=10 * u.angstroms,
switchDistance=8 * u.angstroms)
thermodynamic_state = ThermodynamicState(system,
temperature=temperature)
sampler_state = SamplerState(
positions=inpcrd.getPositions(asNumpy=True),
box_vectors=inpcrd.boxVectors)
return Esmacs(thermodynamic_state, sampler_state, prmtop.topology)
def _system_to_thermodynamic_state(self, system):
"""
Given an OpenMM system object, create a corresponding ThermodynamicState that has the same
temperature and pressure as the current thermodynamic state.
Parameters
----------
system : openmm.System
The OpenMM system for which to create the thermodynamic state
Returns
-------
new_thermodynamic_state : openmmtools.states.ThermodynamicState
The thermodynamic state object representing the given system
"""
return ThermodynamicState(system, temperature=self._temperature, pressure=self._pressure)
def compare_energy_components(rest_system,
other_system,
positions,
platform=REFERENCE_PLATFORM):
"""
Get energy components of a given system
"""
platform = configure_platform(platform)
# Create thermodynamic state and sampler state for non-rest system
thermostate_other = ThermodynamicState(system=other_system,
temperature=temperature)
# Create context for non-rest system
integrator_other = openmm.VerletIntegrator(1.0 * unit.femtosecond)
context_other = thermostate_other.create_context(integrator_other)
context_other.setPositions(positions)
# Get energy components for non-rest system
components_other = [
component[1]
for component in compute_potential_components(context_other, beta=beta)
]
# Create thermodynamic state for rest_system
thermostate_rest = ThermodynamicState(system=rest_system,
temperature=temperature)
# Create context forrest system
integrator_rest = openmm.VerletIntegrator(1.0 * unit.femtosecond)
context_rest = thermostate_rest.create_context(integrator_rest)
context_rest.setPositions(positions)
# Get energy components for rest system
components_rest = [
component[1]
for component in compute_potential_components(context_rest, beta=beta)
]
# Check that bond, angle, and torsion energies match
for other, rest in zip(components_other[:3], components_rest[:3]):
assert np.isclose(
[other], [rest]
), f"The energies do not match for the {other[0]}: {other[1]} (other system) vs. {rest[1]} (REST system)"
# Check that nonbonded energies match
print(components_rest)
nonbonded_other = np.array(components_other[3:]).sum()
nonbonded_rest = np.array(components_rest[3:]).sum()
assert np.isclose(
[nonbonded_other], [nonbonded_rest]
), f"The energies do not match for the NonbondedForce: {nonbonded_other} (other system) vs. {nonbonded_rest} (REST system)"
def __init__(self,
system,
alchemical_composability=AlchemicalState,
temperature=300 * unit.kelvin,
pressure=1.0 * unit.atmosphere,
**kwargs):
"""
Subclass of coddiwomple.states.PDFState that specifically handles openmmtools.states.CompoundThermodynamicStates
Init method does the following:
1. create a openmmtools.states.ThermodynamicState with the given system, temperature, and pressure
2. create an openmmtools.alchemy.AlchemicalState from the given system
3. create an openmtools.states.CompoundThermodynamicState 1 and 2
arguments
system : openmm.System
system object to wrap
alchemical_composability : openmmtools.alchemy.AlchemicalState (super), default openmmtools.alchemy.AlchemicalState
the class holding the method `from_system` which can compose an alchemical state with exposed parameters from a system
temperature : float * unit.kelvin (or temperature units), default 300.0 * unit.kelvin
temperature of the system
pressure : float * unit.atmosphere (or pressure units), default 1.0 * unit.atmosphere
pressure of the system
init method is adapted from https://github.com/choderalab/openmmtools/blob/110524bc5079af77d31f5fab464edd7b668ff5ac/openmmtools/states.py#L2766-L2783
"""
from openmmtools.alchemy import AlchemicalState
openmm_pdf_state = ThermodynamicState(system, temperature, pressure)
alchemical_state = alchemical_composability.from_system(
system, **kwargs)
assert isinstance(
alchemical_state, IComposableState
), f"alchemical state is not an instance of IComposableState"
self.__dict__ = openmm_pdf_state.__dict__
self._composable_states = [alchemical_state]
self.set_system(self._standard_system, fix_state=True)
#class create an internal context
integrator = get_dummy_integrator()
platform = configure_platform(utils.get_fastest_platform().getName())
self._internal_context = self.create_context(integrator,
platform=platform)
_logger.debug(
f"successfully instantiated OpenMMPDFState equipped with the following parameters: {self._parameters}"
)
def test_create_endstates():
"""
test the creation of unsampled endstates
"""
from pkg_resources import resource_filename
smiles_filename = resource_filename("perses",
os.path.join("data", "test.smi"))
fe_setup = RelativeFEPSetup(ligand_input=smiles_filename,
old_ligand_index=0,
new_ligand_index=1,
forcefield_files=[],
small_molecule_forcefield='gaff-2.11',
phases=['vacuum'])
hybrid_factory = HybridTopologyFactory(
topology_proposal=fe_setup._vacuum_topology_proposal,
current_positions=fe_setup._vacuum_positions_old,
new_positions=fe_setup._vacuum_positions_new,
neglected_new_angle_terms=fe_setup._vacuum_forward_neglected_angles,
neglected_old_angle_terms=fe_setup._vacuum_reverse_neglected_angles,
softcore_LJ_v2=True,
interpolate_old_and_new_14s=False)
zero_state_error, one_state_error = validate_endstate_energies(
fe_setup._vacuum_topology_proposal,
hybrid_factory,
added_energy=fe_setup._vacuum_added_valence_energy,
subtracted_energy=fe_setup._vacuum_subtracted_valence_energy,
beta=beta,
platform=openmm.Platform.getPlatformByName('Reference'),
ENERGY_THRESHOLD=ENERGY_THRESHOLD)
lambda_alchemical_state = RelativeAlchemicalState.from_system(
hybrid_factory.hybrid_system)
lambda_protocol = LambdaProtocol(functions='default')
lambda_alchemical_state.set_alchemical_parameters(0.0, lambda_protocol)
thermodynamic_state = CompoundThermodynamicState(
ThermodynamicState(hybrid_factory.hybrid_system,
temperature=temperature),
composable_states=[lambda_alchemical_state])
zero_endstate = copy.deepcopy(thermodynamic_state)
one_endstate = copy.deepcopy(thermodynamic_state)
one_endstate.set_alchemical_parameters(1.0, lambda_protocol)
new_endstates = create_endstates(zero_endstate, one_endstate)
def test_minimizer_all_testsystems():
# testsystem_classes = testsystems.TestSystem.__subclasses__()
testsystem_classes = [testsystems.AlanineDipeptideVacuum]
for testsystem_class in testsystem_classes:
class_name = testsystem_class.__name__
logging.info("Testing minimization with testsystem %s" % class_name)
testsystem = testsystem_class()
sampler_state = SamplerState(testsystem.positions)
thermodynamic_state = ThermodynamicState(testsystem.system, 300*unit.kelvin)
# Create sampler for minimization.
sampler = MCMCSampler(thermodynamic_state, sampler_state, move=None)
sampler.minimize(max_iterations=0)
# Check if NaN.
err_msg = 'Minimization of system {} yielded NaN'.format(class_name)
assert not sampler_state.has_nan(), err_msg
def create_langevin_integrator(htf, constraint_tol):
"""
create lambda alchemical states, thermodynamic states, sampler states, integrator, and return context, thermostate, sampler_state, integrator
"""
fast_lambda_alchemical_state = RelativeAlchemicalState.from_system(
htf.hybrid_system)
fast_lambda_alchemical_state.set_alchemical_parameters(
0.0, LambdaProtocol(functions='default'))
fast_thermodynamic_state = CompoundThermodynamicState(
ThermodynamicState(htf.hybrid_system, temperature=temperature),
composable_states=[fast_lambda_alchemical_state])
fast_sampler_state = SamplerState(
positions=htf._hybrid_positions,
box_vectors=htf.hybrid_system.getDefaultPeriodicBoxVectors())
integrator_1 = integrators.LangevinIntegrator(
temperature=temperature,
timestep=4.0 * unit.femtoseconds,
splitting='V R O R V',
measure_shadow_work=False,
measure_heat=False,
constraint_tolerance=constraint_tol,
collision_rate=5.0 / unit.picoseconds)
# mcmc_moves=mcmc.LangevinSplittingDynamicsMove(timestep = 4.0 * unit.femtoseconds,
# collision_rate=5.0 / unit.picosecond,
# n_steps=1,
# reassign_velocities=False,
# n_restart_attempts=20,
# splitting="V R R R O R R R V",
# constraint_tolerance=constraint_tol)
#print(integrator_1.getConstraintTolerance())
fast_context, fast_integrator = cache.global_context_cache.get_context(
fast_thermodynamic_state, integrator_1)
fast_sampler_state.apply_to_context(fast_context)
return fast_context, fast_thermodynamic_state, fast_sampler_state, fast_integrator
def compute_reduced_potential(thermodynamic_state: states.ThermodynamicState,
sampler_state: states.SamplerState) -> float:
"""
Compute the reduced potential of the given SamplerState under the given ThermodynamicState.
Parameters
----------
thermodynamic_state : openmmtools.states.ThermodynamicState
The thermodynamic state under which to compute the reduced potential
sampler_state : openmmtools.states.SamplerState
The sampler state for which to compute the reduced potential
Returns
-------
reduced_potential : float
unitless reduced potential (kT)
"""
context, integrator = cache.global_context_cache.get_context(
thermodynamic_state)
sampler_state.apply_to_context(context, ignore_velocities=True)
return thermodynamic_state.reduced_potential(context)
def test_barostat_move_frequency():
"""MonteCarloBarostatMove restore barostat's frequency afterwards."""
# Get periodic test case.
for test_case in analytical_testsystems:
testsystem = test_case[1]
if testsystem.system.usesPeriodicBoundaryConditions():
break
assert testsystem.system.usesPeriodicBoundaryConditions(), "Can't find periodic test case!"
sampler_state = SamplerState(testsystem.positions)
thermodynamic_state = ThermodynamicState(testsystem.system, 298*unit.kelvin,
1*unit.atmosphere)
move = MonteCarloBarostatMove(n_attempts=5)
# Test-precondition: the frequency must be different than 1 or it
# will never change during the application of the MCMC move.
old_frequency = thermodynamic_state.barostat.getFrequency()
assert old_frequency != 1
move.apply(thermodynamic_state, sampler_state)
assert thermodynamic_state.barostat.getFrequency() == old_frequency
def test_create_endstates():
"""
test the creation of unsampled endstates
"""
fe_setup = RelativeFEPSetup(ligand_input=f"{os.getcwd()}/test.smi",
old_ligand_index=0,
new_ligand_index=1,
forcefield_files=['gaff.xml'],
phases=['vacuum'])
hybrid_factory = HybridTopologyFactory(
topology_proposal=fe_setup._vacuum_topology_proposal,
current_positions=fe_setup._vacuum_positions_old,
new_positions=fe_setup._vacuum_positions_new,
neglected_new_angle_terms=fe_setup._vacuum_forward_neglected_angles,
neglected_old_angle_terms=fe_setup._vacuum_reverse_neglected_angles,
softcore_LJ_v2=True,
interpolate_old_and_new_14s=False)
zero_state_error, one_state_error = validate_endstate_energies(
fe_setup._vacuum_topology_proposal,
hybrid_factory,
added_energy=fe_setup._vacuum_added_valence_energy,
subtracted_energy=fe_setup._vacuum_subtracted_valence_energy,
beta=beta,
ENERGY_THRESHOLD=ENERGY_THRESHOLD)
lambda_alchemical_state = RelativeAlchemicalState.from_system(
hybrid_factory.hybrid_system)
lambda_protocol = LambdaProtocol(functions='default')
lambda_alchemical_state.set_alchemical_parameters(0.0, lambda_protocol)
thermodynamic_state = CompoundThermodynamicState(
ThermodynamicState(hybrid_factory.hybrid_system,
temperature=temperature),
composable_states=[lambda_alchemical_state])
zero_endstate = copy.deepcopy(thermodynamic_state)
one_endstate = copy.deepcopy(thermodynamic_state)
one_endstate.set_alchemical_parameters(1.0, lambda_protocol)
new_endstates = create_endstates(zero_endstate, one_endstate)
def compute_reduced_potential(thermodynamic_state: states.ThermodynamicState, sampler_state: states.SamplerState) -> float:
"""
Compute the reduced potential of the given SamplerState under the given ThermodynamicState.
Parameters
----------
thermodynamic_state : openmmtools.states.ThermodynamicState
The thermodynamic state under which to compute the reduced potential
sampler_state : openmmtools.states.SamplerState
The sampler state for which to compute the reduced potential
Returns
-------
reduced_potential : float
unitless reduced potential (kT)
"""
if type(cache.global_context_cache) == cache.DummyContextCache:
integrator = openmm.VerletIntegrator(1.0) #we won't take any steps, so use a simple integrator
context, integrator = cache.global_context_cache.get_context(thermodynamic_state, integrator)
else:
context, integrator = cache.global_context_cache.get_context(thermodynamic_state)
sampler_state.apply_to_context(context, ignore_velocities=True)
return thermodynamic_state.reduced_potential(context)
def from_testsystem( # pylint: disable=too-many-arguments
self,
test,
reference_state,
thermodynamic_states,
pressure=None,
storage=None,
target_state=0,
metadata=None,
**kwargs):
"""Initialize sampler from TestSystem object."""
if not isinstance(thermodynamic_states, Sequence): # a scalar
thermodynamic_states = [thermodynamic_states]
# check if temp or thermodynamic states or something else
if not isinstance(thermodynamic_states[0], ThermodynamicState):
# as temperatures
temperatures = thermodynamic_states
if not isinstance(temperatures[0], unit.Quantity): # no units
temperatures = [t * unit.kelvin for t in temperatures]
thermodynamic_states = [
ThermodynamicState(
system=test.system,
temperature=t,
pressure=pressure,
) for t in temperatures
]
sampler_states = SamplerState(positions=test.positions,
box_vectors=test.default_box_vectors)
self.create(reference_state,
thermodynamic_states,
sampler_states,
test.topology,
target_state=target_state,
storage=storage,
metadata=metadata,
**kwargs)
def create_endstates(first_thermostate, last_thermostate):
"""
utility function to generate unsampled endstates
1. move all alchemical atom LJ parameters from CustomNonbondedForce to NonbondedForce
2. delete the CustomNonbondedForce
3. set PME tolerance to 1e-5
4. enable LJPME to handle long range dispersion corrections in a physically reasonable manner
Arguments
---------
first_thermostate : openmmtools.states.CompoundThermodynamicState
the first thermodynamic state for which an unsampled endstate will be created
last_thermostate : openmmtools.states.CompoundThermodynamicState
the last thermodynamic state for which an unsampled endstate will be created
Returns
-------
unsampled_endstates : list of openmmtools.states.CompoundThermodynamicState
the corrected unsampled endstates
"""
unsampled_endstates = []
for master_lambda, endstate in zip([0., 1.],
[first_thermostate, last_thermostate]):
dispersion_system = endstate.get_system()
energy_unit = unit.kilocalories_per_mole
# Find the NonbondedForce (there must be only one)
forces = {
force.__class__.__name__: force
for force in dispersion_system.getForces()
}
# Set NonbondedForce to use LJPME
forces['NonbondedForce'].setNonbondedMethod(
openmm.NonbondedForce.LJPME)
# Set tight PME tolerance
TIGHT_PME_TOLERANCE = 1.0e-5
forces['NonbondedForce'].setEwaldErrorTolerance(TIGHT_PME_TOLERANCE)
# Move alchemical LJ sites from CustomNonbondedForce back to NonbondedForce
for particle_index in range(
forces['NonbondedForce'].getNumParticles()):
charge, sigma, epsilon = forces[
'NonbondedForce'].getParticleParameters(particle_index)
sigmaA, epsilonA, sigmaB, epsilonB, unique_old, unique_new = forces[
'CustomNonbondedForce'].getParticleParameters(particle_index)
if (epsilon / energy_unit == 0.0) and ((epsilonA > 0.0) or
(epsilonB > 0.0)):
sigma = (1 - master_lambda) * sigmaA + master_lambda * sigmaB
epsilon = (1 -
master_lambda) * epsilonA + master_lambda * epsilonB
forces['NonbondedForce'].setParticleParameters(
particle_index, charge, sigma, epsilon)
# Delete the CustomNonbondedForce since we have moved all alchemical particles out of it
for force_index, force in enumerate(list(
dispersion_system.getForces())):
if force.__class__.__name__ == 'CustomNonbondedForce':
custom_nonbonded_force_index = force_index
break
dispersion_system.removeForce(custom_nonbonded_force_index)
# Set all parameters to master lambda
for force_index, force in enumerate(list(
dispersion_system.getForces())):
if hasattr(force, 'getNumGlobalParameters'):
for parameter_index in range(force.getNumGlobalParameters()):
if force.getGlobalParameterName(
parameter_index)[0:7] == 'lambda_':
force.setGlobalParameterDefaultValue(
parameter_index, master_lambda)
# Store the unsampled endstate
unsampled_endstates.append(
ThermodynamicState(dispersion_system,
temperature=endstate.temperature))
return unsampled_endstates
def run_equilibrium(equilibrium_result: EquilibriumResult, thermodynamic_state: states.ThermodynamicState,
nsteps_equil: int, topology: md.Topology, n_iterations : int,
atom_indices_to_save: List[int] = None, trajectory_filename: str = None, splitting: str="V R O R V", timestep: unit.Quantity=1.0*unit.femtoseconds) -> EquilibriumResult:
"""
Run nsteps of equilibrium sampling at the specified thermodynamic state and return the final sampler state
as well as a trajectory of the positions after each application of an MCMove. This means that if the MCMove
is configured to run 1000 steps of dynamics, and n_iterations is 100, there will be 100 frames in the resulting
trajectory; these are the result of 100,000 steps (1000*100) of dynamics.
Parameters
----------
equilibrium_result : EquilibriumResult
EquilibriumResult namedtuple containing the information necessary to resume
thermodynamic_state : openmmtools.states.ThermodynamicState
The thermodynamic state (including context parameters) that should be used
nsteps_equil : int
The number of equilibrium steps that a move should make when apply is called
topology : mdtraj.Topology
an MDTraj topology object used to construct the trajectory
n_iterations : int
The number of times to apply the move. Note that this is not the number of steps of dynamics; it is
n_iterations*n_steps (which is set in the MCMove).
splitting: str, default "V R O H R V"
The splitting string for the dynamics
atom_indices_to_save : list of int, default None
list of indices to save (when excluding waters, for instance). If None, all indices are saved.
trajectory_filename : str, optional, default None
Full filepath of trajectory files. If none, trajectory files are not written.
splitting: str, default "V R O H R V"
The splitting string for the dynamics
Returns
-------
equilibrium_result : EquilibriumResult
Container namedtuple that has the SamplerState for resuming, an MDTraj trajectory, and the reduced potential of the
final frame.
"""
sampler_state = equilibrium_result.sampler_state
#get the atom indices we need to subset the topology and positions
if atom_indices_to_save is None:
atom_indices = list(range(topology.n_atoms))
subset_topology = topology
else:
subset_topology = topology.subset(atom_indices_to_save)
atom_indices = atom_indices_to_save
n_atoms = subset_topology.n_atoms
#construct the MCMove:
mc_move = mcmc.LangevinSplittingDynamicsMove(n_steps=nsteps_equil, splitting=splitting)
mc_move.n_restart_attempts = 10
#create a numpy array for the trajectory
trajectory_positions = np.zeros([n_iterations, n_atoms, 3])
trajectory_box_lengths = np.zeros([n_iterations, 3])
trajectory_box_angles = np.zeros([n_iterations, 3])
#loop through iterations and apply MCMove, then collect positions into numpy array
for iteration in range(n_iterations):
mc_move.apply(thermodynamic_state, sampler_state)
trajectory_positions[iteration, :] = sampler_state.positions[atom_indices, :].value_in_unit_system(unit.md_unit_system)
#get the box lengths and angles
a, b, c, alpha, beta, gamma = mdtrajutils.unitcell.box_vectors_to_lengths_and_angles(*sampler_state.box_vectors)
trajectory_box_lengths[iteration, :] = [a, b, c]
trajectory_box_angles[iteration, :] = [alpha, beta, gamma]
#construct trajectory object:
trajectory = md.Trajectory(trajectory_positions, subset_topology, unitcell_lengths=trajectory_box_lengths, unitcell_angles=trajectory_box_angles)
#get the reduced potential from the final frame for endpoint perturbations
reduced_potential_final_frame = thermodynamic_state.reduced_potential(sampler_state)
#construct equilibrium result object
equilibrium_result = EquilibriumResult(sampler_state, reduced_potential_final_frame)
#If there is a trajectory filename passed, write out the results here:
if trajectory_filename is not None:
write_equilibrium_trajectory(equilibrium_result, trajectory, trajectory_filename)
return equilibrium_result
def setup(self,
n_states,
temperature,
storage_file,
minimisation_steps=100,
n_replicas=None,
lambda_schedule=None,
lambda_protocol=LambdaProtocol(),
endstates=True):
from perses.dispersed import feptasks
hybrid_system = self._factory.hybrid_system
positions = self._factory.hybrid_positions
lambda_zero_alchemical_state = RelativeAlchemicalState.from_system(
hybrid_system)
thermostate = ThermodynamicState(hybrid_system,
temperature=temperature)
compound_thermodynamic_state = CompoundThermodynamicState(
thermostate, composable_states=[lambda_zero_alchemical_state])
thermodynamic_state_list = []
sampler_state_list = []
context_cache = cache.ContextCache()
if n_replicas is None:
_logger.info(
f'n_replicas not defined, setting to match n_states, {n_states}'
)
n_replicas = n_states
elif n_replicas > n_states:
_logger.warning(
f'More sampler states: {n_replicas} requested greater than number of states: {n_states}. Setting n_replicas to n_states: {n_states}'
)
n_replicas = n_states
# TODO this feels like it should be somewhere else... just not sure where. Maybe into lambda_protocol
if lambda_schedule is None:
lambda_schedule = np.linspace(0., 1., n_states)
else:
assert (
len(lambda_schedule) == n_states
), 'length of lambda_schedule must match the number of states, n_states'
assert (
lambda_schedule[0] == 0.), 'lambda_schedule must start at 0.'
assert (
lambda_schedule[-1] == 1.), 'lambda_schedule must end at 1.'
difference = np.diff(lambda_schedule)
assert (all(i >= 0. for i in difference)
), 'lambda_schedule must be monotonicly increasing'
#starting with the initial positions generated py geometry.py
sampler_state = SamplerState(
positions,
box_vectors=hybrid_system.getDefaultPeriodicBoxVectors())
for lambda_val in lambda_schedule:
compound_thermodynamic_state_copy = copy.deepcopy(
compound_thermodynamic_state)
compound_thermodynamic_state_copy.set_alchemical_parameters(
lambda_val, lambda_protocol)
thermodynamic_state_list.append(compound_thermodynamic_state_copy)
# now generating a sampler_state for each thermodyanmic state, with relaxed positions
context, context_integrator = context_cache.get_context(
compound_thermodynamic_state_copy)
feptasks.minimize(compound_thermodynamic_state_copy, sampler_state)
sampler_state_list.append(copy.deepcopy(sampler_state))
reporter = storage_file
# making sure number of sampler states equals n_replicas
if len(sampler_state_list) != n_replicas:
# picking roughly evenly spaced sampler states
# if n_replicas == 1, then it will pick the first in the list
idx = np.round(
np.linspace(0,
len(sampler_state_list) - 1,
n_replicas)).astype(int)
sampler_state_list = [
state for i, state in enumerate(sampler_state_list) if i in idx
]
assert len(sampler_state_list) == n_replicas
if endstates:
# generating unsampled endstates
_logger.info('Generating unsampled endstates.')
unsampled_dispersion_endstates = create_endstates(
copy.deepcopy(thermodynamic_state_list[0]),
copy.deepcopy(thermodynamic_state_list[-1]))
self.create(
thermodynamic_states=thermodynamic_state_list,
sampler_states=sampler_state_list,
storage=reporter,
unsampled_thermodynamic_states=unsampled_dispersion_endstates)
else:
self.create(thermodynamic_states=thermodynamic_state_list,
sampler_states=sampler_state_list,
storage=reporter)
prmtop = utils.get_data_filename('blues',
'tests/data/eqToluene.prmtop') #TOL-parm
inpcrd = utils.get_data_filename('blues', 'tests/data/eqToluene.inpcrd')
tol = parmed.load_file(prmtop, xyz=inpcrd)
tol.system = tol.createSystem(nonbondedMethod=openmm.app.PME,
nonbondedCutoff=10 * unit.angstrom,
constraints=openmm.app.HBonds,
hydrogenMass=3.024 * unit.dalton,
rigidWater=True,
removeCMMotion=True,
flexibleConstraints=True,
splitDihedrals=False)
# Create our State objects
sampler_state = SamplerState(positions=tol.positions)
thermodynamic_state = ThermodynamicState(system=tol.system,
temperature=temperature)
md_reporter = NetCDF4Storage(outfname + '.nc', reportInterval)
state_reporter = BLUESStateDataStorage(reportInterval=reportInterval,
title='md',
step=True,
speed=True,
progress=True,
totalSteps=int(n_steps * nIter))
ncmc_state_reporter = BLUESStateDataStorage(reportInterval=reportInterval,
title='ncmc',
step=True,
speed=True,
progress=True,
def ANI_endstate_sampler(
system,
system_subset,
subset_indices_map,
positions_cache_filename,
md_topology,
index_to_run,
steps_per_application,
integrator_kwargs = {'temperature': 300.0 * unit.kelvin,
'collision_rate': 1.0 / unit.picoseconds,
'timestep': 2.0 * unit.femtoseconds,
'splitting': "V R O R F",
'constraint_tolerance': 1e-6,
'pressure': 1.0 * unit.atmosphere},
position_extractor = None
):
"""
conduct ani endstate sampling
arguments
system : openmm.System
system
system_subset : openmm.System
subset system
subset_indices_map : dict
dict of {openmm_pdf_state atom_index : openmm_pdf_state_subset atom index}
positions_cache_filename : str
path to the cache positions
index_to_run : int
index of the positions to anneal
number_of_applications : int
number of applications of the propagator
steps_per_application : int
number of integration steps per application
integrator_kwargs : dict, see default
kwargs to pass to OMMLIAIS integrator
"""
from coddiwomple.particles import Particle
from coddiwomple.openmm.states import OpenMMParticleState, OpenMMPDFState
#load the endstate cache
traj = np.load(positions_cache_filename)
positions = traj['positions'][index_to_run,:,:] * unit.nanometers
if position_extractor is not None:
positions = position_extractor(_positions)
try:
box_vectors = traj['box_vectors'][index_to_run,:,:] * unit.nanometers
except Exception as e:
box_vectors = None
species_str = ''.join([atom.element.symbol for atom in md_topology.subset(list(subset_indices_map.keys())).atoms])
_logger.info(f"species string: {species_str}")
ani_handler = ANI1_force_and_energy(model = torchani.models.ANI1ccx(),
atoms=species_str,
platform='cpu',
temperature=integrator_kwargs['temperature'])
#make thermostates
pressure = integrator_kwargs['pressure'] if box_vectors is not None else None
pdf_state = ThermodynamicState(system = system, temperature = integrator_kwargs['temperature'], pressure=pressure)
pdf_state_subset = ThermodynamicState(system = system_subset, temperature = integrator_kwargs['temperature'], pressure = None)
#make an integrator
integrator = Integrator(**integrator_kwargs)
#make a propagator
propagator = ANIPropagator(openmm_pdf_state = pdf_state,
openmm_pdf_state_subset = pdf_state_subset,
subset_indices_map = subset_indices_map,
integrator = integrator,
ani_handler = ani_handler,
context_cache=None,
reassign_velocities=True,
n_restart_attempts=0,
reporter = None)
particle = Particle(0)
particle_state = OpenMMParticleState(positions = positions, box_vectors = box_vectors)
particle.update_state(particle_state)
particle_state, _return_dict = propagator.apply(particle_state, n_steps = steps_per_application, reset_integrator=True, apply_pdf_to_context=True)
if box_vectors is None:
particle_state.box_vectors=None
return particle_state, np.array(propagator.state_works[0])
def annealed_importance_sampling(direction,
system,
system_subset,
subset_indices_map,
positions_cache_filename,
index_to_run,
directory_name,
trajectory_prefix,
md_topology,
steps_per_application,
integrator_kwargs = {'temperature': 300.0 * unit.kelvin,
'collision_rate': 1.0 / unit.picoseconds,
'timestep': 1.0 * unit.femtoseconds,
'splitting': "V R O R F",
'constraint_tolerance': 1e-6,
'pressure': 1.0 * unit.atmosphere},
save_indices = None,
position_extractor = None,
write_trajectory_interval=1
):
"""
conduct annealed importance sampling in the openmm regime; will write the accumulated work dictionary after each application
arguments
direction : str
forward or backward
system : openmm.System
system
system_subset : openmm.System
subset system
subset_indices_map : dict
dict of {openmm_pdf_state atom_index : openmm_pdf_state_subset atom index}
positions_cache_filename : str
path to the cache positions
index_to_run : int
index of the positions to anneal
directory_name : str
directory that will be written to
trajectory_prefix : str
.pdb prefix
md_topology : mdtraj.Topology
topology that will write the trajectory
save_indices : list
list of indices that will be saved
number_of_applications : int
number of applications of the propagator
steps_per_application : int
number of integration steps per application
integrator_kwargs : dict, see default
kwargs to pass to OMMLIAIS integrator
save_indices : list(int)
list of indices of md_topology atoms to save to disk
position_extractor : function, default None
function to extract appropriate positons from the cache
write_trajectory_interval : int
frequency with which to write trajectory to disk
"""
from coddiwomple.particles import Particle
from coddiwomple.openmm.states import OpenMMParticleState, OpenMMPDFState
#load the endstate cache
traj = np.load(positions_cache_filename)
positions = traj['positions'][index_to_run,:,:] * unit.nanometers
if position_extractor is not None:
positions = position_extractor(_positions)
try:
box_vectors = traj['box_vectors'][index_to_run,:,:] * unit.nanometers
except Exception as e:
box_vectors = None
assert direction in ['forward', 'backward']
#make a handle object for ANI
species_str = ''.join([atom.element.symbol for atom in md_topology.subset(list(subset_indices_map.keys())).atoms])
_logger.info(f"species string: {species_str}")
ani_handler = ANI1_force_and_energy(model = torchani.models.ANI1ccx(),
atoms=species_str,
platform='cpu',
temperature=integrator_kwargs['temperature'])
#make thermostates
pressure = integrator_kwargs['pressure'] if box_vectors is not None else None
pdf_state = ThermodynamicState(system = system, temperature = integrator_kwargs['temperature'], pressure=pressure)
pdf_state_subset = ThermodynamicState(system = system_subset, temperature = integrator_kwargs['temperature'], pressure = None)
#make a reporter
saveable_topology = md_topology.subset(save_indices)
reporter = OpenMMReporter(directory_name, trajectory_prefix, saveable_topology, subset_indices = save_indices)
#make an integrator
integrator = Integrator(**integrator_kwargs)
#make a propagator
if direction == 'forward':
propagator = Propagator(openmm_pdf_state = pdf_state,
openmm_pdf_state_subset = pdf_state_subset,
subset_indices_map = subset_indices_map,
integrator = integrator,
ani_handler = ani_handler,
context_cache=None,
reassign_velocities=True,
n_restart_attempts=0,
reporter = reporter,
write_trajectory_interval = write_trajectory_interval)
else:
propagator = BackwardPropagator(openmm_pdf_state = pdf_state,
openmm_pdf_state_subset = pdf_state_subset,
subset_indices_map = subset_indices_map,
integrator = integrator,
ani_handler = ani_handler,
context_cache=None,
reassign_velocities=True,
n_restart_attempts=0,
reporter = reporter,
write_trajectory_interval = write_trajectory_interval)
particle = Particle(0)
particle_state = OpenMMParticleState(positions = positions, box_vectors = box_vectors)
particle.update_state(particle_state)
particle_state, _return_dict = propagator.apply(particle_state, n_steps = steps_per_application, reset_integrator=True, apply_pdf_to_context=True)
if box_vectors is None:
particle_state.box_vectors=None
return particle_state, np.array(propagator.state_works[0])
def subtest_mcmc_expectation(testsystem, move):
if debug:
print(testsystem.__class__.__name__)
print(str(move))
# Retrieve system and positions.
[system, positions] = [testsystem.system, testsystem.positions]
# Test settings.
temperature = 298.0 * unit.kelvin
niterations = 500 # number of production iterations
if system.usesPeriodicBoundaryConditions():
pressure = 1.0 * unit.atmosphere
else:
pressure = None
# Compute properties.
kB = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA
kT = kB * temperature
ndof = 3 * system.getNumParticles() - system.getNumConstraints()
# Create sampler and thermodynamic state.
sampler_state = SamplerState(positions=positions)
thermodynamic_state = ThermodynamicState(system=system,
temperature=temperature,
pressure=pressure)
# Create MCMC sampler
sampler = MCMCSampler(thermodynamic_state, sampler_state, move=move)
# Accumulate statistics.
x_n = np.zeros(
[niterations], np.float64
) # x_n[i] is the x position of atom 1 after iteration i, in angstroms
potential_n = np.zeros(
[niterations], np.float64
) # potential_n[i] is the potential energy after iteration i, in kT
kinetic_n = np.zeros(
[niterations], np.float64
) # kinetic_n[i] is the kinetic energy after iteration i, in kT
temperature_n = np.zeros(
[niterations], np.float64
) # temperature_n[i] is the instantaneous kinetic temperature from iteration i, in K
volume_n = np.zeros(
[niterations],
np.float64) # volume_n[i] is the volume from iteration i, in K
for iteration in range(niterations):
# Update sampler state.
sampler.run(1)
# Get statistics.
potential_energy = sampler.sampler_state.potential_energy
kinetic_energy = sampler.sampler_state.kinetic_energy
instantaneous_temperature = kinetic_energy * 2.0 / ndof / kB
volume = sampler.sampler_state.volume
# Accumulate statistics.
x_n[iteration] = sampler_state.positions[0, 0] / unit.angstroms
potential_n[iteration] = potential_energy / kT
kinetic_n[iteration] = kinetic_energy / kT
temperature_n[iteration] = instantaneous_temperature / unit.kelvin
volume_n[iteration] = volume / (unit.nanometers**3)
# Compute expected statistics.
if (hasattr(testsystem, 'get_potential_expectation') and
testsystem.get_potential_standard_deviation(thermodynamic_state) /
kT.unit != 0.0):
assert potential_n.std(
) != 0.0, 'Test {} shows no potential fluctuations'.format(
testsystem.__class__.__name__)
potential_expectation = testsystem.get_potential_expectation(
thermodynamic_state) / kT
[t0, g, Neff_max] = timeseries.detectEquilibration(potential_n)
potential_mean = potential_n[t0:].mean()
dpotential_mean = potential_n[t0:].std() / np.sqrt(Neff_max)
potential_error = potential_mean - potential_expectation
nsigma = abs(potential_error) / dpotential_mean
err_msg = (
'Potential energy expectation\n'
'observed {:10.5f} +- {:10.5f}kT | expected {:10.5f} | '
'error {:10.5f} +- {:10.5f} ({:.1f} sigma) | t0 {:5d} | g {:5.1f} | Neff {:8.1f}\n'
'----------------------------------------------------------------------------'
).format(potential_mean, dpotential_mean, potential_expectation,
potential_error, dpotential_mean, nsigma, t0, g, Neff_max)
assert nsigma <= NSIGMA_CUTOFF, err_msg.format()
if debug:
print(err_msg)
elif debug:
print('Skipping potential expectation test.')
if (hasattr(testsystem, 'get_volume_expectation')
and testsystem.get_volume_standard_deviation(thermodynamic_state) /
(unit.nanometers**3) != 0.0):
assert volume_n.std(
) != 0.0, 'Test {} shows no volume fluctuations'.format(
testsystem.__class__.__name__)
volume_expectation = testsystem.get_volume_expectation(
thermodynamic_state) / (unit.nanometers**3)
[t0, g, Neff_max] = timeseries.detectEquilibration(volume_n)
volume_mean = volume_n[t0:].mean()
dvolume_mean = volume_n[t0:].std() / np.sqrt(Neff_max)
volume_error = volume_mean - volume_expectation
nsigma = abs(volume_error) / dvolume_mean
err_msg = (
'Volume expectation\n'
'observed {:10.5f} +- {:10.5f}kT | expected {:10.5f} | '
'error {:10.5f} +- {:10.5f} ({:.1f} sigma) | t0 {:5d} | g {:5.1f} | Neff {:8.1f}\n'
'----------------------------------------------------------------------------'
).format(volume_mean, dvolume_mean, volume_expectation, volume_error,
dvolume_mean, nsigma, t0, g, Neff_max)
assert nsigma <= NSIGMA_CUTOFF, err_msg.format()
if debug:
print(err_msg)
elif debug:
print('Skipping volume expectation test.')
def integrate(self, topology_proposal, initial_sampler_state, proposed_sampler_state, iteration=None):
"""
Performs NCMC switching to either delete or insert atoms according to the provided `topology_proposal`.
For `delete`, the system is first modified from fully interacting to alchemically modified, and then NCMC switching is used to eliminate atoms.
For `insert`, the system begins with eliminated atoms in an alchemically noninteracting form and NCMC switching is used to turn atoms on, followed by making system real.
Parameters
----------
topology_proposal : TopologyProposal
Contains old/new Topology and System objects and atom mappings.
initial_sampler_state : openmmtools.states.SamplerState representing the initial (old) system
Configurational properties of the atoms at the beginning of the NCMC switching.
proposed_sampler_state : openmmtools.states.SamplerState representing the proposed (post-geometry new) system
Configurational properties new system atoms at beginning of NCMC switching
iteration : int, optional, default=None
Iteration number, for storage purposes.
Returns
-------
final_old_sampler_state : openmmtools.State.SamplerState
The final configurational properties of the old system after hybrid alchemical switching
final_sampler_state : openmmtools.states.SamplerState
The final configurational properties after `nsteps` steps of alchemical switching, and reversion to the nonalchemical system
logP_work : float
The NCMC work contribution to the log acceptance probability (Eqs. 62 and 63)
logP_initial : float
The initial logP of the hybrid configuration
logP_final : float
The final logP of the hybrid configuration
"""
assert not initial_sampler_state.has_nan() and not proposed_sampler_state.has_nan()
#generate or retrieve the hybrid topology factory:
hybrid_factory = self.make_alchemical_system(topology_proposal, initial_sampler_state.positions, proposed_sampler_state.positions)
if hybrid_factory is None:
_logger.warning("Unable to construct hybrid system for {} -> {}".format(topology_proposal.old_chemical_state_key, topology_proposal.new_chemical_state_key))
return initial_sampler_state, proposed_sampler_state, -np.inf, 0.0, 0.0
topology = hybrid_factory.hybrid_topology
#generate the corresponding thermodynamic and sampler states so that we can use the NonequilibriumSwitchingMove:
#First generate the thermodynamic state:
hybrid_system = hybrid_factory.hybrid_system
hybrid_thermodynamic_state = ThermodynamicState(hybrid_system, temperature=self._temperature, pressure=self._pressure)
#Now create an RelativeAlchemicalState from the hybrid system:
alchemical_state = RelativeAlchemicalState.from_system(hybrid_system)
alchemical_state.set_alchemical_parameters(0.0)
#Now create a compound thermodynamic state that combines the hybrid thermodynamic state with the alchemical state:
compound_thermodynamic_state = CompoundThermodynamicState(hybrid_thermodynamic_state, composable_states=[alchemical_state])
#construct a sampler state from the hybrid positions and the box vectors of the initial sampler state:
initial_hybrid_positions = hybrid_factory.hybrid_positions
initial_hybrid_box_vectors = initial_sampler_state.box_vectors
initial_hybrid_sampler_state = SamplerState(initial_hybrid_positions, box_vectors=initial_hybrid_box_vectors)
final_hybrid_sampler_state = copy.deepcopy(initial_hybrid_sampler_state)
#create the nonequilibrium move:
#ne_move = NonequilibriumSwitchingMove(self._functions, self._integrator_splitting, self._temperature, self._nsteps, self._timestep,
# work_save_interval=self._write_ncmc_interval, top=topology,subset_atoms=None,
# save_configuration=self._save_configuration, measure_shadow_work=self._measure_shadow_work)
ne_move = ExternalNonequilibriumSwitchingMove(self._functions, nsteps_neq=self._nsteps,
timestep=self._timestep, temperature=self._temperature,
work_configuration_save_interval=self._work_save_interval,
splitting="V R O R V")
#run the NCMC protocol
try:
ne_move.apply(compound_thermodynamic_state, final_hybrid_sampler_state)
except Exception as e:
_logger.warn("NCMC failed because {}; rejecting.".format(str(e)))
logP_work = -np.inf
return [initial_sampler_state, proposed_sampler_state, -np.inf, 0.0, 0.0]
#get the total work:
logP_work = - ne_move.cumulative_work[-1]
# Compute contribution of transforming to and from the hybrid system:
context, integrator = global_context_cache.get_context(hybrid_thermodynamic_state)
#set all alchemical parameters to zero:
for parameter in self._functions.keys():
context.setParameter(parameter, 0.0)
initial_hybrid_sampler_state.apply_to_context(context, ignore_velocities=True)
initial_reduced_potential = hybrid_thermodynamic_state.reduced_potential(context)
#set all alchemical parameters to one:
for parameter in self._functions.keys():
context.setParameter(parameter, 1.0)
final_hybrid_sampler_state.apply_to_context(context, ignore_velocities=True)
final_reduced_potential = hybrid_thermodynamic_state.reduced_potential(context)
#reset the parameters back to zero just in case
for parameter in self._functions.keys():
context.setParameter(parameter, 0.0)
#compute the output SamplerState, which has the atoms only for the new system post-NCMC:
new_positions = hybrid_factory.new_positions(final_hybrid_sampler_state.positions)
new_box_vectors = final_hybrid_sampler_state.box_vectors
final_sampler_state = SamplerState(new_positions, box_vectors=new_box_vectors)
#compute the output SamplerState for the atoms only in the old system (required for geometry_logP_reverse)
old_positions = hybrid_factory.old_positions(final_hybrid_sampler_state.positions)
old_box_vectors = copy.deepcopy(new_box_vectors) #these are the same as the new system
final_old_sampler_state = SamplerState(old_positions, box_vectors=old_box_vectors)
#extract the trajectory and box vectors from the move:
trajectory = ne_move.trajectory[::-self._write_ncmc_interval, :, :][::-1]
topology = hybrid_factory.hybrid_topology
position_varname = "ncmcpositions"
nframes = np.shape(trajectory)[0]
#extract box vectors:
box_vec_varname = "ncmcboxvectors"
box_lengths = ne_move.box_lengths[::-self._write_ncmc_interval, :][::-1]
box_angles = ne_move.box_angles[::-self._write_ncmc_interval, :][::-1]
box_lengths_and_angles = np.stack([box_lengths, box_angles])
#write out the positions of the topology
if self._storage:
for frame in range(nframes):
self._storage.write_configuration(position_varname, trajectory[frame, :, :], topology, iteration=iteration, frame=frame, nframes=nframes)
#write out the periodict box vectors:
if self._storage:
self._storage.write_array(box_vec_varname, box_lengths_and_angles, iteration=iteration)
#retrieve the protocol work and write that out too:
protocol_work = ne_move.cumulative_work
if self._storage:
self._storage.write_array("protocolwork", protocol_work, iteration=iteration)
# Return
return [final_old_sampler_state, final_sampler_state, logP_work, -initial_reduced_potential, -final_reduced_potential]
def run_equilibrium(
equilibrium_result: EquilibriumResult,
thermodynamic_state: states.ThermodynamicState,
nsteps_equil: int,
topology: md.Topology,
n_iterations: int,
atom_indices_to_save: List[int] = None,
trajectory_filename: str = None,
splitting: str = "V R O R V",
timestep: unit.Quantity = 1.0 * unit.femtoseconds
) -> EquilibriumResult:
"""
Run nsteps of equilibrium sampling at the specified thermodynamic state and return the final sampler state
as well as a trajectory of the positions after each application of an MCMove. This means that if the MCMove
is configured to run 1000 steps of dynamics, and n_iterations is 100, there will be 100 frames in the resulting
trajectory; these are the result of 100,000 steps (1000*100) of dynamics.
Parameters
----------
equilibrium_result : EquilibriumResult
EquilibriumResult namedtuple containing the information necessary to resume
thermodynamic_state : openmmtools.states.ThermodynamicState
The thermodynamic state (including context parameters) that should be used
nsteps_equil : int
The number of equilibrium steps that a move should make when apply is called
topology : mdtraj.Topology
an MDTraj topology object used to construct the trajectory
n_iterations : int
The number of times to apply the move. Note that this is not the number of steps of dynamics; it is
n_iterations*n_steps (which is set in the MCMove).
splitting: str, default "V R O H R V"
The splitting string for the dynamics
atom_indices_to_save : list of int, default None
list of indices to save (when excluding waters, for instance). If None, all indices are saved.
trajectory_filename : str, optional, default None
Full filepath of trajectory files. If none, trajectory files are not written.
splitting: str, default "V R O H R V"
The splitting string for the dynamics
Returns
-------
equilibrium_result : EquilibriumResult
Container namedtuple that has the SamplerState for resuming, an MDTraj trajectory, and the reduced potential of the
final frame.
"""
sampler_state = equilibrium_result.sampler_state
#get the atom indices we need to subset the topology and positions
if atom_indices_to_save is None:
atom_indices = list(range(topology.n_atoms))
subset_topology = topology
else:
subset_topology = topology.subset(atom_indices_to_save)
atom_indices = atom_indices_to_save
n_atoms = subset_topology.n_atoms
#construct the MCMove:
mc_move = mcmc.LangevinSplittingDynamicsMove(n_steps=nsteps_equil,
splitting=splitting)
mc_move.n_restart_attempts = 10
#create a numpy array for the trajectory
trajectory_positions = np.zeros([n_iterations, n_atoms, 3])
trajectory_box_lengths = np.zeros([n_iterations, 3])
trajectory_box_angles = np.zeros([n_iterations, 3])
#loop through iterations and apply MCMove, then collect positions into numpy array
for iteration in range(n_iterations):
mc_move.apply(thermodynamic_state, sampler_state)
trajectory_positions[iteration, :] = sampler_state.positions[
atom_indices, :].value_in_unit_system(unit.md_unit_system)
#get the box lengths and angles
a, b, c, alpha, beta, gamma = mdtrajutils.unitcell.box_vectors_to_lengths_and_angles(
*sampler_state.box_vectors)
trajectory_box_lengths[iteration, :] = [a, b, c]
trajectory_box_angles[iteration, :] = [alpha, beta, gamma]
#construct trajectory object:
trajectory = md.Trajectory(trajectory_positions,
subset_topology,
unitcell_lengths=trajectory_box_lengths,
unitcell_angles=trajectory_box_angles)
#get the reduced potential from the final frame for endpoint perturbations
reduced_potential_final_frame = thermodynamic_state.reduced_potential(
sampler_state)
#construct equilibrium result object
equilibrium_result = EquilibriumResult(sampler_state,
reduced_potential_final_frame)
#If there is a trajectory filename passed, write out the results here:
if trajectory_filename is not None:
write_equilibrium_trajectory(equilibrium_result, trajectory,
trajectory_filename)
return equilibrium_result
