2.5D Gaussian Solver

docs/source/usage/parameters.rst

@@ -882,9 +882,23 @@ See there ``nslice`` option on lattice elements for slicing.
   * ``"Gauss3D"`: Calculate 3D space charge forces as if the beam was a Gaussian distribution.
 
     This model is supported only in particle tracking mode (when ``algo.track = "particles"``).
-    Ref.: J. Qiang et al., "Two-and-a-half dimensional symplectic space-charge solver", LBNL Report Number: LBNL-2001674 (2025).
+    Ref.: J. Qiang, "Two-and-a-half dimensional symplectic space-charge solver", LBNL Report Number: LBNL-2001674 (2025).
     (This reference describes both 3D and 2.5D models.)
 
+    This model supports the following sub-option:
+
+    * ``algo.space_charge.gauss_nint`` (``int``, default: ``101``)
+
+      Number of steps for computing the integrals (Eqs. 45-47 in the above paper).
+
+    * ``algo.space_charge.gauss_taylor_delta`` (``float``, default: ``0.01``)
+
+      Initial integral region to avoid divergence of integrand at 0.
+
+    * ``algo.space_charge.gauss_charge_z_bins`` (``int``, default: ``129``)
+
+      Number of bins for longitudinal line density deposition.
+
 * ``amr.n_cell`` (3 integers) optional (default: 1 `blocking_factor <https://amrex-codes.github.io/amrex/docs_html/GridCreation.html>`__ per MPI process)
 
   The number of grid points along each direction (on the **coarsest level**)

docs/source/usage/python.rst

