Update Reduce Beam Diagnostics

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

Author: ax3l

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:
 
@@ -66,8 +69,6 @@ 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)
 * ``x_mean/min/max``, ``y_mean/min/max``, ``t_mean/min/max``
     Average / minimum / maximum beam particle position 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 = np.sqrt(1 + impactx_surrogate_ref_particle["ref_beta_gamma"] ** 2)
 beam_gamma = (
     ref_gamma
     - impactx_surrogate_reduced_diags["pt_mean"]
-    * impactx_surrogate_reduced_diags["ref_beta_gamma"]
+    * impactx_surrogate_ref_particle["ref_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" << " "
                              << "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,15 @@ 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") << " "
+            file_handler << step << " "
                          << 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.cpp

@@ -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.cpp

@@ -257,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 );