serialization, docstrings, mdtraj integration

dwhswenson · Jul 28, 2021 · Jul 31, 2021 · Jul 31, 2021 · Jul 31, 2021 · Jul 31, 2021
commit 4290267c34e82162c8f17380000dac2b2e9d68d4
diff --git a/openpathsampling/engines/openmm/mcengine.py b/openpathsampling/engines/openmm/mcengine.py
@@ -1,9 +1,11 @@
 import openpathsampling as paths
 from openpathsampling.engines.openmm import Snapshot
+import copy
 
 try:
     from openmmtools import mcmc
     from openmmtools import states
+    import openmmtools
 except ImportError:
     HAS_OPENMMTOOLS=False
 else:
@@ -16,17 +18,51 @@ class EngineNotInitializedError(Exception):
                                                  "before using engine.")
 
 def snapshot_from_sampler_state(sampler_state, engine=None):
+    """Generate an OPS SnapShot from a OpenMMTools SamplerState
+
+    Parameters
+    ----------
+    sampler_state: openmmtools.states.SamplerState
+        The ``SamplerState`` containing the positions, velocities, and box
+        vectors to define the snapshot.
+    engine: :class:`.DynamicsEngine`, optional
+        The engine to associate with this snapshot; default is ``None``,
+        which will make some functionality (e.g, MDTraj integration)
+        impossible.
+    """
     return Snapshot.construct(
-        coordinates=sampler_state.positions.copy(),
-        velocities=sampler_state.velocities.copy(),
-        box_vectors=sampler_state.box_vectors.copy(),
+        coordinates=copy.copy(sampler_state.positions),
+        velocities=copy.copy(sampler_state.velocities),
+        box_vectors=copy.copy(sampler_state.box_vectors),
         engine=engine
     )
 
 class OpenMMToolsMCEngine(paths.engines.DynamicsEngine):
-    """Engine for the MCMC utilities in OpenMMTools
+    """Engine for the OpenMMTools MCMC package.
+
+    This takes user-provided ``thermodynamics_state`` and ``move``, and will
+    internally create a ``openmmtools.mcmc.MCMCSampler``, which will be used
+    to generate "trajectories."
+
+    Parameters
+    ----------
+    thermodynamic_state : openmmtools.states.thermodynamic_state
+        the thermodynamic state to sample
+    move : openmmtools.mcmc.MCMCMove
+        the Monte Carlo move to use to propagate the system
+    options : Dict
+        Dictionary with additional options. Allowed options are:
+
+            'n_steps_per_frame': int, default: 10
+                the number of individual steps per snapshot
+            'n_frames_max': int, default: 5000
+                maximum number of frames before a trajectory aborts
+
+    topology : :class:`.Topology`, optional
+        Topology information for this sysetm. If an :class:`.MDTrajTopology`
+        is provided, this will enable MDTraj integration.
     """
-    def __init__(self, thermodynamic_state, move, options):
+    def __init__(self, thermodynamic_state, move, options, topology=None):
         # we assume that we have standard n_atoms * 3 dimensions
         descriptor = paths.engines.SnapshotDescriptor.construct(
             Snapshot,
@@ -37,6 +73,7 @@ def __init__(self, thermodynamic_state, move, options):
         self.thermodynamic_state = thermodynamic_state
         self.move = move
         self.options = options
+        self.topology = topology
         # _sampler is the OpenMMTools MCMCSampler
         self._sampler = None
         self._current_snapshot = None
@@ -92,4 +129,26 @@ def generate_next_frame(self):
 
         return self.current_snapshot
 
+    def to_dict(self):
+        serialize = openmmtools.utils.serialize
+        return {
+            'thermodynamic_state': serialize(self.thermodynamic_state),
+            'move': serialize(self.move),
+            'options': self.options,
+            'topology': self.topology,
+        }
+
+    @classmethod
+    def from_dict(cls, dct):
+        deserialize = openmmtools.utils.deserialize
+        dct['thermodynamic_state'] = deserialize(dct['thermodynamic_state'])
+        dct['move'] = deserialize(dct['move'])
+        return cls(**dct)
+
+    @property
+    def mdtraj_topology(self):
+        if self.topology is None:
+            raise RuntimeError("Can not convert to MDTraj: no topology "
+                               "associated with this engine.")
+        return self.topology.mdtraj