Skip to content

Commit fa4c54d

Browse files
qianglblax3lcemitch99
authored
Gauss 3D Space Charge Pusher (#1127)
* Gauss 3D SC Pusher Add the new Gauss 3D space charge pusher. * add reference for 3D Gaussian SC solver * Fix Compilation Together * fixed a defect in 3D Gaussian SC * added FODO test example with SC from a 3D Gaussian distribution * added reference for the 3D SC Gaussian distribution solver. * User-Facing Docs (Manual) * Example: README & CMake Update - replace with new file name in doc - move into own section in CMake * Source: Cleaning, Formatting, TODOs * Update analysis * GPU Support, Performance Opt - avoid temporary arrays - integrate in one go * Envelope: Abort on Gauss3D Only used for particles. Ref orbit asserts well, too. * Python Bindings: New Allowed Value * Python Example * Cleanup * GPU Kernel: `AMREX_GPU_DEVICE` * Apply suggestions from code review --------- Co-authored-by: Axel Huebl <[email protected]> Co-authored-by: Chad Mitchell <[email protected]>
1 parent 771000b commit fa4c54d

File tree

15 files changed

+596
-44
lines changed

15 files changed

+596
-44
lines changed

‎docs/source/usage/parameters.rst‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -879,6 +879,12 @@ See there ``nslice`` option on lattice elements for slicing.
879879

880880
When running in envelope mode (when ``algo.track = "envelope"``), this model currently assumes that ``<xy> = <yt> = <tx> = 0``.
881881

882+
* ``"Gauss3D"`: Calculate 3D space charge forces as if the beam was a Gaussian distribution.
883+
884+
This model is supported only in particle tracking mode (when ``algo.track = "particles"``).
885+
Ref.: J. Qiang et al., "Two-and-a-half dimensional symplectic space-charge solver", LBNL Report Number: LBNL-2001674 (2025).
886+
(This reference describes both 3D and 2.5D models.)
887+
882888
* ``amr.n_cell`` (3 integers) optional (default: 1 `blocking_factor <https://amrex-codes.github.io/amrex/docs_html/GridCreation.html>`__ per MPI process)
883889

884890
The number of grid points along each direction (on the **coarsest level**)

‎docs/source/usage/python.rst‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,11 @@ Collective Effects & Overall Simulation Parameters
7676

7777
When running in envelope mode (when ``algo.track = "envelope"``), this model currently assumes that ``<xy> = <yt> = <tx> = 0``.
7878

79+
* ``"Gauss3D"`: Calculate 3D space charge forces as if the beam was a Gaussian distribution.
80+
81+
This model is supported only in particle tracking mode (when ``algo.track = "particles"``).
82+
Ref.: J. Qiang et al., "Two-and-a-half dimensional symplectic space-charge solver", LBNL Report Number: LBNL-2001674 (2025).
83+
(This reference describes both 3D and 2.5D models.)
7984
.. py:property:: poisson_solver
8085
8186
The numerical solver to solve the Poisson equation when calculating space charge effects.

‎examples/CMakeLists.txt‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1388,6 +1388,21 @@ add_impactx_test(FODO.envelope.sc.py
13881388
OFF # no plot script yet
13891389
)
13901390

1391+
# FODO cell with Gaussian 3D space charge using particle tracking #############
1392+
#
1393+
add_impactx_test(FODO.Gauss3d.sc
1394+
examples/fodo_space_charge/input_fodo_Gauss3D_sc.in
1395+
ON # ImpactX MPI-parallel
1396+
examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py
1397+
OFF # no plot script yet
1398+
)
1399+
add_impactx_test(FODO.Gauss3d.sc.py
1400+
examples/fodo_space_charge/run_fodo_Gauss3D_sc.py
1401+
ON # ImpactX MPI-parallel
1402+
examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py
1403+
OFF # no plot script yet
1404+
)
1405+
13911406
# A single bend with incoherent synchrotron radiation ######################
13921407
#
13931408
# no space charge

‎examples/fodo_space_charge/README.rst‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,3 +51,56 @@ We run the following script to analyze correctness:
5151
.. literalinclude:: analysis_fodo_envelope_sc.py
5252
:language: python3
5353
:caption: You can copy this file from ``examples/fodo_space_charge/analysis_fodo_envelope_sc.py``.
54+
55+
56+
.. _examples-fodo-envelope-sc-gaussian:
57+
58+
FODO Cell with 3D Gaussian Space Charge Using Particle Tracking
59+
===============================================================
60+
61+
This example illustrates a 1 nC electron beam with a kinetic energy of 100 MeV in a FODO cell,
62+
with 3D Gaussian space charge included. The parameters are those described in:
63+
64+
The purpose of this example is to illustrate the use of particle tracking mode with 3D space charge from a Gaussian density
65+
distribution.
66+
67+
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 with small noticeable growth.
68+
69+
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.
70+
71+
72+
Run
73+
---
74+
75+
This example can be run **either** as:
76+
77+
* **Python** script: ``python3 run_fodo_Gauss3D_sc.py`` or
78+
* ImpactX **executable** using an input file: ``impactx input_fodo_Gauss3D_sc.in``
79+
80+
For `MPI-parallel <https://www.mpi-forum.org>`__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system.
81+
82+
.. tab-set::
83+
84+
.. tab-item:: Python: Script
85+
86+
.. literalinclude:: run_fodo_Gauss3D_sc.py
87+
:language: python3
88+
:caption: You can copy this file from ``examples/fodo_space_charge/run_fodo_Gauss3D_sc.py``.
89+
90+
.. tab-item:: Executable: Input File
91+
92+
.. literalinclude:: input_fodo_Gauss3D_sc.in
93+
:language: ini
94+
:caption: You can copy this file from ``examples/fodo_space_charge/input_fodo_Gauss3D_sc.in``.
95+
96+
97+
Analyze
98+
-------
99+
100+
We run the following script to analyze correctness:
101+
102+
.. dropdown:: Script ``analysis_fodo_Gauss3D_sc.py``
103+
104+
.. literalinclude:: analysis_fodo_Gauss3D_sc.py
105+
:language: python3
106+
:caption: You can copy this file from ``examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py``.

‎examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
#!/usr/bin/env python3
2+
#
3+
# Copyright 2022-2023 ImpactX contributors
4+
# Authors: Axel Huebl, Ji Qiang
5+
# License: BSD-3-Clause-LBNL
6+
#
7+
8+
import numpy as np
9+
import openpmd_api as io
10+
from scipy.stats import moment
11+
12+
13+
def get_moments(beam):
14+
"""Calculate standard deviations of beam position & momenta
15+
and emittance values
16+
17+
Returns
18+
-------
19+
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t
20+
"""
21+
sigx = moment(beam["position_x"], moment=2) ** 0.5 # variance -> std dev.
22+
sigpx = moment(beam["momentum_x"], moment=2) ** 0.5
23+
sigy = moment(beam["position_y"], moment=2) ** 0.5
24+
sigpy = moment(beam["momentum_y"], moment=2) ** 0.5
25+
sigt = moment(beam["position_t"], moment=2) ** 0.5
26+
sigpt = moment(beam["momentum_t"], moment=2) ** 0.5
27+
28+
epstrms = beam.cov(ddof=0)
29+
emittance_x = (sigx**2 * sigpx**2 - epstrms["position_x"]["momentum_x"] ** 2) ** 0.5
30+
emittance_y = (sigy**2 * sigpy**2 - epstrms["position_y"]["momentum_y"] ** 2) ** 0.5
31+
emittance_t = (sigt**2 * sigpt**2 - epstrms["position_t"]["momentum_t"] ** 2) ** 0.5
32+
33+
return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t)
34+
35+
36+
# initial/final beam
37+
series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
38+
last_step = list(series.iterations)[-1]
39+
initial = series.iterations[1].particles["beam"].to_df()
40+
beam_final = series.iterations[last_step].particles["beam"]
41+
final = beam_final.to_df()
42+
43+
# compare number of particles
44+
num_particles = 10000
45+
assert num_particles == len(initial)
46+
assert num_particles == len(final)
47+
48+
print("Initial Beam:")
49+
sig_xi, sig_yi, sig_ti, emittance_xi, emittance_yi, emittance_ti = get_moments(initial)
50+
print(f" sigx={sig_xi:e} sigy={sig_yi:e} sigt={sig_ti:e}")
51+
print(
52+
f" emittance_x={emittance_xi:e} emittance_y={emittance_yi:e} emittance_t={emittance_ti:e}"
53+
)
54+
55+
atol = 0.0 # ignored
56+
rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution
57+
print(f" rtol={rtol} (ignored: atol~={atol})")
58+
59+
assert np.allclose(
60+
[sig_xi, sig_yi, sig_ti, emittance_xi, emittance_yi, emittance_ti],
61+
[7.515765e-05, 7.511883e-05, 9.997395e-4, 2.001510e-09, 1.999755e-09, 1.999289e-06],
62+
rtol=rtol,
63+
atol=atol,
64+
)
65+
66+
67+
print("")
68+
print("Final Beam:")
69+
sig_xf, sig_yf, sig_tf, emittance_xf, emittance_yf, emittance_tf = get_moments(initial)
70+
print(f" sigx={sig_xf:e} sigy={sig_yf:e} sigt={sig_tf:e}")
71+
print(
72+
f" emittance_x={emittance_xf:e} emittance_y={emittance_yf:e} emittance_t={emittance_tf:e}"
73+
)
74+
75+
atol = 0.0 # ignored
76+
rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution
77+
print(f" rtol={rtol} (ignored: atol~={atol})")
78+
79+
assert np.allclose(
80+
[sig_xf, sig_yf, sig_tf, emittance_xf, emittance_yf, emittance_tf],
81+
[
82+
7.51576586332169e-05,
83+
7.511883208451813e-05,
84+
0.0009997395499750136,
85+
2.0015106608723994e-09,
86+
1.999755254276969e-09,
87+
1.9992898444562777e-06,
88+
],
89+
rtol=rtol,
90+
atol=atol,
91+
)

