Apochromatic focusing example (#487)

* 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.

* Update input_kurth_10nC.in

Selected numerical values for amr.n_cell, lattice.nslice, and geometry.prob_relative.

* Add draft of input files.

* Delete examples/kurth/input_kurth_10nC.in

Not part of this PR.

* Comment Formatting

* Merge fix.

* Reduce energy spread to 1%.

* Update input_apochromatic.in

Change to a Gaussian distribution.

* Update run_apochromatic.py

Change to a gaussian distribution.

* Fix merge

* Add to CMakeLists and docs.

* Update CMakeLists.txt

Correct CMakeLists specification.

* Update CMakeLists.txt

Correct misspelling.

* Update run_apochromatic.py

Update Python to 10^5 particles

* Update run_apochromatic.py

Python distribution type should be Gaussian.

* Formatting: Title Underline

docs/source/usage/examples.rst

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

examples/CMakeLists.txt

@@ -711,3 +711,21 @@ add_impactx_test(aperture.py
     examples/aperture/analysis_aperture.py
     OFF  # no plot script yet
 )
+
+# Apochromat example ########################################################
+#
+# w/o space charge
+add_impactx_test(apochromat
+    examples/apochromatic/input_apochromatic.in
+      ON  # ImpactX MPI-parallel
+      OFF  # ImpactX Python interface
+    examples/apochromatic/analysis_apochromatic.py
+    OFF  # no plot script yet
+)
+add_impactx_test(apochromat.py
+    examples/apochromatic/run_apochromatic.py
+      OFF  # ImpactX MPI-parallel
+      ON   # ImpactX Python interface
+    examples/apochromatic/analysis_apochromatic.py
+    OFF  # no plot script yet
+)

examples/apochromatic/README.rst

@@ -0,0 +1,58 @@
+.. _examples-apochromat:
+
+Apochromatic Drift-Quadrupole Beamline
+======================================
+
+Electron beam matched to the 1st-order apochromatic drift-quadrupole beamline appearing in Fig. 4a of:
+C. A. Lindstrom and E. Adli, "Design of general apochromatic drift-quadrupole beam lines," Phys. Rev. Accel. Beams 19, 071002 (2016).
+
+The matched Twiss parameters at entry are:
+
+* :math:`\beta_\mathrm{x} = 0.325` m
+* :math:`\alpha_\mathrm{x} = 0`
+* :math:`\beta_\mathrm{y} = 0.325` m
+* :math:`\alpha_\mathrm{y} = 0`
+
+We use a 100 GeV electron beam with an initially 6D Gaussian distribution of normalized rms emittance 1 micron and relative energy spread of 1%.
+
+The second moments of the particle distribution after the focusing beamline should coincide with the second moments of the particle distribution before the beamline, to within the level expected due to noise due to statistical sampling.
+The emittance growth due to chromatic effects remain below 1%.  In the absence of chromatic correction, the projected emittance growth is near 10%.
+
+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_apochromatic.py`` or
+* ImpactX **executable** using an input file: ``impactx input_apochromatic.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_apochromatic.py
+          :language: python3
+          :caption: You can copy this file from ``examples/apochromatic/run_apochromatic.py``.
+
+   .. tab-item:: Executable: Input File
+
+       .. literalinclude:: input_apochromatic.in
+          :language: ini
+          :caption: You can copy this file from ``examples/apochromatic/input_apochromatic.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_apochromatic.py``
+
+   .. literalinclude:: analysis_apochromatic.py
+      :language: python3
+      :caption: You can copy this file from ``examples/apochromatic/analysis_apochromatic.py``.

examples/apochromatic/analysis_apochromatic.py

@@ -0,0 +1,126 @@
+#!/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["divergence_x"], moment=2) ** 0.5
+    sigy = moment(beam["position_y"], moment=2) ** 0.5
+    sigpy = moment(beam["divergence_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 = 100000
+assert num_particles == len(initial)
+assert num_particles == len(final)
+
+scale = (
+    (1.0 - initial.momentum_t) ** 2
+    + (initial.momentum_x) ** 2
+    + (initial.momentum_y) ** 2
+)
+xp = initial.momentum_x / np.sqrt(scale)
+initial["divergence_x"] = xp
+yp = initial.momentum_y / np.sqrt(scale)
+initial["divergence_y"] = yp
+
+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 = 3.0 * 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.288697604e-6,
+        1.288697604e-6,
+        1.0e-6,
+        5.10997388810014764e-12,
+        5.10997388810014764e-12,
+        1.0e-8,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+
+scale = (
+    (1.0 - final.momentum_t) ** 2 + (final.momentum_x) ** 2 + (final.momentum_y) ** 2
+)
+xp = final.momentum_x / np.sqrt(scale)
+final["divergence_x"] = xp
+yp = final.momentum_y / np.sqrt(scale)
+final["divergence_y"] = yp
+
+print("")
+print("Final Beam:")
+sigx, sigy, sigt, emittance_xf, emittance_yf, emittance_tf = get_moments(final)
+demittance_x = 100 * (emittance_xf - emittance_x) / emittance_x
+demittance_y = 100 * (emittance_yf - emittance_y) / emittance_y
+demittance_t = 100 * (emittance_tf - emittance_t) / emittance_t
+
+print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
+print(
+    f"  emittance change x (%)={demittance_x:e} emittance change y (%)={demittance_y:e} emittance change t (%)={demittance_t:e}"
+)
+
+atol = 0.0  # ignored
+rtol = 15.0 * num_particles**-0.5  # from random sampling of a smooth distribution
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+assert np.allclose(
+    [sigx, sigy, sigt, demittance_x, demittance_y, emittance_t],
+    [
+        1.245e-6,
+        1.245e-6,
+        1.0e-6,
+        0.94,
+        0.94,
+        1.0e-8,
+    ],
+    rtol=rtol,
+    atol=atol,
+)

examples/apochromatic/input_apochromatic.in

@@ -0,0 +1,77 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 100000
+beam.units = static
+beam.kin_energy = 100.0e3  # 100 GeV nominal energy
+beam.charge = 1.0e-9
+beam.particle = electron
+beam.distribution = gaussian
+beam.sigmaX = 1.288697604e-6
+beam.sigmaY = 1.288697604e-6
+beam.sigmaT = 1.0e-6
+beam.sigmaPx = 3.965223396e-6
+beam.sigmaPy = 3.965223396e-6
+beam.sigmaPt = 0.01  #1% energy spread
+beam.muxpx = 0.0
+beam.muypy = 0.0
+beam.mutpt = 0.0
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor dr1 q1 q2 q3 dr2 q4 q5 dr2 q6 q7 q8 dr1 monitor
+lattice.nslice = 1
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+dr1.type = drift_chromatic
+dr1.ds = 1.0
+
+dr2.type = drift_chromatic
+dr2.ds = 10.0
+
+q1.type = quad_chromatic
+q1.ds = 1.2258333333
+q1.k = 0.5884
+
+q2.type = quad_chromatic
+q2.ds = 1.5677083333
+q2.k = -0.7525
+
+q3.type = quad_chromatic
+q3.ds = 1.205625
+q3.k = 0.5787
+
+q4.type = quad_chromatic
+q4.ds = 1.2502083333
+q4.k = -0.6001
+
+q5.type = quad_chromatic
+q5.ds = 1.2502083333
+q5.k = 0.6001
+
+q6.type = quad_chromatic
+q6.ds = 1.205625
+q6.k = -0.5787
+
+q7.type = quad_chromatic
+q7.ds = 1.5677083333
+q7.k = 0.7525
+
+q8.type = quad_chromatic
+q8.ds = 1.2258333333
+q8.k = -0.5884
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = false
+
+
+###############################################################################
+# Diagnostics
+###############################################################################
+diag.slice_step_diagnostics = true

examples/apochromatic/run_apochromatic.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+import amrex.space3d as amr
+from impactx import ImpactX, RefPart, distribution, elements
+
+sim = ImpactX()
+
+# set numerical parameters and IO control
+sim.particle_shape = 2  # B-spline order
+sim.space_charge = False
+# sim.diagnostics = False  # benchmarking
+sim.slice_step_diagnostics = True
+
+# domain decomposition & space charge mesh
+sim.init_grids()
+
+# load a 2 GeV electron beam with an initial
+# unnormalized rms emittance of 2 nm
+kin_energy_MeV = 100.0e3  # reference energy
+bunch_charge_C = 1.0e-9  # used with space charge
+npart = 100000  # number of macro particles
+
+#   reference particle
+ref = sim.particle_container().ref_particle()
+ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV)
+
+#   particle bunch
+distr = distribution.Gaussian(
+    sigmaX=1.288697604e-6,
+    sigmaY=1.288697604e-6,
+    sigmaT=1.0e-6,
+    sigmaPx=3.965223396e-6,
+    sigmaPy=3.965223396e-6,
+    sigmaPt=0.01,  # 1% energy spread
+    muxpx=0.0,
+    muypy=0.0,
+    mutpt=0.0,
+)
+sim.add_particles(bunch_charge_C, distr, npart)
+
+# add beam diagnostics
+monitor = elements.BeamMonitor("monitor", backend="h5")
+
+# design the accelerator lattice)
+ns = 25  # number of slices per ds in the element
+
+# Drift elements
+dr1 = elements.ChrDrift(ds=1.0, nslice=ns)
+dr2 = elements.ChrDrift(ds=10.0, nslice=ns)
+
+# Quad elements
+q1 = elements.ChrQuad(ds=1.2258333333, k=0.5884, nslice=ns)
+q2 = elements.ChrQuad(ds=1.5677083333, k=-0.7525, nslice=ns)
+q3 = elements.ChrQuad(ds=1.205625, k=0.5787, nslice=ns)
+q4 = elements.ChrQuad(ds=1.2502083333, k=-0.6001, nslice=ns)
+q5 = elements.ChrQuad(ds=1.2502083333, k=0.6001, nslice=ns)
+q6 = elements.ChrQuad(ds=1.205625, k=-0.5787, nslice=ns)
+q7 = elements.ChrQuad(ds=1.5677083333, k=0.7525, nslice=ns)
+q8 = elements.ChrQuad(ds=1.2258333333, k=-0.5884, nslice=ns)
+
+lattice_line = [monitor, dr1, q1, q2, q3, dr2, q4, q5, dr2, q6, q7, q8, dr1, monitor]
+
+# define the lattice
+sim.lattice.extend(lattice_line)
+
+# run simulation
+sim.evolve()
+
+# clean shutdown
+del sim
+amr.finalize()