Source code for openpathsampling.engines.openmm.tools

import numpy as np

from openpathsampling.integration_tools import (
    openmm, md, error_if_no_mdtraj, unit, error_if_no_simtk_unit, HAS_OPENMM
)

if HAS_OPENMM:
    # this is in case we directly import tools (e.g., for
    # trajectory_to/from_mdtraj) when we don't have OpenMM installed. In
    # that case, we skip these imports (engines/openmm/__init__.py prevents
    # them from being made)
    from .snapshot import Snapshot
    from openpathsampling.engines.topology import Topology, MDTrajTopology
    try:
        # openmm >= 7.6
        from openmm.app import internal as _internal
    except ImportError:
        # openmm < 7.6
        from simtk.openmm.app import internal as _internal

    reducePeriodicBoxVectors = _internal.unitcell.reducePeriodicBoxVectors


from openpathsampling.engines import Trajectory, NoEngine, SnapshotDescriptor


__author__ = 'Jan-Hendrik Prinz'


class TopologyEngine(NoEngine):
    _default_options = {}

    def __init__(self, topology):

        descriptor = SnapshotDescriptor.construct(
            Snapshot,
            {
                'n_atoms': topology.n_atoms,
                'n_spatial': topology.n_spatial
            }
        )

        super(NoEngine, self).__init__(
            descriptor=descriptor
        )

        self.topology = topology

    @property
    def mdtraj_topology(self):
        return self.topology.mdtraj

    def to_dict(self):
        return {
            'topology': self.topology,
        }


class FileEngine(TopologyEngine):

    _default_options = {}

    def __init__(self, topology, filename):
        super(FileEngine, self).__init__(
            topology=topology
        )

        self.filename = filename

    def to_dict(self):
        return {
            'topology': self.topology,
            'filename': self.filename
        }


class OpenMMToolsTestsystemEngine(TopologyEngine):

    _default_options = {}

    def __init__(self, topology, testsystem_name):

        super(OpenMMToolsTestsystemEngine, self).__init__(
            topology=topology
        )

        self.testsystem_name = testsystem_name

    def to_dict(self):
        return {
            'topology': self.topology,
            'testsystem_name': self.testsystem_name
        }


