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