Add function to get total weight of all particles in WarpXParticleCon…

…tainer (#4795)
ECP-WarpX · Mar 21, 2024 · 2b50441 · 2b50441
1 parent 1bdd450
commit 2b50441
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 22 deletions.
diff --git a/Source/Diagnostics/ReducedDiags/ParticleNumber.cpp b/Source/Diagnostics/ReducedDiags/ParticleNumber.cpp
@@ -112,27 +112,13 @@ void ParticleNumber::ComputeDiags (int step)
     for (int i_s = 0; i_s < nSpecies; ++i_s)
     {
         // get WarpXParticleContainer class object
-        const auto & myspc = mypc.GetParticleContainer(i_s);
+        auto & myspc = mypc.GetParticleContainer(i_s);
 
         // Save total number of macroparticles for this species
         m_data[idx_first_species_macroparticles + i_s] = myspc.TotalNumberOfParticles();
 
-        using PType = typename WarpXParticleContainer::SuperParticleType;
-
-        // Reduction to compute sum of weights for this species
-        auto Wtot = ReduceSum( myspc,
-        [=] AMREX_GPU_HOST_DEVICE (const PType& p) -> amrex::Real
-        {
-            return p.rdata(PIdx::w);
-        });
-
-        // MPI reduction
-        amrex::ParallelDescriptor::ReduceRealSum
-            (Wtot, amrex::ParallelDescriptor::IOProcessorNumber());
-
         // Save sum of particles weight for this species
-        m_data[idx_first_species_sum_weight + i_s] = Wtot;
-
+        m_data[idx_first_species_sum_weight + i_s] = myspc.sumParticleWeight(false);
 
         // Increase total number of macroparticles and total weight (all species)
         m_data[idx_total_macroparticles] += m_data[idx_first_species_macroparticles + i_s];

diff --git a/Source/Particles/WarpXParticleContainer.H b/Source/Particles/WarpXParticleContainer.H
@@ -280,6 +280,7 @@ public:
     /// This is needed when solving Poisson's equation with periodic boundary conditions.
     ///
     amrex::ParticleReal sumParticleCharge(bool local = false);
+    amrex::ParticleReal sumParticleWeight(bool local = false);
 
     std::array<amrex::ParticleReal, 3> meanParticleVelocity(bool local = false);
 

diff --git a/Source/Particles/WarpXParticleContainer.cpp b/Source/Particles/WarpXParticleContainer.cpp
@@ -1306,9 +1306,9 @@ WarpXParticleContainer::GetChargeDensity (int lev, bool local)
     return rho;
 }
 
-amrex::ParticleReal WarpXParticleContainer::sumParticleCharge(bool local) {
+amrex::ParticleReal WarpXParticleContainer::sumParticleWeight(bool local) {
 
-    amrex::ParticleReal total_charge = 0.0;
+    amrex::ParticleReal total_weight = 0.0;
     ReduceOps<ReduceOpSum> reduce_op;
     ReduceData<ParticleReal> reduce_data(reduce_op);
 
@@ -1328,11 +1328,15 @@ amrex::ParticleReal WarpXParticleContainer::sumParticleCharge(bool local) {
         }
     }
 
-    total_charge = get<0>(reduce_data.value());
+    total_weight = get<0>(reduce_data.value());
+
+    if (!local) { ParallelDescriptor::ReduceRealSum(total_weight); }
+    return total_weight;
+}
+
+amrex::ParticleReal WarpXParticleContainer::sumParticleCharge(bool local) {
 
-    if (!local) { ParallelDescriptor::ReduceRealSum(total_charge); }
-    total_charge *= this->charge;
-    return total_charge;
+    return this->sumParticleWeight(local) * this->charge;
 }
 
 std::array<ParticleReal, 3> WarpXParticleContainer::meanParticleVelocity(bool local) {

diff --git a/Source/Python/Particles/WarpXParticleContainer.cpp b/Source/Python/Particles/WarpXParticleContainer.cpp
@@ -109,6 +109,10 @@ void init_WarpXParticleContainer (py::module& m)
             &WarpXParticleContainer::TotalNumberOfParticles,
             py::arg("valid_particles_only"), py::arg("local")
         )
+        .def("sum_particle_weight",
+            &WarpXParticleContainer::sumParticleWeight,
+            py::arg("local")
+        )
         .def("sum_particle_charge",
             &WarpXParticleContainer::sumParticleCharge,
             py::arg("local")