From aaa148cadfd88a870efdc8fbb2d24f68a95fdec1 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Wed, 22 May 2024 21:22:01 -0500 Subject: [PATCH] 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 * Formatting * Text Output of RBC: Keep `s` column Position on the beamline/ring is important for quick processing. --------- Co-authored-by: Ryan Sandberg --- docs/source/dataanalysis/dataanalysis.rst | 11 ++++++---- .../visualize_ml_surrogate_15_stage.py | 20 +++++++++---------- .../diagnostics/DiagnosticOutput.cpp | 12 ++++++++--- .../diagnostics/ReducedBeamCharacteristics.H | 4 +++- .../ReducedBeamCharacteristics.cpp | 10 ++++------ src/particles/elements/diagnostics/openPMD.H | 7 +++++++ .../elements/diagnostics/openPMD.cpp | 16 +++++++++++++-- 7 files changed, 54 insertions(+), 26 deletions(-) diff --git a/docs/source/dataanalysis/dataanalysis.rst b/docs/source/dataanalysis/dataanalysis.rst index 22df2b073..1f2f00641 100644 --- a/docs/source/dataanalysis/dataanalysis.rst +++ b/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 `. 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`` diff --git a/examples/pytorch_surrogate_model/visualize_ml_surrogate_15_stage.py b/examples/pytorch_surrogate_model/visualize_ml_surrogate_15_stage.py index 54b94fc51..a7ffd5097 100644 --- a/examples/pytorch_surrogate_model/visualize_ml_surrogate_15_stage.py +++ b/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 diff --git a/src/particles/diagnostics/DiagnosticOutput.cpp b/src/particles/diagnostics/DiagnosticOutput.cpp index 1a82e4857..1b0c6190b 100644 --- a/src/particles/diagnostics/DiagnosticOutput.cpp +++ b/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,6 +80,7 @@ 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) @@ -84,7 +88,9 @@ namespace impactx::diagnostics std::unordered_map 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") << " " diff --git a/src/particles/diagnostics/ReducedBeamCharacteristics.H b/src/particles/diagnostics/ReducedBeamCharacteristics.H index dcabb4e05..5fddac0a2 100644 --- 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 reduced_beam_characteristics (ImpactXParticleContainer const & pc); diff --git a/src/particles/diagnostics/ReducedBeamCharacteristics.cpp b/src/particles/diagnostics/ReducedBeamCharacteristics.cpp index 471ebd818..c0ea9ed75 100644 --- 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(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 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; } diff --git a/src/particles/elements/diagnostics/openPMD.H b/src/particles/elements/diagnostics/openPMD.H index c340bf855..04eb1a925 100644 --- a/src/particles/elements/diagnostics/openPMD.H +++ b/src/particles/elements/diagnostics/openPMD.H @@ -18,6 +18,7 @@ #include #include +#include #include @@ -161,6 +162,12 @@ namespace detail */ std::vector m_offset; + /** Reduced Beam Characteristics + * + * In situ calculated particle bunch moments. + */ + std::unordered_map m_rbc; + }; /** Calculate additional particle properties. diff --git a/src/particles/elements/diagnostics/openPMD.cpp b/src/particles/elements/diagnostics/openPMD.cpp index c98202666..10813c4e0 100644 --- 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 #include @@ -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 real_soa_names = pc.RealSoA_names(); std::vector int_soa_names = pc.intSoA_names();