Skip to content

Commit 771000b

Browse files
kyrsjopre-commit-ci[bot]cemitch99
authored
ChrPlasmaLens element: Support k<0, and fix uninitialized-variable bug in tracking of reference particle (#1030)
* Compiles? * Fixed the math, no we can defocus with k<0! * Initial add of defocusing ChrPlasmaLens * Fix use of uninitialized variable for reference particle tracking in ChrPlasmaLens * Remove forgotten debugging comment * Splelling * Fix tabs/spaces * Fix uninitialized variable giving bad results for envelope tracking with Quad when k=0 * Handle m_g=0 correctly in ChrPlasmaLens * Use right litteral type in comparisons * Fix indentation * litteral types for assignments * Rollback the changes in ConstF to what it looks in development branch; these changes will be handled in another PR. * Fix apochromatic_pl test as suggested by @cemitch99 * Chromatic effect in drift as well, as in ChrQuad and in ChrDrift * Fix time propagation for g=0 in ChrPlasmaLens * Fix time dependency for negative focusing * Start adding tests for APLs * Activate the APL test * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Refactoring * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Improve ChrPlasmaLens tests and add plot * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Test the ChrPlasmaLens, not drift * Coding conventions * Cleanup * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add focusing test to ChrPlasmaLens * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add defocusing test * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cleanup and add to CMakeLists * Fix analytical estimate, generalize run script for test * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix error in CMakeLists * Rename run_APL_ChrPlasmaLens.py -> run_APL.py * Swap reference values to analytical values * Catch error * Add README.rst for active_plasma_lens_test * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Hopefully fixing REST syntax * Update examples/active_plasma_lens/plot_APL_ChrPlasmaLens_analytical.py * Link README to examples documentation. * Update docs/source/usage/examples.rst * Apply suggestions from code review --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Co-authored-by: Chad Mitchell <ChadMitchell@lbl.gov>
1 parent 1077e91 commit 771000b

19 files changed

Lines changed: 1005 additions & 42 deletions

‎docs/source/usage/examples.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ Single Particle Dynamics
4141
examples/reversibility/README.rst
4242
examples/charge_sign/README.rst
4343
examples/symplectic_integration/README.rst
44+
examples/active_plasma_lens/README.rst
4445

4546

4647
Collective Effects

‎examples/CMakeLists.txt‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1604,6 +1604,53 @@ add_impactx_test(fodo-pals.py
16041604
OFF # no plot script yet
16051605
)
16061606

1607+
# Active Plasma Lens tests ####################################
1608+
#
1609+
# ChrPlasmaLens, k = 0 (acting like a drift)
1610+
file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/run_APL.py
1611+
DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_zero.py)
1612+
file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens.py
1613+
DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_zero.py)
1614+
file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/plot_APL_ChrPlasmaLens.py
1615+
DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_zero.py)
1616+
1617+
add_impactx_test(APL_ChrPlasmaLens_zero.py
1618+
examples/active_plasma_lens/run_APL_ChrPlasmaLens_zero.py
1619+
OFF # No MPI needed
1620+
examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_zero.py
1621+
examples/active_plasma_lens/plot_APL_ChrPlasmaLens_zero.py
1622+
)
1623+
1624+
# ChrPlasmaLens, mg = -1000T/m (focusing)
1625+
file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/run_APL.py
1626+
DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_focusing.py)
1627+
file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens.py
1628+
DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_focusing.py)
1629+
file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/plot_APL_ChrPlasmaLens.py
1630+
DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_focusing.py)
1631+
1632+
add_impactx_test(APL_ChrPlasmaLens_focusing.py
1633+
examples/active_plasma_lens/run_APL_ChrPlasmaLens_focusing.py
1634+
OFF # No MPI needed
1635+
examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_focusing.py
1636+
examples/active_plasma_lens/plot_APL_ChrPlasmaLens_focusing.py
1637+
)
1638+
1639+
# ChrPlasmaLens, mg = 1000T/m (defocusing)
1640+
file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/run_APL.py
1641+
DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_defocusing.py)
1642+
file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens.py
1643+
DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_defocusing.py)
1644+
file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/plot_APL_ChrPlasmaLens.py
1645+
DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_defocusing.py)
1646+
1647+
add_impactx_test(APL_ChrPlasmaLens_defocusing.py
1648+
examples/active_plasma_lens/run_APL_ChrPlasmaLens_defocusing.py
1649+
OFF # No MPI needed
1650+
examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_defocusing.py
1651+
examples/active_plasma_lens/plot_APL_ChrPlasmaLens_defocusing.py
1652+
)
1653+
16071654
# Exactly-solvable (non-uniform) soft-edge solenoid ##############
16081655
#
16091656
# copy PALS lattice file

