Merge pull request #1433 from glotzerlab/refactor/rigid-bodies

Refactor Rigid Bodies Body Definition

hoomd/md/ForceComposite.cc

@@ -8,6 +8,8 @@
 #include <sstream>
 #include <string.h>
 
+#include <pybind11/stl.h>
+
 /*! \file ForceComposite.cc
     \brief Contains code for the ForceComposite class
 */
@@ -54,9 +56,6 @@ ForceComposite::ForceComposite(std::shared_ptr<SystemDefinition> sysdef)
         h_body_len.data[i] = 0;
         }
 
-    m_body_charge.resize(m_pdata->getNTypes());
-    m_body_diameter.resize(m_pdata->getNTypes());
-
     m_d_max.resize(m_pdata->getNTypes(), Scalar(0.0));
     m_d_max_changed.resize(m_pdata->getNTypes(), false);
 
@@ -96,15 +95,11 @@ ForceComposite::~ForceComposite()
 void ForceComposite::setParam(unsigned int body_typeid,
                               std::vector<unsigned int>& type,
                               std::vector<Scalar3>& pos,
-                              std::vector<Scalar4>& orientation,
-                              std::vector<Scalar>& charge,
-                              std::vector<Scalar>& diameter)
+                              std::vector<Scalar4>& orientation)
     {
     assert(m_body_types.getPitch() >= m_pdata->getNTypes());
     assert(m_body_pos.getPitch() >= m_pdata->getNTypes());
     assert(m_body_orientation.getPitch() >= m_pdata->getNTypes());
-    assert(m_body_charge.size() >= m_pdata->getNTypes());
-    assert(m_body_diameter.size() >= m_pdata->getNTypes());
 
     if (body_typeid >= m_pdata->getNTypes())
         {
@@ -118,21 +113,6 @@ void ForceComposite::setParam(unsigned int body_typeid,
                   << " (position, orientation, type) are of unequal length.";
         throw std::runtime_error(error_msg.str());
         }
-    if (charge.size() && charge.size() != pos.size())
-        {
-        std::ostringstream error_msg;
-        error_msg << "Error initializing ForceComposite: Charges are non-empty but of different "
-                  << "length than the positions.";
-        throw std::runtime_error(error_msg.str());
-        }
-    if (diameter.size() && diameter.size() != pos.size())
-        {
-        std::ostringstream error_msg;
-        error_msg << "Error initializing ForceComposite: Diameters are non-empty but of different "
-                  << "length than the positions.";
-        throw std::runtime_error(error_msg.str());
-        }
-
     bool body_updated = false;
 
     bool body_len_changed = false;
@@ -200,18 +180,12 @@ void ForceComposite::setParam(unsigned int body_typeid,
                                                     access_location::host,
                                                     access_mode::readwrite);
 
-            m_body_charge[body_typeid].resize(type.size());
-            m_body_diameter[body_typeid].resize(type.size());
-
             // store body data in GlobalArray
             for (unsigned int i = 0; i < type.size(); ++i)
                 {
                 h_body_type.data[m_body_idx(body_typeid, i)] = type[i];
                 h_body_pos.data[m_body_idx(body_typeid, i)] = pos[i];
                 h_body_orientation.data[m_body_idx(body_typeid, i)] = orientation[i];
-
-                m_body_charge[body_typeid][i] = charge[i];
-                m_body_diameter[body_typeid][i] = diameter[i];
                 }
             }
         m_bodies_changed = true;
@@ -513,7 +487,39 @@ void ForceComposite::validateRigidBodies()
     m_particles_added_removed = false;
     }
 
