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

gmso/external/convert_mbuild.py

@@ -285,6 +285,7 @@ def _parse_molecule_residue(site_map, compound):
     connected_subgraph = compound.bond_graph.connected_components()
     molecule_tracker = dict()
     residue_tracker = dict()
+    total_molecule_count = 0
     for molecule in connected_subgraph:
         if len(molecule) == 1:
             ancestors = [molecule[0]]
@@ -305,8 +306,9 @@ def _parse_molecule_residue(site_map, compound):
         if molecule_tag.name in molecule_tracker:
             molecule_tracker[molecule_tag.name] += 1
         else:
-            molecule_tracker[molecule_tag.name] = 0
+            molecule_tracker[molecule_tag.name] = 0 + total_molecule_count
         molecule_number = molecule_tracker[molecule_tag.name]
+        total_molecule_count += 1
         """End of molecule parsing"""
 
         for particle in molecule:

gmso/formats/gsd.py

@@ -17,7 +17,6 @@ def write_gsd(
     top,
     filename,
     base_units=None,
-    rigid_bodies=None,
     shift_coords=True,
     write_special_pairs=True,
 ):
@@ -33,11 +32,6 @@ def write_gsd(
         gmso.Topology object
     filename : str
         Path of the output file.
-    rigid_bodies : list of int, optional, default=None
-        List of rigid body information. An integer value is required for each
-        atom corresponding to the index of the rigid body the particle is to be
-        associated with. A value of None indicates the atom is not part of a
-        rigid body.
     shift_coords : bool, optional, default=True
         Shift coordinates from (0, L) to (-L/2, L/2) if necessary.
     write_special_pairs : bool, optional, default=True
@@ -54,7 +48,6 @@ def write_gsd(
     gsd_snapshot = to_gsd_snapshot(
         top=top,
         base_units=base_units,
-        rigid_bodies=rigid_bodies,
         shift_coords=shift_coords,
         parse_special_pairs=write_special_pairs,
     )[0]

gmso/tests/test_convert_mbuild.py

@@ -132,7 +132,7 @@ def test_uneven_hierarchy(self):
         for site in top.sites:
             if site.name == "particle2":
                 assert site.group == "mid"
-                assert site.molecule == ("particle2", 0)
+                assert site.molecule == ("particle2", 1)
             elif site.name == "particle1":
                 assert site.molecule == ("particle1", 0)
 

gmso/tests/test_hoomd.py

@@ -1,11 +1,16 @@
 import forcefield_utilities as ffutils
 import hoomd
+import numpy as np
 import pytest
 import unyt as u
 
 from gmso import ForceField
 from gmso.external import from_mbuild
-from gmso.external.convert_hoomd import to_hoomd_forcefield, to_hoomd_snapshot
+from gmso.external.convert_hoomd import (
+    to_gsd_snapshot,
+    to_hoomd_forcefield,
+    to_hoomd_snapshot,
+)
 from gmso.parameterization import apply
 from gmso.tests.base_test import BaseTest
 from gmso.tests.utils import get_path
@@ -24,7 +29,7 @@ def run_hoomd_nvt(snapshot, forces):
     sim = hoomd.Simulation(device=cpu)
     sim.create_state_from_snapshot(snapshot)
 
-    integrator = hoomd.md.Integrator(dt=0.001)
+    integrator = hoomd.md.Integrator(dt=0.0001)
     integrator.forces = list(set().union(*forces.values()))
 
     temp = 300 * u.K
@@ -45,6 +50,129 @@ def run_hoomd_nvt(snapshot, forces):
     return sim
 
 
+@pytest.mark.skipif(not has_mbuild, reason="mbuild not installed")
+class TestGsd(BaseTest):
+    def test_rigid_bodies(self):
+        ethane = mb.lib.molecules.Ethane()
+        box = mb.fill_box(ethane, n_compounds=2, box=[2, 2, 2])
+        top = from_mbuild(box)
+        for site in top.sites:
+            site.molecule.isrigid = True
+        top.identify_connections()
+
+        top_no_rigid = from_mbuild(box)
+        top_no_rigid.identify_connections()
+
+        rigid_ids = [site.molecule.number for site in top.sites]
+        assert set(rigid_ids) == {0, 1}
+
+        snapshot, refs, rigid = to_gsd_snapshot(top)
+        snapshot.validate()
+        snapshot_no_rigid, refs = to_gsd_snapshot(top_no_rigid)
+        # Check that snapshot has rigid particles added
+        assert "Ethane" in snapshot.particles.types
+        assert "Ethane" not in snapshot_no_rigid.particles.types
+        assert snapshot.particles.N - 2 == snapshot_no_rigid.particles.N
+        assert np.array_equal(snapshot.particles.typeid[:2], np.array([0, 0]))
+        assert np.array_equal(
+            snapshot.particles.body[2:10], np.array([0] * ethane.n_particles)
+        )
+        assert np.array_equal(
+            snapshot.particles.body[10:], np.array([1] * ethane.n_particles)
+        )
+        assert np.allclose(snapshot.particles.mass[0], ethane.mass, atol=1e-2)
+        # Check that topology isn't changed by using rigid bodies
+        assert snapshot.bonds.N == snapshot_no_rigid.bonds.N
+        assert snapshot.angles.N == snapshot_no_rigid.angles.N
+        assert snapshot.dihedrals.N == snapshot_no_rigid.dihedrals.N
+        # Check if particle indices in snapshot groups are adjusted by N rigid bodies
+        for group1, group2 in zip(snapshot.bonds.group, snapshot_no_rigid.bonds.group):
+            assert np.array_equal(np.array(group1), np.array(group2) + 2)
+        for group1, group2 in zip(
+            snapshot.angles.group, snapshot_no_rigid.angles.group
+        ):
+            assert np.array_equal(np.array(group1), np.array(group2) + 2)
+        for group1, group2 in zip(
+            snapshot.dihedrals.group, snapshot_no_rigid.dihedrals.group
+        ):
+            assert np.array_equal(np.array(group1), np.array(group2) + 2)
+        for group1, group2 in zip(
+            snapshot.impropers.group, snapshot_no_rigid.impropers.group
+        ):
+            assert np.array_equal(np.array(group1), np.array(group2) + 2)
+
+        for site in top.iter_sites_by_molecule("Ethane"):
+            assert site.name in rigid.body["Ethane"]["constituent_types"]
+
+    def test_multiple_rigid_bodies(self, gaff_forcefield):
+        ethane = mb.lib.molecules.Ethane()
+        benzene = mb.load("c1ccccc1", smiles=True)
+        benzene.name = "Benzene"
+        box = mb.fill_box([ethane, benzene], n_compounds=[1, 1], box=[2, 2, 2])
+        top = from_mbuild(box)
+        for site in top.sites:
+            site.molecule.isrigid = True
+        apply(top, gaff_forcefield, identify_connections=True)
+
+        rigid_ids = [site.molecule.number for site in top.sites]
+        assert set(rigid_ids) == {0, 1}
+
+        snapshot, refs, rigid = to_gsd_snapshot(top)
+        snapshot.validate()
+
+        for site in top.iter_sites_by_molecule("Ethane"):
+            assert site.atom_type.name in rigid.body["Ethane"]["constituent_types"]
+
+        for site in top.iter_sites_by_molecule("Benzene"):
+            assert site.atom_type.name in rigid.body["Benzene"]["constituent_types"]
+
+        assert np.array_equal(np.zeros(8), snapshot.particles.body[2:10])
+        assert np.array_equal(np.ones(12), snapshot.particles.body[10:])
+        assert round(float(snapshot.particles.mass[0]), 4) == round(
+            float(ethane.mass), 4
+        )
+        assert round(float(snapshot.particles.mass[1]), 4) == round(
+            float(benzene.mass), 4
+        )
+
+    def test_rigid_bodies_mix(self):
+        ethane = mb.lib.molecules.Ethane()
+        ethane.name = "ethane"
+        benzene = mb.load("c1ccccc1", smiles=True)
+        benzene.name = "benzene"
+        box = mb.fill_box([benzene, ethane], n_compounds=[1, 1], box=[2, 2, 2])
+        top = from_mbuild(box)
+        for site in top.sites:
+            if site.molecule.name == "benzene":
+                site.molecule.isrigid = True
+
+        snapshot, refs, rigid = to_gsd_snapshot(top)
+        snapshot.validate()
+        assert snapshot.particles.typeid[0] == 0
+        assert snapshot.particles.N == top.n_sites + 1
+        assert np.array_equal(
+            snapshot.particles.body[-ethane.n_particles :],
+            np.array([-1] * ethane.n_particles),
+        )
+        assert np.array_equal(
+            snapshot.particles.body[1 : benzene.n_particles + 1],
+            np.array([0] * benzene.n_particles),
+        )
+
+    def test_rigid_bodies_bad_hierarchy(self):
+        ethane = mb.lib.molecules.Ethane()
+        ethane.name = "ethane"
+        benzene = mb.load("c1ccccc1", smiles=True)
+        benzene.name = "benzene"
+        box = mb.fill_box([ethane, benzene], n_compounds=[1, 1], box=[2, 2, 2])
+        top = from_mbuild(box)
+        for site in top.sites:
+            if site.molecule.name == "benzene":
+                site.molecule.isrigid = True
+        with pytest.raises(RuntimeError):
+            to_gsd_snapshot(top)
+
+
 class TestHoomd(BaseTest):
     def test_hoomd_simulation(self):
         compound = mb.load("CCC", smiles=True)

gmso/tests/test_utils.py

@@ -1,8 +1,10 @@
+import numpy as np
 import pytest
 import unyt as u
 
 from gmso.core.atom import Atom
 from gmso.core.dihedral import Dihedral
+from gmso.utils.geometry import moment_of_inertia
 from gmso.utils.io import run_from_ipython
 from gmso.utils.misc import unyt_to_hashable
 from gmso.utils.sorting import sort_connection_members, sort_connection_strings
@@ -55,3 +57,27 @@ def test_sorting():
         "atom2",
         "atom4",
     )
+
+
+def test_moment_of_inertia():
+    xyz = [
+        np.array([-0.5, -0.5, 0]),
+        np.array([0.5, -0.5, 0]),
+        np.array([0.5, 0.5, 0]),
+        np.array([-0.5, 0.5, 0]),
+    ]
+    tensor = moment_of_inertia(xyz=xyz, masses=np.array([1.0 for i in xyz]))
+    assert np.array_equal(tensor, np.array([1, 1, 2]))
+
+    xyz = [
+        np.array([0, 0, 0]),
+        np.array([1, 0, 0]),
+        np.array([1, 1, 0]),
+        np.array([0, 1, 0]),
+    ]
+    tensor = moment_of_inertia(
+        xyz=xyz,
+        center=np.array((0.5, 0.5, 0)),
+        masses=np.array([1.0 for i in xyz]),
+    )
+    assert np.array_equal(tensor, np.array([1, 1, 2]))

gmso/utils/geometry.py

@@ -29,3 +29,34 @@ def coord_shift(xyz, box_lengths):
         xyz += box_max
 
     return xyz
+
+
+def moment_of_inertia(xyz, masses, center=np.zeros(3)):
+    """Find the moment of inertia tensor given a set of
+    particle coordinates and their corresponding masses.
+
+    This method is used in setting rigid body moments
+    of inertia.
+
+    Parameters
+    ----------
+    xyz : numpy.ndarray (N,3)
+        Coordinates of the particles
+    masses : numpy.ndarray (N,)
+        Masses of the particles
+    center : numpy.ndarray (3,), default (0,0,0)
+        Coordinates of the particle's center
+
+    Returns
+    -------
+    numpy.ndarray (3,)
+        Diagonal components of the moment of inertia tensor.
+    """
+    xyz -= np.asarray(center)
+    x = xyz[:, 0]
+    y = xyz[:, 1]
+    z = xyz[:, 2]
+    Ixx = np.sum((y**2 + z**2) * masses)
+    Iyy = np.sum((x**2 + z**2) * masses)
+    Izz = np.sum((x**2 + y**2) * masses)
+    return np.array((Ixx, Iyy, Izz))