Merge pull request #65 from marjanAlbouye/main

Save angle and dihedrals in CG snapshots

environment.yml

@@ -18,3 +18,4 @@ dependencies:
   - pytest
   - pytest-cov
   - rdkit
+  - cmeutils >= 1.2.0

grits/coarsegrain.py

@@ -9,6 +9,7 @@
 import freud
 import gsd.hoomd
 import numpy as np
+from cmeutils.gsd_utils import identify_snapshot_connections
 from ele import element_from_symbol
 from mbuild import Compound, clone
 from mbuild.utils.io import run_from_ipython
@@ -760,7 +761,7 @@ def save(self, cg_gsdfile, start=0, stop=None):
                 new_snap.particles.N = len(typeid)
                 new_snap.particles.position = position
                 new_snap.particles.image = images
-                new_snap.particles.typeid = typeid
+                new_snap.particles.typeid = typeid.astype(int)
                 new_snap.particles.types = types
                 new_snap.particles.mass = mass
 
@@ -775,4 +776,5 @@ def save(self, cg_gsdfile, start=0, stop=None):
                     new_snap.bonds.type_shapes = bond_type_shapes
                     if self.aniso_beads:
                         new_snap.particles.orientation = orientation
+                    new_snap = identify_snapshot_connections(new_snap)
                 new.append(new_snap)

grits/tests/test_coarsegrain.py

@@ -278,6 +278,22 @@ def test_pps(self, tmp_path):
             )
             assert len(snap.bonds.types) == 1
 
+            assert (
+                len(snap.angles.typeid)
+                == len(snap.angles.group)
+                == snap.angles.N
+            )
+            assert snap.angles.N == 15 * 13
+            assert snap.angles.types == ["_B-_B-_B"]
+
+            assert (
+                len(snap.dihedrals.typeid)
+                == len(snap.dihedrals.group)
+                == snap.dihedrals.N
+            )
+            assert snap.dihedrals.N == 15 * 12
+            assert snap.dihedrals.types == ["_B-_B-_B-_B"]
+
         cg_json = tmp_path / "cg-pps.json"
         system.save_mapping(cg_json)
 
@@ -303,6 +319,21 @@ def test_pps_noH(self, tmp_path):
                 len(snap.bonds.typeid) == len(snap.bonds.group) == snap.bonds.N
             )
             assert len(snap.bonds.types) == 1
+            assert (
+                len(snap.angles.typeid)
+                == len(snap.angles.group)
+                == snap.angles.N
+            )
+            assert snap.angles.N == 15 * 13
+            assert snap.angles.types == ["_B-_B-_B"]
+
+            assert (
+                len(snap.dihedrals.typeid)
+                == len(snap.dihedrals.group)
+                == snap.dihedrals.N
+            )
+            assert snap.dihedrals.N == 15 * 12
+            assert snap.dihedrals.types == ["_B-_B-_B-_B"]
 
         cg_json = tmp_path / "cg-pps.json"
         system.save_mapping(cg_json)