-void ForceComposite::createRigidBodies()
+void ForceComposite::pyCreateRigidBodies(pybind11::dict charges)
+    {
+    if (pybind11::len(charges) == 0)
+        {
+        createRigidBodies(std::unordered_map<unsigned int, std::vector<Scalar>>());
+        return;
+        }
+    ArrayHandle<unsigned int> h_body_len(m_body_len, access_location::host, access_mode::read);
+    std::unordered_map<unsigned int, std::vector<Scalar>> charges_map;
+    for (const auto& item : charges)
+        {
+        const auto type = m_pdata->getTypeByName(item.first.cast<std::string>());
+        if (h_body_len.data[type] == 0)
+            {
+            throw std::runtime_error("Charge provided for non-central particle type.");
+            }
+        const auto charges_list = item.second.cast<pybind11::list>();
+        if (pybind11::len(charges_list) != h_body_len.data[type])
+            {
+            throw std::runtime_error("Charges provided not consistent with rigid body size.");
+            }
+        std::vector<Scalar> charges_vector;
+        for (auto& charge : charges_list)
+            {
+            charges_vector.emplace_back(charge.cast<Scalar>());
+            }
+        charges_map.insert({type, charges_vector});
+        }
+    createRigidBodies(charges_map);
+    }
+
+void ForceComposite::createRigidBodies(
+    const std::unordered_map<unsigned int, std::vector<Scalar>> charges)
     {
     SnapshotParticleData<Scalar> snap;
 
@@ -522,6 +528,7 @@ void ForceComposite::createRigidBodies()
     bool remove_existing_bodies = false;
     unsigned int n_constituent_particles = 0;
     unsigned int n_free_bodies = 0;
+    if (m_exec_conf->getRank() == 0)
         {
         // We restrict the scope of h_body_len to ensure if we remove_existing_bodies or in any way
         // reallocated m_body_len to a new memory location that we will be forced to reaquire the
@@ -604,12 +611,12 @@ void ForceComposite::createRigidBodies()
                 snap.type[constituent_particle_tag]
                     = h_body_type.data[m_body_idx(body_type, current_body_index)];
                 snap.body[constituent_particle_tag] = particle_tag;
-                snap.charge[constituent_particle_tag]
-                    = m_body_charge[body_type][current_body_index];
-                snap.diameter[constituent_particle_tag]
-                    = m_body_diameter[body_type][current_body_index];
                 snap.pos[constituent_particle_tag] = snap.pos[particle_tag];
-
+                if (!charges.empty())
+                    {
+                    snap.charge[constituent_particle_tag]
+                        = charges.at(body_type)[current_body_index];
+                    }
                 // Since the central particle tags here will be [0, n_central_particles), we know
                 // that the molecule number will be the same as the central particle tag.
                 molecule_tag[constituent_particle_tag] = particle_tag;
@@ -1029,7 +1036,7 @@ void export_ForceComposite(pybind11::module& m)
         .def("setBody", &ForceComposite::setBody)
         .def("getBody", &ForceComposite::getBody)
         .def("validateRigidBodies", &ForceComposite::validateRigidBodies)
-        .def("createRigidBodies", &ForceComposite::createRigidBodies)
+        .def("createRigidBodies", &ForceComposite::pyCreateRigidBodies)
         .def("updateCompositeParticles", &ForceComposite::updateCompositeParticles);
     }
 

hoomd/md/ForceComposite.h

@@ -62,9 +62,7 @@ class PYBIND11_EXPORT ForceComposite : public MolecularForceCompute
     virtual void setParam(unsigned int body_typeid,
                           std::vector<unsigned int>& type,
                           std::vector<Scalar3>& pos,
-                          std::vector<Scalar4>& orientation,
-                          std::vector<Scalar>& charge,
-                          std::vector<Scalar>& diameter);
+                          std::vector<Scalar4>& orientation);
 
     //! Returns true because we compute the torque on the central particle
     virtual bool isAnisotropic()
@@ -86,7 +84,11 @@ class PYBIND11_EXPORT ForceComposite : public MolecularForceCompute
     virtual void validateRigidBodies();
 
     //! Create rigid body constituent particles
-    virtual void createRigidBodies();
+    void pyCreateRigidBodies(pybind11::dict charges);
+
+    //! Create rigid body constituent particles
+    virtual void
+    createRigidBodies(const std::unordered_map<unsigned int, std::vector<Scalar>> charges);
 
     /// Construct from a Python dictionary
     void setBody(std::string typ, pybind11::object v)
@@ -98,11 +100,9 @@ class PYBIND11_EXPORT ForceComposite : public MolecularForceCompute
         pybind11::list types = v["constituent_types"];
         pybind11::list positions = v["positions"];
         pybind11::list orientations = v["orientations"];
-        pybind11::list charges = v["charges"];
-        pybind11::list diameters = v["diameters"];
         auto N = pybind11::len(positions);
         // Ensure proper list lengths
