2D Space Charge in Particle Tracking

cmake/dependencies/ABLASTR.cmake

@@ -178,7 +178,7 @@ set(ImpactX_openpmd_src ""
 set(ImpactX_ablastr_repo "https://github.com/BLAST-WarpX/warpx.git"
     CACHE STRING
     "Repository URI to pull and build ABLASTR from if(ImpactX_ablastr_internal)")
-set(ImpactX_ablastr_branch "a76eecf1553f5d0e72ec948c9a67db4111a1544f"
+set(ImpactX_ablastr_branch "5dc12e9de9fc1428d4ad4d6c132043abedc8c531"
     CACHE STRING
     "Repository branch for ImpactX_ablastr_repo if(ImpactX_ablastr_internal)")
 

cmake/dependencies/pyAMReX.cmake

@@ -74,7 +74,7 @@ option(ImpactX_pyamrex_internal "Download & build pyAMReX" ON)
 set(ImpactX_pyamrex_repo "https://github.com/AMReX-Codes/pyamrex.git"
     CACHE STRING
     "Repository URI to pull and build pyamrex from if(ImpactX_pyamrex_internal)")
-set(ImpactX_pyamrex_branch "25.09"
+set(ImpactX_pyamrex_branch "de24942c76ebc61ffaa184204c2491aebc357a07"
     CACHE STRING
     "Repository branch for ImpactX_pyamrex_repo if(ImpactX_pyamrex_internal)")
 

docs/source/usage/how_to_run.rst

@@ -16,7 +16,7 @@ tracking of the beam envelope (6x6 covariance matrix) through linearized transpo
 ================== =============== =============== ====================
 Mode               Use Case        Generality      Space Charge Effects
 ================== =============== =============== ====================
-Particle Tracking  Full Dynamics   Most general    Supported (3D only)
+Particle Tracking  Full Dynamics   Most general    Supported (2D or 3D)
 Envelope Tracking  Rapid Scans     Linearized      Supported (2D or 3D)
 Reference Tracking Early Design    Reference orbit No
 ================== =============== =============== ====================

docs/source/usage/parameters.rst

@@ -873,8 +873,6 @@ See there ``nslice`` option on lattice elements for slicing.
 
   * ``"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.
 
     When running in envelope mode (when ``algo.track = "envelope"``), this model currently assumes that ``<xy> = <yt> = <tx> = 0``.
@@ -940,8 +938,8 @@ See there ``nslice`` option on lattice elements for slicing.
 * ``algo.poisson_solver`` (``string``, optional, default: ``"fft"``)
 
   The numerical solver to solve the Poisson equation when calculating space charge effects.
-  Currently, this is a 3D solver.
-  An additional `2D/2.5D solver <https://github.com/BLAST-ImpactX/impactx/issues/401>`__ will be added in the near future.
+  Currently, the multigrid solver supports only 3D space charge.  The fft solver supports either 2D or 3D space charge.
+  An additional `2.5D solver <https://github.com/BLAST-ImpactX/impactx/issues/401>`__ will be added in the near future.
 
   Options:
 

docs/source/usage/python.rst

@@ -70,8 +70,6 @@ Collective Effects & Overall Simulation Parameters
 
       * ``"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.
 
         When running in envelope mode (when ``algo.track = "envelope"``), this model currently assumes that ``<xy> = <yt> = <tx> = 0``.
@@ -86,8 +84,8 @@ Collective Effects & Overall Simulation Parameters
       The numerical solver to solve the Poisson equation when calculating space charge effects.
       Either ``"fft"`` (default) or ``"multigrid"``.
 
-      Currently, this is a 3D solver.
-      An additional `2D/2.5D solver <https://github.com/BLAST-ImpactX/impactx/issues/401>`__ will be added in the near future.
+      Currently, the multigrid solver supports only 3D space charge.  The fft solver supports either 2D or 3D space charge.
+      An additional `2.5D solver <https://github.com/BLAST-ImpactX/impactx/issues/401>`__ will be added in the near future.
 
       * ``fft``: Poisson's equation is solved using an Integrated Green Function method (which requires FFT calculations).
         See these references for more details `Qiang et al. (2006) <https://doi.org/10.1103/PhysRevSTAB.9.044204>`__ (+ `Erratum <https://doi.org/10.1103/PhysRevSTAB.10.129901>`__).

examples/CMakeLists.txt

@@ -1675,3 +1675,35 @@ add_impactx_test(solenoid-softedge-solvable.py
     examples/solenoid_softedge/analysis_solenoid_softedge_solvable.py
     OFF  # no plot script yet
 )
+
+# Expanding unbunched beam in free space with 2D space charge ######################
+#
+# with space charge
+add_impactx_test(expanding-fft-2d
+    examples/expanding_beam/input_expanding_fft_2D.in
+    ON   # ImpactX MPI-parallel
+    examples/expanding_beam/analysis_expanding_fft_2D.py
+    OFF  # no plot script yet
+)
+add_impactx_test(expanding-fft-2d.py
+    examples/expanding_beam/run_expanding_fft_2D.py
+    OFF   # ImpactX MPI-parallel
+    examples/expanding_beam/analysis_expanding_fft_2D.py
+    OFF  # no plot script yet
+)
+
+# FODO cell with 2D space charge using particle tracking ######################
+#
+# with space charge
+add_impactx_test(fodo-2d-sc
+    examples/fodo_space_charge/input_fodo_2d_sc.in
+    ON   # ImpactX MPI-parallel
+    examples/fodo_space_charge/analysis_fodo_2d_sc.py
+    OFF  # no plot script yet
+)
+add_impactx_test(fodo-2d-sc.py
+    examples/fodo_space_charge/run_fodo_2d_sc.py
+    OFF   # ImpactX MPI-parallel
+    examples/fodo_space_charge/analysis_fodo_2d_sc.py
+    OFF  # no plot script yet
+)

examples/expanding_beam/README.rst

@@ -1,7 +1,8 @@
+
 .. _examples-expanding:
 
-Expanding Beam in Free Space
-============================
+Expanding Beam in Free Space with 3D Space Charge
+==================================================
 
 A coasting bunch expanding freely in free space under its own space charge.
 
@@ -69,3 +70,58 @@ We run the following script to analyze correctness:
    .. literalinclude:: analysis_expanding.py
       :language: python3
       :caption: You can copy this file from ``examples/expanding/analysis_expanding.py``.
+
+
+
+.. _examples-expanding-fft-2d:
+
+Expanding Beam in Free Space with 2D Space Charge
+==================================================
+
+A long, coasting unbunched beam expanding freely in free space under its own 2D space charge.
+
+We use a cold (zero emittance) 250 MeV electron bunch whose
+initial distribution is a uniformly-populated cylinder of radius R0 = 1 mm.
+
+In the laboratory frame, the beam expands to twice its original transverse size.  This is tested using the second moments of the distribution.
+
+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.
+
+
+Run
+---
+
+This example can be run **either** as:
+
+* **Python** script: ``python3 run_expanding_fft_2D.py`` or
+* ImpactX **executable** using an input file: ``impactx input_expanding_fft_2D.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.
+
+We also provide the same example with the multi-grid (MLMG) Poisson solver.
+
+.. tab-set::
+
+   .. tab-item:: Python: Script (FFT)
+
+      .. literalinclude:: run_expanding_fft_2D.py
+         :language: python3
+         :caption: You can copy this file from ``examples/expanding/run_expanding_fft_2D.py``.
+
+   .. tab-item:: Executable: Input File (FFT)
+
+       .. literalinclude:: input_expanding_fft_2D.in
+          :language: ini
+          :caption: You can copy this file from ``examples/expanding/input_expanding_fft_2D.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_expanding_fft_2D.py``
+
+   .. literalinclude:: analysis_expanding_fft_2D.py
+      :language: python3
+      :caption: You can copy this file from ``examples/expanding/analysis_expanding_fft_2D.py``.

examples/expanding_beam/analysis_expanding_fft_2D.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell
+# 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()
+final = series.iterations[last_step].particles["beam"].to_df()
+
+# compare number of particles
+num_particles = 10000
+assert num_particles == len(initial)
+assert num_particles == len(final)
+
+print("Initial Beam:")
+sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(initial)
+print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
+print(
+    f"  emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
+)
+
+atol = 0.0  # ignored
+rtol = 1.5 * num_particles**-0.5  # from random sampling of a smooth distribution
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+assert np.allclose(
+    [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
+    [
+        5.0e-004,
+        5.0e-004,
+        1.0e-003,
+        0.0e-006,
+        0.0e-006,
+        0.0e-006,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+
+print("")
+print("Final Beam:")
+sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final)
+print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
+print(
+    f"  emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
+)
+
+atol = 0.0  # ignored
+rtol = 2.0 * num_particles**-0.5  # from random sampling of a smooth distribution
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+assert np.allclose(
+    [sigx, sigy, sigt],
+    [
+        1.0e-003,
+        1.0e-003,
+        1.0e-003,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+atol = 1.0e-8
+rtol = 0.0  # ignored
+assert np.allclose(
+    [emittance_x, emittance_y, emittance_t],
+    [
+        0.0,
+        0.0,
+        0.0,
+    ],
+    rtol=rtol,
+    atol=atol,
+)

examples/expanding_beam/input_expanding_fft_2D.in

@@ -0,0 +1,49 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 10000  # outside tests, use 1e5 or more
+beam.units = static
+beam.kin_energy = 250.0
+beam.current = 0.15
+beam.particle = proton
+beam.distribution = kvdist
+beam.lambdaX = 5.0e-4
+beam.lambdaY = beam.lambdaX
+beam.lambdaT = 1.0e-3  #result should not depend on this value
+beam.lambdaPx = 0.0
+beam.lambdaPy = 0.0
+beam.lambdaPt = 0.0
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor drift1 monitor
+lattice.nslice = 100
+
+drift1.type = drift
+drift1.ds = 10.612823669911099  #doubling-distance
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+#algo.track = "envelope"
+algo.space_charge = 2D
+algo.poisson_solver = "fft"
+
+amr.n_cell = 32 32 1
+amr.blocking_factor_x = 16
+amr.blocking_factor_y = 16
+amr.blocking_factor_z = 1
+
+geometry.prob_relative = 1.1
+
+###############################################################################
+# Diagnostics
+###############################################################################
+diag.slice_step_diagnostics = true

examples/expanding_beam/run_expanding_fft_2D.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+from impactx import ImpactX, distribution, elements
+
+sim = ImpactX()
+
+# set numerical parameters and IO control
+sim.max_level = 0
+sim.n_cell = [32, 32, 1]
+sim.blocking_factor_x = [16]
+sim.blocking_factor_y = [16]
+sim.blocking_factor_z = [1]
+
+sim.particle_shape = 2  # B-spline order
+sim.space_charge = "2D"
+sim.poisson_solver = "fft"
+sim.dynamic_size = True
+sim.prob_relative = [1.1]
+
+# beam diagnostics
+# sim.diagnostics = False  # benchmarking
+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 = 250  # reference energy
+beam_current_A = 0.15  # beam current
+npart = 10000  # number of macro particles (outside tests, use 1e5 or more)
+
+#   reference particle
+ref = sim.particle_container().ref_particle()
+ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
+
+#   particle bunch
+distr = distribution.KVdist(
+    lambdaX=5.0e-4,
+    lambdaY=5.0e-4,
+    lambdaT=1.0e-3,
+    lambdaPx=0.0,
+    lambdaPy=0.0,
+    lambdaPt=0.0,
+)
+sim.add_particles(beam_current_A, distr, npart)
+
+# add beam diagnostics
+monitor = elements.BeamMonitor("monitor", backend="h5")
+
+# design the accelerator lattice
+doubling_distance = 10.612823669911099
+
+sim.lattice.extend(
+    [monitor, elements.Drift(name="d1", ds=doubling_distance, nslice=100), monitor]
+)
+
+# run simulation
+sim.track_particles()
+
+# clean shutdown
+sim.finalize()