Diagnostics: NLL Invariants to Monitor (#565)

* Diagnostics: NLL Invariants to Monitor

The calculation of the invariants of motion H and I for the
nonlinear lens were still calculated in a slow, inefficient way using
per-particle ASCII output.

This moves the calculation and output to the beam monitor, where
it can be optionally enabled and has significantly more performant
I/O throughput using openPMD.

docs/source/usage/parameters.rst

@@ -536,6 +536,21 @@ Lattice Elements
                 openPMD `iteration encoding <https://openpmd-api.readthedocs.io/en/0.14.0/usage/concepts.html#iteration-and-series>`__: (v)ariable based, (f)ile based, (g)roup based (default)
                 variable based is an `experimental feature with ADIOS2 <https://openpmd-api.readthedocs.io/en/0.14.0/backends/adios2.html#experimental-new-adios2-schema>`__.
 
+            * ``<element_name>.nonlinear_lens_invariants`` (``boolean``, default value: ``false``)
+
+                Compute and output the invariants H and I within the nonlinear magnetic insert element (see: ``nonlinear_lens``).
+                Invariants associated with the nonlinear magnetic insert described by V. Danilov and S. Nagaitsev, PRSTAB 13, 084002 (2010), Sect. V.A.
+
+                * ``<element_name>.alpha`` (``float``, unitless) Twiss alpha of the bare linear lattice at the location of output for the nonlinear IOTA invariants H and I.
+                  Horizontal and vertical values must be equal.
+
+                * ``<element_name>.beta`` (``float``, meters) Twiss beta of the bare linear lattice at the location of output for the nonlinear IOTA invariants H and I.
+                  Horizontal and vertical values must be equal.
+
+                * ``<element_name>.tn`` (``float``, unitless) dimensionless strength of the IOTA nonlinear magnetic insert element used for computing H and I.
+
+                * ``<element_name>.cn`` (``float``, meters^(1/2)) scale factor of the IOTA nonlinear magnetic insert element used for computing H and I.
+
         * ``line`` a sub-lattice (line) of elements to append to the lattice.
 
             * ``<element_name>.elements`` (``list of strings``) optional (default: no elements)
@@ -694,26 +709,6 @@ 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.
 
-.. _running-cpp-parameters-diagnostics-reduced:
-
-Reduced Diagnostics
-^^^^^^^^^^^^^^^^^^^
-
-Reduced diagnostics allow the user to compute some reduced quantity (invariants of motion, particle temperature, max of a field, ...) and write a small amount of data to text files.
-Reduced diagnostics are run *in situ* with the simulation.
-
-Diagnostics related to integrable optics in the IOTA nonlinear magnetic insert element:
-
-* ``diag.alpha`` (``float``, unitless) Twiss alpha of the bare linear lattice at the location of output for the nonlinear IOTA invariants H and I.
-  Horizontal and vertical values must be equal.
-
-* ``diag.beta`` (``float``, meters) Twiss beta of the bare linear lattice at the location of output for the nonlinear IOTA invariants H and I.
-  Horizontal and vertical values must be equal.
-
-* ``diag.tn`` (``float``, unitless) dimensionless strength of the IOTA nonlinear magnetic insert element used for computing H and I.
-
-* ``diag.cn`` (``float``, meters^(1/2)) scale factor of the IOTA nonlinear magnetic insert element used for computing H and I.
-
 
 .. _running-cpp-parameters-diagnostics-insitu:
 

docs/source/usage/python.rst

@@ -116,15 +116,6 @@ General
       The minimum number of digits (default: ``6``) used for the step
       number appended to the diagnostic file names.
 
-   .. py:method:: set_diag_iota_invariants(alpha, beta, tn, cn)
-
-      Set the parameters of the IOTA nonlinear lens invariants diagnostics.
-
-      :param float alpha: Twiss alpha
-      :param float beta: Twiss beta (m)
-      :param float tn: dimensionless strength of the nonlinear insert
-      :param float cn: scale parameter of the nonlinear insert (m^[1/2])
-
    .. py:property:: particle_lost_diagnostics_backend
 
       Diagnostics for particles lost in apertures.
@@ -610,6 +601,32 @@ This module provides elements for the accelerator lattice.
    :param backend: I/O backend, e.g., ``bp``, ``h5``, ``json``
    :param encoding: openPMD iteration encoding: (v)ariable based, (f)ile based, (g)roup based (default)
 
+   .. py:property:: name
+
+      name of the series
+
+   .. py:property:: nonlinear_lens_invariants
+
+      Compute and output the invariants H and I within the nonlinear magnetic insert element
+
+   .. py:property:: alpha
+
+      Twiss alpha of the bare linear lattice at the location of output for the nonlinear IOTA invariants H and I.
+      Horizontal and vertical values must be equal.
+
+   .. py:property:: beta
+
+      Twiss beta (in meters) of the bare linear lattice at the location of output for the nonlinear IOTA invariants H and I.
+      Horizontal and vertical values must be equal.
+
+   .. py:property:: tn
+
+      Dimensionless strength of the IOTA nonlinear magnetic insert element used for computing H and I.
+
+   .. py:property:: cn
+
+      Scale factor (in meters^(1/2)) of the IOTA nonlinear magnetic insert element used for computing H and I.
+
 .. py:class:: impactx.elements.Programmable
 
    A programmable beam optics element.

examples/iota_lattice/analysis_iotalattice_sdep.py

@@ -4,10 +4,8 @@
 # Authors: Axel Huebl, Chad Mitchell
 # License: BSD-3-Clause-LBNL
 #
-import glob
-
 import numpy as np
-import pandas as pd
+import openpmd_api as io
 from scipy.stats import moment
 
 
@@ -25,26 +23,11 @@ def get_moments(beam):
     return (meanH, sigH, meanI, sigI)
 
 
-def read_all_files(file_pattern):
-    """Read in all CSV files from each MPI rank (and potentially OpenMP
-    thread). Concatenate into one Pandas dataframe.
-    Returns
-    -------
-    pandas.DataFrame
-    """
-    return pd.concat(
-        (
-            pd.read_csv(filename, delimiter=r"\s+")
-            for filename in glob.glob(file_pattern)
-        ),
-        axis=0,
-        ignore_index=True,
-    ).set_index("id")
-
-
 # initial/final beam
-initial = read_all_files("diags/nonlinear_lens_invariants_000000.*")
-final = read_all_files("diags/nonlinear_lens_invariants_final.*")
+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

examples/iota_lattice/input_iotalattice.in

@@ -214,6 +214,7 @@ qe3.k = -6.69148177
 # Beam Monitor: Diagnostics
 monitor.type = beam_monitor
 monitor.backend = h5
+monitor.nonlinear_lens_invariants = true
 
 
 ###############################################################################

examples/iota_lattice/input_iotalattice_sdep.in

@@ -271,6 +271,11 @@ qe3.k = -6.69148177
 # Beam Monitor: Diagnostics
 monitor.type = beam_monitor
 monitor.backend = h5
+monitor.nonlinear_lens_invariants = true
+monitor.alpha = 1.376381920471173
+monitor.beta = 1.892632003628881
+monitor.tn = 0.4
+monitor.cn = 0.01
 
 
 ###############################################################################
@@ -284,8 +289,3 @@ algo.space_charge = false
 # Diagnostics
 ###############################################################################
 diag.slice_step_diagnostics = true
-
-diag.alpha = 1.376381920471173
-diag.beta = 1.892632003628881
-diag.tn = 0.4
-diag.cn = 0.01

examples/iota_lattice/run_iotalattice_sdep.py

@@ -21,11 +21,6 @@
 # domain decomposition & space charge mesh
 sim.init_grids()
 
-# diagnostics: IOTA nonlinear lens invariants calculation
-sim.set_diag_iota_invariants(
-    alpha=1.376381920471173, beta=1.892632003628881, tn=0.4, cn=0.01
-)
-
 # init particle beam
 energy_MeV = 2.5
 bunch_charge_C = 1.0e-9  # used with space charge
@@ -51,6 +46,11 @@
 
 # add beam diagnostics
 monitor = elements.BeamMonitor("monitor", backend="h5")
+monitor.nonlinear_lens_invariants = True
+monitor.alpha = 1.376381920471173
+monitor.beta = 1.892632003628881
+monitor.tn = 0.4
+monitor.cn = 0.01
 
 # init accelerator lattice
 ns = 10  # number of slices per ds in the element

examples/iota_lens/analysis_iotalens.py

@@ -4,10 +4,9 @@
 # Authors: Axel Huebl, Chad Mitchell
 # License: BSD-3-Clause-LBNL
 #
-import glob
 
 import numpy as np
-import pandas as pd
+import openpmd_api as io
 from scipy.stats import moment
 
 
@@ -25,26 +24,11 @@ def get_moments(beam):
     return (meanH, sigH, meanI, sigI)
 
 
-def read_all_files(file_pattern):
-    """Read in all CSV files from each MPI rank (and potentially OpenMP
-    thread). Concatenate into one Pandas dataframe.
-    Returns
-    -------
-    pandas.DataFrame
-    """
-    return pd.concat(
-        (
-            pd.read_csv(filename, delimiter=r"\s+")
-            for filename in glob.glob(file_pattern)
-        ),
-        axis=0,
-        ignore_index=True,
-    ).set_index("id")
-
-
 # initial/final beam
-initial = read_all_files("diags/nonlinear_lens_invariants_000000.*")
-final = read_all_files("diags/nonlinear_lens_invariants_final.*")
+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

examples/iota_lens/analysis_iotalens_sdep.py

@@ -4,10 +4,8 @@
 # Authors: Axel Huebl, Chad Mitchell
 # License: BSD-3-Clause-LBNL
 #
-import glob
-
 import numpy as np
-import pandas as pd
+import openpmd_api as io
 from scipy.stats import moment
 
 
@@ -25,26 +23,11 @@ def get_moments(beam):
     return (meanH, sigH, meanI, sigI)
 
 
-def read_all_files(file_pattern):
-    """Read in all CSV files from each MPI rank (and potentially OpenMP
-    thread). Concatenate into one Pandas dataframe.
-    Returns
-    -------
-    pandas.DataFrame
-    """
-    return pd.concat(
-        (
-            pd.read_csv(filename, delimiter=r"\s+")
-            for filename in glob.glob(file_pattern)
-        ),
-        axis=0,
-        ignore_index=True,
-    ).set_index("id")
-
-
 # initial/final beam
-initial = read_all_files("diags/nonlinear_lens_invariants_000000.*")
-final = read_all_files("diags/nonlinear_lens_invariants_final.*")
+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

examples/iota_lens/input_iotalens.in

@@ -21,7 +21,7 @@ beam.mutpt = 0.0
 ###############################################################################
 # Beamline: lattice elements and segments
 ###############################################################################
-lattice.elements = const_end nllens foclens const_end
+lattice.elements = monitor const_end nllens foclens const_end monitor
 
 foclens.type = line
 foclens.elements = const nllens
@@ -43,18 +43,17 @@ const.kx = 1.0
 const.ky = 1.0
 const.kt = 1.0e-12
 
+monitor.type = beam_monitor
+monitor.backend = h5
+monitor.nonlinear_lens_invariants = true
+monitor.alpha = 0.0
+monitor.beta = 1.0
+monitor.tn = 0.4
+monitor.cn = 0.01
+
 
 ###############################################################################
 # Algorithms
 ###############################################################################
 algo.particle_shape = 2
 algo.space_charge = false
-
-
-###############################################################################
-# Diagnostics
-###############################################################################
-diag.alpha = 0.0
-diag.beta = 1.0
-diag.tn = 0.4
-diag.cn = 0.01

examples/iota_lens/input_iotalens_sdep.in

@@ -92,18 +92,15 @@ const.nslice = 1
 # Beam Monitor: Diagnostics
 monitor.type = beam_monitor
 monitor.backend = h5
+monitor.nonlinear_lens_invariants = true
+monitor.alpha = 1.376381920471173
+monitor.beta = 1.892632003628881
+monitor.tn = 0.4
+monitor.cn = 0.01
+
 
 ###############################################################################
 # Algorithms
 ###############################################################################
 algo.particle_shape = 2
 algo.space_charge = false
-
-
-###############################################################################
-# Diagnostics
-###############################################################################
-diag.alpha = 1.376381920471173
-diag.beta = 1.892632003628881
-diag.tn = 0.4
-diag.cn = 0.01

examples/iota_lens/run_iotalens.py

@@ -19,9 +19,6 @@
 # domain decomposition & space charge mesh
 sim.init_grids()
 
-# diagnostics: IOTA nonlinear lens invariants calculation
-sim.set_diag_iota_invariants(alpha=0.0, beta=1.0, tn=0.4, cn=0.01)
-
 # load a 2.5 MeV proton beam
 kin_energy_MeV = 2.5  # reference energy
 bunch_charge_C = 1.0e-9  # used with space charge
@@ -42,13 +39,25 @@
 )
 sim.add_particles(bunch_charge_C, distr, npart)
 
+# add beam diagnostics
+monitor = elements.BeamMonitor("monitor", backend="h5")
+monitor.nonlinear_lens_invariants = True
+monitor.alpha = 0.0
+monitor.beta = 1.0
+monitor.tn = 0.4
+monitor.cn = 0.01
+
 # design the accelerator lattice
 constEnd = elements.ConstF(ds=0.05, kx=1.0, ky=1.0, kt=1.0e-12)
 nllens = elements.NonlinearLens(knll=4.0e-6, cnll=0.01)
 const = elements.ConstF(ds=0.1, kx=1.0, ky=1.0, kt=1.0e-12)
 
 num_lenses = 18
-nllens_lattice = [constEnd] + [nllens, const] * (num_lenses - 1) + [nllens, constEnd]
+nllens_lattice = (
+    [monitor, constEnd]
+    + [nllens, const] * (num_lenses - 1)
+    + [nllens, constEnd, monitor]
+)
 
 # add elements to the lattice segment
 sim.lattice.extend(nllens_lattice)