Added cold expanding beam benchmark example (with space charge). (#244)

* Added cold expanding beam benchmark example (with space charge).

* Fix rst formatting

* Cheaper resolution converges 🎉

* Fix atol

* Add in CMakeLists.txt & Manual

* Analysis: Emittance Absolute to Zero

* Expanding Beam: Python Version

* Slightly Increase Tolerance

* Disable Analysis for now

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

docs/source/usage/examples.rst

@@ -11,6 +11,7 @@ This section allows you to **download input files** that correspond to different
    examples/fodo/README.rst
    examples/chicane/README.rst
    examples/cfchannel/README.rst
+   examples/expanding_beam/README.rst
    examples/kurth/README.rst
    examples/fodo_rf/README.rst
    examples/multipole/README.rst

examples/CMakeLists.txt

@@ -280,6 +280,26 @@ add_impactx_test(multipole
     OFF  # no plot script yet
 )
 
+# Expanding Beam Test #################################################################
+#
+add_impactx_test(expanding_beam
+    examples/expanding_beam/input_expanding.in
+      OFF  # ImpactX MPI-parallel
+      OFF  # ImpactX Python interface
+    OFF  # examples/expanding_beam/analysis_expanding.py
+    OFF  # no plot script yet
+)
+
+# Python: Expanding Beam Test #########################################################
+#
+add_impactx_test(expanding_beam.py
+    examples/expanding_beam/run_expanding.py
+      OFF  # ImpactX MPI-parallel
+      ON   # ImpactX Python interface
+    OFF  # examples/expanding_beam/analysis_expanding.py
+    OFF  # no plot script yet
+)
+
 # Python: Chain of Multipoles Test #########################################################
 #
 add_impactx_test(multipole.py

examples/expanding_beam/README.rst

@@ -0,0 +1,48 @@
+.. _examples-expanding:
+
+Expanding Beam in Free Space
+============================
+
+A coasting bunch expanding freely in free space under its own space charge.
+
+We use a cold (zero emittance) 250 MeV electron bunch whose
+initial distribution is a uniformly-populated 3D ball of radius R0 = 1 mm when viewed in the bunch rest
+frame.
+
+In the laboratory frame, the bunch is a uniformly-populated ellipsoid, which
+expands to twice its original 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 as a Python script (``python3 run_expanding.py``) or with an app with an input file (``impactx input_expanding.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_expanding.py
+         :language: python3
+         :caption: You can copy this file from ``examples/expanding/run_expanding.py``.
+
+   .. tab-item:: App Input File
+
+       .. literalinclude:: input_expanding.in
+          :language: ini
+          :caption: You can copy this file from ``examples/expanding/input_expanding.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_expanding.py``
+
+   .. literalinclude:: analysis_expanding.py
+      :language: python3
+      :caption: You can copy this file from ``examples/expanding/analysis_expanding.py``.

examples/expanding_beam/analysis_expanding.py

@@ -0,0 +1,124 @@
+#!/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 = 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],
+    [
+        4.4721359550e-004,
+        4.4721359550e-004,
+        9.1224186858e-007,
+        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 = 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],
+    [
+        8.9442719100e-004,
+        8.9442719100e-004,
+        1.8244837370e-006,
+    ],
+    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.in

@@ -0,0 +1,38 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 10000
+beam.units = static
+beam.energy = 250.0
+beam.charge = 1.0e-9
+beam.particle = electron
+beam.distribution = waterbag
+beam.sigmaX = 4.472135955e-4
+beam.sigmaY = 4.472135955e-4
+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
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = drift1
+lattice.nslice = 30
+
+drift1.type = drift
+drift1.ds = 6.0
+
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = true
+
+amr.n_cell = 40 40 32
+geometry.prob_relative = 1.0

examples/expanding_beam/run_expanding.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022 ImpactX contributors
+# Authors: 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", [40, 40, 32])
+
+sim = ImpactX()
+
+# set numerical parameters and IO control
+sim.set_particle_shape(2)  # B-spline order
+sim.set_space_charge(True)
+sim.dynamic_size = True
+sim.prob_relative = 1.0
+
+# beam diagnostics
+# sim.set_diagnostics(False)  # benchmarking
+sim.set_slice_step_diagnostics(False)
+
+# domain decomposition & space charge mesh
+sim.init_grids()
+
+# load a 2 GeV electron beam with an initial
+# unnormalized rms emittance of 2 nm
+energy_MeV = 250  # 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_energy_MeV(energy_MeV)
+
+#   particle bunch
+distr = distribution.Waterbag(
+    sigmaX=4.472135955e-4,
+    sigmaY=4.472135955e-4,
+    sigmaT=9.12241869e-7,
+    sigmaPx=0.0,
+    sigmaPy=0.0,
+    sigmaPt=0.0,
+)
+sim.add_particles(bunch_charge_C, distr, npart)
+
+# design the accelerator lattice
+sim.lattice.append(elements.Drift(ds=6.0, nslice=30))
+
+# run simulation
+sim.evolve()
+
+# clean shutdown
+del sim
+amrex.finalize()