Add short linac segment example.

docs/source/usage/examples.rst

@@ -104,6 +104,7 @@ Virtual Test Stands
    examples/pytorch_surrogate_model/README.rst
    examples/apochromatic/README.rst
    examples/fodo_tune/README.rst
+   examples/linac_segment/README.rst
 
 
 Unit tests

examples/CMakeLists.txt

@@ -1040,3 +1040,19 @@ add_impactx_test(rotation-xy.py
     examples/rotation/analysis_rotation_xy.py
     OFF  # no plot script yet
 )
+
+# PIP-II linac segment #########################################################
+#
+# with space charge
+add_impactx_test(linac-segment
+    examples/linac_segment/input_linac_segment.in
+    ON   # ImpactX MPI-parallel
+    examples/linac_segment/analysis_linac_segment.py
+    OFF  # no plot script yet
+)
+add_impactx_test(linac-segment.py
+    examples/linac_segment/run_linac_segment.py
+    OFF  # ImpactX MPI-parallel
+    examples/linac_segment/analysis_linac_segment.py
+    OFF  # no plot script yet
+)

examples/linac_segment/README.rst

@@ -0,0 +1,50 @@
+.. _examples-linac-segment:
+
+PIP-II Linac Segment
+=======================
+
+This example illustrates a small segment of the PIP-II linac, which includes the last RF cavity in the MEBT section (a quarter-wave resonator cavity for longitudinal bunching) and the first 
+RF cavity in the succeeding cryomodule (a superconducting half-wave resonator for acceleration).
+
+For simplicity, aperture restrictions have been removed, and there are no active kickers.  RF cavity and solenoid field maps are represented using the Fourier coefficients of the longitudinal fields on-axis.
+
+The nominal beam current is 4.845 mA (29.83 pC at 162.5 MHz), which assumes ~3% beam loss relative to the 5 mA current entering the MEBT.
+
+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_linac_segment.py`` or
+* ImpactX **executable** using an input file: ``impactx input_linac_segment.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.
+
+.. tab-set::
+
+   .. tab-item:: Python: Script
+
+       .. literalinclude:: run_linac_segment.py
+          :language: python3
+          :caption: You can copy this file from ``examples/linac_segment/run_linac_segment.py``.
+
+   .. tab-item:: Executable: Input File
+
+       .. literalinclude:: input_linac_segment.in
+          :language: ini
+          :caption: You can copy this file from ``examples/linac_segment/input_linac_segment.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_linac_segment.py``
+
+   .. literalinclude:: analysis_linac_segment.py
+      :language: python3
+      :caption: You can copy this file from ``examples/linac_segment/analysis_linac_segment.py``.

examples/linac_segment/analysis_linac_segment.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 = 2.2e12 * 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.5451170454175073e-005,
+        7.5441588239210947e-005,
+        9.9775878164077539e-004,
+        1.9959540393751392e-009,
+        2.0175015289132990e-009,
+        2.0013820193294972e-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.2e12 * 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.4790118496224206e-005,
+        7.5357525169680140e-005,
+        9.9775879288128088e-004,
+        1.9959539836392703e-009,
+        2.0175014668882125e-009,
+        2.0013820380883801e-006,
+    ],
+    rtol=rtol,
+    atol=atol,
+)