Added examples for 3D space charge benchmarking.

examples/CMakeLists.txt

@@ -183,34 +183,34 @@ add_impactx_test(cfchannel
     OFF  # no plot script yet
 )
 
-# Python: Constant Focusing Channel ###########################################
+# Constant Focusing Channel with Space Charge #################################
 #
-add_impactx_test(cfchannel.py
-    examples/cfchannel/run_cfchannel.py
+add_impactx_test(cfchannel_spacecharge
+    examples/cfchannel/input_cfchannel_10nC.in
       OFF  # ImpactX MPI-parallel
-      ON   # ImpactX Python interface
-    examples/cfchannel/analysis_cfchannel.py
+      OFF  # ImpactX Python interface
+    examples/cfchannel/analysis_cfchannel_10nC.py
     OFF  # no plot script yet
 )
 
-# Kurth Distribution Test #####################################################
+# Python: Constant Focusing Channel with Space Charge #########################
 #
-add_impactx_test(kurth
-    examples/kurth/input_kurth.in
+add_impactx_test(cfchannel_spacecharge.py
+    examples/cfchannel/run_cfchannel_10nC.py
       OFF  # ImpactX MPI-parallel
-      OFF  # ImpactX Python interface
-    examples/kurth/analysis_kurth.py
+      ON   # ImpactX Python interface
+    examples/cfchannel/analysis_cfchannel_10nC.py
     OFF  # no plot script yet
 )
 
-# Python: Kurth Distribution Test #############################################
+# Python: Constant Focusing Channel ###########################################
 #
-add_impactx_test(kurth.py
-    examples/kurth/run_kurth.py
+add_impactx_test(cfchannel.py
+    examples/cfchannel/run_cfchannel.py
       OFF  # ImpactX MPI-parallel
-      ON  # ImpactX Python interface
-    examples/kurth/analysis_kurth.py
-    OFF  # not plotting script yet
+      ON   # ImpactX Python interface
+    examples/cfchannel/analysis_cfchannel.py
+    OFF  # no plot script yet
 )
 
 # 6D Gaussian Distribution Test ###############################################
@@ -343,7 +343,6 @@ add_impactx_test(iotalattice.MPI
     OFF  # no plot script yet
 )
 
-
 # Python: IOTA Linear Lattice Test ############################################
 #
 add_impactx_test(iotalattice.py.MPI
@@ -354,6 +353,39 @@ add_impactx_test(iotalattice.py.MPI
     OFF  # no plot script yet
 )
 
+# Kurth Beam in a Periodic Channel Test ###################################################
+#
+# w/o space charge
+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_periodic.py
+    examples/kurth/run_kurth_periodic.py
+      OFF  # ImpactX MPI-parallel
+      ON   # ImpactX Python interface
+    examples/kurth/analysis_kurth_periodic.py
+    OFF  # no plot script yet
+)
+
+# w/ space charge
+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_periodic.py
+    examples/kurth/run_kurth_10nC_periodic.py
+      OFF  # ImpactX MPI-parallel
+      ON   # ImpactX Python interface
+    examples/kurth/analysis_kurth_10nC_periodic.py
+    OFF  # no plot script yet
+)
 
 # Acceleration by RF Cavities #################################################
 #
@@ -365,7 +397,6 @@ add_impactx_test(rfcavity.MPI
     OFF  # no plot script yet
 )
 
-
 # Python: Acceleration by RF Cavities #########################################
 #
 add_impactx_test(rfcavity.py.MPI

examples/cfchannel/README.rst

@@ -53,3 +53,61 @@ 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 a Python script (``python3 run_cfchannel_10nC.py``) or  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
+
+       .. literalinclude:: run_cfchannel_10nC.py
+          :language: python3
+          :caption: You can copy this file from ``examples/cfchannel/run_cfchannel_10nC.py``.
+
+   .. 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/beam_000000.*")
+final = read_all_files("diags/beam_final.*")
+
+# 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],
+    [
+        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 = 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],
+    [
+        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,41 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 10000
+#beam.npart = 100000  # optional for increased precision
+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
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = constf1
+lattice.nslice = 50
+#lattice.nslice = 60 # optional for increased precision
+
+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
+
+amr.n_cell = 48 48 40
+#amr.n_cell = 72 72 64 #optional for increased precision
+geometry.prob_relative = 1.0

examples/cfchannel/run_cfchannel_10nC.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022 ImpactX contributors
+# Authors: Marco Garten, Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+import amrex
+from impactx import ImpactX, RefPart, distribution, elements
+
+pp_amr = amrex.ParmParse("amr")
+pp_amr.addarr("n_cell", [48, 48, 40])  # [72, 72, 64] for increased precision
+
+sim = ImpactX()
+
+# set numerical parameters and IO control
+sim.particle_shape = 2  # B-spline order
+sim.space_charge = True
+sim.prob_relative = 1.0
+# sim.diagnostics = False  # benchmarking
+sim.slice_step_diagnostics = True
+
+# domain decomposition & space charge mesh
+sim.init_grids()
+
+# load a 2 GeV proton beam with an initial
+# normalized transverse rms emittance of 1 um
+energy_MeV = 2.0e3  # reference energy
+bunch_charge_C = 1.0e-8  # used with space charge
+npart = 10000  # number of macro particles; use 1e5 for increased precision
+
+#   reference particle
+ref = sim.particle_container().ref_particle()
+ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_energy_MeV(energy_MeV)
+
+#   particle bunch
+distr = distribution.Waterbag(
+    sigmaX=1.2154443728379865788e-3,
+    sigmaY=1.2154443728379865788e-3,
+    sigmaT=4.0956844276541331005e-4,
+    sigmaPx=8.2274435782286157175e-4,
+    sigmaPy=8.2274435782286157175e-4,
+    sigmaPt=2.4415943602685364584e-3,
+)
+sim.add_particles(bunch_charge_C, distr, npart)
+
+# design the accelerator lattice
+nslice = 50  # use 1e5 for increased precision
+sim.lattice.append(elements.ConstF(ds=2.0, kx=1.0, ky=1.0, kt=1.0, nslice=nslice))
+
+# run simulation
+sim.evolve()
+
+# clean shutdown
+del sim
+amrex.finalize()

examples/expanding_beam/input_expanding.in

@@ -13,9 +13,6 @@ beam.sigmaT = 9.12241869e-7
 beam.sigmaPx = 0.0
 beam.sigmaPy = 0.0
 beam.sigmaPt = 0.0
-beam.muxpx = 0.0
-beam.muypy = 0.0
-beam.mutpt = 0.0
 
 
 ###############################################################################