Skip to content

Commit

Permalink
Merge pull request #74 from cmelab/add-stride
Browse files Browse the repository at this point in the history
Add stride parameter when saving coarse-grained trajectories
  • Loading branch information
chrisjonesBSU authored May 1, 2024
2 parents d0c5816 + 29e570d commit ce2cf97
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 2 deletions.
6 changes: 4 additions & 2 deletions grits/coarsegrain.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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]
Expand Down Expand Up @@ -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 = []
Expand Down
13 changes: 13 additions & 0 deletions grits/tests/test_coarsegrain.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down

0 comments on commit ce2cf97

Please sign in to comment.