Beispiel #1
0
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
Beispiel #2
0
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)
Beispiel #4
0
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
        )
Beispiel #5
0
    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]
                )