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 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 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])