diff --git a/src/elements/CFbend.H b/src/elements/CFbend.H index 4fcd35787..dad8cb2b6 100644 --- a/src/elements/CFbend.H +++ b/src/elements/CFbend.H @@ -81,7 +81,56 @@ namespace impactx::elements /** Push all particles */ using BeamOptic::operator(); + /** 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); + + // length of the current slice + m_slice_ds = m_ds / nslice(); + + // find beta*gamma^2, beta + amrex::ParticleReal const betgam2 = std::pow(refpart.pt, 2) - 1.0_prt; + amrex::ParticleReal const bet = std::sqrt(betgam2 / (1.0_prt + betgam2)); + m_ibetgam2 = 1.0_prt / betgam2; + amrex::ParticleReal const b2rc2 = std::pow(bet, 2) * std::pow(m_rc, 2); + m_igobr = 1.0_prt / ( m_gx * m_omega_x * b2rc2 ); + + // update horizontal and longitudinal phase space variables + m_gx = m_k + std::pow(m_rc,-2); + m_omega_x = std::sqrt(std::abs(m_gx)); + + // update vertical phase space variables + m_gy = -m_k; + m_omega_y = std::sqrt(std::abs(m_gy)); + + // trigonometry + auto const [sinx, cosx] = amrex::Math::sincos(m_omega_x * m_slice_ds); + m_sinx = sinx; + m_cosx = cosx; + m_sinhx = std::sinh(m_omega_x * m_slice_ds); + m_coshx = std::cosh(m_omega_x * m_slice_ds); + auto const [siny, cosy] = amrex::Math::sincos(m_omega_y * m_slice_ds); + m_siny = siny; + m_cosy = cosy; + m_sinhy = std::sinh(m_omega_y * m_slice_ds); + m_coshy = std::cosh(m_omega_y * m_slice_ds); + + m_rgbrc = 1.0_prt / ( m_gx * bet * m_rc ); + m_robrc = m_omega_x * bet * m_rc; + } + /** This is a cfbend functor, so that a variable of this type can be used like a cfbend function. + * + * The @see compute_constants method must be called before pushing particles through this operator. * * @param x particle position in x * @param y particle position in y @@ -90,7 +139,7 @@ namespace impactx::elements * @param py particle momentum in y * @param pt particle momentum in t * @param idcpu particle global index - * @param refpart reference particle + * @param refpart reference particle (unused) */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator() ( @@ -101,7 +150,7 @@ namespace impactx::elements amrex::ParticleReal & AMREX_RESTRICT py, amrex::ParticleReal & AMREX_RESTRICT pt, uint64_t & AMREX_RESTRICT idcpu, - RefPart const & refpart + [[maybe_unused]] RefPart const & AMREX_RESTRICT refpart ) const { using namespace amrex::literals; // for _rt and _prt @@ -119,67 +168,44 @@ namespace impactx::elements amrex::ParticleReal pyout = py; amrex::ParticleReal ptout = pt; - // length of the current slice - amrex::ParticleReal const slice_ds = m_ds / nslice(); - - // access reference particle values to find beta*gamma^2 - amrex::ParticleReal const pt_ref = refpart.pt; - amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1.0_prt; - amrex::ParticleReal const bet = std::sqrt(betgam2/(1.0_prt + betgam2)); - // update horizontal and longitudinal phase space variables - amrex::ParticleReal const gx = m_k + std::pow(m_rc,-2); - amrex::ParticleReal const omegax = std::sqrt(std::abs(gx)); - - if(gx > 0.0) { - // calculate expensive terms once - auto const [sinx, cosx] = amrex::Math::sincos(omegax * slice_ds); - amrex::ParticleReal const r56 = slice_ds/betgam2 - + (sinx - omegax*slice_ds)/(gx*omegax * std::pow(bet,2) * std::pow(m_rc,2)); + if (m_gx > 0.0) + { + amrex::ParticleReal const r56 = m_slice_ds * m_ibetgam2 + + (m_sinx - m_omega_x * m_slice_ds) * m_igobr; // advance position and momentum (focusing) - x = cosx*xout + sinx/omegax*px - (1.0_prt - cosx)/(gx*bet*m_rc)*pt; - pxout = -omegax*sinx*xout + cosx*px - sinx/(omegax*bet*m_rc)*pt; + x = m_cosx * xout + m_sinx / m_omega_x * px - (1.0_prt - m_cosx) * m_rgbrc * pt; + pxout = -m_omega_x * m_sinx * xout + m_cosx * px - m_sinx * m_robrc * pt; - y = sinx/(omegax*bet*m_rc)*xout + (1.0_prt - cosx)/(gx*bet*m_rc)*px - + tout + r56*pt; + y = m_sinx * m_robrc * xout + (1.0_prt - m_cosx) * m_rgbrc * px + + tout + r56 * pt; ptout = pt; - } else { - // calculate expensive terms once - amrex::ParticleReal const sinhx = std::sinh(omegax * slice_ds); - amrex::ParticleReal const coshx = std::cosh(omegax * slice_ds); - amrex::ParticleReal const r56 = slice_ds/betgam2 - + (sinhx - omegax*slice_ds)/(gx*omegax * std::pow(bet,2) * std::pow(m_rc,2)); + } else + { + amrex::ParticleReal const r56 = m_slice_ds * m_ibetgam2 + + (m_sinhx - m_omega_x * m_slice_ds) * m_igobr; // advance position and momentum (defocusing) - x = coshx*xout + sinhx/omegax*px - (1.0_prt - coshx)/(gx*bet*m_rc)*pt; - pxout = omegax*sinhx*xout + coshx*px - sinhx/(omegax*bet*m_rc)*pt; + x = m_coshx * xout + m_sinhx / m_omega_x * px - (1.0_prt - m_coshx) * m_rgbrc * pt; + pxout = m_omega_x * m_sinhx * xout + m_coshx * px - m_sinhx * m_robrc * pt; - t = sinhx/(omegax*bet*m_rc)*xout + (1.0_prt - coshx)/(gx*bet*m_rc)*px - + tout + r56*pt; + t = m_sinhx * m_robrc * xout + (1.0_prt - m_coshx) * m_rgbrc * px + + tout + r56 * pt; ptout = pt; } // update vertical phase space variables - amrex::ParticleReal const gy = -m_k; - amrex::ParticleReal const omegay = std::sqrt(std::abs(gy)); - - if(gy > 0.0) { - // calculate expensive terms once - auto const [siny, cosy] = amrex::Math::sincos(omegay * slice_ds); - + if (m_gy > 0.0) + { // advance position and momentum (focusing) - y = cosy*yout + siny/omegay*py; - pyout = -omegay*siny*yout + cosy*py; - - } else { - // calculate expensive terms once - amrex::ParticleReal const sinhy = std::sinh(omegay * slice_ds); - amrex::ParticleReal const coshy = std::cosh(omegay * slice_ds); - + y = m_cosy * yout + m_siny / m_omega_y * py; + pyout = -m_omega_y * m_siny * yout + m_cosy * py; + } else + { // advance position and momentum (defocusing) - y = coshy*yout + sinhy/omegay*py; - pyout = omegay*sinhy*yout + coshy*py; + y = m_coshy * yout + m_sinhy / m_omega_y * py; + pyout = m_omega_y * m_sinhy * yout + m_coshy * py; } // assign updated momenta @@ -257,6 +283,27 @@ namespace impactx::elements amrex::ParticleReal m_rc; //! bend radius in m amrex::ParticleReal m_k; //! quadrupole strength in m^(-2) + + private: + // constants that are independent of the individually tracked particle, + // see: compute_constants() to refresh + amrex::ParticleReal m_slice_ds; //! m_ds / nslice(); + amrex::ParticleReal m_ibetgam2; //! 1 / (beta*gamma^2) + amrex::ParticleReal m_gx; + amrex::ParticleReal m_gy; + amrex::ParticleReal m_omega_x; + amrex::ParticleReal m_omega_y; + amrex::ParticleReal m_igobr; + amrex::ParticleReal m_sinx; //! sin(omegax*x) + amrex::ParticleReal m_cosx; //! cos(omegax*x) + amrex::ParticleReal m_sinhx; //! sinh(omegax*x) + amrex::ParticleReal m_coshx; //! cosh(omegax*x) + amrex::ParticleReal m_siny; //! sin(omegay*y) + amrex::ParticleReal m_cosy; //! cos(omegay*y) + amrex::ParticleReal m_sinhy; //! sinh(omegay*y) + amrex::ParticleReal m_coshy; //! cosh(omegay*y) + amrex::ParticleReal m_rgbrc; + amrex::ParticleReal m_robrc; }; } // namespace impactx