Add beam eigenemittances to reduced diagnostics.

docs/source/dataanalysis/dataanalysis.rst

@@ -93,6 +93,8 @@ The code writes out the values in an ASCII file prefixed ``reduced_beam_characte
     Horizontal and vertical dispersion (unit: meter)
 * ``dispersion_px``, ``dispersion_py``
     Derivative of horizontal and vertical dispersion (unit: dimensionless)
+* ``emittance_1``, ``emittance_2``, ``emittance_3``
+    Normalized rms beam eigenemittances (aka mode emittances) (unit: meter)
 * ``charge``
     Total beam charge (unit: Coulomb)
 

docs/source/usage/examples.rst

@@ -33,6 +33,8 @@ Single Particle Dynamics
    examples/achromatic_spectrometer/README.rst
    examples/fodo_programmable/README.rst
    examples/dogleg/README.rst
+   examples/coupled_optics/README.rst
+
 
 Collective Effects
 ------------------

examples/CMakeLists.txt

@@ -959,3 +959,19 @@ add_impactx_test(dogleg.py
     examples/dogleg/analysis_dogleg.py
     OFF  # no plot script yet
 )
+
+# Coupled Optics #############################################################
+#
+# w/o space charge
+add_impactx_test(coupled-optics
+    examples/coupled_optics/input_coupled_optics.in
+    ON   # ImpactX MPI-parallel
+    examples/coupled_optics/analysis_coupled_optics.py
+    OFF  # no plot script yet
+)
+add_impactx_test(coupled-optics.py
+    examples/coupled_optics/run_coupled_optics.py
+    OFF  # ImpactX MPI-parallel
+    examples/coupled_optics/analysis_coupled_optics.py
+    OFF  # no plot script yet
+)

examples/coupled_optics/README.rst

