Add chromatic plasma lens model (#514)

* Examples for 3D space charge benchmarking - Modified the initial beam size in the IOTA lens benchmark example. - Added 2 benchmarks of 3D space charge for initial testing. - Add documentation for 2 benchmarks with space charge. - Add a benchmark example with space charge and periodic s-dependent focusing. - Added an s-dependent example using a Kurth beam without space charge. - Modified tolerance for IOTA lens benchmark example. Reduced tolerance to account for smaller initial beam size and improved preservation of invariants of motion. - Modified tolerances of space charge examples to allow CI tests to pass when space charge is not active. - Modified tolerance for space charge examples. These should fail unless space charge is turned on. * Update input_kurth_10nC.in Selected numerical values for amr.n_cell, lattice.nslice, and geometry.prob_relative. * Add chromatic plasma lens model. * Delete examples/kurth/input_kurth_10nC.in This is not part of this PR. * Add Python and element docs. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove unused variable, add README. * Add draft of example input files. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add test and modify input field strength units to T/m. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix duplicate test name (apochromat). * Add draft of anlysis script. * Update analysis_apochromatic_pl.py Update prediction for relative emittance growth. * Update analysis_apochromatic_pl.py Modify tolerance evaluation for relative emittance growth. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Formatting, Includes, Copy-Paste Leftovers * Modify input parameter units to allow two options. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Correct some text in comments. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Axel Huebl <[email protected]>
ECP-WarpX · Feb 2, 2024 · 4a8357d · 4a8357d
1 parent 6e21c86
commit 4a8357d
Show file tree

Hide file tree

Showing 11 changed files with 702 additions and 1 deletion.
diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst
@@ -362,6 +362,22 @@ Lattice Elements
                (default: ``1``)
             * ``<element_name>.nslice`` (``integer``) number of slices used for the application of space charge (default: ``1``)
 
+        * ``plasma_lens_chromatic`` for 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.
+           This requires these additional parameters:
+
+            * ``<element_name>.ds`` (``float``, in meters) the segment length
+            * ``<element_name>.k`` (``float``, in inverse meters squared OR in T/m) the plasma lens focusing strength
+                = (azimuthal magnetic field gradient in T/m) / (magnetic rigidity in T-m) - if units = 0
+
+             OR = azimuthal magnetic field gradient in T/m - if units = 1
+
+            * ``<element_name>.units`` (``integer``) specification of units (default: ``0``)
+            * ``<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
+            * ``<element_name>.nslice`` (``integer``) number of slices used for the application of space charge (default: ``1``)
+
         * ``sbend`` for a bending magnet. This requires these additional parameters:
 
             * ``<element_name>.ds`` (``float``, in meters) the segment length

diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst
@@ -671,6 +671,30 @@ This module provides elements for the accelerator lattice.
 
       unit specification for quad strength
 
+.. py:class:: impactx.elements.ChrPlasmaLens(ds, g, dx=0, dy=0, rotation=0, nslice=1)
+
+   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 ds: Segment length in m.
+   :param k:  focusing strength in m^(-2) (if units = 0)
+              = (azimuthal magnetic field gradient in T/m) / (rigidity in T-m)
+          OR  azimuthal magnetic field gradient in T/m (if units = 1)
+   :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]
+   :param nslice: number of slices used for the application of space charge
+
+   .. py:property:: k
+
+      plasma lens focusing strength in 1/m^2 (or T/m)
+
+   .. py:property:: units
+
+      unit specification for plasma lens focusing strength
+
 .. py:class:: impactx.elements.ChrAcc(ds, ez, bz, dx=0, dy=0, rotation=0, nslice=1)
 
    Acceleration in a uniform field Ez, with a uniform solenoidal field Bz.

diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -763,7 +763,7 @@ add_impactx_test(aperture.py
     OFF  # no plot script yet
 )
 
