Examples for 3D space charge benchmarking

- Modified the initial beam size in the IOTA lens benchmark example.
- Added 2 benchmarks of 3D space charge for initial testing.
- Add documentation for 2 benchmarks with space charge.
- Add a benchmark example with space charge and periodic s-dependent focusing.
- Added an s-dependent example using a Kurth beam without space charge.
- Modified tolerance for IOTA lens benchmark example.
  Reduced tolerance to account for smaller initial beam size and
  improved preservation of invariants of motion.
- Modified tolerances of space charge examples to allow CI tests to
  pass when space charge is not active.

- Modified tolerance for space charge examples.
  These should fail unless space charge is turned on.

examples/CMakeLists.txt

@@ -175,6 +175,16 @@ add_impactx_test(cfchannel
     OFF  # no plot script yet
 )
 
+# Constant Focusing Channel with Space Charge #################################
+#
+add_impactx_test(cfchannel_spacecharge
+    examples/cfchannel/input_cfchannel_10nC.in
+      OFF  # ImpactX MPI-parallel
+      OFF  # ImpactX Python interface
+    examples/cfchannel/analysis_cfchannel_10nC.py
+    OFF  # no plot script yet
+)
+
 # Python: Constant Focusing Channel ###########################################
 #
 add_impactx_test(cfchannel.py
@@ -315,7 +325,6 @@ add_impactx_test(iotalattice.MPI
     OFF  # no plot script yet
 )
 
-
 # Python: IOTA Linear Lattice Test ############################################
 #
 add_impactx_test(iotalattice.py.MPI
@@ -325,3 +334,29 @@ add_impactx_test(iotalattice.py.MPI
     examples/iota_lattice/analysis_iotalattice.py
     OFF  # no plot script yet
 )
+
+# Kurth Beam in a Periodic Channel Test ###################################################
+#
+add_impactx_test(kurth_periodic
+    examples/kurth/input_kurth_periodic.in
+      OFF  # ImpactX MPI-parallel
+      OFF  # ImpactX Python interface
+    examples/kurth/analysis_kurth_periodic.py
+    OFF  # no plot script yet
+)
+
+add_impactx_test(kurth_10nC_periodic
+    examples/kurth/input_kurth_10nC_periodic.in
+      OFF  # ImpactX MPI-parallel
+      OFF  # ImpactX Python interface
+    examples/kurth/analysis_kurth_10nC_periodic.py
+    OFF  # no plot script yet
+)
+
+add_impactx_test(kurth_10nC
+    examples/kurth/input_kurth_10nC.in
+      OFF  # ImpactX MPI-parallel
+      OFF  # ImpactX Python interface
+    examples/kurth/analysis_kurth_10nC.py
+    OFF  # no plot script yet
+)

examples/cfchannel/README.rst

@@ -53,3 +53,59 @@ We run the following script to analyze correctness:
    .. literalinclude:: analysis_cfchannel.py
       :language: python3
       :caption: You can copy this file from ``examples/cfchannel/analysis_cfchannel.py``.
+
+
+.. _examples-cfchannel-10nC:
+
+Constant Focusing Channel with Space Charge
+===========================================
+
+RMS-matched beam in a constant focusing channel with space charge.
+
+The matched Twiss parameters at entry are:
+
+* :math:`\beta_\mathrm{x} = 1.477305` m
+* :math:`\alpha_\mathrm{x} = 0.0`
+* :math:`\beta_\mathrm{y} = 1.477305` m
+* :math:`\alpha_\mathrm{y} = 0.0`
+
+We use a 2 GeV proton beam with initial unnormalized rms emittance of 1 um.
+The longitudinal beam parameters are chosen so that the bunch has radial symmetry when viewed in the beam rest frame.
+The bunch charge is set to 10 nC, resulting in a transverse tune depression ratio of 0.67.
+The initial distribution used is a 6D waterbag.
+
+The beam second moments should remain nearly unchanged, except for some small emittance growth due to nonlinear space charge.
+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:`
+
+
+Run
+---
+
+This example can be run as an app with an input file (``impactx input_cfchannel_10nC.in``).
+Each can also be prefixed with an `MPI executor <https://www.mpi-forum.org>`__, such as ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system.
+
+.. tab-set::
+
+   .. tab-item:: Python Script (TODO)
+
+       This one is to-do.
+
+   .. tab-item:: App Input File
+
+       .. literalinclude:: input_cfchannel_10nC.in
+          :language: ini
+          :caption: You can copy this file from ``examples/cfchannel/input_cfchannel_10nC.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_cfchannel_10nC.py``
+
+   .. literalinclude:: analysis_cfchannel_10nC.py
+      :language: python3
+      :caption: You can copy this file from ``examples/cfchannel/analysis_cfchannel_10nC.py``.

examples/cfchannel/analysis_cfchannel_10nC.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+
+import glob
+
+import numpy as np
+import pandas as pd
+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["x"], moment=2) ** 0.5  # variance -> std dev.
+    sigpx = moment(beam["px"], moment=2) ** 0.5
+    sigy = moment(beam["y"], moment=2) ** 0.5
+    sigpy = moment(beam["py"], moment=2) ** 0.5
+    sigt = moment(beam["t"], moment=2) ** 0.5
+    sigpt = moment(beam["pt"], moment=2) ** 0.5
+
+    epstrms = beam.cov(ddof=0)
+    emittance_x = (sigx**2 * sigpx**2 - epstrms["x"]["px"] ** 2) ** 0.5
+    emittance_y = (sigy**2 * sigpy**2 - epstrms["y"]["py"] ** 2) ** 0.5
+    emittance_t = (sigt**2 * sigpt**2 - epstrms["t"]["pt"] ** 2) ** 0.5
+
+    return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t)
+
+
+def read_all_files(file_pattern):
+    """Read in all CSV files from each MPI rank (and potentially OpenMP
+    thread). Concatenate into one Pandas dataframe.
+
+    Returns
+    -------
+    pandas.DataFrame
+    """
+    return pd.concat(
+        (
+            pd.read_csv(filename, delimiter=r"\s+")
+            for filename in glob.glob(file_pattern)
+        ),
+        axis=0,
+        ignore_index=True,
+    ).set_index("id")
+
+
+# initial/final beam on rank zero
+initial = read_all_files("diags/initial_beam.txt.*")
+final = read_all_files("diags/output_beam.txt.*")
+
+# 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 = 1.0  # a big number
+rtol = 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],
+    [
+        1.2154443728379865788e-003,
+        1.2154443728379865788e-003,
+        4.0956844276541331005317e-004,
+        1.000000000e-006,
+        1.000000000e-006,
+        1.000000000e-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 = 1.0  # a big number
+rtol = 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],
+    [
+        1.2154443728379865788e-003,
+        1.2154443728379865788e-003,
+        4.0956844276541331005317e-004,
+        1.000000000e-006,
+        1.000000000e-006,
+        1.000000000e-006,
+    ],
+    rtol=rtol,
+    atol=atol,
+)

