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