@@ -0,0 +1,51 @@
+.. _examples-coupled-optics:
+
+Coupled Optics
+==============
+
+This is a lattice illustrating fully coupled 6D transport.  It is obtained from the example "dogleg" by adding a solenoid after the first bending dipole.
+The solenoid is identical to that found in the example "solenoid".
+
+Its primary purpose is to benchmark the calculation of the three beam eigenemittances (mode emittances).
+
+In this test, the initial and final values of :math:`\lambda_x`, :math:`\lambda_y`, :math:`\lambda_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must
+agree with nominal values.
+
+In addition, the initial and final values of :math:`emittance_1`, :math:`emittance_2`, :math:`emittance_3` must coincide.
+
+
+Run
+---
+
+This example can be run **either** as:
+
+* **Python** script: ``python3 run_coupled_optics.py`` or
+* ImpactX **executable** using an input file: ``impactx input_coupled_optics.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_coupled_optics.py
+          :language: python3
+          :caption: You can copy this file from ``examples/coupled_optics/run_coupled_optics.py``.
+
+   .. tab-item:: Executable: Input File
+
+       .. literalinclude:: input_coupled_optics.in
+          :language: ini
+          :caption: You can copy this file from ``examples/coupled_optics/input_coupled_optics.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_coupled_optics.py``
+
+   .. literalinclude:: analysis_coupled_optics.py
+      :language: python3
+      :caption: You can copy this file from ``examples/coupled_optics/analysis_coupled_optics.py``.

examples/coupled_optics/analysis_coupled_optics.py

@@ -0,0 +1,142 @@
+#!/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)
+
+
+def get_eigenemittances(openpmd_beam):
+    """Return eigenemittances from an openPMD particle species
+
+    Returns
+    -------
+    emittance_1, emittance_2, emittance_3
+    """
+    emittance_1 = openpmd_beam.get_attribute("emittance_1")
+    emittance_2 = openpmd_beam.get_attribute("emittance_2")
+    emittance_3 = openpmd_beam.get_attribute("emittance_3")
+
+    return (emittance_1, emittance_2, emittance_3)
+
+
+# initial/final beam
+series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
+last_step = list(series.iterations)[-1]
+initial_beam = series.iterations[1].particles["beam"]
+initial = initial_beam.to_df()
+final_beam = series.iterations[last_step].particles["beam"]
+final = final_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.2 * 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],
+    [
+        6.4214719960819659e-005,
+        3.6603372435649773e-005,
+        1.9955175623579313e-004,
+        1.0198263116327677e-010,
+        1.0308359092878036e-010,
+        4.0035161705244885e-010,
+    ],
+    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],
+    [
+        1.922660e-03,
+        2.166654e-05,
+        1.101353e-04,
+        8.561046e-09,
+        1.020439e-10,
+        8.569865e-09,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+print("")
+print("Initial eigenemittances:")
+emittance_1i, emittance_2i, emittance_3i = get_eigenemittances(initial_beam)
+print(
+    f"  emittance_1={emittance_1i:e} emittance_2={emittance_2i:e} emittance_3={emittance_3i:e}"
+)
+
+print("")
+print("Final eigenemittances:")
+emittance_1f, emittance_2f, emittance_3f = get_eigenemittances(final_beam)
+print(
+    f"  emittance_1={emittance_1f:e} emittance_2={emittance_2f:e} emittance_3={emittance_3f:e}"
+)
+
+atol = 0.0  # ignored
+rtol = 3.5 * num_particles**-0.5  # from random sampling of a smooth distribution
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+assert np.allclose(
+    [emittance_1f, emittance_2f, emittance_3f],
+    [
+        emittance_1i,
+        emittance_2i,
+        emittance_3i,
+    ],
+    rtol=rtol,
+    atol=atol,
+)

examples/coupled_optics/input_coupled_optics.in

@@ -0,0 +1,72 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 10000
+beam.units = static
+beam.kin_energy = 5.0e3
+beam.charge = 1.0e-9
+beam.particle = electron
+beam.distribution = waterbag
+beam.lambdaX = 2.2951017632e-5
+beam.lambdaY = 1.3084093142e-5
+beam.lambdaT = 5.5555553e-8
+beam.lambdaPx = 1.598353425e-6
+beam.lambdaPy = 2.803697378e-6
+beam.lambdaPt = 2.000000000e-6
+beam.muxpx = 0.933345606203060
+beam.muypy = 0.933345606203060
+beam.mutpt = 0.999999961419755
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor sbend1 dipedge1 sol drift1 dipedge2 sbend2 drift2 monitor
+lattice.nslice = 25
+
+sbend1.type = sbend
+sbend1.ds = 0.500194828041958       # projected length 0.5 m, angle 2.77 deg
+sbend1.rc = -10.3462283686195526
+
+drift1.type = drift
+drift1.ds = 5.0058489435  # projected length 5 m
+
+sbend2.type = sbend
+sbend2.ds = 0.500194828041958       # projected length 0.5 m, angle 2.77 deg
+sbend2.rc = 10.3462283686195526
+
+drift2.type = drift
+drift2.ds = 0.5
+
+dipedge1.type = dipedge   # dipole edge focusing
+dipedge1.psi = -0.048345620280243
+dipedge1.rc = -10.3462283686195526
+dipedge1.g = 0.0
+dipedge1.K2 = 0.0
+
+dipedge2.type = dipedge
+dipedge2.psi = 0.048345620280243
+dipedge2.rc = 10.3462283686195526
+dipedge2.g = 0.0
+dipedge2.K2 = 0.0
+
+sol.type = solenoid
+sol.ds = 3.820395
+sol.ks = 0.8223219329893234
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = false
+
+
+###############################################################################
+# Diagnostics
+###############################################################################
+diag.slice_step_diagnostics = true
+diag.eigenemittances = true

examples/coupled_optics/run_coupled_optics.py

@@ -0,0 +1,81 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Marco Garten, Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+from impactx import ImpactX, 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
+sim.eigenemittances = True
+
+# domain decomposition & space charge mesh
+sim.init_grids()
+
+# load a 5 GeV electron beam with an initial
+# normalized transverse rms emittance of 1 um
+kin_energy_MeV = 5.0e3  # 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_kin_energy_MeV(kin_energy_MeV)
+
+#   particle bunch
+distr = distribution.Waterbag(
+    lambdaX=2.2951017632e-5,
+    lambdaY=1.3084093142e-5,
+    lambdaT=5.5555553e-8,
+    lambdaPx=1.598353425e-6,
+    lambdaPy=2.803697378e-6,
+    lambdaPt=2.000000000e-6,
+    muxpx=0.933345606203060,
+    muypy=0.933345606203060,
+    mutpt=0.999999961419755,
+)
+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
+rc = 10.3462283686195526  # bend radius (meters)
+psi = 0.048345620280243  # pole face rotation angle (radians)
+lb = 0.500194828041958  # bend arc length (meters)
+
+# Drift elements
+dr1 = elements.Drift(ds=5.0058489435, nslice=ns)
+dr2 = elements.Drift(ds=0.5, nslice=ns)
+
+# Bend elements
+sbend1 = elements.Sbend(ds=lb, rc=-rc, nslice=ns)
+sbend2 = elements.Sbend(ds=lb, rc=rc, nslice=ns)
+
+# Dipole Edge Focusing elements
+dipedge1 = elements.DipEdge(psi=-psi, rc=-rc, g=0.0, K2=0.0)
+dipedge2 = elements.DipEdge(psi=psi, rc=rc, g=0.0, K2=0.0)
+
+# Solenoid element
+sol = elements.Sol(ds=3.820395, ks=0.8223219329893234)
+
+lattice_coupled = [sbend1, dipedge1, sol, dr1, dipedge2, sbend2, dr2]
+
+sim.lattice.append(monitor)
+sim.lattice.extend(lattice_coupled)
+sim.lattice.append(monitor)
+
+# run simulation
+sim.evolve()
+
+# clean shutdown
+sim.finalize()

src/particles/diagnostics/CMakeLists.txt

@@ -2,4 +2,5 @@ target_sources(lib
   PRIVATE
     ReducedBeamCharacteristics.cpp
     DiagnosticOutput.cpp
+    EmittanceInvariants.cpp
 )