ChrPlasmaLens element: Support k<0, and fix uninitialized-variable bug in tracking of reference particle

src/elements/ChrPlasmaLens.H

@@ -49,7 +49,7 @@ namespace impactx::elements
          * @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)
+         *        OR azimuthal magnetic field gradient in T/m
          * @param unit  Unit specification
          *           unit = 0 focusing strength in m^(-2)
          *           unit = 1 focusing strength in T/m
@@ -152,21 +152,43 @@ namespace impactx::elements
             amrex::ParticleReal pyout = py;
             amrex::ParticleReal const ptout = pt;
 
-            // paceholder variables
+            // placeholder variables
             amrex::ParticleReal q1 = x;
             amrex::ParticleReal q2 = y;
             amrex::ParticleReal p1 = px;
             amrex::ParticleReal p2 = py;
 
-            auto const [sin_ods, cos_ods] = amrex::Math::sincos(omega*m_slice_ds);
-
-            // advance transverse position and momentum (focusing)
-            xout = cos_ods * x + sin_ods / (omega * delta1) * px;
-            pxout = -omega * delta1 * sin_ods * x + cos_ods * px;
-
-            yout = cos_ods * y + sin_ods / (omega * delta1) * py;
-            pyout = -omega * delta1 * sin_ods * y + cos_ods * py;
-
+            if (m_g > 0_prt)
+            {
+                auto const [sin_ods, cos_ods] = amrex::Math::sincos(omega*m_slice_ds);
+
+                // advance transverse position and momentum (focusing)
+                xout = cos_ods * x + sin_ods / (omega * delta1) * px;
+                pxout = -omega * delta1 * sin_ods * x + cos_ods * px;
+
+                yout = cos_ods * y + sin_ods / (omega * delta1) * py;
+                pyout = -omega * delta1 * sin_ods * y + cos_ods * py;
+
+            }
+            else if (m_g < 0_prt)
+            {
+                auto const sinh_ods = std::sinh(omega*m_slice_ds);
+                auto const cosh_ods = std::cosh(omega*m_slice_ds);
+
+                // advance transverse position and momentum (defocusing)
+                xout = cosh_ods * x + sinh_ods / (omega * delta1) * px;
+                pxout = +omega * delta1 * sinh_ods * x + cosh_ods * px;
+
+                yout = cosh_ods * y + sinh_ods / (omega * delta1) * py;
+                pyout = +omega * delta1 * sinh_ods * y + cosh_ods * py;
+            }
+            else {
+                xout = x + px*m_slice_ds;
+                pxout = px;
+
+                yout = y + py*m_slice_ds;
+                pyout = py;
+            }
             // advance longitudinal position and momentum
 
             // the corresponding symplectic update to t
@@ -232,7 +254,7 @@ namespace impactx::elements
             refpart.t = t - step*pt;
 
             // advance integrated path length
-            refpart.s = s + m_slice_ds;
+            refpart.s = s + slice_ds;
         }
 
         /** This pushes the covariance matrix. */

src/elements/ConstF.H

@@ -99,19 +99,60 @@ namespace impactx::elements
             amrex::ParticleReal const betgam2 = amrex::Math::powi<2>(refpart.pt) - 1.0_prt;
 
             // trigo
-            auto const [sin_kxds, cos_kxds] = amrex::Math::sincos(m_kx * slice_ds);
-            m_cos_kxds = cos_kxds;
-            m_const_x = -m_kx * sin_kxds;
-            auto const [sin_kyds, cos_kyds] = amrex::Math::sincos(m_ky * slice_ds);
-            m_cos_kyds = cos_kyds;
-            m_const_y = -m_ky * sin_kyds;
+            // Note: The convention here is m_kx[1/m] = sign(k)*sqrt(abs(k))
+            // where k[1/m^2] is in the usual MADX convention
+            if (m_kx > 0_prt)
+            {
+                auto const [sin_kxds, cos_kxds] = amrex::Math::sincos(m_kx * slice_ds);
+                m_cos_kxds = cos_kxds;
+                m_const_x  = -m_kx * sin_kxds;
+                m_sincx    = sin_kxds / m_kx;
+            }
+            else if (m_kx < 0_prt)
+            {
+                auto const sinh_kxds = std::sinh(-m_kx*slice_ds);
+                auto const cosh_kxds = std::cosh(-m_kx*slice_ds);
+                m_cos_kxds = cosh_kxds;
+                m_const_x  = -m_kx * sinh_kxds;
+                m_sincx    = -sinh_kxds / m_kx;
+            }
+            else //m_kx == 0
+            {
+                m_cos_kxds = 1.0_prt;
+                m_const_x  = 0.0_prt;
+                m_sincx    = slice_ds;
+            }
+
+            if (m_ky > 0_prt)
+            {
+                auto const [sin_kyds, cos_kyds] = amrex::Math::sincos(m_ky * slice_ds);
+                m_cos_kyds = cos_kyds;
+                m_const_y  = -m_ky * sin_kyds;
+                m_sincy    = sin_kyds / m_ky;
+            }
+            else if (m_ky < 0_prt)
+            {
+                auto const sinh_kyds = std::sinh(-m_ky*slice_ds);
+                auto const cosh_kyds = std::cosh(-m_ky*slice_ds);
+                m_cos_kyds = cosh_kyds;
+                m_const_y  = -m_ky * sinh_kyds;
+                m_sincy    = -sinh_kyds / m_ky;
+            }
+            else //m_ky == 0
+            {
+                m_cos_kyds = 1.0_prt;
+                m_const_y  = 0.0_prt;
+                m_sincy    = slice_ds;
+            }
+
+            if (m_kt < 0_prt)
+            {
+                throw std::runtime_error(std::string(type) + ": must have kt >= 0!");
+            }
             auto const [sin_ktds, cos_ktds] = amrex::Math::sincos(m_kt * slice_ds);
             m_cos_ktds = cos_ktds;
             m_const_t = -m_kt * betgam2 * sin_ktds;
