Add beam eigenemittances to reduced diagnostics. (ECP-WarpX#702)

* Add eigenemittance calculation.

docs/source/dataanalysis/dataanalysis.rst

@@ -84,7 +84,7 @@ The code writes out the values in an ASCII file prefixed ``reduced_beam_characte
 * ``sig_px``, ``sig_py``, ``sig_pt``
     Standard deviation of the particle momentum deviations (energy difference for ``pt``) normalized by the magnitude of the reference particle momentum (unit: dimensionless)
 * ``emittance_x``, ``emittance_y``, ``emittance_t``
-    Normalized rms beam emittance (unit: meter)
+    Unnormalized rms beam emittances (unit: meter)
 * ``alpha_x``, ``alpha_y``, ``alpha_t``
     Courant-Snyder (Twiss) alpha (unit: dimensionless).  Transverse Twiss functions are calculated after removing correlations with particle energy.
 * ``beta_x``, ``beta_y``, ``beta_t``
@@ -93,6 +93,11 @@ 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_xn``, ``emittance_yn``, ``emittance_tn``
+    Normalized rms beam emittances (unit: meter)
+* ``emittance_1``, ``emittance_2``, ``emittance_3``
+    Normalized rms beam eigenemittances (aka mode emittances) (unit: meter)
+    These three diagnostics are written optionally if the flag eigenemittances = True.
 * ``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
 ------------------

docs/source/usage/parameters.rst

@@ -688,8 +688,8 @@ Diagnostics and output
   This option is ignored for the openPMD output elements (remove them from the lattice to disable).
 
 * ``diag.slice_step_diagnostics`` (``boolean``, optional, default: ``false``)
-  By default, diagnostics is performed at the beginning and end of the simulation.
-  Enabling this flag will write diagnostics every step and slice step
+  By default, diagnostics are computed and written at the beginning and end of the simulation.
+  Enabling this flag will write diagnostics at every step and slice step.
 
 * ``diag.file_min_digits`` (``integer``, optional, default: ``6``)
     The minimum number of digits used for the step number appended to the diagnostic file names.
@@ -699,6 +699,10 @@ Diagnostics and output
   Diagnostics for particles lost in apertures, stored as ``diags/openPMD/particles_lost.*`` at the end of the simulation.
   See the ``beam_monitor`` element for backend values.
 
+* ``diag.eigenemittances`` (``boolean``, optional, default: ``false``)
+  If this flag is enabled, the 3 eigenemittances of the 6D beam distribution are computed and written as diagnostics.
+  This flag is disabled by default to reduce computational cost.
+
 
 .. _running-cpp-parameters-diagnostics-insitu:
 

docs/source/usage/python.rst

@@ -156,6 +156,13 @@ Collective Effects & Overall Simulation Parameters
       Diagnostics for particles lost in apertures.
       See the ``BeamMonitor`` element for backend values.
 
+   .. py:property:: eigenemittances
+
+      Enable (``True``) or disable (``False``) output of eigenemittances at every slice step in elements  (default: ``False``).
+
+      If this flag is enabled, the 3 eigenemittances of the 6D beam distribution are computed and written as diagnostics.
+      This flag is disabled by default to reduce computational cost.
+
    .. py:method:: init_grids()
 
       Initialize AMReX blocks/grids for domain decomposition & space charge mesh.

examples/CMakeLists.txt

@@ -976,3 +976,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