Add Total Bunch Attributes to Monitor (#584)

* 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).

* Update Reduce Beam Diagnostics

Remove ref particle quantities.
Update examples and docs.
Add beta and gamma consistently to ref particle output.

* Finalize: Vis Surrogate Updates

Co-authored-by: Ryan Sandberg <RSandberg@lbl.gov>

* Formatting

* Text Output of RBC: Keep `s` column

Position on the beamline/ring is important for quick processing.

---------

Co-authored-by: Ryan Sandberg <RSandberg@lbl.gov>

docs/source/dataanalysis/dataanalysis.rst

@@ -27,6 +27,7 @@ Reference particle:
 
 * ``beta_ref`` reference particle normalized velocity :math:`\beta = v/c`
 * ``gamma_ref`` reference particle Lorentz factor :math:`\gamma = 1/\sqrt{1-\beta^2}`
+* ``beta_gamma_ref`` reference particle momentum normalized to rest mass :math:`\beta\gamma = p/(mc)`
 * ``s_ref`` integrated orbit path length, in meters
 * ``x_ref`` horizontal position x, in meters
 * ``y_ref`` vertical position y, in meters
@@ -36,8 +37,10 @@ Reference particle:
 * ``py_ref`` momentum in y, normalized to mass*c, :math:`p_y = \gamma \beta_y`
 * ``pz_ref`` momentum in z, normalized to mass*c, :math:`p_z = \gamma \beta_z`
 * ``pt_ref`` energy, normalized by rest energy, :math:`p_t = -\gamma`
-* ``mass`` reference rest mass, in kg
-* ``charge`` reference charge, in C
+* ``mass_ref`` reference rest mass, in kg
+* ``charge_ref`` reference charge, in C
+
+Bunch properties: all properties listed in :ref:`Reduced Beam Characteristics <dataanalysis-beam-characteristics>`.
 
 Example to print the integrated orbit path length ``s`` at each beam monitor position:
 
@@ -70,8 +73,8 @@ The code writes out the values in an ASCII file prefixed ``reduced_beam_characte
 
 * ``step``
     Iteration within the simulation
-* ``s``, ``ref_beta_gamma``
-    Reference particle coordinate ``s`` (unit: meter) and relativistic momentum normalized by the particle mass and the speed of light (unit: dimensionless)
+* ``s``
+    Reference particle coordinate ``s`` (unit: meter)
 * ``x_mean/min/max``, ``y_mean/min/max``, ``t_mean/min/max``
     Average / minimum / maximum particle displacement with respect to the reference particle in the dimensions of ``x``, ``y`` (transverse coordinates, unit: meter), and ``t`` (normalized time difference :math:`ct`, unit: meter)
 * ``sig_x``, ``sig_y``, ``sig_t``

examples/pytorch_surrogate_model/visualize_ml_surrogate_15_stage.py

@@ -291,14 +291,15 @@ def plot_beam_df(
 )
 args = parser.parse_args()
 
+impactx_surrogate_ref_particle = read_time_series("diags/ref_particle.*")
 impactx_surrogate_reduced_diags = read_time_series(
     "diags/reduced_beam_characteristics.*"
 )
-ref_gamma = np.sqrt(1 + impactx_surrogate_reduced_diags["ref_beta_gamma"] ** 2)
+ref_gamma = impactx_surrogate_ref_particle["gamma"]
 beam_gamma = (
     ref_gamma
     - impactx_surrogate_reduced_diags["pt_mean"]
-    * impactx_surrogate_reduced_diags["ref_beta_gamma"]
+    * impactx_surrogate_ref_particle["beta_gamma"]
 )
 beam_u = np.sqrt(beam_gamma**2 - 1)
 emit_x = impactx_surrogate_reduced_diags["emittance_x"]
