ChrPlasmaLens element: Support k<0, and fix uninitialized-variable bug in tracking of reference particle

examples/CMakeLists.txt

@@ -1590,6 +1590,53 @@ add_impactx_test(fodo-pals.py
     OFF  # no plot script yet
 )
 
+# Active Plasma Lens tests ####################################
+#
+# ChrPlasmaLens, k = 0 (acting like a drift)
+file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/run_APL.py
+     DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_zero.py)
+file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens.py
+  DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_zero.py)
+file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/plot_APL_ChrPlasmaLens.py
+  DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_zero.py)
+
+add_impactx_test(APL_ChrPlasmaLens_zero.py
+  examples/active_plasma_lens/run_APL_ChrPlasmaLens_zero.py
+  OFF # No MPI needed
+  examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_zero.py
+  examples/active_plasma_lens/plot_APL_ChrPlasmaLens_zero.py
+)
+
+# ChrPlasmaLens, mg = -1000T/m (focusing)
+file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/run_APL.py
+     DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_focusing.py)
+file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens.py
+  DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_focusing.py)
+file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/plot_APL_ChrPlasmaLens.py
+  DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_focusing.py)
+
+add_impactx_test(APL_ChrPlasmaLens_focusing.py
+  examples/active_plasma_lens/run_APL_ChrPlasmaLens_focusing.py
+  OFF # No MPI needed
+  examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_focusing.py
+  examples/active_plasma_lens/plot_APL_ChrPlasmaLens_focusing.py
+)
+
+# ChrPlasmaLens, mg = 1000T/m (defocusing)
+file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/run_APL.py
+     DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_defocusing.py)
+file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens.py
+  DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_defocusing.py)
+file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/plot_APL_ChrPlasmaLens.py
+  DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_defocusing.py)
+
+add_impactx_test(APL_ChrPlasmaLens_defocusing.py
+  examples/active_plasma_lens/run_APL_ChrPlasmaLens_defocusing.py
+  OFF # No MPI needed
+  examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_defocusing.py
+  examples/active_plasma_lens/plot_APL_ChrPlasmaLens_defocusing.py
+  )
+
 # Exactly-solvable (non-uniform) soft-edge solenoid ##############
 #
 # copy PALS lattice file

examples/active_plasma_lens/README.rst

@@ -0,0 +1,45 @@
+.. _examples-active_plasma_lens:
+
+Active Plasma Lens
+==================
+
+These examples demonstrates the of the effect of an Active Plasma Lens (APL) on the beam.
+The lattice contains this element and nothing else.
+The length of the element is 20 mm, and it can be run in no-field, focusing, and defocusing mode.
+
+We use a 200 MeV electron beam with an initial normalized rms emittance of 10 um.
+The beam is set to have :math:`\alpha = 0` in the middle of the lens in the case of no field.
+The beam size in the middle of the lens is set to 10 µm for the no-field examples (in order to have a strongly parabolic :math:`\beta`-function within the lens), and 100 µm for the focusing and defocusing examples.
+A :math:`\sigma_{pt} = 10^{-3}` is also assumed.
+Before the simulations, this beam is back-propagate to the lens entry asuming no field.
+
+Run
+---
+
+This example can be run as
+* ``python3 run_APL_ChrPlasmaLens_zero.py`` (no field, ``ChrPlasmaLens``, tracking)
+* ``python3 run_APL_ChrPlasmaLens_focusing.py`` (focusing field, ``ChrPlasmaLens``, tracking)
+* ``python3 run_APL_ChrPlasmaLens_defocusing.py`` (defocusing field, ``ChrPlasmaLens``, tracking)
+
+These all use the library ``run_APL.py`` internally to create the simulations.
+
+Analyze
+-------
+
+We run the following scripts to analyze correctness of the output:
+* ``python3 analysis_APL_ChrPlasmaLens_zero.py`` (no field, ``ChrPlasmaLens``, tracking)
+* ``python3 analysis_APL_ChrPlasmaLens_focusing.py`` (focusing field, ``ChrPlasmaLens``, tracking)
+* ``python3 analysis_APL_ChrPlasmaLens_defocusing.py`` (defocusing field, ``ChrPlasmaLens``, tracking)
+
+These all use the library ``analysis_APL_ChrPlasmaLens.py`` internally.
+
+Visualize
+---------
+You can run the following scripts to visualize the beam evolution over time (e.g. :math:`s`):
+* ``python3 s_APL_ChrPlasmaLens_zero.py`` (no field, ``ChrPlasmaLens``, tracking)plot
+* ``python3 plot_APL_ChrPlasmaLens_focusing.py`` (focusing field, ``ChrPlasmaLens``, tracking)
+* ``python3 plot_APL_ChrPlasmaLens_defocusing.py`` (defocusing field, ``ChrPlasmaLens``, tracking)
+
+These all use the library ``plot_APL_ChrPlasmaLens.py`` internally.
+
+Additionally, it is also possible to run ``python3 plot_APL_ChrPlasmaLens_analytical.py``, which plots the expected Twiss :math:`\alpha` and :math:`\beta` functions at the end of the lens as a function of the lens gradient. This uses the stand-alone Twiss propagation function ``analytic_final_estimate()`` from ``run_APL.py``.