‎examples/fodo_space_charge/input_fodo_Gauss3D_sc.in‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
###############################################################################
2+
# Particle Beam(s)
3+
###############################################################################
4+
beam.npart = 10000
5+
beam.units = static
6+
beam.kin_energy = 1.0e2
7+
beam.charge = 1.0e-9
8+
beam.particle = electron
9+
beam.distribution = gaussian
10+
beam.lambdaX = 3.9984884770e-5
11+
beam.lambdaY = beam.lambdaX
12+
beam.lambdaT = 1.0e-3
13+
beam.lambdaPx = 2.6623538760e-5
14+
beam.lambdaPy = beam.lambdaPx
15+
beam.lambdaPt = 2.0e-3
16+
beam.muxpx = -0.846574929020762
17+
beam.muypy = -beam.muxpx
18+
beam.mutpt = 0.0
19+
20+
21+
###############################################################################
22+
# Beamline: lattice elements and segments
23+
###############################################################################
24+
lattice.elements = monitor drift1 monitor quad1 monitor drift2 monitor quad2 monitor drift3 monitor
25+
lattice.nslice = 25
26+
27+
monitor.type = beam_monitor
28+
monitor.backend = h5
29+
30+
drift1.type = drift
31+
drift1.ds = 0.25
32+
33+
quad1.type = quad
34+
quad1.ds = 1.0
35+
quad1.k = 1.0
36+
37+
drift2.type = drift
38+
drift2.ds = 0.5
39+
40+
quad2.type = quad
41+
quad2.ds = 1.0
42+
quad2.k = -1.0
43+
44+
drift3.type = drift
45+
drift3.ds = 0.25
46+
47+
48+
###############################################################################
49+
# Algorithms
50+
###############################################################################
51+
algo.particle_shape = 2
52+
algo.space_charge = Gauss3D
53+
54+
55+
###############################################################################
56+
# Diagnostics
57+
###############################################################################
58+
diag.slice_step_diagnostics = true