‎examples/active_plasma_lens/README.rst‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
.. _examples-active_plasma_lens:
2+
3+
Active Plasma Lens
4+
==================
5+
6+
These examples demonstrate the effect of an Active Plasma Lens (APL) on the beam.
7+
The lattice contains this element and nothing else.
8+
The length of the element is 20 mm, and it can be run in no-field, focusing, and defocusing mode.
9+
10+
We use a 200 MeV electron beam with an initial normalized rms emittance of 10 um.
11+
The beam is set to have :math:`\alpha = 0` in the middle of the lens in the case of no field.
12+
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.
13+
A :math:`\sigma_{pt} = 10^{-3}` is also assumed.
14+
Before the simulation, this beam is back-propagated to the lens entry assuming zero field.
15+
16+
Run
17+
---
18+
19+
This example can be run as
20+
* ``python3 run_APL_ChrPlasmaLens_zero.py`` (no field, ``ChrPlasmaLens``, tracking)
21+
* ``python3 run_APL_ChrPlasmaLens_focusing.py`` (focusing field, ``ChrPlasmaLens``, tracking)
22+
* ``python3 run_APL_ChrPlasmaLens_defocusing.py`` (defocusing field, ``ChrPlasmaLens``, tracking)
23+
24+
These all use the library ``run_APL.py`` internally to create the simulations.
25+
26+
Analyze
27+
-------
28+
29+
We run the following scripts to analyze correctness of the output:
30+
* ``python3 analysis_APL_ChrPlasmaLens_zero.py`` (no field, ``ChrPlasmaLens``, tracking)
31+
* ``python3 analysis_APL_ChrPlasmaLens_focusing.py`` (focusing field, ``ChrPlasmaLens``, tracking)
32+
* ``python3 analysis_APL_ChrPlasmaLens_defocusing.py`` (defocusing field, ``ChrPlasmaLens``, tracking)
33+
34+
These all use the library ``analysis_APL_ChrPlasmaLens.py`` internally.
35+
36+
Visualize
37+
---------
38+
You can run the following scripts to visualize the beam evolution over time (e.g. :math:`s`):
39+
* ``python3 s_APL_ChrPlasmaLens_zero.py`` (no field, ``ChrPlasmaLens``, tracking)plot
40+
* ``python3 plot_APL_ChrPlasmaLens_focusing.py`` (focusing field, ``ChrPlasmaLens``, tracking)
41+
* ``python3 plot_APL_ChrPlasmaLens_defocusing.py`` (defocusing field, ``ChrPlasmaLens``, tracking)
42+
43+
These all use the library ``plot_APL_ChrPlasmaLens.py`` internally.
44+
45+
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‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
#!/usr/bin/env python3
2+
#
3+
# Copyright 2022-2025 ImpactX contributors
4+
# Authors: Axel Huebl, Chad Mitchell, Kyrre Sjobak
5+
# License: BSD-3-Clause-LBNL
6+
#
7+
# -*- coding: utf-8 -*-
8+
9+
import openpmd_api as io
10+
from scipy.stats import moment
11+
12+
13+
def get_moments(beam):
14+
"""Calculate standard deviations of beam position & momenta
15+
and emittance values
16+
17+
Returns
18+
-------
19+
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t
20+
"""
21+
sigx = moment(beam["position_x"], moment=2) ** 0.5 # variance -> std dev.
22+
sigpx = moment(beam["momentum_x"], moment=2) ** 0.5
23+
sigy = moment(beam["position_y"], moment=2) ** 0.5
24+
sigpy = moment(beam["momentum_y"], moment=2) ** 0.5
25+
sigt = moment(beam["position_t"], moment=2) ** 0.5
26+
sigpt = moment(beam["momentum_t"], moment=2) ** 0.5
27+
28+
epstrms = beam.cov(ddof=0)
29+
emittance_x = (sigx**2 * sigpx**2 - epstrms["position_x"]["momentum_x"] ** 2) ** 0.5
30+
emittance_y = (sigy**2 * sigpy**2 - epstrms["position_y"]["momentum_y"] ** 2) ** 0.5
31+
emittance_t = (sigt**2 * sigpt**2 - epstrms["position_t"]["momentum_t"] ** 2) ** 0.5
32+
33+
return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t)
34+
35+
36+
def get_twiss(beam):
37+
"Calculate the beam Twiss parameters from position and momenta values"
38+
39+
epstrms = beam.cov(ddof=0)
40+
41+
sigx2 = epstrms["position_x"]["position_x"]
42+
sigpx2 = epstrms["momentum_x"]["momentum_x"]
43+
emittance_x = (sigx2 * sigpx2 - epstrms["position_x"]["momentum_x"] ** 2) ** 0.5
44+
beta_x = sigx2 / emittance_x
45+
alpha_x = -epstrms["position_x"]["momentum_x"] / emittance_x
46+
47+
sigy2 = epstrms["position_y"]["position_y"]
48+
sigpy2 = epstrms["momentum_y"]["momentum_y"]
49+
emittance_y = (sigy2 * sigpy2 - epstrms["position_y"]["momentum_y"] ** 2) ** 0.5
50+
beta_y = sigy2 / emittance_y
51+
alpha_y = -epstrms["position_y"]["momentum_y"] / emittance_y
52+
53+
return (beta_x, beta_y, alpha_x, alpha_y)
54+
55+
56+
def get_beams():
57+
"Load the initial and final beam from last simulation"
58+
series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
59+
last_step = list(series.iterations)[-1]
60+
initial = series.iterations[1].particles["beam"].to_df()
61+
beam_final = series.iterations[last_step].particles["beam"]
62+
final = beam_final.to_df()
63+
64+
return (initial, beam_final, final)