-        for (const auto& list : {types, orientations, charges, diameters})
+        for (const auto& list : {types, orientations})
             {
             if (pybind11::len(list) != N)
                 {
@@ -113,8 +113,6 @@ class PYBIND11_EXPORT ForceComposite : public MolecularForceCompute
         // extract the data from the python lists
         std::vector<Scalar3> pos_vector;
         std::vector<Scalar4> orientation_vector;
-        std::vector<Scalar> charge_vector;
-        std::vector<Scalar> diameter_vector;
         std::vector<unsigned int> type_vector;
 
         for (size_t i(0); i < N; ++i)
@@ -130,17 +128,10 @@ class PYBIND11_EXPORT ForceComposite : public MolecularForceCompute
                                                          orientation_i[2].cast<Scalar>(),
                                                          orientation_i[3].cast<Scalar>()));
 
-            charge_vector.emplace_back(charges[i].cast<Scalar>());
-            diameter_vector.emplace_back(diameters[i].cast<Scalar>());
             type_vector.emplace_back(m_pdata->getTypeByName(types[i].cast<std::string>()));
             }
 
-        setParam(m_pdata->getTypeByName(typ),
-                 type_vector,
-                 pos_vector,
-                 orientation_vector,
-                 charge_vector,
-                 diameter_vector);
+        setParam(m_pdata->getTypeByName(typ), type_vector, pos_vector, orientation_vector);
         }
 
     /// Convert parameters to a python dictionary
@@ -166,8 +157,6 @@ class PYBIND11_EXPORT ForceComposite : public MolecularForceCompute
         pybind11::list positions;
         pybind11::list orientations;
         pybind11::list types;
-        pybind11::list charges;
-        pybind11::list diameters;
 
         for (unsigned int i = 0; i < N; i++)
             {
@@ -181,15 +170,11 @@ class PYBIND11_EXPORT ForceComposite : public MolecularForceCompute
                                      static_cast<Scalar>(h_body_orientation.data[index].z),
                                      static_cast<Scalar>(h_body_orientation.data[index].w)));
             types.append(m_pdata->getNameByType(h_body_types.data[index]));
-            charges.append(m_body_charge[body_type_id][i]);
-            diameters.append(m_body_diameter[body_type_id][i]);
             }
         pybind11::dict v;
         v["constituent_types"] = types;
         v["positions"] = positions;
         v["orientations"] = orientations;
-        v["charges"] = charges;
-        v["diameters"] = diameters;
         return std::move(v);
         }
 
@@ -202,9 +187,7 @@ class PYBIND11_EXPORT ForceComposite : public MolecularForceCompute
     GlobalArray<Scalar4> m_body_orientation; //!< Constituent particle orientations per type id (2D)
     GlobalArray<unsigned int> m_body_len;    //!< Length of body per type id
 
-    std::vector<std::vector<Scalar>> m_body_charge;   //!< Constituent particle charges
-    std::vector<std::vector<Scalar>> m_body_diameter; //!< Constituent particle diameters
-    Index2D m_body_idx;                               //!< Indexer for body parameters
+    Index2D m_body_idx; //!< Indexer for body parameters
 
     std::vector<Scalar> m_d_max;       //!< Maximum body diameter per constituent particle type
     std::vector<bool> m_d_max_changed; //!< True if maximum body diameter changed (per type)

hoomd/md/constrain.py

@@ -290,19 +290,20 @@ def __init__(self):
                 'constituent_types': [str],
                 'positions': [(float,) * 3],
                 'orientations': [(float,) * 4],
-                'charges': [float],
-                'diameters': [float]
             }),
                                      allow_none=True),
                               len_keys=1))
         self._add_typeparam(body)
         self.body.default = None
 
-    def create_bodies(self, state):
+    def create_bodies(self, state, charges=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.
 
         `create_bodies` removes any existing constituent particles and adds new
         ones based on the body definitions in `body`. It overwrites all existing
@@ -312,7 +313,7 @@ def create_bodies(self, state):
             raise RuntimeError(
                 "Cannot call create_bodies after running simulation.")
         super()._attach(state._simulation)
-        self._cpp_obj.createRigidBodies()
+        self._cpp_obj.createRigidBodies({} if charges is None else charges)
         # Restore previous state
         self._detach()
 

hoomd/md/pytest/test_filter_md.py

@@ -36,8 +36,6 @@ def test_rigid_filter(make_filter_snapshot, simulation_factory):
             [0, 1, 1 / (2**(1. / 2.))],
         ],
         "orientations": [(1.0, 0.0, 0.0, 0.0)] * 4,
