def minimize(thermodynamic_state: states.ThermodynamicState,
sampler_state: states.SamplerState,
max_iterations: int = 100) -> states.SamplerState:
"""
Minimize the given system and state, up to a maximum number of steps.
This does not return a copy of the samplerstate; it is an update-in-place.
Parameters
----------
thermodynamic_state : openmmtools.states.ThermodynamicState
The state at which the system could be minimized
sampler_state : openmmtools.states.SamplerState
The starting state at which to minimize the system.
max_iterations : int, optional, default 100
The maximum number of minimization steps. Default is 100.
Returns
-------
sampler_state : openmmtools.states.SamplerState
The posititions and accompanying state following minimization
"""
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)
_logger.debug(f"using dummy context cache")
else:
_logger.debug(f"using global context cache")
context, integrator = cache.global_context_cache.get_context(
thermodynamic_state)
sampler_state.apply_to_context(context, ignore_velocities=True)
openmm.LocalEnergyMinimizer.minimize(context, maxIterations=max_iterations)
sampler_state.update_from_context(context)
def _update_particle_state_substate(self,
particle_state,
new_state_subset=False):
"""
update the particle state from the context, create a particle substate and update from context
"""
#update the particle state and the particle state subset
particle_state.update_from_context(
self.context, ignore_velocities=True
) #update the particle state from the context
if new_state_subset:
self.particle_state_subset = SamplerState(
positions=particle_state.positions[list(
self._subset_indices_map.keys(
))]) #create a particle state from the subset context
else:
self.particle_state_subset.positions = particle_state.positions[
list(self._subset_indices_map.keys()
)] #update the particle subset positions appropriately
self.particle_state_subset.apply_to_context(
self.context_subset, ignore_velocities=True
) #apply the subset particle state to its context
self.particle_state_subset.update_from_context(
self.context_subset, ignore_velocities=True
) #update the subset particle state from its context to updated the potential energy
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 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 __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 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 _geometry_forward(self, topology_proposal, old_sampler_state):
"""
Run geometry engine to propose new positions and compute logP
Parameters
----------
topology_proposal : TopologyProposal
Contains old/new Topology and System objects and atom mappings.
old_sampler_state : openmmtools.states.SamplerState
Configurational properties of the old system atoms.
Returns
-------
new_sampler_state : openmmtools.states.SamplerState
Configurational properties of new atoms proposed by geometry engine calculation.
geometry_logp_propose : float
The log probability of the forward-only proposal
"""
if self.verbose: print("Geometry engine proposal...")
# Generate coordinates for new atoms and compute probability ratio of old and new probabilities.
initial_time = time.time()
new_positions, geometry_logp_propose = self.geometry_engine.propose(topology_proposal, old_sampler_state.positions, self.sampler.thermodynamic_state.beta)
if self.verbose: print('proposal took %.3f s' % (time.time() - initial_time))
if self.geometry_pdbfile is not None:
print("Writing proposed geometry...")
from simtk.openmm.app import PDBFile
PDBFile.writeFile(topology_proposal.new_topology, new_positions, file=self.geometry_pdbfile)
self.geometry_pdbfile.flush()
new_sampler_state = SamplerState(new_positions, box_vectors=old_sampler_state.box_vectors)
return new_sampler_state, geometry_logp_propose
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_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 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 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 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_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 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_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 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 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_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 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 create( # pylint: disable=too-many-arguments
self,
reference_state,
thermodynamic_states,
sampler_states,
top,
target_state=0,
initial_thermodynamic_states=None,
unsampled_thermodynamic_states=None,
storage=None,
metadata=None,
stride=1):
"""Create new multistate sampler simulation.
Override MultistateSampler create.
Parameters
----------
thermodynamic_states : list of states.ThermodynamicState
Thermodynamic states to simulate, where one replica is allocated
per state.
Each state must have a system with the same number of atoms.
sampler_states : states.SamplerState or list
One or more sets of initial sampler states.
The number of replicas is taken to be the number of sampler states
provided. If the sampler states do not have box_vectors attached
and the system is periodic, an exception will be thrown.
top : Topology or Topography object, optional
target_state : int, optional
The indef of the reference thermodynamic state. Defaults to 0.
initial_thermodynamic_states : None or list or
array-like of int of length len(sampler_states), optional,
Default: None.
Initial thermodynamic_state index for each sampler_state.
If no initial distribution is chosen, ``sampler_states`` are
distributed between the ``thermodynamic_states`` following these
rules:
* If ``len(thermodynamic_states) == len(sampler_states)``:
1-to-1 distribution
* If ``len(thermodynamic_states) > len(sampler_states)``:
First and last state distributed first
Remaining ``sampler_states`` spaced evenly by index until
``sampler_states`` are depleted.
If there is only one ``sampler_state``, then the only first
``thermodynamic_state`` will be chosen
* If ``len(thermodynamic_states) < len(sampler_states)``:
each ``thermodynamic_state`` receives an equal number of
``sampler_states`` until there are insufficient number of
``sampler_states`` remaining to give each
``thermodynamic_state`` an equal number. Then the rules from
the previous point are followed.
unsampled_thermodynamic_states : list of states.ThermodynamicState,
optional, default=None
These are ThermodynamicStates that are not propagated, but their
reduced potential is computed at each iteration for each replica.
These energy can be used as data for reweighting schemes (default
is None).
storage : str or instanced Reporter
If str: the path to the storage file.
Default checkpoint options from Reporter class are used.
If Reporter: Uses the reporter options and storage path
In the future this will be able to take a Storage class as well.
metadata : dict, optional, default=None
Simulation metadata to be stored in the file.
"""
if isinstance(thermodynamic_states, ThermodynamicState):
# a single state
thermodynamic_states = [thermodynamic_states]
if not isinstance(sampler_states,
(SamplerState, Sequence)): # TODO: check this
# as positions
sampler_states = SamplerState(sampler_states)
# Do not modify passed ref thermodynamic state.
self._reference_thermodynamic_state = copy.deepcopy(reference_state)
thermodynamic_state = copy.deepcopy(
self._reference_thermodynamic_state)
self._reference_system = thermodynamic_state.system
if storage is None:
storage = tempfile.NamedTemporaryFile(delete=False).name + '.nc'
self.reporter = self._define_reporter(storage, stride)
# set self.topography
if isinstance(top, Topography):
self.topography = top
else:
self.topography = Topography(top)
self.ref_state = target_state
if metadata is None:
metadata = {}
sampler_full_name = mmtools.utils.typename(self.__class__)
metadata['title'] = 'Created using %s on %s' % (
sampler_full_name, time.asctime(time.localtime()))
metadata['sampler_full_name'] = sampler_full_name
metadata['topography'] = mmtools.utils.serialize(self.topography)
metadata['reference_state'] = mmtools.utils.serialize(
thermodynamic_state) # the ref thermodynamic state
super().create(
thermodynamic_states=thermodynamic_states,
sampler_states=sampler_states,
storage=self.reporter,
initial_thermodynamic_states=initial_thermodynamic_states,
unsampled_thermodynamic_states=unsampled_thermodynamic_states,
metadata=metadata)
def compare_energies(REST_system, other_system, positions, rest_atoms, T_min, T):
# Create thermodynamic state
lambda_zero_alchemical_state = RESTState.from_system(REST_system)
thermostate = ThermodynamicState(REST_system, temperature=T_min)
compound_thermodynamic_state = CompoundThermodynamicState(thermostate,
composable_states=[lambda_zero_alchemical_state])
# Set alchemical parameters
beta_0 = 1 / (kB * T_min)
beta_m = 1 / (kB * T)
compound_thermodynamic_state.set_alchemical_parameters(beta_0, beta_m)
# Minimize and save energy
integrator = openmm.VerletIntegrator(1.0 * unit.femtosecond)
context = compound_thermodynamic_state.create_context(integrator)
context.setPositions(positions)
sampler_state = SamplerState.from_context(context)
REST_energy = compound_thermodynamic_state.reduced_potential(sampler_state)
# Compute energy for non-RESTified system
# Determine regions and scaling factors
solute = rest_atoms
solvent = [i for i in range(other_system.getNumParticles()) if i not in solute]
solute_scaling = beta_m / beta_0
inter_scaling = np.sqrt(beta_m / beta_0)
# Scale the terms in the bond force appropriately
bond_force = other_system.getForce(0)
for bond in range(bond_force.getNumBonds()):
p1, p2, length, k = bond_force.getBondParameters(bond)
if p1 in solute and p2 in solute:
bond_force.setBondParameters(bond, p1, p2, length, k * solute_scaling)
elif (p1 in solute and p2 in solvent) or (p1 in solvent and p2 in solute):
bond_force.setBondParameters(bond, p1, p2, length, k * inter_scaling)
# Scale the terms in the angle force appropriately
angle_force = other_system.getForce(1)
for angle_index in range(angle_force.getNumAngles()):
p1, p2, p3, angle, k = angle_force.getAngleParameters(angle_index)
if p1 in solute and p2 in solute and p3 in solute:
angle_force.setAngleParameters(angle_index, p1, p2, p3, angle, k * solute_scaling)
elif set([p1, p2, p3]).intersection(set(solute)) != set() and set([p1, p2, p3]).intersection(
set(solvent)) != set():
angle_force.setAngleParameters(angle_index, p1, p2, p3, angle, k * inter_scaling)
# Scale the terms in the torsion force appropriately
torsion_force = other_system.getForce(2)
for torsion_index in range(torsion_force.getNumTorsions()):
p1, p2, p3, p4, periodicity, phase, k = torsion_force.getTorsionParameters(torsion_index)
if p1 in solute and p2 in solute and p3 in solute and p4 in solute:
torsion_force.setTorsionParameters(torsion_index, p1, p2, p3, p4, periodicity, phase, k * solute_scaling)
elif set([p1, p2, p3, p4]).intersection(set(solute)) != set() and set([p1, p2, p3, p4]).intersection(
set(solvent)) != set():
torsion_force.setTorsionParameters(torsion_index, p1, p2, p3, p4, periodicity, phase, k * inter_scaling)
# Scale the exceptions in the nonbonded force appropriately
nb_force = other_system.getForce(3)
for nb_index in range(nb_force.getNumExceptions()):
p1, p2, chargeProd, sigma, epsilon = nb_force.getExceptionParameters(nb_index)
if p1 in solute and p2 in solute:
nb_force.setExceptionParameters(nb_index, p1, p2, solute_scaling * chargeProd, sigma, solute_scaling * epsilon)
elif (p1 in solute and p2 in solvent) or (p1 in solvent and p2 in solute):
nb_force.setExceptionParameters(nb_index, p1, p2, inter_scaling * chargeProd, sigma, inter_scaling * epsilon)
# Scale nonbonded interactions for solute-solute region by adding exceptions for all pairs of atoms
exception_pairs = [tuple(sorted([nb_force.getExceptionParameters(nb_index)[0], nb_force.getExceptionParameters(nb_index)[1]])) for nb_index in range(nb_force.getNumExceptions())]
solute_pairs = set([tuple(sorted(pair)) for pair in list(itertools.product(solute, solute))])
for pair in list(solute_pairs):
p1 = pair[0]
p2 = pair[1]
p1_charge, p1_sigma, p1_epsilon = nb_force.getParticleParameters(p1)
p2_charge, p2_sigma, p2_epsilon = nb_force.getParticleParameters(p2)
if p1 != p2:
if pair not in exception_pairs:
nb_force.addException(p1, p2, p1_charge * p2_charge * solute_scaling, 0.5 * (p1_sigma + p2_sigma),
np.sqrt(p1_epsilon * p2_epsilon) * solute_scaling)
# Scale nonbonded interactions for inter region by adding exceptions for all pairs of atoms
for pair in list(itertools.product(solute, solvent)):
p1 = pair[0]
p2 = int(pair[1]) # otherwise, will be a numpy int
p1_charge, p1_sigma, p1_epsilon = nb_force.getParticleParameters(p1)
p2_charge, p2_sigma, p2_epsilon = nb_force.getParticleParameters(p2)
if tuple(sorted(pair)) not in exception_pairs:
nb_force.addException(p1, p2, p1_charge * p2_charge * inter_scaling, 0.5 * (p1_sigma + p2_sigma), np.sqrt(p1_epsilon * p2_epsilon) * inter_scaling)
# Get energy
thermostate = ThermodynamicState(other_system, temperature=T_min)
integrator = openmm.VerletIntegrator(1.0 * unit.femtosecond)
context = thermostate.create_context(integrator)
context.setPositions(positions)
sampler_state = SamplerState.from_context(context)
nonREST_energy = thermostate.reduced_potential(sampler_state)
assert REST_energy - nonREST_energy < 1, f"The energy of the REST system ({REST_energy}) does not match " \
f"that of the non-REST system with terms manually scaled according to REST2({nonREST_energy})."
def HybridTopologyFactory_energies(
current_mol='toluene',
proposed_mol='1,2-bis(trifluoromethyl) benzene'):
"""
Test whether the difference in the nonalchemical zero and alchemical zero states is the forward valence energy. Also test for the one states.
"""
from perses.tests.utils import generate_solvated_hybrid_test_topology, generate_endpoint_thermodynamic_states
import openmmtools.cache as cache
#Just test the solvated system
top_proposal, old_positions, _ = generate_solvated_hybrid_test_topology(
current_mol_name=current_mol, proposed_mol_name=proposed_mol)
#remove the dispersion correction
top_proposal._old_system.getForce(3).setUseDispersionCorrection(False)
top_proposal._new_system.getForce(3).setUseDispersionCorrection(False)
# run geometry engine to generate old and new positions
_geometry_engine = FFAllAngleGeometryEngine(metadata=None,
use_sterics=False,
n_bond_divisions=100,
n_angle_divisions=180,
n_torsion_divisions=360,
verbose=True,
storage=None,
bond_softening_constant=1.0,
angle_softening_constant=1.0,
neglect_angles=False)
_new_positions, _lp = _geometry_engine.propose(top_proposal, old_positions,
beta)
_lp_rev = _geometry_engine.logp_reverse(top_proposal, _new_positions,
old_positions, beta)
# make the hybrid system, reset the CustomNonbondedForce cutoff
HTF = HybridTopologyFactory(top_proposal, old_positions, _new_positions)
hybrid_system = HTF.hybrid_system
nonalch_zero, nonalch_one, alch_zero, alch_one = generate_endpoint_thermodynamic_states(
hybrid_system, top_proposal)
# compute reduced energies
#for the nonalchemical systems...
attrib_list = [(nonalch_zero, old_positions,
top_proposal._old_system.getDefaultPeriodicBoxVectors()),
(alch_zero, HTF._hybrid_positions,
hybrid_system.getDefaultPeriodicBoxVectors()),
(alch_one, HTF._hybrid_positions,
hybrid_system.getDefaultPeriodicBoxVectors()),
(nonalch_one, _new_positions,
top_proposal._new_system.getDefaultPeriodicBoxVectors())]
rp_list = []
for (state, pos, box_vectors) in attrib_list:
context, integrator = cache.global_context_cache.get_context(state)
samplerstate = SamplerState(positions=pos, box_vectors=box_vectors)
samplerstate.apply_to_context(context)
rp = state.reduced_potential(context)
rp_list.append(rp)
#valence energy definitions
forward_added_valence_energy = _geometry_engine.forward_final_context_reduced_potential - _geometry_engine.forward_atoms_with_positions_reduced_potential
reverse_subtracted_valence_energy = _geometry_engine.reverse_final_context_reduced_potential - _geometry_engine.reverse_atoms_with_positions_reduced_potential
nonalch_zero_rp, alch_zero_rp, alch_one_rp, nonalch_one_rp = rp_list[
0], rp_list[1], rp_list[2], rp_list[3]
# print(f"Difference between zeros: {nonalch_zero_rp - alch_zero_rp}; forward added: {forward_added_valence_energy}")
# print(f"Difference between ones: {nonalch_zero_rp - alch_zero_rp}; forward added: {forward_added_valence_energy}")
assert abs(
nonalch_zero_rp - alch_zero_rp + forward_added_valence_energy
) < ENERGY_THRESHOLD, f"The zero state alchemical and nonalchemical energy absolute difference {abs(nonalch_zero_rp - alch_zero_rp + forward_added_valence_energy)} is greater than the threshold of {ENERGY_THRESHOLD}."
assert abs(
nonalch_one_rp - alch_one_rp + reverse_subtracted_valence_energy
) < ENERGY_THRESHOLD, f"The one state alchemical and nonalchemical energy absolute difference {abs(nonalch_one_rp - alch_one_rp + reverse_subtracted_valence_energy)} is greater than the threshold of {ENERGY_THRESHOLD}."
print(
f"Abs difference in zero alchemical vs nonalchemical systems: {abs(nonalch_zero_rp - alch_zero_rp + forward_added_valence_energy)}"
)
print(
f"Abs difference in one alchemical vs nonalchemical systems: {abs(nonalch_one_rp - alch_one_rp + reverse_subtracted_valence_energy)}"
)
def run_hybrid_endpoint_overlap(topology_proposal, current_positions,
new_positions):
"""
Test that the variance of the perturbation from lambda={0,1} to the corresponding nonalchemical endpoint is not
too large.
Parameters
----------
topology_proposal : perses.rjmc.TopologyProposal
TopologyProposal object describing the transformation
current_positions : np.array, unit-bearing
Positions of the initial system
new_positions : np.array, unit-bearing
Positions of the new system
Returns
-------
hybrid_endpoint_results : list
list of [df, ddf, N_eff] for 1 and 0
"""
#create the hybrid system:
#hybrid_factory = HybridTopologyFactory(topology_proposal, current_positions, new_positions, use_dispersion_correction=True)
hybrid_factory = HybridTopologyFactory(
topology_proposal,
current_positions,
new_positions,
use_dispersion_correction=False) # DEBUG
#get the relevant thermodynamic states:
nonalchemical_zero_thermodynamic_state, nonalchemical_one_thermodynamic_state, lambda_zero_thermodynamic_state, lambda_one_thermodynamic_state = utils.generate_endpoint_thermodynamic_states(
hybrid_factory.hybrid_system, topology_proposal)
nonalchemical_thermodynamic_states = [
nonalchemical_zero_thermodynamic_state,
nonalchemical_one_thermodynamic_state
]
alchemical_thermodynamic_states = [
lambda_zero_thermodynamic_state, lambda_one_thermodynamic_state
]
#create an MCMCMove, BAOAB with default parameters (but don't restart if we encounter a NaN)
mc_move = mcmc.LangevinDynamicsMove(n_restart_attempts=0, n_steps=100)
initial_sampler_state = SamplerState(
hybrid_factory.hybrid_positions,
box_vectors=hybrid_factory.hybrid_system.getDefaultPeriodicBoxVectors(
))
hybrid_endpoint_results = []
all_results = []
for lambda_state in (0, 1):
result, non, hybrid = run_endpoint_perturbation(
alchemical_thermodynamic_states[lambda_state],
nonalchemical_thermodynamic_states[lambda_state],
initial_sampler_state,
mc_move,
100,
hybrid_factory,
lambda_index=lambda_state)
all_results.append(non)
all_results.append(hybrid)
print('lambda {} : {}'.format(lambda_state, result))
hybrid_endpoint_results.append(result)
calculate_cross_variance(all_results)
return hybrid_endpoint_results
def run_endpoint_perturbation(lambda_thermodynamic_state,
nonalchemical_thermodynamic_state,
initial_hybrid_sampler_state,
mc_move,
n_iterations,
factory,
lambda_index=0,
print_work=False,
write_system=False,
write_state=False,
write_trajectories=False):
"""
Parameters
----------
lambda_thermodynamic_state : ThermodynamicState
The thermodynamic state corresponding to the hybrid system at a lambda endpoint
nonalchemical_thermodynamic_state : ThermodynamicState
The nonalchemical thermodynamic state for the relevant endpoint
initial_hybrid_sampler_state : SamplerState
Starting positions for the sampler. Must be compatible with lambda_thermodynamic_state
mc_move : MCMCMove
The MCMove that will be used for sampling at the lambda endpoint
n_iterations : int
The number of iterations
factory : HybridTopologyFactory
The hybrid topology factory
lambda_index : int, optional, default=0
The index, 0 or 1, at which to retrieve nonalchemical positions
print_work : bool, optional, default=False
If True, will print work values
write_system : bool, optional, default=False
If True, will write alchemical and nonalchemical System XML files
write_state : bool, optional, default=False
If True, write alchemical (hybrid) State XML files each iteration
write_trajectories : bool, optional, default=False
If True, will write trajectories
Returns
-------
df : float
Free energy difference between alchemical and nonalchemical systems, estimated with EXP
ddf : float
Standard deviation of estimate, corrected for correlation, from EXP estimator.
"""
import mdtraj as md
#run an initial minimization:
mcmc_sampler = mcmc.MCMCSampler(lambda_thermodynamic_state,
initial_hybrid_sampler_state, mc_move)
mcmc_sampler.minimize(max_iterations=20)
new_sampler_state = mcmc_sampler.sampler_state
if write_system:
with open(f'hybrid{lambda_index}-system.xml', 'w') as outfile:
outfile.write(
openmm.XmlSerializer.serialize(
lambda_thermodynamic_state.system))
with open(f'nonalchemical{lambda_index}-system.xml', 'w') as outfile:
outfile.write(
openmm.XmlSerializer.serialize(
nonalchemical_thermodynamic_state.system))
#initialize work array
w = np.zeros([n_iterations])
non_potential = np.zeros([n_iterations])
hybrid_potential = np.zeros([n_iterations])
#run n_iterations of the endpoint perturbation:
hybrid_trajectory = unit.Quantity(
np.zeros([
n_iterations,
lambda_thermodynamic_state.system.getNumParticles(), 3
]), unit.nanometers) # DEBUG
nonalchemical_trajectory = unit.Quantity(
np.zeros([
n_iterations,
nonalchemical_thermodynamic_state.system.getNumParticles(), 3
]), unit.nanometers) # DEBUG
for iteration in range(n_iterations):
# Generate a new sampler state for the hybrid system
mc_move.apply(lambda_thermodynamic_state, new_sampler_state)
# Compute the hybrid reduced potential at the new sampler state
hybrid_context, integrator = cache.global_context_cache.get_context(
lambda_thermodynamic_state)
new_sampler_state.apply_to_context(hybrid_context,
ignore_velocities=True)
hybrid_reduced_potential = lambda_thermodynamic_state.reduced_potential(
hybrid_context)
if write_state:
state = hybrid_context.getState(getPositions=True,
getParameters=True)
state_xml = openmm.XmlSerializer.serialize(state)
with open(f'state{iteration}_l{lambda_index}.xml', 'w') as outfile:
outfile.write(state_xml)
# Construct a sampler state for the nonalchemical system
if lambda_index == 0:
nonalchemical_positions = factory.old_positions(
new_sampler_state.positions)
elif lambda_index == 1:
nonalchemical_positions = factory.new_positions(
new_sampler_state.positions)
else:
raise ValueError(
"The lambda index needs to be either one or zero for this to be meaningful"
)
nonalchemical_sampler_state = SamplerState(
nonalchemical_positions, box_vectors=new_sampler_state.box_vectors)
if write_trajectories:
state = hybrid_context.getState(getPositions=True)
hybrid_trajectory[iteration, :, :] = state.getPositions(
asNumpy=True)
nonalchemical_trajectory[iteration, :, :] = nonalchemical_positions
# Compute the nonalchemical reduced potential
nonalchemical_context, integrator = cache.global_context_cache.get_context(
nonalchemical_thermodynamic_state)
nonalchemical_sampler_state.apply_to_context(nonalchemical_context,
ignore_velocities=True)
nonalchemical_reduced_potential = nonalchemical_thermodynamic_state.reduced_potential(
nonalchemical_context)
# Compute and store the work
w[iteration] = nonalchemical_reduced_potential - hybrid_reduced_potential
non_potential[iteration] = nonalchemical_reduced_potential
hybrid_potential[iteration] = hybrid_reduced_potential
if print_work:
print(
f'{iteration:8d} {hybrid_reduced_potential:8.3f} {nonalchemical_reduced_potential:8.3f} => {w[iteration]:8.3f}'
)
if write_trajectories:
if lambda_index == 0:
nonalchemical_mdtraj_topology = md.Topology.from_openmm(
factory._topology_proposal.old_topology)
elif lambda_index == 1:
nonalchemical_mdtraj_topology = md.Topology.from_openmm(
factory._topology_proposal.new_topology)
md.Trajectory(
hybrid_trajectory / unit.nanometers,
factory.hybrid_topology).save(f'hybrid{lambda_index}.pdb')
md.Trajectory(nonalchemical_trajectory / unit.nanometers,
nonalchemical_mdtraj_topology).save(
f'nonalchemical{lambda_index}.pdb')
# Analyze data and return results
[t0, g, Neff_max] = timeseries.detectEquilibration(w)
w_burned_in = w[t0:]
[df, ddf] = pymbar.EXP(w_burned_in)
ddf_corrected = ddf * np.sqrt(g)
results = [df, ddf_corrected, t0, Neff_max]
return results, non_potential, hybrid_potential
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,
class Propagator(OMMBIP):
"""
Propagator pseudocode:
Step 1: initialization--
set iteration = 0, n_iterations = n_iterations, lambda = 0 (i.e. iteration / n_iterations); work_accumulated = 0.0
generate sample x_0 ~ e^(-p(x))
evaluate work_incremental = 0 (i.e. u_mm(x_0) - g(x_0), but we presume that g = u_mm(.))
work_accumulated <- work_accumulated + work_incremental
x' = x_0
Step 2: sampling
for increment in range(n_iterations):
x = x'
ante_perturbation_potential = (1 - lambda) * u_mm(x) + lambda * u_ani_mm_mix(x)
set iteration <- iteration + 1.0; lambda <- iteration / n_iterations
evaluate work_incremental = [(1 - lambda) * u_mm(x) + lambda * u_ani_mm_mix(x)] - ante_perturbation_potential
work_accumulated <- work_accumulated + work_incremental
create a modified force: modified_f = (1 - lambda) * f_mm + lambda * f_ani_mm_mix (where f_. = -grad(u_.) )
x' = V R O R (where V deterministic update is according to modified_f defined above) w.r.t x
NOTE: in this regime, the last x' is propagated w.r.t. a propagator whose invariant distribution respects u_ani_mm_mix;
this should _not_ be the case. There should be an exception in the Step 2 for loop that breaks once the final work_incremental is computed and updated
to the work_accumulated. Regardless, the distribution of accumulated works is unaffected by this 'bug'; only expectations (as a function of x) w.r.t. these
weights may be affected.
See: 3.1.1. of https://www.stats.ox.ac.uk/~doucet/delmoral_doucet_jasra_sequentialmontecarlosamplersJRSSB.pdf (esp. Remark 1.)
"""
def __init__(self,
openmm_pdf_state,
openmm_pdf_state_subset,
subset_indices_map,
integrator,
ani_handler,
context_cache=None,
reassign_velocities=True,
n_restart_attempts=0,
reporter=None,
write_trajectory_interval = 1,
**kwargs):
"""
arguments
openmm_pdf_state : openmmtools.states.ThermodynamicState
the pdf state of the propagator
openmm_pdf_state_subset : openmmtools.states.ThermodynamicState
the pdf state of the atom subset
subset_indices_map : dict
dict of {openmm_pdf_state atom_index : openmm_pdf_state_subset atom index}
integrator : openmm.Integrator
integrator of dynamics
ani_handler : ANI1_force_and_energy
handler for ani forces and potential energy
context_cache : openmmtools.cache.ContextCache, optional default:None
The ContextCache to use for Context creation. If None, the global cache
openmmtools.cache.global_context_cache is used.
reassign_velocities : bool, optional default:False
If True, the velocities will be reassigned from the Maxwell-Boltzmann
distribution at the beginning of the move.
n_restart_attempts : int, optional default:0
When greater than 0, if after the integration there are NaNs in energies,
the move will restart. When the integrator has a random component, this
may help recovering. On the last attempt, the ``Context`` is
re-initialized in a slower process, but better than the simulation
crashing. An IntegratorMoveError is raised after the given number of
attempts if there are still NaNs.
reporter : coddiwomple.openmm.reporter.OpenMMReporter, default None
a reporter object to write trajectories
write_trajectory_interval : int
frequency of writing trajectory
"""
super().__init__(openmm_pdf_state,
integrator,
context_cache,
reassign_velocities,
n_restart_attempts)
#create a pdf state for the subset indices (usually a vacuum system)
self.pdf_state_subset = openmm_pdf_state_subset
assert self.pdf_state_subset.temperature == self.pdf_state.temperature, f"the temperatures of the pdf states do not match"
#create a dictionary for subset indices
self._subset_indices_map = subset_indices_map
#create an ani handler attribute that can be referenced
self.ani_handler = ani_handler
#create a context for the subset atoms that can be referenced
self.context_subset, _ = cache.global_context_cache.get_context(self.pdf_state_subset)
#create a reporter for the accumulated works
self._state_works = {}
self._state_works_counter = 0
#create a reporter
self._write_trajectory = False if reporter is None else True
self.reporter=reporter
if self._write_trajectory:
from coddiwomple.particles import Particle
self.particle = Particle(0)
self.write_trajectory_interval=write_trajectory_interval
else:
self.particle = None
self.write_trajectory_interval=None
def _initialize_state_works(self):
"""
initialize an empty list and add 0.0 to it (state works)
"""
self._current_state_works = [] #define an interim (auxiliary) list that will track the thermodynamic work of the current application
self._current_state_works.append(0.0) #the first incremental work is always 0 since the importance function is identical to the first target distribution (i.e. fully interacting MM)
def _initialize_iterations(self, n_iterations):
"""
initialize the iteration counter
"""
self._iteration = 0.0 #define the first iteration as 0
self._n_iterations = n_iterations #the number of iterations in the protocol is equal to the number of steps in the application
def _update_particle_state_substate(self, particle_state, new_state_subset=False):
"""
update the particle state from the context, create a particle substate and update from context
"""
#update the particle state and the particle state subset
particle_state.update_from_context(self.context, ignore_velocities=True) #update the particle state from the context
if new_state_subset:
self.particle_state_subset = SamplerState(positions = particle_state.positions[list(self._subset_indices_map.keys())]) #create a particle state from the subset context
else:
self.particle_state_subset.positions = particle_state.positions[list(self._subset_indices_map.keys())] #update the particle subset positions appropriately
self.particle_state_subset.apply_to_context(self.context_subset, ignore_velocities=True) #apply the subset particle state to its context
self.particle_state_subset.update_from_context(self.context_subset, ignore_velocities=True) #update the subset particle state from its context to updated the potential energy
def _update_current_state_works(self, particle_state):
"""
update the current state and associated works
"""
#get the reduced potential
reduced_potential = self._compute_hybrid_potential(_lambda = self._iteration / self._n_iterations, particle_state = particle_state)
perturbed_reduced_potential = self._compute_hybrid_potential(_lambda = (self._iteration + 1.0) / self._n_iterations, particle_state = particle_state)
self._current_state_works.append(self._current_state_works[-1] + (perturbed_reduced_potential - reduced_potential))
def _update_force(self, particle_state):
"""
update the force
"""
mm_force_matrix = self._compute_hybrid_forces(_lambda = (self._iteration + 1.0) / self._n_iterations, particle_state = particle_state).value_in_unit_system(unit.md_unit_system)
self.integrator.setPerDofVariableByName('modified_force', mm_force_matrix)
def _before_integration(self, *args, **kwargs):
particle_state = args[0] #define the particle state
n_iterations = args[1] #define the number of iterations
self._initialize_state_works()
self._initialize_iterations(n_iterations)
#update the particle state and the particle state subset
self._update_particle_state_substate(particle_state, new_state_subset=True)
self._update_current_state_works(particle_state)
self._update_force(particle_state)
#report
if self._write_trajectory: # the first state is always saved for processing purposes
self.particle.update_state(particle_state)
self.reporter.record([self.particle])
def _during_integration(self, *args, **kwargs):
particle_state = args[0]
self._iteration += 1.0
self._update_particle_state_substate(particle_state)
#get the reduced potential
if self._iteration < self._n_iterations:
self._update_current_state_works(particle_state)
self._update_force(particle_state)
else:
#we are done
pass
if self._write_trajectory and int(self._iteration) % self.write_trajectory_interval == 0:
self.particle.update_state(particle_state)
if self._iteration == self._n_iterations:
self.reporter.record([self.particle], save_to_disk=True)
else:
self.reporter.record([self.particle], save_to_disk=False)
def _after_integration(self, *args, **kwargs):
self._state_works[self._state_works_counter] = deepcopy(self._current_state_works)
self._state_works_counter += 1
if self._write_trajectory:
self.reporter.reset()
#self._log_context_parameters()
def _compute_hybrid_potential(self,_lambda, particle_state):
"""
function to compute the hybrid reduced potential defined as follows:
U(x_rec, x_lig) = u_mm,rec(x_rec) - lambda*u_mm,lig(x_lig) + lambda*u_ani,lig(x_lig)
"""
reduced_potential = (self.pdf_state.reduced_potential(particle_state)
- _lambda * self.pdf_state_subset.reduced_potential(self.particle_state_subset)
+ _lambda * self.ani_handler.calculate_energy(self.particle_state_subset.positions) * self.pdf_state.beta)
return reduced_potential
def _compute_hybrid_forces(self, _lambda, particle_state):
"""
function to compute a hybrid force matrix of shape num_particles x 3
in the spirit of the _compute_hybrid_potential, we compute the forces in the following way
F(x_rec, x_lig) = F_mm(x_rec, x_lig) - lambda * F_mm(x_lig) + lambda * F_ani(x_lig)
"""
# get the complex mm forces
state = self.context.getState(getForces=True)
mm_force_matrix = state.getForces(asNumpy=True) # returns forces in kJ/(nm mol)
# get the ligand mm forces
subset_state = self.context_subset.getState(getForces=True)
mm_force_matrix_subset = subset_state.getForces(asNumpy=True)
# get the ligand ani forces
coords = self.particle_state_subset.positions
subset_ani_force_matrix, energie = self.ani_handler.calculate_force(coords) # returns force in kJ/(A mol)
#print(f"ani force matrix head: ",subset_ani_force_matrix[0])
# now combine the ligand forces
subset_force_matrix = _lambda * (subset_ani_force_matrix - mm_force_matrix_subset) #we are adding two Quantities with different units, but they are compatible
#print(f"mm subset force matrix head", mm_force_matrix_subset[0])
# and append to the complex forces...
#print(f"mm force matrix head", mm_force_matrix[0])
mm_force_matrix[list(self._subset_indices_map.keys()), :] += subset_force_matrix #and same, here...
#print(f"mm force matrix head (after ani modification)", mm_force_matrix[0])
return mm_force_matrix
def _get_context_subset_parameters(self):
"""
return a dictionary of the self.context_subset's parameters
returns
context_parameters : dict
{parameter name <str> : parameter value value <float>}
"""
swig_parameters = self.context_subset.getParameters()
context_parameters = {q: swig_parameters[q] for q in swig_parameters}
return context_parameters
def _log_context_parameters(self):
"""
log the context and context subset parameters
"""
context_parameters = self._get_context_parameters()
context_subset_parameters = self._get_context_subset_parameters()
_logger.debug(f"\tcontext_parameters during integration:")
for key, val in context_parameters.items():
_logger.debug(f"\t\t{key}: {val}")
_logger.debug(f"\tcontext subset parameters during integration:")
for key, val in context_subset_parameters:
_logger.debug(f"\t\t{key}: {val}")
@property
def state_works(self):
return self._state_works
# alanine = testsystems.AlanineDipeptideVacuum()
# system = alanine.system
# positions = alanine.positions
# b) create & load our own
pdb = PDBFile(os.path.dirname(os.path.realpath(__file__)) +
'/input/alanine-dipeptide-nowater.pdb')
forcefield = ForceField('amber10.xml')
system = forcefield.createSystem(pdb.topology, nonbondedMethod=PME,
nonbondedCutoff=1*unit.nanometer, constraints=HBonds)
thermodynamic_state = ThermodynamicState(
system=system, temperature=298.15*unit.kelvin)
positions = pdb.getPositions()
# finally, start with sampling
sampler_state = SamplerState(positions=positions)
# TODO: decide on nr. of timesteps to go before recalculation
V_calculator = TorsionNeglectingPotentialEnergyCalculator(timesteps=5)
move = DeterministicRotateDisplaceMove(
displacement_sigma=5.0*unit.nanometer,
potential_energy_calculator=V_calculator,
atom_subset=5,
rng_seed=42
)
sampler = MCMCSampler(thermodynamic_state, sampler_state, move=move)
# run the sampler
sampler.minimize()
sampler.run(n_iterations=400) # TODO: increase iterations for final run
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)
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 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_endpoint_perturbation(lambda_thermodynamic_state, nonalchemical_thermodynamic_state, initial_hybrid_sampler_state, mc_move, n_iterations, factory,
lambda_index=0, print_work=False, write_system=False, write_state=False, write_trajectories=False):
"""
Parameters
----------
lambda_thermodynamic_state : ThermodynamicState
The thermodynamic state corresponding to the hybrid system at a lambda endpoint
nonalchemical_thermodynamic_state : ThermodynamicState
The nonalchemical thermodynamic state for the relevant endpoint
initial_hybrid_sampler_state : SamplerState
Starting positions for the sampler. Must be compatible with lambda_thermodynamic_state
mc_move : MCMCMove
The MCMove that will be used for sampling at the lambda endpoint
n_iterations : int
The number of iterations
factory : HybridTopologyFactory
The hybrid topology factory
lambda_index : int, optional, default=0
The index, 0 or 1, at which to retrieve nonalchemical positions
print_work : bool, optional, default=False
If True, will print work values
write_system : bool, optional, default=False
If True, will write alchemical and nonalchemical System XML files
write_state : bool, optional, default=False
If True, write alchemical (hybrid) State XML files each iteration
write_trajectories : bool, optional, default=False
If True, will write trajectories
Returns
-------
df : float
Free energy difference between alchemical and nonalchemical systems, estimated with EXP
ddf : float
Standard deviation of estimate, corrected for correlation, from EXP estimator.
"""
import mdtraj as md
#run an initial minimization:
mcmc_sampler = mcmc.MCMCSampler(lambda_thermodynamic_state, initial_hybrid_sampler_state, mc_move)
mcmc_sampler.minimize(max_iterations=20)
new_sampler_state = mcmc_sampler.sampler_state
if write_system:
with open(f'hybrid{lambda_index}-system.xml', 'w') as outfile:
outfile.write(openmm.XmlSerializer.serialize(lambda_thermodynamic_state.system))
with open(f'nonalchemical{lambda_index}-system.xml', 'w') as outfile:
outfile.write(openmm.XmlSerializer.serialize(nonalchemical_thermodynamic_state.system))
#initialize work array
w = np.zeros([n_iterations])
non_potential = np.zeros([n_iterations])
hybrid_potential = np.zeros([n_iterations])
#run n_iterations of the endpoint perturbation:
hybrid_trajectory = unit.Quantity(np.zeros([n_iterations, lambda_thermodynamic_state.system.getNumParticles(), 3]), unit.nanometers) # DEBUG
nonalchemical_trajectory = unit.Quantity(np.zeros([n_iterations, nonalchemical_thermodynamic_state.system.getNumParticles(), 3]), unit.nanometers) # DEBUG
for iteration in range(n_iterations):
# Generate a new sampler state for the hybrid system
mc_move.apply(lambda_thermodynamic_state, new_sampler_state)
# Compute the hybrid reduced potential at the new sampler state
hybrid_context, integrator = cache.global_context_cache.get_context(lambda_thermodynamic_state)
new_sampler_state.apply_to_context(hybrid_context, ignore_velocities=True)
hybrid_reduced_potential = lambda_thermodynamic_state.reduced_potential(hybrid_context)
if write_state:
state = hybrid_context.getState(getPositions=True, getParameters=True)
state_xml = openmm.XmlSerializer.serialize(state)
with open(f'state{iteration}_l{lambda_index}.xml', 'w') as outfile:
outfile.write(state_xml)
# Construct a sampler state for the nonalchemical system
if lambda_index == 0:
nonalchemical_positions = factory.old_positions(new_sampler_state.positions)
elif lambda_index == 1:
nonalchemical_positions = factory.new_positions(new_sampler_state.positions)
else:
raise ValueError("The lambda index needs to be either one or zero for this to be meaningful")
nonalchemical_sampler_state = SamplerState(nonalchemical_positions, box_vectors=new_sampler_state.box_vectors)
if write_trajectories:
state = hybrid_context.getState(getPositions=True)
hybrid_trajectory[iteration,:,:] = state.getPositions(asNumpy=True)
nonalchemical_trajectory[iteration,:,:] = nonalchemical_positions
# Compute the nonalchemical reduced potential
nonalchemical_context, integrator = cache.global_context_cache.get_context(nonalchemical_thermodynamic_state)
nonalchemical_sampler_state.apply_to_context(nonalchemical_context, ignore_velocities=True)
nonalchemical_reduced_potential = nonalchemical_thermodynamic_state.reduced_potential(nonalchemical_context)
# Compute and store the work
w[iteration] = nonalchemical_reduced_potential - hybrid_reduced_potential
non_potential[iteration] = nonalchemical_reduced_potential
hybrid_potential[iteration] = hybrid_reduced_potential
if print_work:
print(f'{iteration:8d} {hybrid_reduced_potential:8.3f} {nonalchemical_reduced_potential:8.3f} => {w[iteration]:8.3f}')
if write_trajectories:
if lambda_index == 0:
nonalchemical_mdtraj_topology = md.Topology.from_openmm(factory._topology_proposal.old_topology)
elif lambda_index == 1:
nonalchemical_mdtraj_topology = md.Topology.from_openmm(factory._topology_proposal.new_topology)
md.Trajectory(hybrid_trajectory / unit.nanometers, factory.hybrid_topology).save(f'hybrid{lambda_index}.pdb')
md.Trajectory(nonalchemical_trajectory / unit.nanometers, nonalchemical_mdtraj_topology).save(f'nonalchemical{lambda_index}.pdb')
# Analyze data and return results
[t0, g, Neff_max] = timeseries.detectEquilibration(w)
w_burned_in = w[t0:]
[df, ddf] = pymbar.EXP(w_burned_in)
ddf_corrected = ddf * np.sqrt(g)
results = [df, ddf_corrected, t0, Neff_max]
return results, non_potential, hybrid_potential
