BLAST-ImpactX · ax3l · Feb 24, 2025 · Feb 11, 2025 · Feb 11, 2025 · Feb 11, 2025
diff --git 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/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 <https://github.com/ECP-WarpX/impactx/issues/826>`__
+================== =============== =============== ===================
+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:

diff --git 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 ``algo.space_charge = "2D"``
+
 * ``beam.particle`` (``string``)
   particle type: currently either ``electron``, ``positron`` or ``proton``
 
@@ -714,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)
+
+  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.
+
+  * ``"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"``).
 
-  Whether to calculate space charge effects.
+  * ``"3D"``: Space charge forces are computed in three dimensions, assuming the beam is bunched.
 
-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.
+    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)
 

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:
+
+      * ``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
 
@@ -180,13 +190,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=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) for 3D or beam current (A) for 2D space charge
 
    .. py:method:: particle_container()
 

diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -1269,3 +1269,19 @@ 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/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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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/fodo_space_charge/README.rst b/examples/fodo_space_charge/README.rst
@@ -0,0 +1,53 @@
+.. _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," <https://accelconf.web.cern.ch/e04/PAPERS/WEPLT047.PDF>`__
+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_envelope_sc.py`` or
+* ImpactX **executable** using an input file: ``impactx input_fodo_envelope_sc.in``
+
+For `MPI-parallel <https://www.mpi-forum.org>`__ 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_envelope_sc.py
+          :language: python3
+          :caption: You can copy this file from ``examples/fodo_space_charge/run_fodo_envelope_sc.py``.
+
+   .. tab-item:: Executable: Input File
+
+       .. literalinclude:: input_fodo_envelope_sc.in
+          :language: ini
+          :caption: You can copy this file from ``examples/fodo_space_charge/input_fodo_envelope_sc.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_fodo_envelope_sc.py``
+
+   .. literalinclude:: analysis_fodo_envelope_sc.py
+      :language: python3
+      :caption: You can copy this file from ``examples/fodo_space_charge/analysis_fodo_envelope_sc.py``.