-        "charges": [0.0, 1.0, 2.0, 3.5],
-        "diameters": [1.0, 1.5, 0.5, 1.0]
     }
 
     snapshot = make_filter_snapshot(n=100, particle_types=["A", "B", "C"])

hoomd/md/pytest/test_rigid.py

@@ -30,8 +30,6 @@ def valid_body_definition():
             [0, 1, 1 / (2**(1. / 2.))],
         ],
         "orientations": [(1.0, 0.0, 0.0, 0.0)] * 4,
-        "charges": [0.0, 1.0, 2.0, 3.5],
-        "diameters": [1.0, 1.5, 0.5, 1.0]
     }
 
 
@@ -41,8 +39,6 @@ def test_body_setting(valid_body_definition):
         "positions": [[(1, 2)], [(1.0, 4.0, "foo")], 1.0, "hello"],
         "orientations": [[(1, 2, 3)], [(1.0, 4.0, 5.0, "foo")], [1.0], 1.0,
                          "foo"],
-        "charges": [0.0, ["foo"]],
-        "diameters": [1.0, "foo", ["foo"]]
     }
 
     rigid = md.constrain.Rigid()
@@ -67,7 +63,7 @@ def test_body_setting(valid_body_definition):
         current_body_definition[key] = valid_body_definition[key]
 
 
-def check_bodies(snapshot, definition):
+def check_bodies(snapshot, definition, charges=None):
     """Non-general assumes a snapshot from two_particle_snapshot_factory.
 
     This is just to prevent duplication of code from test_create_bodies and
@@ -82,14 +78,10 @@ def check_bodies(snapshot, definition):
     assert all(snapshot.particles.body[6:] == 1)
 
     # check charges
-    for i in range(4):
-        assert snapshot.particles.charge[i + 2] == definition["charges"][i]
-        assert snapshot.particles.charge[i + 6] == definition["charges"][i]
-
-    # check diameters
-    for i in range(4):
-        assert snapshot.particles.diameter[i + 2] == definition["diameters"][i]
-        assert snapshot.particles.diameter[i + 6] == definition["diameters"][i]
+    if charges is not None:
+        for i in range(4):
+            assert snapshot.particles.charge[i + 2] == charges[i]
+            assert snapshot.particles.charge[i + 6] == charges[i]
 
     particle_one = (snapshot.particles.position[0],
                     snapshot.particles.orientation[0])
@@ -135,10 +127,11 @@ def test_create_bodies(simulation_factory, two_particle_snapshot_factory,
         initial_snapshot.particles.types = ["A", "B"]
     sim = simulation_factory(initial_snapshot)
 
-    rigid.create_bodies(sim.state)
+    charges = [1.0, 2.0, 3.0, 4.0]
+    rigid.create_bodies(sim.state, charges={"A": charges})
     snapshot = sim.state.get_snapshot()
     if snapshot.communicator.rank == 0:
-        check_bodies(snapshot, valid_body_definition)
+        check_bodies(snapshot, valid_body_definition, charges)
 
     sim.operations.integrator = hoomd.md.Integrator(dt=0.005, rigid=rigid)
     # Ensure validate bodies passes
@@ -246,12 +239,13 @@ def test_running_simulation(simulation_factory, two_particle_snapshot_factory,
     sim = simulation_factory(initial_snapshot)
     sim.seed = 5
 
-    rigid.create_bodies(sim.state)
+    charges = [1.0, 2.0, 3.0, 4.0]
+    rigid.create_bodies(sim.state, charges={"A": charges})
     sim.operations += integrator
     sim.run(5)
     snapshot = sim.state.get_snapshot()
     if sim.device.communicator.rank == 0:
-        check_bodies(snapshot, valid_body_definition)
+        check_bodies(snapshot, valid_body_definition, charges)
 
     autotuned_kernel_parameter_check(instance=rigid,
                                      activate=lambda: sim.run(1))