Bmad standard

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

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

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)

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