def __init__(self,
system,
alchemical_composability=AlchemicalState,
temperature=300 * unit.kelvin,
pressure=1.0 * unit.atmosphere,
**kwargs):
"""
Subclass of coddiwomple.states.PDFState that specifically handles openmmtools.states.CompoundThermodynamicStates
Init method does the following:
1. create a openmmtools.states.ThermodynamicState with the given system, temperature, and pressure
2. create an openmmtools.alchemy.AlchemicalState from the given system
3. create an openmtools.states.CompoundThermodynamicState 1 and 2
arguments
system : openmm.System
system object to wrap
alchemical_composability : openmmtools.alchemy.AlchemicalState (super), default openmmtools.alchemy.AlchemicalState
the class holding the method `from_system` which can compose an alchemical state with exposed parameters from a system
temperature : float * unit.kelvin (or temperature units), default 300.0 * unit.kelvin
temperature of the system
pressure : float * unit.atmosphere (or pressure units), default 1.0 * unit.atmosphere
pressure of the system
init method is adapted from https://github.com/choderalab/openmmtools/blob/110524bc5079af77d31f5fab464edd7b668ff5ac/openmmtools/states.py#L2766-L2783
"""
from openmmtools.alchemy import AlchemicalState
openmm_pdf_state = ThermodynamicState(system, temperature, pressure)
alchemical_state = alchemical_composability.from_system(
system, **kwargs)
assert isinstance(
alchemical_state, IComposableState
), f"alchemical state is not an instance of IComposableState"
self.__dict__ = openmm_pdf_state.__dict__
self._composable_states = [alchemical_state]
self.set_system(self._standard_system, fix_state=True)
#class create an internal context
integrator = get_dummy_integrator()
platform = configure_platform(utils.get_fastest_platform().getName())
self._internal_context = self.create_context(integrator,
platform=platform)
_logger.debug(
f"successfully instantiated OpenMMPDFState equipped with the following parameters: {self._parameters}"
)
def test_HybridSystemFactory():
"""
run the `make_HybridSystemFactory`
"""
from qmlify.openmm_torch.utils import configure_platform
from openmmtools import utils
hsf, testsystem_class = make_HybridSystemFactory()
platform = configure_platform(platform_name=utils.get_fastest_platform(),
fallback_platform_name='CPU',
precision='mixed')
# non/alchemical integrators
from openmmtools.integrators import LangevinIntegrator
nonalch_int = LangevinIntegrator(temperature=DEFAULT_TEMPERATURE)
alch_int = LangevinIntegrator(temperature=DEFAULT_TEMPERATURE)
system, alch_system = hsf._old_system, hsf.system
nonalch_context, alch_context = openmm.Context(system, nonalch_int,
platform), openmm.Context(
alch_system, alch_int,
platform)
for context in [nonalch_context, alch_context]:
context.setPositions(testsystem_class.positions)
context.setPeriodicBoxVectors(*system.getDefaultPeriodicBoxVectors())
nonalch_energy = nonalch_context.getState(
getEnergy=True).getPotentialEnergy().value_in_unit_system(
unit.md_unit_system)
alch_energy = alch_context.getState(
getEnergy=True).getPotentialEnergy().value_in_unit_system(
unit.md_unit_system)
assert abs(
alch_energy - nonalch_energy
) < ENERGY_DIFFERENCE_TOLERANCE, f"the nonalchemical energy of {nonalch_energy} and the alchemical energy (at lambda=0) of {alch_energy} has a difference that is greater than {ENERGY_DIFFERENCE_TOLERANCE}"
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()}")
except:
print(
"Warning: Returning {} platform instead of requested platform {}".
format(fallback_platform_name, platform_name))
platform = fallback_platform
print(
f"conducting subsequent work with the following platform: {platform.getName()}"
)
return platform
#########
cache.global_context_cache.platform = configure_platform(
utils.get_fastest_platform().getName())
#########
#smc functions
def compute_survival_rate(sMC_particle_ancestries):
"""
compute the time-series survival rate as a function of resamples
Arguments
---------
sMC_particle_ancestries : dict of {_direction : list(np.array(ints))}
dict of the particle ancestor indices
Returns
################################################################################
# IMPORTS
################################################################################
from simtk import openmm, unit
from simtk.openmm import app
import mdtraj as md
import numpy as np
from openmmtools import testsystems
import copy
import time
from openmmtools.states import SamplerState, ThermodynamicState, CompoundThermodynamicState, group_by_compatibility
from openmmtools.multistate import sams, replicaexchange
from openmmtools import cache, utils
from perses.dispersed.utils import configure_platform
cache.global_context_cache.platform = configure_platform(utils.get_fastest_platform().getName())
from perses.annihilation.ncmc_switching import NCMCEngine
from perses.dispersed import feptasks
from perses.storage import NetCDFStorageView
from perses.utils.openeye import smiles_to_oemol
################################################################################
# LOGGER
################################################################################
import logging
_logger = logging.getLogger()
_logger.setLevel(logging.INFO)
_logger = logging.getLogger("samplers")
def prepare_ml_system(positions,
topology,
system,
residue_indices,
model_name='ani2x',
save_filename='animodel.pt',
torch_scale_name='torch_scale',
torch_scale_default_value=0.,
HybridSystemFactory_kwargs={},
minimizer_kwargs={'maxIterations': 1000}):
"""
prepare an ani-force-compatible system with built-in lambda assertions and energy compatibility assertions
"""
from qmlify.openmm_torch.torchforce_generator import torch_alchemification_wrapper
from openmmtools import utils
from openmmtools.integrators import LangevinIntegrator
from openmmtools.constants import kB
from simtk.openmm import LocalEnergyMinimizer
import numpy as np
DEFAULT_TEMPERATURE = 300.0 * unit.kelvin
ENERGY_DIFFERENCE_TOLERANCE = 1e-2
_logger.info("preparing ML system and initializing assertions...")
# make ml system and hybrid factory
_logger.info(
f"executing torch alchemification wrapper to make ml_system and hybrid_factory"
)
ml_system, hybrid_factory = torch_alchemification_wrapper(
topology,
system,
residue_indices,
model_name,
save_filename,
torch_scale_name,
torch_scale_default_value,
)
# get platform
platform = configure_platform(
platform_name=utils.get_fastest_platform().getName())
beta = 1. / (kB * DEFAULT_TEMPERATURE)
# get integrators
old_mm_int = LangevinIntegrator(temperature=DEFAULT_TEMPERATURE)
mm_int = LangevinIntegrator(temperature=DEFAULT_TEMPERATURE)
ml_int = LangevinIntegrator(temperature=DEFAULT_TEMPERATURE)
# make mm contexts at lambda 0
mm_context = openmm.Context(hybrid_factory.system, mm_int, platform)
mm_context.setPositions(positions)
mm_context.setPeriodicBoxVectors(
*hybrid_factory.system.getDefaultPeriodicBoxVectors())
# get the swig parameters and check the alchemical mm system
_logger.debug(
f"ensuring appropriate lambda initialization at lambda0 for alchemical system..."
)
mm_swig_params = mm_context.getParameters()
for name in mm_swig_params:
assert DEFAULT_LAMBDA0s[name] == mm_swig_params[
name], f"swig parameter {name} is {mm_swig_params[name]} but should be {DEFAULT_LAMBDA0s[name]}"
# minimize mm context
LocalEnergyMinimizer.minimize(mm_context, **minimizer_kwargs)
# apply the positions to the ml context
ml_context = openmm.Context(ml_system, ml_int, platform)
# check the ml context swig parameters
ml_context.setPositions(
mm_context.getState(getPositions=True).getPositions(asNumpy=True))
ml_context.setPeriodicBoxVectors(
*hybrid_factory.system.getDefaultPeriodicBoxVectors())
# get the swig parameters and check the alchemical ml system
ml_swig_params = ml_context.getParameters()
torch_parameters_lambda0 = {
torch_scale_name: torch_scale_default_value,
f'auxiliary_{torch_scale_name}': 1.
} #this is hard coded...want this?
_logger.debug(
f"ensuring appropriate lambda initialization at lambda0 for ml alchemical system..."
)
for name in ml_swig_params:
if name in list(DEFAULT_LAMBDA0s.keys()):
assert DEFAULT_LAMBDA0s[name] == ml_swig_params[
name], f"swig parameter {name} is {ml_swig_params[name]} but should be {DEFAULT_LAMBDA0s[name]}"
else: #it is a special torch parameter
assert ml_swig_params[name] == torch_parameters_lambda0[name]
# build the old (nonalch) system
old_mm_context = openmm.Context(hybrid_factory._old_system, old_mm_int,
platform)
old_mm_context.setPositions(
mm_context.getState(getPositions=True).getPositions(asNumpy=True))
old_mm_context.setPeriodicBoxVectors(
*hybrid_factory.system.getDefaultPeriodicBoxVectors())
# now check energy by components
_logger.debug(
f"computing potential components of _all_ contexts...standby.")
old_mm_potential_components = compute_potential_components(
old_mm_context, beta, platform)
mm_potential_components = compute_potential_components(
mm_context, beta, platform)
# ml_potential_components = compute_potential_components(ml_context, beta, platform) #we can't do this right now since there is a bug...
sum_old_mm_potential_components = np.sum(
[tup[1] for tup in old_mm_potential_components])
sum_mm_potential_components = np.sum(
[tup[1] for tup in mm_potential_components])
# sum_ml_potential_components = np.sum(list(ml_potential_components.values()))
mm_difference = abs(sum_old_mm_potential_components -
sum_mm_potential_components)
ml_difference = abs(
sum_mm_potential_components -
ml_context.getState(getEnergy=True).getPotentialEnergy() * beta)
try:
_logger.info(f"checking mm bookkeeping energies...")
assert mm_difference < ENERGY_DIFFERENCE_TOLERANCE
except Exception as e:
_logger.warning(
f"{e}; difference between energies of the lambda0 alchemical mm and nonalchemical mm energy is {mm_difference}, which is higher than the tolerance of {ENERGY_DIFFERENCE_TOLERANCE}"
)
try:
_logger.info(f"checking mm bookkeeping energies...")
ml_difference < ENERGY_DIFFERENCE_TOLERANCE
except Exception as e:
_logger.warning(
f"{e}; difference between energies of the lambda0 alchemical mm and ml energy is {mm_difference}, which is higher than the tolerance of {ENERGY_DIFFERENCE_TOLERANCE}"
)
# we cannot do the following...
# for key, val in DEFAULT_LAMBDA1s:
# mm_context.setParameter(key, val)
# mm_final_potential_components = compute_potential_components(mm_context, beta, platform)
#
# try:
# _logger.info(f"checking ml bookkeeping energies...")
# """
# here, we are making sure that the alchemical forces starting with `Custom` are all zero and that the other components are unchanged
# """
# for forcename, energy in mm_final_potential_components.items():
# if forcename in [torch_scale_name, f'auxiliary_{torch_scale_name}']:
# # don't check the torch force...at least not yet
# pass
# elif forcename[:7] == 'Custom':
# assert np.isclose(energy, 0.), f"the energy of {forcename} at lambda 1 is {energy} when it should be 0."
# else:
# lambda0_energy = mm_potential_components[forcename]
# assert np.isclose(energy, lambda0_energy), f"the energy of {forcename} at lambda 1 is {energy} when it should be {lambda0_energy}"
# except Exception as e:
# _logger.warning(f"{e}; there is an issue associated with the lambda1 endstate energy bookkeeping. see above for which assertion failed.")
# TODO : add a test for scaling lambdas?
# remove the contexts and integrators used for testing (this will shore up some memory)...
_logger.debug(f"removing contexts...")
for context in [old_mm_context, mm_context, ml_context]:
del context
_logger.debug(f"removing integrators...")
for integrator in [old_mm_int, mm_int, ml_int]:
del integrator
return ml_system, hybrid_factory
else:
from cStringIO import StringIO
from commands import getstatusoutput
from openmmtools.constants import kB
from openmmtools import alchemy, states
import contextlib
from openmmtools import utils
################################################################################
# CONSTANTS
################################################################################
temperature = 300.0 * unit.kelvin
kT = kB * temperature
beta = 1.0 / kT
ENERGY_THRESHOLD = 1e-1
DEFAULT_PLATFORM = utils.get_fastest_platform()
################################################################################
# UTILITIES
################################################################################]
import logging
_logger = logging.getLogger("tests-utils")
_logger.setLevel(logging.INFO)
@contextlib.contextmanager
def enter_temp_directory():
"""Create and enter a temporary directory; used as context manager."""
import tempfile
temp_dir = tempfile.mkdtemp()
nonalch_system = testsystem_class.system
# In[8]:
nonalch_int = LangevinIntegrator()
ml_int = LangevinIntegrator(splitting= 'V0 V1 R O R V1 V0')
# In[9]:
print("getting platform")
platform = configure_platform(utils.get_fastest_platform().getName())
# In[10]:
print(f"getting contexts")
nonalch_context = openmm.Context(nonalch_system, nonalch_int, platform)
ml_context = openmm.Context(ml_system, ml_int, platform)
# In[13]:
print("setting positions...")
for context in [nonalch_context, ml_context]:
context.setPositions(testsystem_class.positions)
context.setPeriodicBoxVectors(*testsystem_class.system.getDefaultPeriodicBoxVectors())