From e35d285bf755fc66889b7537aa7869d7636b2ad7 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Fri, 24 Apr 2026 11:52:59 -0700 Subject: [PATCH 1/5] Add TaperedPL spin push. --- src/elements/TaperedPL.H | 123 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 123 insertions(+) diff --git a/src/elements/TaperedPL.H b/src/elements/TaperedPL.H index de485d506..c37026061 100644 --- a/src/elements/TaperedPL.H +++ b/src/elements/TaperedPL.H @@ -15,6 +15,7 @@ #include "mixin/beamoptic.H" #include "mixin/thin.H" #include "mixin/lineartransport.H" +#include "mixin/spintransport.H" #include "mixin/named.H" #include "mixin/nofinalize.H" @@ -35,6 +36,7 @@ namespace impactx::elements public mixin::LinearTransport, public mixin::Thin, public mixin::Alignment, + public mixin::SpinTransport, public mixin::NoFinalize // At least on Intel AVX512, there is a small overhead to vectorize this element, see // https://github.com/BLAST-ImpactX/impactx/pull/1002 @@ -76,6 +78,32 @@ namespace impactx::elements /** Reverse the element (negate strength) */ void reverse () { m_k = -m_k; } + /** 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 + + Alignment::compute_constants(refpart); + + // normalize focusing strength units to MAD-X convention if needed + if (m_unit == 1) { + m_g = m_k / refpart.rigidity_Tm(); + } + + // additional constants needed for spin push + m_gamma_ref = refpart.gamma(); + m_gyro_anomaly = refpart.gyromagnetic_anomaly; + m_beta = refpart.beta(); + m_ibeta = 1_prt / m_beta; + + } + /** Push all particles */ using BeamOptic::operator(); @@ -150,6 +178,88 @@ namespace impactx::elements /** This pushes the reference particle. */ using Thin::operator(); + + /** 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 sx particle spin x-component + * @param sy particle spin y-component + * @param sz particle spin z-component + * @param idcpu particle global index + * @param refpart reference particle (unused) + */ + template + 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 + { + using namespace amrex::literals; // for _rt and _prt + using amrex::Math::powi; + + // initialize the three components of the axis-angle vector + T_Real lambdax = 0_prt; + T_Real lambday = 0_prt; + T_Real lambdaz = 0_prt; + + // normalize focusing strength units to MAD-X convention if needed + // amrex::ParticleReal g = m_k; + // if (m_unit == 1) { + // g = m_k / refpart.rigidity_Tm(); + // } + + // access reference beta*gamma + //amrex::ParticleReal beta_gamma_ref = refpart.beta_gamma(); + + // compute multipole magnetic field normalized by q/mc: + T_Real const Bx = m_g * ( y + m_taper * x * y ) * m_gamma_ref * m_beta; + T_Real const By = -m_g * ( x + m_taper*0.5_prt * (powi<2>(x) + powi<2>(y)) ) * m_gamma_ref * m_beta; + T_Real const Bz = 0_prt; + + // Electric field normalized by q/mc^2: + T_Real const Ex = 0.0_prt; + T_Real const Ey = 0.0_prt; + T_Real const Ez = 0.0_prt; + + // Quantities required to evaluate the full Thomas-BMT precession vector. + T_Real const Pnorm = sqrt(1_prt - 2_prt * pt * m_ibeta + pt*pt); + T_Real const iPnorm = 1_prt/Pnorm; + T_Real const Ps = sqrt(Pnorm*Pnorm - px*px - py*py); + T_Real const gamma = m_gamma_ref * (1_prt - pt*m_beta); + T_Real const ux = px * iPnorm; + T_Real const uy = py * iPnorm; + T_Real const uz = Ps * iPnorm; + + // Zero curvature in a quadrupole + amrex::ParticleReal const h = 0.0_prt; + + // Evaluation of the full Thomas-BMT precession vector. + tbmt_precession_vector(x,ux,uy,uz,gamma,h,m_gyro_anomaly,Bx,By,Bz,Ex,Ey,Ez,lambdax,lambday,lambdaz); + + // push the spin vector using the generator just determined + rotate_spin(lambdax,lambday,lambdaz,sx,sy,sz); + + // phase space push + (*this)(x, y, t, px, py, pt, idcpu, refpart); + + } + + /** This pushes the covariance matrix. */ using LinearTransport::operator(); @@ -209,6 +319,19 @@ for beam-quality preservation between plasma-accelerator stages", amrex::ParticleReal m_k; //! linear focusing strength (field gradient) amrex::ParticleReal m_taper; //! horizontal taper parameter int m_unit; //! units for linear focusing strength (field gradient) + // see: compute_factorial() to refresh + int m_mfactorial; //! factorial of multipole index + // see: compute_constants() to refresh + amrex::GpuComplex m_alpha; + amrex::ParticleReal m_inv_factorial; + + private: + // constants that are independent of the individually tracked particle + amrex::ParticleReal m_gamma_ref; //! relativistic gamma + amrex::ParticleReal m_gyro_anomaly; //! particle gyromagnetic anomaly + amrex::ParticleReal m_beta; //! relativistic beta + amrex::ParticleReal m_ibeta; //! inverse of beta + amrex::ParticleReal m_g; //! quadrupole strength }; } // namespace impactx From 0346f8ed0276fc5b07157aca65771d564f6d75b1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 24 Apr 2026 18:54:46 +0000 Subject: [PATCH 2/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/elements/TaperedPL.H | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/elements/TaperedPL.H b/src/elements/TaperedPL.H index c37026061..afbe2f6ae 100644 --- a/src/elements/TaperedPL.H +++ b/src/elements/TaperedPL.H @@ -85,15 +85,15 @@ namespace impactx::elements * * @param refpart reference particle */ - void compute_constants (RefPart const & refpart) + void compute_constants (RefPart const & refpart) { using namespace amrex::literals; // for _rt and _prt - + Alignment::compute_constants(refpart); - + // normalize focusing strength units to MAD-X convention if needed if (m_unit == 1) { - m_g = m_k / refpart.rigidity_Tm(); + m_g = m_k / refpart.rigidity_Tm(); } // additional constants needed for spin push @@ -101,7 +101,7 @@ namespace impactx::elements m_gyro_anomaly = refpart.gyromagnetic_anomaly; m_beta = refpart.beta(); m_ibeta = 1_prt / m_beta; - + } /** Push all particles */ @@ -218,7 +218,7 @@ namespace impactx::elements T_Real lambdaz = 0_prt; // normalize focusing strength units to MAD-X convention if needed - // amrex::ParticleReal g = m_k; + // amrex::ParticleReal g = m_k; // if (m_unit == 1) { // g = m_k / refpart.rigidity_Tm(); // } From 916ad30c705d0615efd63c6a4798974399a4034e Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Fri, 24 Apr 2026 12:35:33 -0700 Subject: [PATCH 3/5] Delete unused commented lines. --- src/elements/TaperedPL.H | 9 --------- 1 file changed, 9 deletions(-) diff --git a/src/elements/TaperedPL.H b/src/elements/TaperedPL.H index afbe2f6ae..5be40a3f3 100644 --- a/src/elements/TaperedPL.H +++ b/src/elements/TaperedPL.H @@ -217,15 +217,6 @@ namespace impactx::elements T_Real lambday = 0_prt; T_Real lambdaz = 0_prt; - // normalize focusing strength units to MAD-X convention if needed - // amrex::ParticleReal g = m_k; - // if (m_unit == 1) { - // g = m_k / refpart.rigidity_Tm(); - // } - - // access reference beta*gamma - //amrex::ParticleReal beta_gamma_ref = refpart.beta_gamma(); - // compute multipole magnetic field normalized by q/mc: T_Real const Bx = m_g * ( y + m_taper * x * y ) * m_gamma_ref * m_beta; T_Real const By = -m_g * ( x + m_taper*0.5_prt * (powi<2>(x) + powi<2>(y)) ) * m_gamma_ref * m_beta; From 58421f95a4ac446e3bb14c80a0c2501c95fae973 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Tue, 28 Apr 2026 12:25:11 -0700 Subject: [PATCH 4/5] Apply suggestions from review. --- src/elements/TaperedPL.H | 38 ++++++++++++++++++-------------------- 1 file changed, 18 insertions(+), 20 deletions(-) diff --git a/src/elements/TaperedPL.H b/src/elements/TaperedPL.H index 5be40a3f3..6cb8f8903 100644 --- a/src/elements/TaperedPL.H +++ b/src/elements/TaperedPL.H @@ -97,8 +97,6 @@ namespace impactx::elements } // additional constants needed for spin push - m_gamma_ref = refpart.gamma(); - m_gyro_anomaly = refpart.gyromagnetic_anomaly; m_beta = refpart.beta(); m_ibeta = 1_prt / m_beta; @@ -191,22 +189,22 @@ namespace impactx::elements * @param sy particle spin y-component * @param sz particle spin z-component * @param idcpu particle global index - * @param refpart reference particle (unused) + * @param refpart reference particle */ template 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, + 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_Real & AMREX_RESTRICT sx, + T_Real & AMREX_RESTRICT sy, + T_Real & AMREX_RESTRICT sz, [[maybe_unused]] T_IdCpu & AMREX_RESTRICT idcpu, - [[maybe_unused]] RefPart const & AMREX_RESTRICT refpart + RefPart const & AMREX_RESTRICT refpart ) const { using namespace amrex::literals; // for _rt and _prt @@ -218,8 +216,9 @@ namespace impactx::elements T_Real lambdaz = 0_prt; // compute multipole magnetic field normalized by q/mc: - T_Real const Bx = m_g * ( y + m_taper * x * y ) * m_gamma_ref * m_beta; - T_Real const By = -m_g * ( x + m_taper*0.5_prt * (powi<2>(x) + powi<2>(y)) ) * m_gamma_ref * m_beta; + amrex::ParticleReal gamma_ref = refpart.gamma(); + T_Real const Bx = m_g * ( y + m_taper * x * y ) * gamma_ref * m_beta; + T_Real const By = -m_g * ( x + m_taper*0.5_prt * (powi<2>(x) + powi<2>(y)) ) * gamma_ref * m_beta; T_Real const Bz = 0_prt; // Electric field normalized by q/mc^2: @@ -231,7 +230,7 @@ namespace impactx::elements T_Real const Pnorm = sqrt(1_prt - 2_prt * pt * m_ibeta + pt*pt); T_Real const iPnorm = 1_prt/Pnorm; T_Real const Ps = sqrt(Pnorm*Pnorm - px*px - py*py); - T_Real const gamma = m_gamma_ref * (1_prt - pt*m_beta); + T_Real const gamma = gamma_ref * (1_prt - pt*m_beta); T_Real const ux = px * iPnorm; T_Real const uy = py * iPnorm; T_Real const uz = Ps * iPnorm; @@ -240,10 +239,11 @@ namespace impactx::elements amrex::ParticleReal const h = 0.0_prt; // Evaluation of the full Thomas-BMT precession vector. - tbmt_precession_vector(x,ux,uy,uz,gamma,h,m_gyro_anomaly,Bx,By,Bz,Ex,Ey,Ez,lambdax,lambday,lambdaz); + amrex::ParticleReal gyro_anomaly = refpart.gyromagnetic_anomaly; + tbmt_precession_vector(x, ux, uy, uz, gamma, h, gyro_anomaly, Bx, By, Bz, Ex, Ey, Ez, lambdax, lambday, lambdaz); // push the spin vector using the generator just determined - rotate_spin(lambdax,lambday,lambdaz,sx,sy,sz); + rotate_spin(lambdax, lambday, lambdaz, sx, sy, sz); // phase space push (*this)(x, y, t, px, py, pt, idcpu, refpart); @@ -318,8 +318,6 @@ for beam-quality preservation between plasma-accelerator stages", private: // constants that are independent of the individually tracked particle - amrex::ParticleReal m_gamma_ref; //! relativistic gamma - amrex::ParticleReal m_gyro_anomaly; //! particle gyromagnetic anomaly amrex::ParticleReal m_beta; //! relativistic beta amrex::ParticleReal m_ibeta; //! inverse of beta amrex::ParticleReal m_g; //! quadrupole strength From 4d4058432854be7b43bb408a4b41967e13f71818 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Wed, 6 May 2026 16:43:55 -0700 Subject: [PATCH 5/5] Apply suggestions from code review Co-authored-by: Axel Huebl --- src/elements/TaperedPL.H | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/src/elements/TaperedPL.H b/src/elements/TaperedPL.H index 6cb8f8903..115c816bf 100644 --- a/src/elements/TaperedPL.H +++ b/src/elements/TaperedPL.H @@ -92,9 +92,7 @@ namespace impactx::elements Alignment::compute_constants(refpart); // normalize focusing strength units to MAD-X convention if needed - if (m_unit == 1) { - m_g = m_k / refpart.rigidity_Tm(); - } + m_g = m_unit == 1 ? m_k / refpart.rigidity_Tm() : m_k; // additional constants needed for spin push m_beta = refpart.beta(); @@ -310,11 +308,6 @@ for beam-quality preservation between plasma-accelerator stages", amrex::ParticleReal m_k; //! linear focusing strength (field gradient) amrex::ParticleReal m_taper; //! horizontal taper parameter int m_unit; //! units for linear focusing strength (field gradient) - // see: compute_factorial() to refresh - int m_mfactorial; //! factorial of multipole index - // see: compute_constants() to refresh - amrex::GpuComplex m_alpha; - amrex::ParticleReal m_inv_factorial; private: // constants that are independent of the individually tracked particle