[WIP] Linear Transport: All Maps & Tests

src/diagnostics/LinearMap.cpp

@@ -62,8 +62,17 @@ namespace
         }
         else
         {
-            // Cold path: only instantiated for element types whose
-            // transport_map is known to throw. No try/catch needed.
+            // Cold path: element has no compile-time-guaranteed linear
+            // map. Some elements (e.g. @c Programmable with an
+            // optional user-supplied hook) can still provide one at
+            // runtime -- try them first and fall through to the
+            // @p on_missing policy only if they throw.
+            try {
+                return element.transport_map(ref);
+            } catch (std::exception const &) {
+                // Fall through to the policy branch below.
+            }
+
             std::string const type_name{T_Element::type};
             switch (on_missing)
             {

src/elements/ChrUniformAcc.H

@@ -33,7 +33,7 @@ namespace impactx::elements
     struct ChrAcc
     : public mixin::Named,
       public mixin::BeamOptic<ChrAcc>,
-      public mixin::LinearTransport<ChrAcc, false>,
+      public mixin::LinearTransport<ChrAcc>,
       public mixin::Thick,
       public mixin::Alignment,
       public mixin::PipeAperture,
@@ -289,15 +289,126 @@ namespace impactx::elements
 
         /** This function returns the linear transport map.
          *
-         * @param[in] refpart reference particle
+         * Jacobian at the reference trajectory of the ChrAcc push (see
+         * @see operator()). The nonlinear push composes, in order:
+         *   (1) static->dynamic unit conversion: diag(1, bgi, 1, bgi, 1, bgi)
+         *       acting on (x, px, y, py, t, pt);
+         *   (2) solenoid-like focusing by phase advance theta (cf. Sol);
+         *   (3) rotation by the same theta in the (x, y) plane;
+         *   (4) dynamic->static unit conversion: diag(1, 1/bgf, 1, 1/bgf,
+         *       1, 1/bgf) on the output;
+         * plus a symplectic correction to t. Here
+         *   theta = (alpha/ez) * log((pzf_ref - ptf_ref)/(pzi_ref - pti_ref))
+         * with alpha = Bz/2, ez the normalized longitudinal electric field,
+         * pzi_ref = sqrt(pti_ref^2 - 1), pzf_ref = sqrt(ptf_ref^2 - 1), and
+         * (pti_ref, ptf_ref) = (initial, final) reference pt on this slice.
+         *
+         * This routine is called after the reference particle has been
+         * advanced by one slice (see diagnostics/LinearMap.cpp), so
+         * refpart.pt holds ptf_ref and pti_ref = ptf_ref + ez * slice_ds.
+         *
+         * The closed-form linear map was derived with SymPy by Jacobian
+         * of the exact push at (x, px, y, py, t, pt) = 0. Writing
+         * c = cos(theta0), s = sin(theta0) with theta0 the reference
+         * phase advance above, the nonzero entries are
+         *   R(1,1) = c^2,                R(1,2) = bgi * c*s / alpha,
+         *   R(1,3) = c*s,                R(1,4) = bgi * s^2 / alpha,
+         *   R(2,1) = -alpha * c*s / bgf, R(2,2) = (bgi/bgf) * c^2,
+         *   R(2,3) = -alpha * s^2 / bgf, R(2,4) = (bgi/bgf) * c*s,
+         *   R(3,1) = -c*s,               R(3,2) = -bgi * s^2 / alpha,
+         *   R(3,3) = c^2,                R(3,4) = bgi * c*s / alpha,
+         *   R(4,1) =  alpha * s^2 / bgf, R(4,2) = -(bgi/bgf) * c*s,
+         *   R(4,3) = -alpha * c*s / bgf, R(4,4) = (bgi/bgf) * c^2,
+         *   R(5,5) = 1,
+         *   R(5,6) = (bgi/ez) * (ptf_ref/pzf_ref - pti_ref/pzi_ref),
+         *   R(6,6) = bgi/bgf.
+         * All other entries are zero. The quadratic-in-(px, py, x, y)
+         * part of the t correction vanishes at linear order.
+         *
+         * References:
+         *   - A. J. Dragt, Lie Methods for Nonlinear Dynamics with
+         *     Applications to Accelerator Physics, Univ. of Maryland,
+         *     2020 (MaryLie documentation), chapter on uniform-Ez
+         *     acceleration with solenoidal focusing;
+         *   - H. Wiedemann, Particle Accelerator Physics, 4th ed.,
+         *     Springer 2015, chapters on electrostatic acceleration and
+         *     solenoid focusing;
+         *   - The Hamiltonian and nonlinear push follow Chad Mitchell's
+         *     formulation already implemented in @see operator().
+         *
+         * Ideal behavior only: alignment and rotation errors and
+         * feed-down from an off-axis reference orbit are ignored.
+         *
+         * @param[in] refpart reference particle after this slice's
+         *                    reference advance (pt = ptf_ref)
          * @returns 6x6 transport matrix
          */
         AMREX_GPU_HOST AMREX_FORCE_INLINE
         Map6x6
-        transport_map ([[maybe_unused]] RefPart const & AMREX_RESTRICT refpart) const
+        transport_map (RefPart const & AMREX_RESTRICT refpart) const
         {
-            throw std::runtime_error(std::string(type) + ": Envelope tracking is not yet implemented!");
-            return Map6x6::Identity();
+            using namespace amrex::literals; // for _rt and _prt
+            using amrex::Math::powi;
+
+            // slice length and initial/final reference pt:
+            // transport_map is called after the reference advance, so
+            // refpart.pt holds the post-slice (final) pt.
+            amrex::ParticleReal const slice_ds = m_ds / nslice();
+            amrex::ParticleReal const ptf_ref = refpart.pt;
+            amrex::ParticleReal const pti_ref = ptf_ref + m_ez * slice_ds;
+
+            amrex::ParticleReal const pzf_ref = std::sqrt(powi<2>(ptf_ref) - 1.0_prt);
+            amrex::ParticleReal const pzi_ref = std::sqrt(powi<2>(pti_ref) - 1.0_prt);
+            amrex::ParticleReal const bgf = pzf_ref;  // |beta*gamma|_final
+            amrex::ParticleReal const bgi = pzi_ref;  // |beta*gamma|_initial
+
+            // Solenoid focusing constant alpha = Bz / 2. The "log"
+            // factor is the integrated drift length in dynamic units and
+            // is well-defined even when alpha = 0 (no solenoid).
+            amrex::ParticleReal const alpha = 0.5_prt * m_bz;
+            amrex::ParticleReal const drift_log =
+                std::log((pzf_ref - ptf_ref) / (pzi_ref - pti_ref)) / m_ez;
+            amrex::ParticleReal const rbg = bgi / bgf;
+
+            Map6x6 R = Map6x6::Identity();
+
+            // t couples to pt only: symplectic correction to the
+            // time-of-flight, linearized at pt = 0.
+            R(5,6) = bgi * (ptf_ref / pzf_ref - pti_ref / pzi_ref) / m_ez;
+            // pt scales by the static<->dynamic unit change bgi/bgf.
+            R(6,6) = rbg;
+
+            if (alpha == 0.0_prt)
+            {
+                // No solenoid: pure accelerating "drift" with adiabatic
+                // damping on the transverse momenta.
+                R(1,2) = bgi * drift_log;
+                R(2,2) = rbg;
+                R(3,4) = bgi * drift_log;
+                R(4,4) = rbg;
+            }
+            else
+            {
+                // Solenoid + acceleration: rotation * focusing, bracketed
+                // by static<->dynamic unit conversions on (px, py).
+                amrex::ParticleReal const theta0 = alpha * drift_log;
+                auto const [s, c] = amrex::Math::sincos(theta0);
+
+                amrex::ParticleReal const c2 = c * c;
+                amrex::ParticleReal const s2 = s * s;
+                amrex::ParticleReal const cs = c * s;
+
+                R(1,1) =  c2;                R(1,2) =  bgi * cs / alpha;
+                R(1,3) =  cs;                R(1,4) =  bgi * s2 / alpha;
+                R(2,1) = -alpha * cs / bgf;  R(2,2) =  rbg * c2;
+                R(2,3) = -alpha * s2 / bgf;  R(2,4) =  rbg * cs;
+                R(3,1) = -cs;                R(3,2) = -bgi * s2 / alpha;
+                R(3,3) =  c2;                R(3,4) =  bgi * cs / alpha;
+                R(4,1) =  alpha * s2 / bgf;  R(4,2) = -rbg * cs;
+                R(4,3) = -alpha * cs / bgf;  R(4,4) =  rbg * c2;
+            }
+
+            return R;
         }
 
         amrex::ParticleReal m_ez; //! electric field strength in 1/m

src/elements/Multipole.H

@@ -32,7 +32,7 @@ namespace impactx::elements
     struct Multipole
     : public mixin::Named,
       public mixin::BeamOptic<Multipole>,
-      public mixin::LinearTransport<Multipole, false>,
+      public mixin::LinearTransport<Multipole>,
       public mixin::Thin,
       public mixin::Alignment,
       public mixin::NoFinalize
@@ -195,15 +195,58 @@ namespace impactx::elements
 
         /** This function returns the linear transport map.
          *
-         * @param[in] refpart reference particle
+         * The thin multipole kick implemented by @see operator() is
+         *   dpx - i*dpy = -(K_normal + i*K_skew) * (x + i*y)^(m-1) / (m-1)!
+         * matching the MAD-X MULTIPOLE thin kick (MAD-X User Guide,
+         * "Generalisation to normal and skew components"),
+         *   dPx - i*dPy = -sum_n (KN_n + i*KS_n) * (x + i*y)^n / n!,
+         * with the index mapping n = m_multipole - 1 and
+         * K_normal <-> KNL[n], K_skew <-> KSL[n].
+         *
+         * The linear map is the Jacobian of this kick around the reference
+         * trajectory, i.e. at (x, y) = (0, 0). Ideal behavior only: we
+         * ignore alignment and rotation errors, and we ignore feed-down
+         * from an off-axis reference orbit (the magnetic center and the
+         * reference trajectory are assumed coincident).
+         *
+         * Three cases:
+         *  - m_multipole == 1 (thin dipole kick): constant, yields the
+         *    identity linear map.
+         *  - m_multipole == 2 (thin quadrupole): the only nontrivial case;
+         *    R(2,1) = -K_normal, R(2,3) = +K_skew,
+         *    R(4,1) = +K_skew,   R(4,3) = +K_normal.
+         *    (Wiedemann, Particle Accelerator Physics, 3rd ed., Sect. 6.3;
+         *    Dragt, Lie Methods for Nonlinear Dynamics with Applications
+         *    to Accelerator Physics, Sect. 7.3.)
+         *  - m_multipole >= 3 (sextupole and higher): the kick vanishes at
+         *    least quadratically at the origin, yielding the identity.
+         *
+         * Sign/units convention: ImpactX phase-space ordering is
+         * (x, px, y, py, t, pt) with 1-based matrix indices, matching the
+         * other linear elements in this directory (e.g. Drift, Quad).
+         *
+         * @param[in] refpart reference particle (unused)
          * @returns 6x6 transport matrix
          */
         AMREX_GPU_HOST AMREX_FORCE_INLINE
         Map6x6
         transport_map ([[maybe_unused]] RefPart const & AMREX_RESTRICT refpart) const
         {
-            throw std::runtime_error(std::string(type) + ": Envelope tracking is not yet implemented!");
-            return Map6x6::Identity();
+            using namespace amrex::literals; // for _rt and _prt
+
+            Map6x6 R = Map6x6::Identity();
+
+            // Only m_multipole == 2 (thin quadrupole) has a nonzero Jacobian
+            // at the reference trajectory; m_multipole == 1 is a constant
+            // kick and m_multipole >= 3 vanishes at the origin.
+            if (m_multipole == 2) {
+                R(2,1) = -m_Kn;
+                R(2,3) = +m_Ks;
+                R(4,1) = +m_Ks;
+                R(4,3) = +m_Kn;
+            }
+
+            return R;
         }
 
         int m_multipole; //! multipole index

src/elements/NonlinearLens.H

@@ -35,7 +35,7 @@ namespace impactx::elements
     struct NonlinearLens
     : public mixin::Named,
       public mixin::BeamOptic<NonlinearLens>,
-      public mixin::LinearTransport<NonlinearLens, false>,
+      public mixin::LinearTransport<NonlinearLens>,
       public mixin::Thin,
       public mixin::Alignment,
       public mixin::NoFinalize
@@ -196,15 +196,47 @@ namespace impactx::elements
 
         /** This function returns the linear transport map.
          *
-         * @param[in] refpart reference particle
+         * The Danilov-Nagaitsev kick implemented by @see operator() is
+         *   dpx - i*dpy = -(knll / cnll) * F'(zeta),
+         * with zeta = (x + i*y) / cnll and
+         *   F'(zeta) = zeta / (1 - zeta^2) + arcsin(zeta) / (1 - zeta^2)^(3/2).
+         * The Taylor series around zeta = 0 starts at
+         *   F'(zeta) = 2*zeta + (8/3)*zeta^3 + O(zeta^5),
+         * so the linearized map at the reference trajectory is a thin
+         * quadrupole with equal and opposite focusing strengths in the
+         * two transverse planes:
+         *   R(2,1) = -2 * knll / cnll^2,
+         *   R(4,3) = +2 * knll / cnll^2.
+         *
+         * Ideal behavior only: alignment and rotation errors, as well as
+         * feed-down from an off-axis reference orbit, are ignored; the
+         * map is evaluated at (x, y) = (0, 0).
+         *
+         * This element appears in MAD-X as type NLLENS. The linear limit
+         * above is the MAD-X linear quadrupole equivalent of the NLLENS
+         * thin kick (cf. MAD-X Physical Methods Manual, nonlinear lens).
+         *
+         * Reference: V. Danilov and S. Nagaitsev, Phys. Rev. ST Accel.
+         * Beams 13, 084002 (2010), Sect. V.A (see also the class-level
+         * docstring above, which cites the same reference).
+         *
+         * @param[in] refpart reference particle (unused)
          * @returns 6x6 transport matrix
          */
         AMREX_GPU_HOST AMREX_FORCE_INLINE
         Map6x6
         transport_map ([[maybe_unused]] RefPart const & AMREX_RESTRICT refpart) const
         {
-            throw std::runtime_error(std::string(type) + ": Envelope tracking is not yet implemented!");
-            return Map6x6::Identity();
+            using namespace amrex::literals; // for _rt and _prt
+            using amrex::Math::powi;
+
+            Map6x6 R = Map6x6::Identity();
+            // Leading 2*zeta term of F'(zeta); higher orders vanish at zero.
+            amrex::ParticleReal const k_lin = 2.0_prt * m_knll / powi<2>(m_cnll);
+            R(2,1) = -k_lin;
+            R(4,3) = +k_lin;
+
+            return R;
         }
 
         amrex::ParticleReal m_knll; //! integrated strength of the nonlinear lens (m)

src/elements/PRot.H

@@ -33,7 +33,7 @@ namespace impactx::elements
     struct PRot
     : public mixin::Named,
       public mixin::BeamOptic<PRot>,
-      public mixin::LinearTransport<PRot, false>,
+      public mixin::LinearTransport<PRot>,
       public mixin::Thin,
       public mixin::NoFinalize,
       public amrex::simd::Vectorized<amrex::simd::native_simd_size_particlereal>
@@ -168,15 +168,65 @@ namespace impactx::elements
 
         /** This function returns the linear transport map.
          *
-         * @param[in] refpart reference particle
+         * Jacobian at the reference trajectory of the exact pole-face
+         * rotation implemented by @see operator(). With net rotation
+         * angle theta = phi_out - phi_in, the nonzero entries are
+         *   R(1,1) = cos(phi_in) / cos(phi_out),
+         *   R(2,2) = cos(phi_out) / cos(phi_in),
+         *   R(2,6) = -sin(theta) / (beta * cos(phi_in)),
+         *   R(5,1) = +sin(theta) / (beta * cos(phi_out)),
+         * and all other diagonal entries are 1. The derivation uses the
+         * trig identity
+         *   cos(phi_in) * cos(phi_in - phi_out)
+         *     + sin(phi_in) * sin(phi_in - phi_out) = cos(phi_out),
+         * which simplifies the x-denominator pz|_0 -> cos(phi_out).
+         * Symplecticity of the reduced (x, px, t, pt) block was
+         * verified with SymPy (R^T J R = J).
+         *
+         * The (y, py) sub-block is the identity: pole-face rotation in
+         * the x-z plane does not mix the vertical plane at linear order.
+         *
+         * Reference: MaryLie 3.0 Users' Manual, Sect. 6.4 ("General
+         * Bends, Geometric Corrections"); A. J. Dragt, Lie Methods for
+         * Nonlinear Dynamics with Applications to Accelerator Physics,
+         * Sect. 12.5 (exact bend geometry / pole-face rotations).
+         *
+         * Sign convention: with ImpactX's pt = -Delta(gamma)/(beta0
+         * gamma0) (opposite sign to MAD-X's PT = +Delta(E)/(p0*c)), the
+         * R(2,6) and R(5,1) signs agree with the sibling Sbend element
+         * in ImpactX but differ overall from the MAD-X PROT formulas.
+         *
+         * Ideal behavior only: alignment and rotation errors are
+         * ignored; the linearization is at (x, y, t) = (0, 0, 0) and
+         * (px, py, pt) = (0, 0, 0).
+         *
+         * @param[in] refpart reference particle (used for beta)
          * @returns 6x6 transport matrix
          */
         AMREX_GPU_HOST AMREX_FORCE_INLINE
         Map6x6
-        transport_map ([[maybe_unused]] RefPart const & AMREX_RESTRICT refpart) const
+        transport_map (RefPart const & AMREX_RESTRICT refpart) const
         {
-            throw std::runtime_error(std::string(type) + ": Envelope tracking is not yet implemented!");
-            return Map6x6::Identity();
+            using namespace amrex::literals; // for _rt and _prt
+
+            amrex::ParticleReal const beta = refpart.beta();
+
+            // Precompute trigonometric quantities.
+            amrex::ParticleReal const theta = m_phi_out - m_phi_in;
+            auto const [sin_theta, cos_theta] = amrex::Math::sincos(theta);
+            (void) cos_theta;  // not needed for the linear map
+            auto const [sin_phi_in, cos_phi_in] = amrex::Math::sincos(m_phi_in);
+            (void) sin_phi_in;
+            auto const [sin_phi_out, cos_phi_out] = amrex::Math::sincos(m_phi_out);
+            (void) sin_phi_out;
+
+            Map6x6 R = Map6x6::Identity();
+            R(1,1) = cos_phi_in / cos_phi_out;
+            R(2,2) = cos_phi_out / cos_phi_in;
+            R(2,6) = -sin_theta / (beta * cos_phi_in);
+            R(5,1) = +sin_theta / (beta * cos_phi_out);
+
+            return R;
         }
 
         amrex::ParticleReal m_phi_in; //! normalized (max) RF voltage drop.

src/elements/PolygonAperture.H

@@ -52,7 +52,7 @@ namespace impactx::elements
     struct PolygonAperture
     : public mixin::Named,
       public mixin::BeamOptic<PolygonAperture>,
-      public mixin::LinearTransport<PolygonAperture, false>,
+      public mixin::LinearTransport<PolygonAperture>,
       public mixin::Thin,
       public mixin::Alignment,
       public mixin::NoFinalize
@@ -276,14 +276,17 @@ namespace impactx::elements
 
         /** This function returns the linear transport map.
          *
-         * @param[in] refpart reference particle
+         * A polygon aperture (potentially) removes particles; it does not
+         * deflect or drift the beam, so the linear transport map is the
+         * identity.
+         *
+         * @param[in] refpart reference particle (unused)
          * @returns 6x6 transport matrix
          */
         AMREX_GPU_HOST AMREX_FORCE_INLINE
         Map6x6
         transport_map ([[maybe_unused]] RefPart const & AMREX_RESTRICT refpart) const
         {
-            throw std::runtime_error(std::string(type) + ": Envelope tracking is not yet implemented!");
             return Map6x6::Identity();
         }