Add linear return maps for additional elements (#923)

src/elements/Buncher.H

@@ -149,8 +149,23 @@ namespace impactx::elements
         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
+
+            // find beta*gamma^2
+            amrex::ParticleReal const betgam2 = std::pow(refpart.pt, 2) - 1.0_prt;
+
+            amrex::ParticleReal const neg_kV = -m_k * m_V;
+            amrex::ParticleReal const kV_r2bg2 = -neg_kV / (2.0_prt * betgam2);
+
+            // initialize linear map matrix elements
+            Map6x6 R = Map6x6::Identity();
+
+            // off-identity matrix elements
+            R(2,1) = kV_r2bg2;
+            R(4,3) = kV_r2bg2;
+            R(6,5) = neg_kV;
+
+            return R;
         }
 
         amrex::ParticleReal m_V; //! normalized (max) RF voltage drop.

src/elements/CFbend.H

@@ -261,8 +261,57 @@ namespace impactx::elements
         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
+
+            // length of the current slice
+            amrex::ParticleReal const slice_ds = m_ds / nslice();
+
+            // find beta*gamma^2, beta
+            amrex::ParticleReal const betgam2 = std::pow(refpart.pt, 2) - 1_prt;
+            amrex::ParticleReal const bet = refpart.beta();
+            amrex::ParticleReal const ibetgam2 = 1_prt / betgam2;
+            amrex::ParticleReal const b2rc2 = std::pow(bet, 2) * std::pow(m_rc, 2);
+
+            // update horizontal and longitudinal phase space variables
+            amrex::ParticleReal const gx = m_k + std::pow(m_rc,-2);
+            amrex::ParticleReal const omega_x = std::sqrt(std::abs(gx));
+
+            // update vertical phase space variables
+            amrex::ParticleReal const gy = -m_k;
+            amrex::ParticleReal const omega_y = std::sqrt(std::abs(gy));
+
+            // trigonometry
+            auto const [sinx, cosx] = amrex::Math::sincos(omega_x * slice_ds);
+            amrex::ParticleReal const sinhx = std::sinh(omega_x * slice_ds);
+            amrex::ParticleReal const coshx = std::cosh(omega_x * slice_ds);
+            auto const [siny, cosy] = amrex::Math::sincos(omega_y * slice_ds);
+            amrex::ParticleReal const sinhy = std::sinh(omega_y * slice_ds);
+            amrex::ParticleReal const coshy = std::cosh(omega_y * slice_ds);
+
+            amrex::ParticleReal const igbrc = 1_prt / (gx * bet * m_rc);
+            amrex::ParticleReal const iobrc = 1_prt / (omega_x * bet * m_rc);
+            amrex::ParticleReal const igobr = 1_prt / (gx * omega_x * b2rc2);
+
+            // initialize linear map matrix elements
+            Map6x6 R = Map6x6::Identity();
+
+            R(1,1) = gx > 0_prt ? cosx : coshx;
+            R(1,2) = gx > 0_prt ? sinx / omega_x : sinhx / omega_x;
+            R(1,6) = gx > 0_prt ? -(1_prt - cosx) * igbrc : -(1_prt - coshx) * igbrc;
+            R(2,1) = gx > 0_prt ? -omega_x * sinx : omega_x * sinhx;
+            R(2,2) = gx > 0_prt ? cosx : coshx;
+            R(2,6) = gx > 0_prt ? -sinx * iobrc : -sinhx * iobrc;
+            R(3,3) = gy > 0_prt ? cosy : coshy;
+            R(3,4) = gy > 0_prt ? siny / omega_y : sinhy / omega_y;
+            R(4,3) = gy > 0_prt ? -omega_y * siny : omega_y * sinhy;
+            R(4,4) = gy > 0_prt ? cosy : coshy;
+            R(5,1) = gx > 0_prt ? sinx * iobrc : sinhx * iobrc;
+            R(5,2) = gx > 0_prt ? (1_prt - cosx) * igbrc : (1_prt - coshx) * igbrc;
+            R(5,6) = gx > 0_prt ?
+                slice_ds * ibetgam2 + (sinx  - omega_x * slice_ds) * igobr :
+                slice_ds * ibetgam2 + (sinhx - omega_x * slice_ds) * igobr;
+
+            return R;
         }
 
         amrex::ParticleReal m_rc; //! bend radius in m

src/elements/DipEdge.H

@@ -161,15 +161,15 @@ namespace impactx::elements
             Map6x6 R = Map6x6::Identity();
 
             // edge focusing matrix elements (zero gap)
-            amrex::ParticleReal const R21 = std::tan(m_psi)/m_rc;
-            amrex::ParticleReal R43 = -R21;
-            amrex::ParticleReal vf = 0;
+            amrex::ParticleReal const R21 = std::tan(m_psi) / m_rc;
 
             // first-order effect of nonzero gap
             auto const [sin_psi, cos_psi] = amrex::Math::sincos(m_psi);
