Add thin nonlinear plasma lens with transverse taper

docs/source/usage/examples.rst

@@ -30,6 +30,7 @@ Single Particle Dynamics
    examples/thin_dipole/README.rst
    examples/aperture/README.rst
    examples/iota_lens/README.rst
+   examples/achromatic_spectrometer/README.rst
 
 
 Space Charge

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,33 @@ 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)
+
+   A thin nonlinear plasma lens with transverse taper.
+
+   :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,54 @@
+.. _examples-achromatic-spectrometer:
+
+Achromatic Spectrometer
+=======================
+
+A spectrometer beamline using a bending dipole.
+A transversely-tapered plasma lens is added for chromatic correction.
+The tapered plasma lens design is based on:
+
+C. A. Lindstrom, presentation at the EuroNNAc Special Topics Workshop 2022, `slides <https://agenda.infn.it/event/28376/contributions/178724/attachments/96899/133588/Lindstr%C3%B8m,%20EuroNNAc%20workshop,%2022%20Sep%202022.pdf>`__
+"Solutions and challenges for a multi-stage plasma accelerator",
+
+https://agenda.infn.it/event/28376/contributions/178724/attachments/96899/133588/Lindstr%C3%B8m,%20EuroNNAc%20workshop,%2022%20Sep%202022.pdf
+
+We use a 1 GeV electron beam with initial normalized rms emittance of 2 microns and 2% rms relative energy spread.
+
+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_spectrometer.py`` or
+* ImpactX **executable** using an input file: ``impactx input_spectrometer.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_spectrometer.py
+          :language: python3
+          :caption: You can copy this file from ``examples/achromatic_spectrometer/run_spectrometer.py``.
+
+   .. tab-item:: Executable: Input File
+
+       .. literalinclude:: input_spectrometer.in
+          :language: ini
+          :caption: You can copy this file from ``examples/achromatic_spectrometer/input_spectrometer.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_spectrometer.py``
+
+   .. literalinclude:: analysis_spectrometer.py
+      :language: python3
+      :caption: You can copy this file from ``examples/achromatic_spectrometer/analysis_spectrometer.py``.

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.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],
+    [
+        3.164175e-06,
+        3.160807e-06,
+        1.001072e-03,
+        9.992214e-10,
+        9.985766e-10,
+        2.004245e-05,
+    ],
+    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.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],
+    [
+        1.723892e-03,
+        7.252619e-06,
+        1.006631e-03,
+        5.548917e-07,
+        2.225222e-09,
+        2.004943e-05,
+    ],
+    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

examples/achromatic_spectrometer/run_spectrometer.py

@@ -0,0 +1,84 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: 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
+
+# domain decomposition & space charge mesh
+sim.init_grids()
+
+# load a 1 GeV electron beam
+kin_energy_MeV = 1.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=3.162277660e-6,
+    lambdaY=3.162277660e-6,
+    lambdaT=1.0e-3,
+    lambdaPx=3.16227766017e-4,
+    lambdaPy=3.16227766017e-4,
+    lambdaPt=2.0e-2,
+    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
+
+# specify thick tapered plasma lens element
+lens_length = 0.02  # length in m
+num_lenses = 10
+focal_length = 0.5  # focal length in m
+dtaper = 11.488289081903567  # 1/(horizontal dispersion in m)
+ds = lens_length / num_lenses
+dk = 1.0 / (focal_length * num_lenses)
+
+# drifts appearing the drift-kick sequence
+ds_half = ds / 2.0
+dr = elements.Drift(ds=ds_half, nslice=ns)
+
+# define the lens segments
+thick_lens = []
+for _ in range(0, num_lenses):
+    pl = elements.TaperedPL(k=dk, taper=dtaper, units=0)
+    segment = [dr, pl, dr]
+    thick_lens.extend(segment)
+
+bend = elements.ExactSbend(ds=1.0, phi=10.0, B=0.0, nslice=ns)
+drift = elements.Drift(ds=1.0, nslice=ns)
+
+# specify the lattice sequence
+sim.lattice.append(monitor)
+sim.lattice.append(bend)
+sim.lattice.extend(thick_lens)
+sim.lattice.append(drift)
+sim.lattice.append(monitor)
+
+# run simulation
+sim.evolve()
+
+# clean shutdown
+sim.finalize()