‎examples/fodo_space_charge/run_fodo_Gauss3D_sc.py‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
#!/usr/bin/env python3
2+
#
3+
# Copyright 2022-2023 ImpactX contributors
4+
# Authors: Axel Huebl, Chad Mitchell
5+
# License: BSD-3-Clause-LBNL
6+
#
7+
# -*- coding: utf-8 -*-
8+
9+
from impactx import ImpactX, distribution, elements
10+
11+
sim = ImpactX()
12+
13+
# set numerical parameters and IO control
14+
sim.particle_shape = 2 # B-spline order
15+
sim.space_charge = "Gauss3D"
16+
sim.slice_step_diagnostics = True
17+
18+
# domain decomposition & space charge mesh
19+
sim.init_grids()
20+
21+
# load a 2 GeV electron beam with an initial
22+
# unnormalized rms emittance of 2 nm
23+
kin_energy_MeV = 100 # reference energy
24+
bunch_charge_C = 1.0e-9 # used with space charge
25+
npart = 10000 # number of macro particles
26+
27+
# reference particle
28+
ref = sim.particle_container().ref_particle()
29+
ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV)
30+
31+
# particle bunch
32+
distr = distribution.Gaussian(
33+
lambdaX=3.9984884770e-5,
34+
lambdaY=3.9984884770e-5,
35+
lambdaT=1.0e-3,
36+
lambdaPx=2.6623538760e-5,
37+
lambdaPy=2.6623538760e-5,
38+
lambdaPt=2.0e-3,
39+
muxpx=-0.846574929020762,
40+
muypy=0.846574929020762,
41+
mutpt=0.0,
42+
)
43+
sim.add_particles(bunch_charge_C, distr, npart)
44+
45+
# add beam diagnostics
46+
monitor = elements.BeamMonitor("monitor", backend="h5")
47+
48+
# design the accelerator lattice)
49+
ns = 25 # number of slices per ds in the element
50+
fodo = [
51+
monitor,
52+
elements.ChrDrift(name="drift1", ds=0.25, nslice=ns),
53+
monitor,
54+
elements.ChrQuad(name="quad1", ds=1.0, k=1.0, nslice=ns),
55+
monitor,
56+
elements.ChrDrift(name="drift2", ds=0.5, nslice=ns),
57+
monitor,
58+
elements.ChrQuad(name="quad2", ds=1.0, k=-1.0, nslice=ns),
59+
monitor,
60+
elements.ChrDrift(name="drift3", ds=0.25, nslice=ns),
61+
monitor,
62+
]
63+
# assign a fodo segment
64+
sim.lattice.extend(fodo)
65+
66+
# run simulation
67+
sim.track_particles()
68+
69+
# clean shutdown
70+
sim.finalize()

‎src/initialization/Algorithms.H‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,8 @@ namespace impactx
2121
AMREX_ENUM(SpaceChargeAlgo,
2222
False, /**< Disabled: no space charge calculation */
2323
True_3D, /**< 3D beam distribution */
24-
True_2D /**< averaged 2D transverse beam distribution with a current along s */
24+
Gauss3D, /**< Assume a 3D Gaussian beam distribution */
25+
True_2D /**< Averaged 2D transverse beam distribution with a current along s */
2526
);
2627

2728
/** Return the currently active space charge algorithm */

‎src/initialization/Algorithms.cpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,10 @@ namespace impactx
6262
{
6363
return SpaceChargeAlgo::True_3D;
6464
}
65+
else if (space_charge == "Gauss3D")
66+
{
67+
return SpaceChargeAlgo::Gauss3D;
68+
}
6569
else if (space_charge == "2D")
6670
{
6771
return SpaceChargeAlgo::True_2D;

‎src/particles/spacecharge/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
target_sources(lib
22
PRIVATE
33
ForceFromSelfFields.cpp
4+
Gauss3dPush.cpp
45
GatherAndPush.cpp
56
HandleSpacecharge.cpp
67
PoissonSolve.cpp

0 commit comments

Comments
 (0)