From 15afcde488deddca24e396ce8fa4c788d0217d22 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Feb 2025 15:40:10 -0800 Subject: [PATCH 01/10] [pre-commit.ci] pre-commit autoupdate (#839) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit updates: - [github.com/astral-sh/ruff-pre-commit: v0.9.4 → v0.9.6](https://github.com/astral-sh/ruff-pre-commit/compare/v0.9.4...v0.9.6) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5b21c78f1..624de45d0 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -66,7 +66,7 @@ repos: # Python: Ruff linter & formatter # https://docs.astral.sh/ruff/ - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.9.4 + rev: v0.9.6 hooks: # Run the linter - id: ruff From f8f7af01296441c7e8767fa92ae955bdddc1d133 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 10 Feb 2025 15:40:28 -0800 Subject: [PATCH 02/10] Doc: How-To Update (#838) Move "workflows" to "how-to guides" as in WarpX. Clean up some wording. --- docs/source/index.rst | 2 +- docs/source/usage/howto.rst | 11 ++++++++ .../{workflows => howto}/add_element.rst | 26 ++++++++++--------- docs/source/usage/workflows.rst | 15 ----------- examples/fodo_userdef/README.rst | 2 +- 5 files changed, 27 insertions(+), 29 deletions(-) create mode 100644 docs/source/usage/howto.rst rename docs/source/usage/{workflows => howto}/add_element.rst (83%) delete mode 100644 docs/source/usage/workflows.rst diff --git a/docs/source/index.rst b/docs/source/index.rst index 967641494..1491d4529 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -80,7 +80,7 @@ Usage usage/python usage/parameters usage/dashboard - usage/workflows + usage/howto Data Analysis ------------- diff --git a/docs/source/usage/howto.rst b/docs/source/usage/howto.rst new file mode 100644 index 000000000..9148b4173 --- /dev/null +++ b/docs/source/usage/howto.rst @@ -0,0 +1,11 @@ +.. _usage-howto: + +How-To Guides +============= + +This section collects typical user workflows and best practices for ImpactX. + +.. toctree:: + :maxdepth: 1 + + howto/add_element diff --git a/docs/source/usage/workflows/add_element.rst b/docs/source/usage/howto/add_element.rst similarity index 83% rename from docs/source/usage/workflows/add_element.rst rename to docs/source/usage/howto/add_element.rst index e557f1647..3a6d92202 100644 --- a/docs/source/usage/workflows/add_element.rst +++ b/docs/source/usage/howto/add_element.rst @@ -1,4 +1,4 @@ -.. _usage-workflows-add-element: +.. _usage-howto-add-element: Add New Beamline Elements ========================= @@ -10,27 +10,27 @@ The workflows described here apply both for thin kicks or thick elements. Thick elements can also use soft-edged fringe fields (see `existing soft-edged elements for implementation details `__). -.. _usage-workflows-add-element-linmap: +.. _usage-howto-add-element-linmap: Linear Map ---------- A custom linear element can be provided by specifying the 6x6 linear transport matrix :math:`R` as an input. -See the :ref:`example ` for Python and inputs file syntax to specify a custom linear element. +See the :ref:`FODO cell example ` for a demonstration of the custom linear element. The matrix elements :math:`R(i,j)` are indexed beginning with 1, so that :math:`i,j=1,2,3,4,5,6`. -The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m -and momenta :math:`(px,py,pt)` are dimensionless. So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m. +The matrix :math:`R` multiplies the phase space vector :math:`(x,p_x,y,p_y,t,p_t)`, where coordinates :math:`(x,y,t)` have units of m and momenta :math:`(p_x,p_y,p_t)` are dimensionless. +So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m. .. note:: If a user-provided linear map is used, it is up to the user to ensure that the 6x6 transport matrix is symplectic. - If a more general form of user-defined transport is needed, the :ref:`Python Programmable Element ` and the :ref:`C++ Element ` provide a more general approach. + If a more general form of user-defined transport is needed, the :ref:`Python Programmable Element ` and the :ref:`C++ Element ` provide a more general approach. -.. _usage-workflows-add-element-python: +.. _usage-howto-add-element-python: Python Programmable Element --------------------------- @@ -52,7 +52,7 @@ Detailed examples that show usage of the programmable element are: Detailed particle computing interfaces are presented in the `pyAMReX examples `__. -.. _usage-workflows-add-element-cxx: +.. _usage-howto-add-element-cxx: C++ Element ----------- @@ -72,10 +72,12 @@ To simplify the logic, we use so-called `mixin classes `__ over time. diff --git a/examples/fodo_userdef/README.rst b/examples/fodo_userdef/README.rst index 657c94071..b6e3d5d71 100644 --- a/examples/fodo_userdef/README.rst +++ b/examples/fodo_userdef/README.rst @@ -10,7 +10,7 @@ However, in the example here we define *additional user-defined, custom linear e Note that generally, if a user-provided linear map is used, the beam transport may not be symplectic. For an even more general, user-defined element, see :ref:`the FODO Cell example that uses a Programmable Element `. - For more details, see :ref:`this section `. + For more details, see :ref:`this section `. The matched Twiss parameters at entry are: From 001c192f0214797f9db2a5836576989d3e3ccd64 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 10 Feb 2025 17:03:43 -0800 Subject: [PATCH 03/10] Fix `beam.units = static` (#840) Currently, we only support (and document correctly) "static" units, but the user was able to also pass "dynamic", upon which we did print "Dynamic units" but then still used static ones. This now errors out correctly for anything else than "static" until we implement another variant. --- src/initialization/InitDistribution.cpp | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 8935daf3d..95341fb3a 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -452,10 +452,8 @@ namespace impactx if (unit_type == "static") { amrex::Print() << "Static units" << std::endl; - } else if (unit_type == "dynamic") { - amrex::Print() << "Dynamic units" << std::endl; } else { - throw std::runtime_error("Unknown units (static/dynamic): " + unit_type); + throw std::runtime_error("Unknown units (use 'static'): " + unit_type); } amrex::Print() << "Initialized beam distribution parameters" << std::endl; From 8f0350362ebe72ae572b0893e20ff657eec926b6 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 11 Feb 2025 10:14:13 -0800 Subject: [PATCH 04/10] 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__", From cdf408b29a1542433e60127d84114dbcf8739b26 Mon Sep 17 00:00:00 2001 From: Parthib Roy <159463257+proy30@users.noreply.github.com> Date: Tue, 11 Feb 2025 10:19:47 -0800 Subject: [PATCH 05/10] Dashboard: Add tooltips (#843) * update dashboard __init__.py import * Add tooltips functionality to non-nested states * Add tooltips for the custom vselects * update type hints --- .../impactx/dashboard/Input/components.py | 49 ++++++++++++------- .../impactx/dashboard/Input/defaults.py | 20 ++++++++ .../dashboard/Input/defaults_helper.py | 37 ++++++++++++++ src/python/impactx/dashboard/__init__.py | 2 + 4 files changed, 89 insertions(+), 19 deletions(-) create mode 100644 src/python/impactx/dashboard/Input/defaults_helper.py diff --git a/src/python/impactx/dashboard/Input/components.py b/src/python/impactx/dashboard/Input/components.py index f6eb75b5f..13884270e 100644 --- a/src/python/impactx/dashboard/Input/components.py +++ b/src/python/impactx/dashboard/Input/components.py @@ -1,6 +1,7 @@ from typing import Optional -from .. import setup_server, vuetify +from .. import html, setup_server, vuetify +from .defaults import TooltipDefaults from .generalFunctions import generalFunctions server, state, ctrl = setup_server() @@ -91,13 +92,18 @@ def select( generalFunctions.get_default(f"{v_model_name}_list", "default_values"), ) - return vuetify.VSelect( - label=label, - v_model=(v_model_name,), - items=items, - dense=True, - **kwargs, - ) + with vuetify.VTooltip(**TooltipDefaults.TOOLTIP_STYLE): + with vuetify.Template(v_slot_activator="{ on, attrs }"): + vuetify.VSelect( + label=label, + v_model=(v_model_name,), + items=items, + dense=True, + **kwargs, + v_on="on", + v_bind="attrs", + ) + html.Span(TooltipDefaults.TOOLTIP.get(v_model_name)) @staticmethod def text_field( @@ -121,17 +127,22 @@ def text_field( if v_model_name is None: v_model_name = label.lower().replace(" ", "_") - return vuetify.VTextField( - label=label, - v_model=(v_model_name,), - error_messages=(f"{v_model_name}_error_message", []), - type="number", - step=generalFunctions.get_default(f"{v_model_name}", "steps"), - suffix=generalFunctions.get_default(f"{v_model_name}", "units"), - __properties=["step"], - dense=True, - **kwargs, - ) + with vuetify.VTooltip(**TooltipDefaults.TOOLTIP_STYLE): + with vuetify.Template(v_slot_activator="{ on, attrs }"): + vuetify.VTextField( + label=label, + v_model=(v_model_name,), + error_messages=(f"{v_model_name}_error_message", []), + type="number", + step=generalFunctions.get_default(f"{v_model_name}", "steps"), + suffix=generalFunctions.get_default(f"{v_model_name}", "units"), + __properties=["step"], + dense=True, + v_on="on", + v_bind="attrs", + **kwargs, + ) + html.Span(TooltipDefaults.TOOLTIP.get(v_model_name)) class NavigationComponents: diff --git a/src/python/impactx/dashboard/Input/defaults.py b/src/python/impactx/dashboard/Input/defaults.py index f452470a3..398f63f60 100644 --- a/src/python/impactx/dashboard/Input/defaults.py +++ b/src/python/impactx/dashboard/Input/defaults.py @@ -1,3 +1,8 @@ +from impactx.impactx_pybind import ImpactX, RefPart + +from .defaults_helper import InputDefaultsHelper + + class DashboardDefaults: """ Defaults for input parameters in the ImpactX dashboard. @@ -113,3 +118,18 @@ class DashboardDefaults: "beta": "m", "emitt": "m", } + + +class TooltipDefaults: + """ + Defaults for input toolips in the ImpactX dashboard. + """ + + TOOLTIP_STYLE = { + "bottom": True, + "nudge_top": "20", + } + + TOOLTIP = InputDefaultsHelper.get_docstrings( + [RefPart, ImpactX], DashboardDefaults.DEFAULT_VALUES + ) diff --git a/src/python/impactx/dashboard/Input/defaults_helper.py b/src/python/impactx/dashboard/Input/defaults_helper.py new file mode 100644 index 000000000..69e09e605 --- /dev/null +++ b/src/python/impactx/dashboard/Input/defaults_helper.py @@ -0,0 +1,37 @@ +import inspect +from typing import Dict, List, Type + + +class InputDefaultsHelper: + """ + Methods in this class are used to dynamically parse + core ImpactX data (default values, docstrings, etc.) + """ + + @staticmethod + def get_docstrings( + class_names: List[Type], default_list: Dict[str, any] + ) -> Dict[str, str]: + """ + Retrieves docstrings for each method and property + in the provided clases. + + :param classes: The class names to parse docstrings with. + :param defaults_list: The dictionary of defaults value. + """ + + docstrings = {} + + for each_class in class_names: + for name, attribute in inspect.getmembers(each_class): + if name not in default_list: + continue + + is_method = inspect.isfunction(attribute) + is_property = inspect.isdatadescriptor(attribute) + + if is_method or is_property: + docstring = inspect.getdoc(attribute) or "" + docstrings[name] = docstring + + return docstrings diff --git a/src/python/impactx/dashboard/__init__.py b/src/python/impactx/dashboard/__init__.py index 875ecc666..0b73e2e2b 100644 --- a/src/python/impactx/dashboard/__init__.py +++ b/src/python/impactx/dashboard/__init__.py @@ -1,3 +1,4 @@ +from trame.widgets import html from trame.widgets import vuetify as vuetify # isort: off @@ -18,6 +19,7 @@ __all__ = [ + "html", "JupyterApp", "setup_server", "vuetify", From 6f46573e973030a774f140f73cb6ffe96464e453 Mon Sep 17 00:00:00 2001 From: ax3l <1353258+ax3l@users.noreply.github.com> Date: Tue, 11 Feb 2025 18:33:38 +0000 Subject: [PATCH 06/10] Update Stub Files --- .../impactx/impactx_pybind/__init__.pyi | 1 + .../impactx_pybind/elements/__init__.pyi | 48 +++++++++++++++++++ 2 files changed, 49 insertions(+) diff --git a/src/python/impactx/impactx_pybind/__init__.pyi b/src/python/impactx/impactx_pybind/__init__.pyi index 9c552fdb8..fbfe9b7f3 100644 --- a/src/python/impactx/impactx_pybind/__init__.pyi +++ b/src/python/impactx/impactx_pybind/__init__.pyi @@ -724,6 +724,7 @@ def push( | elements.SoftSolenoid | elements.SoftQuadrupole | elements.Sol + | elements.Source | elements.TaperedPL | elements.ThinDipole, step: int = 0, diff --git a/src/python/impactx/impactx_pybind/elements/__init__.pyi b/src/python/impactx/impactx_pybind/elements/__init__.pyi index 4a3b14159..444aefb82 100644 --- a/src/python/impactx/impactx_pybind/elements/__init__.pyi +++ b/src/python/impactx/impactx_pybind/elements/__init__.pyi @@ -42,6 +42,7 @@ __all__ = [ "SoftQuadrupole", "SoftSolenoid", "Sol", + "Source", "TaperedPL", "ThinDipole", "mixin", @@ -882,6 +883,7 @@ class KnownElementsList: | SoftSolenoid | SoftQuadrupole | Sol + | Source | TaperedPL | ThinDipole, ) -> None: ... @@ -919,6 +921,7 @@ class KnownElementsList: | SoftSolenoid | SoftQuadrupole | Sol + | Source | TaperedPL | ThinDipole ]: ... @@ -957,6 +960,7 @@ class KnownElementsList: | SoftSolenoid | SoftQuadrupole | Sol + | Source | TaperedPL | ThinDipole, ) -> None: @@ -1710,6 +1714,50 @@ class Sol(mixin.Named, mixin.Thick, mixin.Alignment, mixin.PipeAperture): @ks.setter def ks(self, arg1: float) -> None: ... +class Source(mixin.Named, mixin.Thin): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... + def __init__( + self, distribution: str, openpmd_path: str, name: str | None = None + ) -> None: + """ + A particle source. + """ + def __repr__(self) -> str: ... + @typing.overload + def push( + self, + pc: impactx.impactx_pybind.ImpactXParticleContainer, + step: int = 0, + period: int = 0, + ) -> None: + """ + Push first the reference particle, then all other particles. + """ + @typing.overload + def push( + self, + cm: amrex.space3d.amrex_3d_pybind.SmallMatrix_6x6_F_SI1_double, + ref: impactx.impactx_pybind.RefPart, + ) -> None: + """ + Linear push of the covariance matrix through an element. Expects that the reference particle was advanced first. + """ + @property + def distribution(self) -> str: + """ + Distribution type of particles in the source + """ + @distribution.setter + def distribution(self, arg1: str) -> None: ... + @property + def series_name(self) -> str: + """ + Path to openPMD series as accepted by openPMD_api.Series + """ + @series_name.setter + def series_name(self, arg1: str) -> None: ... + class TaperedPL(mixin.Named, mixin.Thin, mixin.Alignment): @staticmethod def _pybind11_conduit_v1_(*args, **kwargs): ... From 5f2e98dba4c0fd2f8fa48780046b546053afef48 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 11 Feb 2025 14:31:51 -0800 Subject: [PATCH 07/10] Release: 25.02 (#845) --- CMakeLists.txt | 2 +- cmake/dependencies/ABLASTR.cmake | 4 ++-- cmake/dependencies/pyAMReX.cmake | 4 ++-- docs/source/conf.py | 4 ++-- setup.py | 2 +- 5 files changed, 8 insertions(+), 8 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 308e7bc1c..4e9814c8a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,7 +1,7 @@ # Preamble #################################################################### # cmake_minimum_required(VERSION 3.24.0) -project(ImpactX VERSION 25.01) +project(ImpactX VERSION 25.02) include(${ImpactX_SOURCE_DIR}/cmake/ImpactXFunctions.cmake) diff --git a/cmake/dependencies/ABLASTR.cmake b/cmake/dependencies/ABLASTR.cmake index 87ffab855..93d3110bf 100644 --- a/cmake/dependencies/ABLASTR.cmake +++ b/cmake/dependencies/ABLASTR.cmake @@ -128,7 +128,7 @@ macro(find_ablastr) set(COMPONENT_DIM 3D) set(COMPONENT_PRECISION ${ImpactX_PRECISION} P${ImpactX_PRECISION}) - find_package(ABLASTR 25.01 CONFIG REQUIRED COMPONENTS ${COMPONENT_DIM}) + find_package(ABLASTR 25.02 CONFIG REQUIRED COMPONENTS ${COMPONENT_DIM}) message(STATUS "ABLASTR: Found version '${ABLASTR_VERSION}'") endif() @@ -162,7 +162,7 @@ set(ImpactX_openpmd_src "" set(ImpactX_ablastr_repo "https://github.com/ECP-WarpX/WarpX.git" CACHE STRING "Repository URI to pull and build ABLASTR from if(ImpactX_ablastr_internal)") -set(ImpactX_ablastr_branch "958a39463c9e3fea0bbe1da0306104ccf9a2164c" +set(ImpactX_ablastr_branch "25.02" CACHE STRING "Repository branch for ImpactX_ablastr_repo if(ImpactX_ablastr_internal)") diff --git a/cmake/dependencies/pyAMReX.cmake b/cmake/dependencies/pyAMReX.cmake index 8ab600e0b..96a3d5a1c 100644 --- a/cmake/dependencies/pyAMReX.cmake +++ b/cmake/dependencies/pyAMReX.cmake @@ -59,7 +59,7 @@ function(find_pyamrex) endif() elseif(NOT ImpactX_pyamrex_internal) # TODO: MPI control - find_package(pyAMReX 25.01 CONFIG REQUIRED) + find_package(pyAMReX 25.02 CONFIG REQUIRED) message(STATUS "pyAMReX: Found version '${pyAMReX_VERSION}'") endif() endfunction() @@ -74,7 +74,7 @@ option(ImpactX_pyamrex_internal "Download & build pyAMReX" ON) set(ImpactX_pyamrex_repo "https://github.com/AMReX-Codes/pyamrex.git" CACHE STRING "Repository URI to pull and build pyamrex from if(ImpactX_pyamrex_internal)") -set(ImpactX_pyamrex_branch "458c9ae7ab3cd4ca4e4e9736e82c60f9a7e7606c" +set(ImpactX_pyamrex_branch "25.02" CACHE STRING "Repository branch for ImpactX_pyamrex_repo if(ImpactX_pyamrex_internal)") diff --git a/docs/source/conf.py b/docs/source/conf.py index d6ac1b8b6..fdd748cd5 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -81,9 +81,9 @@ # built documents. # # The short X.Y version. -version = "25.01" +version = "25.02" # The full version, including alpha/beta/rc tags. -release = "25.01" +release = "25.02" # The language for content autogenerated by Sphinx. Refer to documentation # for a list of supported languages. diff --git a/setup.py b/setup.py index fa49892d6..91643702e 100644 --- a/setup.py +++ b/setup.py @@ -230,7 +230,7 @@ def build_extension(self, ext): setup( name="impactx", # note PEP-440 syntax: x.y.zaN but x.y.z.devN - version="25.01", + version="25.02", packages=["impactx"], # Python sources: package_dir={"": "src/python"}, From 7b22e71b236ec32d269ef2716c40a65eb3adaf2e Mon Sep 17 00:00:00 2001 From: ax3l <1353258+ax3l@users.noreply.github.com> Date: Tue, 11 Feb 2025 22:45:13 +0000 Subject: [PATCH 08/10] Update Stub Files --- src/python/impactx/__init__.pyi | 2 +- src/python/impactx/impactx_pybind/__init__.pyi | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/python/impactx/__init__.pyi b/src/python/impactx/__init__.pyi index 9a9bee145..16d67fb57 100644 --- a/src/python/impactx/__init__.pyi +++ b/src/python/impactx/__init__.pyi @@ -81,7 +81,7 @@ __author__: str = ( "Axel Huebl, Chad Mitchell, Ryan Sandberg, Marco Garten, Ji Qiang, et al." ) __license__: str = "BSD-3-Clause-LBNL" -__version__: str = "25.01" +__version__: str = "25.02" s: impactx_pybind.CoordSystem # value = t: impactx_pybind.CoordSystem # value = cxx = impactx_pybind diff --git a/src/python/impactx/impactx_pybind/__init__.pyi b/src/python/impactx/impactx_pybind/__init__.pyi index fbfe9b7f3..6181c1b1e 100644 --- a/src/python/impactx/impactx_pybind/__init__.pyi +++ b/src/python/impactx/impactx_pybind/__init__.pyi @@ -738,6 +738,6 @@ __author__: str = ( "Axel Huebl, Chad Mitchell, Ryan Sandberg, Marco Garten, Ji Qiang, et al." ) __license__: str = "BSD-3-Clause-LBNL" -__version__: str = "25.01" +__version__: str = "25.02" s: CoordSystem # value = t: CoordSystem # value = From 220531e599088f4bdfe41d4152fc9af6f6fd931b Mon Sep 17 00:00:00 2001 From: Parthib Roy <159463257+proy30@users.noreply.github.com> Date: Tue, 11 Feb 2025 23:29:34 -0800 Subject: [PATCH 09/10] Dashboard: Add documentation sidebar (#836) * Add sidebar for documentation Instead of opening a new tab to the documentation when users click on the info button on a section header, the documentation is integrated onto the dashboard and is shown through a drawer style ui * update documentation still may be wrong link * remove unused styling * fix importing issues * move function to GeneralFunctions Future pr needs to split up generalFunctions * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../impactx/dashboard/Input/components.py | 23 ++++++++++++- .../impactx/dashboard/Input/defaults.py | 8 +++++ .../dashboard/Input/generalFunctions.py | 32 +++++++------------ src/python/impactx/dashboard/__init__.py | 1 + src/python/impactx/dashboard/__main__.py | 1 + 5 files changed, 43 insertions(+), 22 deletions(-) diff --git a/src/python/impactx/dashboard/Input/components.py b/src/python/impactx/dashboard/Input/components.py index 13884270e..3d9d6b498 100644 --- a/src/python/impactx/dashboard/Input/components.py +++ b/src/python/impactx/dashboard/Input/components.py @@ -6,6 +6,9 @@ server, state, ctrl = setup_server() +state.documentation_drawer_open = False +state.documentation_url = "" + class CardComponents: """ @@ -41,7 +44,7 @@ def documentation_icon(section_name: str) -> vuetify.VIcon: return vuetify.VIcon( "mdi-information", style="color: #00313C;", - click=lambda: generalFunctions.documentation(section_name), + click=lambda: generalFunctions.open_documentation(section_name), ) @staticmethod @@ -187,3 +190,21 @@ def create_dialog_tabs(name: str, num_tabs: int, tab_names: list[str]) -> None: for tab_name in tab_names: vuetify.VTab(tab_name) vuetify.VDivider() + + @staticmethod + def create_documentation_drawer(): + with vuetify.VNavigationDrawer( + v_model=("documentation_drawer_open",), + absolute=True, + right=True, + hide_overlay=True, + style="width: 30vw; top: 64px !important; position: fixed;", + ): + with vuetify.VContainer( + fluid=True, + classes="pa-0 fill-height", + ): + html.Iframe( + src=("documentation_url",), + style="width: 100%; height: 100%; border: none;", + ) diff --git a/src/python/impactx/dashboard/Input/defaults.py b/src/python/impactx/dashboard/Input/defaults.py index 398f63f60..1266efd9f 100644 --- a/src/python/impactx/dashboard/Input/defaults.py +++ b/src/python/impactx/dashboard/Input/defaults.py @@ -119,6 +119,14 @@ class DashboardDefaults: "emitt": "m", } + DOCUMENTATION = { + "input_parameters": "https://impactx.readthedocs.io/en/latest/usage/python.html#impactx.ImpactX", + "lattice_configuration": "https://impactx.readthedocs.io/en/latest/usage/python.html#lattice-elements", + "distribution_parameters": "https://impactx.readthedocs.io/en/latest/usage/python.html#initial-beam-distributions", + "space_charge": "https://impactx.readthedocs.io/en/latest/usage/parameters.html#space-charge", + "csr": "https://impactx.readthedocs.io/en/latest/usage/parameters.html#coherent-synchrotron-radiation-csr", + } + class TooltipDefaults: """ diff --git a/src/python/impactx/dashboard/Input/generalFunctions.py b/src/python/impactx/dashboard/Input/generalFunctions.py index b02ebd91d..a58507f33 100644 --- a/src/python/impactx/dashboard/Input/generalFunctions.py +++ b/src/python/impactx/dashboard/Input/generalFunctions.py @@ -7,10 +7,7 @@ """ import inspect -import os import re -import subprocess -import webbrowser from .. import setup_server from .defaults import DashboardDefaults @@ -24,27 +21,20 @@ class generalFunctions: @staticmethod - def documentation(section_name): + def open_documentation(section_name): """ - Opens a tab to the specified section link in the documentation. - :param section_name (str): The name of the documentation section to open. + Retrieves the documentation link with the provided section_name + and opens the documentation sidebar on the dashoard. + + :param section_name: The name for the input section. """ - url_dict = { - "input_parameters": "https://impactx.readthedocs.io/en/latest/usage/python.html#general", - "lattice_configuration": "https://impactx.readthedocs.io/en/latest/usage/python.html#lattice-elements", - "distribution_parameters": "https://impactx.readthedocs.io/en/latest/usage/python.html#initial-beam-distributions", - "space_charge": "https://impactx.readthedocs.io/en/latest/usage/parameters.html#space-charge", - "csr": "https://impactx.readthedocs.io/en/latest/usage/parameters.html#coherent-synchrotron-radiation-csr", - } - - url = url_dict.get(section_name) - if url is None: - raise ValueError(f"Invalid section name: {section_name}") - - if "WSL_DISTRO_NAME" in os.environ: - subprocess.run(["explorer.exe", url]) + + new_url = DashboardDefaults.DOCUMENTATION.get(section_name) + if state.documentation_drawer_open and state.documentation_url == new_url: + state.documentation_drawer_open = False else: - webbrowser.open_new_tab(url) + state.documentation_url = new_url + state.documentation_drawer_open = True @staticmethod def get_default(parameter, type): diff --git a/src/python/impactx/dashboard/__init__.py b/src/python/impactx/dashboard/__init__.py index 0b73e2e2b..fe6d85916 100644 --- a/src/python/impactx/dashboard/__init__.py +++ b/src/python/impactx/dashboard/__init__.py @@ -22,6 +22,7 @@ "html", "JupyterApp", "setup_server", + "html", "vuetify", "AnalyzeSimulation", "NavigationComponents", diff --git a/src/python/impactx/dashboard/__main__.py b/src/python/impactx/dashboard/__main__.py index b31b89358..b66624ed7 100644 --- a/src/python/impactx/dashboard/__main__.py +++ b/src/python/impactx/dashboard/__main__.py @@ -90,6 +90,7 @@ def application(): NavigationComponents.create_route("Analyze", "mdi-chart-box-multiple") with layout.content: + NavigationComponents.create_documentation_drawer() router.RouterView() init_terminal() return layout From 5816accb6f0dc58ec039730caaae8a450be33acf Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Wed, 12 Feb 2025 12:02:40 -0800 Subject: [PATCH 10/10] Introduce `HandleSpacecharge` (#842) * Introduce `HandleSpacecharge` Similar to our `HandleWakefield` logic for particle tracking, this simplifies the central particle tracking loop by cleanly hiding implementation details in a separate function. * Fix Doxygen String Co-authored-by: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> --------- Co-authored-by: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> --- src/particles/spacecharge/CMakeLists.txt | 1 + .../spacecharge/ForceFromSelfFields.H | 4 +- .../spacecharge/ForceFromSelfFields.cpp | 4 +- src/particles/spacecharge/GatherAndPush.H | 4 +- src/particles/spacecharge/GatherAndPush.cpp | 4 +- src/particles/spacecharge/HandleSpacecharge.H | 38 ++++++ .../spacecharge/HandleSpacecharge.cpp | 108 ++++++++++++++++++ src/particles/spacecharge/PoissonSolve.H | 4 +- src/particles/spacecharge/PoissonSolve.cpp | 4 +- src/tracking/particles.cpp | 72 +----------- 10 files changed, 162 insertions(+), 81 deletions(-) create mode 100644 src/particles/spacecharge/HandleSpacecharge.H create mode 100644 src/particles/spacecharge/HandleSpacecharge.cpp diff --git a/src/particles/spacecharge/CMakeLists.txt b/src/particles/spacecharge/CMakeLists.txt index 30e575b43..38ad2ea66 100644 --- a/src/particles/spacecharge/CMakeLists.txt +++ b/src/particles/spacecharge/CMakeLists.txt @@ -2,5 +2,6 @@ target_sources(lib PRIVATE ForceFromSelfFields.cpp GatherAndPush.cpp + HandleSpacecharge.cpp PoissonSolve.cpp ) diff --git a/src/particles/spacecharge/ForceFromSelfFields.H b/src/particles/spacecharge/ForceFromSelfFields.H index 48ff60670..aeca5e252 100644 --- a/src/particles/spacecharge/ForceFromSelfFields.H +++ b/src/particles/spacecharge/ForceFromSelfFields.H @@ -19,7 +19,7 @@ #include -namespace impactx::spacecharge +namespace impactx::particles::spacecharge { /** Calculate the space charge force field from the electric potential * @@ -36,6 +36,6 @@ namespace impactx::spacecharge const amrex::Vector& geom ); -} // namespace impactx +} // namespace impactx::particles::spacecharge #endif // IMPACTX_FORCEFROMSELFFIELDS_H diff --git a/src/particles/spacecharge/ForceFromSelfFields.cpp b/src/particles/spacecharge/ForceFromSelfFields.cpp index b46e74680..943d6a477 100644 --- a/src/particles/spacecharge/ForceFromSelfFields.cpp +++ b/src/particles/spacecharge/ForceFromSelfFields.cpp @@ -14,7 +14,7 @@ #include // for AMREX_D_DECL -namespace impactx::spacecharge +namespace impactx::particles::spacecharge { void ForceFromSelfFields ( std::unordered_map > & space_charge_field, @@ -60,4 +60,4 @@ namespace impactx::spacecharge } } } -} // namespace impactx::spacecharge +} // namespace impactx::particles::spacecharge diff --git a/src/particles/spacecharge/GatherAndPush.H b/src/particles/spacecharge/GatherAndPush.H index 2fa700d0b..e6c567f4a 100644 --- a/src/particles/spacecharge/GatherAndPush.H +++ b/src/particles/spacecharge/GatherAndPush.H @@ -20,7 +20,7 @@ #include -namespace impactx::spacecharge +namespace impactx::particles::spacecharge { /** Gather force fields and push particles in x,y,z * @@ -41,6 +41,6 @@ namespace impactx::spacecharge amrex::ParticleReal slice_ds ); -} // namespace impactx +} // namespace impactx::particles::spacecharge #endif // IMPACTX_GATHER_AND_PUSH_H diff --git a/src/particles/spacecharge/GatherAndPush.cpp b/src/particles/spacecharge/GatherAndPush.cpp index fe0b7454f..1566a5aff 100644 --- a/src/particles/spacecharge/GatherAndPush.cpp +++ b/src/particles/spacecharge/GatherAndPush.cpp @@ -16,7 +16,7 @@ #include // for AMREX_D_DECL -namespace impactx::spacecharge +namespace impactx::particles::spacecharge { void GatherAndPush ( ImpactXParticleContainer & pc, @@ -105,4 +105,4 @@ namespace impactx::spacecharge } // end loop over all particle boxes } // env mesh-refinement level loop } -} // namespace impactx::spacecharge +} // namespace impactx::particles::spacecharge diff --git a/src/particles/spacecharge/HandleSpacecharge.H b/src/particles/spacecharge/HandleSpacecharge.H new file mode 100644 index 000000000..a35c6b2d8 --- /dev/null +++ b/src/particles/spacecharge/HandleSpacecharge.H @@ -0,0 +1,38 @@ +/* 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, Chad Mitchell + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_HANDLE_SPACECHARGE_H +#define IMPACTX_HANDLE_SPACECHARGE_H + +#include "initialization/AmrCoreData.H" + +#include + +#include +#include + + +namespace impactx::particles::spacecharge +{ + /** Function to handle space charge including charge deposition, Poisson solve, + * field calculation, interpolation, and particle push. + * + * @param[in,out] amr_data AMR data like particle container and fields + * @param[in] ResizeMesh function to call to resize the mesh + * @param[in] slice_ds slice spacing along s + */ + void HandleSpacecharge ( + std::unique_ptr & amr_data, + std::function ResizeMesh, + amrex::Real slice_ds + ); + +} // namespace impactx::particles::spacecharge + +#endif // IMPACTX_HANDLE_SPACECHARGE_H diff --git a/src/particles/spacecharge/HandleSpacecharge.cpp b/src/particles/spacecharge/HandleSpacecharge.cpp new file mode 100644 index 000000000..4cc13a73a --- /dev/null +++ b/src/particles/spacecharge/HandleSpacecharge.cpp @@ -0,0 +1,108 @@ +/* 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, Chad Mitchell + * License: BSD-3-Clause-LBNL + */ +#include "HandleSpacecharge.H" + +#include "initialization/AmrCoreData.H" +#include "particles/ImpactXParticleContainer.H" +#include "particles/spacecharge/ForceFromSelfFields.H" +#include "particles/spacecharge/GatherAndPush.H" +#include "particles/spacecharge/PoissonSolve.H" +#include "particles/transformation/CoordinateTransformation.H" + +#include +#include + +#include + + +namespace impactx::particles::spacecharge +{ + void HandleSpacecharge ( + std::unique_ptr & amr_data, + std::function ResizeMesh, + amrex::Real slice_ds + ) + { + BL_PROFILE("impactx::particles::wakefields::HandleSpacecharge") + + amrex::ParmParse const pp_algo("algo"); + bool space_charge = false; + pp_algo.query("space_charge", space_charge); + + // turn off if disabled by user + if (!space_charge) { return; } + + // turn off if less than 2 particles + if (amr_data->track_particles.m_particle_container->TotalNumberOfParticles(true, false) < 2) { return; } + + // transform from x',y',t to x,y,z + transformation::CoordinateTransformation( + *amr_data->track_particles.m_particle_container, + CoordSystem::t + ); + + // Note: The following operation assume that + // the particles are in x, y, z coordinates. + + // Resize the mesh, based on `amr_data->track_particles.m_particle_container` extent + ResizeMesh(); + + // Redistribute particles in the new mesh in x, y, z + amr_data->track_particles.m_particle_container->Redistribute(); + + // charge deposition + amr_data->track_particles.m_particle_container->DepositCharge( + amr_data->track_particles.m_rho, + amr_data->refRatio() + ); + + // poisson solve in x,y,z + spacecharge::PoissonSolve( + *amr_data->track_particles.m_particle_container, + amr_data->track_particles.m_rho, + amr_data->track_particles.m_phi, + amr_data->refRatio() + ); + + // calculate force in x,y,z + spacecharge::ForceFromSelfFields( + amr_data->track_particles.m_space_charge_field, + amr_data->track_particles.m_phi, + amr_data->Geom() + ); + + // gather and space-charge push in x,y,z , assuming the space-charge + // field is the same before/after transformation + // TODO: This is currently using linear order. + spacecharge::GatherAndPush( + *amr_data->track_particles.m_particle_container, + amr_data->track_particles.m_space_charge_field, + amr_data->Geom(), + slice_ds + ); + + // transform from x,y,z to x',y',t + transformation::CoordinateTransformation( + *amr_data->track_particles.m_particle_container, + CoordSystem::s + ); + + // for later: original Impact implementation as an option + // Redistribute particles in x',y',t + // TODO: only needed if we want to gather and push space charge + // in x',y',t + // TODO: change geometry beforehand according to transformation + //m_amr_data->track_particles.m_particle_container->Redistribute(); + // + // in original Impact, we gather and space-charge push in x',y',t , + // assuming that the distribution did not change + } + +} // namespace impactx::particles::spacecharge diff --git a/src/particles/spacecharge/PoissonSolve.H b/src/particles/spacecharge/PoissonSolve.H index d6772f2a1..d6fb8e906 100644 --- a/src/particles/spacecharge/PoissonSolve.H +++ b/src/particles/spacecharge/PoissonSolve.H @@ -17,7 +17,7 @@ #include -namespace impactx::spacecharge +namespace impactx::particles::spacecharge { /** Calculate the electric potential from charge density * @@ -36,6 +36,6 @@ namespace impactx::spacecharge amrex::Vector rel_ref_ratio ); -} // namespace impactx +} // namespace impactx::particles::spacecharge #endif // IMPACTX_POISSONSOLVE_H diff --git a/src/particles/spacecharge/PoissonSolve.cpp b/src/particles/spacecharge/PoissonSolve.cpp index 82db71b90..23fb5f381 100644 --- a/src/particles/spacecharge/PoissonSolve.cpp +++ b/src/particles/spacecharge/PoissonSolve.cpp @@ -21,7 +21,7 @@ #include -namespace impactx::spacecharge +namespace impactx::particles::spacecharge { void PoissonSolve ( ImpactXParticleContainer const & pc, @@ -121,4 +121,4 @@ namespace impactx::spacecharge phi_at_level.FillBoundary(pc.GetParGDB()->Geom()[lev].periodicity()); } } -} // impactx::spacecharge +} // impactx::particles::spacecharge diff --git a/src/tracking/particles.cpp b/src/tracking/particles.cpp index 944fcdc7a..655944caf 100644 --- a/src/tracking/particles.cpp +++ b/src/tracking/particles.cpp @@ -13,10 +13,7 @@ #include "particles/ImpactXParticleContainer.H" #include "particles/Push.H" #include "diagnostics/DiagnosticOutput.H" -#include "particles/spacecharge/ForceFromSelfFields.H" -#include "particles/spacecharge/GatherAndPush.H" -#include "particles/spacecharge/PoissonSolve.H" -#include "particles/transformation/CoordinateTransformation.H" +#include "particles/spacecharge/HandleSpacecharge.H" #include "particles/wakefields/HandleWakefield.H" #include @@ -114,71 +111,8 @@ namespace impactx // Wakefield calculation: call wakefield function to apply wake effects particles::wakefields::HandleWakefield(*amr_data->track_particles.m_particle_container, element_variant, slice_ds); - // Space-charge calculation: turn off if there is only 1 particle - if (space_charge && - amr_data->track_particles.m_particle_container->TotalNumberOfParticles(true, false)) { - - // transform from x',y',t to x,y,z - transformation::CoordinateTransformation( - *amr_data->track_particles.m_particle_container, - CoordSystem::t); - - // Note: The following operation assume that - // the particles are in x, y, z coordinates. - - // Resize the mesh, based on `m_particle_container` extent - ResizeMesh(); - - // Redistribute particles in the new mesh in x, y, z - amr_data->track_particles.m_particle_container->Redistribute(); - - // charge deposition - amr_data->track_particles.m_particle_container->DepositCharge( - amr_data->track_particles.m_rho, - amr_data->refRatio() - ); - - // poisson solve in x,y,z - spacecharge::PoissonSolve( - *amr_data->track_particles.m_particle_container, - amr_data->track_particles.m_rho, - amr_data->track_particles.m_phi, - amr_data->refRatio() - ); - - // calculate force in x,y,z - spacecharge::ForceFromSelfFields( - amr_data->track_particles.m_space_charge_field, - amr_data->track_particles.m_phi, - amr_data->Geom() - ); - - // gather and space-charge push in x,y,z , assuming the space-charge - // field is the same before/after transformation - // TODO: This is currently using linear order. - spacecharge::GatherAndPush( - *amr_data->track_particles.m_particle_container, - amr_data->track_particles.m_space_charge_field, - amr_data->Geom(), - slice_ds - ); - - // transform from x,y,z to x',y',t - transformation::CoordinateTransformation( - *amr_data->track_particles.m_particle_container, - CoordSystem::s - ); - } - - // for later: original Impact implementation as an option - // Redistribute particles in x',y',t - // TODO: only needed if we want to gather and push space charge - // in x',y',t - // TODO: change geometry beforehand according to transformation - //m_particle_container->Redistribute(); - // - // in original Impact, we gather and space-charge push in x',y',t , - // assuming that the distribution did not change + // Space-charge calculation + particles::spacecharge::HandleSpacecharge(amr_data, [this](){ this->ResizeMesh(); }, slice_ds); // push all particles with external maps Push(*amr_data->track_particles.m_particle_container, element_variant, step, period);