‎examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_defocusing.py‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
#!/usr/bin/env python3
2+
#
3+
# Copyright 2022-2025 ImpactX contributors
4+
# Authors: Axel Huebl, Chad Mitchell, Kyrre Sjobak
5+
# License: BSD-3-Clause-LBNL
6+
#
7+
# -*- coding: utf-8 -*-
8+
9+
import numpy as np
10+
from analysis_APL_ChrPlasmaLens import get_beams, get_moments, get_twiss
11+
12+
# initial/final beam
13+
(initial, beam_final, final) = get_beams()
14+
15+
# compare number of particles
16+
num_particles = 10000
17+
assert num_particles == len(initial)
18+
assert num_particles == len(final)
19+
20+
print("Initial Beam:")
21+
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(initial)
22+
print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
23+
print(
24+
f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
25+
)
26+
27+
(betax, betay, alphax, alphay) = get_twiss(initial)
28+
print(f" betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}")
29+
30+
atol = 0.0 # ignored
31+
rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution
32+
print(f" rtol={rtol} (ignored: atol~={atol})")
33+
34+
# Compare to analytical values
35+
assert np.allclose(
36+
[sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
37+
[
38+
0.00010003246877770656,
39+
0.00010003246877770656,
40+
0.001,
41+
2.548491664266332e-08,
42+
2.548491664266332e-08,
43+
1e-06,
44+
],
45+
rtol=rtol,
46+
atol=atol,
47+
)
48+
49+
50+
print("")
51+
print("Final Beam:")
52+
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final)
53+
s_ref = beam_final.get_attribute("s_ref")
54+
gamma_ref = beam_final.get_attribute("gamma_ref")
55+
print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
56+
print(
57+
f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}\n"
58+
f" s_ref={s_ref:e} gamma_ref={gamma_ref:e}"
59+
)
60+
61+
(betax, betay, alphax, alphay) = get_twiss(final)
62+
print(f" betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}")
63+
64+
atol = 0.0 # ignored
65+
# rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution
66+
rtol = (
67+
2.9 * num_particles**-0.5
68+
) # from random sampling of a smooth distribution -- tolerance increased here
69+
print(f" rtol={rtol} (ignored: atol~={atol})")
70+
71+
# Compare to analytical values
72+
assert np.allclose(
73+
[sigx, sigy, sigt, emittance_x, emittance_y, emittance_t, s_ref, gamma_ref],
74+
[
75+
0.0001314429025974998,
76+
0.0001314429025974998,
77+
0.001,
78+
2.514662e-08,
79+
2.514662e-08,
80+
1e-06,
81+
20e-3,
82+
3.923902e02,
83+
],
84+
rtol=rtol,
85+
atol=atol,
86+
)

