Make MPCD collision methods compatible with rigid bodies

CHANGELOG.rst

@@ -9,6 +9,8 @@ Next release
 
 *Added*
 
+* Implemented rigid body compatibility with MPCD without MPI
+  (`#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>`__).
 

hoomd/md/ForceComposite.cc

@@ -47,6 +47,12 @@ ForceComposite::ForceComposite(std::shared_ptr<SystemDefinition> sysdef)
         h_body_len.data[i] = 0;
         }
 
+    GPUVector<unsigned int> rigid_center(m_exec_conf);
+    m_rigid_center.swap(rigid_center);
+
+    GPUVector<unsigned int> lookup_center(m_exec_conf);
+    m_lookup_center.swap(lookup_center);
+
     m_d_max.resize(m_pdata->getNTypes(), Scalar(0.0));
     m_d_max_changed.resize(m_pdata->getNTypes(), false);
 
@@ -1048,6 +1054,45 @@ void ForceComposite::updateCompositeParticles(uint64_t timestep)
         }
     }
 
+void ForceComposite::findRigidCenters()
+    {
+    ArrayHandle<unsigned int> h_tag(m_pdata->getTags(), access_location::host, access_mode::read);
+    ArrayHandle<unsigned int> h_rtag(m_pdata->getRTags(), access_location::host, access_mode::read);
+    ArrayHandle<unsigned int> h_body(m_pdata->getBodies(),
+                                     access_location::host,
+                                     access_mode::read);
+
+    m_rigid_center.resize(m_pdata->getN() + m_pdata->getNGhosts());
+    m_lookup_center.resize(m_pdata->getN() + m_pdata->getNGhosts());
+    ArrayHandle<unsigned int> h_rigid_center(m_rigid_center,
+                                             access_location::host,
+                                             access_mode::overwrite);
+    ArrayHandle<unsigned int> h_lookup_center(m_lookup_center,
+                                              access_location::host,
+                                              access_mode::overwrite);
+    unsigned int num_centers = 0;
+    const unsigned int n_particles_local = m_pdata->getN() + m_pdata->getNGhosts();
+    for (unsigned int idx = 0; idx < n_particles_local; ++idx)
+        {
+        unsigned int lookup_center = NO_BODY;
+
+        // if particle is in a rigid body, find its center
+        const unsigned int central_tag = h_body.data[idx];
+        if (central_tag < MIN_FLOPPY)
+            {
+            lookup_center = h_rtag.data[central_tag];
+            // also record rigid particles that are centers
+            if (lookup_center == idx)
+                {
+                h_rigid_center.data[num_centers++] = lookup_center;
+                }
+            }
+
+        h_lookup_center.data[idx] = lookup_center;
+        }
+    m_n_rigid = num_centers;
+    }
+
 namespace detail
     {
 void export_ForceComposite(pybind11::module& m)

hoomd/md/ForceComposite.h

@@ -87,7 +87,7 @@ class PYBIND11_EXPORT ForceComposite : public MolecularForceCompute
 
 #ifdef ENABLE_MPI
     //! Get ghost particle fields requested by this pair potential
-    virtual CommFlags getRequestedCommFlags(uint64_t timestep);
+    CommFlags getRequestedCommFlags(uint64_t timestep) override;
 #endif
 
     /// Update the constituent particles of a composite particle using the position, velocity
@@ -199,6 +199,56 @@ class PYBIND11_EXPORT ForceComposite : public MolecularForceCompute
         return m_n_free_particles_global;
         }
 
+    /// Get constituent particle types per type id
+    const GPUArray<unsigned int>& getBodyTypes() const
+        {
+        return m_body_types;
+        }
+
+    /// Get constituent particle offsets per type id
+    const GPUArray<Scalar3>& getBodyOffsets() const
+        {
+        return m_body_pos;
+        }
+
+    /// Get constituent particle orientations per type id
+    const GPUArray<Scalar4>& getBodyOrientations() const
+        {
+        return m_body_orientation;
+        }
+
+    /// Get length of body per type id
+    const GPUArray<unsigned int>& getBodyLengths() const
+        {
+        return m_body_len;
+        }
+
+    /// Get body parameter indexer
+    const Index2D& getBodyIndexer()
+        {
+        return m_body_idx;
+        }
+
+    //! Get rigid centers
+    GPUVector<unsigned int>& getRigidCenters()
+        {
+        checkParticlesSorted();
+        return m_rigid_center;
+        }
+
+    //! Get lookup centers
+    GPUVector<unsigned int>& getLookupCenters()
+        {
+        checkParticlesSorted();
+        return m_lookup_center;
+        }
+
+    //! Get number of local rigid bodies
+    const unsigned int getNLocal() const
+        {
+        return m_n_rigid;
+        }
+
     protected:
     bool m_bodies_changed;          //!< True if constituent particles have changed
     bool m_particles_added_removed; //!< True if particles have been added or removed
@@ -216,6 +266,10 @@ class PYBIND11_EXPORT ForceComposite : public MolecularForceCompute
     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)
 
+    unsigned int m_n_rigid;                  //!< Number of rigid bodies on the local rank.
+    GPUVector<unsigned int> m_rigid_center;  //!< Local particle indices of all central particles
+    GPUVector<unsigned int> m_lookup_center; //!< Lookup particle index -> central particle index
+
 #ifdef ENABLE_MPI
     /// The system's communicator.
     std::shared_ptr<Communicator> m_comm;
@@ -227,11 +281,26 @@ class PYBIND11_EXPORT ForceComposite : public MolecularForceCompute
         m_particles_added_removed = true;
         }
 