-            vf = (1.0_prt + std::pow(sin_psi, 2)) / std::pow(cos_psi, 3);
-            vf *= m_g * m_K2/(std::pow(m_rc,2));
-            R43 += vf;
+            amrex::ParticleReal const vf = (1.0_prt + std::pow(sin_psi, 2))
+                / std::pow(cos_psi, 3)
+                * m_g * m_K2 / std::pow(m_rc, 2);
+
+            amrex::ParticleReal const R43 = vf - R21;
 
             // set non-identity matrix elements
             R(2,1) = R21;

src/elements/Kicker.H

@@ -170,8 +170,10 @@ namespace impactx::elements
         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();
+            // an ideal kick does not affect the linear map
+            Map6x6 R = Map6x6::Identity();
+
+            return R;
         }
 
         amrex::ParticleReal m_xkick; //! horizontal kick strength

src/elements/LinearMap.H

@@ -186,8 +186,10 @@ namespace impactx::elements
         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();
+            Map6x6 R = Map6x6::Identity();
+            R = m_transport_map;
+
+            return R;
         }
 
         Map6x6 m_transport_map;  // 6x6 transport map

src/elements/PlaneXYRot.H

@@ -160,8 +160,21 @@ namespace impactx::elements
         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();
+            // store sin(phi) and cos(phi)
+            auto const [sin_phi, cos_phi] = amrex::Math::sincos(m_phi);
+
+            // assign linear map matrix elements
+            Map6x6 R = Map6x6::Identity();
+            R(1,1) = cos_phi;
+            R(1,3) = -sin_phi;
+            R(2,2) = cos_phi;
+            R(2,4) = -sin_phi;
+            R(3,1) = sin_phi;
+            R(3,3) = cos_phi;
+            R(4,2) = sin_phi;
+            R(4,4) = cos_phi;
+
+            return R;
         }
 
         amrex::ParticleReal m_phi; //! angle of rotation (rad)

src/elements/SoftQuad.H

@@ -342,8 +342,11 @@ namespace SoftQuadrupoleData
         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();
+
+            Map6x6 R = Map6x6::Identity();
+            R = refpart.map;
+
+            return R;
         }
 
         /** This evaluates the on-axis magnetic field Bz at a fixed location

src/elements/SoftSol.H

@@ -351,8 +351,11 @@ namespace SoftSolenoidData
         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();
+
+            Map6x6 R = Map6x6::Identity();
+            R = refpart.map;
+
+            return R;
         }
 
         /** This evaluates the on-axis magnetic field Bz at a fixed location

src/elements/Sol.H

@@ -239,8 +239,51 @@ namespace impactx::elements
         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
+
+            // length of the current slice
+            amrex::ParticleReal const slice_ds = m_ds / nslice();
+
+            // initialize linear map matrix elements
+            Map6x6 R = Map6x6::Identity();
+            Map6x6 Rrotation = Map6x6::Identity();
+            Map6x6 Rfocusing = Map6x6::Identity();
+
+            // access reference particle values to find beta*gamma^2
+            amrex::ParticleReal const betgam2 = std::pow(refpart.pt, 2) - 1.0_prt;
+
+            // compute phase advance per unit length (in rad/m) and
+            // rotation angle (in rad)
+            amrex::ParticleReal const alpha = m_ks / 2.0_prt;
+            amrex::ParticleReal const theta = alpha*slice_ds;
+            auto const [sin_theta, cos_theta] = amrex::Math::sincos(theta);
+
+            amrex::ParticleReal const sin_theta_ialpha = sin_theta / alpha;
+            amrex::ParticleReal const neg_alpha_sin_theta = -alpha * sin_theta;
+            amrex::ParticleReal const ds_bg2 = slice_ds / betgam2;
+
+            Rfocusing(1,1) = cos_theta;
+            Rfocusing(1,2) = sin_theta_ialpha;
+            Rfocusing(2,1) = neg_alpha_sin_theta;
+            Rfocusing(2,2) = cos_theta;
+            Rfocusing(3,3) = cos_theta;
+            Rfocusing(3,4) = sin_theta_ialpha;
+            Rfocusing(4,3) = neg_alpha_sin_theta;
+            Rfocusing(4,4) = cos_theta;
+            Rfocusing(5,6) = ds_bg2;
+
+            Rrotation(1,1) = cos_theta;
+            Rrotation(1,3) = sin_theta;
+            Rrotation(2,2) = cos_theta;
+            Rrotation(2,4) = sin_theta;
+            Rrotation(3,1) = -sin_theta;
+            Rrotation(3,3) = cos_theta;
+            Rrotation(4,2) = -sin_theta;
+            Rrotation(4,4) = cos_theta;
+
+            R = Rrotation * Rfocusing;
+
+            return R;
         }
 
         amrex::ParticleReal m_ks; //! solenoid strength in 1/m