def get_potential_energy(item, coordinates=None, platform_name='CUDA'): from simtk.openmm import LangevinIntegrator, Platform, Context from simtk import unit import numpy as np integrator = LangevinIntegrator(0.0 * unit.kelvin, 0.0 / unit.picoseconds, 2.0 * unit.femtoseconds) platform = Platform.getPlatformByName(platform_name) context = Context(item.system, integrator, platform) if coordinates is None: context.setPositions(item.coordinates) else: context.setPositions(coordinates) if item.box is not None: context.setPeriodicBoxVectors(item.box[0], item.box[1], item.box[2]) state = context.getState(getEnergy=True) potential_energy = state.getPotentialEnergy() return potential_energy
class OpenMMStateBuilder(StateBuilder): """Build an OpenMM "state" that can be sent to a device to simulate. """ def __init__(self, system, integrator=None): # if strings are passed in, assume that they are paths to # xml files on disk if isinstance(system, basestring): with open(system) as f: system = XmlSerializer.deserialize(f.read()) if isinstance(integrator, basestring): with open(integrator) as f: integrator = XmlSerializer.deserialize(f.read()) if integrator is None: # this integrator isn't really necessary, but it has to be something # for the openmm API to let us serialize the state integrator = VerletIntegrator(2 * femtoseconds) self.context = Context(system, integrator, Platform.getPlatformByName('Reference')) def build(self, trajectory): """Create a serialized XML state from the first frame in a trajectory Parameteters ------------ trajectory : mdtraj.trajectory.Trajectory The trajectory to take the frame from. We'll use both the the positions and the box vectors (if you're using periodic boundary conditions) """ periodic = False if trajectory.unitcell_vectors is not None: a, b, c = trajectory.unitcell_lengths[0] np.testing.assert_array_almost_equal(trajectory.unitcell_angles[0], np.ones(3) * 90) self.context.setPeriodicBoxVectors([a, 0, 0] * nanometers, [0, b, 0] * nanometers, [0, 0, c] * nanometers) periodic = True self.context.setPositions(trajectory.openmm_positions(0)) state = self.context.getState(getPositions=True, getVelocities=True, getForces=True, getEnergy=True, getParameters=True, enforcePeriodicBox=periodic) return XmlSerializer.serialize(state)
class OpenMMStateBuilder(StateBuilder): """Build an OpenMM "state" that can be sent to a device to simulate. """ def __init__(self, system, integrator=None): # if strings are passed in, assume that they are paths to # xml files on disk if isinstance(system, basestring): with open(system) as f: system = XmlSerializer.deserialize(f.read()) if isinstance(integrator, basestring): with open(integrator) as f: integrator = XmlSerializer.deserialize(f.read()) if integrator is None: # this integrator isn't really necessary, but it has to be something # for the openmm API to let us serialize the state integrator = VerletIntegrator(2*femtoseconds) self.context = Context(system, integrator, Platform.getPlatformByName('Reference')) def build(self, trajectory): """Create a serialized XML state from the first frame in a trajectory Parameteters ------------ trajectory : mdtraj.trajectory.Trajectory The trajectory to take the frame from. We'll use both the the positions and the box vectors (if you're using periodic boundary conditions) """ periodic = False if trajectory.unitcell_vectors is not None: a, b, c = trajectory.unitcell_lengths[0] np.testing.assert_array_almost_equal(trajectory.unitcell_angles[0], np.ones(3)*90) self.context.setPeriodicBoxVectors([a, 0, 0] * nanometers, [0, b, 0] * nanometers, [0, 0, c] * nanometers) periodic = True self.context.setPositions(trajectory.openmm_positions(0)) state = self.context.getState(getPositions=True, getVelocities=True, getForces=True, getEnergy=True, getParameters=True, enforcePeriodicBox=periodic) return XmlSerializer.serialize(state)
class SingleContext: """Mimics the MultiContext API but does not spawn worker processes. Parameters: ----------- n_workers : int Needs to be 1. system : simtk.openmm.System The system that contains all forces. integrator : simtk.openmm.Integrator An OpenMM integrator. platform_name : str The name of an OpenMM platform ('Reference', 'CPU', 'CUDA', or 'OpenCL') platform_properties : dict, optional A dictionary of platform properties. """ def __init__(self, n_workers, system, integrator, platform_name, platform_properties={}): """Set up workers and queues.""" from simtk.openmm import Platform, Context assert n_workers == 1 openmm_platform = Platform.getPlatformByName(platform_name) self._openmm_context = Context(system, integrator, openmm_platform, platform_properties) def evaluate( self, positions, box_vectors=None, evaluate_energy=True, evaluate_force=True, evaluate_positions=False, evaluate_path_probability_ratio=False, err_handling="warning", n_simulation_steps=0 ): """Compute energies and/or forces. Parameters: ----------- positions : numpy.ndarray The particle positions in nanometer; its shape is (batch_size, num_particles, 3). box_vectors : numpy.ndarray, optional The periodic box vectors in nanometer; its shape is (batch_size, 3, 3). If not specified, don't change the box vectors. evaluate_energy : bool, optional Whether to compute energies. evaluate_force : bool, optional Whether to compute forces. evaluate_positions : bool, optional Whether to return positions. evaluate_path_probability_ratio : bool, optional Whether to compute the log path probability ratio. Makes only sense for PathProbabilityIntegrator instances. _err_handling : str, optional How to handle infinite energies (one of {"warning", "ignore", "exception"}). n_simulation_steps : int, optional If > 0, perform a number of simulation steps and compute energy and forces for the resulting state. Returns: -------- energies : np.ndarray or None The energies in units of kilojoule/mole; its shape is (len(positions), ) forces : np.ndarray or None The forces in units of kilojoule/mole/nm; its shape is (len(positions), num_particles, 3) new_positions : np.ndarray or None The positions in units of nm; its shape is (len(positions), num_particles, 3) log_path_probability_ratio : np.ndarray or None The logarithmic path probability ratios; its shape is (len(positions), ) """ from simtk import unit assert box_vectors is None or len(box_vectors) == len(positions), \ "box_vectors and positions have to be the same length" box_vectors = [None for _ in positions] if box_vectors is None else box_vectors forces = np.zeros_like(positions) energies = np.zeros_like(positions[:,0,0]) new_positions = np.zeros_like(positions) log_path_probability_ratios = np.zeros_like(positions[:,0,0]) for i, (p, bv) in enumerate(zip(positions, box_vectors)): try: # initialize state self._openmm_context.setPositions(p) if bv is not None: self._openmm_context.setPeriodicBoxVectors(bv) log_path_probability_ratio = self._openmm_context.getIntegrator().step(n_simulation_steps) # compute energy and forces state = self._openmm_context.getState( getEnergy=evaluate_energy, getForces=evaluate_force, getPositions=evaluate_positions ) energy = state.getPotentialEnergy().value_in_unit(unit.kilojoule_per_mole) if evaluate_energy else None force = ( state.getForces(asNumpy=True).value_in_unit(unit.kilojoule_per_mole / unit.nanometer) if evaluate_force else None ) new_pos = state.getPositions().value_in_unit(unit.nanometers) if evaluate_positions else None energies[i] = energy if evaluate_energy else 0.0 forces[i,:,:] = force if evaluate_force else 0.0 new_positions[i,:,:] = new_pos if evaluate_positions else 0.0 log_path_probability_ratios[i] = log_path_probability_ratio if evaluate_path_probability_ratio else 0.0 except Exception as e: if err_handling == "warning": warnings.warn("Suppressed exception: {}".format(e)) elif err_handling == "exception": raise e return ( energies if evaluate_energy else None, forces if evaluate_force else None, new_positions if evaluate_positions else None, log_path_probability_ratios if evaluate_path_probability_ratio else None )
class Worker(mp.Process): """A worker process that computes energies in its own context. Parameters: ----------- task_queue : multiprocessing.Queue The queue that the MultiContext pushes tasks to. result_queue : multiprocessing.Queue The queue that the MultiContext receives results from. system : simtk.openmm.System The system that contains all forces. integrator : simtk.openmm.Integrator An OpenMM integrator. platform_name : str The name of an OpenMM platform ('Reference', 'CPU', 'CUDA', or 'OpenCL') platform_properties : dict A dictionary of platform properties. """ def __init__(self, task_queue, result_queue, system, integrator, platform_name, platform_properties): super(MultiContext.Worker, self).__init__() self._task_queue = task_queue self._result_queue = result_queue self._openmm_system = system self._openmm_integrator = pickle.loads( pickle.dumps(integrator)) self._openmm_platform_name = platform_name self._openmm_platform_properties = platform_properties self._openmm_context = None def run(self): """Run the process: set positions and compute energies and forces. Positions and box vectors are received from the task_queue in units of nanometers. Energies and forces are pushed to the result_queue in units of kJ/mole and kJ/mole/nm, respectively. """ from simtk import unit from simtk.openmm import Platform, Context # create the context # it is crucial to do that in the run function and not in the constructor # for some reason, the CPU platform hangs if the context is created in the constructor # see also https://github.com/openmm/openmm/issues/2602 openmm_platform = Platform.getPlatformByName(self._openmm_platform_name) self._openmm_context = Context( self._openmm_system, self._openmm_integrator, openmm_platform, self._openmm_platform_properties ) self._openmm_context.reinitialize(preserveState=True) # get tasks from the task queue for task in iter(self._task_queue.get, None): (index, positions, box_vectors, evaluate_energy, evaluate_force, evaluate_positions, evaluate_path_probability_ratio, err_handling, n_simulation_steps) = task try: # initialize state self._openmm_context.setPositions(positions) if box_vectors is not None: self._openmm_context.setPeriodicBoxVectors(box_vectors) log_path_probability_ratio = self._openmm_integrator.step(n_simulation_steps) # compute energy and forces state = self._openmm_context.getState( getEnergy=evaluate_energy, getForces=evaluate_force, getPositions=evaluate_positions ) energy = state.getPotentialEnergy().value_in_unit(unit.kilojoule_per_mole) if evaluate_energy else None forces = ( state.getForces(asNumpy=True).value_in_unit(unit.kilojoule_per_mole / unit.nanometer) if evaluate_force else None ) new_positions = state.getPositions().value_in_unit(unit.nanometers) if evaluate_positions else None except Exception as e: if err_handling == "warning": warnings.warn("Suppressed exception: {}".format(e)) elif err_handling == "exception": raise e # push energies and forces to the results queue self._result_queue.put( [index, energy, forces, new_positions, log_path_probability_ratio] )