Merge pull request #74 from cmelab/add-stride

Add stride parameter when saving coarse-grained trajectories

grits/coarsegrain.py

@@ -667,7 +667,7 @@ def save_mapping(self, filename):
             json.dump(self.mapping, f, cls=NumpyEncoder)
         print(f"Mapping saved to {filename}")
 
-    def save(self, cg_gsdfile, start=0, stop=None):
+    def save(self, cg_gsdfile, start=0, stop=None, stride=1):
         """Save the coarse-grain system to a gsd file.
 
         Does not calculate the image of the coarse-grain bead.
@@ -688,6 +688,8 @@ def save(self, cg_gsdfile, start=0, stop=None):
         stop : int, default None
             Where to stop reading the gsd trajectory the system was created
             with. If None, will stop at the last frame.
+        stride : int, default 1
+            The step size to use when iterating through start:stop
         """
         typeid = []
         types = [i.split("...")[0] for i in self.mapping]
@@ -722,7 +724,7 @@ def save(self, cg_gsdfile, start=0, stop=None):
         ) as old:
             # stop being None is fine; slicing [0:None] gives whole array,
             # even in edge case where there's only one or two frames
-            for s in old[start:stop]:
+            for s in old[start:stop:stride]:
                 new_snap = gsd.hoomd.Frame()
                 position = []
                 mass = []

grits/tests/test_coarsegrain.py

@@ -213,6 +213,19 @@ def test_benzene(self, tmp_path):
         cg_json = tmp_path / "cg-benzene.json"
         system.save_mapping(cg_json)
 
+    def test_stride(self, tmp_path):
+        gsdfile = path.join(asset_dir, "benzene-aa.gsd")
+        system = CG_System(
+            gsdfile,
+            beads={"_B": "c1ccccc1"},
+            conversion_dict=amber_dict,
+            mass_scale=12.011,
+        )
+        cg_gsd = tmp_path / "cg-benzene.gsd"
+        system.save(cg_gsdfile=cg_gsd, start=0, stop=-1, stride=2)
+        with gsd.hoomd.open(cg_gsd) as f:
+            assert len(f) == 3
+
     def test_anisobenzene(self, tmp_path):
         gsdfile = path.join(asset_dir, "benzene-aa.gsd")
         system = CG_System(