@@ -315,13 +316,13 @@ def plot_beam_df(
 ax = axT[0][0]
 scale = 1e6
 ax.plot(
-    impactx_surrogate_reduced_diags["s"][ix_slice],
+    impactx_surrogate_ref_particle["s"][ix_slice],
     emit_nx[ix_slice] * scale,
     "bo",
     label="x",
 )
 ax.plot(
-    impactx_surrogate_reduced_diags["s"][ix_slice],
+    impactx_surrogate_ref_particle["s"][ix_slice],
     emit_ny[ix_slice] * scale,
     "r",
     marker=ymarker,
@@ -335,7 +336,7 @@ def plot_beam_df(
 ax = axT[0][1]
 scale = m_e * c**2 / e * 1e-9
 ax.plot(
-    impactx_surrogate_reduced_diags["s"][ix_slice],
+    impactx_surrogate_ref_particle["s"][ix_slice],
     beam_gamma[ix_slice] * scale,
     "go",
 )
@@ -346,13 +347,13 @@ def plot_beam_df(
 ax = axT[1][0]
 scale = 1e6
 ax.plot(
-    impactx_surrogate_reduced_diags["s"][ix_slice],
+    impactx_surrogate_ref_particle["s"][ix_slice],
     impactx_surrogate_reduced_diags["sig_x"][ix_slice] * scale,
     "bo",
     label="x",
 )
 ax.plot(
-    impactx_surrogate_reduced_diags["s"][ix_slice],
+    impactx_surrogate_ref_particle["s"][ix_slice],
     impactx_surrogate_reduced_diags["sig_y"][ix_slice] * scale,
     "r",
     marker=ymarker,
@@ -367,13 +368,13 @@ def plot_beam_df(
 ax = axT[1][1]
 scale = 1e3
 ax.semilogy(
-    impactx_surrogate_reduced_diags["s"][ix_slice],
+    impactx_surrogate_ref_particle["s"][ix_slice],
     impactx_surrogate_reduced_diags["sig_px"][ix_slice] * scale,
     "bo",
     label="x",
 )
 ax.semilogy(
-    impactx_surrogate_reduced_diags["s"][ix_slice],
+    impactx_surrogate_ref_particle["s"][ix_slice],
     impactx_surrogate_reduced_diags["sig_py"][ix_slice] * scale,
     "r",
     marker=ymarker,
@@ -397,7 +398,6 @@ def plot_beam_df(
     "diags/openPMD/monitor.bp", io.Access.read_only
 )
 impactx_surrogate_steps = list(beam_impactx_surrogate_series.iterations)
-impactx_surrogate_ref_particle = read_time_series("diags/ref_particle.*")
 
 millimeter = 1.0e3
 micron = 1.0e6

src/particles/diagnostics/DiagnosticOutput.cpp

@@ -41,9 +41,9 @@ namespace impactx::diagnostics
         // write file header per MPI RANK
         if (!append) {
             if (otype == OutputType::PrintRefParticle) {
-                file_handler << "step s x y z t px py pz pt\n";
+                file_handler << "step s beta gamma beta_gamma x y z t px py pz pt\n";
             } else if (otype == OutputType::PrintReducedBeamCharacteristics) {
-                file_handler << "step" << " " << "s" << " " << "ref_beta_gamma" << " "
+                file_handler << "step" << " " << "s" << " "
                              << "x_mean" << " " << "x_min" << " " << "x_max" << " "
                              << "y_mean" << " " << "y_min" << " " << "y_max" << " "
                              << "t_mean" << " " << "t_min" << " " << "t_max" << " "
@@ -65,6 +65,9 @@ namespace impactx::diagnostics
             RefPart const ref_part = pc.GetRefParticle();
 
             amrex::ParticleReal const s = ref_part.s;
+            amrex::ParticleReal const beta = ref_part.beta();
+            amrex::ParticleReal const gamma = ref_part.gamma();
+            amrex::ParticleReal const beta_gamma = ref_part.beta_gamma();
             amrex::ParticleReal const x = ref_part.x;
             amrex::ParticleReal const y = ref_part.y;
             amrex::ParticleReal const z = ref_part.z;
@@ -77,14 +80,17 @@ namespace impactx::diagnostics
             // write particle data to file
             file_handler
                     << step << " " << s << " "
+                    << beta << " " << gamma << " " << beta_gamma << " "
                     << x << " " << y << " " << z << " " << t << " "
                     << px << " " << py << " " << pz << " " << pt << "\n";
         } // if( otype == OutputType::PrintRefParticle)
         else if (otype == OutputType::PrintReducedBeamCharacteristics) {
             std::unordered_map<std::string, amrex::ParticleReal> const rbc =
                 diagnostics::reduced_beam_characteristics(pc);
 
-            file_handler << step << " " << rbc.at("s") << " " << rbc.at("ref_beta_gamma") << " "
+            amrex::ParticleReal const s = pc.GetRefParticle().s;
+
+            file_handler << step << " " << s << " "
                          << rbc.at("x_mean") << " " << rbc.at("x_min") << " " << rbc.at("x_max") << " "
                          << rbc.at("y_mean") << " " << rbc.at("y_min") << " " << rbc.at("y_max") << " "
                          << rbc.at("t_mean") << " " << rbc.at("t_min") << " " << rbc.at("t_max") << " "

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);
 

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)
-        amrex::ParallelDescriptor::ReduceRealSum(
+        // reduced sum over mpi ranks (allreduce)
+        amrex::ParallelAllReduce::Sum(
             values_per_rank_2nd.data(),
             values_per_rank_2nd.size(),
-            amrex::ParallelDescriptor::IOProcessorNumber()
+            amrex::ParallelDescriptor::Communicator()
         );
 
         // minimum values
@@ -259,8 +259,6 @@ namespace impactx::diagnostics
         amrex::ParticleReal const alpha_t = - tpt / emittance_t;
 
         std::unordered_map<std::string, amrex::ParticleReal> data;
-        data["s"] = ref_part.s;  // TODO: remove when the output gets rerouted to openPMD
-        data["ref_beta_gamma"] = ref_part.beta_gamma();  // TODO: remove when the output gets rerouted to openPMD
         data["x_mean"] = x_mean;
         data["x_min"] = x_min;
         data["x_max"] = x_max;
@@ -294,7 +292,7 @@ namespace impactx::diagnostics
         data["beta_x"] = beta_x;
         data["beta_y"] = beta_y;
         data["beta_t"] = beta_t;
-        data["charge_C"] = charge;  // TODO: remove when the output gets rerouted to openPMD
+        data["charge_C"] = charge;
 
         return data;
     }

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.

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>
@@ -256,6 +257,7 @@ namespace detail
         // reference particle information
         beam.setAttribute( "beta_ref", ref_part.beta() );
         beam.setAttribute( "gamma_ref", ref_part.gamma() );
+        beam.setAttribute( "beta_gamma_ref", ref_part.beta_gamma() );
         beam.setAttribute( "s_ref", ref_part.s );
         beam.setAttribute( "x_ref", ref_part.x );
         beam.setAttribute( "y_ref", ref_part.y );
@@ -265,8 +267,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 +321,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();