diff --git a/ext/BeamTrackingBeamlinesExt/BeamTrackingBeamlinesExt.jl b/ext/BeamTrackingBeamlinesExt/BeamTrackingBeamlinesExt.jl index af96ae46..555d231e 100644 --- a/ext/BeamTrackingBeamlinesExt/BeamTrackingBeamlinesExt.jl +++ b/ext/BeamTrackingBeamlinesExt/BeamTrackingBeamlinesExt.jl @@ -1,7 +1,7 @@ module BeamTrackingBeamlinesExt using Beamlines, BeamTracking, GTPSA, StaticArrays, KernelAbstractions using Beamlines: isactive, deval, unsafe_getparams, o2i, BitsBeamline, BitsLineElement, isnullspecies -using BeamTracking: get_N_particle, R_to_gamma, R_to_pc, runkernels!, +using BeamTracking: get_N_particle, R_to_gamma, R_to_beta_gamma, R_to_pc, runkernels!, @makekernel, Coords, KernelCall, KernelChain, push import BeamTracking: track!, C_LIGHT, chargeof, massof @@ -122,6 +122,7 @@ include("scibmadstandard.jl") include("linear.jl") include("exact.jl") include("integration.jl") +include("bmadstandard.jl") end \ No newline at end of file diff --git a/ext/BeamTrackingBeamlinesExt/bmadstandard.jl b/ext/BeamTrackingBeamlinesExt/bmadstandard.jl new file mode 100644 index 00000000..63b618dd --- /dev/null +++ b/ext/BeamTrackingBeamlinesExt/bmadstandard.jl @@ -0,0 +1,161 @@ +@inline function cavity(tm::BmadStandard, bunch, rfparams, L) + V = rfparams.voltage + φ0 = rfparams.phi0 + tilde_m, gamsqr_0, β0 = ExactTracking.drift_params(bunch.species, bunch.R_ref) + if rfparams.harmon_master == true + wave_number = rfparams.harmonic_number / rfparams.L_ring + else + wave_number = rfparams.rf_frequency / (β0 * C_LIGHT) + end + p0c = R_to_pc(bunch.species, bunch.R_ref) + return KernelCall(BmadStandardTracking.bmad_cavity!, (V, wave_number, φ0, β0, gamsqr_0, tilde_m, p0c, L)) +end + +@inline function pure_patch(tm::BmadStandard, bunch, patchparams, L) + tilde_m, gamsqr_0, beta_0 = ExactTracking.drift_params(bunch.species, bunch.R_ref) + winv = ExactTracking.w_inv_quaternion(patchparams.dx_rot, patchparams.dy_rot, patchparams.dz_rot) + return KernelCall(ExactTracking.patch!, (beta_0, gamsqr_0, tilde_m, patchparams.dt, patchparams.dx, patchparams.dy, patchparams.dz, winv, L)) +end + +@inline function misalign(tm::BmadStandard, bunch, ap, in) + tilde_m = massof(bunch.species)/R_to_pc(bunch.species, bunch.R_ref) + sign = 2*in - 1 + dx = sign * ap.x_offset + dy = sign * ap.y_offset + dz = sign * ap.z_offset + if ap.x_rot != 0 || ap.y_rot != 0 || ap.tilt != 0 + @warn "Rotational misalignments are ignored (currently not supported)" + end + winv = @SMatrix [1.0 0 0; 0 1.0 0; 0 0 1.0] + return KernelCall(ExactTracking.patch!, (tilde_m, 0, dx, dy, dz, winv, 0)) +end + +@inline function drift(tm::BmadStandard, bunch, L) + tilde_m, gamsqr_0, β0 = ExactTracking.drift_params(bunch.species, bunch.R_ref) + return KernelCall(BmadStandardTracking.magnus_drift!, (β0, gamsqr_0, tilde_m, L)) +end + +@inline function thick_pure_bsolenoid(tm::BmadStandard, bunch, bm0, L) + Ks = get_strengths(bm0, L, bunch.R_ref) + tilde_m, gamsqr_0, β0 = ExactTracking.drift_params(bunch.species, bunch.R_ref) + G = BeamTracking.anom(bunch.species) + return KernelCall(BmadStandardTracking.magnus_solenoid!, (Ks, β0, gamsqr_0, tilde_m, G, L)) +end + +@inline function thick_pure_bdipole(tm::BmadStandard, bunch, bm1, L) + Kn0, Ks0 = get_strengths(bm1, L, bunch.R_ref) + K0 = sqrt(Kn0^2 + Ks0^2) + tilt = atan(Ks0, Kn0) + w = ExactTracking.w_quaternion(0,0,tilt) + w_inv = ExactTracking.w_inv_quaternion(0,0,tilt) + βγ0 = R_to_beta_gamma(bunch.species, bunch.R_ref) + γ0 = sqrt(1 + βγ0^2) + G = BeamTracking.anom(bunch.species) + return KernelCall(BmadStandardTracking.magnus_sbend!, (0, 0, 0, K0, w, w_inv, γ0, βγ0, G, L)) +end + +@inline function thick_bend_pure_bdipole(tm::BmadStandard, bunch, bendparams, bm1, L) + g = bendparams.g_ref + e1 = bendparams.e1 + e2 = bendparams.e2 + tilt = bendparams.tilt_ref + Kn0, Ks0 = get_strengths(bm1, L, bunch.R_ref) + Ks0 ≈ 0 || error("A skew dipole field cannot yet be used in a bend") + w = ExactTracking.w_quaternion(0,0,tilt) + w_inv = ExactTracking.w_inv_quaternion(0,0,tilt) + βγ0 = R_to_beta_gamma(bunch.species, bunch.R_ref) + γ0 = sqrt(1 + βγ0^2) + G = BeamTracking.anom(bunch.species) + return KernelCall(BmadStandardTracking.magnus_sbend!, (e1, e2, g, Kn0, w, w_inv, γ0, βγ0, G, L)) +end + +@inline function thick_bdipole(tm::BmadStandard, bunch, bendparams, bdict, L) + if any(b -> (b.order > 2 || b.order == 0), bdict) + error("BmadStandard does not support thick combined dipoles with higher order multipoles") + end + kn, ks = get_strengths(bdict, L, bunch.R_ref) + mm = bdict.order + k0 = sign(kn[1])*sqrt(kn[1]^2 + ks[1]^2) * (mm[1] == 1) + k1 = sign(kn[2])*sqrt(kn[2]^2 + ks[2]^2) * (mm[2] == 2) + if k1 ≈ 0 + thick_pure_bdipole(tm::BmadStandard, bunch, first(bm), L) + end + if !(ks[1] ≈ 0 && ks[2] ≈ 0) + error("BmadStandard does not support tilted combined function magnets") + end + tilde_m = massof(bunch.species)/R_to_pc(bunch.species, bunch.R_ref) + G = BeamTracking.anom(bunch.species) + return KernelCall(BmadStandardTracking.magnus_combined_func!, (0, k0, k1, tilde_m, G, L)) +end + +@inline function thick_bend_bdipole(tm::BmadStandard, bunch, bendparams, bdict, L) + if any(b -> (b.order > 2 || b.order == 0), bdict) + error("BmadStandard does not support thick combined dipoles with higher order multipoles") + end + g = bendparams.g_ref + kn, ks = get_strengths(bdict, L, bunch.R_ref) + mm = bdict.order + k0 = sign(kn[1])*sqrt(kn[1]^2 + ks[1]^2) * (mm[1] == 1) + k1 = sign(kn[2])*sqrt(kn[2]^2 + ks[2]^2) * (mm[2] == 2) + if k1 ≈ 0 + thick_bend_pure_bdipole(tm::BmadStandard, bunch, bendparams, first(bm), L) + end + if !(ks[1] ≈ 0 && ks[2] ≈ 0) + error("BmadStandard does not support tilted combined function magnets") + end + tilde_m = massof(bunch.species)/R_to_pc(bunch.species, bunch.R_ref) + G = BeamTracking.anom(bunch.species) + return KernelCall(BmadStandardTracking.magnus_combined_func!, (g, k0, k1, tilde_m, G, L)) +end + +@inline function bend_fringe(tm::BmadStandard, bunch, bendparams, bm1, L, upstream) + e = upstream * bendparams.e1 + (1 - upstream) * bendparams.e2 + g = bendparams.g_ref + K0 = get_strengths(bm1, L, bunch.R_ref) + G = BeamTracking.anom(bunch.species) + βγ0 = R_to_beta_gamma(bunch.species, bunch.R_ref) + return KernelCall(BmadStandardTracking.hwang_edge!, (e, g, K0, 0, G, βγ0, upstream)) +end + +@inline function bend_fringe(tm::BmadStandard, bunch, bendparams, bm1, bm2, L, upstream) + e = upstream * bendparams.e1 + (1 - upstream) * bendparams.e2 + g = bendparams.g_ref + K0 = get_strengths(bm1, L, bunch.R_ref) + K1 = get_strengths(bm2, L, bunch.R_ref) + G = BeamTracking.anom(bunch.species) + βγ0 = R_to_beta_gamma(bunch.species, bunch.R_ref) + return KernelCall(BmadStandardTracking.hwang_edge!, (e, g, K0, K1, G, βγ0, upstream)) +end + +@inline function thick_pure_bquadrupole(tm::BmadStandard, bunch, bm2, L) + Kn1, Ks1 = get_strengths(bm2, L, bunch.R_ref) + K1 = sign(Kn1)*sqrt(Kn1^2 + Ks1^2) + if abs(K1) < 1e-10 + tilde_m, gamsqr_0, β0 = ExactTracking.drift_params(bunch.species, bunch.R_ref) + return KernelCall(BmadStandardTracking.magnus_drift!, (β0, gamsqr_0, tilde_m, L)) + else + βγ0 = R_to_beta_gamma(bunch.species, bunch.R_ref) + tilde_m = 1/βγ0 + G = BeamTracking.anom(bunch.species) + return KernelCall(BmadStandardTracking.magnus_quadrupole!, (K1, βγ0, tilde_m, G, L)) + end +end + +@inline function thick_pure_bmultipole(tm::BmadStandard, bunch, bmn, L) + if bmn.order == 3 # Pure sextupole (No tilt) + Kn2, Ks2 = get_strengths(bmn, L, bunch.R_ref) + K2 = sign(Kn2)*sqrt(Kn2^2 + Ks2^2) + tilde_m, gamsqr_0, β0 = ExactTracking.drift_params(bunch.species, bunch.R_ref) + G = BeamTracking.anom(bunch.species) + return KernelCall(BmadStandardTracking.magnus_thick_sextupole!, (K2, β0, gamsqr_0, tilde_m, G, L)) + + elseif bmn.order == 4 # Pure octupole (No spin) + Kn3, Ks3 = get_strengths(bmn, L, bunch.R_ref) + tilde_m, gamsqr_0, β0 = ExactTracking.drift_params(bunch.species, bunch.R_ref) + return KernelCall(BmadStandardTracking.thick_octupole!, (Kn3, Ks3, β0, gamsqr_0, tilde_m, L)) + + else + error("BmadStandard does not support thick multipoles of order greater than octupole") + end + +end \ No newline at end of file diff --git a/ext/BeamTrackingBeamlinesExt/unpack.jl b/ext/BeamTrackingBeamlinesExt/unpack.jl index 44d8db4e..2eac1b98 100644 --- a/ext/BeamTrackingBeamlinesExt/unpack.jl +++ b/ext/BeamTrackingBeamlinesExt/unpack.jl @@ -14,9 +14,10 @@ function _track!( bm = deval(ele.BMultipoleParams) pp = deval(ele.PatchParams) dp = deval(ele.ApertureParams) + rf = ele isa BitsLineElement ? nothing : deval(ele.RFParams) # Function barrier - universal!(i, coords, tm, bunch, L, ap, bp, bm, pp, dp; kwargs...) + universal!(i, coords, tm, bunch, L, ap, bp, bm, pp, dp, rf; kwargs...) end # Step 2: Push particles through ----------------------------------------- @@ -30,7 +31,8 @@ function universal!( bendparams, bmultipoleparams, patchparams, - apertureparams; + apertureparams, + rfparams; kwargs... ) kc = KernelChain(Val{1}()) @@ -133,6 +135,9 @@ function universal!( kc = push(kc, @inline(bmultipole(tm, bunch, bmultipoleparams, L))) end end + elseif isactive(rfparams) # RF cavity + # RF cavity + kc = push(kc, @inline(cavity(tm, bunch, rfparams, L))) elseif !(L ≈ 0) kc = push(kc, @inline(drift(tm, bunch, L))) end @@ -171,6 +176,7 @@ end @inline thick_pure_bmultipole(tm, bunch, bmk, L) = error("Undefined for tracking method $tm") @inline thick_bmultipole(tm, bunch, bmultipoleparams, L) = error("Undefined for tracking method $tm") +@inline cavity(tm, bunch, rfparams, L) = error("Undefined for tracking method $tm") # === Elements with curving coordinate system "bend" === # # "Bend" means ONLY a coordinate system curvature through the element. @@ -221,3 +227,6 @@ end @inline bend_pure_bmultipole(tm, bunch, bendparams, bmk, L) = L ≈ 0 ? thin_bend_pure_bmultipole(tm, bunch, bendparams, bmk) : thick_bend_pure_bmultipole(tm, bunch, bendparams, bmk, L) @inline bend_bmultipole(tm, bunch, bendparams, bmultipoleparams, L) = L ≈ 0 ? thin_bend_bmultipole(tm, bunch, bendparams, bmultipoleparams) : thick_bend_bmultipole(tm, bunch, bendparams, bmultipoleparams, L) +# === Fringe fields === # +@inline bend_fringe(tm, bunch, bendparams, bm1, L, upstream) = drift(tm, bunch, L) +@inline bend_fringe(tm, bunch, bendparams, bm1, bm2, L, upstream) = drift(tm, bunch, L) diff --git a/src/BeamTracking.jl b/src/BeamTracking.jl index 60962207..83669784 100644 --- a/src/BeamTracking.jl +++ b/src/BeamTracking.jl @@ -25,6 +25,7 @@ export Bunch, State, ParticleView, sincu, sinhcu, sincuc, expq, quat_mul, atan2 export LinearTracking, Linear export ExactTracking, Exact export IntegrationTracking, SplitIntegration, DriftKick, BendKick, SolenoidKick, MatrixKick +export BmadStandardTracking, BmadStandard export track! include("utils.jl") @@ -36,6 +37,7 @@ include("kernel.jl") include("modules/ExactTracking.jl") #; TRACKING_METHOD(::ExactTracking) = Exact include("modules/LinearTracking.jl") #; TRACKING_METHOD(::LinearTracking) = Linear include("modules/IntegrationTracking.jl") #; TRACKING_METHOD(::LinearTracking) = SplitIntegration, DriftKick, BendKick, SolenoidKick, MatrixKick +include("modules/BmadStandardTracking.jl") #; TRACKING_METHOD(::BmadStandardTracking) = BmadStandard # Empty tracking method to be imported+implemented by package extensions diff --git a/src/modules/BmadStandardTracking.jl b/src/modules/BmadStandardTracking.jl new file mode 100644 index 00000000..46525a91 --- /dev/null +++ b/src/modules/BmadStandardTracking.jl @@ -0,0 +1,780 @@ +struct BmadStandard end + +module BmadStandardTracking +using ..GTPSA, ..BeamTracking, ..StaticArrays, ..KernelAbstractions +using ..BeamTracking: XI, PXI, YI, PYI, ZI, PZI, @makekernel, Coords, quat_mul! +const TRACKING_METHOD = BmadStandard + + +@makekernel fastgtpsa=true function bmad_cavity!(i, coords::Coords, V, wave_number, φ0, β0, gamsqr_0, tilde_m, p0c, L) + v = coords.v + z = v[i, ZI] + # ============= Kick - Drift - Kick scheme ============= + # First kick + rel_p = 1 + v[i, PZI] + φ = φ0 - 2π * wave_number * v[i, ZI] * sqrt(rel_p^2 + tilde_m^2) / (rel_p * sqrt(1 + tilde_m^2)) - π + dE = V * sin(φ) / 2 + pc = rel_p * p0c + E_old = sqrt(pc^2 + (tilde_m * p0c)^2) + tmp = (2 * E_old * dE + dE^2) / pc^2 + v[i, PZI] += rel_p * tmp / (1 + sqrt(1 + tmp)) + v[i, ZI] *= (1 + v[i, PZI]) * E_old / (rel_p * (E_old + dE)) + + # Drift + z = v[i, ZI] + ExactTracking.exact_drift!(i, coords::Coords, β0, gamsqr_0, tilde_m, L) + rel_p = 1 + v[i, PZI] + φ -= 2π * wave_number * (v[i, ZI] - z) * sqrt(rel_p^2 + tilde_m^2) / (rel_p * sqrt(1 + tilde_m^2)) + + # Second kick + dE = V * sin(φ) / 2 + pc = rel_p * p0c + E_old = sqrt(pc^2 + (tilde_m * p0c)^2) + tmp = (2 * E_old * dE + dE^2) / pc^2 + v[i, PZI] += rel_p * tmp / (1 + sqrt(1 + tmp)) + v[i, ZI] *= (1 + v[i, PZI]) * E_old / (rel_p * (E_old + dE)) +end + +# Drift - no spin rotation +@makekernel fastgtpsa=true function magnus_drift!(i, coords::Coords, β0, gamsqr_0, tilde_m, L) + ExactTracking.exact_drift!(i, coords::Coords, β0, gamsqr_0, tilde_m, L) +end + +# Solenoid -- second order Magnus expansion +@makekernel fastgtpsa=true function magnus_solenoid!(i, coords::Coords, Ks, β0, gamsqr_0, tilde_m, G, L) + v = coords.v + + if !isnothing(coords.q) + # Constants + rel_p = 1 + v[i,PZI] + γ = sqrt(1 + (rel_p / tilde_m)^2) + χ = 1 + G * γ + ξ = G * (γ - 1) + ι = 2 * γ - 1 + + # Compute effective longitudinal momentum + pr = sqrt(rel_p^2 - (v[i,PXI] + v[i,YI]*Ks/2)^2 - (v[i,PYI] - v[i,XI]*Ks/2)^2) + θ = Ks * L / pr + c = cos(θ) + s = sin(θ) + + # --- Spin rotation coefficients --- + # A: x-plane spin rotation + a1 = 2 * pr * (c - 1) * v[i,PYI] + a2 = -2 * pr * s * v[i,PXI] + a3 = Ks^2 * L * v[i,YI] + a4 = -Ks * (2 * L * v[i,PXI] + (c - 1) * pr * v[i,XI] + pr * s * v[i,YI]) + A = (a1 + a2 + a3 + a4) * ξ / (8.0 * rel_p^2) + + # B: y-plane spin rotation + b1 = Ks * L * (2 * v[i,PYI] + Ks * v[i,XI]) + b2 = pr * (2 * (c - 1) * v[i,PXI] + 2 * s * v[i,PYI] + c * Ks * v[i,YI] - Ks * (s * v[i,XI] + v[i,YI])) + B = -(b1 + b2) * ξ / (8.0 * rel_p^2) + + # CC: z-plane spin rotation + cc1 = 4.0 * pr * (v[i,PXI]^2 + v[i,PYI]^2) * s * (1 + G * ι) + cc2 = Ks^3 * L * (v[i,XI]^2 + v[i,YI]^2) * (1 + G * ι) + cc3 = -Ks^2 * pr * s * (v[i,XI]^2 + v[i,YI]^2) * (1 + G * ι) + cc4 = 4.0 * Ks * ((c - 1) * pr * (v[i,PXI] * v[i,XI] + v[i,PYI] * v[i,YI]) * (1 + G * ι) + + L * (v[i,PXI]^2 + v[i,PYI]^2 + 4.0 * rel_p^2 + G * (4.0 * rel_p^2 + (v[i,PXI]^2 + v[i,PYI]^2) * ι))) + cc5 = 2 * Ks^2 * L * ( + v[i,XI] * (-2 * v[i,PYI] + 2 * v[i,PXI] * s + Ks * v[i,XI]) + + 2 * (v[i,PXI] + v[i,PYI] * s) * v[i,YI] + + Ks * v[i,YI]^2 + ) + cc6 = pr * ( + 4.0 * s * (v[i,PXI]^2 + v[i,PYI]^2) + + 4.0 * Ks * v[i,PYI] * (s * v[i,XI] - 2 * v[i,YI]) - + 4.0 * Ks * v[i,PXI] * (2 * v[i,XI] + s * v[i,YI]) - + 3.0 * Ks^2 * s * (v[i,XI]^2 + v[i,YI]^2) + ) + cc7 = c * Ks * ( + 8.0 * pr * (v[i,PXI] * v[i,XI] + v[i,PYI] * v[i,YI]) + + L * (-4.0 * (v[i,PXI]^2 + v[i,PYI]^2) + Ks^2 * (v[i,XI]^2 + v[i,YI]^2)) + ) + CC = (cc1 + cc2 + cc3 + cc4) / (32.0 * rel_p^3) + CC -= (cc5 + cc6 + cc7) * pr * ξ^2 / (64.0 * rel_p^4) + + # Quaternion update for spin + ζ = sqrt(A^2 + B^2 + CC^2) + sc = sincu(ζ) + quat_mul!(@SVector[-cos(ζ), A*sc, B*sc, CC*sc], @view coords.q[i,:]) + end + + # Update coordinates + ExactTracking.exact_solenoid!(i, coords::Coords, Ks, β0, gamsqr_0, tilde_m, L) +end + +# SBend +@makekernel fastgtpsa=true function magnus_sbend!(i, coords::Coords, e1, e2, g, K0, w, w_inv, γ0, βγ0, G, L) + + if !isnothing(coords.q) + v = coords.v + rel_p = 1 + v[i,PZI] + γ = sqrt(1 + (βγ0 * rel_p)^2) + χ = 1 + G * γ + ξ = G * (γ - 1) + kx = g * K0 + ωx = sqrt(abs(kx))/sqrt(rel_p) + xc = abs(kx) > 0 ? (g*rel_p - K0)/kx : 0.0 + xd = v[i,XI] - xc + if kx > 0 + cx = cos(ωx*L) + sx = sin(ωx*L)/ωx + cx2 = cos(ωx*L/2) + sx2 = sin(ωx*L/2) + τx = -1 + elseif kx < 0 + cx = cosh(ωx*L) + sx = sinh(ωx*L)/ωx + cx2 = cosh(ωx*L/2) + sx2 = sinh(ωx*L/2) + τx = 1 + else + cx = 1 + sx = 0.0 + τx = 0.0 + end + # Compute auxiliary variables for spin rotation + υ = v[i, PYI]^2 * ξ - rel_p^2 * χ + pry = sqrt(rel_p^2 - v[i, PYI]^2) + + if g == 0 + # Straight dipole + A = -0.5 * K0 * v[i, PYI] * ξ * (L * v[i, PXI] - 0.5 * K0 * L^2) / (rel_p^2 * pry) + B = -0.5 * K0 * (L * v[i, PXI]^2 - K0 * L^2 * v[i, PXI] + K0^2 * L^3 / 3.0) * υ / (2 * rel_p^2 * pry^3) - + 0.5 * K0 * L * υ / (rel_p^2 * pry) + CC = -0.5 * K0 * v[i, PYI] * ξ * L / (rel_p^2) + elseif K0 == 0 + # Bend no field + A = 0.0 + B = -0.5 * g * L + CC = 0.0 + else + # Bend with field + σ = rel_p * sx * (v[i, XI] - 3.0 * xc) * xd + + L * (-(rel_p * xd * (v[i, XI] - (2 + cx) * xc)) + + v[i, PXI] * (sx * xc + τx * v[i, PXI] / (rel_p * ωx^2))) - + v[i, PXI] * (v[i, PXI] * sx + 2 * (cx - 1) * rel_p * xc) * τx / (rel_p * ωx^2) + + # A: x-plane spin rotation + A = -K0 * v[i, PYI] * sx2 * ξ * + (g * v[i, PXI] * sx * ωx + rel_p * (2 + g * (v[i, XI] * (1 + cx) + xc * (1 - cx))) * ωx) * + (cx2 * v[i, PXI] + rel_p * sx2 * xd * τx * ωx) / + (2 * rel_p^3 * ωx^2 * pry) + + # B: y-plane spin rotation + B = -0.5 * K0 * L * υ / (rel_p^2 * pry) + B += -g * K0 * υ * (sx * xd + L * xc + (cx - 1) * v[i, PXI] * τx / (rel_p * ωx^2)) / (2 * rel_p^2 * pry) + B += -K0 * υ * ( + cx^2 * v[i, PXI] * rel_p * xd + + v[i, PXI] * rel_p * xd * τx * (-τx + sx^2 * ωx^2) + + L * (v[i, PXI]^2 - rel_p^2 * xd^2 * τx * ωx^2) + + cx * sx * (v[i, PXI]^2 + rel_p^2 * xd^2 * τx * ωx^2) + ) / (8.0 * rel_p^2 * pry^3) + B += σ * g * K0^2 * v[i, PYI]^2 * ξ^2 / (4.0 * rel_p^4 * pry) + B += -0.5 * g * L + + # CC: z-plane spin rotation + CC = -K0 * v[i, PYI] * ξ * (L + g * sx * xd + g * L * xc + (τx * (cx - 1) * g * v[i, PXI] / (rel_p * ωx^2))) / (2 * rel_p^2) + CC += -σ * g * K0^2 * v[i, PYI] * ξ * υ / (4.0 * rel_p^4 * pry^2) + end + + ζ = sqrt(A^2 + B^2 + CC^2) + sc = sincu(ζ) + quat_mul!(@SVector[-cos(ζ), A*sc, B*sc, CC*sc], @view coords.q[i,:]) + end + + # Update coordinates + ExactTracking.exact_bend!(i, coords::Coords, e1, e2, g * L, g, K0, w, w_inv, 1/βγ0, βγ0/γ0, L)end + +# Quadrupole +@makekernel fastgtpsa=true function magnus_quadrupole!(i, coords::Coords, K1, βγ0, tilde_m, G, L) + v = coords.v + rel_p = 1 + v[i,PZI] + if !isnothing(coords.q) + γ = sqrt(1 + (βγ0 * rel_p)^2) + χ = 1 + G * γ + ξ = G * (γ - 1) + k1_norm = K1 / rel_p + ω = sqrt(abs(k1_norm)) + s = sin(ω*L) + sh = sinh(ω*L) + c = cos(ω*L) + ch = cosh(ω*L) + + if k1_norm > 0 + cx = c + sx = s/ω + cy = ch + sy = sh/ω + τy = 1 + else + cy = c + sy = s/ω + cx = ch + sx = sh/ω + τy = -1 + end + + # Compute spin rotation coefficients for quadrupole + # A: x-plane spin rotation + A = 0.5 * χ * k1_norm * ( + sy * v[i, YI] + + τy * (cy - 1) / (ω^2) * v[i, PYI] / rel_p + ) + + # B: y-plane spin rotation + B = 0.5 * χ * k1_norm * ( + sx * v[i, XI] + + τy * (1 - cx) / (ω^2) * v[i, PXI] / rel_p + ) + + # CC: z-plane spin rotation + CC = v[i, PXI] * sx * (v[i, PYI] * sy + cy * rel_p * v[i, YI]) + + rel_p * v[i, XI] * (cx * v[i, PYI] * sy + rel_p * (cx * cy - 1) * v[i, YI]) + CC *= (-0.5 * ξ * k1_norm / rel_p^2) + CC += 0.25 * χ^2 / rel_p^2 * ( + (cx - cy) * v[i, PXI] * v[i, PYI] + + τy * ω^2 * ( + v[i, PXI] * v[i, PYI] * sx * sy + + v[i, PXI] * rel_p * (cy * sx - sy) * v[i, YI] + + rel_p * v[i, XI] * ( + v[i, PYI] * (cx * sy - sx) + + (cx * cy - 1) * rel_p * v[i, YI] + ) + ) + ) + + # Quaternion update for spin + ζ = sqrt(A^2 + B^2 + CC^2) + sc = sincu(ζ) + quat_mul!(@SVector[-cos(ζ), A*sc, B*sc, CC*sc], @view coords.q[i,:]) + end + + # Update coordinates + IntegrationTracking.quadrupole_matrix!(i, coords::Coords, K1, L) + # beta != beta_0 correction + rel_e = sqrt(rel_p^2 + tilde_m^2) + v[i,ZI] += L * ( tilde_m^2 * v[i,PZI] * (2 + v[i,PZI]) / ( rel_e * ( rel_p * sqrt(1 + tilde_m^2) + rel_e ) ) ) +end + +@makekernel fastgtpsa=true function magnus_combined_func!(i, coords::Coords, g, k0, k1, tilde_m, G, L) + v = coords.v + rel_p = 1 + v[i, PZI] + inv_rel_p = 1 / rel_p + kx = k1 + g * k0 + xc = (g * rel_p - k0) / kx + + ωx = sqrt(abs(kx)) * sqrt(inv_rel_p) + ωy = sqrt(abs(k1)) * sqrt(inv_rel_p) + + + arg = ωx * L + if arg < 1e-6 + sx = (1 - sign(kx) * arg^2 / 6) * L # s_x + cx = 1 - sign(kx) * arg^2 / 2 # c_x + z2 = g * L^2 / (2 * rel_p) # z2 + elseif kx > 0 + sx = sin(arg) / ωx # s_x + cx = cos(arg) # c_x + z2 = -sign(kx) * g * (1 - cx) / (rel_p * ωx^2)# z2 + else + sx = sinh(arg) / ωx # s_x + cx = cosh(arg) # c_x + z2 = -sign(kx) * g * (1 - cx) / (rel_p * ωx^2)# z2 + end + + arg = ωy * L + if arg < 1e-6 + sy = (1 + sign(k1) * arg^2 / 6) * L # s_y + cy = 1 + sign(k1) * arg^2 / 2 # c_y + elseif k1 < 0 + sy = sin(arg) / ωy # s_y + cy = cos(arg) # c_y + else + sy = sinh(arg) / ωy # s_y + cy = cosh(arg) # c_y + end + + x0 = v[i, XI] + px0 = v[i, PXI] + y0 = v[i, YI] + py0 = v[i, PYI] + + + if !isnothing(coords.q) + γ = sqrt(1 + (rel_p / tilde_m)^2) + χ = 1 + G * γ + ξ = G * (γ - 1) + ν = 2 * (1 + G) - χ + + η = ωx^2 + ωy^2 + μ = ωx^2 - ωy^2 + + c = cx * cy - 1 + xd = x0 - xc + + a1 = b1 = cc1 = a2 = b2 = cc2 = 0.0 + + if kx > 0 && k1 > 0 + println("kx > 0 && k1 > 0") + a1 = k0 * ξ * ωy^2 * ( + py0 * ωx^2 * (-cy * px0 * sx - c * rel_p * xd * ωx) - + (cx * px0 * py0 * sy + c * px0 * rel_p * y0 + px0 * rel_p * sx * sy * y0 * ωx^2 - + rel_p * xd * (py0 * sx * sy + rel_p * (cy * sx - cx * sy) * y0) * ωx^3) * ωy^2 + ) + a1 += k1 * χ * ( + (cy - 1) * py0 * rel_p * (1 + g * xc) * η + + g * py0 * (px0 * (cy * sx - cx * sy) + cx * cy * rel_p * xd + rel_p * xd * (sx * sy * ωx^2 - 1)) * ωy^2 + + rel_p * y0 * ωy^2 * ( + rel_p * sy * η + + g * px0 * (sx * sy * ωy^2 - c) + + g * rel_p * (sy * xc * η + cy * sx * xd * ωx^2 + cx * sy * xd * ωy^2) + ) + ) + a1 *= 0.5 / (rel_p^3 * η * ωy^2) + b1 = 2 * g * px0 * ( + k1 * L * px0 - 2 * (cx - 1) * k0 * rel_p + + k1 * (rel_p * (x0 + 3 * xc) - cx * (px0 * sx + 4 * rel_p * xc) - cx^2 * rel_p * xd) + ) * χ + b1 += k0 * py0 * (L * py0 + cy * py0 * sy + (cy^2 - 1) * y0) * ν * ωx^2 + b1 -= 2 * g * rel_p * ( + -sx * xd * (2 * k0 * rel_p + k1 * px0 * sx + k1 * rel_p * (4 * xc + cx * xd)) * χ + + L * rel_p * (2 * rel_p - (2 * k0 * xc + k1 * (x0^2 - 2 * x0 * xc + 3 * xc^2)) * χ) + ) * ωx^2 + b1 += 4 * k1 * rel_p * χ * (px0 * (1 - cx) + rel_p * (L * xc + sx * xd) * ωx^2) + b1 += k0 * ωx^2 * ( + L * (4 + px0^2) * χ + χ * ( + 8 * L * (rel_p - 1) + 4 * L * (rel_p - 1)^2 + + cx * px0^2 * sx + (cx^2 - 1) * px0 * rel_p * xd * ωx - + px0 * rel_p * sx^2 * xd * ωx^3 + rel_p^3 * (L - cx * sx) * xd^2 * ωx^4 + ) - rel_p^2 * (L - cy * sy) * y0^2 * ν * ωy^2 + + py0 * y0 * ν * ((cy^2 - 1) * (rel_p - 1) + rel_p * sy^2 * ωy^2) + ) + b1 /= 8 * rel_p^3 * ωx^2 + cc1 = k0 * rel_p * (1 + g * xc) * (py0 * sy + (cy - 1) * rel_p * y0) * η + cc1 += kx * py0 * (-c * px0 + cy * rel_p * sx * xd * ωx^2) + cc1 += kx * ( + px0 * py0 * sx * sy + px0 * rel_p * (cy * sx - cx * sy) * y0 + + rel_p * xd * (cx * py0 * sy + cx * cy * rel_p * y0 + rel_p * y0 * (sx * sy * ωx^2 - 1)) + ) * ωy^2 + cc1 += k1 * ( + px0 * ( + py0 * (c + sx * sy * ωx^2) + rel_p * y0 * (cy * sx * ωx^2 + cx * sy * ωy^2) + ) + + rel_p * ( + py0 * (sy * xc * η - cy * sx * xd * ωx^3 + cx * sy * xd * ωx^3) + + rel_p * y0 * ((cy - 1) * xc * η + xd * ωx^3 * (c - sx * sy * ωy^2)) + ) + ) + cc1 *= -0.5 * ξ / (rel_p^3 * η) + + a2 = rel_p * y0 * ωy^2 * ( + (px0 - rel_p * (L * xc - 2 * sy * xc + sx * xd) * η) * ωx^2 + + cx * px0 * (η - cy * ωx^2) + + px0 * (-1 + 2 * sx * sy * ωx^2) * ωy^2 + + cx * (cy * px0 + 2 * rel_p * sy * xd * ωx^2) * ωy^2 + + cy * (-px0 * η + rel_p * ωx^2 * (-L * xc * η + sx * xd * μ)) + ) + a2 += py0 * ( + -rel_p * xc * ωx^2 * (η - 2 * cy * η + ωx^2 + L * sy * η * ωy^2) + + ωy^2 * ( + -rel_p * ωx^2 * (x0 + xd - 2 * cx * cy * xd - sx * sy * xd * μ) + + px0 * (2 * cy * sx * ωx^2 - sy * (η + cx * μ)) + ) + ) + a2 *= 0.25 * k0 * kx * ξ * χ / (rel_p^4 * η * ωx^2 * ωy^2) + + b2 = k0 * k1 * (L - sy) * ξ * χ * (py0^2 - rel_p^2 * y0^2 * ωy^2) / (4 * rel_p^4 * ωy^2) + + cc2 = px0 * ( + py0 * ((cx - cy - 1) * η + cx * cy * ωy^2 + ωx^2 * (2 - cx * cy + 2 * sx * sy * ωy^2)) + + rel_p * y0 * ωy^2 * (2 * cy * sx * ωx^2 - sy * (η + cx * μ)) + ) + cc2 -= rel_p * ωx^2 * ( + L * xc * η * ((cy + 1) * py0 + rel_p * sy * y0 * ωy^2) + + py0 * (-2 * sy * (xc * η + cx * xd * ωy^2) + sx * xd * (η - cy * μ)) + + rel_p * y0 * ( + 2 * xc * ωx^2 - 2 * cy * (xc * η + cx * xd * ωy^2) + ωy^2 * (2 * x0 - sx * sy * xd * μ) + ) + ) + cc2 *= 0.25 * k1 * kx * χ^2 / (rel_p^4 * η * ωx^2 * ωy^2) + + elseif kx > 0 && k1 < 0 + if abs(μ) > 1e-20 + a1 = k0 * ξ * ωy^2 * ( + py0 * ωx^2 * (-cy * px0 * sx - c * rel_p * xd * ωx) + + (cx * px0 * py0 * sy + c * px0 * rel_p * y0 + px0 * rel_p * sx * sy * y0 * ωx^2 - + rel_p * xd * (py0 * sx * sy + rel_p * (cy * sx - cx * sy) * y0) * ωx^3) * ωy^2 + ) + a1 += k1 * χ * ( + -((cy - 1) * py0 * rel_p * (1 + g * xc) * ωx^2) + + (rel_p * y0 * (-c * g * px0 + rel_p * (sy + g * sy * xc + cy * g * sx * xd) * ωx^2) + + py0 * (-rel_p - cx * g * px0 * sy + cy * (rel_p + g * (rel_p - 1) * (xc + cx * xd) + g * (px0 * sx + xc + cx * xd)) - + g * rel_p * (xc + xd - sx * sy * xd * ωx^2))) * ωy^2 - + rel_p * sy * (rel_p + g * (rel_p - 1) * (xc + cx * xd) + g * (px0 * sx + xc + cx * xd)) * y0 * ωy^4 + ) + a1 *= 0.5 / (rel_p^3 * μ * ωy^2) + + cc1 = -k0 * rel_p * (1 + g * xc) * (py0 * sy + (cy - 1) * rel_p * y0) * μ + cc1 += kx * (c * px0 * py0 + px0 * (py0 * sx * sy + rel_p * (cy * sx - cx * sy) * y0) * ωy^2 + + rel_p * xd * (-cy * py0 * sx * ωx^2 + (cx * py0 * sy + cx * cy * rel_p * y0 + + rel_p * y0 * (-1 + sx * sy * ωx^2)) * ωy^2)) + cc1 -= k1 * ( + px0 * (py0 * (c + sx * sy * ωx^2) + rel_p * y0 * (cy * sx * ωx^2 - cx * sy * ωy^2)) + + rel_p * ( + py0 * ωx^2 * (-cy * sx * xd * ωx + sy * (xc + cx * xd * ωx)) - py0 * sy * xc * ωy^2 + + rel_p * y0 * ((cy - 1) * xc * μ + xd * ωx^3 * (c - sx * sy * ωy^2)) + ) + ) + cc1 *= 0.5 * ξ / (rel_p^3 * μ) + + a2 = -rel_p * y0 * ωy^2 * ( + (1 + cx) * (cy - 1) * px0 * ωx^2 + px0 * ((cx - 1) * (1 + cy) + 2 * sx * sy * ωx^2) * ωy^2 + + rel_p * ωx^2 * ((1 + cy) * L * xc - 2 * sy * xc + sx * xd - cy * sx * xd) * ωx^2 - + ((1 + cy) * L * xc + (1 + cy) * sx * xd - 2 * sy * (xc + cx * xd)) * ωy^2 + ) + a2 += py0 * ( + -rel_p * xc * ωx^2 * μ * (-2 + 2 * cy + L * sy * ωy^2) + + ωy^2 * (2 * cy * (px0 * sx + cx * rel_p * xd) * ωx^2 - + px0 * sy * ((1 + cx) * ωx^2 + (cx - 1) * ωy^2) + + rel_p * xd * ωx^2 * (-2 + sx * sy * η)) + ) + a2 *= 0.25 * k0 * kx * ξ * χ / (rel_p^4 * ωx^2 * ωy^2 * μ) + + cc2 = px0 * ((1 + cx) * (cy - 1) * py0 * ωx^2 + py0 * ((cx - 1) * (1 + cy) + 2 * sx * sy * ωx^2) * ωy^2 + + rel_p * y0 * ωy^2 * ((2 * cy * sx - (1 + cx) * sy) * ωx^2 - (1 - cx) * sy * ωy^2)) + cc2 -= rel_p * ωx^2 * ( + py0 * (2 * sy * xc + (cy - 1) * sx * xd) * ωx^2 + py0 * ((1 + cy) * sx * xd - 2 * sy * (xc + cx * xd)) * ωy^2 + + L * xc * μ * (-((1 + cy) * py0) + rel_p * sy * y0 * ωy^2) + + rel_p * y0 * (2 * (cy - 1) * xc * ωx^2 + + (x0 + xc - 2 * cy * xc + xd - 2 * cx * cy * xd - sx * sy * xd * ωx^2) * ωy^2 - + sx * sy * xd * ωy^4) + ) + cc2 *= 0.25 * k1 * kx * χ^2 / (rel_p^4 * ωx^2 * ωy^2 * μ) + println("We are here") + else + error("μ ≈ 0 should not happen for k0 > 0 and k1 < 0") + end + elseif kx < 0 && k1 > 0 + if abs(μ) > 1e-20 + a1 = k0 * ξ * ωy^2 * ( + py0 * ωx^2 * (-cy * px0 * sx - c * rel_p * xd * ωx) + + (cx * px0 * py0 * sy + c * px0 * rel_p * y0 - px0 * rel_p * sx * sy * y0 * ωx^2 + + rel_p * xd * (py0 * sx * sy + rel_p * (cy * sx - cx * sy) * y0) * ωx^3) * ωy^2 + ) + a1 += k1 * χ * ( + (cy - 1) * py0 * rel_p * (1 + g * xc) * ωx^2 + + (rel_p * y0 * (c * g * px0 + rel_p * (sy + g * sy * xc + cy * g * sx * xd) * ωx^2) + + py0 * (rel_p - cy * (rel_p + g * px0 * sx + g * rel_p * xc + cx * g * rel_p * xd) + + g * (cx * px0 * sy + rel_p * (xc + xd + sx * sy * xd * ωx^2)))) * ωy^2 - + rel_p * sy * (rel_p + g * px0 * sx + g * rel_p * xc + cx * g * rel_p * xd) * y0 * ωy^4 + ) + a1 *= 0.5 / (rel_p^3 * μ * ωy^2) + + cc1 = kx * py0 * (-px0 * c - cy * rel_p * sx * xd * ωx^2) + cc1 += kx * ( + px0 * py0 * sx * sy + px0 * rel_p * (cy * sx - cx * sy) * y0 + + rel_p * xd * (cx * py0 * sy + cx * cy * rel_p * y0 - rel_p * y0 * (1 + sx * sy * ωx^2)) + ) * ωy^2 + cc1 -= k0 * rel_p * (1 + g * xc) * (py0 * sy + (cy - 1) * rel_p * y0) * μ + cc1 += k1 * ( + px0 * (py0 * (c - sx * sy * ωx^2) - rel_p * y0 * (cy * sx * ωx^2 - cx * sy * ωy^2)) + + rel_p * ( + py0 * ωx^2 * (cy * sx * xd * ωx - sy * (xc + cx * xd * ωx)) + py0 * sy * xc * ωy^2 + + rel_p * y0 * (-(cy - 1) * xc * μ + xd * ωx^3 * (c - sx * sy * ωy^2)) + ) + ) + cc1 *= 0.5 * ξ / (rel_p^3 * μ) + + a2 = rel_p * y0 * ωy^2 * ( + (1 + cx) * (cy - 1) * px0 * ωx^2 + px0 * ((cx - 1) * (1 + cy) - 2 * sx * sy * ωx^2) * ωy^2 + + rel_p * ωx^2 * ( + 2 * sy * xc + (cy - 1) * sx * xd) * ωx^2 + ((1 + cy) * sx * xd - 2 * sy * (xc + cx * xd)) * ωy^2 - + (1 + cy) * L * xc * μ + ) + a2 += py0 * ( + -rel_p * xc * ωx^2 * μ * (2 - 2 * cy - L * sy * ωy^2) + + ωy^2 * ( + -2 * cy * (px0 * sx + cx * rel_p * xd) * ωx^2 + + px0 * sy * ((1 + cx) * ωx^2 + (cx - 1) * ωy^2) + + rel_p * xd * ωx^2 * (2 + sx * sy * η) + ) + ) + a2 *= 0.25 * k0 * kx * ξ * χ / (rel_p^4 * ωx^2 * ωy^2 * μ) + + cc2 = px0 * ( + (1 + cx) * (cy - 1) * py0 * ωx^2 + ((cx - 1) * (1 + cy) * py0 + (-2 * py0 * sx * sy + + rel_p * (-2 * cy * sx + sy + cx * sy) * y0) * ωx^2) * ωy^2 + (-1 + cx) * rel_p * sy * y0 * ωy^4 + ) + cc2 += rel_p * ωx^2 * ( + py0 * (2 * sy * xc + (cy - 1) * sx * xd) * ωx^2 + py0 * ((1 + cy) * sx * xd - 2 * sy * (xc + cx * xd)) * ωy^2 - + L * xc * μ * ((1 + cy) * py0 + rel_p * sy * y0 * ωy^2) + + rel_p * y0 * ( + 2 * (cy - 1) * xc * ωx^2 + + (x0 + xc - 2 * cy * xc + xd - 2 * cx * cy * xd + sx * sy * xd * ωx^2) * ωy^2 + + sx * sy * xd * ωy^4 + ) + ) + cc2 *= 0.25 * k1 * kx * χ^2 / (rel_p^4 * ωx^2 * ωy^2 * μ) + else + error("μ ≈ 0 should not happen for k0 < 0 and k1 > 0") + end + else + # kx < 0 && k1 < 0 + a1 = k0 * ξ * ωy^2 * ( + py0 * ωx^2 * (-cy * px0 * sx - c * rel_p * xd * ωx) - + (cx * px0 * py0 * sy + c * px0 * rel_p * y0 - px0 * rel_p * sx * sy * y0 * ωx^2 + + rel_p * xd * (py0 * sx * sy + rel_p * (cy * sx - cx * sy) * y0) * ωx^3) * ωy^2 + ) + a1 += k1 * χ * ( + -((cy - 1) * py0 * rel_p * (1 + g * xc) * η) + + g * py0 * (cx * px0 * sy - cy * (px0 * sx + cx * rel_p * xd) + rel_p * xd * (1 + sx * sy * ωx^2)) * ωy^2 + + rel_p * y0 * ωy^2 * ( + rel_p * sy * η + g * px0 * (c + sx * sy * ωy^2) + + g * rel_p * (sy * xc * η + cy * sx * xd * ωx^2 + cx * sy * xd * ωy^2) + ) + ) + a1 *= 0.5 / (rel_p^3 * η * ωy^2) + + b1 = 2 * g * px0 * ( + 2 * (cx - 1) * k0 * rel_p - k1 * (L * px0 - cx * px0 * sx + x0 + (rel_p - 1) * x0) + + (-3 + 4 * cx) * k1 * rel_p * xc + cx^2 * k1 * rel_p * xd + ) * χ + b1 += k0 * py0 * (L * py0 + cy * py0 * sy + (cy^2 - 1) * y0) * ν * ωx^2 + b1 -= 2 * g * rel_p * ( + -sx * xd * (2 * k0 * rel_p + k1 * px0 * sx + k1 * rel_p * (4 * xc + cx * xd)) * χ + + L * rel_p * (2 * rel_p - (2 * k0 * xc + k1 * (x0^2 - 2 * x0 * xc + 3 * xc^2)) * χ) + ) * ωx^2 + b1 += 4 * k1 * rel_p * χ * ((-1 + cx) * px0 + rel_p * (L * xc + sx * xd) * ωx^2) + b1 += k0 * ωx^2 * ( + L * (4 + px0^2) * χ + χ * ( + 8 * L * (rel_p - 1) + 4 * L * (rel_p - 1)^2 + cx * px0^2 * sx + + (cx^2 - 1) * px0 * rel_p * xd * ωx + px0 * rel_p * sx^2 * xd * ωx^3 - + rel_p^2 * (L - cx * sx) * xd^2 * ωx^4 + ) + ) + b1 += rel_p^2 * (L - cy * sy) * y0^2 * ν * ωy^2 + b1 += py0 * y0 * ν * ((cy^2 - 1) * (rel_p - 1) - rel_p * sy^2 * ωy^2) + b1 /= 8 * rel_p^3 * ωx^2 + + cc1 = -k0 * rel_p * (1 + g * xc) * (py0 * sy + (cy - 1) * rel_p * y0) * η + cc1 -= kx * ( + c * px0 * py0 + px0 * (py0 * sx * sy + rel_p * (cy * sx - cx * sy) * y0) * ωy^2 + + rel_p * xd * (cy * py0 * sx * ωx^2 + (cx * py0 * sy + cx * cy * rel_p * y0 - + rel_p * y0 * (1 + sx * sy * ωx^2)) * ωy^2) + ) + cc1 += k1 * ( + px0 * (py0 * (c - sx * sy * ωx^2) - rel_p * y0 * (cy * sx * ωx^2 + cx * sy * ωy^2)) + + rel_p * ( + -py0 * sy * xc * η + py0 * (cy * sx - cx * sy) * xd * ωx^3 - + rel_p * y0 * ((cy - 1) * xc * η + xd * ωx^3 * (c + sx * sy * ωy^2)) + ) + ) + cc1 *= 0.5 * ξ / (rel_p^3 * η) + + a2 = rel_p * y0 * ωy^2 * ( + px0 * ((1 + cx - cy) * η - cx * cy * ωx^2 + (-2 + cx * cy - 2 * sx * sy * ωx^2) * ωy^2) + + rel_p * ωx^2 * ((1 + cy) * L * xc * η + sx * xd * (η - cy * μ) - 2 * sy * (xc * η + cx * xd * ωy^2)) + ) + a2 += py0 * ( + rel_p * xc * η * ωx^2 * (-2 + 2 * cy + L * sy * ωy^2) + + ωy^2 * (2 * cy * px0 * sx * ωx^2 - px0 * sy * (η + cx * μ) + rel_p * xd * ωx^2 * (2 * c - sx * sy * μ)) + ) + a2 *= -0.25 * k0 * kx * ξ * χ / (rel_p^4 * η * ωx^2 * ωy^2) + + b2 = -k0 * k1 * (L - sy) * ξ * χ * (py0^2 + rel_p^2 * y0^2 * ωy^2) / (4 * rel_p^4 * ωy^2) + + cc2 = px0 * ( + py0 * ((1 + cx - cy) * η - cx * cy * ωx^2 + (-2 + cx * cy - 2 * sx * sy * ωx^2) * ωy^2) + + rel_p * y0 * ωy^2 * (-2 * cy * sx * ωx^2 + sy * (η + cx * μ)) + ) + cc2 += rel_p * ωx^2 * ( + L * xc * η * ((1 + cy) * py0 - rel_p * sy * y0 * ωy^2) + + py0 * (sx * xd * (η - cy * ωx^2 + cy * ωy^2) - 2 * sy * (xc * η + cx * xd * ωy^2)) - + rel_p * y0 * (2 * (cy - 1) * xc * η + xd * ωy^2 * (2 * c - sx * sy * μ)) + ) + cc2 *= 0.25 * k1 * kx * χ^2 / (rel_p^4 * η * ωx^2 * ωy^2) + end + + # Final sum + A = a1 + a2 + B = b1 + b2 + CC = cc1 + cc2 + + ζ = sqrt(A^2 + B^2 + CC^2) + sc = sincu(ζ) + quat_mul!(@SVector[-cos(ζ), A*sc, B*sc, CC*sc], @view coords.q[i,:]) + end + + x0 -= xc + + # Update transverse + v[i, XI] = cx * x0 + sx * px0 * inv_rel_p + xc + v[i, PXI] = -sign(kx) * ωx^2 * rel_p * sx * x0 + cx * px0 + v[i, YI] = cy * y0 + sy * py0 * inv_rel_p + v[i, PYI] = sign(k1) * ωy^2 * rel_p * sy * y0 + cy * py0 + + # Longitudinal update + rel_e = sqrt(rel_p^2 + tilde_m^2) + v[i, ZI] += L * ( tilde_m^2 * v[i,PZI] * (2 + v[i,PZI]) / ( rel_e * ( rel_p * sqrt(1 + tilde_m^2) + rel_e ) ) ) + + (-g * xc * L) + + (-g * sx) * x0 + + z2 * px0 + + (-sign(kx) * ωx^2 * (L - cx * sx) / 4) * x0^2 + + (sign(kx) * ωx^2 * sx^2 / (2 * rel_p)) * x0 * px0 + + (-(L + cx * sx) / (4 * rel_p^2)) * px0^2 + + (sign(k1) * ωy^2 * (L - cy * sy) / 4) * y0^2 + + (-sign(k1) * ωy^2 * sy^2 / (2 * rel_p)) * y0 * py0 + + (-(L + cy * sy) / (4 * rel_p^2)) * py0^2 +end + +@makekernel fastgtpsa=true function thin_snake!(i, coords::Coords, axis, angle) + quat_mul!(SVector{4}([cos(angle/2), sin(angle/2)*axis...]), @view coords.q[i,:]) +end + +# Sextupole +@makekernel fastgtpsa=true function magnus_thick_sextupole!(i, coords::Coords, K2, β0, gamsqr_0, tilde_m, G, L) + mm = [3] + K2N = [K2 * L/2] + K2S = [0] + + # First kick of a kick-drift-kick split + ExactTracking.multipole_kick!(i, coords, mm, K2N, K2S, -1) + + # ========== Magnus spin rotation ========== + if !isnothing(coords.q) + v = coords.v + rel_p = 1 + v[i,PZI] + γ = sqrt(1 + (rel_p / tilde_m)^2) + χ = 1 + G * γ + ξ = G * (γ - 1) + pl2 = rel_p^2 - v[i,PXI]^2 - v[i,PYI]^2 + if pl2 <= 0 + b.state[i] = State.Lost + @warn "Particle lost in sextupole (transverse momentum too large)" + else + pl = sqrt(pl2) + + # Common terms for reuse + L2 = L^2 + pl2 = pl^2 + pl3 = pl^3 + rel_p2 = rel_p^2 + + # Convenience pointers + Px = v[i,PXI] + Py = v[i,PYI] + X = v[i,XI] + Y = v[i,YI] + + # A + term_a1 = L2 * (-3Px^2 * Py + Py^3) + term_a2 = 3L * pl * (2Px * Py * X + Y * (Px^2 - Py^2)) + term_a3 = 3pl2 * (2Px * X * Y + Py * (X^2 - Y^2)) + term_a4 = rel_p2 * χ * (L * Py * (2L * Px + 3pl * X) + 3pl * (L * Px + 2pl * X) * Y) + + A = Px * ξ * (term_a1 - term_a2 - term_a3) + term_a4 + A *= K2 * L / (12.0 * pl3 * rel_p2) + + # B + term_b1 = L2 * (-3Px^2 * Py + Py^3) + term_b2 = 3L * pl * (2Px * Py * X + Y * (Px^2 - Py^2)) + term_b3 = 3pl2 * (2Px * X * Y + Py * (X^2 - Y^2)) + term_b4 = rel_p2 * χ * (L2 * (Px^2 - Py^2) + 3pl2 * (X^2 - Y^2) + 3L * pl * (Px * X - Py * Y)) + + B = Py * ξ * (term_b1 - term_b2 - term_b3) + term_b4 + B *= K2 * L / (12.0 * pl3 * rel_p2) + + # CC + term_c1 = L2 * (-3Px^2 * Py + Py^3) + term_c2 = 3L * pl * (2Px * Py * X + Y * (Px^2 - Py^2)) + term_c3 = 3pl2 * (2Px * X * Y + Py * (X^2 - Y^2)) + + CC = (term_c1 - term_c2 - term_c3) * K2 * L * ξ / (12.0 * pl2 * rel_p2) + + # Quaternion update + ζ = sqrt(A^2 + B^2 + CC^2) + sc = sincu(ζ) + quat_mul!(@SVector[-cos(ζ), A*sc, B*sc, CC*sc], @view coords.q[i,:]) + end + end + # ========== End of Magnus spin rotation ========== + + # Drift and second kick of a kick-drift-kick split + ExactTracking.exact_drift!( i, coords, β0, gamsqr_0, tilde_m, L) + ExactTracking.multipole_kick!(i, coords, mm, K2N, K2S, -1) +end + + +# Octupole +@makekernel fastgtpsa=true function thick_octupole!(i, coords::Coords, K3N, K3S, β0, gamsqr_0, tilde_m, L) + mm = [4] + K3N = [K3N * L/2] + K3S = [K3S * L/2] + + # Kick-drift-kick split + ExactTracking.multipole_kick!(i, coords, mm, K3N, K3S, -1) + ExactTracking.exact_drift!( i, coords, β0, gamsqr_0, tilde_m, L) + ExactTracking.multipole_kick!(i, coords, mm, K3N, K3S, -1) +end + +@makekernel fastgtpsa=true function hwang_edge!(i, coords::Coords, e, g, k0, k1, G, βγ0, upstream) + cos_e = cos(e); + sin_e = sin(e); + tan_e = sin_e / cos_e; + sec_e = 1 / cos_e + gt = k0 * tan_e + t2 = tan_e * tan_e + gt2 = k0 * t2 + gs2 = k0 * sec_e^2 + k1_tane = k1 * tan_e + + s = 2 * upstream - 1 + + v = coords.v + v1_2 = v[i, XI] * v[i, XI] + v13 = v[i, XI] * v[i, YI] + v3_2 = v[i, YI] * v[i, YI] + e_factor = 1 / (1 + v[i, PZI]) + fg_factor = 0 + + dx = s * (-gt2 * v1_2 + gs2 * v3_2) * e_factor / 2 + + dpx = e_factor * (s * gt2 * (v[i, XI] * v[i, PXI] - v[i, YI] * v[i, PYI]) + + k1_tane * (v1_2 - v3_2) + k0 * gt * ( + 0.5 * upstream * (1 + 2 * t2) * v3_2 + + 0.25 * (s - 1) * t2 * (v1_2 + v3_2) + ) + ) + + dy = s * gt2 * v13 * e_factor + + dpy = (fg_factor * v[i, YI] - s * (gt2 * v[i, XI] * v[i, PYI] + (k0 + gt2) * v[i, PXI] * v[i, YI]) + + ((1 - upstream) * gt * gs2 - 2 * k1_tane) * v13) * e_factor + + + dz = e_factor^2 * 0.5 * (v3_2 * fg_factor + + v[i, XI]^3 * (4.0 * k1_tane - gt * gt2) / 6.0 + 0.5 * v[i, XI]*v3_2 * (-4.0 * k1_tane + gt * gs2) + + s * ((v[i, XI]^2*v[i, PXI] - 2.0 * v13*v[i, PYI]) * gt2 - v[i, PXI]*v3_2 * gs2)) + + + v[i, PXI] += dpx + gt * v[i, XI] + v[i, PYI] += dpy - gt * v[i, YI] + v[i, XI] += dx + v[i, YI] += dy + v[i, ZI] += dz + + # spin tracking should be symmetrized, and exit edge face e2 does not reproduce Bmad + if !isnothing(coords.q) + rel_p = 1 + v[i, PZI] + γ = sqrt(1 + (rel_p * βγ0)^2) + B = k0 .* [-sin_e * v[i, YI], -tan_e * v[i, XI], cos_e * v[i, YI]] + #B -= k1 * tan_e .* [v[i, XI] * v13, v[i, XI] * (v1_2 - v3_2), 0.0] + B[3] *= s + quat_mul!(TBMT_quat(G, 0, 0, γ, 1, B, coords.v, i), @view coords.q[i,:]) + end +end + +end diff --git a/src/utils.jl b/src/utils.jl index b81c577d..a080a418 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -93,6 +93,38 @@ function sincos_quaternion(x) c_out = vifelse(x > threshold, c, 1-x/2) return s_out, c_out end +""" +returns instantaneous T_BMT quaternion in a straight coordinate system for particle `i` in a bunch +""" +function TBMT_quat(G, gx, gy, γ, Bρ, B::Vector{<:Number}, bv::Matrix{<:Number}, i) + v = bv[i,:] + rel_p = 1 + v[6] + ps = sqrt(rel_p^2 - v[2]^2 - v[4]^2) + β_hat = [v[2], v[4], ps] / rel_p + B_para = (B' * β_hat) * β_hat + B_perp = B - B_para + # T-BMT precession in straight coordinate system + Ω = -((1+G*γ) * B_perp .+ (1+G) * B_para) / (Bρ * ps) + θ = norm(Ω) + Ω /= (θ + (θ==0)) # if θ==0, denominator is 1, so Ω unchanged; else Ω /= θ + return SVector(cos(θ/2), sin(θ/2)*Ω...) +end + +function quat_mul!(q1, q2) + # In-place quaternion multiplication: q2 := q1 * q2 + w1, x1, y1, z1 = q1 + w2, x2, y2, z2 = q2 + q2[1] = w1*w2 - x1*x2 - y1*y2 - z1*z2 + q2[2] = w1*x2 + x1*w2 + y1*z2 - z1*y2 + q2[3] = w1*y2 - x1*z2 + y1*w2 + z1*x2 + q2[4] = w1*z2 + x1*y2 - y1*x2 + z1*w2 +end + +function quat_inv(q) + w, x, y, z = q + norm_sq = w^2 + x^2 + y^2 + z^2 + return ([w, -x, -y, -z] ./ norm_sq) +end """ diff --git a/test/ExactTracking.jl b/test/ExactTracking.jl index cf4d5f64..5822dd92 100644 --- a/test/ExactTracking.jl +++ b/test/ExactTracking.jl @@ -897,6 +897,7 @@ zf_mn4 = [ 0., 3.140908277834687e-8, -3.1503450227072763e-8, 3.140908186274627e tilde_m = mc2/p0c gamsqr_0 = 1 + 1/tilde_m^2 beta_0 = 1/sqrt(1 + tilde_m^2) + winv = @SArray T[1 0 0; 0 1 0; 0 0 1] dt = T(4e-9) dx = T(1) dy = T(2)