+    //! Helper function to check if particles have been sorted and rebuild indices if necessary
+    void checkParticlesSorted() override
+        {
+        if (m_rebuild_molecules)
+            // identify center particles for use in GPU kernel
+            findRigidCenters();
+
+        // Must be called second since the method sets m_rebuild_molecules
+        // to false if it is true.
+        MolecularForceCompute::checkParticlesSorted();
+        }
+
+    //! Helper kernel to sort rigid bodies by their center particles
+    virtual void findRigidCenters();
+
     /// Return the requested minimum ghost layer width for a body's central particle.
     virtual Scalar requestBodyGhostLayerWidth(unsigned int type, Scalar* h_r_ghost);
 
     //! Compute the forces and torques on the central particle
-    virtual void computeForces(uint64_t timestep);
+    void computeForces(uint64_t timestep) override;
     };
 
     } // end namespace md

hoomd/md/ForceCompositeGPU.h

@@ -36,26 +36,14 @@ class PYBIND11_EXPORT ForceCompositeGPU : public ForceComposite
      * \param remote If true, consider remote bodies, otherwise bodies
      *        with a local central particle
      */
-    virtual void updateCompositeParticles(uint64_t timestep);
+    void updateCompositeParticles(uint64_t timestep) override;
 
     protected:
     //! Compute the forces and torques on the central particle
-    virtual void computeForces(uint64_t timestep);
+    void computeForces(uint64_t timestep) override;
 
     //! Helper kernel to sort rigid bodies by their center particles
-    virtual void findRigidCenters();
-
-    //! Helper function to check if particles have been sorted and rebuild indices if necessary
-    virtual void checkParticlesSorted()
-        {
-        if (m_rebuild_molecules)
-            // identify center particles for use in GPU kernel
-            findRigidCenters();
-
-        // Must be called second since the method sets m_rebuild_molecules
-        // to false if it is true.
-        MolecularForceCompute::checkParticlesSorted();
-        }
+    void findRigidCenters() override;
 
     /// Autotuner for block size and threads per body.
     std::shared_ptr<Autotuner<2>> m_tuner_force;
@@ -67,12 +55,6 @@ class PYBIND11_EXPORT ForceCompositeGPU : public ForceComposite
     std::shared_ptr<Autotuner<1>> m_tuner_update;
 
     GPUArray<uint2> m_flag; //!< Flag to read out error condition
-
-    /// Number of rigid bodies on the local rank.
-    unsigned int m_n_rigid;
-
-    GPUVector<unsigned int> m_rigid_center;  //!< Contains particle indices of all central particles
-    GPUVector<unsigned int> m_lookup_center; //!< Lookup particle index -> central particle index
     };
 
     } // end namespace md

hoomd/md/IntegratorTwoStep.h

@@ -102,7 +102,7 @@ class PYBIND11_EXPORT IntegratorTwoStep : public Integrator
         return m_rigid_bodies;
         }
 
-    void setRigid(std::shared_ptr<ForceComposite> new_rigid)
+    virtual void setRigid(std::shared_ptr<ForceComposite> new_rigid)
         {
         m_rigid_bodies = new_rigid;
         }

hoomd/mpcd/ATCollisionMethod.cc

@@ -17,8 +17,10 @@ mpcd::ATCollisionMethod::ATCollisionMethod(std::shared_ptr<SystemDefinition> sys
                                            uint64_t period,
                                            int phase,
                                            std::shared_ptr<Variant> T)
-    : mpcd::CollisionMethod(sysdef, cur_timestep, period, phase), m_T(T)
+    : mpcd::CollisionMethod(sysdef, cur_timestep, period, phase)
     {
+    setTemperature(T);
+    requireTemperature();
     m_exec_conf->msg->notice(5) << "Constructing MPCD AT collision method" << std::endl;
     }
 
@@ -278,10 +280,7 @@ void export_ATCollisionMethod(pybind11::module& m)
                             uint64_t,
                             uint64_t,
                             int,
-                            std::shared_ptr<Variant>>())
-        .def_property("kT",
-                      &mpcd::ATCollisionMethod::getTemperature,
-                      &mpcd::ATCollisionMethod::setTemperature);
+                            std::shared_ptr<Variant>>());
     }
     } // namespace detail
     } // namespace mpcd

hoomd/mpcd/ATCollisionMethod.h

@@ -43,22 +43,9 @@ class PYBIND11_EXPORT ATCollisionMethod : public mpcd::CollisionMethod
 
     void setCellList(std::shared_ptr<mpcd::CellList> cl) override;
 
-    //! Get the temperature
-    std::shared_ptr<Variant> getTemperature() const
-        {
-        return m_T;
-        }
-
-    //! Set the temperature
-    void setTemperature(std::shared_ptr<Variant> T)
-        {
-        m_T = T;
-        }
-
     protected:
     std::shared_ptr<mpcd::CellThermoCompute> m_thermo;      //!< Cell thermo
     std::shared_ptr<mpcd::CellThermoCompute> m_rand_thermo; //!< Cell thermo for random velocities
-    std::shared_ptr<Variant> m_T;                           //!< Temperature for thermostat
 
     //! Implementation of the collision rule
     virtual void rule(uint64_t timestep) override;

hoomd/mpcd/CMakeLists.txt

@@ -93,6 +93,7 @@ if(ENABLE_HIP)
         CellThermoComputeGPU.h
         CellListGPU.cuh
         CellListGPU.h
+        CollisionMethod.cuh
         CommunicatorGPU.cuh
         CommunicatorGPU.h
         ParallelPlateGeometryFillerGPU.cuh
@@ -114,6 +115,7 @@ if(ENABLE_HIP)
         BounceBackNVEGPU.cu
         CellThermoComputeGPU.cu
         CellListGPU.cu
+        CollisionMethod.cu
         CommunicatorGPU.cu
         ParticleData.cu
         ParallelPlateGeometryFillerGPU.cu