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 compute_potential_components(context,
beta=beta,
platform=DEFAULT_PLATFORM):
"""
Compute potential energy, raising an exception if it is not finite.
Parameters
----------
context : simtk.openmm.Context
The context from which to extract, System, parameters, and positions.
"""
# Make a deep copy of the system.
import copy
from perses.dispersed.utils import configure_platform
platform = configure_platform(platform.getName(),
fallback_platform_name='Reference',
precision='double')
system = context.getSystem()
system = copy.deepcopy(system)
# Get positions.
positions = context.getState(getPositions=True).getPositions(asNumpy=True)
# Get Parameters
parameters = context.getParameters()
# Segregate forces.
for index in range(system.getNumForces()):
force = system.getForce(index)
force.setForceGroup(index)
# Create new Context.
integrator = openmm.VerletIntegrator(1.0 * unit.femtoseconds)
context = openmm.Context(system, integrator, platform)
context.setPositions(positions)
for (parameter, value) in parameters.items():
context.setParameter(parameter, value)
energy_components = list()
for index in range(system.getNumForces()):
force = system.getForce(index)
forcename = force.__class__.__name__
groups = 1 << index
potential = beta * context.getState(
getEnergy=True, groups=groups).getPotentialEnergy()
energy_components.append((forcename, potential))
del context, integrator
return energy_components
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 __init__(self,
openmm_pdf_state,
integrator,
context_cache=None,
reassign_velocities=False,
n_restart_attempts=0):
"""
call the BaseIntegratorMove init method
arguments
openmm_pdf_state : coddiwomple.openmm_adapters.OpenMMPDFState
the pdf state of the propagator
integrator : openmm.Integrator
integrator of dynamics
context_cache : openmmtools.cache.ContextCache, optional
The ContextCache to use for Context creation. If None, the global cache
openmmtools.cache.global_context_cache is used (default is None).
reassign_velocities : bool, optional
If True, the velocities will be reassigned from the Maxwell-Boltzmann
distribution at the beginning of the move (default is False).
n_restart_attempts : int, optional
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.
attributes
pdf_state : coddiwomple.openmm_adapters.OpenMMPDFState
integrator : openmm.Integrator
context_cache : openmmtools.cache.ContextCache
reassign_velocities : bool
n_restart_attempts : int or None
"""
super().__init__(n_steps = None,
context_cache = context_cache,
reassign_velocities = reassign_velocities,
n_restart_attempts = n_restart_attempts)
_logger.debug(f"successfully executed {BaseIntegratorMove.__class__.__name__} init.")
import openmmtools.cache as cache
from perses.dispersed.utils import check_platform, configure_platform
from openmmtools.utils import get_fastest_platform
try:
cache.global_context_cache.platform = configure_platform(get_fastest_platform().getName())
except Exception as e:
_logger.warning(f"platform configuration error: {e}")
self.pdf_state = openmm_pdf_state
# Check if we have to use the global cache.
if self.context_cache is None:
self._context_cache = cache.global_context_cache
else:
self._context_cache = self.context_cache
# Create context and reset integrator for good measure
self.context, self.integrator = self._context_cache.get_context(self.pdf_state, integrator)
self.integrator.reset()
_logger.debug(f"successfully equipped integrator: {self.integrator.__class__.__name__}")
_logger.debug(f"integrator printable: {self.integrator.pretty_print()}")
#############################################################################
# HYBRID SYSTEM SAMPLERS
#############################################################################
from perses.annihilation.lambda_protocol import RelativeAlchemicalState, LambdaProtocol
from openmmtools.multistate import sams, replicaexchange
from openmmtools import cache, utils
from perses.dispersed.utils import configure_platform
from openmmtools import cache
cache.global_context_cache.platform = configure_platform(
utils.get_fastest_platform().getName())
from openmmtools.states import *
from perses.dispersed.utils import create_endstates
import numpy as np
import copy
import logging
_logger = logging.getLogger()
_logger.setLevel(logging.INFO)
_logger = logging.getLogger("multistate")
class HybridCompatibilityMixin(object):
"""
Mixin that allows the MultistateSampler to accommodate the situation where
unsampled endpoints have a different number of degrees of freedom.
"""
def __init__(self, *args, hybrid_factory=None, **kwargs):
self._hybrid_factory = hybrid_factory
def validate_endstate_energies(topology_proposal,
htf,
added_energy,
subtracted_energy,
beta=1.0 / kT,
ENERGY_THRESHOLD=1e-6,
platform=DEFAULT_PLATFORM,
trajectory_directory=None):
"""
Function to validate that the difference between the nonalchemical versus alchemical state at lambda = 0,1 is
equal to the difference in valence energy (forward and reverse).
Parameters
----------
topology_proposal : perses.topology_proposal.TopologyProposal object
top_proposal for relevant transformation
htf : perses.new_relative.HybridTopologyFactory object
hybrid top factory for setting alchemical hybrid states
added_energy : float
reduced added valence energy
subtracted_energy: float
reduced subtracted valence energy
Returns
-------
zero_state_energy_difference : float
reduced potential difference of the nonalchemical and alchemical lambda = 0 state (corrected for valence energy).
one_state_energy_difference : float
reduced potential difference of the nonalchemical and alchemical lambda = 1 state (corrected for valence energy).
"""
import copy
#import openmmtools.cache as cache
#context_cache = cache.global_context_cache
from perses.dispersed.utils import configure_platform
from perses.utils import data
platform = configure_platform(platform.getName(),
fallback_platform_name='Reference',
precision='double')
#create copies of old/new systems and set the dispersion correction
top_proposal = copy.deepcopy(topology_proposal)
forces = {
top_proposal._old_system.getForce(index).__class__.__name__:
top_proposal._old_system.getForce(index)
for index in range(top_proposal._old_system.getNumForces())
}
forces['NonbondedForce'].setUseDispersionCorrection(False)
forces = {
top_proposal._new_system.getForce(index).__class__.__name__:
top_proposal._new_system.getForce(index)
for index in range(top_proposal._new_system.getNumForces())
}
forces['NonbondedForce'].setUseDispersionCorrection(False)
#create copy of hybrid system, define old and new positions, and turn off dispersion correction
hybrid_system = copy.deepcopy(htf.hybrid_system)
hybrid_system_n_forces = hybrid_system.getNumForces()
for force_index in range(hybrid_system_n_forces):
forcename = hybrid_system.getForce(force_index).__class__.__name__
if forcename == 'NonbondedForce':
hybrid_system.getForce(force_index).setUseDispersionCorrection(
False)
old_positions, new_positions = htf._old_positions, htf._new_positions
#generate endpoint thermostates
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 = [('real-old', nonalch_zero, old_positions,
top_proposal._old_system.getDefaultPeriodicBoxVectors()),
('hybrid-old', alch_zero, htf._hybrid_positions,
hybrid_system.getDefaultPeriodicBoxVectors()),
('hybrid-new', alch_one, htf._hybrid_positions,
hybrid_system.getDefaultPeriodicBoxVectors()),
('real-new', nonalch_one, new_positions,
top_proposal._new_system.getDefaultPeriodicBoxVectors())]
rp_list = []
for (state_name, state, pos, box_vectors) in attrib_list:
integrator = openmm.VerletIntegrator(1.0 * unit.femtoseconds)
context = state.create_context(integrator, platform)
samplerstate = states.SamplerState(positions=pos,
box_vectors=box_vectors)
samplerstate.apply_to_context(context)
rp = state.reduced_potential(context)
rp_list.append(rp)
energy_comps = compute_potential_components(context)
for name, force in energy_comps:
print("\t\t\t{}: {}".format(name, force))
_logger.debug(
f'added forces:{sum([energy for name, energy in energy_comps])}')
_logger.debug(f'rp: {rp}')
if trajectory_directory is not None:
_logger.info(
f'Saving {state_name} state xml to {trajectory_directory}/{state_name}-state.gz'
)
state = context.getState(getPositions=True,
getVelocities=True,
getForces=True,
getEnergy=True,
getParameters=True)
data.serialize(state,
f'{trajectory_directory}-{state_name}-state.gz')
del context, integrator
nonalch_zero_rp, alch_zero_rp, alch_one_rp, nonalch_one_rp = rp_list[
0], rp_list[1], rp_list[2], rp_list[3]
ratio = abs((nonalch_zero_rp - alch_zero_rp + added_energy) /
(nonalch_zero_rp + alch_zero_rp + added_energy))
assert ratio < ENERGY_THRESHOLD, f"The ratio in energy difference for the ZERO state is {ratio}.\n This is greater than the threshold of {ENERGY_THRESHOLD}.\n real-zero: {nonalch_zero_rp} \n alc-zero: {alch_zero_rp} \nadded-valence: {added_energy}"
ratio = abs((nonalch_one_rp - alch_one_rp + subtracted_energy) /
(nonalch_one_rp + alch_one_rp + subtracted_energy))
assert ratio < ENERGY_THRESHOLD, f"The ratio in energy difference for the ONE state is {ratio}.\n This is greater than the threshold of {ENERGY_THRESHOLD}.\n real-one: {nonalch_one_rp} \n alc-one: {alch_one_rp} \nsubtracted-valence: {subtracted_energy}"
return abs(nonalch_zero_rp - alch_zero_rp +
added_energy), abs(nonalch_one_rp - alch_one_rp +
subtracted_energy)