Initialize rigid body constituent particle mass

CHANGELOG.rst

@@ -13,6 +13,8 @@ Next release
   (`#2136 <https://github.com/glotzerlab/hoomd-blue/pull/2136>`__).
 * Add thermodynamically consistent frictional contact forces: ``hoomd.md.pair.friction``
   (`#2116 <https://github.com/glotzerlab/hoomd-blue/pull/2116>`__).
+* ``create_bodies`` method takes optional ``masses`` argument to set masses
+  (`#2169 <https://github.com/glotzerlab/hoomd-blue/pull/2169>`__).
 
 *Fixed*
 

hoomd/md/ForceComposite.cc

@@ -434,15 +434,10 @@ void ForceComposite::validateRigidBodies()
     m_particles_added_removed = false;
     }
 
-void ForceComposite::pyCreateRigidBodies(pybind11::dict charges)
+void ForceComposite::pyCreateRigidBodies(pybind11::dict charges, pybind11::dict masses)
     {
-    if (pybind11::len(charges) == 0)
-        {
-        createRigidBodies(std::unordered_map<unsigned int, std::vector<Scalar>>());
-        return;
-        }
-
     std::unordered_map<unsigned int, std::vector<Scalar>> charges_map;
+    std::unordered_map<unsigned int, std::vector<Scalar>> masses_map;
         {
         ArrayHandle<unsigned int> h_body_len(m_body_len, access_location::host, access_mode::read);
         for (const auto& item : charges)
@@ -464,13 +459,33 @@ void ForceComposite::pyCreateRigidBodies(pybind11::dict charges)
                 }
             charges_map.insert({type, charges_vector});
             }
+        for (const auto& item : masses)
+            {
+            const auto type = m_pdata->getTypeByName(item.first.cast<std::string>());
+            if (h_body_len.data[type] == 0)
+                {
+                throw std::runtime_error("Masses provided for non-central particle type.");
+                }
+            const auto masses_list = item.second.cast<pybind11::list>();
+            if (pybind11::len(masses_list) != h_body_len.data[type])
+                {
+                throw std::runtime_error("Masses provided not consistent with rigid body size.");
+                }
+            std::vector<Scalar> masses_vector;
+            for (auto& mass : masses_list)
+                {
+                masses_vector.emplace_back(mass.cast<Scalar>());
+                }
+            masses_map.insert({type, masses_vector});
+            }
         }
 
