Add thin nonlinear plasma lens with transverse taper

docs/source/usage/parameters.rst

@@ -516,6 +516,22 @@ Lattice Elements
             * ``<element_name>.dy`` (``float``, in meters) vertical translation error
             * ``<element_name>.rotation`` (``float``, in degrees) rotation error in the transverse plane
 
+        * ``tapered_pl`` for a thin nonlinear (tapered) plasma lens.
+          This requires these additional parameters:
+
+            * ``<element_name>.k`` (``float``, in inverse meters OR in T) the integrated plasma lens focusing strength
+                = (length in m) * (azimuthal magnetic field gradient in T/m) / (magnetic rigidity in T-m) - if units = 0
+
+             OR = (length in m) * (azimuthal magnetic field gradient in T/m) - if units = 1
+
+            * ``<element_name>.units`` (``integer``) specification of units (default: ``0``)
+            * ``<element_name>.taper`` (``float``, in 1/meters) horizontal taper parameter
+                = 1 / (target horizontal dispersion in m)
+
+            * ``<element_name>.dx`` (``float``, in meters) horizontal translation error
+            * ``<element_name>.dy`` (``float``, in meters) vertical translation error
+            * ``<element_name>.rotation`` (``float``, in degrees) rotation error in the transverse plane
+
         * ``beam_monitor`` a beam monitor, writing all beam particles at fixed ``s`` to openPMD files.
           If the same element name is used multiple times, then an output series is created with multiple outputs.
 

docs/source/usage/python.rst

@@ -916,6 +916,35 @@ This module provides elements for the accelerator lattice.
 
    * G. Ripken and F. Schmidt, Thin-Lens Formalism for Tracking, CERN/SL/95-12 (AP), 1995.
 
+.. py:class:: impactx.elements.TaperedPL(k, taper, units, dx=0, dy=0, rotation=0)
+
+   An active cylindrically symmetric plasma lens, with chromatic effects included.
+   The Hamiltonian is expanded through second order in the transverse variables
+   (x,px,y,py), with the exact pt dependence retained.
+
+   :param k:  integrated focusing strength in m^(-1) (if units = 0)
+              = (length in m) * (azimuthal magnetic field gradient in T/m) / (rigidity in T-m)
+          OR  integrated focusing strength in T (if units = 1)
+              = (length in m) * (azimuthal magnetic field gradient in T/m)
+   :param taper: horizontal taper parameter in m^(-1)
+              = 1 / (target horizontal dispersion in m)
+   :param units: specification of units for plasma lens focusing strength
+   :param dx: horizontal translation error in m
+   :param dy: vertical translation error in m
+   :param rotation: rotation error in the transverse plane [degrees]
+
+   .. py:property:: k
+
+      integrated plasma lens focusing strength in 1/m (or T)
+
+   .. py:property:: taper
+
+      horizontal taper parameter in 1/m
+
+   .. py:property:: units
+
+      unit specification for plasma lens focusing strength
+
 
 Coordinate Transformation
 -------------------------

examples/CMakeLists.txt

@@ -917,3 +917,21 @@ add_impactx_test(initialize_from_array.py
     examples/initialize_from_array/analyze_from_array.py
     examples/initialize_from_array/visualize_from_array.py
 )
+
+# Achromatic Spectrometer ########################################################
+#
+# w/o space charge
+add_impactx_test(spectrometer
+    examples/achromatic_spectrometer/input_spectrometer.in
+      ON   # ImpactX MPI-parallel
+      OFF  # ImpactX Python interface
+    examples/achromatic_spectrometer/analysis_spectrometer.py
+    OFF  # no plot script yet
+)
+add_impactx_test(spectrometer.py
+    examples/achromatic_spectrometer/run_spectrometer.py
+      OFF  # ImpactX MPI-parallel
+      ON   # ImpactX Python interface
+    examples/achromatic_spectrometer/analysis_spectrometer.py
+    OFF  # no plot script yet
+)

