Add New Beamline Elements

In ImpactX, one can easily add new beamline elements as a user. There are multiple ways to add new elements to ImpactX, you can pick the one that fits your needs best.

The workflows described here apply both for thin kicks or thick elements. Thick elements can also use soft-edged fringe fields (see existing soft-edged elements for implementation details).

Linear Map

A custom linear element can be provided by specifying the 6x6 linear transport matrix \(R\) as an input. See the FODO cell example for a demonstration of the custom linear element.

The matrix elements \(R(i,j)\) are indexed beginning with 1, so that \(i,j=1,2,3,4,5,6\).

The matrix \(R\) multiplies the phase space vector \((x,p_x,y,p_y,t,p_t)\), where coordinates \((x,y,t)\) have units of m and momenta \((p_x,p_y,p_t)\) are dimensionless. So, for example, \(R(1,1)\) is dimensionless, and \(R(1,2)\) has units of m.

Note

If a user-provided linear map is used, it is up to the user to ensure that the 6x6 transport matrix is symplectic. If a more general form of user-defined transport is needed, the Python Programmable Element and the C++ Element provide a more general approach.

Python Programmable Element

Using the ImpactX Python interface, a custom element named impactx.elements.Programmable can be defined to advance particles using NumPy, CuPy, Numba, PyTorch or any other compatible Python library.

The Programmable element can implement a custom element in two ways:

• Push the whole container, by assigning a push function or

• Push the reference particle and beam particles in two individual functions (beam_particles and ref_particle).

Per ImpactX convention, the reference particle is updated before the beam particles are pushed.

Detailed examples that show usage of the programmable element are:

• FODO cell: implements a user-defined drift

• 15 stage laser-plasma accelerator: implements a user-defined LPA accelerator element using a neural network surrogate via PyTorch

Detailed particle computing interfaces are presented in the pyAMReX examples.

C++ Element

Adding a new beamline element directly to the C++ code base of ImpactX is straight forward and described in the following.

We store all beamline elements under src/elements/.

Let’s take a look at an example, the Drift implementation. To simplify the logic, we use so-called mixin classes, which provide commonly used logic for parallelization, thin/thick elements, alignment error support, etc.

