Backwards Compatible algo.space_charge evolution

BLAST-ImpactX · Feb 21, 2025 · e45cafd · e45cafd
1 parent 865cbd1
commit e45cafd
Show file tree

Hide file tree

Showing 16 changed files with 201 additions and 83 deletions.
diff --git 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 <running-cpp-parameters-lattice>`.
 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 <https://amrex-codes.github.io/amrex/docs_html/GridCreation.html>`__ 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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 <AMReX_Enum.H>
+
+
+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
@@ -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 <ablastr/warn_manager/WarnManager.H>
+
+#include <AMReX_ParmParse.H>
+
+#include <algorithm>  // for std::transform
+#include <stdexcept>
+#include <string>
+
+
+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
@@ -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
@@ -47,7 +47,7 @@ namespace impactx::initialization
     Envelope
     create_envelope (
         distribution::KnownDistributions const & distr,
-        std::optional<amrex::ParticleReal> intensity
+        std::optional<amrex::ParticleReal> 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
@@ -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
@@ -9,6 +9,8 @@
  */
 #pragma once
 
+#include "initialization/Algorithms.H"
+
 #include <ablastr/warn_manager/WarnManager.H>
 
 #include <AMReX.H>
@@ -32,16 +34,15 @@ 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);
 
         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
@@ -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
@@ -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; }