-# Apochromat example ##########################################################
+# Apochromat drift-quad example ##########################################################
 #
 # w/o space charge
 add_impactx_test(apochromat
@@ -781,6 +781,24 @@ add_impactx_test(apochromat.py
     OFF  # no plot script yet
 )
 
+# Apochromat drift-plasma lens example ###################################################
+#
+# w/o space charge
+add_impactx_test(apochromat_pl
+    examples/apochromatic/input_apochromatic_pl.in
+      ON  # ImpactX MPI-parallel
+      OFF  # ImpactX Python interface
+    examples/apochromatic/analysis_apochromatic_pl.py
+    OFF  # no plot script yet
+)
+add_impactx_test(apochromat_pl.py
+    examples/apochromatic/run_apochromatic_pl.py
+      OFF  # ImpactX MPI-parallel
+      ON   # ImpactX Python interface
+    examples/apochromatic/analysis_apochromatic_pl.py
+    OFF  # no plot script yet
+)
+
 # Misalignment Quad ###########################################################
 #
 # w/o space charge

diff --git a/examples/apochromatic/README.rst b/examples/apochromatic/README.rst
@@ -56,3 +56,63 @@ We run the following script to analyze correctness:
    .. literalinclude:: analysis_apochromatic.py
       :language: python3
       :caption: You can copy this file from ``examples/apochromatic/analysis_apochromatic.py``.
+
+
+.. _examples-apochromat_pl:
+
+Apochromatic Drift-Plasma Lens Beamline
+=======================================
+
+Electron beam matched to the 3rd-order apochromatic drift-plasma lens beamline appearing in Fig. 4b of:
+C. A. Lindstrom and E. Adli, "Design of general apochromatic drift-quadrupole beam lines," Phys. Rev. Accel. Beams 19, 071002 (2016).
+
+The matched Twiss parameters at entry are:
+
+* :math:`\beta_\mathrm{x} = 0.325` m
+* :math:`\alpha_\mathrm{x} = 0`
+* :math:`\beta_\mathrm{y} = 0.325` m
+* :math:`\alpha_\mathrm{y} = 0`
+
+We use a 100 GeV electron beam with an initially 6D Gaussian distribution of normalized rms emittance 1 micron and relative energy spread of 1%.
+
+The second moments of the particle distribution after the focusing beamline should coincide with the second moments of the particle distribution before the beamline, to within the level expected due to noise due to statistical sampling.
+The emittance growth due to chromatic effects remain below 1%.  In the absence of chromatic correction, the projected emittance growth is near 10%.
+
+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_apochromatic_pl.py`` or
+* ImpactX **executable** using an input file: ``impactx input_apochromatic_pl.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_apochromatic_pl.py
+          :language: python3
+          :caption: You can copy this file from ``examples/apochromatic/run_apochromatic_pl.py``.
+
+   .. tab-item:: Executable: Input File
+
+       .. literalinclude:: input_apochromatic_pl.in
+          :language: ini
+          :caption: You can copy this file from ``examples/apochromatic/input_apochromatic_pl.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_apochromatic_pl.py``
+
+   .. literalinclude:: analysis_apochromatic_pl.py
+      :language: python3
+      :caption: You can copy this file from ``examples/apochromatic/analysis_apochromatic_pl.py``.
diff --git a/examples/apochromatic/analysis_apochromatic_pl.py b/examples/apochromatic/analysis_apochromatic_pl.py
@@ -0,0 +1,129 @@
+#!/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["divergence_x"], moment=2) ** 0.5
+    sigy = moment(beam["position_y"], moment=2) ** 0.5
+    sigpy = moment(beam["divergence_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 = 100000
+assert num_particles == len(initial)
+assert num_particles == len(final)
+
+scale = (
+    (1.0 - initial.momentum_t) ** 2
+    + (initial.momentum_x) ** 2
+    + (initial.momentum_y) ** 2
+)
+xp = initial.momentum_x / np.sqrt(scale)
+initial["divergence_x"] = xp
+yp = initial.momentum_y / np.sqrt(scale)
+initial["divergence_y"] = yp
+
+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 = 3.0 * 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.288697604e-6,
+        1.288697604e-6,
+        1.0e-6,
+        5.10997388810014764e-12,
+        5.10997388810014764e-12,
+        1.0e-8,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+
+scale = (
+    (1.0 - final.momentum_t) ** 2 + (final.momentum_x) ** 2 + (final.momentum_y) ** 2
+)
+xp = final.momentum_x / np.sqrt(scale)
+final["divergence_x"] = xp
+yp = final.momentum_y / np.sqrt(scale)
+final["divergence_y"] = yp
+
+print("")
+print("Final Beam:")
+sigx, sigy, sigt, emittance_xf, emittance_yf, emittance_tf = get_moments(final)
+demittance_x = 100 * abs(emittance_xf - emittance_x) / emittance_x
+demittance_y = 100 * abs(emittance_yf - emittance_y) / emittance_y
+demittance_t = 100 * abs(emittance_tf - emittance_t) / emittance_t
+
+print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
+print(
+    f"  emittance change x (%)={demittance_x:e} emittance change y (%)={demittance_y:e} emittance change t (%)={demittance_t:e}"
+)
+
+atol = 0.0  # ignored
+rtol = 19.0 * num_particles**-0.5  # from random sampling of a smooth distribution
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+assert np.allclose(
+    [sigx, sigy, sigt],
+    [
+        1.283476e-06,
+        1.291507e-06,
+        1.0e-6,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+atol = 2.0e-3  # from random sampling of a smooth distribution
+print(f"  atol={atol} %")
+
+assert np.allclose(
+    [demittance_x, demittance_y, demittance_t],
+    [
+        0.0,
+        0.0,
+        0.0,
+    ],
+    atol=atol,
+)
diff --git a/examples/apochromatic/input_apochromatic_pl.in b/examples/apochromatic/input_apochromatic_pl.in
@@ -0,0 +1,89 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 100000
+beam.units = static
+beam.kin_energy = 100.0e3  # 100 GeV nominal energy
+beam.charge = 1.0e-9
+beam.particle = electron
+beam.distribution = gaussian
+beam.sigmaX = 1.288697604e-6
+beam.sigmaY = 1.288697604e-6
+beam.sigmaT = 1.0e-6
+beam.sigmaPx = 3.965223396e-6
+beam.sigmaPy = 3.965223396e-6
+beam.sigmaPt = 0.01  #1% energy spread
+beam.muxpx = 0.0
+beam.muypy = 0.0
+beam.mutpt = 0.0
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor dr1 q1 dr2 q2 dr2 q3 dr2 q4 dr2 q5 dr2 q6 dr2 q7 dr1 monitor
+lattice.nslice = 1
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+dr1.type = drift_chromatic
+dr1.ds = 1.0
+
+dr2.type = drift_chromatic
+dr2.ds = 2.0
+
+q1.type = plasma_lens_chromatic
+q1.ds = 0.331817852986604588
+q1.k = 996.147787384348956
+q1.units = 1
+#q1.k = 2.98636067687944129
+
+q2.type = plasma_lens_chromatic
+q2.ds = 0.176038957633108457
+q2.k = 528.485181135649785
+q2.units = 1
+#q2.k = 1.584350618697976110
+
+q3.type = plasma_lens_chromatic
+q3.ds = 1.041842576046930486
+q3.k = 3127.707468391874166
+q3.units = 1
+#q3.k = 9.37658318442237437
+
+q4.type = plasma_lens_chromatic
+q4.ds = 0.334367090894399520
+q4.k = 501.900417308233112
+q4.units = 1
+#q4.k = 1.50465190902479784
+
+q5.type = plasma_lens_chromatic
+q5.ds = 1.041842576046930486
+q5.k = 3127.707468391874166
+q5.units = 1
+#q5.k = 9.37658318442237437
+
+q6.type = plasma_lens_chromatic
+q6.ds = 0.176038957633108457
+q6.k = 528.485181135649785
+q6.units = 1
+#q6.k = 1.584350618697976110
+
+q7.type = plasma_lens_chromatic
+q7.ds = 0.331817852986604588
+q7.k = 996.147787384348956
+q7.units = 1
+#q7.k = 2.98636067687944129
+
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = false
+
+
+###############################################################################
+# Diagnostics
+###############################################################################
+diag.slice_step_diagnostics = true