Add fnal booster input files

docs/source/usage/examples.rst

@@ -35,6 +35,7 @@ This section allows you to **download input files** that correspond to different
    examples/pytorch_surrogate_model/README.rst
    examples/apochromatic/README.rst
    examples/fodo_tune/README.rst
+   examples/fnal_booster/README.rst
 
 
 Unit tests

examples/CMakeLists.txt

@@ -855,6 +855,21 @@ add_impactx_test(alignment.py
     OFF  # no plot script yet
 )
 
+# Fermilab Booster thin-kick model ################################################
+#
+# w/o space charge
+# copy lattice file & parser
+file(COPY ${ImpactX_SOURCE_DIR}/examples/fnal_booster/fermi-booster-madx-3_no-sc.sxf
+     DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/booster.py)
+add_impactx_test(booster.py
+    examples/fnal_booster/fnal_booster_no-sc.py
+      OFF  # ImpactX MPI-parallel
+      ON   # ImpactX Python interface
+    examples/fnal_booster/analysis_booster.py
+    OFF  # no plot script yet
+)
+label_impactx_test(booster.py slow)
+
 # Tune Calculation ###########################################################
 #
 # w/o space charge

examples/fnal_booster/README.rst

@@ -0,0 +1,45 @@
+.. _examples-booster:
+
+FNAL Booster Model
+===================
+
+This example is based on a thin-kick model of the Fermilab Booster ring originally created in MAD-X,
+and contributed by F. Schmidt (CERN).  The lattice file is expressed in SXF format, and it is parsed and
+executed in ImpactX using a Python script.  A complete set of input and output files can be found on Zenodo:
+
+https://doi.org/10.5281/zenodo.11645618
+
+We use a proton beam with total energy of 1.338272088 GeV and an initial normalized rms emittance of 1.6 microns.
+
+The second moments of the particle distribution after one turn should coincide with the initial second moments of the particle distribution, to within the level expected due to noise due to statistical sampling.
+
+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:
+
+* **Python** script: ``python3 fnal_booster_no-sc.py`` or
+
+For `MPI-parallel <https://www.mpi-forum.org>`__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system.
+
+.. tab-set::
+
+   .. tab-item:: Python: Script
+
+       .. literalinclude:: fnal_booster_no-sc.py
+          :language: python3
+          :caption: You can copy this file from ``examples/fnal_booster/fnal_booster_no-sc.py``.
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_booster.py``
+
+   .. literalinclude:: analysis_booster.py
+      :language: python3
+      :caption: You can copy this file from ``examples/fnal_booster/analysis_booster.py``.

examples/fnal_booster/analysis_booster.py

@@ -0,0 +1,98 @@
+#!/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],
+    [
+        7.5519124163421e-3,
+        2.7653008349980e-3,
+        1.0000000000000e-6,
+        1.5731883856314145e-6,
+        1.5731883856314145e-6,
+        1.000000000000000e-22,
+    ],
+    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],
+    [
+        7.5519124163421e-3,
+        2.7653008349980e-3,
+        1.0000000000000e-6,
+        1.5731883856314145e-6,
+        1.5731883856314145e-6,
+        1.000000000000000e-22,
+    ],
+    rtol=rtol,
+    atol=atol,
+)