[docs]def snapshot_from_pdb(pdb_file, simple_topology=False): """ Construct a Snapshot from the first frame in a pdb file without velocities Parameters ---------- pdb_file : str The filename of the .pdb file to be used simple_topology : bool if `True` only a simple topology with n_atoms will be created. This cannot be used with complex CVs but loads and stores very fast Returns ------- :class:`openpathsampling.engines.Snapshot` the constructed Snapshot """ snap = ops_load_trajectory(pdb_file)[0] if simple_topology: topology = Topology(*pdb.xyz[0].shape) else: topology = snap.topology snapshot = Snapshot.construct( coordinates=snap.coordinates, box_vectors=snap.box_vectors, velocities=snap.velocities, engine=FileEngine(topology, pdb_file) ) return snapshot
def topology_from_pdb(pdb_file, simple_topology=False): """ Construct a Topology from the first frame in a pdb file without velocities Parameters ---------- pdb_file : str The filename of the .pdb file to be used simple_topology : bool if `True` only a simple topology with n_atoms will be created. This cannot be used with complex CVs but loads and stores very fast Returns ------- :class:`openpathsampling.engines.Snapshot` the constructed Snapshot """ pdb = md.load(pdb_file) if simple_topology: topology = Topology(*pdb.xyz[0].shape) else: topology = MDTrajTopology(pdb.topology) return topology
[docs]def snapshot_from_testsystem(testsystem, simple_topology=False, periodic=True): """ Construct a Snapshot from openmm topology and state objects Parameters ---------- testsystem : openmmtools.Topology The filename of the .pdb file to be used simple_topology : bool if `True` only a simple topology with n_atoms will be created. This cannot be used with complex CVs but loads and stores very fast periodic : bool True (default) if system is periodic; if False, box vectors are None Returns ------- :class:`openpathsampling.engines.Snapshot` the constructed Snapshot """ error_if_no_simtk_unit("snapshot_from_testsystem") u_nm = unit.nanometers u_ps = unit.picoseconds velocities = unit.Quantity(np.zeros(testsystem.positions.shape), u_nm / u_ps) if simple_topology: topology = Topology(*testsystem.positions.shape) else: topology = MDTrajTopology(md.Topology.from_openmm(testsystem.topology)) if periodic: sys_box_vectors = testsystem.system.getDefaultPeriodicBoxVectors() box_vectors = np.array([v / u_nm for v in sys_box_vectors]) * u_nm else: box_vectors = None snapshot = Snapshot.construct( coordinates=testsystem.positions, box_vectors=box_vectors, velocities=velocities, engine=OpenMMToolsTestsystemEngine(topology, testsystem.name) ) return snapshot
[docs]def trajectory_from_mdtraj(mdtrajectory, simple_topology=False, velocities=None): """ Construct a Trajectory object from an mdtraj.Trajectory object Parameters ---------- mdtrajectory : mdtraj.Trajectory Input mdtraj.Trajectory simple_topology : bool if `True` only a simple topology with n_atoms will be created. This cannot be used with complex CVs but loads and stores very fast velocities : np.array velocities in units of nm/ps Returns ------- openpathsampling.engines.Trajectory the constructed Trajectory instance """ error_if_no_simtk_unit("trajectory_from_mdtraj") trajectory = Trajectory() u_nm = unit.nanometer u_ps = unit.picosecond vel_unit = u_nm / u_ps if simple_topology: topology = Topology(*mdtrajectory.xyz[0].shape) else: topology = MDTrajTopology(mdtrajectory.topology) if velocities is None: empty_vel = unit.Quantity(np.zeros(mdtrajectory.xyz[0].shape), vel_unit) if mdtrajectory.unitcell_vectors is not None: box_vects = unit.Quantity(mdtrajectory.unitcell_vectors, unit.nanometers) else: box_vects = [None] * len(mdtrajectory) engine = TopologyEngine(topology) for frame_num in range(len(mdtrajectory)): # mdtraj trajectories only have coordinates and box_vectors coord = unit.Quantity(mdtrajectory.xyz[frame_num], u_nm) if velocities is not None: vel = unit.Quantity(velocities[frame_num], vel_unit) else: vel = empty_vel box_v = box_vects[frame_num] statics = Snapshot.StaticContainer( coordinates=coord, box_vectors=box_v, engine=engine ) kinetics = Snapshot.KineticContainer(velocities=vel, engine=engine) snap = Snapshot( statics=statics, kinetics=kinetics, engine=engine ) trajectory.append(snap) return trajectory
[docs]def empty_snapshot_from_openmm_topology(topology, simple_topology=False): """ Return an empty snapshot from an openmm.Topology object Velocities will be set to zero. Parameters ---------- topology : openmm.Topology the topology representing the structure and number of atoms simple_topology : bool if `True` only a simple topology with n_atoms will be created. This cannot be used with complex CVs but loads and stores very fast Returns ------- openpathsampling.engines.Snapshot the complete snapshot with zero coordinates and velocities """ error_if_no_simtk_unit("empty_snapshot_from_openmm_topology") u_nm = unit.nanometers u_ps = unit.picoseconds n_atoms = topology.n_atoms if simple_topology: topology = Topology(n_atoms, 3) else: error_if_no_mdtraj("empty_snaphsot_from_openmm_topology") topology = MDTrajTopology(md.Topology.from_openmm(topology)) snapshot = Snapshot.construct( coordinates=unit.Quantity(np.zeros((n_atoms, 3)), u_nm), box_vectors=unit.Quantity(topology.setUnitCellDimensions(), u_nm), velocities=unit.Quantity(np.zeros((n_atoms, 3)), u_nm / u_ps), engine=TopologyEngine(topology) ) return snapshot
[docs]def to_openmm_topology(obj): """ Contruct an openmm.Topology file out of a Snapshot or Configuration object. This uses the mdtraj.Topology in the Configuration as well as the box_vectors. Parameters ---------- obj : openpathsampling.engines.BaseSnapshot or Configuration the object to be used in the topology construction Returns ------- openmm.Topology an object representing an openmm.Topology """ if obj.topology is not None: if hasattr(obj.topology, 'mdtraj'): openmm_topology = obj.topology.mdtraj.to_openmm() box_size_dimension = np.linalg.norm( obj.box_vectors.value_in_unit(unit.nanometer), axis=1) openmm_topology.setUnitCellDimensions(box_size_dimension) return openmm_topology else: return None
def trajectory_to_mdtraj(trajectory, md_topology=None): """ Construct a `mdtraj.Trajectory` object from an :obj:`Trajectory` object Parameters ---------- trajectory : :obj:`openpathsampling.engines.Trajectory` Input Trajectory Returns ------- :obj:`mdtraj.Trajectory` the constructed Trajectory instance """ if not hasattr(trajectory, 'to_mdtraj'): try: _ = len(trajectory) except TypeError: trajectory = Trajectory([trajectory]) else: trajectory = Trajectory(trajectory) # For now, let's keep all the code in one place, and better for # engines.openmm.tools to require engines.trajectory than vice versa return trajectory.to_mdtraj(md_topology) def ops_load_trajectory(filename, **kwargs): error_if_no_mdtraj("ops_load_trajectory") return trajectory_from_mdtraj(md.load(filename, **kwargs)) # ops_load_trajectory and the mdtraj stuff is not OpenMM-specific def reduced_box_vectors(snapshot): """Reduced box vectors for a snapshot (with units) See also -------- reduce_trajectory_box_vectors Parameters ---------- snapshot : :class:`.Snapshot` input snapshot Returns ------- :class:`.Snapshot` snapshot with correctly reduced box vectors """ nm = unit.nanometer return np.array( reducePeriodicBoxVectors(snapshot.box_vectors).value_in_unit(nm) ) * nm def reduce_trajectory_box_vectors(trajectory): """Trajectory with reduced box vectors. OpenMM has strict requirements on the box vectors describing the unit cell. In some cases, such as trajectories loaded from files that have rounded the box vectors, these conditions might not be satisfied. This method forces the box vectors to meet OpenMM's criteria. Parameters ---------- trajectory : :class:`.Trajectory` input trajectory Returns ------- :class:`.Trajectory` trajectory with correctly reduced box vectors """ return Trajectory([ snap.copy_with_replacement(box_vectors=reduced_box_vectors(snap)) for snap in trajectory ]) def load_trr(trr_file, top, velocities=True): """Load a TRR file, ready for use as input to an OpenMMEngine. This is a single method to handle several peculiarities of both the TRR format (which rounds some values) and OpenMM (which has certain requirements of box vectors), plus the possibility that you'll want velocities. Parameters ---------- trr_file : string name of TRR file to load top : string name of topology (e.g., ``.gro``) file to use. See MDTraj documentation on md.load. velocities : bool whether to also load velocities from the TRR file; default ``True`` Return ------ :class:`.Trajectory` the OPS trajectory, with OpenMM-reduced box vectors and velocities (if requested) """ mdt = md.load(trr_file, top=top) trr = md.formats.TRRTrajectoryFile(trr_file) if velocities: vel = trr._read(n_frames=len(mdt), atom_indices=None, get_velocities=True)[5] else: vel = None traj = trajectory_from_mdtraj(mdt, velocities=vel) return reduce_trajectory_box_vectors(traj) def n_dofs_from_system(system): """Get the number of degrees of freedom from an Openmm System Parameters ---------- system : :class:`simtk.openmm.System` object describing the system Returns ------- int : number of degrees of freedom """ # dof calculation based on OpenMM's StateDataReporter n_spatial = 3 n_particles = system.getNumParticles() dofs_particles = sum([n_spatial for i in range(n_particles) if system.getParticleMass(i) > 0*unit.dalton]) dofs_constaints = system.getNumConstraints() dofs_motion_removers = 0 has_cm_motion_remover = any( type(system.getForce(i)) == openmm.CMMotionRemover for i in range(system.getNumForces()) ) if has_cm_motion_remover: dofs_motion_removers += 3 dofs = dofs_particles - dofs_constaints - dofs_motion_removers return dofs def has_constraints_from_system(system): """Get the number of degrees of freedom from an Openmm System Parameters ---------- system : :class:`openmm.System` object describing the system Returns ------- bool : whether there are constraints """ return system.getNumConstraints() > 0