def test_OpenMMReporter():
"""
test the OpenMMReporter object for its ability to make appropriate trajectory writes for particles.
use the harmonic oscillator testsystem
NOTE : this class will conduct dynamics on 5 particles defined by the harmonic oscillator testsystem in accordance with the coddiwomple.openmm.propagators.OMMBIP
equipped with the coddiwomple.openmm.integrators.OMMLI integrator, but will NOT explicitly conduct a full test on the propagators or integrators.
"""
from coddiwomple.openmm.propagators import OMMBIP
from coddiwomple.openmm.integrators import OMMLI
temperature = 300 * unit.kelvin
pressure = None
from coddiwomple.openmm.states import OpenMMPDFState, OpenMMParticleState
from coddiwomple.particles import Particle
from coddiwomple.openmm.reporters import OpenMMReporter
import shutil
#create the default get_harmonic_testsystem
testsystem, period, collision_rate, timestep, alchemical_functions = get_harmonic_testsystem(temperature = temperature)
#create a particle state and 5 particles
particles = []
for i in range(5):
particle_state = OpenMMParticleState(positions = testsystem.positions) #make a particle state
particle = Particle(index = i, record_state=False, iteration = 0)
particle.update_state(particle_state)
particles.append(particle)
#since we are copying over the positions, we need a simple assert statement to make sure that the id(hex(particle_state.positions)) are separate in memory
position_hexes = [hex(id(particle.state.positions)) for particle in particles]
assert len(position_hexes) == len(list(set(position_hexes))), f"positions are copied identically; this is a problem"
#create a pdf_state
pdf_state = OpenMMPDFState(system = testsystem.system, alchemical_composability = HarmonicAlchemicalState, temperature = temperature, pressure = pressure)
#create an integrator
integrator = OMMLI(temperature=temperature, collision_rate=collision_rate, timestep=timestep)
#create a propagator
propagator = OMMBIP(openmm_pdf_state = pdf_state, integrator = integrator)
steps_per_application = 100
#the only thing we want to do here is to run independent md for each of the particles and save trajectories; at the end, we will delete the directory and the traj files
temp_traj_dir, temp_traj_prefix = os.path.join(os.getcwd(), 'test_dir'), 'traj_prefix'
reporter = OpenMMReporter(trajectory_directory = 'test_dir', trajectory_prefix='traj_prefix', md_topology=testsystem.mdtraj_topology)
assert reporter.write_traj
num_applications=10
for application_index in range(num_applications):
returnables = [propagator.apply(particle.state, n_steps=100, reset_integrator=True, apply_pdf_to_context=True, randomize_velocities=True) for particle in particles]
_save=True if application_index == num_applications-1 else False
reporter.record(particles, save_to_disk=_save)
assert reporter.hex_counter == len(reporter.hex_dict)
assert os.path.exists(temp_traj_dir)
assert os.listdir(temp_traj_dir) is not None
#then we can delete
shutil.rmtree(temp_traj_dir)
def test_OpenMMParticleState():
"""
conduct a class-wide test on coddiwomple.openmm.states.OpenMMParticleState with the `get_harmonic_testsystem` testsystem
this will assert successes on __init__, as well as _all_ methods in the coddiwomple.particles.Particle class
"""
temperature = 300 * unit.kelvin
pressure = None
from coddiwomple.openmm.states import OpenMMPDFState, OpenMMParticleState
from coddiwomple.particles import Particle
#create the default get_harmonic_testsystem
testsystem, period, collision_rate, timestep, alchemical_functions = get_harmonic_testsystem(temperature = temperature)
#test __init__ method
particle_state = OpenMMParticleState(positions = testsystem.positions) #make a particle state
particle = Particle(index = 0, record_state=False, iteration = 0)
#test update_state
assert particle.state is None
assert not particle._record_states
particle.update_state(particle_state)
assert particle.state is not None
#test update_iteration
assert particle.iteration == 0
particle.update_iteration()
assert particle.iteration == 1
#test update ancestry
assert particle.ancestry == [0]
particle.update_ancestry(1)
assert particle.ancestry == [0,1]
def test_OMMBIP():
"""
test OMMBIP (OpenMMBaseIntegratorPropagator) in the baoab regime on the harmonic test system;
specifically, we validate the init, apply, _get_global_integrator_variables, _get_context_parameters methods.
For the sake of testing all of the internal methods, we equip an OMMLI integrator
"""
from coddiwomple.openmm.propagators import OMMBIP
from coddiwomple.openmm.integrators import OMMLI
temperature = 300 * unit.kelvin
pressure = None
from coddiwomple.openmm.states import OpenMMPDFState, OpenMMParticleState
#create the default get_harmonic_testsystem
testsystem, period, collision_rate, timestep, alchemical_functions = get_harmonic_testsystem(temperature = temperature)
particle_state = OpenMMParticleState(positions = testsystem.positions) #make a particle state
num_applications = 100
#create a pdf_state
pdf_state = OpenMMPDFState(system = testsystem.system, alchemical_composability = HarmonicAlchemicalState, temperature = temperature, pressure = pressure)
#create an integrator
integrator = OMMLI(temperature=temperature, collision_rate=collision_rate, timestep=timestep)
#create a propagator
propagator = OMMBIP(openmm_pdf_state = pdf_state, integrator = integrator)
#check the __init__ method for appropriate equipment
assert hex(id(propagator.pdf_state)) == hex(id(pdf_state)) #the defined pdf state is tethered to the propagator (this is VERY important for SMC)
#conduct null application
prior_reduced_potential = pdf_state.reduced_potential(particle_state)
return_state, proposal_work = propagator.apply(particle_state, n_steps=0)
assert proposal_work == 0. #there is no proposal work if returnable_key is None
assert pdf_state.reduced_potential(particle_state) == prior_reduced_potential
propagator_state = propagator.context.getState(getEnergy=True)
assert np.isclose(propagator_state.getPotentialEnergy()*pdf_state.beta, pdf_state.reduced_potential(particle_state))
#check context update internals
prior_reduced_potential = pdf_state.reduced_potential(particle_state)
parameters = pdf_state.get_parameters() #change an alchemical parameter
parameters['testsystems_HarmonicOscillator_U0'] = 1. #update parameter dict
pdf_state.set_parameters(parameters) #set new params
_ = propagator.apply(particle_state, n_steps=0, apply_pdf_to_context=False) # if we do not apply to context, then the internal_context should not be modified
assert propagator._get_context_parameters()['testsystems_HarmonicOscillator_U0'] == 0.
assert np.isclose(propagator.context.getState(getEnergy=True).getPotentialEnergy()*pdf_state.beta, prior_reduced_potential)
_ = propagator.apply(particle_state, n_steps=0, apply_pdf_to_context=True) # if we do apply to context, then the internal_context should be modified
assert propagator._get_context_parameters()['testsystems_HarmonicOscillator_U0'] == 1.
assert np.isclose(prior_reduced_potential + 1.0 * unit.kilojoules_per_mole * pdf_state.beta, propagator.context.getState(getEnergy=True).getPotentialEnergy()*pdf_state.beta)
#check gettable integrator variables
integrator_vars = propagator._get_global_integrator_variables()
#check propagator stability with integrator reset and velocity randomization
_ = propagator.apply(particle_state, n_steps=1000, reset_integrator=True, apply_pdf_to_context=True, randomize_velocities=True)
def propagator_testprep():
"""
wrapper that outputs all necessary Propagator (and subclass) inputs for testing (the test system is butane solvated in tip3p)
returns
pdf_state : coddiwomple.OpenMMPDFState
subclass of openmmtools.states.ThermodynamicState of the system object
pdf_state_subset : coddiwomple.OpenMMPDFState
subclass of openmmtools.states.ThermodynamicState of the system_subset object
integrator : qmlify.propagation.Integrator
integrator that equips the Propagator
ani_handler : qmlify.ANI_force_and_energy
handler of the ani components
atom_map : dict
index map of {md_top_atom_index : md_subset_top_atom_index} of the matching atoms where key is the index of the md_topology atom and the value is the index of the matching md_subset_topology atom
particle : coddiwomple.particles.Particle
particle containing a coddiwomple.openmm.OpenMMParticleState (subclass of openmmtools.states.SamplerState)
"""
import mdtraj as md
from coddiwomple.particles import Particle
from coddiwomple.openmm.states import OpenMMParticleState
from qmlify.utils import generate_propagator_inputs
vac_sys_pos_top, sol_sys_pos_top = generate_testsystem()
vac_system, vac_positions, vac_topology = vac_sys_pos_top
sol_system, sol_positions, sol_topology = sol_sys_pos_top
md_topology = md.Topology.from_openmm(sol_topology)
md_subset_topology = md.Topology.from_openmm(vac_topology)
pdf_state, pdf_state_subset, integrator, ani_handler, atom_map = generate_propagator_inputs(system = sol_system,
system_subset = vac_system,
md_topology = md_topology,
md_subset_topology = md_subset_topology)
particle = Particle(0)
box_vectors = sol_system.getDefaultPeriodicBoxVectors()
particle_state = OpenMMParticleState(positions = sol_positions, box_vectors = box_vectors)
particle.update_state(particle_state)
return pdf_state, pdf_state_subset, integrator, ani_handler, atom_map, particle
def test_OpenMMPDFState():
"""
conduct a class-wide test on coddiwomple.openmm.states.OpenMMPDFState with the `get_harmonic_testsystem` testsystem
this will assert successes on __init__, set_parameters, get_parameters, reduced_potential methods
"""
temperature = 300 * unit.kelvin
pressure = None
from coddiwomple.openmm.states import OpenMMPDFState, OpenMMParticleState
#create the default get_harmonic_testsystem
testsystem, period, collision_rate, timestep, alchemical_functions = get_harmonic_testsystem(temperature = temperature)
#test init method
pdf_state = OpenMMPDFState(system = testsystem.system, alchemical_composability = HarmonicAlchemicalState, temperature = temperature, pressure = pressure)
assert isinstance(pdf_state._internal_context, openmm.Context)
print(f"pdf_state parameters: {pdf_state._parameters}")
#test set_parameters
new_parameters = {key : 1.0 for key, val in pdf_state._parameters.items() if val is not None}
pdf_state.set_parameters(new_parameters) #this should set the new parameters, but now we have to make sure that the context actually has those parameters bound
swig_parameters = pdf_state._internal_context.getParameters()
context_parameters = {q: swig_parameters[q] for q in swig_parameters}
assert context_parameters['testsystems_HarmonicOscillator_x0'] == 1.
assert context_parameters['testsystems_HarmonicOscillator_U0'] == 1.
#test get_parameters
returnable_parameters = pdf_state.get_parameters()
assert len(returnable_parameters) == 2
assert returnable_parameters['testsystems_HarmonicOscillator_x0'] == 1.
assert returnable_parameters['testsystems_HarmonicOscillator_U0'] == 1.
#test reduced_potential
particle_state = OpenMMParticleState(positions = testsystem.positions) #make a particle state so that we can compute a reduced potential
reduced_potential = pdf_state.reduced_potential(particle_state)
externally_computed_reduced_potential = pdf_state._internal_context.getState(getEnergy=True).getPotentialEnergy()*pdf_state.beta
assert np.isclose(reduced_potential, externally_computed_reduced_potential)
def run(setup_dict):
"""
execute a Propagator
"""
import torchani
from simtk import unit
import sys
import numpy as np
import mdtraj as md
from coddiwomple.particles import Particle
from coddiwomple.openmm.states import OpenMMParticleState
from qmlify.utils import load_yaml, deserialize_xml, position_extractor, generate_propagator_inputs, depickle
#pull systems
system = deserialize_xml(setup_dict['system'])
system_subset = deserialize_xml(setup_dict['subset_system'])
#load topologies
md_topology = md.Topology.from_openmm(depickle(setup_dict['topology']))
md_subset_topology = md.Topology.from_openmm(
depickle(setup_dict['subset_topology']))
#load positions and box vectors
positions, box_vectors = position_extractor(
positions_cache_filename=setup_dict['positions_cache_filename'],
index_to_extract=setup_dict['position_extraction_index'])
positions *= unit.nanometers
if box_vectors is not None: box_vectors *= unit.nanometers
#integrator integrator_kwargs
default_integrator_kwargs = {
'temperature': 300.0 * unit.kelvin,
'collision_rate': 1.0 / unit.picoseconds,
'timestep': 1.0 * unit.femtoseconds,
'splitting': "V R O R F",
'constraint_tolerance': 1e-6,
'pressure': 1.0 * unit.atmosphere
}
if 'integrator_kwargs' in setup_dict.keys():
integrator_kwargs = setup_dict['integrator_kwargs']
if integrator_kwargs is not None:
if 'temperature' in integrator_kwargs.keys():
integrator_kwargs['temperature'] *= unit.kelvin
if 'collision_rate' in integrator_kwargs.keys():
integrator_kwargs['collision_rate'] /= unit.picoseconds
if 'timestep' in integrator_kwargs.keys():
integrator_kwargs['timestep'] *= unit.femtoseconds
if 'pressure' in integrator_kwargs.keys(
) and integrator_kwargs['pressure'] is not None:
integrator_kwargs['pressure'] *= unit.atmosphere
default_integrator_kwargs.update(integrator_kwargs)
pdf_state, pdf_state_subset, integrator, ani_handler, atom_map = generate_propagator_inputs(
system=system,
system_subset=system_subset,
md_topology=md_topology,
md_subset_topology=md_subset_topology,
ani_model=torchani.models.ANI2x(),
integrator_kwargs=default_integrator_kwargs)
if setup_dict['direction'] == 'forward':
from qmlify.propagation import Propagator
prop = Propagator
elif setup_dict['direction'] == 'backward':
from qmlify.propagation import BackwardPropagator
prop = BackwardPropagator
elif setup_dict['direction'] == 'ani_endstate':
from qmlify.propagation import ANIPropagator
prop = ANIPropagator
else:
raise Exception(
f"{setup_dict['direction']} is not valid. allowed directions are 'forward', 'backward', 'ani_endstate'"
)
propagator = prop(openmm_pdf_state=pdf_state,
openmm_pdf_state_subset=pdf_state_subset,
subset_indices_map=atom_map,
integrator=integrator,
ani_handler=ani_handler,
context_cache=None,
reassign_velocities=True,
n_restart_attempts=0)
particle = Particle(0)
particle_state = OpenMMParticleState(positions=positions,
box_vectors=box_vectors)
particle.update_state(particle_state)
particle_state, _return_dict = propagator.apply(
particle_state,
n_steps=setup_dict['num_steps'],
reset_integrator=True,
apply_pdf_to_context=True)
if box_vectors is None:
particle_state.box_vectors = None
work_array = np.array(propagator.state_works[0])
if particle_state.box_vectors is not None:
np.savez(setup_dict['out_positions_npz'],
positions=np.array([
particle_state.positions.value_in_unit_system(
unit.md_unit_system)
]),
box_vectors=np.array([
particle_state.box_vectors.value_in_unit_system(
unit.md_unit_system)
]))
else:
np.savez(setup_dict['out_positions_npz'],
positions=particle_state.positions.value_in_unit_system(
unit.md_unit_system))
np.savez(setup_dict['out_works_npz'], works=work_array)
def ANI_endstate_sampler(
system,
system_subset,
subset_indices_map,
positions_cache_filename,
md_topology,
index_to_run,
steps_per_application,
integrator_kwargs = {'temperature': 300.0 * unit.kelvin,
'collision_rate': 1.0 / unit.picoseconds,
'timestep': 2.0 * unit.femtoseconds,
'splitting': "V R O R F",
'constraint_tolerance': 1e-6,
'pressure': 1.0 * unit.atmosphere},
position_extractor = None
):
"""
conduct ani endstate sampling
arguments
system : openmm.System
system
system_subset : openmm.System
subset system
subset_indices_map : dict
dict of {openmm_pdf_state atom_index : openmm_pdf_state_subset atom index}
positions_cache_filename : str
path to the cache positions
index_to_run : int
index of the positions to anneal
number_of_applications : int
number of applications of the propagator
steps_per_application : int
number of integration steps per application
integrator_kwargs : dict, see default
kwargs to pass to OMMLIAIS integrator
"""
from coddiwomple.particles import Particle
from coddiwomple.openmm.states import OpenMMParticleState, OpenMMPDFState
#load the endstate cache
traj = np.load(positions_cache_filename)
positions = traj['positions'][index_to_run,:,:] * unit.nanometers
if position_extractor is not None:
positions = position_extractor(_positions)
try:
box_vectors = traj['box_vectors'][index_to_run,:,:] * unit.nanometers
except Exception as e:
box_vectors = None
species_str = ''.join([atom.element.symbol for atom in md_topology.subset(list(subset_indices_map.keys())).atoms])
_logger.info(f"species string: {species_str}")
ani_handler = ANI1_force_and_energy(model = torchani.models.ANI1ccx(),
atoms=species_str,
platform='cpu',
temperature=integrator_kwargs['temperature'])
#make thermostates
pressure = integrator_kwargs['pressure'] if box_vectors is not None else None
pdf_state = ThermodynamicState(system = system, temperature = integrator_kwargs['temperature'], pressure=pressure)
pdf_state_subset = ThermodynamicState(system = system_subset, temperature = integrator_kwargs['temperature'], pressure = None)
#make an integrator
integrator = Integrator(**integrator_kwargs)
#make a propagator
propagator = ANIPropagator(openmm_pdf_state = pdf_state,
openmm_pdf_state_subset = pdf_state_subset,
subset_indices_map = subset_indices_map,
integrator = integrator,
ani_handler = ani_handler,
context_cache=None,
reassign_velocities=True,
n_restart_attempts=0,
reporter = None)
particle = Particle(0)
particle_state = OpenMMParticleState(positions = positions, box_vectors = box_vectors)
particle.update_state(particle_state)
particle_state, _return_dict = propagator.apply(particle_state, n_steps = steps_per_application, reset_integrator=True, apply_pdf_to_context=True)
if box_vectors is None:
particle_state.box_vectors=None
return particle_state, np.array(propagator.state_works[0])
def annealed_importance_sampling(direction,
system,
system_subset,
subset_indices_map,
positions_cache_filename,
index_to_run,
directory_name,
trajectory_prefix,
md_topology,
steps_per_application,
integrator_kwargs = {'temperature': 300.0 * unit.kelvin,
'collision_rate': 1.0 / unit.picoseconds,
'timestep': 1.0 * unit.femtoseconds,
'splitting': "V R O R F",
'constraint_tolerance': 1e-6,
'pressure': 1.0 * unit.atmosphere},
save_indices = None,
position_extractor = None,
write_trajectory_interval=1
):
"""
conduct annealed importance sampling in the openmm regime; will write the accumulated work dictionary after each application
arguments
direction : str
forward or backward
system : openmm.System
system
system_subset : openmm.System
subset system
subset_indices_map : dict
dict of {openmm_pdf_state atom_index : openmm_pdf_state_subset atom index}
positions_cache_filename : str
path to the cache positions
index_to_run : int
index of the positions to anneal
directory_name : str
directory that will be written to
trajectory_prefix : str
.pdb prefix
md_topology : mdtraj.Topology
topology that will write the trajectory
save_indices : list
list of indices that will be saved
number_of_applications : int
number of applications of the propagator
steps_per_application : int
number of integration steps per application
integrator_kwargs : dict, see default
kwargs to pass to OMMLIAIS integrator
save_indices : list(int)
list of indices of md_topology atoms to save to disk
position_extractor : function, default None
function to extract appropriate positons from the cache
write_trajectory_interval : int
frequency with which to write trajectory to disk
"""
from coddiwomple.particles import Particle
from coddiwomple.openmm.states import OpenMMParticleState, OpenMMPDFState
#load the endstate cache
traj = np.load(positions_cache_filename)
positions = traj['positions'][index_to_run,:,:] * unit.nanometers
if position_extractor is not None:
positions = position_extractor(_positions)
try:
box_vectors = traj['box_vectors'][index_to_run,:,:] * unit.nanometers
except Exception as e:
box_vectors = None
assert direction in ['forward', 'backward']
#make a handle object for ANI
species_str = ''.join([atom.element.symbol for atom in md_topology.subset(list(subset_indices_map.keys())).atoms])
_logger.info(f"species string: {species_str}")
ani_handler = ANI1_force_and_energy(model = torchani.models.ANI1ccx(),
atoms=species_str,
platform='cpu',
temperature=integrator_kwargs['temperature'])
#make thermostates
pressure = integrator_kwargs['pressure'] if box_vectors is not None else None
pdf_state = ThermodynamicState(system = system, temperature = integrator_kwargs['temperature'], pressure=pressure)
pdf_state_subset = ThermodynamicState(system = system_subset, temperature = integrator_kwargs['temperature'], pressure = None)
#make a reporter
saveable_topology = md_topology.subset(save_indices)
reporter = OpenMMReporter(directory_name, trajectory_prefix, saveable_topology, subset_indices = save_indices)
#make an integrator
integrator = Integrator(**integrator_kwargs)
#make a propagator
if direction == 'forward':
propagator = Propagator(openmm_pdf_state = pdf_state,
openmm_pdf_state_subset = pdf_state_subset,
subset_indices_map = subset_indices_map,
integrator = integrator,
ani_handler = ani_handler,
context_cache=None,
reassign_velocities=True,
n_restart_attempts=0,
reporter = reporter,
write_trajectory_interval = write_trajectory_interval)
else:
propagator = BackwardPropagator(openmm_pdf_state = pdf_state,
openmm_pdf_state_subset = pdf_state_subset,
subset_indices_map = subset_indices_map,
integrator = integrator,
ani_handler = ani_handler,
context_cache=None,
reassign_velocities=True,
n_restart_attempts=0,
reporter = reporter,
write_trajectory_interval = write_trajectory_interval)
particle = Particle(0)
particle_state = OpenMMParticleState(positions = positions, box_vectors = box_vectors)
particle.update_state(particle_state)
particle_state, _return_dict = propagator.apply(particle_state, n_steps = steps_per_application, reset_integrator=True, apply_pdf_to_context=True)
if box_vectors is None:
particle_state.box_vectors=None
return particle_state, np.array(propagator.state_works[0])
def endstate_equilibration(system,
endstate_positions,
box_vectors,
directory_name,
trajectory_prefix,
md_topology,
number_of_applications,
steps_per_application,
endstate_parameters = 0.0,
alchemical_composability = RelativeAlchemicalState):
"""
conduct an endstate equilibration with pdf_state parameters defined by `endstate_parameters` with decorrelation
that will be written to disk
arguments
system : openmm.System
parameterizable system
endstate_positions : np.ndarray(N,3) * unit.nanometers (or length units)
starting positions that will be minimized and simulated
box_vectors : np.ndarray(3,3) * unit.nanometers (or length units)
starting box vectors
directory_name : str
directory that will be written to
trajectory_prefix : str
.pdb prefix
md_topology : mdtraj.Topology
topology that will write the trajectory
number_of_applications : int
number of applications of the propagator
steps_per_application : int
number of integration steps per application
endstate_parameters : float
endstate parameters
alchemical_composability : openmmtools.alchemy.AlchemicalState
composer for alchemical composability creation
"""
from perses.dispersed.feptasks import minimize
from perses.dispersed.utils import compute_timeseries
import mdtraj.utils as mdtrajutils
#determine pressure
forces = {type(force).__name__: force for force in system.getForces()}
if "MonteCarloBarostat" in list(forces.keys()):
pressure = 1.0 * unit.atmosphere
else:
pressure = None
print(f"pressure: {pressure}")
particle_state = OpenMMParticleState(positions = endstate_positions,
box_vectors = np.array(system.getDefaultPeriodicBoxVectors()),
)
pdf_state = OpenMMPDFState(system = system,
alchemical_composability = alchemical_composability,
pressure=pressure)
#set the pdf_state endstate parameters
pdf_state_parameters = pdf_state.get_parameters()
reset_pdf_state_parameters = {key: endstate_parameters for key in pdf_state_parameters.keys()}
pdf_state.set_parameters(reset_pdf_state_parameters)
langevin_integrator = OMMLI(temperature = pdf_state.temperature,
timestep = 2.0 * unit.femtoseconds)
endstate_propagator = OMMBIP(openmm_pdf_state = pdf_state,
integrator = langevin_integrator)
reporter = OpenMMReporter(directory_name, trajectory_prefix, md_topology = md_topology)
particle = Particle(0)
particle.update_state(particle_state)
reporter.record([particle])
minimize(endstate_propagator.pdf_state, particle_state)
potentials = []
for i in tqdm.trange(number_of_applications):
particle_state, _return_dict = endstate_propagator.apply(particle_state, n_steps = steps_per_application)
reporter.record([particle])
potentials.append(_return_dict['new_pe'])
[t0, g, Neff_max, A_t, uncorrelated_indices] = compute_timeseries(np.array(potentials))
print(f"t0: {t0}, g: {g}, Neff_max: {Neff_max}, A_t: {A_t}, uncorrelated_indices:{uncorrelated_indices}")
particle_hex = hex(id(particle))
reporter.hex_dict[particle_hex] = [[reporter.hex_dict[particle_hex][0][q] for q in uncorrelated_indices],
[reporter.hex_dict[particle_hex][1][q] for q in uncorrelated_indices],
[reporter.hex_dict[particle_hex][2][q] for q in uncorrelated_indices]]
try:
reporter._write_trajectory(particle_hex, f"{reporter.neq_traj_filename}.{format(reporter.hex_to_index[particle_hex], '04')}.pdb")
except Exception as e:
print(e)
def mcmc_smc_resampler(system,
endstate_cache_filename,
directory_name,
trajectory_prefix,
md_topology,
number_of_particles,
parameter_sequence,
observable = nESS,
threshold = vanilla_floor_threshold,
alchemical_composability = RelativeAlchemicalState,
integrator_kwargs = {'temperature': 300.0 * unit.kelvin,
'collision_rate': 1.0 / unit.picoseconds,
'timestep': 2.0 * unit.femtoseconds,
'splitting': "V R O R V",
'constraint_tolerance': 1e-6},
record_state_work_interval = 1,
trajectory_write_interval = 1):
"""
Conduct a single-direction annealing protocol in the AIS regime (wherein particle incremental weights are computed according to Eq. 31 of https://www.stats.ox.ac.uk/~doucet/delmoral_doucet_jasra_sequentialmontecarlosamplersJRSSB.pdf)
with resampling.
arguments
system : openmm.System
parameterizable system
endstate_cache_filename : str
path to the endstate cache pdb
directory_name : str
directory that will be written to
trajectory_prefix : str
.pdb prefix
md_topology : mdtraj.Topology
topology that will write the trajectory
number_of_particles : int
number of particles that will be used
parameter_sequence : list of dict
list of dictionary of parameters to be set
observable : function
function to compute observables on particles
threshold : function
function to compute a threshold on an observable
alchemical_composability : openmmtools.alchemy.AlchemicalState
composer for alchemical composability creation
integrator_kwargs : dict, see default
kwargs to pass to OMMLIAIS integrator
"""
pressure = handle_pressure(system)
print(f"pressure: {pressure}")
traj = md.Trajectory.load(endstate_cache_filename)
num_frames = traj.n_frames
proposal_pdf_state = OpenMMPDFState(system = system, alchemical_composability = alchemical_composability, pressure=pressure)
target_pdf_state = OpenMMPDFState(system = system, alchemical_composability = alchemical_composability, pressure = pressure)
reporter = OpenMMReporter(directory_name, trajectory_prefix, md_topology = md_topology)
integrator = OMMLI(**integrator_kwargs)
propagator = OMMBIP(proposal_pdf_state, integrator)
proposal_factory = ProposalFactory(parameter_sequence, propagator)
target_factory = TargetFactory(target_pdf_state, parameter_sequence, termination_parameters = None)
resampler = MultinomialResampler()
particles = [Particle(index = idx) for idx in range(number_of_particles)]
for idx in range(number_of_particles):
random_frame = np.random.choice(range(num_frames))
positions = traj.xyz[random_frame] * unit.nanometers
box_vectors = traj.unitcell_vectors[random_frame]*unit.nanometers
particle_state = OpenMMParticleState(positions=positions, box_vectors = box_vectors)
proposal_factory.equip_initial_sample(particle = particles[idx],
initial_particle_state = particle_state,
generation_pdf = None)
reporter.record(particles)
while True: #be careful...
try:
[particle.update_iteration() for particle in particles]#update the iteration
incremental_works = np.array([target_factory.compute_incremental_work(particle, neglect_proposal_work = True) for particle in particles]) #compute incremental works
randomize_velocities = True if particles[0].iteration == 1 else False #whether to randomize velocities in the first iteration
[proposal_factory.propagate(particle, randomize_velocities=randomize_velocities, apply_pdf_to_context=True) for particle in particles] #propagate particles (be sure to save to context)
resampler.resample(particles, incremental_works, observable = observable, threshold = threshold, update_particle_indices = True) #attempt to resample particles
#ask to terminate
if all(target_factory.terminate(particle) for particle in particles):
reporter.record(particles, save_to_disk=True)
break
else:
#update the auxiliary works...
[particle.zero_auxiliary_work() for particle in particles]
[particle.update_auxiliary_work(-target_factory.pdf_state.reduced_potential(particle.state)) for particle in particles]
reporter.record(particles)
except Exception as e:
print(f"error raised in iteration {particles[0].iteration}: {e}")
break
return particles
def annealed_importance_sampling(system,
endstate_cache_filename,
directory_name,
trajectory_prefix,
md_topology,
alchemical_functions,
number_of_applications,
steps_per_application,
endstate_parameters,
alchemical_composability = RelativeAlchemicalState,
integrator_kwargs = {'temperature': 300.0 * unit.kelvin,
'collision_rate': 1.0 / unit.picoseconds,
'timestep': 2.0 * unit.femtoseconds,
'splitting': "P I H N S V R O R V",
'constraint_tolerance': 1e-6},
record_state_work_interval = 1,
trajectory_write_interval = 1,
):
"""
conduct annealed importance sampling in the openmm regime
arguments
system : openmm.System
parameterizable system
endstate_cache_filename : str
path to the endstate cache pdb
directory_name : str
directory that will be written to
trajectory_prefix : str
.pdb prefix
md_topology : mdtraj.Topology
topology that will write the trajectory
alchemical_functions : dict
{pdf_parameter <str>: function <str>, lepton-readable}
number_of_applications : int
number of applications of the propagator
steps_per_application : int
number of integration steps per application
endstate_parameters : float
endstate parameters
alchemical_composability : openmmtools.alchemy.AlchemicalState
composer for alchemical composability creation
integrator_kwargs : dict, see default
kwargs to pass to OMMLIAIS integrator
"""
#determine pressure
pressure = handle_pressure(system)
print(f"pressure: {pressure}")
traj = md.Trajectory.load(endstate_cache_filename)
num_frames = traj.n_frames
pdf_state = OpenMMPDFState(system = system, alchemical_composability = alchemical_composability, pressure=pressure)
#set the pdf_state endstate parameters
pdf_state_parameters = pdf_state.get_parameters()
reset_pdf_state_parameters = {key: endstate_parameters for key in pdf_state_parameters.keys()}
pdf_state.set_parameters(reset_pdf_state_parameters)
reporter = OpenMMReporter(directory_name, trajectory_prefix, md_topology = md_topology)
ais_integrator = OMMLIAIS(
alchemical_functions,
steps_per_application,
**integrator_kwargs)
ais_propagator = OMMAISPR(openmm_pdf_state = pdf_state,
integrator = ais_integrator,
record_state_work_interval = record_state_work_interval,
reporter = reporter,
trajectory_write_interval = trajectory_write_interval,
context_cache=None,
reassign_velocities=True,
n_restart_attempts=0)
frames = np.random.choice(range(num_frames), number_of_applications)
particle = Particle(0)
for i in tqdm.trange(number_of_applications):
pdf_state.set_parameters(reset_pdf_state_parameters)
particle_state = OpenMMParticleState(positions = traj.xyz[frames[i]] * unit.nanometers, box_vectors = traj.unitcell_vectors[frames[i]]*unit.nanometers)
particle.update_state(particle_state)
try:
_, _return_dict = ais_propagator.apply(particle_state, n_steps = steps_per_application, reset_integrator=True, apply_pdf_to_context=True)
except Exception as e:
print(e)
return ais_propagator.state_works
def test_OMMLI():
"""
test OMMLI (OpenMMLangevinIntegrator) in the baoab regime on the harmonic test system;
Specifically, we run MD to convergence and assert that the potential energy of the system and the standard
deviation thereof is within a specified threshold.
We also check the accumulation of shadow, proposal works, as well as the ability to reset, initialize, and subsume the integrator into an OMMBIP propagator
"""
from coddiwomple.openmm.propagators import OMMBIP
from coddiwomple.openmm.integrators import OMMLI
import tqdm
temperature = 300 * unit.kelvin
pressure = None
from coddiwomple.openmm.states import OpenMMPDFState, OpenMMParticleState
from coddiwomple.particles import Particle
#create the default get_harmonic_testsystem
testsystem, period, collision_rate, timestep, alchemical_functions = get_harmonic_testsystem(temperature = temperature)
particle_state = OpenMMParticleState(positions = testsystem.positions) #make a particle state
particle = Particle(index = 0, record_state=False, iteration = 0)
particle.update_state(particle_state)
num_applications = 100
#create a pdf_state
pdf_state = OpenMMPDFState(system = testsystem.system, alchemical_composability = HarmonicAlchemicalState, temperature = temperature, pressure = pressure)
#create an integrator
integrator = OMMLI(temperature=temperature, collision_rate=collision_rate, timestep=timestep)
#create a propagator
propagator = OMMBIP(openmm_pdf_state = pdf_state, integrator = integrator)
#expected reduced potential
mean_reduced_potential = testsystem.get_potential_expectation(pdf_state) * pdf_state.beta
std_dev_reduced_potential = testsystem.get_potential_standard_deviation(pdf_state) * pdf_state.beta
reduced_pe = []
#some sanity checks for propagator:
global_integrator_variables_before_integration = propagator._get_global_integrator_variables()
print(f"starting integrator variables: {global_integrator_variables_before_integration}")
#some sanity checks for integrator:
start_proposal_work = propagator.integrator._get_energy_with_units('proposal_work', dimensionless=True)
start_shadow_work = propagator.integrator._get_energy_with_units('shadow_work', dimensionless=True)
assert start_proposal_work == global_integrator_variables_before_integration['proposal_work']
assert start_shadow_work == global_integrator_variables_before_integration['shadow_work']
for app_num in tqdm.trange(num_applications):
particle_state, proposal_work = propagator.apply(particle_state, n_steps=20, reset_integrator=False, apply_pdf_to_context=False, randomize_velocities=True)
assert proposal_work==0. #this must be the case since we did not pass a 'returnable_key'
#sanity checks for inter-application methods
assert propagator.integrator._get_energy_with_units('proposal_work', dimensionless=True) != 0. #this cannot be zero after a step of MD without resets
assert propagator.integrator._get_energy_with_units('shadow_work', dimensionless=True) != 0. #this cannot be zero after a step of MD without resets
reduced_pe.append(pdf_state.reduced_potential(particle_state))
tol=6 * std_dev_reduced_potential
calc_mean_reduced_pe = np.mean(reduced_pe)
calc_stddev_reduced_pe = np.std(reduced_pe)
assert calc_mean_reduced_pe < mean_reduced_potential + tol and calc_mean_reduced_pe > mean_reduced_potential - tol, f"the mean reduced energy and standard deviation ({calc_mean_reduced_pe}, {calc_stddev_reduced_pe}) is outside the tolerance \
of a theoretical mean potential energy of {mean_reduced_potential} +/- {tol}"
print(f"the mean reduced energy/standard deviation is {calc_mean_reduced_pe, calc_stddev_reduced_pe} and the theoretical mean reduced energy and stddev are {mean_reduced_potential}")
#some cleanup of the integrator
propagator.integrator.reset() #this should reset proposal, shadow, and ghmc staticstics (we omit ghmc stats)
assert propagator.integrator._get_energy_with_units('proposal_work', dimensionless=True) == 0. #this should be zero after a reset
assert propagator.integrator._get_energy_with_units('shadow_work', dimensionless=True) == 0. #this should be zero after a reset