From e45cafd815648c39d138fda56e7b3b165ce059bf Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Fri, 21 Feb 2025 00:12:36 -0800 Subject: [PATCH] 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 | 9 +-- src/tracking/particles.cpp | 8 +- 16 files changed, 201 insertions(+), 83 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..5e891272e 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: + + * ``"off"`` (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..c90aa513b 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: + + * ``"off"`` (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..83c41bdeb --- /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, + off, + on_3D, + on_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..cf2546d41 --- /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::off; } + + // 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::on_3D; + } + if (space_charge_lower == "0" || space_charge_lower == "false") + { + 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 = off instead.", + ablastr::warn_manager::WarnPriority::high + ); + return SpaceChargeAlgo::off; + } + + if (space_charge == "off") + { + return SpaceChargeAlgo::off; + } + else if (space_charge == "3D") + { + return SpaceChargeAlgo::on_3D; + } + else if (space_charge == "2D") + { + return SpaceChargeAlgo::on_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..24697ca10 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::off) { // 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::on_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::on_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..f51da5994 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::off) 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::off && 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..3b79a541a 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::off) 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..64b94eb0c 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::off) { 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..3a614b9dd 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)) { + // TODO: bool is deprecated since 25.03, remove some time after + amrex::ParmParse pp_algo("algo"); + if (std::get(space_charge_v)) { + py::print("sim.space_charge = True is deprecated, please use space_charge = \"3D\""); + pp_algo.add("space_charge", std::string("3D")); + } else { + py::print("sim.space_charge = False is deprecated, please use space_charge = \"off\" (or skip for default)"); + pp_algo.add("space_charge", std::string("off")); + } + } + else + { + std::string const space_charge = std::get(space_charge_v); + if (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..632ee9a69 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,12 +79,10 @@ 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"; } bool csr = false; @@ -126,7 +125,7 @@ namespace impactx << " slice_step=" << slice_step << "\n"; } - if (space_charge) + if (space_charge != SpaceChargeAlgo::off) { // 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..2a3a6c921 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,12 @@ 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"; } + amrex::ParmParse const pp_algo("algo"); bool csr = false; pp_algo.query("csr", csr); if (verbose > 0) {