From 8f0350362ebe72ae572b0893e20ff657eec926b6 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 11 Feb 2025 10:14:13 -0800 Subject: [PATCH] openPMD Beam Input via `Source` Element (#820) * `Source` Element for openPMD * Rename Source Files of `BeamMonitor` * Deduplicate openPMD Logic Helpers * Docs * Example * Doc: Fix Copy-Paste --- docs/source/usage/examples.rst | 1 + docs/source/usage/parameters.rst | 14 + docs/source/usage/python.rst | 9 + examples/CMakeLists.txt | 22 + examples/solenoid_restart/README.rst | 50 ++ .../solenoid_restart/analysis_solenoid.py | 96 ++++ examples/solenoid_restart/input_solenoid.in | 44 ++ examples/solenoid_restart/run_solenoid.py | 48 ++ src/elements/All.H | 6 +- src/elements/CMakeLists.txt | 1 + src/elements/Source.H | 93 ++++ src/elements/Source.cpp | 128 +++++ .../diagnostics/AdditionalProperties.cpp | 3 +- src/elements/diagnostics/BeamMonitor.H | 210 +++++++++ src/elements/diagnostics/BeamMonitor.cpp | 432 +++++++++++++++++ src/elements/diagnostics/CMakeLists.txt | 9 +- src/elements/diagnostics/openPMD.H | 211 +-------- src/elements/diagnostics/openPMD.cpp | 440 +----------------- src/initialization/InitDistribution.cpp | 34 +- src/initialization/InitElement.cpp | 11 + src/initialization/Validate.cpp | 15 +- src/python/elements.cpp | 35 ++ 22 files changed, 1275 insertions(+), 637 deletions(-) create mode 100644 examples/solenoid_restart/README.rst create mode 100755 examples/solenoid_restart/analysis_solenoid.py create mode 100644 examples/solenoid_restart/input_solenoid.in create mode 100755 examples/solenoid_restart/run_solenoid.py create mode 100644 src/elements/Source.H create mode 100644 src/elements/Source.cpp create mode 100644 src/elements/diagnostics/BeamMonitor.H create mode 100644 src/elements/diagnostics/BeamMonitor.cpp diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst index cfc5f5d8e..f038a2bf8 100644 --- a/docs/source/usage/examples.rst +++ b/docs/source/usage/examples.rst @@ -83,6 +83,7 @@ Channels & Rings :maxdepth: 1 examples/fodo_channel/README.rst + examples/solenoid_restart/README.rst Lattice Design & Optimization diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 3a31c7da5..f2399550b 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -627,6 +627,20 @@ This requires these additional parameters: * ``.nslice`` (``integer``) number of slices used for the application of space charge (default: ``1``) +``source`` +^^^^^^^^^^^ + +``source`` for a particle source. +Typically at the beginning of a beam line. + +Currently, this only supports openPMD files from our ``beam_monitor``. + +* ``.distribution`` (``string``) + Distribution type of particles in the source. currently, only ``"openPMD"`` is supported +* ``.openpmd_path`` (``string``) + path to the openPMD series + + ``tapered_pl`` ^^^^^^^^^^^^^^ diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 242e8abbe..5f3932038 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -738,6 +738,15 @@ This module provides elements for the accelerator lattice. Scale factor (in meters^(1/2)) of the IOTA nonlinear magnetic insert element used for computing H and I. +.. py:class:: impactx.elements.Source(distribution, openpmd_path, name) + + A particle source. + Currently, this only supports openPMD files from our :py:class:`impactx.elements.BeamMonitor` + + :param distribution: Distribution type of particles in the source. currently, only ``"openPMD"`` is supported + :param openpmd_path: path to the openPMD series + :param name: an optional name for the element + .. py:class:: impactx.elements.Programmable(ds=0.0, nslice=1, name=None) A programmable beam optics element. diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index ddd21a872..370f3e6ce 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -624,6 +624,28 @@ add_impactx_test(solenoid.MADX.py ) +# Ideal, Hard-Edge Solenoid (Restart + 270 degrees) ########################### +# +# w/o space charge +add_impactx_test(solenoid.restart + examples/solenoid_restart/input_solenoid.in + OFF # ImpactX MPI-parallel + examples/solenoid_restart/analysis_solenoid.py + OFF # no plot script yet +) +set_property(TEST solenoid.restart.run APPEND PROPERTY DEPENDS "solenoid.run") + +add_impactx_test(solenoid.restart.py + examples/solenoid_restart/run_solenoid.py + OFF # ImpactX MPI-parallel + examples/solenoid_restart/analysis_solenoid.py + OFF # no plot script yet +) +if(ImpactX_PYTHON) + set_property(TEST solenoid.restart.py.run APPEND PROPERTY DEPENDS "solenoid.py.run") +endif() + + # Pole-face rotation ########################################################## # # w/o space charge diff --git a/examples/solenoid_restart/README.rst b/examples/solenoid_restart/README.rst new file mode 100644 index 000000000..29eb20daf --- /dev/null +++ b/examples/solenoid_restart/README.rst @@ -0,0 +1,50 @@ +.. _examples-solenoid-restart: + +Solenoid channel (Restart) +========================== + +This is the same example as the :ref:`proton beam undergoing 90 deg X-Y rotation in an ideal solenoid channel `, but it restarts the resulting beam and rotates it another 3 channel periods to the initial X-Y conditions. + +The solenoid magnetic field corresponds to B = 2 T. + +The second moments of the transverse particle distribution after the solenoid channel are rotated by 90 (start) + 270 = 360 degrees: the final transverse moments should coincide with the +initial transverse moments to within the level expected due to noise due to statistical sampling. + +In this test, the initial and final values of :math:`\sigma_x`, :math:`\sigma_y`, :math:`\sigma_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values. + + +Run +--- + +This example can be run **either** as: + +* **Python** script: ``python3 run_solenoid.py`` or +* ImpactX **executable** using an input file: ``impactx input_solenoid.in`` + +For `MPI-parallel `__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system. + +.. tab-set:: + + .. tab-item:: Python: Script + + .. literalinclude:: run_solenoid.py + :language: python3 + :caption: You can copy this file from ``examples/solenoid_restart/run_solenoid.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_solenoid.in + :language: ini + :caption: You can copy this file from ``examples/solenoid_restart/input_solenoid.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_solenoid.py`` + + .. literalinclude:: analysis_solenoid.py + :language: python3 + :caption: You can copy this file from ``examples/solenoid_restart/analysis_solenoid.py``. diff --git a/examples/solenoid_restart/analysis_solenoid.py b/examples/solenoid_restart/analysis_solenoid.py new file mode 100755 index 000000000..c5bb2ec55 --- /dev/null +++ b/examples/solenoid_restart/analysis_solenoid.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# + +import numpy as np +import openpmd_api as io +from scipy.stats import moment + + +def get_moments(beam): + """Calculate standard deviations of beam position & momenta + and emittance values + + Returns + ------- + sigx, sigy, sigt, emittance_x, emittance_y, emittance_t + """ + sigx = moment(beam["position_x"], moment=2) ** 0.5 # variance -> std dev. + sigpx = moment(beam["momentum_x"], moment=2) ** 0.5 + sigy = moment(beam["position_y"], moment=2) ** 0.5 + sigpy = moment(beam["momentum_y"], moment=2) ** 0.5 + sigt = moment(beam["position_t"], moment=2) ** 0.5 + sigpt = moment(beam["momentum_t"], moment=2) ** 0.5 + + epstrms = beam.cov(ddof=0) + emittance_x = (sigx**2 * sigpx**2 - epstrms["position_x"]["momentum_x"] ** 2) ** 0.5 + emittance_y = (sigy**2 * sigpy**2 - epstrms["position_y"]["momentum_y"] ** 2) ** 0.5 + emittance_t = (sigt**2 * sigpt**2 - epstrms["position_t"]["momentum_t"] ** 2) ** 0.5 + + return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t) + + +# initial/final beam +series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only) +first_step = list(series.iterations)[0] +last_step = list(series.iterations)[-1] +initial = series.iterations[first_step].particles["beam"].to_df() +final = series.iterations[last_step].particles["beam"].to_df() + +# compare number of particles +num_particles = 10000 +assert num_particles == len(initial) +assert num_particles == len(final) + +print("Initial Beam:") +sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(initial) +print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}") +print( + f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}" +) + +atol = 0.0 # ignored +rtol = 1.3 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t], + [ + 2.205510139392e-3, + 1.559531175539e-3, + 6.404930308742e-3, + 2.0e-6, + 1.0e-6, + 1.0e-6, + ], + rtol=rtol, + atol=atol, +) + +print("") +print("Final Beam:") +sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final) +print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}") +print( + f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}" +) + +atol = 0.0 # ignored +rtol = 2.3 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sigx, sigy, emittance_x, emittance_y, emittance_t], + [ + 1.559531175539e-3, + 2.205510139392e-3, + 1.0e-6, + 2.0e-6, + 1.0e-6, + ], + rtol=rtol, + atol=atol, +) diff --git a/examples/solenoid_restart/input_solenoid.in b/examples/solenoid_restart/input_solenoid.in new file mode 100644 index 000000000..bd2e9df1f --- /dev/null +++ b/examples/solenoid_restart/input_solenoid.in @@ -0,0 +1,44 @@ +############################################################################### +# Reference Particle +############################################################################### +beam.kin_energy = 250.0 +beam.particle = proton + +# ignored +beam.charge = 1.0e-9 +beam.units = static +beam.distribution = empty + + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.elements = source monitor sols monitor +lattice.nslice = 1 + +source.type = source +source.distribution = openPMD +source.openpmd_path = "../solenoid/diags/openPMD/monitor.h5" + +monitor.type = beam_monitor +monitor.backend = h5 + +sol1.type = solenoid +sol1.ds = 3.820395 +sol1.ks = 0.8223219329893234 + +sols.type = line +sols.elements = sol1 +sols.repeat = 3 + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true diff --git a/examples/solenoid_restart/run_solenoid.py b/examples/solenoid_restart/run_solenoid.py new file mode 100755 index 000000000..8627cd60c --- /dev/null +++ b/examples/solenoid_restart/run_solenoid.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Marco Garten, Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from impactx import ImpactX, elements, push + +sim = ImpactX() + +# set numerical parameters and IO control +sim.particle_shape = 2 # B-spline order +sim.space_charge = False +# sim.diagnostics = False # benchmarking +sim.slice_step_diagnostics = True + +# domain decomposition & space charge mesh +sim.init_grids() + +# load a 250 MeV proton beam with an initial +# horizontal rms emittance of 1 um and an +# initial vertical rms emittance of 2 um +kin_energy_MeV = 250.0 # reference energy + +# reference particle +pc = sim.particle_container() +ref = pc.ref_particle() +ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_kin_energy_MeV(kin_energy_MeV) + +# load particle bunch from file +push(pc, elements.Source("openPMD", "../solenoid.py/diags/openPMD/monitor.h5")) + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5") + +# design the accelerator lattice +sol1 = elements.Sol(name="sol1", ds=3.820395, ks=0.8223219329893234) +sim.lattice.append(monitor) +sim.lattice.extend([sol1] * 3) +sim.lattice.append(monitor) + +# run simulation +sim.track_particles() + +# clean shutdown +sim.finalize() diff --git a/src/elements/All.H b/src/elements/All.H index 99dcb3f03..c40cfe924 100644 --- a/src/elements/All.H +++ b/src/elements/All.H @@ -35,12 +35,13 @@ #include "RFCavity.H" #include "Sbend.H" #include "ShortRF.H" -#include "Sol.H" #include "SoftSol.H" #include "SoftQuad.H" +#include "Sol.H" +#include "Source.H" #include "TaperedPL.H" #include "ThinDipole.H" -#include "diagnostics/openPMD.H" +#include "diagnostics/BeamMonitor.H" #include @@ -77,6 +78,7 @@ namespace impactx::elements SoftSolenoid, SoftQuadrupole, Sol, + Source, TaperedPL, ThinDipole >; diff --git a/src/elements/CMakeLists.txt b/src/elements/CMakeLists.txt index dceec99c3..08aaa1844 100644 --- a/src/elements/CMakeLists.txt +++ b/src/elements/CMakeLists.txt @@ -2,6 +2,7 @@ target_sources(lib PRIVATE Aperture.cpp Programmable.cpp + Source.cpp ) add_subdirectory(diagnostics) diff --git a/src/elements/Source.H b/src/elements/Source.H new file mode 100644 index 000000000..817b333c1 --- /dev/null +++ b/src/elements/Source.H @@ -0,0 +1,93 @@ +/* Copyright 2022-2025 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_SOURCE_H +#define IMPACTX_SOURCE_H + +#include "particles/ImpactXParticleContainer.H" +#include "mixin/lineartransport.H" +#include "mixin/named.H" +#include "mixin/nofinalize.H" +#include "mixin/thin.H" + +#include +#include + + +namespace impactx::elements +{ + struct Source + : public mixin::Named, + public mixin::LinearTransport, + public mixin::Thin, + public mixin::NoFinalize + { + static constexpr auto type = "Source"; + using PType = ImpactXParticleContainer::ParticleType; + + /** A particle source + * + * @param distribution Must read "openPMD" for this constructor + * @param openpmd_path path to openPMD series as accepted by openPMD::Series + * @param name a user defined and not necessarily unique name of the element + */ + Source ( + std::string distribution, + std::string openpmd_path, + std::optional name = std::nullopt + ) + : Named(std::move(name)), + m_distribution(distribution), + m_series_name(std::move(openpmd_path)) + { + if (distribution != "openPMD") { + throw std::runtime_error("Only 'openPMD' distribution is supported if openpmd_path is provided!"); + } + } + + /** Initialize/Load particles. + * + * Particles are relative to the reference particle. + * + * @param[in,out] pc particle container to push + * @param[in] step global step for diagnostics + * @param[in] period for periodic lattices, this is the current period (turn or cycle) + */ + void operator() ( + ImpactXParticleContainer & pc, + [[maybe_unused]] int step, + [[maybe_unused]] int period + ); + + /** This function returns the linear transport map. + * + * @param[in] refpart reference particle + * @returns 6x6 transport matrix + */ + AMREX_GPU_HOST AMREX_FORCE_INLINE + Map6x6 + transport_map ([[maybe_unused]] RefPart const & AMREX_RESTRICT refpart) const + { + // nothing to do + return Map6x6::Identity(); + } + + /** This does nothing to the reference particle. */ + using Thin::operator(); + + /** This pushes the covariance matrix. */ + using LinearTransport::operator(); + + std::string m_distribution; //! Distribution type of particles in the source + std::string m_series_name; //! openPMD filename + }; + +} // namespace impactx + +#endif // IMPACTX_SOURCE_H diff --git a/src/elements/Source.cpp b/src/elements/Source.cpp new file mode 100644 index 000000000..2984906ef --- /dev/null +++ b/src/elements/Source.cpp @@ -0,0 +1,128 @@ +/* Copyright 2022-2025 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#include "elements/Source.H" + +#include + +#ifdef ImpactX_USE_OPENPMD +# include "elements/diagnostics/openPMD.H" +# include +namespace io = openPMD; +#else +# include +#endif + + +namespace impactx::elements +{ + void + Source::operator() ( + ImpactXParticleContainer & pc, + int, + int + ) + { +#ifdef ImpactX_USE_OPENPMD + auto series = io::Series(m_series_name, io::Access::READ_ONLY +# if openPMD_HAVE_MPI==1 + , amrex::ParallelDescriptor::Communicator() +# endif + ); + + // read the last iteration (highest openPMD iteration) + // TODO: later we can make this an option + int const read_iteration = std::prev(series.iterations.end())->first; + io::Iteration iteration = series.iterations[read_iteration]; + + // TODO: later we can make the particle species name an option + std::string const species_name = "beam"; + io::ParticleSpecies beam = iteration.particles[species_name]; + // TODO: later we can make the bunch charge an option (i.e., allow rescaling a distribution) + amrex::ParticleReal bunch_charge = beam.getAttribute("charge_C").get(); + + auto const scalar = openPMD::RecordComponent::SCALAR; + auto const getComponentRecord = [&beam](std::string comp_name) { + return diagnostics::detail::get_component_record(beam, std::move(comp_name)); + }; + + int const npart = beam["id"][scalar].getExtent()[0]; // how many particles to read total + + // TODO: read reference particle (optional?) + + // read the particles + + // Logic: We initialize 1/Nth of particles, independent of their + // position, per MPI rank. We then measure the distribution's spatial + // extent, create a grid, resize it to fit the beam, and then + // redistribute particles so that they reside on the correct MPI rank. + int const myproc = amrex::ParallelDescriptor::MyProc(); + int const nprocs = amrex::ParallelDescriptor::NProcs(); + int const navg = npart / nprocs; // note: integer division + int const nleft = npart - navg * nprocs; + std::uint64_t const npart_this_proc = (myproc < nleft) ? navg+1 : navg; // add 1 to each proc until distributed + std::uint64_t npart_before_this_proc = (myproc < nleft) ? (navg+1) * myproc : navg * myproc + nleft; + auto const rel_part_this_proc = + amrex::ParticleReal(npart_this_proc) / amrex::ParticleReal(npart); + + // alloc data for particle attributes + std::map> pinned_SoA; + std::vector real_soa_names = pc.GetRealSoANames(); + for (auto real_idx = 0; real_idx < pc.NumRealComps(); real_idx++) { + pinned_SoA[real_soa_names.at(real_idx)].resize(npart_this_proc); + } + + // read from file + // idcpu: TODO + //amrex::Gpu::PinnedVector pinned_idcpu(npart_this_proc); + //beam["id"][scalar].loadChunkRaw(pinned_idcpu.data(), {npart_before_this_proc}, {npart_this_proc}); + // SoA: Real + { + for (auto real_idx = 0; real_idx < pc.NumRealComps(); real_idx++) { + auto const component_name = real_soa_names.at(real_idx); + getComponentRecord(component_name).loadChunkRaw( + pinned_SoA[component_name].data(), + {npart_before_this_proc}, + {npart_this_proc} + ); + } + } + // SoA: Int + std::vector int_soa_names = pc.GetIntSoANames(); + static_assert(IntSoA::nattribs == 0); // not yet used + if (!int_soa_names.empty()) + throw std::runtime_error("BeamMonitor: int_soa_names output not yet implemented!"); + + series.flush(); + series.close(); + + // copy to device + std::map> d_SoA; + for (auto real_idx = 0; real_idx < pc.NumRealComps(); real_idx++) { + auto const component_name = real_soa_names.at(real_idx); + d_SoA[component_name].resize(npart_this_proc); + amrex::Gpu::copyAsync(amrex::Gpu::hostToDevice, pinned_SoA[component_name].begin(), pinned_SoA[component_name].end(), d_SoA[component_name].begin()); + } + + // finalize distributions and deallocate temporary device global memory + amrex::Gpu::streamSynchronize(); + + // TODO: at this point, we ignore the "id", "qm" and "weighting" in the file. Could be improved + pc.AddNParticles(d_SoA["position_x"], d_SoA["position_y"], d_SoA["position_t"], + d_SoA["momentum_x"], d_SoA["momentum_y"], d_SoA["momentum_t"], + pc.GetRefParticle().qm_ratio_SI(), + bunch_charge * rel_part_this_proc); + +#else // ImpactX_USE_OPENPMD + amrex::ignore_unused(pc); + throw std::runtime_error("BeamMonitor: openPMD not compiled"); +#endif // ImpactX_USE_OPENPMD + } + +} // namespace impactx diff --git a/src/elements/diagnostics/AdditionalProperties.cpp b/src/elements/diagnostics/AdditionalProperties.cpp index ec081622f..db8f2f527 100644 --- a/src/elements/diagnostics/AdditionalProperties.cpp +++ b/src/elements/diagnostics/AdditionalProperties.cpp @@ -8,7 +8,8 @@ * License: BSD-3-Clause-LBNL */ -#include "openPMD.H" +#include "BeamMonitor.H" + #include "diagnostics/NonlinearLensInvariants.H" #include "particles/ImpactXParticleContainer.H" diff --git a/src/elements/diagnostics/BeamMonitor.H b/src/elements/diagnostics/BeamMonitor.H new file mode 100644 index 000000000..a5289c4da --- /dev/null +++ b/src/elements/diagnostics/BeamMonitor.H @@ -0,0 +1,210 @@ +/* Copyright 2022-2023 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_ELEMENTS_DIAGS_BEAMMONITOR_H +#define IMPACTX_ELEMENTS_DIAGS_BEAMMONITOR_H + +#include "elements/mixin/thin.H" +#include "particles/CovarianceMatrix.H" +#include "particles/ImpactXParticleContainer.H" +#include "elements/mixin/lineartransport.H" + +#include +#include + +#include +#include +#include +#include + + +namespace impactx::elements::diagnostics +{ +namespace detail +{ + class ImpactXParticleCounter { + public: + using ParticleContainer = typename ImpactXParticleContainer::ContainerLike; + using ParticleIter = typename ParticleContainer::ParIterType; + + ImpactXParticleCounter (ParticleContainer & pc); + + unsigned long GetTotalNumParticles () { return m_Total; } + + std::vector m_ParticleOffsetAtRank; + std::vector m_ParticleSizeAtRank; + private: + /** get the offset in the overall particle id collection + * + * @param[out] numParticles particles on this processor / amrex fab + * @param[out] offset particle offset over all, mpi-global amrex fabs + * @param[out] sum number of all particles from all amrex fabs + */ + void GetParticleOffsetOfProcessor (const long &numParticles, + unsigned long long &offset, + unsigned long long &sum) const; + + int m_MPIRank = 0; + int m_MPISize = 1; + + unsigned long long m_Total = 0; + + std::vector m_ParticleCounterByLevel; + }; +} // namespace detail + + /** This element writes the particle beam out to openPMD data. + * + * This class behaves like a singleton if constructed with the + * same series name as an existing instance. + */ + struct BeamMonitor + : public mixin::LinearTransport, + public mixin::Thin + { + static constexpr auto type = "BeamMonitor"; + using PType = typename ImpactXParticleContainer::ParticleType; + using PinnedContainer = typename ImpactXParticleContainer::ContainerLike; + + /** This element writes the particle beam out to openPMD data. + * + * Elements with the same series name are identical. + * + * @param series_name name of the data series, usually the element name + * @param backend file format backend for openPMD, e.g., "bp" or "h5" + * @param encoding openPMD iteration encoding: "v"ariable based, "f"ile based, "g"roup based (default) + * @param period_sample_intervals for periodic lattice, only output every Nth period (turn) + */ + BeamMonitor (std::string series_name, std::string backend="default", std::string encoding="g", int period_sample_intervals=1); + + BeamMonitor (BeamMonitor const & other) = default; + BeamMonitor (BeamMonitor && other) = default; + BeamMonitor& operator= (BeamMonitor const & other) = default; + BeamMonitor& operator= (BeamMonitor && other) = default; + + /** Prepare entering the element before starting push logic. + * + * And write reference particle. + * + * @param[in] pc particle container + * @param[in] real_soa_names ParticleReal component names + * @param[in] int_soa_names integer component names + * @param[in] ref_part reference particle + * @param[in] step global step for diagnostics + */ + void prepare ( + PinnedContainer & pc, + std::vector const & real_soa_names, + std::vector const & int_soa_names, + RefPart const & ref_part, + int step + ); + + /** Dump all particles. + * + * Particles are relative to the reference particle. + * + * @param[in,out] pc particle container to push + * @param[in] step global step for diagnostics + * @param[in] period for periodic lattices, this is the current period (turn or cycle) + */ + void operator() ( + ImpactXParticleContainer & pc, + int step, + int period + ); + + /** Write a tile of particles + * + * @param pti particle tile iterator + * @param[in] real_soa_names ParticleReal component names + * @param[in] int_soa_names integer component names + * @param ref_part reference particle + */ + void operator() ( + PinnedContainer::ParIterType & pti, + std::vector const & real_soa_names, + std::vector const & int_soa_names, + RefPart const & ref_part + ); + + /** This does nothing to the reference particle. */ + using Thin::operator(); + + /** This pushes the covariance matrix. */ + using LinearTransport::operator(); + + /** This function returns the linear transport map. + * + * @param[in] refpart reference particle + * @returns 6x6 transport matrix + */ + AMREX_GPU_HOST AMREX_FORCE_INLINE + Map6x6 + transport_map ([[maybe_unused]] RefPart const & AMREX_RESTRICT refpart) const + { + // nothing to do + return Map6x6::Identity(); + } + + /** Get the name of the series + * + * Elements with the same series name are identical. + */ + std::string series_name () const { return m_series_name; } + + /** track all m_series_name instances + * + * Ensure m_series is the same for the same name. + */ + static inline std::map m_unique_series = {}; + + /** Close and deallocate all data series and backends. + */ + void + finalize (); + + private: + std::string m_series_name; //! openPMD filename + std::string m_OpenPMDFileType; //! openPMD backend: usually HDF5 (h5) or ADIOS2 (bp/bp4/bp5) or ADIOS2 SST (sst) + std::any m_series; //! openPMD::Series that holds potentially multiple outputs + int m_step = 0; //! global step for output + + int m_file_min_digits = 6; //! minimum number of digits to iteration number in file name + + int m_period_sample_intervals = 1; //! only output every Nth period (turn or cycle) of periodic lattices + + /** This rank's offset in the MPI-global particle array, by level + * + * This MUST be updated by prepare() before the next step's output. + */ + std::vector m_offset; + + /** Reduced Beam Characteristics + * + * In situ calculated particle bunch moments. + */ + std::unordered_map m_rbc; + + }; + + /** Calculate additional particle properties. + * + * @param[in] element_name name of the element (for input queries) + * @param[inout] pc particle container + */ + void + add_optional_properties ( + std::string const & element_name, + ImpactXParticleContainer & pc + ); + +} // namespace impactx::diagnostics + +#endif // IMPACTX_ELEMENTS_DIAGS_BEAMMONITOR_H diff --git a/src/elements/diagnostics/BeamMonitor.cpp b/src/elements/diagnostics/BeamMonitor.cpp new file mode 100644 index 000000000..328c088e7 --- /dev/null +++ b/src/elements/diagnostics/BeamMonitor.cpp @@ -0,0 +1,432 @@ +/* Copyright 2022-2023 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Axel Huebl + * License: BSD-3-Clause-LBNL + */ + +#include "BeamMonitor.H" + +#include "ImpactXVersion.H" +#include "particles/ImpactXParticleContainer.H" +#include "diagnostics/ReducedBeamCharacteristics.H" + +#include +#include +#include +#include + +#ifdef ImpactX_USE_OPENPMD +# include "elements/diagnostics/openPMD.H" +# include +namespace io = openPMD; +#endif + +#include +#include +#include + + +namespace impactx::elements::diagnostics +{ +namespace detail { + ImpactXParticleCounter::ImpactXParticleCounter (ParticleContainer & pc) + { + m_MPISize = amrex::ParallelDescriptor::NProcs(); + m_MPIRank = amrex::ParallelDescriptor::MyProc(); + + m_ParticleCounterByLevel.resize(pc.finestLevel()+1); + m_ParticleOffsetAtRank.resize(pc.finestLevel()+1); + m_ParticleSizeAtRank.resize(pc.finestLevel()+1); + + for (auto currentLevel = 0; currentLevel <= pc.finestLevel(); currentLevel++) + { + long numParticles = 0; // numParticles in this processor + + for (ParticleIter pti(pc, currentLevel); pti.isValid(); ++pti) { + auto numParticleOnTile = pti.numParticles(); + numParticles += numParticleOnTile; + } + + unsigned long long offset=0; // offset of this level + unsigned long long sum=0; // numParticles in this level (sum from all processors) + + GetParticleOffsetOfProcessor(numParticles, offset, sum); + + m_ParticleCounterByLevel[currentLevel] = sum; + m_ParticleOffsetAtRank[currentLevel] = offset; + m_ParticleSizeAtRank[currentLevel] = numParticles; + + // adjust offset, it should be numbered after particles from previous levels + for (auto lv=0; lv the particles in the comm + // sum of all particles in the comm + // + void + ImpactXParticleCounter::GetParticleOffsetOfProcessor ( + const long& numParticles, + unsigned long long& offset, + unsigned long long& sum + ) const + { + offset = 0; +#if defined(AMREX_USE_MPI) + std::vector result(m_MPISize, 0); + amrex::ParallelGather::Gather (numParticles, result.data(), -1, amrex::ParallelDescriptor::Communicator()); + + sum = 0; + int const num_results = result.size(); + for (int i=0; i(m_series); + series.close(); + m_series.reset(); + } + + // remove from unique series map + if (m_unique_series.count(m_series_name) != 0u) + m_unique_series.erase(m_series_name); +#endif // ImpactX_USE_OPENPMD + } + + BeamMonitor::BeamMonitor (std::string series_name, std::string backend, std::string encoding, int period_sample_intervals) : + m_series_name(std::move(series_name)), m_OpenPMDFileType(std::move(backend)), m_period_sample_intervals(period_sample_intervals) + { +#ifdef ImpactX_USE_OPENPMD + // pick first available backend if default is chosen + if( m_OpenPMDFileType == "default" ) +# if openPMD_HAVE_ADIOS2==1 + m_OpenPMDFileType = "bp"; +# elif openPMD_HAVE_ADIOS1==1 + m_OpenPMDFileType = "bp"; +# elif openPMD_HAVE_HDF5==1 + m_OpenPMDFileType = "h5"; +# else + m_OpenPMDFileType = "json"; +# endif + + // encoding of iterations in the series + openPMD::IterationEncoding series_encoding = openPMD::IterationEncoding::groupBased; + if ( "v" == encoding ) + series_encoding = openPMD::IterationEncoding::variableBased; + else if ( "g" == encoding ) + series_encoding = openPMD::IterationEncoding::groupBased; + else if ( "f" == encoding ) + series_encoding = openPMD::IterationEncoding::fileBased; + + amrex::ParmParse pp_diag("diag"); + // turn filter + pp_diag.queryAddWithParser("period_sample_intervals", m_period_sample_intervals); + // legacy options from other diagnostics + pp_diag.queryAddWithParser("file_min_digits", m_file_min_digits); + + // Ensure m_series is the same for the same names. + if (m_unique_series.count(m_series_name) == 0u) { + std::string filepath = "diags/openPMD/"; + filepath.append(m_series_name); + + if (series_encoding == openPMD::IterationEncoding::fileBased) + { + std::string const fileSuffix = std::string("_%0") + std::to_string(m_file_min_digits) + std::string("T"); + filepath.append(fileSuffix); + } + filepath.append(".").append(m_OpenPMDFileType); + + // transform paths for Windows +# ifdef _WIN32 + filepath = openPMD::auxiliary::replace_all(filepath, "/", "\\"); +# endif + + auto series = io::Series(filepath, io::Access::CREATE +# if openPMD_HAVE_MPI==1 + , amrex::ParallelDescriptor::Communicator() +# endif + , "adios2.engine.usesteps = true" + ); + series.setSoftware("ImpactX", IMPACTX_VERSION); + series.setIterationEncoding( series_encoding ); + m_series = series; + m_unique_series[m_series_name] = series; + } + else { + m_series = m_unique_series[m_series_name]; + } +#else + amrex::AllPrint() << "Warning: openPMD output requested but not compiled for series=" << m_series_name << "\n"; +#endif + } + + void BeamMonitor::prepare ( + PinnedContainer & pc, + std::vector const & real_soa_names, + std::vector const & int_soa_names, + RefPart const & ref_part, + int step + ) { +#ifdef ImpactX_USE_OPENPMD + m_step = step; + + // series & iteration + auto series = std::any_cast(m_series); + io::WriteIterations iterations = series.writeIterations(); + io::Iteration iteration = iterations[m_step]; + io::ParticleSpecies beam = iteration.particles["beam"]; + + // calculate & update particle offset in MPI-global particle array, per level + auto const num_levels = pc.finestLevel() + 1; + m_offset = std::vector(num_levels); + auto counter = detail::ImpactXParticleCounter(pc); + auto const np = counter.GetTotalNumParticles(); + for (auto currentLevel = 0; currentLevel < num_levels; currentLevel++) { + m_offset.at(currentLevel) = static_cast( counter.m_ParticleOffsetAtRank[currentLevel] ); + } + + // helpers to parse strings to openPMD + auto const scalar = openPMD::RecordComponent::SCALAR; + auto const getComponentRecord = [&beam](std::string comp_name) { + return detail::get_component_record(beam, std::move(comp_name)); + }; + + // define data set and metadata + io::Datatype const dtype_fl = io::determineDatatype(); + io::Datatype const dtype_ui = io::determineDatatype(); + auto d_fl = io::Dataset(dtype_fl, {np}); + auto d_ui = io::Dataset(dtype_ui, {np}); + + // 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 ); + beam.setAttribute( "z_ref", ref_part.z ); + beam.setAttribute( "t_ref", ref_part.t ); + beam.setAttribute( "px_ref", ref_part.px ); + 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", 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 + { + + beam["positionOffset"]["x"].resetDataset(d_fl); + beam["positionOffset"]["x"].makeConstant(ref_part.x); + beam["positionOffset"]["y"].resetDataset(d_fl); + beam["positionOffset"]["y"].makeConstant(ref_part.y); + beam["positionOffset"]["t"].resetDataset(d_fl); + beam["positionOffset"]["t"].makeConstant(ref_part.t); + } + + // unique, global particle index + beam["id"][scalar].resetDataset(d_ui); + + // SoA: Real + { + for (auto real_idx = 0; real_idx < pc.NumRealComps(); real_idx++) { + auto const component_name = real_soa_names.at(real_idx); + getComponentRecord(component_name).resetDataset(d_fl); + } + } + // SoA: Int + static_assert(IntSoA::nattribs == 0); // not yet used + if (!int_soa_names.empty()) + throw std::runtime_error("BeamMonitor: int_soa_names output not yet implemented!"); +#else + amrex::ignore_unused(pc, step); +#endif // ImpactX_USE_OPENPMD + } + + void + BeamMonitor::operator() ( + ImpactXParticleContainer & pc, + int step, + int period + ) + { + // filter out this turn? + if (period % m_period_sample_intervals != 0) + return; + +#ifdef ImpactX_USE_OPENPMD + std::string profile_name = "impactx::Push::" + std::string(BeamMonitor::type); + BL_PROFILE(profile_name); + + // preparing to access reference particle data: RefPart + RefPart & ref_part = pc.GetRefParticle(); + + // 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 = impactx::diagnostics::reduced_beam_characteristics(pc); + + // component names + std::vector real_soa_names = pc.GetRealSoANames(); + std::vector int_soa_names = pc.GetIntSoANames(); + + // pinned memory copy + PinnedContainer pinned_pc = pc.make_alike(); + pinned_pc.copyParticles(pc, true); // no filtering + + // TODO: filtering + /* + using SrcData = WarpXParticleContainer::ParticleTileType::ConstParticleTileDataType; + tmp.copyParticles(*pc, + [=] AMREX_GPU_HOST_DEVICE (const SrcData& src, int ip, const amrex::RandomEngine& engine) + { + const SuperParticleType& p = src.getSuperParticle(ip); + return random_filter(p, engine) * uniform_filter(p, engine) + * parser_filter(p, engine) * geometry_filter(p, engine); + }, true); + */ + + // prepare element access & write reference particle + this->prepare(pinned_pc, real_soa_names, int_soa_names, ref_part, step); + + // loop over refinement levels + int const nLevel = pinned_pc.finestLevel(); + for (int lev = 0; lev <= nLevel; ++lev) + { + // loop over all particle boxes + //using ParIt = ImpactXParticleContainer::iterator; + using ParIt = PinnedContainer::ParIterType; + // note: openPMD-api is not thread-safe, so do not run OMP parallel here + for (ParIt pti(pinned_pc, lev); pti.isValid(); ++pti) { + // write beam particles relative to reference particle + this->operator()(pti, real_soa_names, int_soa_names, ref_part); + } // end loop over all particle boxes + } // end mesh-refinement level loop + + auto series = std::any_cast(m_series); + io::WriteIterations iterations = series.writeIterations(); + io::Iteration iteration = iterations[m_step]; + + // close iteration + iteration.close(); +#else + amrex::ignore_unused(pc, step); +#endif // ImpactX_USE_OPENPMD + } + + void + BeamMonitor::operator() ( + PinnedContainer::ParIterType & pti, + std::vector const & real_soa_names, + std::vector const & int_soa_names, + RefPart const & ref_part + ) + { +#ifdef ImpactX_USE_OPENPMD + int const currentLevel = pti.GetLevel(); + + auto & offset = m_offset.at(currentLevel); // ... + + // series & iteration + auto series = std::any_cast(m_series); + io::WriteIterations iterations = series.writeIterations(); + io::Iteration iteration = iterations[m_step]; + + // writing + io::ParticleSpecies beam = iteration.particles["beam"]; + + auto const numParticleOnTile = pti.numParticles(); + uint64_t const numParticleOnTile64 = static_cast( numParticleOnTile ); + + // Do not call storeChunk() with zero-sized particle tiles: + // https://github.com/openPMD/openPMD-api/issues/1147 + //if (numParticleOnTile == 0) continue; + + auto const scalar = openPMD::RecordComponent::SCALAR; + auto const getComponentRecord = [&beam](std::string comp_name) { + return detail::get_component_record(beam, std::move(comp_name)); + }; + + // SoA + auto const& soa = pti.GetStructOfArrays(); + // particle id arrays + { + beam["id"][scalar].storeChunkRaw(soa.GetIdCPUData().data(), {offset}, {numParticleOnTile64}); + } + // SoA floating point (ParticleReal) properties + { + for (auto real_idx=0; real_idx < soa.NumRealComps(); real_idx++) { + auto const component_name = real_soa_names.at(real_idx); + getComponentRecord(component_name).storeChunkRaw( + soa.GetRealData(real_idx).data(), {offset}, {numParticleOnTile64}); + } + } + // SoA integer (int) properties (not yet used) + { + static_assert(IntSoA::nattribs == 0); // not yet used + if (!int_soa_names.empty()) + throw std::runtime_error("BeamMonitor: int_soa_names output not yet implemented!"); + /* + // comment this in once IntSoA::nattribs is > 0 + + std::copy(IntSoA::names_s.begin(), IntSoA::names_s.end(), int_soa_names.begin()); + + for (auto int_idx=0; int_idx < RealSoA::nattribs; int_idx++) { + auto const component_name = int_soa_names.at(int_idx); + getComponentRecord(component_name).storeChunkRaw( + soa.GetIntData(int_idx).data(), {offset}, {numParticleOnTile64}); + } + */ + } + + // TODO + amrex::ignore_unused(ref_part); + + // needs to be higher for next pti; must be reset for next step via prepare + offset += numParticleOnTile64; + + // TODO could be done once after all pti are processed + // TODO at that point, we could also close the iteration/step + series.flush(); +#else + amrex::ignore_unused(pti, ref_part); +#endif // ImpactX_USE_OPENPMD + } + +} // namespace impactx::diagnostics diff --git a/src/elements/diagnostics/CMakeLists.txt b/src/elements/diagnostics/CMakeLists.txt index 7806f3f81..c41a95172 100644 --- a/src/elements/diagnostics/CMakeLists.txt +++ b/src/elements/diagnostics/CMakeLists.txt @@ -1,5 +1,12 @@ target_sources(lib PRIVATE AdditionalProperties.cpp - openPMD.cpp + BeamMonitor.cpp ) + +if(ImpactX_OPENPMD) + target_sources(lib + PRIVATE + openPMD.cpp + ) +endif() diff --git a/src/elements/diagnostics/openPMD.H b/src/elements/diagnostics/openPMD.H index 4561e0cff..494eb4649 100644 --- a/src/elements/diagnostics/openPMD.H +++ b/src/elements/diagnostics/openPMD.H @@ -7,204 +7,37 @@ * Authors: Axel Huebl * License: BSD-3-Clause-LBNL */ -#ifndef IMPACTX_ELEMENTS_DIAGS_OPENPMD_H -#define IMPACTX_ELEMENTS_DIAGS_OPENPMD_H +#ifndef IMPACTX_OPENPMD_H +#define IMPACTX_OPENPMD_H -#include "elements/mixin/thin.H" -#include "particles/CovarianceMatrix.H" -#include "particles/ImpactXParticleContainer.H" -#include "elements/mixin/lineartransport.H" +#ifdef ImpactX_USE_OPENPMD +# include +#endif -#include -#include - -#include #include -#include -#include +#include -namespace impactx::elements::diagnostics -{ -namespace detail +namespace impactx::elements::diagnostics::detail { - class ImpactXParticleCounter { - public: - using ParticleContainer = typename ImpactXParticleContainer::ContainerLike; - using ParticleIter = typename ParticleContainer::ParIterType; - - ImpactXParticleCounter (ParticleContainer & pc); - - unsigned long GetTotalNumParticles () { return m_Total; } - - std::vector m_ParticleOffsetAtRank; - std::vector m_ParticleSizeAtRank; - private: - /** get the offset in the overall particle id collection - * - * @param[out] numParticles particles on this processor / amrex fab - * @param[out] offset particle offset over all, mpi-global amrex fabs - * @param[out] sum number of all particles from all amrex fabs - */ - void GetParticleOffsetOfProcessor (const long &numParticles, - unsigned long long &offset, - unsigned long long &sum) const; - - int m_MPIRank = 0; - int m_MPISize = 1; - - unsigned long long m_Total = 0; - - std::vector m_ParticleCounterByLevel; - }; -} // namespace detail - - /** This element writes the particle beam out to openPMD data. +#ifdef ImpactX_USE_OPENPMD + /** Unclutter a real_names to openPMD record + * + * @TODO move to ABLASTR * - * This class behaves like a singleton if constructed with the - * same series name as an existing instance. + * @param fullName name as in real_names variable + * @return pair of openPMD record and component name */ - struct BeamMonitor - : public mixin::LinearTransport, - public mixin::Thin - { - static constexpr auto type = "BeamMonitor"; - using PType = typename ImpactXParticleContainer::ParticleType; - using PinnedContainer = typename ImpactXParticleContainer::ContainerLike; - - /** This element writes the particle beam out to openPMD data. - * - * Elements with the same series name are identical. - * - * @param series_name name of the data series, usually the element name - * @param backend file format backend for openPMD, e.g., "bp" or "h5" - * @param encoding openPMD iteration encoding: "v"ariable based, "f"ile based, "g"roup based (default) - * @param period_sample_intervals for periodic lattice, only output every Nth period (turn) - */ - BeamMonitor (std::string series_name, std::string backend="default", std::string encoding="g", int period_sample_intervals=1); - - BeamMonitor (BeamMonitor const & other) = default; - BeamMonitor (BeamMonitor && other) = default; - BeamMonitor& operator= (BeamMonitor const & other) = default; - BeamMonitor& operator= (BeamMonitor && other) = default; - - /** Prepare entering the element before starting push logic. - * - * And write reference particle. - * - * @param[in] pc particle container - * @param[in] real_soa_names ParticleReal component names - * @param[in] int_soa_names integer component names - * @param[in] ref_part reference particle - * @param[in] step global step for diagnostics - */ - void prepare ( - PinnedContainer & pc, - std::vector const & real_soa_names, - std::vector const & int_soa_names, - RefPart const & ref_part, - int step - ); - - /** Dump all particles. - * - * Particles are relative to the reference particle. - * - * @param[in,out] pc particle container to push - * @param[in] step global step for diagnostics - * @param[in] period for periodic lattices, this is the current period (turn or cycle) - */ - void operator() ( - ImpactXParticleContainer & pc, - int step, - int period - ); - - /** Write a tile of particles - * - * @param pti particle tile iterator - * @param[in] real_soa_names ParticleReal component names - * @param[in] int_soa_names integer component names - * @param ref_part reference particle - */ - void operator() ( - PinnedContainer::ParIterType & pti, - std::vector const & real_soa_names, - std::vector const & int_soa_names, - RefPart const & ref_part - ); + std::pair< std::string, std::string > + name2openPMD (const std::string& fullName); - /** This does nothing to the reference particle. */ - using Thin::operator(); - /** This pushes the covariance matrix. */ - using LinearTransport::operator(); - - /** This function returns the linear transport map. - * - * @param[in] refpart reference particle - * @returns 6x6 transport matrix - */ - AMREX_GPU_HOST AMREX_FORCE_INLINE - Map6x6 - transport_map ([[maybe_unused]] RefPart const & AMREX_RESTRICT refpart) const - { - // nothing to do - return Map6x6::Identity(); - } - - /** Get the name of the series - * - * Elements with the same series name are identical. - */ - std::string series_name () const { return m_series_name; } - - /** track all m_series_name instances - * - * Ensure m_series is the same for the same name. - */ - static inline std::map m_unique_series = {}; - - /** Close and deallocate all data series and backends. - */ - void - finalize (); - - private: - std::string m_series_name; //! openPMD filename - std::string m_OpenPMDFileType; //! openPMD backend: usually HDF5 (h5) or ADIOS2 (bp/bp4/bp5) or ADIOS2 SST (sst) - std::any m_series; //! openPMD::Series that holds potentially multiple outputs - int m_step = 0; //! global step for output - - int m_file_min_digits = 6; //! minimum number of digits to iteration number in file name - - int m_period_sample_intervals = 1; //! only output every Nth period (turn or cycle) of periodic lattices - - /** This rank's offset in the MPI-global particle array, by level - * - * This MUST be updated by prepare() before the next step's output. - */ - std::vector m_offset; - - /** Reduced Beam Characteristics - * - * In situ calculated particle bunch moments. - */ - std::unordered_map m_rbc; - - }; - - /** Calculate additional particle properties. - * - * @param[in] element_name name of the element (for input queries) - * @param[inout] pc particle container - */ - void - add_optional_properties ( - std::string const & element_name, - ImpactXParticleContainer & pc + //! TODO@ move to ablastr + openPMD::RecordComponent get_component_record ( + openPMD::ParticleSpecies & species, + std::string comp_name ); +#endif +} // namespace impactx::elements::diagnostics::detail -} // namespace impactx::diagnostics - -#endif // IMPACTX_ELEMENTS_DIAGS_OPENPMD_H +#endif // IMPACTX_OPENPMD_H diff --git a/src/elements/diagnostics/openPMD.cpp b/src/elements/diagnostics/openPMD.cpp index 75aab00a7..5cabe38f6 100644 --- a/src/elements/diagnostics/openPMD.cpp +++ b/src/elements/diagnostics/openPMD.cpp @@ -7,111 +7,17 @@ * Authors: Axel Huebl * License: BSD-3-Clause-LBNL */ +#include "elements/diagnostics/openPMD.H" -#include "openPMD.H" -#include "ImpactXVersion.H" -#include "particles/ImpactXParticleContainer.H" -#include "diagnostics/ReducedBeamCharacteristics.H" -#include -#include -#include -#include - -#ifdef ImpactX_USE_OPENPMD -# include -namespace io = openPMD; -#endif - -#include -#include -#include - - -namespace impactx::elements::diagnostics -{ -namespace detail +namespace impactx::elements::diagnostics::detail { - ImpactXParticleCounter::ImpactXParticleCounter (ParticleContainer & pc) - { - m_MPISize = amrex::ParallelDescriptor::NProcs(); - m_MPIRank = amrex::ParallelDescriptor::MyProc(); - - m_ParticleCounterByLevel.resize(pc.finestLevel()+1); - m_ParticleOffsetAtRank.resize(pc.finestLevel()+1); - m_ParticleSizeAtRank.resize(pc.finestLevel()+1); - - for (auto currentLevel = 0; currentLevel <= pc.finestLevel(); currentLevel++) - { - long numParticles = 0; // numParticles in this processor - - for (ParticleIter pti(pc, currentLevel); pti.isValid(); ++pti) { - auto numParticleOnTile = pti.numParticles(); - numParticles += numParticleOnTile; - } - - unsigned long long offset=0; // offset of this level - unsigned long long sum=0; // numParticles in this level (sum from all processors) - - GetParticleOffsetOfProcessor(numParticles, offset, sum); - - m_ParticleCounterByLevel[currentLevel] = sum; - m_ParticleOffsetAtRank[currentLevel] = offset; - m_ParticleSizeAtRank[currentLevel] = numParticles; - - // adjust offset, it should be numbered after particles from previous levels - for (auto lv=0; lv the particles in the comm -// sum of all particles in the comm -// - void - ImpactXParticleCounter::GetParticleOffsetOfProcessor ( - const long& numParticles, - unsigned long long& offset, - unsigned long long& sum - ) const - { - offset = 0; -#if defined(AMREX_USE_MPI) - std::vector result(m_MPISize, 0); - amrex::ParallelGather::Gather (numParticles, result.data(), -1, amrex::ParallelDescriptor::Communicator()); +#ifdef ImpactX_USE_OPENPMD + namespace io = openPMD; - sum = 0; - int const num_results = result.size(); - for (int i=0; i - name2openPMD ( const std::string& fullName ) + std::pair< std::string, std::string > + name2openPMD (const std::string& fullName) { std::string record_name = fullName; std::string component_name = io::RecordComponent::SCALAR; @@ -125,341 +31,15 @@ namespace detail return make_pair(record_name, component_name); } - // TODO: move to ablastr + io::RecordComponent get_component_record ( io::ParticleSpecies & species, std::string comp_name - ) { + ) + { // handle scalar and non-scalar records by name const auto [record_name, component_name] = name2openPMD(std::move(comp_name)); return species[record_name][component_name]; } #endif -} // namespace detail - - void BeamMonitor::finalize () - { -#ifdef ImpactX_USE_OPENPMD - // close shared series alias - if (m_series.has_value()) - { - auto series = std::any_cast(m_series); - series.close(); - m_series.reset(); - } - - // remove from unique series map - if (m_unique_series.count(m_series_name) != 0u) - m_unique_series.erase(m_series_name); -#endif // ImpactX_USE_OPENPMD - } - - BeamMonitor::BeamMonitor (std::string series_name, std::string backend, std::string encoding, int period_sample_intervals) : - m_series_name(std::move(series_name)), m_OpenPMDFileType(std::move(backend)), m_period_sample_intervals(period_sample_intervals) - { -#ifdef ImpactX_USE_OPENPMD - // pick first available backend if default is chosen - if( m_OpenPMDFileType == "default" ) -# if openPMD_HAVE_ADIOS2==1 - m_OpenPMDFileType = "bp"; -# elif openPMD_HAVE_ADIOS1==1 - m_OpenPMDFileType = "bp"; -# elif openPMD_HAVE_HDF5==1 - m_OpenPMDFileType = "h5"; -# else - m_OpenPMDFileType = "json"; -# endif - - // encoding of iterations in the series - openPMD::IterationEncoding series_encoding = openPMD::IterationEncoding::groupBased; - if ( "v" == encoding ) - series_encoding = openPMD::IterationEncoding::variableBased; - else if ( "g" == encoding ) - series_encoding = openPMD::IterationEncoding::groupBased; - else if ( "f" == encoding ) - series_encoding = openPMD::IterationEncoding::fileBased; - - amrex::ParmParse pp_diag("diag"); - // turn filter - pp_diag.queryAddWithParser("period_sample_intervals", m_period_sample_intervals); - // legacy options from other diagnostics - pp_diag.queryAddWithParser("file_min_digits", m_file_min_digits); - - // Ensure m_series is the same for the same names. - if (m_unique_series.count(m_series_name) == 0u) { - std::string filepath = "diags/openPMD/"; - filepath.append(m_series_name); - - if (series_encoding == openPMD::IterationEncoding::fileBased) - { - std::string const fileSuffix = std::string("_%0") + std::to_string(m_file_min_digits) + std::string("T"); - filepath.append(fileSuffix); - } - filepath.append(".").append(m_OpenPMDFileType); - - // transform paths for Windows -# ifdef _WIN32 - filepath = openPMD::auxiliary::replace_all(filepath, "/", "\\"); -# endif - - auto series = io::Series(filepath, io::Access::CREATE -# if openPMD_HAVE_MPI==1 - , amrex::ParallelDescriptor::Communicator() -# endif - , "adios2.engine.usesteps = true" - ); - series.setSoftware("ImpactX", IMPACTX_VERSION); - series.setIterationEncoding( series_encoding ); - m_series = series; - m_unique_series[m_series_name] = series; - } - else { - m_series = m_unique_series[m_series_name]; - } -#else - amrex::AllPrint() << "Warning: openPMD output requested but not compiled for series=" << m_series_name << "\n"; -#endif - } - - void BeamMonitor::prepare ( - PinnedContainer & pc, - std::vector const & real_soa_names, - std::vector const & int_soa_names, - RefPart const & ref_part, - int step - ) { -#ifdef ImpactX_USE_OPENPMD - m_step = step; - - // series & iteration - auto series = std::any_cast(m_series); - io::WriteIterations iterations = series.writeIterations(); - io::Iteration iteration = iterations[m_step]; - io::ParticleSpecies beam = iteration.particles["beam"]; - - // calculate & update particle offset in MPI-global particle array, per level - auto const num_levels = pc.finestLevel() + 1; - m_offset = std::vector(num_levels); - auto counter = detail::ImpactXParticleCounter(pc); - auto const np = counter.GetTotalNumParticles(); - for (auto currentLevel = 0; currentLevel < num_levels; currentLevel++) { - m_offset.at(currentLevel) = static_cast( counter.m_ParticleOffsetAtRank[currentLevel] ); - } - - // helpers to parse strings to openPMD - auto const scalar = openPMD::RecordComponent::SCALAR; - auto const getComponentRecord = [&beam](std::string comp_name) { - return detail::get_component_record(beam, std::move(comp_name)); - }; - - // define data set and metadata - io::Datatype const dtype_fl = io::determineDatatype(); - io::Datatype const dtype_ui = io::determineDatatype(); - auto d_fl = io::Dataset(dtype_fl, {np}); - auto d_ui = io::Dataset(dtype_ui, {np}); - - // 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 ); - beam.setAttribute( "z_ref", ref_part.z ); - beam.setAttribute( "t_ref", ref_part.t ); - beam.setAttribute( "px_ref", ref_part.px ); - 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", 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 - { - - beam["positionOffset"]["x"].resetDataset(d_fl); - beam["positionOffset"]["x"].makeConstant(ref_part.x); - beam["positionOffset"]["y"].resetDataset(d_fl); - beam["positionOffset"]["y"].makeConstant(ref_part.y); - beam["positionOffset"]["t"].resetDataset(d_fl); - beam["positionOffset"]["t"].makeConstant(ref_part.t); - } - - // unique, global particle index - beam["id"][scalar].resetDataset(d_ui); - - // SoA: Real - { - for (auto real_idx = 0; real_idx < pc.NumRealComps(); real_idx++) { - auto const component_name = real_soa_names.at(real_idx); - getComponentRecord(component_name).resetDataset(d_fl); - } - } - // SoA: Int - static_assert(IntSoA::nattribs == 0); // not yet used - if (!int_soa_names.empty()) - throw std::runtime_error("BeamMonitor: int_soa_names output not yet implemented!"); -#else - amrex::ignore_unused(pc, step); -#endif // ImpactX_USE_OPENPMD - } - - void - BeamMonitor::operator() ( - ImpactXParticleContainer & pc, - int step, - int period - ) - { - // filter out this turn? - if (period % m_period_sample_intervals != 0) - return; - -#ifdef ImpactX_USE_OPENPMD - std::string profile_name = "impactx::Push::" + std::string(BeamMonitor::type); - BL_PROFILE(profile_name); - - // preparing to access reference particle data: RefPart - RefPart & ref_part = pc.GetRefParticle(); - - // 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 = impactx::diagnostics::reduced_beam_characteristics(pc); - - // component names - std::vector real_soa_names = pc.GetRealSoANames(); - std::vector int_soa_names = pc.GetIntSoANames(); - - // pinned memory copy - PinnedContainer pinned_pc = pc.make_alike(); - pinned_pc.copyParticles(pc, true); // no filtering - - // TODO: filtering - /* - using SrcData = WarpXParticleContainer::ParticleTileType::ConstParticleTileDataType; - tmp.copyParticles(*pc, - [=] AMREX_GPU_HOST_DEVICE (const SrcData& src, int ip, const amrex::RandomEngine& engine) - { - const SuperParticleType& p = src.getSuperParticle(ip); - return random_filter(p, engine) * uniform_filter(p, engine) - * parser_filter(p, engine) * geometry_filter(p, engine); - }, true); - */ - - // prepare element access & write reference particle - this->prepare(pinned_pc, real_soa_names, int_soa_names, ref_part, step); - - // loop over refinement levels - int const nLevel = pinned_pc.finestLevel(); - for (int lev = 0; lev <= nLevel; ++lev) - { - // loop over all particle boxes - //using ParIt = ImpactXParticleContainer::iterator; - using ParIt = PinnedContainer::ParIterType; - // note: openPMD-api is not thread-safe, so do not run OMP parallel here - for (ParIt pti(pinned_pc, lev); pti.isValid(); ++pti) { - // write beam particles relative to reference particle - this->operator()(pti, real_soa_names, int_soa_names, ref_part); - } // end loop over all particle boxes - } // end mesh-refinement level loop - - auto series = std::any_cast(m_series); - io::WriteIterations iterations = series.writeIterations(); - io::Iteration iteration = iterations[m_step]; - - // close iteration - iteration.close(); -#else - amrex::ignore_unused(pc, step); -#endif // ImpactX_USE_OPENPMD - } - - void - BeamMonitor::operator() ( - PinnedContainer::ParIterType & pti, - std::vector const & real_soa_names, - std::vector const & int_soa_names, - RefPart const & ref_part - ) - { -#ifdef ImpactX_USE_OPENPMD - int const currentLevel = pti.GetLevel(); - - auto & offset = m_offset.at(currentLevel); // ... - - // series & iteration - auto series = std::any_cast(m_series); - io::WriteIterations iterations = series.writeIterations(); - io::Iteration iteration = iterations[m_step]; - - // writing - io::ParticleSpecies beam = iteration.particles["beam"]; - - auto const numParticleOnTile = pti.numParticles(); - uint64_t const numParticleOnTile64 = static_cast( numParticleOnTile ); - - // Do not call storeChunk() with zero-sized particle tiles: - // https://github.com/openPMD/openPMD-api/issues/1147 - //if (numParticleOnTile == 0) continue; - - auto const scalar = openPMD::RecordComponent::SCALAR; - auto const getComponentRecord = [&beam](std::string comp_name) { - return detail::get_component_record(beam, std::move(comp_name)); - }; - - // SoA - auto const& soa = pti.GetStructOfArrays(); - // particle id arrays - { - beam["id"][scalar].storeChunkRaw(soa.GetIdCPUData().data(), {offset}, {numParticleOnTile64}); - } - // SoA floating point (ParticleReal) properties - { - for (auto real_idx=0; real_idx < soa.NumRealComps(); real_idx++) { - auto const component_name = real_soa_names.at(real_idx); - getComponentRecord(component_name).storeChunkRaw( - soa.GetRealData(real_idx).data(), {offset}, {numParticleOnTile64}); - } - } - // SoA integer (int) properties (not yet used) - { - static_assert(IntSoA::nattribs == 0); // not yet used - if (!int_soa_names.empty()) - throw std::runtime_error("BeamMonitor: int_soa_names output not yet implemented!"); - /* - // comment this in once IntSoA::nattribs is > 0 - - std::copy(IntSoA::names_s.begin(), IntSoA::names_s.end(), int_soa_names.begin()); - - for (auto int_idx=0; int_idx < RealSoA::nattribs; int_idx++) { - auto const component_name = int_soa_names.at(int_idx); - getComponentRecord(component_name).storeChunkRaw( - soa.GetIntData(int_idx).data(), {offset}, {numParticleOnTile64}); - } - */ - } - - // TODO - amrex::ignore_unused(ref_part); - - // needs to be higher for next pti; must be reset for next step via prepare - offset += numParticleOnTile64; - - // TODO could be done once after all pti are processed - // TODO at that point, we could also close the iteration/step - series.flush(); -#else - amrex::ignore_unused(pti, ref_part); -#endif // ImpactX_USE_OPENPMD - } - -} // namespace impactx::diagnostics +} // namespace impactx::elements::diagnostics::detail diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 95341fb3a..63d34a2cb 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -154,13 +154,15 @@ namespace impactx sigx, sigy, sigt, sigpx, sigpy, sigpt, muxpx, muypy, mutpt); - } else { + } else if (base_dist_type == "empty") { + dist = distribution::Empty(); + } else + { throw std::runtime_error("Unknown distribution: " + distribution_type); } } - else if (distribution_type == "thermal") - { + else if (distribution_type == "thermal") { amrex::ParticleReal k, kT, kT_halo, normalize, normalize_halo; amrex::ParticleReal halo = 0.0; pp_dist.getWithParser("k", k); @@ -173,6 +175,10 @@ namespace impactx pp_dist.queryWithParser("halo", halo); dist = distribution::Thermal(k, kT, kT_halo, normalize, normalize_halo, halo); + } + else if (distribution_type == "empty") + { + dist = distribution::Empty(); } else { throw std::runtime_error("Unknown distribution: " + distribution_type); } @@ -266,9 +272,9 @@ namespace impactx // redistribute particles so that they reside on the correct MPI rank. int const myproc = amrex::ParallelDescriptor::MyProc(); int const nprocs = amrex::ParallelDescriptor::NProcs(); - int const navg = npart / nprocs; + int const navg = npart / nprocs; // note: integer division int const nleft = npart - navg * nprocs; - int npart_this_proc = (myproc < nleft) ? navg+1 : navg; + int const npart_this_proc = (myproc < nleft) ? navg+1 : navg; // add 1 to each proc until distributed auto const rel_part_this_proc = amrex::ParticleReal(npart_this_proc) / amrex::ParticleReal(npart); @@ -412,13 +418,12 @@ namespace impactx using namespace amrex::literals; // Parse the beam distribution parameters - amrex::ParmParse const pp_dist("beam"); + amrex::ParmParse pp_dist("beam"); amrex::ParmParse pp_algo("algo"); std::string track = "particles"; pp_algo.queryAdd("track", track); - if (track == "particles") - { + if (track == "particles") { // set charge and mass and energy of ref particle RefPart const ref = initialization::read_reference_particle(pp_dist); amr_data->track_particles.m_particle_container->SetRefParticle(ref); @@ -426,14 +431,19 @@ namespace impactx amrex::ParticleReal bunch_charge = 0.0; // Bunch charge (C) pp_dist.getWithParser("charge", bunch_charge); - int npart = 1; // Number of simulation particles - pp_dist.getWithParser("npart", npart); - std::string unit_type; // System of units pp_dist.get("units", unit_type); distribution::KnownDistributions dist = initialization::read_distribution(pp_dist); - add_particles(bunch_charge, dist, npart); + std::string distribution; + pp_dist.get("distribution", distribution); + + int npart = 0; // Number of simulation particles + if (distribution != "empty") + { + pp_dist.getWithParser("npart", npart); + add_particles(bunch_charge, dist, npart); + } // print information on the initialized beam amrex::Print() << "Beam kinetic energy (MeV): " << ref.kin_energy_MeV() << std::endl; diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index 2e4e7e6e7..0e5dceea0 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -503,6 +503,17 @@ namespace detail } m_lattice.emplace_back(diagnostics::BeamMonitor(openpmd_name, openpmd_backend, openpmd_encoding, period_sample_intervals)); + } else if (element_type == "source") + { + std::string distribution, openpmd_path; + pp_element.get("distribution", distribution); + + if (distribution == "openPMD") + { + pp_element.get("openpmd_path", openpmd_path); + } + + m_lattice.emplace_back( Source(distribution, openpmd_path, element_name) ); } else if (element_type == "line") { // Parse the lattice elements for the sub-lattice in the line diff --git a/src/initialization/Validate.cpp b/src/initialization/Validate.cpp index 5b33070c9..020f56752 100644 --- a/src/initialization/Validate.cpp +++ b/src/initialization/Validate.cpp @@ -14,6 +14,7 @@ #include #include +#include namespace impactx @@ -37,9 +38,19 @@ namespace impactx nParticles += amr_data->track_particles.m_particle_container->NumberOfParticlesAtLevel(lev); } if (nParticles == 0) - throw std::runtime_error("No particles found. Cannot run evolve without a beam."); - if (nParticles == 1) + { + // do we have a source element as the first element of the beamline? + auto & first_element = m_lattice.front(); + std::visit([](auto&& element){ + if (std::string_view(element.type) != std::string_view("Source")) { + throw std::runtime_error("No particles found. Cannot run evolve without a beam."); + } + }, first_element); + } + else if (nParticles == 1) + { throw std::runtime_error("Only one particle found. This is not yet supported: https://github.com/ECP-WarpX/impactx/issues/44"); + } } // elements diff --git a/src/python/elements.cpp b/src/python/elements.cpp index b734b7802..253ae70a6 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -1455,6 +1455,41 @@ void init_elements(py::module& m) register_beamoptics_push(py_SoftSolenoid); register_envelope_push(py_SoftSolenoid); + py::class_ py_Source(me, "Source"); + py_Source + .def("__repr__", + [](Source const & src) { + return element_name( + src, + std::make_pair("distribution", src.m_distribution), + std::make_pair("path", src.m_series_name) + ); + } + ) + .def(py::init< + std::string, + std::string, + std::optional + >(), + py::arg("distribution"), + py::arg("openpmd_path"), + py::arg("name") = py::none(), + "A particle source." + ) + .def_property("distribution", + [](Source & src) { return src.m_distribution; }, + [](Source & src, std::string distribution) { src.m_distribution = distribution; }, + "Distribution type of particles in the source" + ) + .def_property("series_name", + [](Source & src) { return src.m_series_name; }, + [](Source & src, std::string series_name) { src.m_series_name = series_name; }, + "Path to openPMD series as accepted by openPMD_api.Series" + ) + ; + register_beamoptics_push(py_Source); + register_envelope_push(py_Source); + py::class_ py_Sol(me, "Sol"); py_Sol .def("__repr__",