examples/achromatic_spectrometer/README.rst

@@ -0,0 +1,78 @@
+.. _examples-fodo:
+
+FODO Cell
+=========
+
+Stable FODO cell with a zero-current phase advance of 67.8 degrees.
+
+The matched Twiss parameters at entry are:
+
+* :math:`\beta_\mathrm{x} = 2.82161941` m
+* :math:`\alpha_\mathrm{x} = -1.59050035`
+* :math:`\beta_\mathrm{y} = 2.82161941` m
+* :math:`\alpha_\mathrm{y} = 1.59050035`
+
+We use a 2 GeV electron beam with initial unnormalized rms emittance of 2 nm.
+
+The second moments of the particle distribution after the FODO cell should coincide with the second moments of the particle distribution before the FODO cell, to within the level expected due to noise due to statistical sampling.
+
+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_fodo.py`` or
+* ImpactX **executable** using an input file: ``impactx input_fodo.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_fodo.py
+          :language: python3
+          :caption: You can copy this file from ``examples/fodo/run_fodo.py``.
+
+   .. tab-item:: Executable: Input File
+
+       .. literalinclude:: input_fodo.in
+          :language: ini
+          :caption: You can copy this file from ``examples/fodo/input_fodo.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_fodo.py``
+
+   .. literalinclude:: analysis_fodo.py
+      :language: python3
+      :caption: You can copy this file from ``examples/fodo/analysis_fodo.py``.
+
+
+Visualize
+---------
+
+You can run the following script to visualize the beam evolution over time:
+
+.. dropdown:: Script ``plot_fodo.py``
+
+   .. literalinclude:: plot_fodo.py
+      :language: python3
+      :caption: You can copy this file from ``examples/fodo/plot_fodo.py``.
+
+.. figure:: https://user-images.githubusercontent.com/1353258/180287840-8561f6fd-278f-4856-abd8-04fbdb78c8ff.png
+   :alt: focusing, defocusing and preserved emittane in our FODO cell benchmark.
+
+   FODO transversal beam width and emittance evolution
+
+.. figure:: https://user-images.githubusercontent.com/1353258/180287845-eb0210a7-2500-4aa9-844c-67fb094329d3.png
+   :alt: focusing, defocusing and phase space rotation in our FODO cell benchmark.
+
+   FODO transversal beam width and phase space evolution

examples/achromatic_spectrometer/analysis_spectrometer.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,
+)

examples/achromatic_spectrometer/input_spectrometer.in

@@ -0,0 +1,63 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 10000
+beam.units = static
+beam.kin_energy = 1.0e3
+beam.charge = 1.0e-9
+beam.particle = electron
+beam.distribution = waterbag
+beam.lambdaX = 3.162277660e-6
+beam.lambdaY = 3.162277660e-6
+beam.lambdaT = 1.0e-3
+beam.lambdaPx = 3.16227766017e-4
+beam.lambdaPy = 3.16227766017e-4
+beam.lambdaPt = 2.0e-2
+beam.muxpx = 0.0
+beam.muypy = 0.0
+beam.mutpt = 0.0
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor bend1 plasma_lens drift1 monitor
+lattice.nslice = 25
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+bend1.type = sbend_exact
+bend1.ds = 1.0
+bend1.phi = 10.0
+bend1.B = 0.0
+
+plasma_lens.type = line
+plasma_lens.elements = drend pl dr pl dr pl dr pl dr pl dr pl dr pl dr pl dr pl dr pl drend
+
+drend.type = drift
+drend.ds = 0.001
+
+pl.type = tapered_plasma_lens
+pl.k = 0.2  #focal length 0.5 m
+pl.taper = 11.488289081903567
+pl.units = 0
+
+dr.type = drift
+dr.ds = 0.002
+
+drift1.type = drift
+drift1.ds = 1.0
+
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = false
+
+
+###############################################################################
+# Diagnostics
+###############################################################################
+diag.slice_step_diagnostics = true