examples/cfchannel/input_cfchannel_10nC.in

@@ -0,0 +1,37 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 10000
+beam.units = static
+beam.energy = 2.0e3
+beam.charge = 1.0e-8
+beam.particle = proton
+beam.distribution = waterbag
+beam.sigmaX = 1.2154443728379865788e-3
+beam.sigmaY = 1.2154443728379865788e-3
+beam.sigmaT = 4.0956844276541331005e-4
+beam.sigmaPx = 8.2274435782286157175e-4
+beam.sigmaPy = 8.2274435782286157175e-4
+beam.sigmaPt = 2.4415943602685364584e-3
+beam.muxpx = 0.0
+beam.muypy = 0.0
+beam.mutpt = 0.0
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+
+lattice.elements = constf1
+
+constf1.type = constf
+constf1.ds = 2.0
+constf1.kx = 1.0
+constf1.ky = 1.0
+constf1.kt = 1.0
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = true

examples/kurth/README.rst

@@ -48,3 +48,54 @@ We run the following script to analyze correctness:
    .. literalinclude:: analysis_kurth.py
       :language: python3
       :caption: You can copy this file from ``examples/kurth/analysis_kurth.py``.
+
+
+.. _examples-kurth-10nC:
+
+Kurth Distribution in a Constant Focusing Channel with Space Charge
+===================================================================
+
+Stationary Kurth distribution in a constant focusing channel with space charge.
+
+The distribution is radially symmetric in (x,y,t) space, and matched to a
+radially symmetric constant linear focusing.
+
+We use a 2 GeV proton beam with initial unnormalized rms emittance of 1 um
+in all three phase planes.  The bunch charge is set to 10 nC, corresponding
+to a transverse tune depression ratio of 0.67.
+
+The particle distribution should remain unchanged, to within the level expected due to numerical particle noise.
+This fact is independent of the length of the channel.  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:`
+
+
+Run
+---
+
+This example can be run as an app with an input file (``impactx input_kurth_periodic.in``).
+Each can also be prefixed with an `MPI executor <https://www.mpi-forum.org>`__, such as ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system.
+
+.. tab-set::
+
+   .. tab-item:: Python Script (TODO)
+
+       This one is to-do.
+
+   .. tab-item:: App Input File
+
+       .. literalinclude:: input_kurth_periodic.in
+          :language: ini
+          :caption: You can copy this file from ``examples/kurth/input_kurth_periodic.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_kurth_periodic.py``
+
+   .. literalinclude:: analysis_kurth_periodic.py
+      :language: python3
+      :caption: You can copy this file from ``examples/kurth/analysis_kurth_periodic.py``.