‎examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_focusing.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
#!/usr/bin/env python3
2+
#
3+
# Copyright 2022-2025 ImpactX contributors
4+
# Authors: Axel Huebl, Chad Mitchell, Kyrre Sjobak
5+
# License: BSD-3-Clause-LBNL
6+
#
7+
# -*- coding: utf-8 -*-
8+
9+
import numpy as np
10+
from analysis_APL_ChrPlasmaLens import get_beams, get_moments, get_twiss
11+
12+
# initial/final beam
13+
(initial, beam_final, final) = get_beams()
14+
15+
# compare number of particles
16+
num_particles = 10000
17+
assert num_particles == len(initial)
18+
assert num_particles == len(final)
19+
20+
print("Initial Beam:")
21+
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(initial)
22+
print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
23+
print(
24+
f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
25+
)
26+
27+
(betax, betay, alphax, alphay) = get_twiss(initial)
28+
print(f" betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}")
29+
30+
atol = 0.0 # ignored
31+
rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution
32+
print(f" rtol={rtol} (ignored: atol~={atol})")
33+
34+
# Compare to analytical values
35+
assert np.allclose(
36+
[sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
37+
[
38+
0.00010003246877770656,
39+
0.00010003246877770656,
40+
0.001,
41+
2.548491664266332e-08,
42+
2.548491664266332e-08,
43+
1e-06,
44+
],
45+
rtol=rtol,
46+
atol=atol,
47+
)
48+
49+
50+
print("")
51+
print("Final Beam:")
52+
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final)
53+
s_ref = beam_final.get_attribute("s_ref")
54+
gamma_ref = beam_final.get_attribute("gamma_ref")
55+
print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
56+
print(
57+
f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}\n"
58+
f" s_ref={s_ref:e} gamma_ref={gamma_ref:e}"
59+
)
60+
61+
(betax, betay, alphax, alphay) = get_twiss(final)
62+
print(f" betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}")
63+
64+
atol = 0.0 # ignored
65+
rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution
66+
print(f" rtol={rtol} (ignored: atol~={atol})")
67+
68+
# Compare to analytical values
69+
assert np.allclose(
70+
[sigx, sigy, sigt, emittance_x, emittance_y, emittance_t, s_ref, gamma_ref],
71+
[
72+
7.161196476484095e-05,
73+
7.161196476484095e-05,
74+
0.001,
75+
2.548491664266332e-08,
76+
2.548491664266332e-08,
77+
1e-06,
78+
20e-3,
79+
3.923902e02,
80+
],
81+
rtol=rtol,
82+
atol=atol,
83+
)

0 commit comments

Comments
 (0)