-    createRigidBodies(charges_map);
+    createRigidBodies(charges_map, masses_map);
     }
 
 void ForceComposite::createRigidBodies(
-    const std::unordered_map<unsigned int, std::vector<Scalar>> charges)
+    const std::unordered_map<unsigned int, std::vector<Scalar>> charges,
+    const std::unordered_map<unsigned int, std::vector<Scalar>> masses)
     {
     SnapshotParticleData<Scalar> snap;
     const BoxDim& global_box = m_pdata->getGlobalBox();
@@ -573,7 +588,10 @@ void ForceComposite::createRigidBodies(
                     snap.charge[constituent_particle_tag]
                         = charges.at(body_type)[current_body_index];
                     }
-
+                if (!masses.empty())
+                    {
+                    snap.mass[constituent_particle_tag] = masses.at(body_type)[current_body_index];
+                    }
                 // Set position and orientation of constituents
                 vec3<Scalar> body_position(snap.pos[particle_tag]);
                 quat<Scalar> body_orientation(snap.orientation[particle_tag]);

hoomd/md/ForceComposite.h

@@ -99,11 +99,12 @@ class PYBIND11_EXPORT ForceComposite : public MolecularForceCompute
     virtual void validateRigidBodies();
 
     //! Create rigid body constituent particles
-    void pyCreateRigidBodies(pybind11::dict charges);
+    void pyCreateRigidBodies(pybind11::dict charges, pybind11::dict masses);
 
     //! Create rigid body constituent particles
     virtual void
-    createRigidBodies(const std::unordered_map<unsigned int, std::vector<Scalar>> charges);
+    createRigidBodies(const std::unordered_map<unsigned int, std::vector<Scalar>> charges,
+                      const std::unordered_map<unsigned int, std::vector<Scalar>> masses);
 
     /// Construct from a Python dictionary
     void setBody(std::string typ, pybind11::object v)

hoomd/md/constrain.py

@@ -331,14 +331,17 @@ def __init__(self):
         self._add_typeparam(body)
         self.body.default = None
 
-    def create_bodies(self, state, charges=None):
+    def create_bodies(self, state, charges=None, masses=None):
         r"""Create rigid bodies from central particles in state.
 
         Args:
             state (hoomd.State): The state in which to create rigid bodies.
             charges (dict[str, list[float]]): (optional) The charges for each of
                 the constituent particles, defaults to ``None``. If ``None``,
                 all charges are zero. The keys should be the central particles.
+            masses (dict[str, list[float]]): (optional) The masses for each of
+                the constituent particles, defaults to ``None``. If ``None``,
+                all masses are one. The keys should be the central particles.
 
         `create_bodies` removes any existing constituent particles and adds new
         ones based on the body definitions in `body`. It overwrites all existing
@@ -347,7 +350,9 @@ def create_bodies(self, state, charges=None):
         if self._attached:
             raise RuntimeError("Cannot call create_bodies after running simulation.")
         super()._attach(state._simulation)
-        self._cpp_obj.createRigidBodies({} if charges is None else charges)
+        self._cpp_obj.createRigidBodies(
+            {} if charges is None else charges, {} if masses is None else masses
+        )
         # Restore previous state
         self._detach()
 

hoomd/md/pytest/test_rigid.py

@@ -66,7 +66,7 @@ def test_body_setting(valid_body_definition):
         current_body_definition[key] = valid_body_definition[key]
 
 
-def check_bodies(snapshot, definition, charges=None):
+def check_bodies(snapshot, definition, charges=None, masses=None):
     """Non-general assumes a snapshot from two_particle_snapshot_factory.
 
     This is just to prevent duplication of code from test_create_bodies and
@@ -86,6 +86,12 @@ def check_bodies(snapshot, definition, charges=None):
             assert snapshot.particles.charge[i + 2] == charges[i]
             assert snapshot.particles.charge[i + 6] == charges[i]
 
+    # check masses
+    if masses is not None:
+        for i in range(4):
+            assert snapshot.particles.mass[i + 2] == masses[i]
+            assert snapshot.particles.mass[i + 6] == masses[i]
+
     particle_one = (snapshot.particles.position[0], snapshot.particles.orientation[0])
     particle_two = (snapshot.particles.position[1], snapshot.particles.orientation[1])
 
@@ -131,11 +137,18 @@ def check_orientation(
 
 
 @skip_rowan
+@pytest.mark.parametrize(
+    "include_charge,include_mass",
+    [(False, False), (True, False), (False, True), (True, True)],
+    ids=["default", "optional-change", "optional-mass", "optional-charge-mass"],
+)
 def test_create_bodies(
     simulation_factory,
     two_particle_snapshot_factory,
     lattice_snapshot_factory,
     valid_body_definition,
+    include_charge,
+    include_mass,
 ):
     rigid = md.constrain.Rigid()
     rigid.body["A"] = valid_body_definition
@@ -145,11 +158,20 @@ def test_create_bodies(
         initial_snapshot.particles.types = ["A", "B"]
     sim = simulation_factory(initial_snapshot)
 
-    charges = [1.0, 2.0, 3.0, 4.0]
-    rigid.create_bodies(sim.state, charges={"A": charges})
+    optional_kwargs = {}
+    if include_charge:
+        optional_kwargs["charges"] = {"A": [1.0, 2.0, 3.0, 4.0]}
+    if include_mass:
+        optional_kwargs["masses"] = {"A": [5.0, 6.0, 7.0, 8.0]}
+    rigid.create_bodies(sim.state, **optional_kwargs)
     snapshot = sim.state.get_snapshot()
     if snapshot.communicator.rank == 0:
-        check_bodies(snapshot, valid_body_definition, charges)
+        check_bodies(
+            snapshot,
+            valid_body_definition,
+            optional_kwargs["charges"]["A"] if include_charge else None,
+            optional_kwargs["masses"]["A"] if include_mass else None,
+        )
 
     sim.operations.integrator = hoomd.md.Integrator(dt=0.005, rigid=rigid)
     # Ensure validate bodies passes
@@ -262,12 +284,13 @@ def test_running_simulation(
     sim.seed = 5
 
     charges = [1.0, 2.0, 3.0, 4.0]
-    rigid.create_bodies(sim.state, charges={"A": charges})
+    masses = [5.0, 6.0, 7.0, 8.0]
+    rigid.create_bodies(sim.state, charges={"A": charges}, masses={"A": masses})
     sim.operations += integrator
     sim.run(5)
     snapshot = sim.state.get_snapshot()
     if sim.device.communicator.rank == 0:
-        check_bodies(snapshot, valid_body_definition, charges)
+        check_bodies(snapshot, valid_body_definition, charges, masses)
 
     autotuned_kernel_parameter_check(instance=rigid, activate=lambda: sim.run(1))
 
@@ -587,16 +610,7 @@ def test_velocity_constituents_constant_angmom(
 
     rigid = md.constrain.Rigid()
     rigid.body["A"] = body_definition
-    rigid.create_bodies(sim.state)
-
-    intermed_snapshot = sim.state.get_snapshot()
-    if intermed_snapshot.communicator.rank == 0:
-        flags = (
-            intermed_snapshot.particles.typeid
-            == intermed_snapshot.particles.types.index("B")
-        )
-        intermed_snapshot.particles.mass[flags] = mass[flags]
-    sim.state.set_snapshot(intermed_snapshot)
+    rigid.create_bodies(sim.state, masses={"A": mass[1:]})
 
     constvol = md.methods.ConstantVolume(filter=hoomd.filter.Rigid())
     integrator = md.Integrator(

hoomd/mpcd/pytest/test_collide.py

@@ -225,29 +225,18 @@ def test_rigid_collide(
             initial_snap.particles.velocity[:] = [velo_rigid]
             initial_snap.particles.angmom[:] = [angmom_rigid]
 
+            # place the mpcd particles on top of constituents
+            initial_snap.mpcd.N = N_mpcd
+            initial_snap.mpcd.types = ["C"]
+            initial_snap.mpcd.position[:] = def_rigid["positions"]
+            initial_snap.mpcd.velocity[:] = velo_mpcd
+
         sim = simulation_factory(initial_snap)
         sim.seed = 5
 
         rigid = hoomd.md.constrain.Rigid()
         rigid.body["A"] = def_rigid
-        rigid.create_bodies(sim.state)
-
-        intermed_snap = sim.state.get_snapshot()
-        if intermed_snap.communicator.rank == 0:
-            # add mass of constituents
-            flags = (
-                intermed_snap.particles.typeid
-                == intermed_snap.particles.types.index("B")
-            )
-            intermed_snap.particles.mass[flags] = properties_rigid["mass"][flags]
-            intermed_snap.wrap()
-
-            # place the mpcd particles on top of constituents
-            intermed_snap.mpcd.N = N_mpcd
-            intermed_snap.mpcd.types = ["C"]
-            intermed_snap.mpcd.position[:] = intermed_snap.particles.position[flags]
-            intermed_snap.mpcd.velocity[:] = velo_mpcd
-        sim.state.set_snapshot(intermed_snap)
+        rigid.create_bodies(sim.state, masses={"A": properties_rigid["mass"][1:]})
 
         sim.operations.integrator = hoomd.mpcd.Integrator(
             dt=0, integrate_rotational_dof=True, rigid=rigid
@@ -364,31 +353,18 @@ def test_rigid_collide_free(
             initial_snap.particles.angmom[:] = [[0, 0, 0, 0], [0, 0, 0, 0]]
             initial_snap.particles.typeid[:] = [0, 2]
 
+            # place the mpcd particles on top of constituents
+            initial_snap.mpcd.N = N_mpcd
+            initial_snap.mpcd.types = ["C"]
+            initial_snap.mpcd.position[:] = def_rigid["positions"]
+            initial_snap.mpcd.velocity[:] = velo_mpcd
+
         sim = simulation_factory(initial_snap)
         sim.seed = 5
 
         rigid = hoomd.md.constrain.Rigid()
         rigid.body["A"] = def_rigid
-        rigid.create_bodies(sim.state)
-
-        intermed_snap = sim.state.get_snapshot()
-        if intermed_snap.communicator.rank == 0:
-            # add mass of constituents
-            flags = np.logical_or(
-                intermed_snap.particles.typeid
-                == intermed_snap.particles.types.index("B"),
-                intermed_snap.particles.typeid
-                == intermed_snap.particles.types.index("D"),
-            )
-            intermed_snap.particles.mass[flags] = properties_rigid["mass"][1:]
-            intermed_snap.wrap()
-
-            # place the mpcd particles on top of constituents
-            intermed_snap.mpcd.N = N_mpcd
-            intermed_snap.mpcd.types = ["C"]
-            intermed_snap.mpcd.position[:] = intermed_snap.particles.position[flags]
-            intermed_snap.mpcd.velocity[:] = velo_mpcd
-        sim.state.set_snapshot(intermed_snap)
+        rigid.create_bodies(sim.state, masses={"A": properties_rigid["mass"][1:]})
 
         sim.operations.integrator = hoomd.mpcd.Integrator(
             dt=0, integrate_rotational_dof=True, rigid=rigid
@@ -538,17 +514,7 @@ def test_rigid_mass_errors(
 
     rigid = hoomd.md.constrain.Rigid()
     rigid.body["A"] = def_rigid
-    rigid.create_bodies(sim.state)
-
-    # add mass of constituents
-    intermed_snap = sim.state.get_snapshot()
-    if intermed_snap.communicator.rank == 0:
-        flags = intermed_snap.particles.typeid == intermed_snap.particles.types.index(
-            "B"
-        )
-        intermed_snap.particles.mass[flags] = masses[1:]
-        intermed_snap.wrap()
-    sim.state.set_snapshot(intermed_snap)
+    rigid.create_bodies(sim.state, masses={"A": masses[1:]})
 
     sim.operations.integrator = hoomd.mpcd.Integrator(dt=0, rigid=rigid)
     sim.operations.integrator.collision_method = (