struct Drift
: public mixin::Named,
  public mixin::BeamOptic<Drift>,
  public mixin::LinearTransport<Drift>,
  public mixin::Thick,
  public mixin::Alignment,
  public mixin::PipeAperture,
  public mixin::SpinTransport,
  public mixin::NoFinalize,
  public amrex::simd::Vectorized<amrex::simd::native_simd_size_particlereal>
{
    static constexpr auto type = "Drift";

After this brief boilerplate, our beamline elements implement these parts:

1. a constructor: storing element options

2. for particle tracking:

   • a single-particle operator: pushing the beam particles

   • a reference-particle operator: pushing the reference particle

3. for envelope tracking: a linear transport map

Example Element: Drift.H

/* 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_DRIFT_H
#define IMPACTX_DRIFT_H

#include "particles/ImpactXParticleContainer.H"
#include "mixin/alignment.H"
#include "mixin/beamoptic.H"
#include "mixin/lineartransport.H"
#include "mixin/spintransport.H"
#include "mixin/named.H"
#include "mixin/nofinalize.H"
#include "mixin/pipeaperture.H"
#include "mixin/thick.H"

#include <AMReX_Extension.H>
#include <AMReX_Math.H>
#include <AMReX_REAL.H>
#include <AMReX_SIMD.H>
#include <AMReX_SmallMatrix.H>

#include <cmath>
#include <utility>


namespace impactx::elements
{
    struct Drift
    : public mixin::Named,
      public mixin::BeamOptic<Drift>,
      public mixin::LinearTransport<Drift>,
      public mixin::Thick,
      public mixin::Alignment,
      public mixin::PipeAperture,
      public mixin::SpinTransport,
      public mixin::NoFinalize,
      public amrex::simd::Vectorized<amrex::simd::native_simd_size_particlereal>
    {
        static constexpr auto type = "Drift";
        using PType = ImpactXParticleContainer::ParticleType;

        /** A drift
         *
         * @param ds Segment length in 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 aperture_x horizontal half-aperture in m
         * @param aperture_y vertical half-aperture in m
         * @param nslice number of slices used for the application of space charge
         * @param name a user defined and not necessarily unique name of the element
         */
        Drift (
            amrex::ParticleReal ds,
            amrex::ParticleReal dx = 0,
            amrex::ParticleReal dy = 0,
            amrex::ParticleReal rotation_degree = 0,
            amrex::ParticleReal aperture_x = 0,
            amrex::ParticleReal aperture_y = 0,
            int nslice = 1,
            std::optional<std::string> name = std::nullopt
        )
        : Named(std::move(name)),
          Thick(ds, nslice),
          Alignment(dx, dy, rotation_degree),
          PipeAperture(aperture_x, aperture_y)
        {
        }

        /** Reverse the element (negate ds) */
        void reverse () { Thick::reverse(); }

        /** Push all particles */
        using BeamOptic::operator();

        /** Compute and cache the constants for the push.
         *
         * In particular, used to pre-compute and cache variables that are
         * independent of the individually tracked particle.
         *
         * @param refpart reference particle
         */
        void compute_constants (RefPart const & refpart)
        {
            using namespace amrex::literals; // for _rt and _prt
            using amrex::Math::powi;

            Alignment::compute_constants(refpart);

            // length of the current slice
            m_slice_ds = m_ds / nslice();

            // find beta*gamma^2
            m_betgam2 = powi<2>(refpart.pt) - 1.0_prt;

            m_slice_bg = m_slice_ds / m_betgam2;
        }

        /** This is a drift functor, so that a variable of this type can be used like a drift function.
         *
         * The @see compute_constants method must be called before pushing particles through this operator.
         *
         * @param x particle position in x
         * @param y particle position in y
         * @param t particle position in t
         * @param px particle momentum in x
         * @param py particle momentum in y
         * @param pt particle momentum in t
         * @param idcpu particle global index
         * @param refpart reference particle (unused)
         */
        template<typename T_Real=amrex::ParticleReal, typename T_IdCpu=uint64_t>
        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        void operator() (
            T_Real & AMREX_RESTRICT x,
            T_Real & AMREX_RESTRICT y,
            T_Real & AMREX_RESTRICT t,
            T_Real & AMREX_RESTRICT px,
            T_Real & AMREX_RESTRICT py,
            T_Real & AMREX_RESTRICT pt,
            T_IdCpu & AMREX_RESTRICT idcpu,
            [[maybe_unused]] RefPart const & AMREX_RESTRICT refpart
        ) const
        {
            using namespace amrex::literals; // for _rt and _prt

            // shift due to alignment errors of the element
            shift_in(x, y, px, py);

            // initialize output values
            T_Real xout = x;
            T_Real yout = y;
            T_Real tout = t;
            T_Real pxout = px;
            T_Real pyout = py;
            T_Real ptout = pt;

            // advance position and momentum (drift)
            xout = x + m_slice_ds * px;
            // pxout = px;
            yout = y + m_slice_ds * py;
            // pyout = py;
            tout = t + m_slice_bg * pt;
            // ptout = pt;

            // assign updated values
            x = xout;
            y = yout;
            t = tout;
            px = pxout;
            py = pyout;
            pt = ptout;

            // apply transverse aperture
            apply_aperture(x, y, idcpu);

            // undo shift due to alignment errors of the element
            shift_out(x, 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
            using amrex::Math::powi;

            // 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 / std::sqrt(powi<2>(pt)-1.0_prt);

            // advance position and momentum (drift)
            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;
        }


        /** This operator pushes the particle spin.
         *
         * @param x particle position in x
         * @param y particle position in y
         * @param t particle position in t
         * @param px particle momentum in x
         * @param py particle momentum in y
         * @param pt particle momentum in t
         * @param idcpu particle global index
         * @param refpart reference particle (unused)
         * @param sx particle spin x-component
         * @param sy particle spin y-component
         * @param sz particle spin z-component
         */
        template<typename T_Real=amrex::ParticleReal, typename T_IdCpu=uint64_t>
        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        void spin_and_phasespace_push (
            [[maybe_unused]] T_Real & AMREX_RESTRICT x,
            [[maybe_unused]] T_Real & AMREX_RESTRICT y,
            [[maybe_unused]] T_Real & AMREX_RESTRICT t,
            [[maybe_unused]] T_Real & AMREX_RESTRICT px,
            [[maybe_unused]] T_Real & AMREX_RESTRICT py,
            [[maybe_unused]] T_Real & AMREX_RESTRICT pt,
            [[maybe_unused]] T_Real & AMREX_RESTRICT sx,
            [[maybe_unused]] T_Real & AMREX_RESTRICT sy,
            [[maybe_unused]] T_Real & AMREX_RESTRICT sz,
            [[maybe_unused]] T_IdCpu & AMREX_RESTRICT idcpu,
            [[maybe_unused]] RefPart const & AMREX_RESTRICT refpart
        ) const
        {
             // spin is unaffected

            // phase space push
            (*this)(x, y, t, px, py, pt, idcpu, refpart);

        }


        /** This pushes the covariance matrix. */
        using LinearTransport::operator();

        /** This function returns the linear transport map.
         *
         * @returns 6x6 transport matrix
         */
        AMREX_GPU_HOST AMREX_FORCE_INLINE
        Map6x6
        transport_map (RefPart const & AMREX_RESTRICT refpart) const
        {
            using namespace amrex::literals; // for _rt and _prt
            using amrex::Math::powi;

            // length of the current slice
            amrex::ParticleReal const slice_ds = m_ds / nslice();

            // access reference particle values to find beta*gamma^2
            amrex::ParticleReal const pt_ref = refpart.pt;
            amrex::ParticleReal const betgam2 = powi<2>(pt_ref) - 1.0_prt;

            // assign linear map matrix elements
            Map6x6 R = Map6x6::Identity();
            R(1,2) = slice_ds;
            R(3,4) = slice_ds;
            R(5,6) = slice_ds / betgam2;

            return R;
        }

    private:
        // constants that are independent of the individually tracked particle,
        // see: compute_constants() to refresh
        amrex::ParticleReal m_slice_ds;  //! m_ds / nslice();
        amrex::ParticleReal m_betgam2;   //! beta*gamma^2
        amrex::ParticleReal m_slice_bg;  //! m_slice_ds / m_betgam2
    };

} // namespace impactx

