def test_sim_factory():
# check sim_factory returns a function when passed a testsystem
construct_sim = utils.sim_factory(AlanineDipeptideVacuum())
assert (callable(construct_sim))
# check that the resulting function returns an app.Simulation when passed an integrator
sim = construct_sim(LangevinIntegrator())
assert (isinstance(sim, app.Simulation))
def test_clone_state():
# construct source system
construct_sim = utils.sim_factory(AlanineDipeptideVacuum())
sim_a = construct_sim(LangevinIntegrator())
sim_b = construct_sim(LangevinIntegrator())
sim_a.step(50)
sim_b.step(50)
utils.clone_state(sim_a, sim_b)
state_a = sim_a.context.getState(getPositions=True, getVelocities=True)
state_b = sim_b.context.getState(getPositions=True, getVelocities=True)
# check that the positions were cloned
assert (np.isclose(state_a.getPositions(asNumpy=True),
state_b.getPositions(asNumpy=True)).all())
# check that the velocities were cloned
assert (np.isclose(state_a.getVelocities(asNumpy=True),
state_b.getVelocities(asNumpy=True)).all())
def test_get_a_paired_kldiv_sample():
# just make sure it returns floats?
testsystem = AlanineDipeptideVacuum(constraints=None)
construct_sim = utils.sim_factory(testsystem)
equilibrium_sim = construct_sim(GHMCIntegrator(timestep=0.25 * unit.femtosecond))
equilibrium_sim.minimizeEnergy()
equilibrium_sim.step(1000)
reference_sim = construct_sim(
LangevinIntegrator(splitting='O V R V O', measure_shadow_work=True, measure_heat=True,
timestep=0.01 * unit.femtosecond))
test_sim = construct_sim(
LangevinIntegrator(splitting='O V R V O', measure_shadow_work=True, measure_heat=True,
timestep=3 * unit.femtosecond))
kldiv_reference, kldiv_test = error.get_a_paired_kldiv_sample(equilibrium_sim, reference_sim, test_sim,
protocol_length=100)
assert (isinstance(kldiv_reference, float))
assert (isinstance(kldiv_test, float))
# more stringent test: check that median of reference here is lower than median of test_sim
reference_samples, test_samples = [], []
for _ in range(100):
equilibrium_sim.step(100)
kldiv_reference, kldiv_test = error.get_a_paired_kldiv_sample(equilibrium_sim, reference_sim, test_sim,
protocol_length=100)
reference_samples.append(kldiv_reference)
test_samples.append(kldiv_test)
assert (np.median(reference_samples) <= np.median(test_samples))
# check that recycle_v works
_ = error.get_a_paired_kldiv_sample(equilibrium_sim, reference_sim, test_sim,
protocol_length=100, recycle_v=True)
# check that it doesn't fail silently when reference_sim and test_sim are at different temperatures
with pytest.raises(ValueError):
test_sim.integrator.setTemperature(200 * unit.kelvin)
error.get_a_paired_kldiv_sample(equilibrium_sim, reference_sim, test_sim)
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 test_measure_shadow_work():
# test that it's consistent with openmmtools
construct_sim = utils.sim_factory(AlanineDipeptideVacuum(constraints=None))
sim = construct_sim(
LangevinIntegrator(splitting='O V R V O',
measure_heat=True,
measure_shadow_work=True,
timestep=2.0 * unit.femtoseconds))
sim.runForClockTime(1 * unit.seconds)
w_shads_here = []
w_shads_openmmtools = []
for _ in range(10):
w_shad_prev = sim.integrator.get_shadow_work(dimensionless=True)
w_shads_here.append(utils.measure_shadow_work(sim, 10))
w_shad_new = sim.integrator.get_shadow_work(dimensionless=True)
w_shads_openmmtools.append(w_shad_new - w_shad_prev)
assert (np.isclose(w_shads_here, w_shads_openmmtools).all())
}
thermodynamic_states.append( ThermodynamicState(system=system, temperature=temperature, pressure=pressure, parameters=parameters) )
# Umbrella on state
alchemical_lambda = 0.0
for theta0 in torsion_umbrella_values:
for r0 in distance_umbrella_values:
parameters = {
'torsion_K' : torsion_K.value_in_unit_system(unit.md_unit_system), 'torsion_theta0' : theta0.value_in_unit_system(unit.md_unit_system), # umbrella parameters
'distance_K' : distance_K.value_in_unit_system(unit.md_unit_system), 'distance_r0' : r0.value_in_unit_system(unit.md_unit_system), # umbrella parameters
}
thermodynamic_states.append( ThermodynamicState(system=system, temperature=temperature, pressure=pressure, parameters=parameters) )
# Select platform automatically; use mixed precision
from openmmtools.integrators import LangevinIntegrator
integrator = LangevinIntegrator(temperature=temperature, collision_rate=collision_rate, timestep=timestep)
context = openmm.Context(system, integrator)
platform = context.getPlatform()
del context
try:
platform.setPropertyDefaultValue('Precision', 'mixed')
platform.setPropertyDefaultValue('DeterministicForces', 'true')
except:
pass
# Minimize
if minimize:
print('Minimizing...')
integrator = openmm.VerletIntegrator(timestep)
context = openmm.Context(system, integrator)
context.setPeriodicBoxVectors(*state.getPeriodicBoxVectors())
nonbondedCutoff=1.2 * unit.nanometers,
constraints=app.HBonds,
rigidWater=True,
ewaldErrorTolerance=0.0005)
#### Thermodynamic State
kB = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA
temperature = 300.0 * unit.kelvin
pressure = 1.0 * unit.atmosphere
#### Integrator
friction = 1.0 / unit.picosecond
step_size = 2.0 * unit.femtoseconds
integrator = LangevinIntegrator(temperature, friction, step_size)
integrator.setConstraintTolerance(0.00001)
#### Barostat
barostat_interval = 25
barostat = mm.MonteCarloBarostat(pressure, temperature, barostat_interval)
system.addForce(barostat)
#### Platform
platform = mm.Platform.getPlatformByName('CUDA')
properties = {'CudaPrecision': 'mixed'}
#### Simulation
system_load_xml_filename = 'system.xml'
state_load_xml_filename = 'state.xml'
# Read in the model
pdb_filename = output_prefix + 'equilibrated.pdb'
print('Loading %s' % pdb_filename)
pdb = app.PDBFile(pdb_filename)
# Load the system
print('Loading OpenMM System...')
with open(output_prefix + system_load_xml_filename, 'r') as infile:
system = openmm.XmlSerializer.deserialize(infile.read())
# Serialize integrator
print('Serializing integrator %s' % output_prefix + integrator_xml_filename)
integrator = LangevinIntegrator(temperature, collision_rate, timestep,
splitting)
with open(output_prefix + integrator_xml_filename, 'w') as outfile:
xml = openmm.XmlSerializer.serialize(integrator)
outfile.write(xml)
print('Loading OpenMM State...')
with open(output_prefix + state_load_xml_filename, 'r') as infile:
state = openmm.XmlSerializer.deserialize(infile.read())
context = openmm.Context(system, integrator)
context.setState(state)
# Equilibrate
print('Equilibrating...')
initial_time = time.time()
for iteration in progressbar.progressbar(range(niterations)):
integrator.step(nsteps)
def make_integrator(splitting):
return LangevinIntegrator(splitting=splitting,
timestep=2 * unit.femtoseconds,
measure_shadow_work=False,
measure_heat=False)
def relax_structure(temperature,
system,
positions,
nequil=1000,
n_steps_per_iteration=250,
platform_name='OpenCL',
timestep=2. * unit.femtosecond,
collision_rate=90. / unit.picosecond,
**kwargs):
"""
arguments
temperature : simtk.unit.Quantity with units compatible with kelvin
temperature of simulation
system : openmm.System
system object for simulation
positions : simtk.unit.Quantity of shape (natoms,3) with units compatible with nanometers
Positions of the atoms in the system
nequil : int, default = 1000
number of equilibration applications
n_steps_per_iteration : int, default = 250
numper of steps per nequil
platform name : str default='OpenCL'
platform to run openmm on. OpenCL is best as this is what is used on FAH
timestep : simtk.unit.Quantity, default = 2*unit.femtosecond
timestep for equilibration NOT for production
collision_rate : simtk.unit.Quantity, default=90./unit.picosecond
return
state : openmm.State
state of simulation (getEnergy=True, getForces=True, getPositions=True, getVelocities=True, getParameters=True)
"""
from openmmtools.integrators import LangevinIntegrator
_logger.info(f'Starting to relax')
integrator = LangevinIntegrator(temperature=temperature,
timestep=timestep,
collision_rate=collision_rate)
platform = openmm.Platform.getPlatformByName(platform_name)
# prepare the plaform
if platform_name in ['CUDA', 'OpenCL']:
platform.setPropertyDefaultValue('Precision', 'mixed')
if platform_name in ['CUDA']:
platform.setPropertyDefaultValue('DeterministicForces', 'true')
context = openmm.Context(system, integrator, platform)
context.setPeriodicBoxVectors(*system.getDefaultPeriodicBoxVectors())
context.setPositions(positions)
_logger.info(f'Starting to minimise')
openmm.LocalEnergyMinimizer.minimize(context)
# Equilibrate
_logger.info(f'set velocities to temperature')
context.setVelocitiesToTemperature(temperature)
_logger.info(
f'Starting to equilibrate for {nequil*n_steps_per_iteration*timestep}')
integrator.step(nequil * n_steps_per_iteration)
context.setVelocitiesToTemperature(temperature)
state = context.getState(getEnergy=True,
getForces=True,
getPositions=True,
getVelocities=True,
getParameters=True)
_logger.info(f'Relax done')
del context, integrator
return state
def _equil_NPT_OpenMM_protocol_0(topology,
positions,
temperature=300.0 * _unit.kelvin,
pressure=1.0 * _unit.atmosphere,
time=1.0 * _unit.nanosecond,
forcefield=None,
verbose=True,
progress_bar=True):
import numpy as np
import openmm.app as app
import openmm as mm
from openmmtools.integrators import LangevinIntegrator, GeodesicBAOABIntegrator
if progress_bar:
from tqdm import tqdm
else:
def tqdm(arg):
return arg
#item needs to be openmm.modeller
forcefield = app.ForceField("amber99sbildn.xml", "tip3p.xml")
topology = item.topology
positions = item.positions
system = forcefield_generator.createSystem(topology,
contraints=app.HBonds,
nonbondedMethod=app.PME,
nonbondedCutoff=1.0 *
_unit.nanometers,
rigidWater=True,
ewaldErrorTolerance=0.0005)
## Thermodynamic State
kB = _unit.BOLTZMANN_CONSTANT_kB * _unit.AVOGADRO_CONSTANT_NA
temperature = temperature
pressure = pressure
## Barostat
barostat_frequency = 25 # steps
barostat = mm.MonteCarloBarostat(pressure, temperature, barostat_frequency)
system.addForce(barostat)
## Integrator
friction = 1.0 / _unit.picosecond
step_size = 2.0 * _unit.femtoseconds
integrator = LangevinIntegrator(temperature, friction, step_size)
integrator.setConstraintTolerance(0.00001)
## Platform
platform = mm.Platform.getPlatformByName('CUDA')
properties = {'CudaPrecision': 'mixed'}
## Simulation
simulation = app.Simulation(topology, system, integrator, platform,
properties)
simulation.context.setPositions(positions)
simulation.context.setVelocitiesToTemperature(temperature)
time_equilibration = time
time_iteration = 0.2 * _unit.picoseconds
number_iterations = int(time_equilibration / time_iteration)
steps_iteration = int(time_iteration / step_size)
steps_equilibration = number_iterations * steps_iteration
## Reporters
net_mass, n_degrees_of_freedom = m3t.get(system,
net_mass=True,
n_degrees_of_freedom=True)
niters = number_iterations
data = dict()
data['time'] = _unit.Quantity(np.zeros([niters], np.float64),
_unit.picoseconds)
data['potential'] = _unit.Quantity(np.zeros([niters], np.float64),
_unit.kilocalories_per_mole)
data['kinetic'] = _unit.Quantity(np.zeros([niters], np.float64),
_unit.kilocalories_per_mole)
data['volume'] = _unit.Quantity(np.zeros([niters], np.float64),
_unit.angstroms**3)
data['density'] = _unit.Quantity(np.zeros([niters], np.float64),
_unit.gram / _unit.centimeters**3)
data['kinetic_temperature'] = unit.Quantity(np.zeros([niters], np.float64),
_unit.kelvin)
for iteration in tqdm(range(number_iterations)):
integrator.step(steps_iteration)
state = simulation.context.getState(getEnergy=True)
time = state.getTime()
potential_energy = state.getPotentialEnergy()
kinetic_energy = state.getKineticEnergy()
volume = state.getPeriodicBoxVolume()
density = (net_mass / volume).in_units_of(unit.gram /
unit.centimeter**3)
kinetic_temperature = (2.0 * kinetic_energy /
kB / n_degrees_of_freedom).in_units_of(
unit.kelvin) # (1/2) ndof * kB * T = KE
data['time'][iteration] = time
data['potential'] = potential_energy
data['kinetic'] = kinetic_energy
data['volume'] = volume
data['density'] = density
data['kinetic_temperature'] = kinetic_temperature
final_state = simulation.context.getState(getPositions=True,
getVelocities=True)
final_positions = final_state.getPositions()
final_velocities = final_state.getVelocities()
return final_positions, final_velocities, data
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
import pytest
from openmmtools.integrators import LangevinIntegrator
from openmmtools.testsystems import AlanineDipeptideVacuum
from simtk import unit
from thresholds import utils, stability
testsystem = AlanineDipeptideVacuum()
construct_sim = utils.sim_factory(testsystem)
tiny_dt = 0.1 * unit.femtoseconds
stable_sim = construct_sim(LangevinIntegrator(timestep=tiny_dt))
huge_dt = 100 * unit.femtoseconds
unstable_sim = construct_sim(LangevinIntegrator(timestep=huge_dt))
def test_check_stability():
# check that stable simulation is labeled stable
assert (stability.check_stability(stable_sim, n_steps=100))
# check that unstable simulation is labeled unstable
assert (not stability.check_stability(unstable_sim, n_steps=100))
def test_stability_oracle_factory():
def set_initial_conditions(sim):
sim.context.setPositions(testsystem.positions)
sim.context.setVelocitiesToTemperature(298 * unit.kelvin)
iterated_stability_oracle = stability.stability_oracle_factory(
# In[6]:
from qmlify.openmm_torch.utils import *
# In[7]:
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)
softcore_alpha_sterics=0.5,
softcore_alpha_electrostatics=0.5)
# grab the modified system and endstate system...
mod_system = hsf.system
endstate_system = hsf.endstate_system
# now that we have the modified system, we want to get the energy at _this_ endstate and make sure the energy is bookkeeping well with the non-alchemically-modified state.
# In[8]:
from openmmtools.integrators import LangevinIntegrator
from simtk import openmm
# In[9]:
nonalch_int = LangevinIntegrator(temperature=T)
alch_int = LangevinIntegrator(temperature=T)
# In[10]:
nonalch_context, alch_context = openmm.Context(system,
nonalch_int), openmm.Context(
mod_system, alch_int)
# In[11]:
for context in [nonalch_context, alch_context]:
context.setPositions(positions)
context.setPeriodicBoxVectors(*system.getDefaultPeriodicBoxVectors())
# In[12]:
