From 4a8357d7f664e78ce634bdd1ddd91946291d94e4 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Fri, 2 Feb 2024 14:10:36 -0800 Subject: [PATCH] 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 --- docs/source/usage/parameters.rst | 16 ++ docs/source/usage/python.rst | 24 ++ examples/CMakeLists.txt | 20 +- examples/apochromatic/README.rst | 60 +++++ .../apochromatic/analysis_apochromatic_pl.py | 129 +++++++++++ .../apochromatic/input_apochromatic_pl.in | 89 ++++++++ examples/apochromatic/run_apochromatic_pl.py | 111 ++++++++++ src/initialization/InitElement.cpp | 11 + src/particles/elements/All.H | 2 + src/particles/elements/ChrPlasmaLens.H | 208 ++++++++++++++++++ src/python/elements.cpp | 33 +++ 11 files changed, 702 insertions(+), 1 deletion(-) create mode 100755 examples/apochromatic/analysis_apochromatic_pl.py create mode 100644 examples/apochromatic/input_apochromatic_pl.in create mode 100644 examples/apochromatic/run_apochromatic_pl.py create mode 100644 src/particles/elements/ChrPlasmaLens.H diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index d1933d537..1f299319b 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -362,6 +362,22 @@ Lattice Elements (default: ``1``) * ``.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: + + * ``.ds`` (``float``, in meters) the segment length + * ``.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 + + * ``.units`` (``integer``) specification of units (default: ``0``) + * ``.dx`` (``float``, in meters) horizontal translation error + * ``.dy`` (``float``, in meters) vertical translation error + * ``.rotation`` (``float``, in degrees) rotation error in the transverse plane + * ``.nslice`` (``integer``) number of slices used for the application of space charge (default: ``1``) + * ``sbend`` for a bending magnet. This requires these additional parameters: * ``.ds`` (``float``, in meters) the segment length diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 4792cdd2d..f852c9601 100644 --- 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 index ef9fbe0bd..3b68c817e 100644 --- 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 index b8798693d..e37f2a4ff 100644 --- 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 `__ 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 new file mode 100755 index 000000000..39acd7b8b --- /dev/null +++ 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 new file mode 100644 index 000000000..8588501f9 --- /dev/null +++ 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 diff --git a/examples/apochromatic/run_apochromatic_pl.py b/examples/apochromatic/run_apochromatic_pl.py new file mode 100644 index 000000000..95b5ee011 --- /dev/null +++ b/examples/apochromatic/run_apochromatic_pl.py @@ -0,0 +1,111 @@ +#!/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, RefPart, 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 2 GeV electron beam with an initial +# unnormalized rms emittance of 2 nm +kin_energy_MeV = 100.0e3 # reference energy +bunch_charge_C = 1.0e-9 # used with space charge +npart = 100000 # 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.Gaussian( + sigmaX=1.288697604e-6, + sigmaY=1.288697604e-6, + sigmaT=1.0e-6, + sigmaPx=3.965223396e-6, + sigmaPy=3.965223396e-6, + sigmaPt=0.01, # 1% energy spread + 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 + +# Drift elements +dr1 = elements.ChrDrift(ds=1.0, nslice=ns) +dr2 = elements.ChrDrift(ds=2.0, nslice=ns) + +# Plasma lens elements +q1 = elements.ChrPlasmaLens( + ds=0.331817852986604588, k=996.147787384348956, units=1, nslice=ns +) +q2 = elements.ChrPlasmaLens( + ds=0.176038957633108457, k=528.485181135649785, units=1, nslice=ns +) +q3 = elements.ChrPlasmaLens( + ds=1.041842576046930486, k=3127.707468391874166, units=1, nslice=ns +) +q4 = elements.ChrPlasmaLens( + ds=0.334367090894399520, k=501.900417308233112, units=1, nslice=ns +) +q5 = elements.ChrPlasmaLens( + ds=1.041842576046930486, k=3127.707468391874166, units=1, nslice=ns +) +q6 = elements.ChrPlasmaLens( + ds=0.176038957633108457, k=528.485181135649785, units=1, nslice=ns +) +q7 = elements.ChrPlasmaLens( + ds=0.331817852986604588, k=996.147787384348956, units=1, nslice=ns +) + +# q1 = elements.ChrPlasmaLens(ds=0.331817852986604588, k=2.98636067687944129, units=0, nslice=ns) +# q2 = elements.ChrPlasmaLens(ds=0.176038957633108457, k=1.584350618697976110, units=0, nslice=ns) +# q3 = elements.ChrPlasmaLens(ds=1.041842576046930486, k=9.37658318442237437, units=0, nslice=ns) +# q4 = elements.ChrPlasmaLens(ds=0.334367090894399520, k=1.50465190902479784, units=0, nslice=ns) +# q5 = elements.ChrPlasmaLens(ds=1.041842576046930486, k=9.37658318442237437, units=0, nslice=ns) +# q6 = elements.ChrPlasmaLens(ds=0.176038957633108457, k=1.584350618697976110, units=0, nslice=ns) +# q7 = elements.ChrPlasmaLens(ds=0.331817852986604588, k=2.98636067687944129, units=0, nslice=ns) + +lattice_line = [ + monitor, + dr1, + q1, + dr2, + q2, + dr2, + q3, + dr2, + q4, + dr2, + q5, + dr2, + q6, + dr2, + q7, + dr1, + monitor, +] + +# define the lattice +sim.lattice.extend(lattice_line) + +# run simulation +sim.evolve() diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index 05439f814..6e949b219 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -294,6 +294,17 @@ namespace detail pp_element.queryAdd("units", units); m_lattice.emplace_back( ChrQuad(ds, k, units, a["dx"], a["dy"], a["rotation_degree"], nslice) ); + } else if (element_type == "plasma_lens_chromatic") + { + auto a = detail::query_alignment(pp_element); + auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default); + + amrex::ParticleReal k; + int units = 0; + pp_element.get("k", k); + pp_element.queryAdd("units", units); + + m_lattice.emplace_back( ChrPlasmaLens(ds, k, units, a["dx"], a["dy"], a["rotation_degree"], nslice) ); } else if (element_type == "drift_exact") { auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default); diff --git a/src/particles/elements/All.H b/src/particles/elements/All.H index 44e6bc7e7..b9ad27ae0 100644 --- a/src/particles/elements/All.H +++ b/src/particles/elements/All.H @@ -14,6 +14,7 @@ #include "Buncher.H" #include "CFbend.H" #include "ChrDrift.H" +#include "ChrPlasmaLens.H" #include "ChrQuad.H" #include "ChrUniformAcc.H" #include "ConstF.H" @@ -49,6 +50,7 @@ namespace impactx CFbend, ChrAcc, ChrDrift, + ChrPlasmaLens, ChrQuad, ConstF, diagnostics::BeamMonitor, diff --git a/src/particles/elements/ChrPlasmaLens.H b/src/particles/elements/ChrPlasmaLens.H new file mode 100644 index 000000000..9d245ce94 --- /dev/null +++ b/src/particles/elements/ChrPlasmaLens.H @@ -0,0 +1,208 @@ +/* Copyright 2022-2023 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Chad Mitchell, Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_CHRPLASMALENS_H +#define IMPACTX_CHRPLASMALENS_H + +#include "particles/ImpactXParticleContainer.H" +#include "mixin/alignment.H" +#include "mixin/beamoptic.H" +#include "mixin/thick.H" +#include "mixin/nofinalize.H" + +#include +#include + +#include + + +namespace impactx +{ + struct ChrPlasmaLens + : public elements::BeamOptic, + public elements::Thick, + public elements::Alignment, + public elements::NoFinalize + { + static constexpr auto name = "ChrPlasmaLens"; + using PType = ImpactXParticleContainer::ParticleType; + + /** An active cylindrically symmetric plasma lens with chromatic focusing + * + * 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 plasma lens focusing strength in m^(-2) + * = (azimuthal field gradient in T/m) / (rigidity in T-m) + * OR azimuthal magnetic field gradient in T/m (k > 0) + * @param unit Unit specification + * unit = 0 focusing strength in m^(-2) + * unit = 1 focusing strength in T/m + * @param dx horizontal translation error in m + * @param dy vertical translation error in m + * @param rotation_degree rotation error in the transverse plane [degrees] + * @param nslice number of slices used for the application of space charge + */ + ChrPlasmaLens ( + amrex::ParticleReal ds, + amrex::ParticleReal k, + int unit, + amrex::ParticleReal dx = 0, + amrex::ParticleReal dy = 0, + amrex::ParticleReal rotation_degree = 0, + int nslice = 1 + ) + : Thick(ds, nslice), + Alignment(dx, dy, rotation_degree), + m_k(k), m_unit(unit) + { + } + + /** Push all particles */ + using BeamOptic::operator(); + + /** This is a chrplasmalens functor, so that a variable of this type can be used like a chrplasmalens function. + * + * @param p Particle AoS data for positions and cpu/id + * @param px particle momentum in x + * @param py particle momentum in y + * @param pt particle momentum in t + * @param refpart reference particle + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + void operator() ( + PType& AMREX_RESTRICT p, + amrex::ParticleReal & AMREX_RESTRICT px, + amrex::ParticleReal & AMREX_RESTRICT py, + amrex::ParticleReal & AMREX_RESTRICT pt, + RefPart const & refpart + ) const + { + using namespace amrex::literals; // for _rt and _prt + + // shift due to alignment errors of the element + shift_in(p.pos(RealAoS::x), p.pos(RealAoS::y), px, py); + + // access AoS data such as positions and cpu/id + amrex::ParticleReal const x = p.pos(RealAoS::x); + amrex::ParticleReal const y = p.pos(RealAoS::y); + amrex::ParticleReal const t = p.pos(RealAoS::t); + + // length of the current slice + amrex::ParticleReal const slice_ds = m_ds / nslice(); + + // access reference particle values to find beta + amrex::ParticleReal const bet = refpart.beta(); + + // normalize focusing strength units to MAD-X convention if needed + amrex::ParticleReal g = m_k; + if (m_unit == 1) { + g = m_k / refpart.rigidity_Tm(); + } + + // compute particle momentum deviation delta + 1 + amrex::ParticleReal delta1; + delta1 = sqrt(1_prt - 2_prt*pt/bet + pow(pt,2)); + amrex::ParticleReal const delta = delta1 - 1_prt; + + // compute phase advance per unit length in s (in rad/m) + // chromatic dependence on delta is included + amrex::ParticleReal const omega = sqrt(std::abs(g)/delta1); + + // intialize output values of momenta + amrex::ParticleReal pxout = px; + amrex::ParticleReal pyout = py; + amrex::ParticleReal const ptout = pt; + + // paceholder variables + amrex::ParticleReal q1 = x; + amrex::ParticleReal q2 = y; + amrex::ParticleReal p1 = px; + amrex::ParticleReal p2 = py; + + // advance transverse position and momentum (focusing) + p.pos(RealAoS::x) = cos(omega*slice_ds)*x + + sin(omega*slice_ds)/(omega*delta1)*px; + pxout = -omega*delta1*sin(omega*slice_ds)*x + cos(omega*slice_ds)*px; + + p.pos(RealAoS::y) = cos(omega*slice_ds)*y + + sin(omega*slice_ds)/(omega*delta1)*py; + pyout = -omega*delta1*sin(omega*slice_ds)*y + cos(omega*slice_ds)*py; + + // advance longitudinal position and momentum + + // the corresponding symplectic update to t + amrex::ParticleReal const term = pt + delta/bet; + amrex::ParticleReal const t0 = t - term*slice_ds/delta1; + + amrex::ParticleReal const w = omega*delta1; + amrex::ParticleReal const term1 = -(pow(p2,2)-pow(q2,2)*pow(w,2))*sin(2_prt*slice_ds*omega); + amrex::ParticleReal const term2 = -(pow(p1,2)-pow(q1,2)*pow(w,2))*sin(2_prt*slice_ds*omega); + amrex::ParticleReal const term3 = -2_prt*q2*p2*w*cos(2_prt*slice_ds*omega); + amrex::ParticleReal const term4 = -2_prt*q1*p1*w*cos(2_prt*slice_ds*omega); + amrex::ParticleReal const term5 = 2_prt*omega*(q1*p1*delta1 + q2*p2*delta1 + -(pow(p1,2)+pow(p2,2))*slice_ds - (pow(q1,2)+pow(q2,2))*pow(w,2)*slice_ds); + p.pos(RealAoS::t) = t0 + (-1_prt+bet*pt)/(8_prt*bet*pow(delta1,3)*omega) + *(term1+term2+term3+term4+term5); + + // ptout = pt; + + // assign updated momenta + px = pxout; + py = pyout; + pt = ptout; + + // undo shift due to alignment errors of the element + shift_out(p.pos(RealAoS::x), p.pos(RealAoS::y), px, py); + } + + /** This pushes the reference particle. + * + * @param[in,out] refpart reference particle + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + void operator() (RefPart & AMREX_RESTRICT refpart) const + { + using namespace amrex::literals; // for _rt and _prt + + // assign input reference particle values + amrex::ParticleReal const x = refpart.x; + amrex::ParticleReal const px = refpart.px; + amrex::ParticleReal const y = refpart.y; + amrex::ParticleReal const py = refpart.py; + amrex::ParticleReal const z = refpart.z; + amrex::ParticleReal const pz = refpart.pz; + amrex::ParticleReal const t = refpart.t; + amrex::ParticleReal const pt = refpart.pt; + amrex::ParticleReal const s = refpart.s; + + // length of the current slice + amrex::ParticleReal const slice_ds = m_ds / nslice(); + + // assign intermediate parameter + amrex::ParticleReal const step = slice_ds / sqrt(pow(pt,2)-1.0_prt); + + // advance position and momentum (straight element) + refpart.x = x + step*px; + refpart.y = y + step*py; + refpart.z = z + step*pz; + refpart.t = t - step*pt; + + // advance integrated path length + refpart.s = s + slice_ds; + } + + amrex::ParticleReal m_k; //! focusing strength in 1/m^2 (or T/m) + int m_unit; //! unit specification for focusing strength + }; + +} // namespace impactx + +#endif // IMPACTX_CHRPLASMALENS_H diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 90d39c001..2a123cb65 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -261,6 +261,39 @@ void init_elements(py::module& m) ; register_beamoptics_push(py_ChrQuad); + py::class_ py_ChrPlasmaLens(me, "ChrPlasmaLens"); + py_ChrPlasmaLens + .def(py::init< + amrex::ParticleReal, + amrex::ParticleReal, + int, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + int + >(), + py::arg("ds"), + py::arg("k"), + py::arg("units") = 0, + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("nslice") = 1, + "An active Plasma Lens with chromatic effects included." + ) + .def_property("k", + [](ChrQuad & cq) { return cq.m_k; }, + [](ChrQuad & cq, amrex::ParticleReal k) { cq.m_k = k; }, + "focusing strength in 1/m^2 (or T/m)" + ) + .def_property("units", + [](ChrQuad & cq) { return cq.m_unit; }, + [](ChrQuad & cq, int unit) { cq.m_unit = unit; }, + "unit specification for focusing strength" + ) + ; + register_beamoptics_push(py_ChrPlasmaLens); + py::class_ py_ChrAcc(me, "ChrAcc"); py_ChrAcc .def(py::init<