IMPACTX_PUSH_EXTERN_TEMPLATE(impactx::elements::Drift)

#endif // IMPACTX_DRIFT_H

As a last step, we expose our C++ beamline elements to Python in src/python/elements.cpp.

Python Binding: Drift

py::class_<Drift, elements::mixin::Named, elements::mixin::Thick, elements::mixin::Alignment, elements::mixin::PipeAperture> py_Drift(me, "Drift");
py_Drift
    .def("__repr__",
         [](Drift const & drift) {
             return element_name(drift);
         }
    )
    .def("to_dict",
         [](Drift const & drift) {
             return element_dict(drift);
         }
    )
    .def(py::init<
            amrex::ParticleReal,
            amrex::ParticleReal,
            amrex::ParticleReal,
            amrex::ParticleReal,
            amrex::ParticleReal,
            amrex::ParticleReal,
            int,
            std::optional<std::string>
         >(),
         py::arg("ds"),
         py::arg("dx") = 0,
         py::arg("dy") = 0,
         py::arg("rotation") = 0,
         py::arg("aperture_x") = 0,
         py::arg("aperture_y") = 0,
         py::arg("nslice") = 1,
         py::arg("name") = py::none(),
         "A drift."
    )
;
register_push(py_Drift);

Pull requests that added a new element and can be taken as examples are:

• Chromatic Plasma Lens

• Thin-Kick Dipole

• Chromatic Elements for Drift, Quad, Uniform Focusing+Solenoid

• Quadrupole with Soft-Edge Fringe Fields

• other pull requests under the component: elements label