examples/active_plasma_lens/analysis_APL_ChrPlasmaLens.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2025 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell, Kyrre Sjobak
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+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_twiss(beam):
+    "Calculate the beam Twiss parameters from position and momenta values"
+
+    epstrms = beam.cov(ddof=0)
+
+    sigx2 = epstrms["position_x"]["position_x"]
+    sigpx2 = epstrms["momentum_x"]["momentum_x"]
+    emittance_x = (sigx2 * sigpx2 - epstrms["position_x"]["momentum_x"] ** 2) ** 0.5
+    beta_x = sigx2 / emittance_x
+    alpha_x = -epstrms["position_x"]["momentum_x"] / emittance_x
+
+    sigy2 = epstrms["position_y"]["position_y"]
+    sigpy2 = epstrms["momentum_y"]["momentum_y"]
+    emittance_y = (sigy2 * sigpy2 - epstrms["position_y"]["momentum_y"] ** 2) ** 0.5
+    beta_y = sigy2 / emittance_y
+    alpha_y = -epstrms["position_y"]["momentum_y"] / emittance_y
+
+    return (beta_x, beta_y, alpha_x, alpha_y)
+
+
+def get_beams():
+    "Load the initial and final beam from last simulation"
+    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()
+    beam_final = series.iterations[last_step].particles["beam"]
+    final = beam_final.to_df()
+
+    return (initial, beam_final, final)

examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_defocusing.py

@@ -0,0 +1,86 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2025 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell, Kyrre Sjobak
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+import numpy as np
+from analysis_APL_ChrPlasmaLens import get_beams, get_moments, get_twiss
+
+# initial/final beam
+(initial, beam_final, final) = get_beams()
+
+# 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}"
+)
+
+(betax, betay, alphax, alphay) = get_twiss(initial)
+print(f"  betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}")
+
+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})")
+
+# Compare to analytical values
+assert np.allclose(
+    [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
+    [
+        0.00010003246877770656,
+        0.00010003246877770656,
+        0.001,
+        2.548491664266332e-08,
+        2.548491664266332e-08,
+        1e-06,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+
+print("")
+print("Final Beam:")
+sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final)
+s_ref = beam_final.get_attribute("s_ref")
+gamma_ref = beam_final.get_attribute("gamma_ref")
+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}\n"
+    f"  s_ref={s_ref:e} gamma_ref={gamma_ref:e}"
+)
+
+(betax, betay, alphax, alphay) = get_twiss(final)
+print(f"  betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}")
+
+atol = 0.0  # ignored
+# rtol = 2.2 * num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = (
+    2.9 * num_particles**-0.5
+)  # from random sampling of a smooth distribution -- tolerance increased here
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+# Compare to analytical values
+assert np.allclose(
+    [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t, s_ref, gamma_ref],
+    [
+        0.0001314429025974998,
+        0.0001314429025974998,
+        0.001,
+        2.514662e-08,
+        2.514662e-08,
+        1e-06,
+        20e-3,
+        3.923902e02,
+    ],
+    rtol=rtol,
+    atol=atol,
+)

examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_focusing.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2025 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell, Kyrre Sjobak
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+import numpy as np
+from analysis_APL_ChrPlasmaLens import get_beams, get_moments, get_twiss
+
+# initial/final beam
+(initial, beam_final, final) = get_beams()
+
+# 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}"
+)
+
+(betax, betay, alphax, alphay) = get_twiss(initial)
+print(f"  betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}")
+
+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})")
+
+# Compare to analytical values
+assert np.allclose(
+    [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
+    [
+        0.00010003246877770656,
+        0.00010003246877770656,
+        0.001,
+        2.548491664266332e-08,
+        2.548491664266332e-08,
+        1e-06,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+
+print("")
+print("Final Beam:")
+sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final)
+s_ref = beam_final.get_attribute("s_ref")
+gamma_ref = beam_final.get_attribute("gamma_ref")
+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}\n"
+    f"  s_ref={s_ref:e} gamma_ref={gamma_ref:e}"
+)
+
+(betax, betay, alphax, alphay) = get_twiss(final)
+print(f"  betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}")
+
+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})")
+
+# Compare to analytical values
+assert np.allclose(
+    [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t, s_ref, gamma_ref],
+    [
+        7.161196476484095e-05,
+        7.161196476484095e-05,
+        0.001,
+        2.548491664266332e-08,
+        2.548491664266332e-08,
+        1e-06,
+        20e-3,
+        3.923902e02,
+    ],
+    rtol=rtol,
+    atol=atol,
+)