@@ -79,8 +79,21 @@ Collective Effects & Overall Simulation Parameters
       * ``"Gauss3D"`: Calculate 3D space charge forces as if the beam was a Gaussian distribution.
 
         This model is supported only in particle tracking mode (when ``algo.track = "particles"``).
-        Ref.: J. Qiang et al., "Two-and-a-half dimensional symplectic space-charge solver", LBNL Report Number: LBNL-2001674 (2025).
+        Ref.: J. Qiang, "Two-and-a-half dimensional symplectic space-charge solver", LBNL Report Number: LBNL-2001674 (2025).
         (This reference describes both 3D and 2.5D models.)
+
+   .. py:property:: space_charge_gauss_nint
+
+      Number of steps for computing the integrals (default: ``101``).
+
+   .. py:property:: space_charge_gauss_taylor_delta
+
+      Initial integral region to avoid integrand divergence at 0 (default: ``0.01``).
+
+   .. py:property:: space_charge_gauss_charge_z_bins
+
+      Number of bins for longitudinal charge density deposition (default: ``129``).
+
    .. py:property:: poisson_solver
 
       The numerical solver to solve the Poisson equation when calculating space charge effects.

examples/fodo_space_charge/README.rst

@@ -53,7 +53,7 @@ We run the following script to analyze correctness:
       :caption: You can copy this file from ``examples/fodo_space_charge/analysis_fodo_envelope_sc.py``.
 
 
-.. _examples-fodo-envelope-sc-gaussian:
+.. _examples-fodo-Gaussian-sc:
 
 FODO Cell with 3D Gaussian Space Charge Using Particle Tracking
 ===============================================================
@@ -104,3 +104,7 @@ We run the following script to analyze correctness:
    .. literalinclude:: analysis_fodo_Gauss3D_sc.py
       :language: python3
       :caption: You can copy this file from ``examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py``.
+
+.. _examples-fodo-2p5dGaussian-sc:
+
+FODO Cell with 2.5D Gaussian Space Charge Using Particle Tracking

examples/fodo_space_charge/analysis_fodo_Gauss2p5D_sc.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Axel Huebl, Ji Qiang
+# 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)
+last_step = list(series.iterations)[-1]
+initial = series.iterations[1].particles["beam"].to_df()
+beam_final = series.iterations[last_step].particles["beam"]
+final = beam_final.to_df()
+
+# compare number of particles
+num_particles = 10000
+assert num_particles == len(initial)
+assert num_particles == len(final)
+
+print("Initial Beam:")
+sig_xi, sig_yi, sig_ti, emittance_xi, emittance_yi, emittance_ti = get_moments(initial)
+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 = 2.2 * num_particles**-0.5  # 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],
+    [7.51e-05, 7.51e-05, 9.99e-4, 1.98e-09, 1.98e-09, 1.97e-06],
+    rtol=rtol,
+    atol=atol,
+)
+
+
+print("")
+print("Final Beam:")
+sig_xf, sig_yf, sig_tf, emittance_xf, emittance_yf, emittance_tf = get_moments(final)
+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 = 2.2 * num_particles**-0.5  # 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],
+    [
+        9.22e-05,
+        8.55e-05,
+        0.000996,
+        2.04e-09,
+        2.01e-09,
+        1.97e-06,
+    ],
+    rtol=rtol,
+    atol=atol,
+)

examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py

@@ -46,19 +46,19 @@ def get_moments(beam):
 assert num_particles == len(final)
 
 print("Initial Beam:")
-sig_xi, sig_yi, sig_ti, emittance_xi, emittance_yi, emittance_ti = get_moments(initial)
+sig_xi, sig_yi, sig_ti, emittance_xi, emittance_yi, emittance_ti = get_moments(final)
 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 = 2.2 * num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 3 * num_particles**-0.5  # 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],
-    [7.515765e-05, 7.511883e-05, 9.997395e-4, 2.001510e-09, 1.999755e-09, 1.999289e-06],
+    [7.51e-05, 7.51e-05, 9.99e-4, 1.98e-09, 1.98e-09, 1.97e-06],
     rtol=rtol,
     atol=atol,
 )
@@ -73,18 +73,18 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = 2.2 * num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 3.5 * num_particles**-0.5  # 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],
     [
-        7.51576586332169e-05,
-        7.511883208451813e-05,
-        0.0009997395499750136,
-        2.0015106608723994e-09,
-        1.999755254276969e-09,
-        1.9992898444562777e-06,
+        9.21e-05,
+        8.54e-05,
+        0.000996,
+        2.04e-09,
+        2.01e-09,
+        1.97e-06,
     ],
     rtol=rtol,
     atol=atol,

examples/fodo_space_charge/input_fodo_Gauss2p5D_sc.in

@@ -0,0 +1,60 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 10000
+beam.units = static
+beam.kin_energy = 1.0e2
+beam.charge = 1.0e-9
+beam.particle = electron
+beam.distribution = gaussian
+beam.lambdaX = 3.9984884770e-5
+beam.lambdaY = beam.lambdaX
+beam.lambdaT = 1.0e-3
+beam.lambdaPx = 2.6623538760e-5
+beam.lambdaPy = beam.lambdaPx
+beam.lambdaPt = 2.0e-3
+beam.muxpx = -0.846574929020762
+beam.muypy = -beam.muxpx
+beam.mutpt = 0.0
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor drift1 monitor quad1 monitor drift2 monitor quad2 monitor drift3 monitor
+lattice.nslice = 25
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+drift1.type = drift
+drift1.ds = 0.25
+
+quad1.type = quad
+quad1.ds = 1.0
+quad1.k = 1.0
+
+drift2.type = drift
+drift2.ds = 0.5
+
+quad2.type = quad
+quad2.ds = 1.0
+quad2.k = -1.0
+
+drift3.type = drift
+drift3.ds = 0.25
+
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = Gauss2p5D
+algo.space_charge.gauss_nint = 101
+algo.space_charge.gauss_charge_z_bins = 129
+algo.space_charge.gauss_taylor_delta = 0.01
+
+###############################################################################
+# Diagnostics
+###############################################################################
+diag.slice_step_diagnostics = true

examples/fodo_space_charge/input_fodo_Gauss3D_sc.in

@@ -50,7 +50,7 @@ drift3.ds = 0.25
 ###############################################################################
 algo.particle_shape = 2
 algo.space_charge = Gauss3D
-
+algo.space_charge.gauss_nint = 101
 
 ###############################################################################
 # Diagnostics

examples/fodo_space_charge/run_fodo_Gauss2p5D_sc.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell, Ji Qiang
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+from impactx import ImpactX, distribution, elements
+
+sim = ImpactX()
+
+# set numerical parameters and IO control
+sim.particle_shape = 2  # B-spline order
+sim.space_charge = "Gauss2p5D"
+sim.space_charge_gauss_nint = 101
+sim.space_charge_gauss_charge_z_bins = 129
+sim.space_charge_gauss_taylor_delta = 0.01
+sim.slice_step_diagnostics = True
+
+# domain decomposition & space charge mesh
+sim.init_grids()
+
+# load a 2 GeV electron beam with an initial
+# unnormalized rms emittance of 2 nm
+kin_energy_MeV = 100  # reference energy
+bunch_charge_C = 1.0e-9  # used with space charge
+npart = 10000  # number of macro particles
+
+#   reference particle
+ref = sim.particle_container().ref_particle()
+ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV)
+
+#   particle bunch
+distr = distribution.Gaussian(
+    lambdaX=3.9984884770e-5,
+    lambdaY=3.9984884770e-5,
+    lambdaT=1.0e-3,
+    lambdaPx=2.6623538760e-5,
+    lambdaPy=2.6623538760e-5,
+    lambdaPt=2.0e-3,
+    muxpx=-0.846574929020762,
+    muypy=0.846574929020762,
+    mutpt=0.0,
+)
+sim.add_particles(bunch_charge_C, distr, npart)
+
+# add beam diagnostics
+monitor = elements.BeamMonitor("monitor", backend="h5")
+
+# design the accelerator lattice)
+ns = 25  # number of slices per ds in the element
+fodo = [
+    monitor,
+    elements.ChrDrift(name="drift1", ds=0.25, nslice=ns),
+    monitor,
+    elements.ChrQuad(name="quad1", ds=1.0, k=1.0, nslice=ns),
+    monitor,
+    elements.ChrDrift(name="drift2", ds=0.5, nslice=ns),
+    monitor,
+    elements.ChrQuad(name="quad2", ds=1.0, k=-1.0, nslice=ns),
+    monitor,
+    elements.ChrDrift(name="drift3", ds=0.25, nslice=ns),
+    monitor,
+]
+# assign a fodo segment
+sim.lattice.extend(fodo)
+
+# run simulation
+sim.track_particles()
+
+# clean shutdown
+sim.finalize()

examples/fodo_space_charge/run_fodo_Gauss3D_sc.py

@@ -13,6 +13,7 @@
 # set numerical parameters and IO control
 sim.particle_shape = 2  # B-spline order
 sim.space_charge = "Gauss3D"
+sim.space_charge_gauss_nint = 101
 sim.slice_step_diagnostics = True
 
 # domain decomposition & space charge mesh

src/initialization/Algorithms.H

@@ -22,6 +22,7 @@ namespace impactx
         False,    /**< Disabled: no space charge calculation */
         True_3D,  /**< 3D beam distribution */
         Gauss3D,  /**< Assume a 3D Gaussian beam distribution */
+        Gauss2p5D,  /**< Assume a transverse 2D Gaussian beam distribution */
         True_2D   /**< Averaged 2D transverse beam distribution with a current along s */
     );
 

src/initialization/Algorithms.cpp

@@ -66,6 +66,10 @@ namespace impactx
         {
             return SpaceChargeAlgo::Gauss3D;
         }
+        else if (space_charge == "Gauss2p5D")
+        {
+            return SpaceChargeAlgo::Gauss2p5D;
+        }
         else if (space_charge == "2D")
         {
             return SpaceChargeAlgo::True_2D;

src/particles/spacecharge/CMakeLists.txt

@@ -2,6 +2,7 @@ target_sources(lib
   PRIVATE
     ForceFromSelfFields.cpp
     Gauss3dPush.cpp
+    Gauss2p5dPush.cpp
     GatherAndPush.cpp
     HandleSpacecharge.cpp
     PoissonSolve.cpp