From a6050811855f9d046a59b968c67d95045246f17f Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Mon, 10 Feb 2025 16:17:33 -0800 Subject: [PATCH 01/47] Initial draft commit. --- src/particles/spacecharge/CMakeLists.txt | 1 + .../spacecharge/EnvelopeSpaceChargePush.H | 45 +++++++++++++ .../spacecharge/EnvelopeSpaceChargePush.cpp | 66 +++++++++++++++++++ src/tracking/envelope.cpp | 14 +++- 4 files changed, 123 insertions(+), 3 deletions(-) create mode 100644 src/particles/spacecharge/EnvelopeSpaceChargePush.H create mode 100644 src/particles/spacecharge/EnvelopeSpaceChargePush.cpp diff --git a/src/particles/spacecharge/CMakeLists.txt b/src/particles/spacecharge/CMakeLists.txt index 30e575b43..51e05b7fd 100644 --- a/src/particles/spacecharge/CMakeLists.txt +++ b/src/particles/spacecharge/CMakeLists.txt @@ -3,4 +3,5 @@ target_sources(lib ForceFromSelfFields.cpp GatherAndPush.cpp PoissonSolve.cpp + EnvelopeSpaceChargePush.cpp ) diff --git a/src/particles/spacecharge/EnvelopeSpaceChargePush.H b/src/particles/spacecharge/EnvelopeSpaceChargePush.H new file mode 100644 index 000000000..b748aa0a5 --- /dev/null +++ b/src/particles/spacecharge/EnvelopeSpaceChargePush.H @@ -0,0 +1,45 @@ +/* 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: Chad Mitchell, Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_ENVELOPESPACECHARGEPUSH_H +#define IMPACTX_ENVELOPESPACECHARGEPUSH_H + +#include "particles/ImpactXParticleContainer.H" +#include "particles/CovarianceMatrix.H" + +#include +#include +#include + +#include + + +namespace impactx::spacecharge +{ + /** This function pushes the 6x6 beam covariance matrix for a slice + * of length ds, using the linear space charge fields in an rms + * equivalent 2D ellipse, as determined from the beam covariance matrix. + * Note: This is a reduced model of 2D space charge. + * + * @param[in] refpart reference particle + * @param[inout] cm covariance matrix + * @param[in] current beam current [A] + * @param[in] ds step size [m] + */ + void + envelope_space_charge2D_push ( + RefPart const & refpart, + Map6x6 & cm, + amrex::ParticleReal & current, + amrex::ParticleReal & ds + ); + +} // namespace impactx + +#endif // IMPACTX_ENVELOPESPACECHARGEPUSH_H diff --git a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp new file mode 100644 index 000000000..302066b58 --- /dev/null +++ b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp @@ -0,0 +1,66 @@ +/* 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: Marco Garten, Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#include "EnvelopeSpaceChargePush.H" + +#include +#include // for Real +#include + +namespace impactx::spacecharge +{ + /** This function returns the linear transport map associated with a + * reduced 2D space charge model, based on the beam covariance matrix. + * + * @param[in] refpart reference particle + * @param[in] cm covariance matrix + * @param[in] current beam current [A] + * @param[in] ds step size [m] + * @returns 6x6 transport matrix + */ + AMREX_GPU_HOST AMREX_FORCE_INLINE + void + envelope_space_charge2D_push ( + RefPart const & refpart, + Map6x6 & cm, + amrex::ParticleReal & current, + amrex::ParticleReal & ds + ) + { + using namespace amrex::literals; + + // initialize the linear transport map + Map6x6 R = Map6x6::Identity(); + + // access reference particle values to find beta*gamma^2 + amrex::ParticleReal const pt_ref = refpart.pt; + amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1.0_prt; + + // evaluate the beam space charge perveance from current + // TODO + amrex::ParticleReal const Kpv=0.0_prt; + + // evaluate the linear transfer map + amrex::ParticleReal const sigma2 = cm(1,1)*cm(3,3)-cm(1,3)*cm(1,3); + amrex::ParticleReal const sigma = std::sqrt(sigma2); + amrex::ParticleReal const D = (sigma + cm(1,1)) * (sigma + cm(3,3)) - cm(1,3)*cm(1,3); + amrex::ParticleReal const coeff = Kpv*ds/(2.0_prt*D); + + R(2,1) = coeff * (sigma + cm(3,3)); + R(2,3) = coeff * (-cm(1,3)); + R(4,1) = R(2,3); + R(4,3) = coeff * (sigma + cm(1,1)); + + // update the beam covariance matrix + cm = R * cm * R.transpose(); + + } + + +} // namespace impactx::spacecharge diff --git a/src/tracking/envelope.cpp b/src/tracking/envelope.cpp index 603c011e3..88ee7dae4 100644 --- a/src/tracking/envelope.cpp +++ b/src/tracking/envelope.cpp @@ -11,6 +11,7 @@ #include "initialization/InitAmrCore.H" #include "particles/ImpactXParticleContainer.H" #include "particles/Push.H" +#include "particles/spacecharge/EnvelopeSpaceChargePush.H" #include "diagnostics/DiagnosticOutput.H" #include @@ -77,7 +78,7 @@ namespace impactx amrex::ParmParse const pp_algo("algo"); bool space_charge = false; pp_algo.query("space_charge", space_charge); - AMREX_ALWAYS_ASSERT_WITH_MESSAGE(!space_charge, "Space charge not yet implemented for envelope tracking."); + //AMREX_ALWAYS_ASSERT_WITH_MESSAGE("2D space charge only is implemented for envelope tracking."); if (verbose > 0) { amrex::Print() << " Space Charge effects: " << space_charge << "\n"; @@ -123,7 +124,7 @@ namespace impactx << " slice_step=" << slice_step << "\n"; } - std::visit([&ref, &cm](auto&& element) + std::visit([&ref, &cm, &slice_ds, &space_charge](auto&& element) { // push reference particle in global coordinates { @@ -131,7 +132,14 @@ namespace impactx element(ref); } - // push Covariance Matrix + if (space_charge) + { + // push Covariance Matrix in 2D space charge fields + amrex::ParticleReal current=0.0; //TODO: This must be set. + spacecharge::envelope_space_charge2D_push(ref,cm,current,slice_ds); + } + + // push Covariance Matrix in external fields element(cm, ref); }, element_variant); From 9b464c161a69ce3c2d8ba4e0fd2325d287a4a039 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 11 Feb 2025 00:20:39 +0000 Subject: [PATCH 02/47] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/particles/spacecharge/EnvelopeSpaceChargePush.H | 4 ++-- src/tracking/envelope.cpp | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/particles/spacecharge/EnvelopeSpaceChargePush.H b/src/particles/spacecharge/EnvelopeSpaceChargePush.H index b748aa0a5..68d97dc64 100644 --- a/src/particles/spacecharge/EnvelopeSpaceChargePush.H +++ b/src/particles/spacecharge/EnvelopeSpaceChargePush.H @@ -23,10 +23,10 @@ namespace impactx::spacecharge { /** This function pushes the 6x6 beam covariance matrix for a slice - * of length ds, using the linear space charge fields in an rms + * of length ds, using the linear space charge fields in an rms * equivalent 2D ellipse, as determined from the beam covariance matrix. * Note: This is a reduced model of 2D space charge. - * + * * @param[in] refpart reference particle * @param[inout] cm covariance matrix * @param[in] current beam current [A] diff --git a/src/tracking/envelope.cpp b/src/tracking/envelope.cpp index 88ee7dae4..828aec744 100644 --- a/src/tracking/envelope.cpp +++ b/src/tracking/envelope.cpp @@ -132,7 +132,7 @@ namespace impactx element(ref); } - if (space_charge) + if (space_charge) { // push Covariance Matrix in 2D space charge fields amrex::ParticleReal current=0.0; //TODO: This must be set. From 862d53e91321ea7aae1b61cddb6710352a310dd3 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Mon, 10 Feb 2025 19:07:24 -0800 Subject: [PATCH 03/47] Update EnvelopeSpaceChargePush.H Co-authored-by: Axel Huebl --- src/particles/spacecharge/EnvelopeSpaceChargePush.H | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/particles/spacecharge/EnvelopeSpaceChargePush.H b/src/particles/spacecharge/EnvelopeSpaceChargePush.H index 68d97dc64..619a01e13 100644 --- a/src/particles/spacecharge/EnvelopeSpaceChargePush.H +++ b/src/particles/spacecharge/EnvelopeSpaceChargePush.H @@ -13,11 +13,7 @@ #include "particles/ImpactXParticleContainer.H" #include "particles/CovarianceMatrix.H" -#include -#include -#include - -#include +#include namespace impactx::spacecharge From 5925584a99c06ac828db954ddabf112ba7846152 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Mon, 10 Feb 2025 19:07:33 -0800 Subject: [PATCH 04/47] Update EnvelopeSpaceChargePush.cpp Co-authored-by: Axel Huebl --- src/particles/spacecharge/EnvelopeSpaceChargePush.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp index 302066b58..c422032a9 100644 --- a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp +++ b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp @@ -13,6 +13,8 @@ #include // for Real #include +#include + namespace impactx::spacecharge { /** This function returns the linear transport map associated with a From 25f9185d0d3eb1725f21f0a3da42059e9d3d27e3 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Mon, 10 Feb 2025 19:07:40 -0800 Subject: [PATCH 05/47] Update EnvelopeSpaceChargePush.cpp Co-authored-by: Axel Huebl --- src/particles/spacecharge/EnvelopeSpaceChargePush.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp index c422032a9..18b29e743 100644 --- a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp +++ b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp @@ -21,7 +21,7 @@ namespace impactx::spacecharge * reduced 2D space charge model, based on the beam covariance matrix. * * @param[in] refpart reference particle - * @param[in] cm covariance matrix + * @param[in,out] cm covariance matrix * @param[in] current beam current [A] * @param[in] ds step size [m] * @returns 6x6 transport matrix From 8130f436710655941a31c5c37460ee8941ce4835 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Tue, 11 Feb 2025 10:01:47 -0800 Subject: [PATCH 06/47] Apply suggestions from code review Co-authored-by: Axel Huebl --- src/particles/spacecharge/EnvelopeSpaceChargePush.H | 4 ++-- src/particles/spacecharge/EnvelopeSpaceChargePush.cpp | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/particles/spacecharge/EnvelopeSpaceChargePush.H b/src/particles/spacecharge/EnvelopeSpaceChargePush.H index 619a01e13..c4c9a7644 100644 --- a/src/particles/spacecharge/EnvelopeSpaceChargePush.H +++ b/src/particles/spacecharge/EnvelopeSpaceChargePush.H @@ -32,8 +32,8 @@ namespace impactx::spacecharge envelope_space_charge2D_push ( RefPart const & refpart, Map6x6 & cm, - amrex::ParticleReal & current, - amrex::ParticleReal & ds + amrex::ParticleReal current, + amrex::ParticleReal ds ); } // namespace impactx diff --git a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp index 18b29e743..6a1628181 100644 --- a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp +++ b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp @@ -31,8 +31,8 @@ namespace impactx::spacecharge envelope_space_charge2D_push ( RefPart const & refpart, Map6x6 & cm, - amrex::ParticleReal & current, - amrex::ParticleReal & ds + amrex::ParticleReal current, + amrex::ParticleReal ds ) { using namespace amrex::literals; From 98a9fb14a4b787fe6c95c9d8bd6f1f585614f81a 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 07/47] [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 cd1d75ccca5c4a0c433a8335ce28e231ab04fd21 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 10 Feb 2025 15:40:28 -0800 Subject: [PATCH 08/47] 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 f43668df7b9910f3c22e5e8339c9ae796cabec7b Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 10 Feb 2025 17:03:43 -0800 Subject: [PATCH 09/47] 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 d37f8108254fef767eb143e289bf879bacc27d4e Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 11 Feb 2025 10:14:13 -0800 Subject: [PATCH 10/47] 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 8f34ad6f13bcbccc905619be84883230850c13b9 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 11/47] 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 41cd4b071857b9980b483aba9314098ce6ab690c 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 12/47] 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 99f69b7cb7082776e22b15bd72b56d7680789e3f Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Tue, 11 Feb 2025 11:54:31 -0800 Subject: [PATCH 13/47] Remove AMREX_FORCE_INLINE --- .../spacecharge/EnvelopeSpaceChargePush.cpp | 20 +++++-------------- src/tracking/envelope.cpp | 16 +++++++-------- 2 files changed, 13 insertions(+), 23 deletions(-) diff --git a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp index 6a1628181..912415f10 100644 --- a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp +++ b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp @@ -17,21 +17,11 @@ namespace impactx::spacecharge { - /** This function returns the linear transport map associated with a - * reduced 2D space charge model, based on the beam covariance matrix. - * - * @param[in] refpart reference particle - * @param[in,out] cm covariance matrix - * @param[in] current beam current [A] - * @param[in] ds step size [m] - * @returns 6x6 transport matrix - */ - AMREX_GPU_HOST AMREX_FORCE_INLINE void envelope_space_charge2D_push ( - RefPart const & refpart, - Map6x6 & cm, - amrex::ParticleReal current, + [[maybe_unused]] RefPart const & refpart, + [[maybe_unused]] Map6x6 & cm, + [[maybe_unused]] amrex::ParticleReal current, amrex::ParticleReal ds ) { @@ -41,8 +31,8 @@ namespace impactx::spacecharge Map6x6 R = Map6x6::Identity(); // access reference particle values to find beta*gamma^2 - amrex::ParticleReal const pt_ref = refpart.pt; - amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1.0_prt; + //amrex::ParticleReal const pt_ref = refpart.pt; + //amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1.0_prt; // evaluate the beam space charge perveance from current // TODO diff --git a/src/tracking/envelope.cpp b/src/tracking/envelope.cpp index 828aec744..1610011bc 100644 --- a/src/tracking/envelope.cpp +++ b/src/tracking/envelope.cpp @@ -124,7 +124,14 @@ namespace impactx << " slice_step=" << slice_step << "\n"; } - std::visit([&ref, &cm, &slice_ds, &space_charge](auto&& element) + if (space_charge) + { + // push Covariance Matrix in 2D space charge fields + amrex::ParticleReal current=0.0; //TODO: This must be set. + spacecharge::envelope_space_charge2D_push(ref,cm,current,slice_ds); + } + + std::visit([&ref, &cm](auto&& element) { // push reference particle in global coordinates { @@ -132,13 +139,6 @@ namespace impactx element(ref); } - if (space_charge) - { - // push Covariance Matrix in 2D space charge fields - amrex::ParticleReal current=0.0; //TODO: This must be set. - spacecharge::envelope_space_charge2D_push(ref,cm,current,slice_ds); - } - // push Covariance Matrix in external fields element(cm, ref); From 1ea852857688ba69f393a0c1e4912bbdbe70b48e Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Tue, 11 Feb 2025 13:45:24 -0800 Subject: [PATCH 14/47] Add calculation of beam perveance. --- src/initialization/InitDistribution.cpp | 8 ++++++++ .../spacecharge/EnvelopeSpaceChargePush.cpp | 17 +++++++++++------ 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 63d34a2cb..11a867d8d 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -471,6 +471,14 @@ namespace impactx } else if (track == "envelope") { + // For treating 3D space charge in bunched beams (not yet implemented in envelope mode) + //amrex::ParticleReal bunch_charge = 0.0; // Bunch charge (C) + //pp_dist.getWithParser("charge", bunch_charge); + + // For treating 2D space charge in unbunched beams + amrex::ParticleReal beam_current = 0.0; // Beam current (A) + pp_dist.getWithParser("current", beam_current); + amr_data->track_envelope.m_ref = initialization::read_reference_particle(pp_dist); auto dist = initialization::read_distribution(pp_dist); amr_data->track_envelope.m_cm = impactx::initialization::create_covariance_matrix(dist); diff --git a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp index 912415f10..3f952f250 100644 --- a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp +++ b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp @@ -20,7 +20,7 @@ namespace impactx::spacecharge void envelope_space_charge2D_push ( [[maybe_unused]] RefPart const & refpart, - [[maybe_unused]] Map6x6 & cm, + Map6x6 & cm, [[maybe_unused]] amrex::ParticleReal current, amrex::ParticleReal ds ) @@ -30,13 +30,18 @@ namespace impactx::spacecharge // initialize the linear transport map Map6x6 R = Map6x6::Identity(); - // access reference particle values to find beta*gamma^2 - //amrex::ParticleReal const pt_ref = refpart.pt; - //amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1.0_prt; + // physical constants and reference quantities + amrex::ParticleReal const c = ablastr::constant::SI::c; + amrex::ParticleReal const ep0 = ablastr::constant::SI::ep0; + amrex::ParticleReal const pi = ablastr::constant::math::pi; + amrex::ParticleReal const mass = refpart.mass; + amrex::ParticleReal const charge = refpart.charge; + amrex::ParticleReal const pt_ref = refpart.pt; + amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1.0_prt; // evaluate the beam space charge perveance from current - // TODO - amrex::ParticleReal const Kpv=0.0_prt; + amrex::ParticleReal const IA = 4.0_prt*pi*ep0*mass*pow(c,3)/charge; + amrex::ParticleReal const Kpv = (current/IA) * 2.0_prt/betgam2; // evaluate the linear transfer map amrex::ParticleReal const sigma2 = cm(1,1)*cm(3,3)-cm(1,3)*cm(1,3); From dbf41d4c34d1b47e8cc4af70c39c6c731c05e504 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Tue, 11 Feb 2025 14:11:21 -0800 Subject: [PATCH 15/47] Fix ParmParse. --- src/particles/spacecharge/EnvelopeSpaceChargePush.cpp | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp index 3f952f250..a4395f25b 100644 --- a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp +++ b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp @@ -12,6 +12,7 @@ #include #include // for Real #include +#include #include @@ -30,6 +31,11 @@ namespace impactx::spacecharge // initialize the linear transport map Map6x6 R = Map6x6::Identity(); + // added temporarily for benchmark testing + amrex::ParmParse pp_dist("dist"); + amrex::ParticleReal beam_current = 0.0; // Beam current (A) + pp_algo.query("current", beam_current); + // physical constants and reference quantities amrex::ParticleReal const c = ablastr::constant::SI::c; amrex::ParticleReal const ep0 = ablastr::constant::SI::ep0; @@ -41,7 +47,7 @@ namespace impactx::spacecharge // evaluate the beam space charge perveance from current amrex::ParticleReal const IA = 4.0_prt*pi*ep0*mass*pow(c,3)/charge; - amrex::ParticleReal const Kpv = (current/IA) * 2.0_prt/betgam2; + amrex::ParticleReal const Kpv = (beam_current/IA) * 2.0_prt/betgam2; // evaluate the linear transfer map amrex::ParticleReal const sigma2 = cm(1,1)*cm(3,3)-cm(1,3)*cm(1,3); From 039758726cafcd6a3927df32612bbc933b7aa434 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Tue, 11 Feb 2025 14:18:20 -0800 Subject: [PATCH 16/47] Update EnvelopeSpaceChargePush.cpp --- src/particles/spacecharge/EnvelopeSpaceChargePush.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp index a4395f25b..b3c3c63dd 100644 --- a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp +++ b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp @@ -34,7 +34,7 @@ namespace impactx::spacecharge // added temporarily for benchmark testing amrex::ParmParse pp_dist("dist"); amrex::ParticleReal beam_current = 0.0; // Beam current (A) - pp_algo.query("current", beam_current); + pp_dist.query("current", beam_current); // physical constants and reference quantities amrex::ParticleReal const c = ablastr::constant::SI::c; From dd171b965b1259238b89b3337b2d3fb64c68d5bd Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Tue, 11 Feb 2025 14:31:57 -0800 Subject: [PATCH 17/47] Update InitDistribution.cpp --- src/initialization/InitDistribution.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 11a867d8d..53ef4bb2a 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -477,7 +477,7 @@ namespace impactx // For treating 2D space charge in unbunched beams amrex::ParticleReal beam_current = 0.0; // Beam current (A) - pp_dist.getWithParser("current", beam_current); + pp_dist.query("current", beam_current); amr_data->track_envelope.m_ref = initialization::read_reference_particle(pp_dist); auto dist = initialization::read_distribution(pp_dist); From 1878ab00ef1a6e3f811ea6fcc89bd1c391dc89d7 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Thu, 13 Feb 2025 19:33:35 -0800 Subject: [PATCH 18/47] Input file for benchmark example. --- examples/fodo/input_fodo_envelope_sc.in | 54 +++++++++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 examples/fodo/input_fodo_envelope_sc.in diff --git a/examples/fodo/input_fodo_envelope_sc.in b/examples/fodo/input_fodo_envelope_sc.in new file mode 100644 index 000000000..c70eb8068 --- /dev/null +++ b/examples/fodo/input_fodo_envelope_sc.in @@ -0,0 +1,54 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.kin_energy = 6.7 +beam.particle = proton +beam.distribution = kvdist_from_twiss +beam.alphaX = 2.4685083 +beam.alphaY = -beam.alphaX +beam.alphaT = 0.0 +beam.betaX = 0.737881 +beam.betaY = 0.737881 +beam.betaT = 0.5 +beam.emittX = 1.0e-6 +beam.emittY = beam.emittX +beam.emittT = 1.0e-12 + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.elements = drift1 quad1 drift2 quad2 drift1 +lattice.nslice = 50 + +monitor.type = beam_monitor +monitor.backend = h5 + +drift1.type = drift +drift1.ds = 7.44e-2 + +quad1.type = quad +quad1.ds = 6.10e-2 +quad1.k = -103.12574100336 + +drift2.type = drift +drift2.ds = 14.88e-2 + +quad2.type = quad +quad2.ds = 6.10e-2 +quad2.k = -quad1.k + +drift3.type = drift +drift3.ds = 0.25 + + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.track = "envelope" +algo.space_charge = true + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true From 7c396624913bf20b179a43f881df9540446a7e8f Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Thu, 13 Feb 2025 19:36:59 -0800 Subject: [PATCH 19/47] Correct perveance calculation. Hard-coded current for testing. --- .../spacecharge/EnvelopeSpaceChargePush.cpp | 12 +++++------- src/tracking/envelope.cpp | 2 +- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp index b3c3c63dd..18b8aa5f7 100644 --- a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp +++ b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp @@ -13,6 +13,7 @@ #include // for Real #include #include +#include #include @@ -22,7 +23,7 @@ namespace impactx::spacecharge envelope_space_charge2D_push ( [[maybe_unused]] RefPart const & refpart, Map6x6 & cm, - [[maybe_unused]] amrex::ParticleReal current, + amrex::ParticleReal current, amrex::ParticleReal ds ) { @@ -31,11 +32,6 @@ namespace impactx::spacecharge // initialize the linear transport map Map6x6 R = Map6x6::Identity(); - // added temporarily for benchmark testing - amrex::ParmParse pp_dist("dist"); - amrex::ParticleReal beam_current = 0.0; // Beam current (A) - pp_dist.query("current", beam_current); - // physical constants and reference quantities amrex::ParticleReal const c = ablastr::constant::SI::c; amrex::ParticleReal const ep0 = ablastr::constant::SI::ep0; @@ -44,10 +40,12 @@ namespace impactx::spacecharge amrex::ParticleReal const charge = refpart.charge; amrex::ParticleReal const pt_ref = refpart.pt; amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1.0_prt; + amrex::ParticleReal const betgam = std::sqrt(betgam2); + amrex::ParticleReal const betgam3 = std::pow(betgam,3); // evaluate the beam space charge perveance from current amrex::ParticleReal const IA = 4.0_prt*pi*ep0*mass*pow(c,3)/charge; - amrex::ParticleReal const Kpv = (beam_current/IA) * 2.0_prt/betgam2; + amrex::ParticleReal const Kpv = std::abs(current/IA) * 2.0_prt/betgam3; // evaluate the linear transfer map amrex::ParticleReal const sigma2 = cm(1,1)*cm(3,3)-cm(1,3)*cm(1,3); diff --git a/src/tracking/envelope.cpp b/src/tracking/envelope.cpp index 1610011bc..dff16a9f3 100644 --- a/src/tracking/envelope.cpp +++ b/src/tracking/envelope.cpp @@ -127,7 +127,7 @@ namespace impactx if (space_charge) { // push Covariance Matrix in 2D space charge fields - amrex::ParticleReal current=0.0; //TODO: This must be set. + amrex::ParticleReal current=0.5; //TODO: This must be set. spacecharge::envelope_space_charge2D_push(ref,cm,current,slice_ds); } From 04ba2ba0f430de1f4e79a177ed45193b8242f07d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 14 Feb 2025 03:37:14 +0000 Subject: [PATCH 20/47] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/particles/spacecharge/EnvelopeSpaceChargePush.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp index 18b8aa5f7..b090932c5 100644 --- a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp +++ b/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp @@ -13,7 +13,7 @@ #include // for Real #include #include -#include +#include #include From 70183351e8f3d264726c0f287649c85a98cd22fd Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Fri, 14 Feb 2025 10:18:27 -0800 Subject: [PATCH 21/47] Update benchmark example input files. --- docs/source/usage/examples.rst | 1 + examples/CMakeLists.txt | 17 +++ examples/fodo_space_charge/README.rst | 75 +++++++++++ .../analysis_fodo_envelope_sc.py | 116 ++++++++++++++++++ .../input_fodo_envelope_sc.in | 0 .../fodo_space_charge/run_fodo_envelope_sc.py | 68 ++++++++++ 6 files changed, 277 insertions(+) create mode 100644 examples/fodo_space_charge/README.rst create mode 100755 examples/fodo_space_charge/analysis_fodo_envelope_sc.py rename examples/{fodo => fodo_space_charge}/input_fodo_envelope_sc.in (100%) create mode 100755 examples/fodo_space_charge/run_fodo_envelope_sc.py diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst index f038a2bf8..71e685f19 100644 --- a/docs/source/usage/examples.rst +++ b/docs/source/usage/examples.rst @@ -52,6 +52,7 @@ Space Charge examples/cfchannel/README.rst examples/kurth/README.rst examples/epac2004_benchmarks/README.rst + examples/fodo_space_charge/README.rst Coherent Synchrotron Radiation (CSR) """""""""""""""""""""""""""""""""""" diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 370f3e6ce..d73fe8b70 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -1269,3 +1269,20 @@ add_impactx_test(examples-scraping.py examples/scraping_beam/analysis_scraping.py OFF # no plot script yet ) + +# FODO cell with 2D space charge using envelope tracking ###################### +# +# with space charge +add_impactx_test(fodo-envelope-sc + examples/fodo_space_charge/input_fodo_envelope_sc.in + ON # ImpactX MPI-parallel + examples/fodo_space_charge/analysis_fodo_envelope_sc.py + OFF # no plot script yet +) +add_impactx_test(fodo-envelope-sc.py + examples/fodo_space_charge/run_fodo_envelope_sc.py + OFF # ImpactX MPI-parallel + examples/fodo_space_charge/analysis_fodo_envelope_sc.py + OFF # no plot script yet +) + diff --git a/examples/fodo_space_charge/README.rst b/examples/fodo_space_charge/README.rst new file mode 100644 index 000000000..29b121ee2 --- /dev/null +++ b/examples/fodo_space_charge/README.rst @@ -0,0 +1,75 @@ +.. _examples-fodo-envelope-sc: + +FODO Cell with 2D Space Charge using Envelope Tracking +======================================================= + +This example illustrates a 0.5 A proton beam with a kinetic energy of 6.7 MeV in a FODO cell, +with 2D space charge included. The parameters are those described in: + +R.D. Ryne et al, "A Test Suite of Space-Charge Problems for Code Benchmarking," +in Proc. EPAC 2004, Lucerne, Switzerland: KV Beam in a FODO Channel + +The purpose of this example is to illustrate the use of envelope tracking mode with 2D space charge. + +The second moments of the particle distribution after the FODO cell should coincide with the second moments of the particle distribution before the FODO cell, to within the level expected due to noise due to statistical sampling. + +In this test, the initial and final values of :math:`\lambda_x`, :math:`\lambda_y`, :math:`\lambda_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_fodo.py`` or +* ImpactX **executable** using an input file: ``impactx input_fodo.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_fodo.py + :language: python3 + :caption: You can copy this file from ``examples/fodo/run_fodo.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_fodo.in + :language: ini + :caption: You can copy this file from ``examples/fodo/input_fodo.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_fodo.py`` + + .. literalinclude:: analysis_fodo.py + :language: python3 + :caption: You can copy this file from ``examples/fodo/analysis_fodo.py``. + + +Visualize +--------- + +You can run the following script to visualize the beam evolution over time: + +.. dropdown:: Script ``plot_fodo.py`` + + .. literalinclude:: plot_fodo.py + :language: python3 + :caption: You can copy this file from ``examples/fodo/plot_fodo.py``. + +.. figure:: https://user-images.githubusercontent.com/1353258/180287840-8561f6fd-278f-4856-abd8-04fbdb78c8ff.png + :alt: focusing, defocusing and preserved emittance in our FODO cell benchmark. + + FODO transversal beam width and emittance evolution + +.. figure:: https://user-images.githubusercontent.com/1353258/180287845-eb0210a7-2500-4aa9-844c-67fb094329d3.png + :alt: focusing, defocusing and phase space rotation in our FODO cell benchmark. + + FODO transversal beam width and phase space evolution diff --git a/examples/fodo_space_charge/analysis_fodo_envelope_sc.py b/examples/fodo_space_charge/analysis_fodo_envelope_sc.py new file mode 100755 index 000000000..da2a8fba8 --- /dev/null +++ b/examples/fodo_space_charge/analysis_fodo_envelope_sc.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# + +import glob +import re + +import numpy as np +import pandas as pd + + +def read_file(file_pattern): + for filename in glob.glob(file_pattern): + df = pd.read_csv(filename, delimiter=r"\s+") + if "step" not in df.columns: + step = int(re.findall(r"[0-9]+", filename)[0]) + df["step"] = step + yield df + + +def read_time_series(file_pattern): + """Read in all CSV files from each MPI rank (and potentially OpenMP + thread). Concatenate into one Pandas dataframe. + + Returns + ------- + pandas.DataFrame + """ + return pd.concat( + read_file(file_pattern), + axis=0, + ignore_index=True, + ) # .set_index('id') + + +# read reduced diagnostics +rbc = read_time_series("diags/reduced_beam_characteristics.*") + +s = rbc["s"] +sig_x = rbc["sig_x"] +sig_y = rbc["sig_y"] +sig_t = rbc["sig_t"] +emittance_x = rbc["emittance_x"] +emittance_y = rbc["emittance_y"] +emittance_t = rbc["emittance_t"] + +sig_xi = sig_x.iloc[0] +sig_yi = sig_y.iloc[0] +sig_ti = sig_t.iloc[0] +emittance_xi = emittance_x.iloc[0] +emittance_yi = emittance_y.iloc[0] +emittance_ti = emittance_t.iloc[0] + +length = len(s) - 1 + +sf = s.iloc[length] +sig_xf = sig_x.iloc[length] +sig_yf = sig_y.iloc[length] +sig_tf = sig_t.iloc[length] +emittance_xf = emittance_x.iloc[length] +emittance_yf = emittance_y.iloc[length] +emittance_tf = emittance_t.iloc[length] + + +print("Initial Beam:") +print(f" sigx={sig_xi:e} sigy={sig_yi:e} sigt={sig_ti:e}") +print( + f" emittance_x={emittance_xi:e} emittance_y={emittance_yi:e} emittance_t={emittance_ti:e}" +) + +atol = 0.0 # ignored +rtol = 1.0e-3 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sig_xi, sig_yi, sig_ti, emittance_xi, emittance_yi, emittance_ti], + [ + 8.590000e-04, + 8.590000e-04, + 7.071068e-07, + 1.000000e-06, + 1.000000e-06, + 1.000000e-12, + ], + rtol=rtol, + atol=atol, +) + + +print("") +print("Final Beam:") +print(f" sigx={sig_xf:e} sigy={sig_yf:e} sigt={sig_tf:e}") +print( + f" emittance_x={emittance_xf:e} emittance_y={emittance_yf:e} emittance_t={emittance_tf:e}" +) + +atol = 0.0 # ignored +rtol = 1.0e-3 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sig_xf, sig_yf, sig_tf, emittance_xf, emittance_yf, emittance_tf], + [ + 8.590000e-04, + 8.590000e-04, + 4.140854e-05, + 1.000000e-06, + 1.000000e-06, + 1.000000e-12, + ], + rtol=rtol, + atol=atol, +) diff --git a/examples/fodo/input_fodo_envelope_sc.in b/examples/fodo_space_charge/input_fodo_envelope_sc.in similarity index 100% rename from examples/fodo/input_fodo_envelope_sc.in rename to examples/fodo_space_charge/input_fodo_envelope_sc.in diff --git a/examples/fodo_space_charge/run_fodo_envelope_sc.py b/examples/fodo_space_charge/run_fodo_envelope_sc.py new file mode 100755 index 000000000..eab3c5577 --- /dev/null +++ b/examples/fodo_space_charge/run_fodo_envelope_sc.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from impactx import ImpactX, distribution, elements, twiss + +sim = ImpactX() + +# set numerical parameters and IO control +sim.particle_shape = 2 # B-spline order +sim.space_charge = True +# sim.diagnostics = False # benchmarking +sim.slice_step_diagnostics = True + +# domain decomposition & space charge mesh +sim.init_grids() + +# model a 2 GeV electron beam with an initial +# unnormalized rms emittance of 2 nm +kin_energy_MeV = 6.7 # reference energy + +# reference particle +ref = sim.particle_container().ref_particle() +ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_kin_energy_MeV(kin_energy_MeV) + +# particle bunch +distr = distribution.KVdist( + **twiss( + beta_x=0.737881, + beta_y=0.737881, + beta_t=0.5, + emitt_x=1.0e-6, + emitt_y=1.0e-6, + emitt_t=1.0e-12, + alpha_x=2.4685083, + alpha_y=-2.4685083, + alpha_t=0.0, + ) +) + +sim.init_envelope(ref, distr) + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5") + +# design the accelerator lattice) +ns = 50 # number of slices per ds in the element +fodo = [ + monitor, + elements.Drift(name="drift1", ds=7.44e-2, nslice=ns), + elements.Quad(name="quad1", ds=6.10e-2, k=-103.12574100336, nslice=ns), + elements.Drift(name="drift2", ds=14.88e-2, nslice=ns), + elements.Quad(name="quad2", ds=6.10e-2, k=103.12574100336, nslice=ns), + elements.Drift(name="drift3", ds=7.44e-2, nslice=ns), + monitor, +] +# assign a fodo segment +sim.lattice.extend(fodo) + +# run simulation +sim.track_envelope() + +# clean shutdown +sim.finalize() From a7ee8b3978500bc3b3ae069f50c8aca5a4c808d7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 14 Feb 2025 18:18:59 +0000 Subject: [PATCH 22/47] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- examples/CMakeLists.txt | 1 - examples/fodo_space_charge/README.rst | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index d73fe8b70..8d232f39b 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -1285,4 +1285,3 @@ add_impactx_test(fodo-envelope-sc.py examples/fodo_space_charge/analysis_fodo_envelope_sc.py OFF # no plot script yet ) - diff --git a/examples/fodo_space_charge/README.rst b/examples/fodo_space_charge/README.rst index 29b121ee2..3ccdb05e6 100644 --- a/examples/fodo_space_charge/README.rst +++ b/examples/fodo_space_charge/README.rst @@ -3,7 +3,7 @@ FODO Cell with 2D Space Charge using Envelope Tracking ======================================================= -This example illustrates a 0.5 A proton beam with a kinetic energy of 6.7 MeV in a FODO cell, +This example illustrates a 0.5 A proton beam with a kinetic energy of 6.7 MeV in a FODO cell, with 2D space charge included. The parameters are those described in: R.D. Ryne et al, "A Test Suite of Space-Charge Problems for Code Benchmarking," From dc32c6d5280fc88923fd696c6ac97d8fb0b2070b Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Fri, 14 Feb 2025 10:46:33 -0800 Subject: [PATCH 23/47] Fix unused ParmParse variable error. --- examples/fodo_space_charge/input_fodo_envelope_sc.in | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/fodo_space_charge/input_fodo_envelope_sc.in b/examples/fodo_space_charge/input_fodo_envelope_sc.in index c70eb8068..d29c99eaf 100644 --- a/examples/fodo_space_charge/input_fodo_envelope_sc.in +++ b/examples/fodo_space_charge/input_fodo_envelope_sc.in @@ -17,7 +17,7 @@ beam.emittT = 1.0e-12 ############################################################################### # Beamline: lattice elements and segments ############################################################################### -lattice.elements = drift1 quad1 drift2 quad2 drift1 +lattice.elements = monitor drift1 quad1 drift2 quad2 drift1 monitor lattice.nslice = 50 monitor.type = beam_monitor From 6effcc91b8dd30bfda1821ccd57098e3594a6ef6 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Fri, 14 Feb 2025 11:13:18 -0800 Subject: [PATCH 24/47] Remove unused drift3. --- examples/fodo_space_charge/input_fodo_envelope_sc.in | 4 ---- 1 file changed, 4 deletions(-) diff --git a/examples/fodo_space_charge/input_fodo_envelope_sc.in b/examples/fodo_space_charge/input_fodo_envelope_sc.in index d29c99eaf..8f94b124f 100644 --- a/examples/fodo_space_charge/input_fodo_envelope_sc.in +++ b/examples/fodo_space_charge/input_fodo_envelope_sc.in @@ -37,10 +37,6 @@ quad2.type = quad quad2.ds = 6.10e-2 quad2.k = -quad1.k -drift3.type = drift -drift3.ds = 0.25 - - ############################################################################### # Algorithms ############################################################################### From 796f19d12462c78e46c847a95ada2cc805c0eaca Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Fri, 14 Feb 2025 15:07:32 -0800 Subject: [PATCH 25/47] Add support for current input, some refactoring. --- .../input_fodo_envelope_sc.in | 1 + .../fodo_space_charge/run_fodo_envelope_sc.py | 5 +- src/initialization/AmrCoreData.H | 2 +- src/initialization/InitDistribution.H | 5 +- src/initialization/InitDistribution.cpp | 21 +++++--- src/particles/CovarianceMatrix.H | 53 ++++++++++++++++++- src/python/ImpactX.cpp | 6 +-- src/python/distribution.cpp | 2 +- src/tracking/envelope.cpp | 7 +-- 9 files changed, 82 insertions(+), 20 deletions(-) diff --git a/examples/fodo_space_charge/input_fodo_envelope_sc.in b/examples/fodo_space_charge/input_fodo_envelope_sc.in index 8f94b124f..8148962fc 100644 --- a/examples/fodo_space_charge/input_fodo_envelope_sc.in +++ b/examples/fodo_space_charge/input_fodo_envelope_sc.in @@ -2,6 +2,7 @@ # Particle Beam(s) ############################################################################### beam.kin_energy = 6.7 +beam.current = 0.5 beam.particle = proton beam.distribution = kvdist_from_twiss beam.alphaX = 2.4685083 diff --git a/examples/fodo_space_charge/run_fodo_envelope_sc.py b/examples/fodo_space_charge/run_fodo_envelope_sc.py index eab3c5577..d711cfbef 100755 --- a/examples/fodo_space_charge/run_fodo_envelope_sc.py +++ b/examples/fodo_space_charge/run_fodo_envelope_sc.py @@ -27,6 +27,9 @@ ref = sim.particle_container().ref_particle() ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_kin_energy_MeV(kin_energy_MeV) +# beam current in A +beam_current_A = 0.5 + # particle bunch distr = distribution.KVdist( **twiss( @@ -42,7 +45,7 @@ ) ) -sim.init_envelope(ref, distr) +sim.init_envelope(ref, distr, beam_current_A) # add beam diagnostics monitor = elements.BeamMonitor("monitor", backend="h5") diff --git a/src/initialization/AmrCoreData.H b/src/initialization/AmrCoreData.H index b4d4a312f..1a457a6c4 100644 --- a/src/initialization/AmrCoreData.H +++ b/src/initialization/AmrCoreData.H @@ -89,7 +89,7 @@ namespace impactx::initialization std::optional m_ref; /** The 6x6 covariance matrix for the beam envelope, relative to the reference particle */ - std::optional m_cm; + std::optional m_env; } track_envelope; diff --git a/src/initialization/InitDistribution.H b/src/initialization/InitDistribution.H index 90e9bb108..6e73a77b3 100644 --- a/src/initialization/InitDistribution.H +++ b/src/initialization/InitDistribution.H @@ -40,8 +40,9 @@ namespace impactx::initialization /** Ignore the shape of a distribution and use the 2nd moments to create a covariance matrix */ - CovarianceMatrix - create_covariance_matrix ( + Envelope + create_envelope ( + amrex::ParticleReal const current, distribution::KnownDistributions const & distr ); diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 53ef4bb2a..af3b321d3 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -185,8 +185,9 @@ namespace impactx return dist; } - CovarianceMatrix - initialization::create_covariance_matrix ( + Envelope + initialization::create_envelope ( + amrex::ParticleReal const current, distribution::KnownDistributions const & distr ) { @@ -233,7 +234,11 @@ namespace impactx } }, distr); - return cv; + Envelope env; + env.set_beam_current_A(current); + env.set_covariance_matrix(cv); + + return env; } void @@ -471,17 +476,17 @@ namespace impactx } else if (track == "envelope") { - // For treating 3D space charge in bunched beams (not yet implemented in envelope mode) - //amrex::ParticleReal bunch_charge = 0.0; // Bunch charge (C) - //pp_dist.getWithParser("charge", bunch_charge); + // relevant for envelope tracking with 3D space charge + amrex::ParticleReal bunch_charge = 0.0; // Bunch charge (C) + pp_dist.query("charge", bunch_charge); - // For treating 2D space charge in unbunched beams + // relevant for envelope tracking with 2D space charge amrex::ParticleReal beam_current = 0.0; // Beam current (A) pp_dist.query("current", beam_current); amr_data->track_envelope.m_ref = initialization::read_reference_particle(pp_dist); auto dist = initialization::read_distribution(pp_dist); - amr_data->track_envelope.m_cm = impactx::initialization::create_covariance_matrix(dist); + amr_data->track_envelope.m_env = impactx::initialization::create_envelope(beam_current,dist); } else if (track == "reference_orbit") { diff --git a/src/particles/CovarianceMatrix.H b/src/particles/CovarianceMatrix.H index b55ed66df..6562edd54 100644 --- a/src/particles/CovarianceMatrix.H +++ b/src/particles/CovarianceMatrix.H @@ -19,9 +19,60 @@ namespace impactx /** this is a 6x6 matrix */ using Map6x6 = amrex::SmallMatrix; - /** the covariance matrix is 6x6 */ + /** this is the 6x6 covariance matrix */ using CovarianceMatrix = Map6x6; + + /** This struct stores the beam envelope attribues, including the 6x6 beam + * covariance matrix. Used during envelope tracking mode. + */ + struct Envelope + { + amrex::ParticleReal bunch_charge = 0.0; /// total charge in C (for 3D space charge) + amrex::ParticleReal beam_current = 0.0; /// current in A (for 2D space charge) + CovarianceMatrix cm; /// the 6x6 beam covariance matrix + + /** Set envelope beam current for 2D space charge + * + * @param beam_current beam current (A) + */ + Envelope & + set_beam_current_A (amrex::ParticleReal const beam_current_A) + { + using namespace amrex::literals; + + beam_current = beam_current_A; + + return *this; + } + + /** Get envelope beam current for 2D space charge + * + * @returns beam current in A + */ + amrex::ParticleReal + beam_current_A () + { + using namespace amrex::literals; + + return beam_current; + } + + /** Set 6x6 covariance matrix for envelope tracking + * + * @param Map6x6 beam 6x6 covariance matrix + */ + Envelope & + set_covariance_matrix (CovarianceMatrix const covariance_matrix) + { + cm = covariance_matrix; + + return *this; + } + + }; + + } // namespace impactx::distribution #endif // IMPACTX_DISTRIBUTION_COVARIANCE_MATRIX_H diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index 01e39021f..2c8782948 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -438,11 +438,11 @@ void init_ImpactX (py::module& m) .def("init_beam_distribution_from_inputs", &ImpactX::initBeamDistributionFromInputs) .def("init_lattice_elements_from_inputs", &ImpactX::initLatticeElementsFromInputs) .def("init_envelope", - [](ImpactX & ix, RefPart ref, distribution::KnownDistributions distr) { + [](ImpactX & ix, RefPart ref, distribution::KnownDistributions distr, amrex::Real current) { ix.amr_data->track_envelope.m_ref = ref; - ix.amr_data->track_envelope.m_cm = initialization::create_covariance_matrix(distr); + ix.amr_data->track_envelope.m_env = initialization::create_envelope(current,distr); }, - py::arg("ref"), py::arg("distr"), + py::arg("ref"), py::arg("distr"), py::arg("current"), "Envelope tracking mode:" "Create a 6x6 covariance matrix from a distribution and then initialize " "the the simulation for envelope tracking relative to a reference particle." diff --git a/src/python/distribution.cpp b/src/python/distribution.cpp index 2659a5f76..fe484547a 100644 --- a/src/python/distribution.cpp +++ b/src/python/distribution.cpp @@ -131,5 +131,5 @@ void init_distribution(py::module& m) "A 6D Waterbag distribution" ); - m.def("create_covariance_matrix", &initialization::create_covariance_matrix); + m.def("create_covariance_matrix", &initialization::create_envelope); } diff --git a/src/tracking/envelope.cpp b/src/tracking/envelope.cpp index dff16a9f3..110d6b89d 100644 --- a/src/tracking/envelope.cpp +++ b/src/tracking/envelope.cpp @@ -47,12 +47,14 @@ namespace impactx { throw std::runtime_error("track_envelope: Reference particle not set."); } - if (!amr_data->track_envelope.m_cm.has_value()) + if (!amr_data->track_envelope.m_env.has_value()) { throw std::runtime_error("track_envelope: Envelope (covariance matrix) not set."); } auto & ref = amr_data->track_envelope.m_ref.value(); - auto & cm = amr_data->track_envelope.m_cm.value(); + auto & env = amr_data->track_envelope.m_env.value(); + auto & cm = env.cm; + auto & current = env.beam_current; // output of init state amrex::ParmParse pp_diag("diag"); @@ -127,7 +129,6 @@ namespace impactx if (space_charge) { // push Covariance Matrix in 2D space charge fields - amrex::ParticleReal current=0.5; //TODO: This must be set. spacecharge::envelope_space_charge2D_push(ref,cm,current,slice_ds); } From 4e7c288d43ccf5ea02fda820a1e5975d9313e265 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Fri, 14 Feb 2025 15:41:13 -0800 Subject: [PATCH 26/47] Make Python current arg optional. --- src/python/ImpactX.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index 2c8782948..51f19048a 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -442,7 +442,7 @@ void init_ImpactX (py::module& m) ix.amr_data->track_envelope.m_ref = ref; ix.amr_data->track_envelope.m_env = initialization::create_envelope(current,distr); }, - py::arg("ref"), py::arg("distr"), py::arg("current"), + py::arg("ref"), py::arg("distr"), py::arg("current") = 0.0, "Envelope tracking mode:" "Create a 6x6 covariance matrix from a distribution and then initialize " "the the simulation for envelope tracking relative to a reference particle." From 8fe40cfdd77f7606a0bf41dd4e3f1f5c21e7d598 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 14 Feb 2025 23:44:22 +0000 Subject: [PATCH 27/47] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/particles/CovarianceMatrix.H | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/particles/CovarianceMatrix.H b/src/particles/CovarianceMatrix.H index 6562edd54..9345ab4fc 100644 --- a/src/particles/CovarianceMatrix.H +++ b/src/particles/CovarianceMatrix.H @@ -40,16 +40,16 @@ namespace impactx set_beam_current_A (amrex::ParticleReal const beam_current_A) { using namespace amrex::literals; - + beam_current = beam_current_A; - + return *this; } /** Get envelope beam current for 2D space charge * * @returns beam current in A - */ + */ amrex::ParticleReal beam_current_A () { @@ -66,7 +66,7 @@ namespace impactx set_covariance_matrix (CovarianceMatrix const covariance_matrix) { cm = covariance_matrix; - + return *this; } From 47a87a799c00ba4dd1068cf6492dc1062a92a085 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Fri, 14 Feb 2025 16:22:43 -0800 Subject: [PATCH 28/47] Fix Doxygen in CovarianceMatrix.H --- src/particles/CovarianceMatrix.H | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/particles/CovarianceMatrix.H b/src/particles/CovarianceMatrix.H index 9345ab4fc..e918d5347 100644 --- a/src/particles/CovarianceMatrix.H +++ b/src/particles/CovarianceMatrix.H @@ -28,9 +28,9 @@ namespace impactx */ struct Envelope { - amrex::ParticleReal bunch_charge = 0.0; /// total charge in C (for 3D space charge) - amrex::ParticleReal beam_current = 0.0; /// current in A (for 2D space charge) - CovarianceMatrix cm; /// the 6x6 beam covariance matrix + amrex::ParticleReal bunch_charge = 0.0; ///< total charge in C (for 3D space charge) + amrex::ParticleReal beam_current = 0.0; ///< current in A (for 2D space charge) + CovarianceMatrix cm; ///< the 6x6 beam covariance matrix /** Set envelope beam current for 2D space charge * From e2b903e33835bd547375b276064f3668db2db504 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Fri, 14 Feb 2025 17:12:13 -0800 Subject: [PATCH 29/47] Pass covariance matrix by reference. --- src/particles/CovarianceMatrix.H | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/particles/CovarianceMatrix.H b/src/particles/CovarianceMatrix.H index e918d5347..7c8b67c95 100644 --- a/src/particles/CovarianceMatrix.H +++ b/src/particles/CovarianceMatrix.H @@ -63,7 +63,7 @@ namespace impactx * @param Map6x6 beam 6x6 covariance matrix */ Envelope & - set_covariance_matrix (CovarianceMatrix const covariance_matrix) + set_covariance_matrix (const CovarianceMatrix& covariance_matrix) { cm = covariance_matrix; From c4812503115e3496081733df4d304ef43fb37859 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Fri, 14 Feb 2025 19:30:38 -0800 Subject: [PATCH 30/47] Fix Doxygen error. --- src/particles/CovarianceMatrix.H | 2 +- src/python/distribution.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/particles/CovarianceMatrix.H b/src/particles/CovarianceMatrix.H index 7c8b67c95..21f96f227 100644 --- a/src/particles/CovarianceMatrix.H +++ b/src/particles/CovarianceMatrix.H @@ -34,7 +34,7 @@ namespace impactx /** Set envelope beam current for 2D space charge * - * @param beam_current beam current (A) + * @param beam_current_A beam current (A) */ Envelope & set_beam_current_A (amrex::ParticleReal const beam_current_A) diff --git a/src/python/distribution.cpp b/src/python/distribution.cpp index fe484547a..198e45fe5 100644 --- a/src/python/distribution.cpp +++ b/src/python/distribution.cpp @@ -131,5 +131,5 @@ void init_distribution(py::module& m) "A 6D Waterbag distribution" ); - m.def("create_covariance_matrix", &initialization::create_envelope); + m.def("create_envelope", &initialization::create_envelope); } From 11fac545a91f60a047deea4a6f4a7e6a56cb9513 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Fri, 14 Feb 2025 19:36:53 -0800 Subject: [PATCH 31/47] Another Doxygen error. --- src/particles/CovarianceMatrix.H | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/particles/CovarianceMatrix.H b/src/particles/CovarianceMatrix.H index 21f96f227..471b4a348 100644 --- a/src/particles/CovarianceMatrix.H +++ b/src/particles/CovarianceMatrix.H @@ -60,7 +60,7 @@ namespace impactx /** Set 6x6 covariance matrix for envelope tracking * - * @param Map6x6 beam 6x6 covariance matrix + * @param CovarianceMatrix beam 6x6 covariance matrix */ Envelope & set_covariance_matrix (const CovarianceMatrix& covariance_matrix) From 8bb3f279dcb0f303c5bc185d57ca2d7a7a1c0abf Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Fri, 14 Feb 2025 19:45:22 -0800 Subject: [PATCH 32/47] Try to fix Doxygen again. --- src/particles/CovarianceMatrix.H | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/particles/CovarianceMatrix.H b/src/particles/CovarianceMatrix.H index 471b4a348..0321bb573 100644 --- a/src/particles/CovarianceMatrix.H +++ b/src/particles/CovarianceMatrix.H @@ -60,7 +60,7 @@ namespace impactx /** Set 6x6 covariance matrix for envelope tracking * - * @param CovarianceMatrix beam 6x6 covariance matrix + * @param covariance_matrix beam 6x6 covariance matrix */ Envelope & set_covariance_matrix (const CovarianceMatrix& covariance_matrix) From 810bdb7c5373e64a099195c6dbf8038590029bd6 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Sun, 16 Feb 2025 11:19:48 -0800 Subject: [PATCH 33/47] Add space_charge_model input. --- .../input_fodo_envelope_sc.in | 1 + .../fodo_space_charge/run_fodo_envelope_sc.py | 2 +- src/initialization/InitDistribution.cpp | 21 ++++++++++++------- src/python/ImpactX.cpp | 14 +++++++++++++ 4 files changed, 29 insertions(+), 9 deletions(-) diff --git a/examples/fodo_space_charge/input_fodo_envelope_sc.in b/examples/fodo_space_charge/input_fodo_envelope_sc.in index 8148962fc..0b471b9d8 100644 --- a/examples/fodo_space_charge/input_fodo_envelope_sc.in +++ b/examples/fodo_space_charge/input_fodo_envelope_sc.in @@ -44,6 +44,7 @@ quad2.k = -quad1.k algo.particle_shape = 2 algo.track = "envelope" algo.space_charge = true +algo.space_charge_model = 2D ############################################################################### # Diagnostics diff --git a/examples/fodo_space_charge/run_fodo_envelope_sc.py b/examples/fodo_space_charge/run_fodo_envelope_sc.py index d711cfbef..fa2a58ba6 100755 --- a/examples/fodo_space_charge/run_fodo_envelope_sc.py +++ b/examples/fodo_space_charge/run_fodo_envelope_sc.py @@ -13,7 +13,7 @@ # set numerical parameters and IO control sim.particle_shape = 2 # B-spline order sim.space_charge = True -# sim.diagnostics = False # benchmarking +sim.space_charge_model = "2D" sim.slice_step_diagnostics = True # domain decomposition & space charge mesh diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index af3b321d3..428a12145 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -476,17 +476,22 @@ namespace impactx } else if (track == "envelope") { - // relevant for envelope tracking with 3D space charge - amrex::ParticleReal bunch_charge = 0.0; // Bunch charge (C) - pp_dist.query("charge", bunch_charge); - - // relevant for envelope tracking with 2D space charge - amrex::ParticleReal beam_current = 0.0; // Beam current (A) - pp_dist.query("current", beam_current); + std::string space_charge_model = "2D"; // 2D space charge as existing default + pp_algo.query("space_charge_model", space_charge_model); + amrex::ParticleReal intensity = 0.0; // bunch charge (C) for 3D model, beam current (A) for 2D model + + if (space_charge_model == "3D") { + pp_dist.query("charge", intensity); + throw std::runtime_error("3D space charge model not yet implemented in envelope mode."); + } else if (space_charge_model == "2D") { + pp_dist.query("current", intensity); + } else { + throw std::runtime_error("Unknown space_charge_model (use '2D' or '3D') "); + } amr_data->track_envelope.m_ref = initialization::read_reference_particle(pp_dist); auto dist = initialization::read_distribution(pp_dist); - amr_data->track_envelope.m_env = impactx::initialization::create_envelope(beam_current,dist); + amr_data->track_envelope.m_env = impactx::initialization::create_envelope(intensity,dist); } else if (track == "reference_orbit") { diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index 51f19048a..12c6306bf 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -236,6 +236,20 @@ void init_ImpactX (py::module& m) }, "Enable or disable space charge calculations (default: enabled)." ) + .def_property("space_charge_model", + [](ImpactX & /* ix */) { + return detail::get_or_throw("algo", "space_charge_model"); + }, + [](ImpactX & /* ix */, std::string const space_charge_model) { + if (space_charge_model != "2D" && space_charge_model != "3D") { + throw std::runtime_error("Space charge model must be 2D or 3D but is: " + space_charge_model); + } + + amrex::ParmParse pp_algo("algo"); + pp_algo.add("space_charge_model", space_charge_model); + }, + "The model to be used when calculating space charge effects. Either 2D or 3D." + ) .def_property("poisson_solver", [](ImpactX & /* ix */) { return detail::get_or_throw("algo", "poisson_solver"); From 0df5bc42ab6986a61b0cee9832b5226431e72ba8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 16 Feb 2025 19:20:08 +0000 Subject: [PATCH 34/47] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/python/ImpactX.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index 12c6306bf..c6d7445a1 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -244,7 +244,7 @@ void init_ImpactX (py::module& m) if (space_charge_model != "2D" && space_charge_model != "3D") { throw std::runtime_error("Space charge model must be 2D or 3D but is: " + space_charge_model); } - + amrex::ParmParse pp_algo("algo"); pp_algo.add("space_charge_model", space_charge_model); }, From eee70e9cb7333f4164c7192d840884c20ab9c694 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Tue, 18 Feb 2025 12:31:58 -0800 Subject: [PATCH 35/47] Move space charge envelope files. --- src/CMakeLists.txt | 1 + src/elements/Aperture.H | 1 - src/envelope/CMakeLists.txt | 5 +++++ src/envelope/spacecharge/CMakeLists.txt | 4 ++++ .../spacecharge/EnvelopeSpaceChargePush.H | 0 .../spacecharge/EnvelopeSpaceChargePush.cpp | 0 src/particles/spacecharge/CMakeLists.txt | 1 - src/tracking/envelope.cpp | 2 +- 8 files changed, 11 insertions(+), 3 deletions(-) create mode 100644 src/envelope/CMakeLists.txt create mode 100644 src/envelope/spacecharge/CMakeLists.txt rename src/{particles => envelope}/spacecharge/EnvelopeSpaceChargePush.H (100%) rename src/{particles => envelope}/spacecharge/EnvelopeSpaceChargePush.cpp (100%) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 3bccb50fe..b20bbe268 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -12,6 +12,7 @@ add_subdirectory(diagnostics) add_subdirectory(elements) add_subdirectory(initialization) add_subdirectory(particles) +add_subdirectory(envelope) add_subdirectory(tracking) diff --git a/src/elements/Aperture.H b/src/elements/Aperture.H index a305a0eee..a94bfdfc1 100644 --- a/src/elements/Aperture.H +++ b/src/elements/Aperture.H @@ -12,7 +12,6 @@ #include "particles/ImpactXParticleContainer.H" -#include "particles/CovarianceMatrix.H" #include "mixin/alignment.H" #include "mixin/beamoptic.H" #include "mixin/thin.H" diff --git a/src/envelope/CMakeLists.txt b/src/envelope/CMakeLists.txt new file mode 100644 index 000000000..58b6ea661 --- /dev/null +++ b/src/envelope/CMakeLists.txt @@ -0,0 +1,5 @@ +target_sources(lib + PRIVATE +) + +add_subdirectory(spacecharge) diff --git a/src/envelope/spacecharge/CMakeLists.txt b/src/envelope/spacecharge/CMakeLists.txt new file mode 100644 index 000000000..dc9cf8000 --- /dev/null +++ b/src/envelope/spacecharge/CMakeLists.txt @@ -0,0 +1,4 @@ +target_sources(lib + PRIVATE + EnvelopeSpaceChargePush.cpp +) diff --git a/src/particles/spacecharge/EnvelopeSpaceChargePush.H b/src/envelope/spacecharge/EnvelopeSpaceChargePush.H similarity index 100% rename from src/particles/spacecharge/EnvelopeSpaceChargePush.H rename to src/envelope/spacecharge/EnvelopeSpaceChargePush.H diff --git a/src/particles/spacecharge/EnvelopeSpaceChargePush.cpp b/src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp similarity index 100% rename from src/particles/spacecharge/EnvelopeSpaceChargePush.cpp rename to src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp diff --git a/src/particles/spacecharge/CMakeLists.txt b/src/particles/spacecharge/CMakeLists.txt index 3daaaa5af..38ad2ea66 100644 --- a/src/particles/spacecharge/CMakeLists.txt +++ b/src/particles/spacecharge/CMakeLists.txt @@ -4,5 +4,4 @@ target_sources(lib GatherAndPush.cpp HandleSpacecharge.cpp PoissonSolve.cpp - EnvelopeSpaceChargePush.cpp ) diff --git a/src/tracking/envelope.cpp b/src/tracking/envelope.cpp index 110d6b89d..7bee8f685 100644 --- a/src/tracking/envelope.cpp +++ b/src/tracking/envelope.cpp @@ -11,7 +11,7 @@ #include "initialization/InitAmrCore.H" #include "particles/ImpactXParticleContainer.H" #include "particles/Push.H" -#include "particles/spacecharge/EnvelopeSpaceChargePush.H" +#include "envelope/spacecharge/EnvelopeSpaceChargePush.H" #include "diagnostics/DiagnosticOutput.H" #include From 9fc5601a1b155bc108d5d7f439bc84a99552bab7 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Tue, 18 Feb 2025 14:50:04 -0800 Subject: [PATCH 36/47] Add basic documentation. --- docs/source/usage/parameters.rst | 17 +++++++++++++++++ docs/source/usage/python.rst | 3 ++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index f2399550b..31f0bcb2a 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -42,6 +42,9 @@ Initial Beam Distributions * ``beam.charge`` (``float``, in C) bunch charge +* ``beam.current`` (``float``, in A) + beam current, used only if space_charge_model = "2D" + * ``beam.particle`` (``string``) particle type: currently either ``electron``, ``positron`` or ``proton`` @@ -824,6 +827,20 @@ Multigrid-specific numerical options: Currently MLMG solver looks for verbosity levels from 0-5. A higher number results in more verbose output. +* ``algo.space_charge_model`` (``string``, optional, default: ``3D``) + + The physical model of space charge used. + + Options: + + * ``2D``: Space charge forces are computed in the plane ``(x,y)`` transverse to the reference particle velocity, assuming the beam is long and unbunched. + + Currently, this model is supported only in evelope mode (when algo.track = "envelope"). + + * ``3D``: Space charge forces are computed in three dimensions, assuming the beam is bunched. + + Currently, this model is supported only in particle mode (when algo.track = "particles"). + .. _running-cpp-parameters-collective-csr: diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 5f3932038..40f56f4dd 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -180,13 +180,14 @@ Collective Effects & Overall Simulation Parameters :param distr: distribution function to draw from (object from :py:mod:`impactx.distribution`) :param int npart: number of particles to draw - .. py:method:: init_envelope(ref, distr) + .. py:method:: init_envelope(ref, distr, intensity) Envelope tracking mode: Create a 6x6 covariance matrix from a distribution and then initialize the the simulation for envelope tracking relative to a reference particle. :param ref: the reference particle (object from :py:class:`impactx.RefPart`) :param distr: distribution function (object from :py:mod:`impactx.distribution`) + :param float intensity: the beam intensity, given as bunch charge (C) or beam current (A) .. py:method:: particle_container() From cfc02c96dc1d888d64238b75ff55cb2967d5648d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 18 Feb 2025 22:50:17 +0000 Subject: [PATCH 37/47] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/source/usage/parameters.rst | 8 ++++---- docs/source/usage/python.rst | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 31f0bcb2a..2f3c57cde 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -44,7 +44,7 @@ Initial Beam Distributions * ``beam.current`` (``float``, in A) beam current, used only if space_charge_model = "2D" - + * ``beam.particle`` (``string``) particle type: currently either ``electron``, ``positron`` or ``proton`` @@ -834,13 +834,13 @@ Multigrid-specific numerical options: Options: * ``2D``: Space charge forces are computed in the plane ``(x,y)`` transverse to the reference particle velocity, assuming the beam is long and unbunched. - + Currently, this model is supported only in evelope mode (when algo.track = "envelope"). - + * ``3D``: Space charge forces are computed in three dimensions, assuming the beam is bunched. Currently, this model is supported only in particle mode (when algo.track = "particles"). - + .. _running-cpp-parameters-collective-csr: diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 40f56f4dd..436947553 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -187,7 +187,7 @@ Collective Effects & Overall Simulation Parameters :param ref: the reference particle (object from :py:class:`impactx.RefPart`) :param distr: distribution function (object from :py:mod:`impactx.distribution`) - :param float intensity: the beam intensity, given as bunch charge (C) or beam current (A) + :param float intensity: the beam intensity, given as bunch charge (C) or beam current (A) .. py:method:: particle_container() From bfa2cdd3646476d31d658abd5ee5d693b42ca435 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Wed, 19 Feb 2025 17:11:32 -0800 Subject: [PATCH 38/47] Update how_to_run table. --- docs/source/usage/how_to_run.rst | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/source/usage/how_to_run.rst b/docs/source/usage/how_to_run.rst index e5ab89196..5cda59106 100644 --- a/docs/source/usage/how_to_run.rst +++ b/docs/source/usage/how_to_run.rst @@ -13,13 +13,13 @@ ImpactX can be run using any of three distinct tracking modes. ImpactX's most p Additionally, ImpactX provides two simplified tracking modes to aid scientists through every step, from beamline inception to operation: tracking of the beam envelope (6x6 covariance matrix) through linearized transport maps, or only tracking of the reference particle orbit. -================== =============== =============== ================== +================== =============== =============== =================== Mode Use Case Generality Collective Effects -================== =============== =============== ================== -Particle Tracking Full Dynamics Most general Supported -Envelope Tracking Rapid Scans Linearized `Soon `__ +================== =============== =============== =================== +Particle Tracking Full Dynamics Most general Supported (3D only) +Envelope Tracking Rapid Scans Linearized Supported (2D only) Reference Tracking Early Design Reference orbit No -================== =============== =============== ================== +================== =============== =============== =================== .. _usage_run-user-interface: From 0a6fd038de929d33f4a21cd075d82792f0e3e4e1 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Thu, 20 Feb 2025 15:23:54 -0800 Subject: [PATCH 39/47] Verbatim --- docs/source/usage/parameters.rst | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 2f3c57cde..5aac9ced0 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -43,7 +43,7 @@ Initial Beam Distributions bunch charge * ``beam.current`` (``float``, in A) - beam current, used only if space_charge_model = "2D" + beam current, used only if ``algo.space_charge_model = "2D"`` * ``beam.particle`` (``string``) particle type: currently either ``electron``, ``positron`` or ``proton`` @@ -835,11 +835,11 @@ Multigrid-specific numerical options: * ``2D``: Space charge forces are computed in the plane ``(x,y)`` transverse to the reference particle velocity, assuming the beam is long and unbunched. - Currently, this model is supported only in evelope mode (when algo.track = "envelope"). + Currently, this model is supported only in envelope mode (when ``algo.track = "envelope"``). * ``3D``: Space charge forces are computed in three dimensions, assuming the beam is bunched. - Currently, this model is supported only in particle mode (when algo.track = "particles"). + Currently, this model is supported only in particle mode (when ``algo.track = "particles"``). .. _running-cpp-parameters-collective-csr: From 4692d4e46e9800beafd2ac4c43b586c55604dbe1 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Thu, 20 Feb 2025 15:31:53 -0800 Subject: [PATCH 40/47] CMake: Rename Examples --- examples/CMakeLists.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 8d232f39b..0c711c446 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -1273,13 +1273,13 @@ add_impactx_test(examples-scraping.py # FODO cell with 2D space charge using envelope tracking ###################### # # with space charge -add_impactx_test(fodo-envelope-sc +add_impactx_test(FODO.envelope.sc examples/fodo_space_charge/input_fodo_envelope_sc.in ON # ImpactX MPI-parallel examples/fodo_space_charge/analysis_fodo_envelope_sc.py OFF # no plot script yet ) -add_impactx_test(fodo-envelope-sc.py +add_impactx_test(FODO.envelope.sc.py examples/fodo_space_charge/run_fodo_envelope_sc.py OFF # ImpactX MPI-parallel examples/fodo_space_charge/analysis_fodo_envelope_sc.py From cc17e5766d5ec450ce8214c436fea88f22edce01 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Thu, 20 Feb 2025 15:32:11 -0800 Subject: [PATCH 41/47] CMake: Remove Empty Command --- src/envelope/CMakeLists.txt | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/envelope/CMakeLists.txt b/src/envelope/CMakeLists.txt index 58b6ea661..cdffe61e8 100644 --- a/src/envelope/CMakeLists.txt +++ b/src/envelope/CMakeLists.txt @@ -1,5 +1 @@ -target_sources(lib - PRIVATE -) - add_subdirectory(spacecharge) From b17c29ac912f4abe8699621acabbf15639ea277d Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Thu, 20 Feb 2025 15:32:32 -0800 Subject: [PATCH 42/47] Readme Formatting & Link --- examples/fodo_space_charge/README.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/fodo_space_charge/README.rst b/examples/fodo_space_charge/README.rst index 3ccdb05e6..627b16a67 100644 --- a/examples/fodo_space_charge/README.rst +++ b/examples/fodo_space_charge/README.rst @@ -1,12 +1,12 @@ .. _examples-fodo-envelope-sc: FODO Cell with 2D Space Charge using Envelope Tracking -======================================================= +====================================================== This example illustrates a 0.5 A proton beam with a kinetic energy of 6.7 MeV in a FODO cell, with 2D space charge included. The parameters are those described in: -R.D. Ryne et al, "A Test Suite of Space-Charge Problems for Code Benchmarking," +R.D. Ryne et al, `"A Test Suite of Space-Charge Problems for Code Benchmarking," `__ in Proc. EPAC 2004, Lucerne, Switzerland: KV Beam in a FODO Channel The purpose of this example is to illustrate the use of envelope tracking mode with 2D space charge. From b922ed413e2fa5deca2db89567e19929900f073f Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Thu, 20 Feb 2025 15:55:11 -0800 Subject: [PATCH 43/47] Apply Changes from Review, Fix Tests --- docs/source/usage/python.rst | 4 +-- .../spacecharge/EnvelopeSpaceChargePush.H | 9 +++--- .../spacecharge/EnvelopeSpaceChargePush.cpp | 31 +++++++++---------- src/initialization/InitDistribution.H | 8 +++-- src/initialization/InitDistribution.cpp | 8 ++--- src/particles/CovarianceMatrix.H | 27 +++++++--------- src/python/ImpactX.cpp | 7 +++-- src/tracking/envelope.cpp | 6 ++-- 8 files changed, 51 insertions(+), 49 deletions(-) diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 436947553..f27f20753 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -180,14 +180,14 @@ Collective Effects & Overall Simulation Parameters :param distr: distribution function to draw from (object from :py:mod:`impactx.distribution`) :param int npart: number of particles to draw - .. py:method:: init_envelope(ref, distr, intensity) + .. py:method:: init_envelope(ref, distr, intensity=None) Envelope tracking mode: Create a 6x6 covariance matrix from a distribution and then initialize the the simulation for envelope tracking relative to a reference particle. :param ref: the reference particle (object from :py:class:`impactx.RefPart`) :param distr: distribution function (object from :py:mod:`impactx.distribution`) - :param float intensity: the beam intensity, given as bunch charge (C) or beam current (A) + :param float intensity: the beam intensity, given as bunch charge (C) for 3D or beam current (A) for 2D space charge .. py:method:: particle_container() diff --git a/src/envelope/spacecharge/EnvelopeSpaceChargePush.H b/src/envelope/spacecharge/EnvelopeSpaceChargePush.H index c4c9a7644..34ba1883b 100644 --- a/src/envelope/spacecharge/EnvelopeSpaceChargePush.H +++ b/src/envelope/spacecharge/EnvelopeSpaceChargePush.H @@ -13,10 +13,11 @@ #include "particles/ImpactXParticleContainer.H" #include "particles/CovarianceMatrix.H" +#include // for AMREX_RESTRICT #include -namespace impactx::spacecharge +namespace impactx::envelope::spacecharge { /** This function pushes the 6x6 beam covariance matrix for a slice * of length ds, using the linear space charge fields in an rms @@ -29,9 +30,9 @@ namespace impactx::spacecharge * @param[in] ds step size [m] */ void - envelope_space_charge2D_push ( - RefPart const & refpart, - Map6x6 & cm, + space_charge2D_push ( + RefPart const & AMREX_RESTRICT refpart, + Map6x6 & AMREX_RESTRICT cm, amrex::ParticleReal current, amrex::ParticleReal ds ); diff --git a/src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp b/src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp index b090932c5..f60932675 100644 --- a/src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp +++ b/src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp @@ -9,20 +9,19 @@ */ #include "EnvelopeSpaceChargePush.H" -#include #include // for Real #include -#include #include #include -namespace impactx::spacecharge + +namespace impactx::envelope::spacecharge { - void - envelope_space_charge2D_push ( - [[maybe_unused]] RefPart const & refpart, - Map6x6 & cm, + void + space_charge2D_push ( + RefPart const & AMREX_RESTRICT refpart, + Map6x6 & AMREX_RESTRICT cm, amrex::ParticleReal current, amrex::ParticleReal ds ) @@ -33,25 +32,25 @@ namespace impactx::spacecharge Map6x6 R = Map6x6::Identity(); // physical constants and reference quantities - amrex::ParticleReal const c = ablastr::constant::SI::c; - amrex::ParticleReal const ep0 = ablastr::constant::SI::ep0; - amrex::ParticleReal const pi = ablastr::constant::math::pi; + using ablastr::constant::SI::c; + using ablastr::constant::SI::ep0; + using ablastr::constant::math::pi; amrex::ParticleReal const mass = refpart.mass; amrex::ParticleReal const charge = refpart.charge; amrex::ParticleReal const pt_ref = refpart.pt; - amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1.0_prt; + amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1_prt; amrex::ParticleReal const betgam = std::sqrt(betgam2); amrex::ParticleReal const betgam3 = std::pow(betgam,3); // evaluate the beam space charge perveance from current - amrex::ParticleReal const IA = 4.0_prt*pi*ep0*mass*pow(c,3)/charge; - amrex::ParticleReal const Kpv = std::abs(current/IA) * 2.0_prt/betgam3; + amrex::ParticleReal const IA = 4_prt * pi * ep0 * mass * std::pow(c,3) / charge; + amrex::ParticleReal const Kpv = std::abs(current / IA) * 2_prt / betgam3; // evaluate the linear transfer map - amrex::ParticleReal const sigma2 = cm(1,1)*cm(3,3)-cm(1,3)*cm(1,3); + amrex::ParticleReal const sigma2 = cm(1,1) * cm(3,3) - cm(1,3) * cm(1,3); amrex::ParticleReal const sigma = std::sqrt(sigma2); - amrex::ParticleReal const D = (sigma + cm(1,1)) * (sigma + cm(3,3)) - cm(1,3)*cm(1,3); - amrex::ParticleReal const coeff = Kpv*ds/(2.0_prt*D); + amrex::ParticleReal const D = (sigma + cm(1,1)) * (sigma + cm(3,3)) - cm(1,3) * cm(1,3); + amrex::ParticleReal const coeff = Kpv * ds / (2_prt * D); R(2,1) = coeff * (sigma + cm(3,3)); R(2,3) = coeff * (-cm(1,3)); diff --git a/src/initialization/InitDistribution.H b/src/initialization/InitDistribution.H index 6e73a77b3..09528a9f5 100644 --- a/src/initialization/InitDistribution.H +++ b/src/initialization/InitDistribution.H @@ -17,6 +17,7 @@ #include // for AMREX_RESTRICT #include +#include // for std::optional #include // for std::move @@ -39,11 +40,14 @@ namespace impactx::initialization read_distribution (amrex::ParmParse const & pp_dist); /** Ignore the shape of a distribution and use the 2nd moments to create a covariance matrix + * + * @param distr the distribution to draw from + * @param intensity the beam charge (in C, 3D space charge) or current (in A, 2D space charge) */ Envelope create_envelope ( - amrex::ParticleReal const current, - distribution::KnownDistributions const & distr + distribution::KnownDistributions const & distr, + std::optional intensity ); /** Initialize a single particle's data using the given distribution diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 428a12145..5b470f6de 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -187,8 +187,8 @@ namespace impactx Envelope initialization::create_envelope ( - amrex::ParticleReal const current, - distribution::KnownDistributions const & distr + distribution::KnownDistributions const & distr, + std::optional intensity ) { // zero out the 6x6 matrix @@ -235,7 +235,7 @@ namespace impactx }, distr); Envelope env; - env.set_beam_current_A(current); + if (intensity) { env.set_beam_intensity(intensity.value()); } env.set_covariance_matrix(cv); return env; @@ -491,7 +491,7 @@ namespace impactx amr_data->track_envelope.m_ref = initialization::read_reference_particle(pp_dist); auto dist = initialization::read_distribution(pp_dist); - amr_data->track_envelope.m_env = impactx::initialization::create_envelope(intensity,dist); + amr_data->track_envelope.m_env = impactx::initialization::create_envelope(dist, intensity); } else if (track == "reference_orbit") { diff --git a/src/particles/CovarianceMatrix.H b/src/particles/CovarianceMatrix.H index 0321bb573..e363158b2 100644 --- a/src/particles/CovarianceMatrix.H +++ b/src/particles/CovarianceMatrix.H @@ -28,34 +28,31 @@ namespace impactx */ struct Envelope { - amrex::ParticleReal bunch_charge = 0.0; ///< total charge in C (for 3D space charge) - amrex::ParticleReal beam_current = 0.0; ///< current in A (for 2D space charge) - CovarianceMatrix cm; ///< the 6x6 beam covariance matrix + CovarianceMatrix m_env; ///< the 6x6 beam covariance matrix + amrex::ParticleReal m_beam_intensity = 0.0; ///< optional: charge in A (for 3D space charge) or current in A (for 2D space charge) - /** Set envelope beam current for 2D space charge + /** Set envelope beam charge/current for 3D/2D space charge * - * @param beam_current_A beam current (A) + * @param intensity beam charge (C) in 3D or beam current (A) in 2D */ Envelope & - set_beam_current_A (amrex::ParticleReal const beam_current_A) + set_beam_intensity (amrex::ParticleReal const intensity) { - using namespace amrex::literals; - - beam_current = beam_current_A; + m_beam_intensity = intensity; return *this; } - /** Get envelope beam current for 2D space charge + /** Get envelope beam charge/current for 3D/2D space charge * - * @returns beam current in A + * @returns 3D: beam charge in C; 2D: beam current in A */ amrex::ParticleReal - beam_current_A () + beam_intensity () const { using namespace amrex::literals; - return beam_current; + return m_beam_intensity; } /** Set 6x6 covariance matrix for envelope tracking @@ -63,9 +60,9 @@ namespace impactx * @param covariance_matrix beam 6x6 covariance matrix */ Envelope & - set_covariance_matrix (const CovarianceMatrix& covariance_matrix) + set_covariance_matrix (CovarianceMatrix const & covariance_matrix) { - cm = covariance_matrix; + m_env = covariance_matrix; return *this; } diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index c6d7445a1..80d53acc4 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -16,6 +16,7 @@ #if defined(AMREX_DEBUG) || defined(DEBUG) # include #endif +#include #include @@ -452,11 +453,11 @@ void init_ImpactX (py::module& m) .def("init_beam_distribution_from_inputs", &ImpactX::initBeamDistributionFromInputs) .def("init_lattice_elements_from_inputs", &ImpactX::initLatticeElementsFromInputs) .def("init_envelope", - [](ImpactX & ix, RefPart ref, distribution::KnownDistributions distr, amrex::Real current) { + [](ImpactX & ix, RefPart ref, distribution::KnownDistributions distr, std::optional intensity) { ix.amr_data->track_envelope.m_ref = ref; - ix.amr_data->track_envelope.m_env = initialization::create_envelope(current,distr); + ix.amr_data->track_envelope.m_env = initialization::create_envelope(distr, intensity); }, - py::arg("ref"), py::arg("distr"), py::arg("current") = 0.0, + py::arg("ref"), py::arg("distr"), py::arg("intensity") = py::none(), "Envelope tracking mode:" "Create a 6x6 covariance matrix from a distribution and then initialize " "the the simulation for envelope tracking relative to a reference particle." diff --git a/src/tracking/envelope.cpp b/src/tracking/envelope.cpp index 7bee8f685..eb6e19f43 100644 --- a/src/tracking/envelope.cpp +++ b/src/tracking/envelope.cpp @@ -53,8 +53,8 @@ namespace impactx } auto & ref = amr_data->track_envelope.m_ref.value(); auto & env = amr_data->track_envelope.m_env.value(); - auto & cm = env.cm; - auto & current = env.beam_current; + auto & cm = env.m_env; + auto & current = env.m_beam_intensity; // output of init state amrex::ParmParse pp_diag("diag"); @@ -129,7 +129,7 @@ namespace impactx if (space_charge) { // push Covariance Matrix in 2D space charge fields - spacecharge::envelope_space_charge2D_push(ref,cm,current,slice_ds); + envelope::spacecharge::space_charge2D_push(ref,cm,current,slice_ds); } std::visit([&ref, &cm](auto&& element) From 865cbd1312c35b0c3fafebd0747cf81632414953 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Thu, 20 Feb 2025 17:14:01 -0800 Subject: [PATCH 44/47] Update benchmark README. --- examples/fodo_space_charge/README.rst | 40 ++++++--------------------- 1 file changed, 9 insertions(+), 31 deletions(-) diff --git a/examples/fodo_space_charge/README.rst b/examples/fodo_space_charge/README.rst index 627b16a67..0486f4c79 100644 --- a/examples/fodo_space_charge/README.rst +++ b/examples/fodo_space_charge/README.rst @@ -21,8 +21,8 @@ Run This example can be run **either** as: -* **Python** script: ``python3 run_fodo.py`` or -* ImpactX **executable** using an input file: ``impactx input_fodo.in`` +* **Python** script: ``python3 run_fodo_envelope_sc.py`` or +* ImpactX **executable** using an input file: ``impactx input_fodo_envelope_sc.in`` For `MPI-parallel `__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system. @@ -30,15 +30,15 @@ For `MPI-parallel `__ runs, prefix these lines with ` .. tab-item:: Python: Script - .. literalinclude:: run_fodo.py + .. literalinclude:: run_fodo_envelope_sc.py :language: python3 - :caption: You can copy this file from ``examples/fodo/run_fodo.py``. + :caption: You can copy this file from ``examples/fodo_space_charge/run_fodo_envelope_sc.py``. .. tab-item:: Executable: Input File - .. literalinclude:: input_fodo.in + .. literalinclude:: input_fodo_envelope_sc.in :language: ini - :caption: You can copy this file from ``examples/fodo/input_fodo.in``. + :caption: You can copy this file from ``examples/fodo_space_charge/input_fodo_envelope_sc.in``. Analyze @@ -46,30 +46,8 @@ Analyze We run the following script to analyze correctness: -.. dropdown:: Script ``analysis_fodo.py`` +.. dropdown:: Script ``analysis_fodo_envelope_sc.py`` - .. literalinclude:: analysis_fodo.py + .. literalinclude:: analysis_fodo_envelope_sc.py :language: python3 - :caption: You can copy this file from ``examples/fodo/analysis_fodo.py``. - - -Visualize ---------- - -You can run the following script to visualize the beam evolution over time: - -.. dropdown:: Script ``plot_fodo.py`` - - .. literalinclude:: plot_fodo.py - :language: python3 - :caption: You can copy this file from ``examples/fodo/plot_fodo.py``. - -.. figure:: https://user-images.githubusercontent.com/1353258/180287840-8561f6fd-278f-4856-abd8-04fbdb78c8ff.png - :alt: focusing, defocusing and preserved emittance in our FODO cell benchmark. - - FODO transversal beam width and emittance evolution - -.. figure:: https://user-images.githubusercontent.com/1353258/180287845-eb0210a7-2500-4aa9-844c-67fb094329d3.png - :alt: focusing, defocusing and phase space rotation in our FODO cell benchmark. - - FODO transversal beam width and phase space evolution + :caption: You can copy this file from ``examples/fodo_space_charge/analysis_fodo_envelope_sc.py``. From 0d04230f1ddc1e5f1117d01dd1697db84daae423 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Fri, 21 Feb 2025 00:50:06 -0800 Subject: [PATCH 45/47] Python: `Envelope` Class --- src/particles/CovarianceMatrix.H | 12 +++++++++--- src/python/distribution.cpp | 7 +++++++ 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/src/particles/CovarianceMatrix.H b/src/particles/CovarianceMatrix.H index e363158b2..a67a489ad 100644 --- a/src/particles/CovarianceMatrix.H +++ b/src/particles/CovarianceMatrix.H @@ -23,7 +23,7 @@ namespace impactx using CovarianceMatrix = Map6x6; - /** This struct stores the beam envelope attribues, including the 6x6 beam + /** This struct stores the beam envelope attributes, including the 6x6 beam * covariance matrix. Used during envelope tracking mode. */ struct Envelope @@ -50,8 +50,6 @@ namespace impactx amrex::ParticleReal beam_intensity () const { - using namespace amrex::literals; - return m_beam_intensity; } @@ -67,6 +65,14 @@ namespace impactx return *this; } + /** Get the 6x6 covariance matrix for envelope tracking + */ + CovarianceMatrix + covariance_matrix () const + { + return m_env; + } + }; diff --git a/src/python/distribution.cpp b/src/python/distribution.cpp index 198e45fe5..80d754b2c 100644 --- a/src/python/distribution.cpp +++ b/src/python/distribution.cpp @@ -131,5 +131,12 @@ void init_distribution(py::module& m) "A 6D Waterbag distribution" ); + py::class_(m, "Envelope") + .def(py::init<>()) + .def(py::init()) + .def_property("envelope", &Envelope::covariance_matrix, &Envelope::set_covariance_matrix) + .def_property("beam_intensity", &Envelope::beam_intensity, &Envelope::set_beam_intensity) + ; + m.def("create_envelope", &initialization::create_envelope); } From 6e658167c7185be5e2473e9fb4192ec63037eace Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Fri, 21 Feb 2025 00:12:36 -0800 Subject: [PATCH 46/47] Backwards Compatible `algo.space_charge` evolution --- docs/source/usage/parameters.rst | 36 +++++---- docs/source/usage/python.rst | 14 +++- examples/CMakeLists.txt | 2 +- .../input_fodo_envelope_sc.in | 3 +- .../fodo_space_charge/run_fodo_envelope_sc.py | 3 +- src/initialization/Algorithms.H | 29 +++++++ src/initialization/Algorithms.cpp | 75 +++++++++++++++++++ src/initialization/CMakeLists.txt | 1 + src/initialization/InitDistribution.H | 2 +- src/initialization/InitDistribution.cpp | 35 +++++---- src/initialization/InitMeshRefinement.H | 9 ++- src/initialization/InitMeshRefinement.cpp | 7 +- .../spacecharge/HandleSpacecharge.cpp | 7 +- src/python/ImpactX.cpp | 44 ++++++----- src/tracking/envelope.cpp | 15 ++-- src/tracking/particles.cpp | 12 ++- src/tracking/reference.cpp | 16 ++++ 17 files changed, 226 insertions(+), 84 deletions(-) create mode 100644 src/initialization/Algorithms.H create mode 100644 src/initialization/Algorithms.cpp diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 5aac9ced0..3243bd52c 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -43,7 +43,7 @@ Initial Beam Distributions bunch charge * ``beam.current`` (``float``, in A) - beam current, used only if ``algo.space_charge_model = "2D"`` + beam current, used only if ``algo.space_charge = "2D"`` * ``beam.particle`` (``string``) particle type: currently either ``electron``, ``positron`` or ``proton`` @@ -717,12 +717,24 @@ Space Charge Space charge kicks are applied in between slices of thick :ref:`lattice elements `. See there ``nslice`` option on lattice elements for slicing. -* ``algo.space_charge`` (``boolean``, optional, default: ``false``) +* ``algo.space_charge`` (``string``, optional) - Whether to calculate space charge effects. + The physical model of space charge used. + + ImpactX uses an AMReX grid of boxes to organize and parallelize space charge simulation domain. + These boxes also contain a field mesh, if space charge calculations are enabled. + + Options: + + * ``"false"`` (default): space charge effects are not calculated. -ImpactX uses an AMReX grid of boxes to organize and parallelize space charge simulation domain. -These boxes also contain a field mesh, if space charge calculations are enabled. + * ``"2D"``: Space charge forces are computed in the plane ``(x,y)`` transverse to the reference particle velocity, assuming the beam is long and unbunched. + + Currently, this model is supported only in envelope mode (when ``algo.track = "envelope"``). + + * ``"3D"``: Space charge forces are computed in three dimensions, assuming the beam is bunched. + + Currently, this model is supported only in particle mode (when ``algo.track = "particles"``). * ``amr.n_cell`` (3 integers) optional (default: 1 `blocking_factor `__ per MPI process) @@ -827,20 +839,6 @@ Multigrid-specific numerical options: Currently MLMG solver looks for verbosity levels from 0-5. A higher number results in more verbose output. -* ``algo.space_charge_model`` (``string``, optional, default: ``3D``) - - The physical model of space charge used. - - Options: - - * ``2D``: Space charge forces are computed in the plane ``(x,y)`` transverse to the reference particle velocity, assuming the beam is long and unbunched. - - Currently, this model is supported only in envelope mode (when ``algo.track = "envelope"``). - - * ``3D``: Space charge forces are computed in three dimensions, assuming the beam is bunched. - - Currently, this model is supported only in particle mode (when ``algo.track = "particles"``). - .. _running-cpp-parameters-collective-csr: diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index f27f20753..57477e48c 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -62,9 +62,19 @@ Collective Effects & Overall Simulation Parameters .. py:property:: space_charge - Enable (``True``) or disable (``False``) space charge calculations (default: ``False``). + The physical model of space charge used. + + Options: + + * ``False`` (default): space charge effects are not calculated. + + * ``"2D"``: Space charge forces are computed in the plane ``(x,y)`` transverse to the reference particle velocity, assuming the beam is long and unbunched. + + Currently, this model is supported only in envelope mode (when ``algo.track = "envelope"``). + + * ``"3D"``: Space charge forces are computed in three dimensions, assuming the beam is bunched. - Whether to calculate space charge effects. + Currently, this model is supported only in particle mode (when ``algo.track = "particles"``). .. py:property:: poisson_solver diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 0c711c446..bb6b74385 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -106,7 +106,7 @@ function(add_impactx_test name input is_mpi analysis_script plot_script) amrex.throw_exception = 1 amrex.signal_handling = 0 impactx.always_warn_immediately=1 - impactx.abort_on_warning_threshold=low + #impactx.abort_on_warning_threshold=low ${INPUTS_ARGS} WORKING_DIRECTORY ${THIS_WORKING_DIR} ) diff --git a/examples/fodo_space_charge/input_fodo_envelope_sc.in b/examples/fodo_space_charge/input_fodo_envelope_sc.in index 0b471b9d8..0ee52f44f 100644 --- a/examples/fodo_space_charge/input_fodo_envelope_sc.in +++ b/examples/fodo_space_charge/input_fodo_envelope_sc.in @@ -43,8 +43,7 @@ quad2.k = -quad1.k ############################################################################### algo.particle_shape = 2 algo.track = "envelope" -algo.space_charge = true -algo.space_charge_model = 2D +algo.space_charge = 2D ############################################################################### # Diagnostics diff --git a/examples/fodo_space_charge/run_fodo_envelope_sc.py b/examples/fodo_space_charge/run_fodo_envelope_sc.py index fa2a58ba6..929a24e43 100755 --- a/examples/fodo_space_charge/run_fodo_envelope_sc.py +++ b/examples/fodo_space_charge/run_fodo_envelope_sc.py @@ -12,8 +12,7 @@ # set numerical parameters and IO control sim.particle_shape = 2 # B-spline order -sim.space_charge = True -sim.space_charge_model = "2D" +sim.space_charge = "2D" sim.slice_step_diagnostics = True # domain decomposition & space charge mesh diff --git a/src/initialization/Algorithms.H b/src/initialization/Algorithms.H new file mode 100644 index 000000000..27b34dc0a --- /dev/null +++ b/src/initialization/Algorithms.H @@ -0,0 +1,29 @@ +/* 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_ALGORITHMS_H +#define IMPACTX_ALGORITHMS_H + +#include + + +namespace impactx +{ + AMREX_ENUM(SpaceChargeAlgo, + False, + True_3D, + True_2D + ); + + SpaceChargeAlgo + get_space_charge_algo (); + +} // namespace impactx + +#endif // IMPACTX_ALGORITHMS_H diff --git a/src/initialization/Algorithms.cpp b/src/initialization/Algorithms.cpp new file mode 100644 index 000000000..6f8d97c5f --- /dev/null +++ b/src/initialization/Algorithms.cpp @@ -0,0 +1,75 @@ +/* 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 "initialization/Algorithms.H" + +#include + +#include + +#include // for std::transform +#include +#include + + +namespace impactx +{ + SpaceChargeAlgo + get_space_charge_algo () + { + amrex::ParmParse const pp_algo("algo"); + std::string space_charge; + bool has_space_charge = pp_algo.query("space_charge", space_charge); + + if (!has_space_charge) { return SpaceChargeAlgo::False; } + + // TODO: at some point, remove backwards compatibility to pre 25.03 + std::string space_charge_lower = space_charge; + std::transform(space_charge.begin(), space_charge.end(), space_charge_lower.begin(), + [](unsigned char c){ return std::tolower(c); }); + if (space_charge_lower == "1" || space_charge_lower == "true" || space_charge_lower == "on") + { + ablastr::warn_manager::WMRecordWarning( + "algo.space_charge", + "The option algo.space_charge = true is deprecated and will be removed in a future version of ImpactX. " + "Please use algo.space_charge = 3D instead.", + ablastr::warn_manager::WarnPriority::high + ); + return SpaceChargeAlgo::True_3D; + } + if (space_charge_lower == "0") + { + ablastr::warn_manager::WMRecordWarning( + "algo.space_charge", + "The option algo.space_charge = 0 is deprecated and will be removed in a future version of ImpactX. " + "Please use algo.space_charge = false instead.", + ablastr::warn_manager::WarnPriority::high + ); + return SpaceChargeAlgo::False; + } + + if (space_charge == "false" || space_charge == "off") + { + return SpaceChargeAlgo::False; + } + else if (space_charge == "3D") + { + return SpaceChargeAlgo::True_3D; + } + else if (space_charge == "2D") + { + return SpaceChargeAlgo::True_2D; + } + else + { + throw std::runtime_error("algo.space_charge = " + space_charge + " is not a valid option"); + } + } + +} // namespace impactx diff --git a/src/initialization/CMakeLists.txt b/src/initialization/CMakeLists.txt index b0a828b52..745af67b3 100644 --- a/src/initialization/CMakeLists.txt +++ b/src/initialization/CMakeLists.txt @@ -1,5 +1,6 @@ target_sources(lib PRIVATE + Algorithms.cpp AmrCoreData.cpp InitAMReX.cpp InitAmrCore.cpp diff --git a/src/initialization/InitDistribution.H b/src/initialization/InitDistribution.H index 09528a9f5..f1f615b42 100644 --- a/src/initialization/InitDistribution.H +++ b/src/initialization/InitDistribution.H @@ -47,7 +47,7 @@ namespace impactx::initialization Envelope create_envelope ( distribution::KnownDistributions const & distr, - std::optional intensity + std::optional intensity = std::nullopt ); /** Initialize a single particle's data using the given distribution diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 5b470f6de..9142f0e68 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -10,6 +10,7 @@ #include "initialization/InitDistribution.H" #include "ImpactX.H" +#include "initialization/Algorithms.H" #include "particles/CovarianceMatrix.H" #include "particles/ImpactXParticleContainer.H" #include "particles/distribution/All.H" @@ -318,13 +319,11 @@ namespace impactx ref.qm_ratio_SI(), bunch_charge * rel_part_this_proc); - bool space_charge = false; - amrex::ParmParse pp_algo("algo"); - pp_algo.queryAdd("space_charge", space_charge); + auto space_charge = get_space_charge_algo(); // For pure tracking simulations, we keep the particles split equally // on all MPI ranks, and ignore spatial "RealBox" extents of grids. - if (space_charge) { + if (space_charge != SpaceChargeAlgo::False) { // Resize the mesh to fit the spatial extent of the beam and then // redistribute particles, so they reside on the MPI rank that is // responsible for the respective spatial particle position. @@ -476,22 +475,26 @@ namespace impactx } else if (track == "envelope") { - std::string space_charge_model = "2D"; // 2D space charge as existing default - pp_algo.query("space_charge_model", space_charge_model); + amr_data->track_envelope.m_ref = initialization::read_reference_particle(pp_dist); + auto dist = initialization::read_distribution(pp_dist); + + amrex::ParticleReal intensity = 0.0; // bunch charge (C) for 3D model, beam current (A) for 2D model - if (space_charge_model == "3D") { - pp_dist.query("charge", intensity); + auto space_charge = get_space_charge_algo(); + if (space_charge == SpaceChargeAlgo::True_3D) + { + //pp_dist.get("charge", intensity); + //amr_data->track_envelope.m_env = impactx::initialization::create_envelope(dist, intensity); throw std::runtime_error("3D space charge model not yet implemented in envelope mode."); - } else if (space_charge_model == "2D") { - pp_dist.query("current", intensity); - } else { - throw std::runtime_error("Unknown space_charge_model (use '2D' or '3D') "); + } else if (space_charge == SpaceChargeAlgo::True_2D) + { + pp_dist.get("current", intensity); + amr_data->track_envelope.m_env = impactx::initialization::create_envelope(dist, intensity); + } else + { + amr_data->track_envelope.m_env = impactx::initialization::create_envelope(dist); } - - amr_data->track_envelope.m_ref = initialization::read_reference_particle(pp_dist); - auto dist = initialization::read_distribution(pp_dist); - amr_data->track_envelope.m_env = impactx::initialization::create_envelope(dist, intensity); } else if (track == "reference_orbit") { diff --git a/src/initialization/InitMeshRefinement.H b/src/initialization/InitMeshRefinement.H index a8b1be634..3b680cc12 100644 --- a/src/initialization/InitMeshRefinement.H +++ b/src/initialization/InitMeshRefinement.H @@ -9,6 +9,8 @@ */ #pragma once +#include "initialization/Algorithms.H" + #include #include @@ -32,8 +34,7 @@ namespace impactx::initialization amrex::ParmParse pp_amr("amr"); amrex::ParmParse pp_geometry("geometry"); - bool space_charge = false; - pp_algo.queryAdd("space_charge", space_charge); + auto space_charge = get_space_charge_algo(); std::string poisson_solver = "multigrid"; pp_algo.queryAdd("poisson_solver", poisson_solver); @@ -41,7 +42,7 @@ namespace impactx::initialization int max_level = 0; pp_amr.queryWithParser("max_level", max_level); - if (max_level > 1 && !space_charge) + if (max_level > 1 && space_charge != SpaceChargeAlgo::False) throw std::runtime_error( "Mesh-refinement (amr.max_level>=0) is only supported with " "space charge modeling (algo.space_charge=1)."); @@ -51,7 +52,7 @@ namespace impactx::initialization prob_relative[0] = 3.0; // top/bottom pad the beam on the lowest level by default by its width pp_geometry.queryarr("prob_relative", prob_relative); - if (prob_relative[0] < 3.0 && space_charge && poisson_solver == "multigrid") + if (prob_relative[0] < 3.0 && space_charge != SpaceChargeAlgo::False && poisson_solver == "multigrid") ablastr::warn_manager::WMRecordWarning( "ImpactX::read_mr_prob_relative", "Dynamic resizing of the mesh uses a geometry.prob_relative " diff --git a/src/initialization/InitMeshRefinement.cpp b/src/initialization/InitMeshRefinement.cpp index 7fd142710..cfddce123 100644 --- a/src/initialization/InitMeshRefinement.cpp +++ b/src/initialization/InitMeshRefinement.cpp @@ -8,6 +8,7 @@ * License: BSD-3-Clause-LBNL */ #include "ImpactX.H" +#include "initialization/Algorithms.H" #include "initialization/InitAmrCore.H" #include "particles/ImpactXParticleContainer.H" #include "particles/distribution/Waterbag.H" @@ -82,10 +83,8 @@ namespace detail BL_PROFILE("ImpactX::ResizeMesh"); { - amrex::ParmParse pp_algo("algo"); - bool space_charge = false; - pp_algo.query("space_charge", space_charge); - if (!space_charge) + auto space_charge = get_space_charge_algo(); + if (space_charge == SpaceChargeAlgo::False) ablastr::warn_manager::WMRecordWarning( "ImpactX::ResizeMesh", "This is a simulation without space charge. " diff --git a/src/particles/spacecharge/HandleSpacecharge.cpp b/src/particles/spacecharge/HandleSpacecharge.cpp index 4cc13a73a..9612b236f 100644 --- a/src/particles/spacecharge/HandleSpacecharge.cpp +++ b/src/particles/spacecharge/HandleSpacecharge.cpp @@ -9,6 +9,7 @@ */ #include "HandleSpacecharge.H" +#include "initialization/Algorithms.H" #include "initialization/AmrCoreData.H" #include "particles/ImpactXParticleContainer.H" #include "particles/spacecharge/ForceFromSelfFields.H" @@ -32,12 +33,10 @@ namespace impactx::particles::spacecharge { BL_PROFILE("impactx::particles::wakefields::HandleSpacecharge") - amrex::ParmParse const pp_algo("algo"); - bool space_charge = false; - pp_algo.query("space_charge", space_charge); + auto space_charge = get_space_charge_algo(); // turn off if disabled by user - if (!space_charge) { return; } + if (space_charge == SpaceChargeAlgo::False) { return; } // turn off if less than 2 particles if (amr_data->track_particles.m_particle_container->TotalNumberOfParticles(true, false) < 2) { return; } diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index 80d53acc4..8ba1fd0fe 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -18,6 +18,8 @@ #endif #include #include +#include +#include namespace py = pybind11; @@ -228,28 +230,32 @@ void init_ImpactX (py::module& m) "Enable or disable eigenemittance diagnostic calculations (default: disabled)." ) .def_property("space_charge", - [](ImpactX & /* ix */) { - return detail::get_or_throw("algo", "space_charge"); - }, - [](ImpactX & /* ix */, bool const enable) { - amrex::ParmParse pp_algo("algo"); - pp_algo.add("space_charge", enable); - }, - "Enable or disable space charge calculations (default: enabled)." - ) - .def_property("space_charge_model", - [](ImpactX & /* ix */) { - return detail::get_or_throw("algo", "space_charge_model"); + [](ImpactX & /* ix */) -> std::string { + return detail::get_or_throw("algo", "space_charge"); }, - [](ImpactX & /* ix */, std::string const space_charge_model) { - if (space_charge_model != "2D" && space_charge_model != "3D") { - throw std::runtime_error("Space charge model must be 2D or 3D but is: " + space_charge_model); + [](ImpactX & /* ix */, std::variant space_charge_v) { + if (std::holds_alternative(space_charge_v)) { + amrex::ParmParse pp_algo("algo"); + if (std::get(space_charge_v)) { + // TODO: boolean True is deprecated since 25.03, remove some time after + py::print("sim.space_charge = True is deprecated, please use space_charge = \"3D\""); + pp_algo.add("space_charge", std::string("3D")); + } else { + // map boolean False to "false" / off + pp_algo.add("space_charge", std::string("false")); + } + } + else + { + std::string const space_charge = std::get(space_charge_v); + if (space_charge != "false" && space_charge != "off" && space_charge != "2D" && space_charge != "3D") { + throw std::runtime_error("Space charge model must be 2D or 3D but is: " + space_charge); + } + amrex::ParmParse pp_algo("algo"); + pp_algo.add("space_charge", space_charge); } - - amrex::ParmParse pp_algo("algo"); - pp_algo.add("space_charge_model", space_charge_model); }, - "The model to be used when calculating space charge effects. Either 2D or 3D." + "The model to be used when calculating space charge effects. Either off, 2D, or 3D." ) .def_property("poisson_solver", [](ImpactX & /* ix */) { diff --git a/src/tracking/envelope.cpp b/src/tracking/envelope.cpp index eb6e19f43..50f77c9bc 100644 --- a/src/tracking/envelope.cpp +++ b/src/tracking/envelope.cpp @@ -8,6 +8,7 @@ * License: BSD-3-Clause-LBNL */ #include "ImpactX.H" +#include "initialization/Algorithms.H" #include "initialization/InitAmrCore.H" #include "particles/ImpactXParticleContainer.H" #include "particles/Push.H" @@ -78,21 +79,23 @@ namespace impactx } amrex::ParmParse const pp_algo("algo"); - bool space_charge = false; - pp_algo.query("space_charge", space_charge); - //AMREX_ALWAYS_ASSERT_WITH_MESSAGE("2D space charge only is implemented for envelope tracking."); + auto space_charge = get_space_charge_algo(); if (verbose > 0) { - amrex::Print() << " Space Charge effects: " << space_charge << "\n"; + amrex::Print() << " Space Charge effects: " << amrex::getEnumNameString(space_charge) << "\n"; + } + if (space_charge == SpaceChargeAlgo::True_3D) + { + throw std::runtime_error("3D space charge effects are not yet implemented for envelope tracking."); } bool csr = false; pp_algo.query("csr", csr); - AMREX_ALWAYS_ASSERT_WITH_MESSAGE(!csr, "CSR not yet implemented for envelope tracking."); if (verbose > 0) { amrex::Print() << " CSR effects: " << csr << "\n"; } + AMREX_ALWAYS_ASSERT_WITH_MESSAGE(!csr, "CSR effects are not yet implemented for envelope tracking."); // periods through the lattice int num_periods = 1; @@ -126,7 +129,7 @@ namespace impactx << " slice_step=" << slice_step << "\n"; } - if (space_charge) + if (space_charge != SpaceChargeAlgo::False) { // push Covariance Matrix in 2D space charge fields envelope::spacecharge::space_charge2D_push(ref,cm,current,slice_ds); diff --git a/src/tracking/particles.cpp b/src/tracking/particles.cpp index 655944caf..a024b3b33 100644 --- a/src/tracking/particles.cpp +++ b/src/tracking/particles.cpp @@ -8,6 +8,7 @@ * License: BSD-3-Clause-LBNL */ #include "ImpactX.H" +#include "initialization/Algorithms.H" #include "initialization/InitAmrCore.H" #include "particles/CollectLost.H" #include "particles/ImpactXParticleContainer.H" @@ -68,13 +69,16 @@ namespace impactx } - amrex::ParmParse const pp_algo("algo"); - bool space_charge = false; - pp_algo.query("space_charge", space_charge); + auto space_charge = get_space_charge_algo(); if (verbose > 0) { - amrex::Print() << " Space Charge effects: " << space_charge << "\n"; + amrex::Print() << " Space Charge effects: " << amrex::getEnumNameString(space_charge) << "\n"; + } + if (space_charge == SpaceChargeAlgo::True_2D) + { + throw std::runtime_error("2D space charge effects are not yet implemented for particle tracking."); } + amrex::ParmParse const pp_algo("algo"); bool csr = false; pp_algo.query("csr", csr); if (verbose > 0) { diff --git a/src/tracking/reference.cpp b/src/tracking/reference.cpp index 372a12917..5e42651c1 100644 --- a/src/tracking/reference.cpp +++ b/src/tracking/reference.cpp @@ -8,6 +8,7 @@ * License: BSD-3-Clause-LBNL */ #include "ImpactX.H" +#include "initialization/Algorithms.H" #include "initialization/InitAmrCore.H" #include "particles/ImpactXParticleContainer.H" #include "particles/Push.H" @@ -20,6 +21,7 @@ #include #include +#include namespace impactx @@ -60,6 +62,20 @@ namespace impactx } + auto space_charge = get_space_charge_algo(); + if (space_charge != SpaceChargeAlgo::False) + { + throw std::runtime_error("Space charge effects cannot be modeled for single particle tracking."); + } + + amrex::ParmParse const pp_algo("algo"); + bool csr = false; + pp_algo.query("csr", csr); + if (!csr) + { + throw std::runtime_error("CSR effects cannot be modeled for single particle tracking."); + } + // periods through the lattice int num_periods = 1; amrex::ParmParse("lattice").queryAddWithParser("periods", num_periods); From b23c9aff4e9a8d143adda5b20ed5a5a25d68d9a1 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Fri, 21 Feb 2025 11:19:51 -0800 Subject: [PATCH 47/47] Update all Examples --- examples/CMakeLists.txt | 2 +- examples/cfchannel/input_cfchannel_10nC_fft.in | 2 +- examples/cfchannel/input_cfchannel_10nC_mlmg.in | 2 +- examples/cfchannel/run_cfchannel_10nC_fft.py | 2 +- examples/cfchannel/run_cfchannel_10nC_mlmg.py | 2 +- examples/epac2004_benchmarks/input_bithermal.in | 2 +- examples/epac2004_benchmarks/input_fodo_rf_SC.in | 2 +- examples/epac2004_benchmarks/input_thermal.in | 2 +- examples/epac2004_benchmarks/run_bithermal.py | 2 +- examples/epac2004_benchmarks/run_fodo_rf_SC.py | 2 +- examples/epac2004_benchmarks/run_thermal.py | 2 +- examples/expanding_beam/input_expanding_fft.in | 2 +- examples/expanding_beam/input_expanding_mlmg.in | 2 +- examples/expanding_beam/run_expanding_fft.py | 2 +- examples/expanding_beam/run_expanding_mlmg.py | 2 +- examples/kurth/input_kurth_10nC_periodic.in | 2 +- examples/kurth/run_kurth_10nC_periodic.py | 2 +- examples/linac_segment/input_linac_segment.in | 2 +- examples/linac_segment/run_linac_segment.py | 2 +- 19 files changed, 19 insertions(+), 19 deletions(-) diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index bb6b74385..0c711c446 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -106,7 +106,7 @@ function(add_impactx_test name input is_mpi analysis_script plot_script) amrex.throw_exception = 1 amrex.signal_handling = 0 impactx.always_warn_immediately=1 - #impactx.abort_on_warning_threshold=low + impactx.abort_on_warning_threshold=low ${INPUTS_ARGS} WORKING_DIRECTORY ${THIS_WORKING_DIR} ) diff --git a/examples/cfchannel/input_cfchannel_10nC_fft.in b/examples/cfchannel/input_cfchannel_10nC_fft.in index 5b93893b9..1fedfc541 100644 --- a/examples/cfchannel/input_cfchannel_10nC_fft.in +++ b/examples/cfchannel/input_cfchannel_10nC_fft.in @@ -37,7 +37,7 @@ constf1.kt = 1.0 # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D algo.poisson_solver = "fft" amr.n_cell = 48 48 40 diff --git a/examples/cfchannel/input_cfchannel_10nC_mlmg.in b/examples/cfchannel/input_cfchannel_10nC_mlmg.in index 36d39811c..a009a33d8 100644 --- a/examples/cfchannel/input_cfchannel_10nC_mlmg.in +++ b/examples/cfchannel/input_cfchannel_10nC_mlmg.in @@ -37,7 +37,7 @@ constf1.kt = 1.0 # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D amr.n_cell = 48 48 40 #amr.n_cell = 72 72 64 # optional for increased precision diff --git a/examples/cfchannel/run_cfchannel_10nC_fft.py b/examples/cfchannel/run_cfchannel_10nC_fft.py index 6349438b4..19f352d8f 100755 --- a/examples/cfchannel/run_cfchannel_10nC_fft.py +++ b/examples/cfchannel/run_cfchannel_10nC_fft.py @@ -13,7 +13,7 @@ # set numerical parameters and IO control sim.n_cell = [48, 48, 40] # [72, 72, 64] for increased precision sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" sim.poisson_solver = "fft" sim.prob_relative = [1.1] # sim.diagnostics = False # benchmarking diff --git a/examples/cfchannel/run_cfchannel_10nC_mlmg.py b/examples/cfchannel/run_cfchannel_10nC_mlmg.py index 45f057edc..dc2506f65 100755 --- a/examples/cfchannel/run_cfchannel_10nC_mlmg.py +++ b/examples/cfchannel/run_cfchannel_10nC_mlmg.py @@ -13,7 +13,7 @@ # set numerical parameters and IO control sim.n_cell = [48, 48, 40] # [72, 72, 64] for increased precision sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" sim.prob_relative = [3.0] # sim.diagnostics = False # benchmarking sim.slice_step_diagnostics = True diff --git a/examples/epac2004_benchmarks/input_bithermal.in b/examples/epac2004_benchmarks/input_bithermal.in index 928b80467..2345b1415 100644 --- a/examples/epac2004_benchmarks/input_bithermal.in +++ b/examples/epac2004_benchmarks/input_bithermal.in @@ -36,7 +36,7 @@ constf1.nslice = 50 # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D #amr.n_cell = 128 128 128 #full resolution amr.n_cell = 64 64 64 diff --git a/examples/epac2004_benchmarks/input_fodo_rf_SC.in b/examples/epac2004_benchmarks/input_fodo_rf_SC.in index 6e6cf533f..67229db67 100644 --- a/examples/epac2004_benchmarks/input_fodo_rf_SC.in +++ b/examples/epac2004_benchmarks/input_fodo_rf_SC.in @@ -117,7 +117,7 @@ gapb1.sin_coefficients = 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 \ # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D amr.n_cell = 56 56 64 geometry.prob_relative = 4.0 diff --git a/examples/epac2004_benchmarks/input_thermal.in b/examples/epac2004_benchmarks/input_thermal.in index d6ba68b6d..10bc41518 100644 --- a/examples/epac2004_benchmarks/input_thermal.in +++ b/examples/epac2004_benchmarks/input_thermal.in @@ -33,7 +33,7 @@ constf1.nslice = 400 #full resolution # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D #amr.n_cell = 128 128 128 #full resolution amr.n_cell = 64 64 64 diff --git a/examples/epac2004_benchmarks/run_bithermal.py b/examples/epac2004_benchmarks/run_bithermal.py index 008a3a7bd..aba07705e 100755 --- a/examples/epac2004_benchmarks/run_bithermal.py +++ b/examples/epac2004_benchmarks/run_bithermal.py @@ -14,7 +14,7 @@ # sim.n_cell = [128, 128, 128] # full resolution sim.n_cell = [64, 64, 64] sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" sim.dynamic_size = True sim.prob_relative = [3.0] diff --git a/examples/epac2004_benchmarks/run_fodo_rf_SC.py b/examples/epac2004_benchmarks/run_fodo_rf_SC.py index 07c25eedb..8a6f2a84d 100755 --- a/examples/epac2004_benchmarks/run_fodo_rf_SC.py +++ b/examples/epac2004_benchmarks/run_fodo_rf_SC.py @@ -13,7 +13,7 @@ # set numerical parameters and IO control sim.n_cell = [56, 56, 64] sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" sim.dynamic_size = True sim.prob_relative = [4.0] diff --git a/examples/epac2004_benchmarks/run_thermal.py b/examples/epac2004_benchmarks/run_thermal.py index c20dc1b8d..f8fe36503 100755 --- a/examples/epac2004_benchmarks/run_thermal.py +++ b/examples/epac2004_benchmarks/run_thermal.py @@ -13,7 +13,7 @@ # set numerical parameters and IO control sim.n_cell = [56, 56, 64] sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" sim.dynamic_size = True sim.prob_relative = [4.0] diff --git a/examples/expanding_beam/input_expanding_fft.in b/examples/expanding_beam/input_expanding_fft.in index 9ba39c481..45676b5de 100644 --- a/examples/expanding_beam/input_expanding_fft.in +++ b/examples/expanding_beam/input_expanding_fft.in @@ -32,7 +32,7 @@ monitor.backend = h5 # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D algo.poisson_solver = "fft" # Space charge solver with one MR level diff --git a/examples/expanding_beam/input_expanding_mlmg.in b/examples/expanding_beam/input_expanding_mlmg.in index a31c6ed97..cb2693e2b 100644 --- a/examples/expanding_beam/input_expanding_mlmg.in +++ b/examples/expanding_beam/input_expanding_mlmg.in @@ -32,7 +32,7 @@ monitor.backend = h5 # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D # Space charge solver with one MR level amr.max_level = 1 diff --git a/examples/expanding_beam/run_expanding_fft.py b/examples/expanding_beam/run_expanding_fft.py index 92fa1269c..e34ef060f 100755 --- a/examples/expanding_beam/run_expanding_fft.py +++ b/examples/expanding_beam/run_expanding_fft.py @@ -18,7 +18,7 @@ sim.blocking_factor_z = [4] sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" sim.poisson_solver = "fft" sim.dynamic_size = True sim.prob_relative = [1.2, 1.1] diff --git a/examples/expanding_beam/run_expanding_mlmg.py b/examples/expanding_beam/run_expanding_mlmg.py index 6485d95eb..a53246e99 100755 --- a/examples/expanding_beam/run_expanding_mlmg.py +++ b/examples/expanding_beam/run_expanding_mlmg.py @@ -18,7 +18,7 @@ sim.blocking_factor_z = [4] sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" sim.dynamic_size = True sim.prob_relative = [3.0, 1.1] diff --git a/examples/kurth/input_kurth_10nC_periodic.in b/examples/kurth/input_kurth_10nC_periodic.in index 20360cea6..dcac80da4 100644 --- a/examples/kurth/input_kurth_10nC_periodic.in +++ b/examples/kurth/input_kurth_10nC_periodic.in @@ -40,7 +40,7 @@ constf1.kt = 0.7 # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D amr.n_cell = 48 48 40 #amr.n_cell = 72 72 72 # optional for increased precision diff --git a/examples/kurth/run_kurth_10nC_periodic.py b/examples/kurth/run_kurth_10nC_periodic.py index 8c94035e3..6c2a92dc5 100755 --- a/examples/kurth/run_kurth_10nC_periodic.py +++ b/examples/kurth/run_kurth_10nC_periodic.py @@ -13,7 +13,7 @@ # set numerical parameters and IO control sim.n_cell = [48, 48, 40] # use [72, 72, 72] for increased precision sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" # sim.diagnostics = False # benchmarking sim.slice_step_diagnostics = True diff --git a/examples/linac_segment/input_linac_segment.in b/examples/linac_segment/input_linac_segment.in index 3a858ce53..217e48d3a 100644 --- a/examples/linac_segment/input_linac_segment.in +++ b/examples/linac_segment/input_linac_segment.in @@ -359,7 +359,7 @@ RF1b.sin_coefficients = \ # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D # Space charge using IGF solver algo.poisson_solver = "fft" diff --git a/examples/linac_segment/run_linac_segment.py b/examples/linac_segment/run_linac_segment.py index 392280056..ab5d76675 100644 --- a/examples/linac_segment/run_linac_segment.py +++ b/examples/linac_segment/run_linac_segment.py @@ -17,7 +17,7 @@ # sim.n_cell = [64, 64, 64] #use this for high-resolution runs sim.n_cell = [32, 32, 32] sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" sim.poisson_solver = "fft" sim.dynamic_size = True sim.prob_relative = [1.1, 1.1]