Quad spin push

src/elements/Quad.H

@@ -18,6 +18,7 @@
 #include "mixin/named.H"
 #include "mixin/nofinalize.H"
 #include "mixin/lineartransport.H"
+#include "mixin/spintransport.H"
 
 #include <AMReX_Extension.H>
 #include <AMReX_Math.H>
@@ -36,6 +37,7 @@ namespace impactx::elements
       public mixin::Thick,
       public mixin::Alignment,
       public mixin::PipeAperture,
+      public mixin::SpinTransport,
       public mixin::NoFinalize
       // At least on Intel AVX512, there is a small overhead to vectorize this element for DP and a gain for SP, see
       // https://github.com/BLAST-ImpactX/impactx/pull/1002
@@ -289,6 +291,84 @@ namespace impactx::elements
             return R;
         }
 
+        /** 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 operator() (
+            [[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_IdCpu & AMREX_RESTRICT idcpu,
+            [[maybe_unused]] RefPart const & AMREX_RESTRICT refpart,
+            T_Real & AMREX_RESTRICT sx,
+            T_Real & AMREX_RESTRICT sy,
+            T_Real & AMREX_RESTRICT sz
+        ) const
+        {
+            using namespace amrex::literals; // for _rt and _prt
+
+            // initialize the three components of the axis-angle vector
+            T_Real lambdax = 0_prt;
+            T_Real lambday = 0_prt;
+            T_Real lambdaz = 0_prt;
+
+            // the value of the gyromagnetic anomaly (TO BE PASSED FROM USER INPUT)
+            // currently, the default value for electrons is used
+            amrex::ParticleReal G = 0.00115965218046;
+            amrex::ParticleReal const gyro_const = 1_prt + G * refpart.gamma();
+
+            // length of the current slice
+            amrex::ParticleReal const slice_ds = m_ds / nslice();
+
+            // compute phase advance per unit length in s (in rad/m)
+            amrex::ParticleReal const omega = std::sqrt(std::abs(m_k));
+
+            // compute trigonometric quantities
+            amrex::ParticleReal const sin_omega_ds = std::sin(omega*slice_ds);
+            amrex::ParticleReal const cos_omega_ds = std::cos(omega*slice_ds);
+            amrex::ParticleReal const sinh_omega_ds = std::sinh(omega*slice_ds);
+            amrex::ParticleReal const cosh_omega_ds = std::cosh(omega*slice_ds);
+
+            if (m_k > 0.0_prt)
+            {
+
+                // horizontally focusing quad case
+                lambdax = -gyro_const * ( py*(cosh_omega_ds - 1_prt) + y*omega*sinh_omega_ds );
+                lambday = -gyro_const * ( px*(1_prt - cos_omega_ds) + x*omega*sin_omega_ds );
+
+            } else if (m_k < 0.0_prt)
+            {
+
+                // horizontally defocusing quad case
+                lambdax = -gyro_const * ( px*(1_prt - cos_omega_ds) + x*omega*sin_omega_ds );
+                lambday = -gyro_const * ( py*(cosh_omega_ds - 1_prt) + y*omega*sinh_omega_ds );
+
+            } else {
+                // treat as a drift
+            }
+
+            // push the spin vector using the generator just determined
+            rotate_spin(lambdax,lambday,lambdaz,sx,sy,sz);
+
+        }
+
+
         /** This pushes the covariance matrix. */
         using LinearTransport::operator();
 

src/elements/mixin/spintransport.H

@@ -0,0 +1,74 @@
+/* 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
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef IMPACTX_ELEMENTS_MIXIN_SPIN_TRANSPORT_H
+#define IMPACTX_ELEMENTS_MIXIN_SPIN_TRANSPORT_H
+
+#include "particles/ImpactXParticleContainer.H"
+
+#include <ablastr/constant.H>
+
+#include <AMReX_Math.H>
+#include <AMReX_Extension.H>
+#include <AMReX_REAL.H>
+
+namespace impactx::elements::mixin
+{
+    /** This is a helper class for applying a spin map generated by a three-component vector
+     *  lambda, whose magnitude defines the angle of rotation and whose direction defines
+     *  the axis of rotation.
+     */
+    struct SpinTransport
+    {
+        /** Rotate the spin vector (sx,sy,sz).
+         *
+         * @param[in] lambdax generator x-component
+         * @param[in] lambday generator y-component
+         * @param[in] lambdaz generator z-component
+         * @param[inout] sx spin vector x-component
+         * @param[inout] sy spin vector y-component
+         * @param[inout] sz spin vector z-component
+         */
+        template<typename T_Real>
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void rotate_spin (
+            amrex::ParticleReal lambdax,
+            amrex::ParticleReal lambday,
+            amrex::ParticleReal lambdaz,
+            amrex::ParticleReal sx,
+            amrex::ParticleReal sy,
+            amrex::ParticleReal sz
+        ) const
+        {
+            using namespace amrex::literals;  // for _prt
+
+            // Compute quaternion coefficients
+            amrex::ParticleReal angle = std::sqrt(lambdax*lambdax+lambday*lambday+lambdaz*lambdaz);
+            auto const [sin_half_angle, cos_half_angle] = amrex::Math::sincos(angle/2_prt);
+            amrex::ParticleReal w0 = cos_half_angle;
+            amrex::ParticleReal w1 = (angle==0_prt) ? 0_prt : -(lambdax/angle)*sin_half_angle;
+            amrex::ParticleReal w2 = (angle==0_prt) ? 0_prt : -(lambday/angle)*sin_half_angle;
+            amrex::ParticleReal w3 = (angle==0_prt) ? 0_prt : -(lambdaz/angle)*sin_half_angle;
+
+            // Apply rotation
+            amrex::ParticleReal sxout = (1_prt-2_prt*(w2*w2+w3*w3))*sx + 2_prt*(w1*w2+w0*w3)*sy + 2_prt*(w1*w3-w0*w2)*sz;
+            amrex::ParticleReal syout = 2_prt*(w1*w2-w0*w3)*sx + (1_prt-2_prt*(w1*w1+w3*w3))*sy + 2_prt*(w0*w1+w2*w3)*sz;
+            amrex::ParticleReal szout = 2_prt*(w0*w2+w1*w3)*sx + 2_prt*(w2*w3-w0*w1)*sy + (1_prt-2_prt*(w1*w1+w2*w2))*sz;
+
+            // Update spin vector
+            sx = sxout;
+            sy = syout;
+            sz = szout;
+        }
+
+    };
+
+} // namespace impactx::elements::mixin
+
+#endif // IMPACTX_ELEMENTS_MIXIN_SPIN_TRANSPORT_H