Improvements to gmx engine tools

* GroFileEngine is now storable * snapshot_from_gro is in the `engines.gromacs` namespace
dwhswenson · Oct 14, 2020 · Jul 21, 2020 · Jul 23, 2020 · Aug 3, 2020 · Aug 6, 2020
commit 8ad72d2418b8845b475738454925f2777801f0d0
diff --git a/openpathsampling/engines/gromacs/__init__.py b/openpathsampling/engines/gromacs/__init__.py
@@ -6,7 +6,9 @@ def requires_mdtraj(*args, **kwargs):  # pragma: no cover
 except ImportError:
     Engine = requires_mdtraj
     ExternalMDSnapshot = requires_mdtraj
+    snapshot_from_gro = requires_mdtraj
 else:
     from .engine import GromacsEngine as Engine
     from .engine import ExternalMDSnapshot
+    from .engine import snapshot_from_gro
 
diff --git a/openpathsampling/engines/gromacs/engine.py b/openpathsampling/engines/gromacs/engine.py
@@ -30,33 +30,40 @@
     _debug_open_files, close_file_descriptors
 )
 
+class _GroFileEngine(ExternalEngine):
+    def __init__(self, gro):
+        self.gro = gro
+        traj = md.load(gro)
+        self.topology = MDTrajTopology(traj.topology)
+        n_atoms = self.topology.n_atoms
+        n_spatial = self.topology.n_spatial
+        descriptor = SnapshotDescriptor.construct(
+            snapshot_class=ExternalMDSnapshot,
+            snapshot_dimensions={'n_spatial': n_spatial,
+                                 'n_atoms': n_atoms}
+        )
+        super(_GroFileEngine, self).__init__(options={},
+                                            descriptor=descriptor,
+                                            template=None)
+
+    def to_dict(self):
+        return {'gro': self.gro}
+
+    @classmethod
+    def from_dict(cls, dct):
+        return cls(dct['gro'])
+
+    def read_frame_data(self, file_name, file_position):
+        traj = md.load(file_name)
+        xyz = traj.xyz[0]
+        vel = np.zeros(shape=xyz.shape)
+        box = traj.unitcell_vectors[0]
+        return (xyz, vel, box)
+
+
 
 def snapshot_from_gro(gro_file):
-    class GroFileEngine(ExternalEngine):
-        def __init__(self, gro):
-            traj = md.load(gro)
-            self.topology = MDTrajTopology(traj.topology)
-            n_atoms = self.topology.n_atoms
-            n_spatial = self.topology.n_spatial
-            descriptor = SnapshotDescriptor.construct(
-                snapshot_class=ExternalMDSnapshot,
-                snapshot_dimensions={'n_spatial': n_spatial,
-                                     'n_atoms': n_atoms}
-            )
-            super(GroFileEngine, self).__init__(options={},
-                                                descriptor=descriptor,
-                                                template=None)
-
-        # read_frame_data = GromacsEngine.read_frame_data
-
-        def read_frame_data(self, file_name, file_position):
-            traj = md.load(file_name)
-            xyz = traj.xyz[0]
-            vel = np.zeros(shape=xyz.shape)
-            box = traj.unitcell_vectors[0]
-            return (xyz, vel, box)
-
-    template_engine = GroFileEngine(gro_file)
+    template_engine = _GroFileEngine(gro_file)
     snapshot = ExternalMDSnapshot(file_name=gro_file,
                                   file_position=0,
                                   engine=template_engine)