-
-            // intermediate quantities - to avoid division by zero
-            m_sincx = m_kx > 0 ? std::sin(m_kx * slice_ds) / m_kx : slice_ds;
-            m_sincy = m_ky > 0 ? std::sin(m_ky * slice_ds) / m_ky : slice_ds;
+            // intermediate quantity - to avoid division by zero
             amrex::ParticleReal const sinct = m_kt > 0 ? std::sin(m_kt * slice_ds) / m_kt : slice_ds;
             m_const_pt = sinct / betgam2;
         }
@@ -237,29 +278,67 @@ namespace impactx::elements
             amrex::ParticleReal const betgam2 = amrex::Math::powi<2>(refpart.pt) - 1.0_prt;
 
             // trigo
-            auto const [sin_kxds, cos_kxds] = amrex::Math::sincos(m_kx * slice_ds);
-            amrex::ParticleReal const_x = -m_kx * sin_kxds;
-            auto const [sin_kyds, cos_kyds] = amrex::Math::sincos(m_ky * slice_ds);
-            amrex::ParticleReal const_y = -m_ky * sin_kyds;
+            // Note: The convention here is m_kx[1/m] = sign(k)*sqrt(abs(k))
+            // where k[1/m^2] is in the usual MADX convention
+
+            amrex::ParticleReal sin_kxds, cos_kxds, const_x, sincx;
+            if (m_kx > 0_prt) {
+                std::tie(sin_kxds, cos_kxds) = amrex::Math::sincos(m_kx * slice_ds);
+                const_x = -m_kx * sin_kxds;
+                sincx   = sin_kxds / m_kx;
+            }
+            else if (m_kx < 0_prt) {
+                sin_kxds = std::sinh(-m_kx*slice_ds);
+                cos_kxds = std::cosh(-m_kx*slice_ds);
+                const_x  = -m_kx * sin_kxds;
+                sincx    = -sin_kxds / m_kx;
+            }
+            else { //m_kx == 0
+                cos_kxds = 1.0_prt;
+                const_x  = 0.0_prt;
+                sincx    = slice_ds;
+            }
+
+            amrex::ParticleReal sin_kyds, cos_kyds, const_y, sincy;
+            if (m_ky > 0_prt) {
+                std::tie(sin_kyds, cos_kyds) = amrex::Math::sincos(m_ky * slice_ds);
+                const_y = -m_ky * sin_kyds;
+                sincy   = sin_kyds / m_ky;
+            }
+            else if (m_ky < 0_prt) {
+                sin_kyds = std::sinh(-m_ky*slice_ds);
+                cos_kyds = std::cosh(-m_ky*slice_ds);
+                const_y  = -m_ky * sin_kyds;
+                sincy    = -sin_kyds / m_ky;
+            }
+            else { //m_ky == 0
+                cos_kyds = 1.0_prt;
+                const_y  = 0.0_prt;
+                sincy    = slice_ds;
+            }
+
+            if (m_kt < 0_prt) {
+                throw std::runtime_error(std::string(type) + ": must have kt >= 0!");
+            }
             auto const [sin_ktds, cos_ktds] = amrex::Math::sincos(m_kt * slice_ds);
             amrex::ParticleReal const_t = -m_kt * betgam2 * sin_ktds;
-
             // intermediate quantities - to avoid division by zero
-            amrex::ParticleReal sincx = m_kx > 0 ? std::sin(m_kx * slice_ds) / m_kx : slice_ds;
-            amrex::ParticleReal sincy = m_ky > 0 ? std::sin(m_ky * slice_ds) / m_ky : slice_ds;
             amrex::ParticleReal const sinct = m_kt > 0 ? std::sin(m_kt * slice_ds) / m_kt : slice_ds;
             amrex::ParticleReal const_pt = sinct / betgam2;
 
             // assign linear map matrix elements
             Map6x6 R = Map6x6::Identity();
+
             R(1,1) = cos_kxds;
             R(1,2) = sincx;
             R(2,1) = const_x;
             R(2,2) = cos_kxds;
+
             R(3,3) = cos_kyds;
             R(3,4) = sincy;
             R(4,3) = const_y;
             R(4,4) = cos_kyds;
+
             R(5,5) = cos_ktds;
             R(5,6) = const_pt;
             R(6,5) = const_t;

src/elements/Quad.H

@@ -271,9 +271,9 @@ namespace impactx::elements
                 R(4,4) = std::cos(omega*slice_ds);
                 R(5,6) = slice_ds/betgam2;
             } else {
-                R(1,2) = m_slice_ds;
-                R(3,4) = m_slice_ds;
-                R(5,6) = m_slice_ds / betgam2;
+                R(1,2) = slice_ds;
+                R(3,4) = slice_ds;
+                R(5,6) = slice_ds / betgam2;
             }
 
             return R;