Add support for writing rigid bodies when converting to HOOMD formats.

gmso/abc/abstract_site.py

@@ -220,6 +220,12 @@ def __repr__(self):
         alias="position",
     )
 
+    rigid_id_: Optional[int] = Field(
+        None,
+        description="Rigid body ID used to assign rigid body contraints in HOOMD-Blue topologies.",
+        alias="rigid_id",
+    )
+
     model_config = ConfigDict(
         alias_to_fields={
             "name": "name_",
@@ -261,6 +267,11 @@ def residue(self):
         """Return the residue assigned to the site."""
         return self.__dict__.get("residue_")
 
+    @property
+    def rigid_id(self):
+        """Return the rigid body ID for this atom."""
+        return self.__dict__.get("rigid_id_")
+
     @field_serializer("position_")
     def serialize_position(self, position_: PositionType):
         return unyt_to_dict(position_)

gmso/external/convert_hoomd.py

@@ -284,8 +284,6 @@ def _parse_particle_information(
         site.name if site.atom_type is None else site.atom_type.name
         for site in top.sites
     ]
-    unique_types = sorted(list(set(types)))
-    typeids = np.array([unique_types.index(t) for t in types])
     masses = list()
     charges = list()
     for site in top.sites:
@@ -295,6 +293,39 @@ def _parse_particle_information(
             else 1 * base_units["mass"]
         )
         charges.append(site.charge if site.charge else 0 * u.elementary_charge)
+    # Check for rigid IDs
+    rigid_ids = [site.rigid_id for site in top.sites]
+    rigid_ids_set = set(rigid_ids)
+    if all(rigid_ids):  # Need to create rigid bodies
+        n_rigid = len(rigid_ids_set)
+        write_rigid = True
+    else:
+        write_rigid = False
+        n_rigid = 0
+    unique_types = sorted(list(set(types)))
+    typeids = np.array([unique_types.index(t) for t in types])
+    if write_rigid:
+        rigid_masses = np.zeros(n_rigid)
+        rigid_xyz = np.zeros(n_rigid)
+        # Rigid particle type defaults to "R"; add to front of list
+        # TODO: Can we always use "R" here? What if an atom_type is "R"?
+        unique_types = ["R"] + unique_types
+        # Rigid particles get type ID 0, move all others up by 1
+        typeids = np.concatenate(([0] * n_rigid), typeids + 1)
+        # Update mass list and position list of Frame
+        for idx, _id in enumerate(rigid_ids_set):
+            group_indices = np.where(rigid_ids == _id)[0]
+            group_positions = xyz[group_indices]
+            group_masses = masses[group_indices]
+            com_xyz = np.sum(group_positions.T * group_masses, axis=1) / sum(
+                group_masses
+            )
+            rigid_masses[idx] = sum(group_masses)
+            rigid_xyz[idx] = com_xyz
+        # Append rigid center mass and xyz to front
+        masses = np.concatenate(rigid_masses, masses)
+        xyz = np.concatenate(rigid_xyz, xyz)
+        rigid_id_tags = np.concatenate(np.arange(n_rigid), np.array(rigid_ids))
 
     """
     Permittivity of free space = 2.39725e-4 e^2/((kcal/mol)(angstrom)),
@@ -309,24 +340,23 @@ def _parse_particle_information(
     ) ** 0.5
 
     if isinstance(snapshot, hoomd.Snapshot):
-        snapshot.particles.N = top.n_sites
+        snapshot.particles.N = top.n_sites + n_rigid
         snapshot.particles.types = unique_types
         snapshot.particles.position[0:] = xyz
         snapshot.particles.typeid[0:] = typeids
         snapshot.particles.mass[0:] = masses
         snapshot.particles.charge[0:] = charges / charge_factor
+        if write_rigid:
+            snapshot.particles.body[0:] = rigid_id_tags
     elif isinstance(snapshot, gsd.hoomd.Frame):
         snapshot.particles.N = top.n_sites
         snapshot.particles.types = unique_types
         snapshot.particles.position = xyz
         snapshot.particles.typeid = typeids
         snapshot.particles.mass = masses
         snapshot.particles.charge = charges / charge_factor
-    if rigid_bodies:
-        warnings.warn(
-            "Rigid bodies detected, but not yet implemented for GSD",
-            NotYetImplementedWarning,
-        )
+        if write_rigid:
+            snapshot.particles.body = rigid_id_tags
 
 
 def _parse_pairs_information(

gmso/external/convert_mbuild.py

@@ -251,6 +251,7 @@ def _parse_particle(particle_map, site):
         charge=charge,
         mass=mass,
     )
+    particle.rigid_id = site.rigid_id
     particle_map[site] = particle
     return particle
 
@@ -272,6 +273,7 @@ def _parse_site(site_map, particle, search_method, infer_element=False):
         element=ele,
         charge=charge,
         mass=mass,
+        rigid_id=particle.rigid_id,
         molecule=site_map[particle]["molecule"],
         residue=site_map[particle]["residue"],
         group=site_map[particle]["group"],

gmso/tests/test_convert_mbuild.py

@@ -259,3 +259,14 @@ def test_nontop_level_compound(self, mb_ethane):
         cpd.add(mb_ethane)
         with pytest.raises(AssertionError):
             from_mbuild(mb_ethane)
+
+    def test_rigid_ids(self, mb_ethane):
+        box = mb.fill_box(mb_ethane, n_compounds=3, box=[2, 2, 2])
+        for i, child in enumerate(box.children):
+            for p in child.particles():
+                p.rigid_id = i
+
+        top = from_mbuild(box)
+        assert set([site.rigid_id for site in top.sites]) == {0, 1, 2}
+        for particle, site in zip(box.particles(), top.sites):
+            assert particle.rigid_id == site.rigid_id