Add Total Bunch Attributes to Monitor

Add our reduced beam diagnostics (moments of the particle bunch) as metadata to openPMD files (the beam monitor element).
ECP-WarpX · Apr 18, 2024 · 5895560 · 5895560
1 parent 231d6bb
commit 5895560
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 5 deletions.
diff --git a/src/particles/diagnostics/ReducedBeamCharacteristics.H b/src/particles/diagnostics/ReducedBeamCharacteristics.H
@@ -21,7 +21,9 @@
 namespace impactx::diagnostics
 {
     /** Compute momenta of the beam distribution
-      */
+     *
+     * This uses an MPI Allreduce and returns a result on all ranks.
+     */
     std::unordered_map<std::string, amrex::ParticleReal>
     reduced_beam_characteristics (ImpactXParticleContainer const & pc);
 

diff --git a/src/particles/diagnostics/ReducedBeamCharacteristics.cpp b/src/particles/diagnostics/ReducedBeamCharacteristics.cpp
@@ -205,11 +205,11 @@ namespace impactx::diagnostics
             values_per_rank_2nd[i] = amrex::get<i>(r2);
         });
 
-        // reduced sum over mpi ranks (reduce to IO rank)
+        // reduced sum over mpi ranks (allreduce)
         amrex::ParallelDescriptor::ReduceRealSum(
             values_per_rank_2nd.data(),
             values_per_rank_2nd.size(),
-            amrex::ParallelDescriptor::IOProcessorNumber()
+            amrex::ParallelDescriptor::Communicator()
         );
 
         // minimum values

diff --git a/src/particles/elements/diagnostics/openPMD.H b/src/particles/elements/diagnostics/openPMD.H
@@ -18,6 +18,7 @@
 
 #include <any>
 #include <string>
+#include <unordered_map>
 #include <vector>
 
 
@@ -161,6 +162,12 @@ namespace detail
          */
         std::vector<uint64_t> m_offset;
 
+        /** Reduced Beam Characteristics
+         *
+         * In situ calculated particle bunch moments.
+         */
+        std::unordered_map<std::string, amrex::ParticleReal> m_rbc;
+
     };
 
     /** Calculate additional particle properties.

diff --git a/src/particles/elements/diagnostics/openPMD.cpp b/src/particles/elements/diagnostics/openPMD.cpp
@@ -11,6 +11,7 @@
 #include "openPMD.H"
 #include "ImpactXVersion.H"
 #include "particles/ImpactXParticleContainer.H"
+#include "particles/diagnostics/ReducedBeamCharacteristics.H"
 
 #include <AMReX.H>
 #include <AMReX_BLProfiler.H>
@@ -265,8 +266,14 @@ namespace detail
         beam.setAttribute( "py_ref", ref_part.py );
         beam.setAttribute( "pz_ref", ref_part.pz );
         beam.setAttribute( "pt_ref", ref_part.pt );
-        beam.setAttribute( "mass", ref_part.mass );
-        beam.setAttribute( "charge", ref_part.charge );
+        beam.setAttribute( "mass_ref", ref_part.mass );
+        beam.setAttribute( "charge_ref", ref_part.charge );
+
+        // total particle bunch information
+        //   @see impactx::diagnostics::reduced_beam_characteristics
+        for(const auto &kv : m_rbc) {
+            beam.setAttribute(kv.first, kv.second);
+        }
 
         // openPMD coarse position: for global coordinates
         {
@@ -313,6 +320,10 @@ namespace detail
         // optional: add and calculate additional particle properties
         add_optional_properties(m_series_name, pc);
 
+        // optional: calculate total particle bunch information
+        m_rbc.clear();
+        m_rbc = diagnostics::reduced_beam_characteristics(pc);
+
         // component names
         std::vector<std::string> real_soa_names = pc.RealSoA_names();
         std::vector<std::string> int_soa_names = pc.intSoA_names();