Batch3

.github/workflows/CI.yml

@@ -40,6 +40,16 @@ jobs:
             arch: aarch64
           - os: macos-latest
             arch: x64
+        include:
+          - os: ubuntu-24.04-arm
+            version: 'lts'
+            arch: aarch64
+          - os: ubuntu-24.04-arm
+            version: '1.11'
+            arch: aarch64
+          - os: ubuntu-24.04-arm
+            version: '1.12'
+            arch: aarch64
     steps:
       - uses: actions/checkout@v6
       - uses: julia-actions/setup-julia@v2

.ipynb_checkpoints/Untitled-checkpoint.ipynb

@@ -0,0 +1,6 @@
+{
+ "cells": [],
+ "metadata": {},
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

Project.toml

@@ -1,12 +1,13 @@
 name = "BeamTracking"
 uuid = "8ef5c10a-4ca3-437f-8af5-b84d8af36df0"
 authors = ["mattsignorelli <mgs255@cornell.edu> and contributors"]
-version = "0.5.5"
+version = "0.5.6"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 AtomicAndPhysicalConstants = "5c0d271c-5419-4163-b387-496237733d8b"
+EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8"
 HostCPUFeatures = "3e5b6fbb-0976-4d2c-9146-d79de83f2fb0"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
@@ -21,15 +22,18 @@ Unrolled = "9602ed7d-8fef-5bc8-8597-8f21381861e8"
 
 [weakdeps]
 Beamlines = "5bb90b03-0719-46b8-8ce4-1ef3afd3cd4b"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 
 [extensions]
 BeamTrackingBeamlinesExt = "Beamlines"
+BeamTrackingCUDAExt = "CUDA"
 
 [compat]
 Accessors = "0.1.42"
 Adapt = "4.3.0"
 AtomicAndPhysicalConstants = "0.8.0"
 Beamlines = "0.8.0"
+CUDA = "5.9"
 GTPSA = "1.5.3"
 HostCPUFeatures = "0.1"
 KernelAbstractions = "0.9.35"
@@ -48,8 +52,9 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "Distributions", "JET", "GTPSA", "BenchmarkTools", "Beamlines", "StaticArrays"]
+test = ["Test", "Distributions", "JET", "GTPSA", "BenchmarkTools", "Beamlines", "StaticArrays", "LinearAlgebra"]

ext/BeamTrackingBeamlinesExt/exact.jl

@@ -76,8 +76,16 @@ end
 @inline function thick_bend_pure_bdipole(tm::Exact, bunch, bendparams, bm1, L)
   g = bendparams.g_ref
   tilt = bendparams.tilt_ref
-  e1 = bendparams.e1
-  e2 = bendparams.e2
+  if tm.fringe_at == Fringe.BothEnds || tm.fringe_at == Fringe.EntranceEnd
+    e1 = bendparams.e1
+  else
+    e1 = 0
+  end
+  if tm.fringe_at == Fringe.BothEnds || tm.fringe_at == Fringe.ExitEnd
+    e2 = bendparams.e2
+  else
+    e2 = 0
+  end
   w = rot_quaternion(0,0,-tilt)
   w_inv = inv_rot_quaternion(0,0,-tilt)
   theta = g * L

ext/BeamTrackingBeamlinesExt/utils.jl

@@ -77,31 +77,34 @@ This works for both BMultipole and BMultipoleParams. Branchless bc SIMD -> basic
 no loss in computing both but benefit of branchless.
 """
 @inline function get_strengths(bm, L, p_over_q_ref)
-  if isconcretetype(eltype(bm.n))
-    T = promote_type(eltype(bm.n),
+  bmn = getfield(bm, :n)
+  bms = getfield(bm, :s)
+  bmtilt = getfield(bm, :tilt)
+  if isconcretetype(eltype(bmn))
+    T = promote_type(eltype(bmn),
                     typeof(L), typeof(p_over_q_ref)
     )
   else
-    if bm.n isa AbstractArray
-      T = promote_type(reduce(promote_type, typeof.(bm.n)), 
-                      reduce(promote_type, typeof.(bm.s)),
-                      reduce(promote_type, typeof.(bm.tilt)),
+    if bmn isa AbstractArray
+      T = promote_type(reduce(promote_type, typeof.(bmn)), 
+                      reduce(promote_type, typeof.(bms)),
+                      reduce(promote_type, typeof.(bmtilt)),
                       typeof(L), typeof(p_over_q_ref)
       )
     else
-      T = promote_type(typeof(bm.n), 
-                      typeof(bm.s),
-                      typeof(bm.tilt),
+      T = promote_type(typeof(bmn), 
+                      typeof(bms),
+                      typeof(bmtilt),
                       typeof(L), typeof(p_over_q_ref)
       )
     end
   end
-  n = T.(make_static(bm.n))
-  s = T.(make_static(bm.s))
-  tilt = T.(make_static(bm.tilt))
-  order = bm.order
-  normalized = bm.normalized
-  integrated = bm.integrated
+  n = T.(make_static(bmn))
+  s = T.(make_static(bms))
+  tilt = T.(make_static(bmtilt))
+  order = getfield(bm, :order)
+  normalized = getfield(bm, :normalized)
+  integrated = getfield(bm, :integrated)
   np = @. n*cos(order*tilt) + s*sin(order*tilt)
   sp = @. -n*sin(order*tilt) + s*cos(order*tilt)
   np = @. ifelse(!normalized, np/p_over_q_ref, np) 
@@ -112,31 +115,34 @@ no loss in computing both but benefit of branchless.
 end
 
 @inline function get_integrated_strengths(bm, L, p_over_q_ref)
-  if isconcretetype(eltype(bm.n))
-    T = promote_type(eltype(bm.n),
+  bmn = getfield(bm, :n)
+  bms = getfield(bm, :s)
+  bmtilt = getfield(bm, :tilt)
+  if isconcretetype(eltype(bmn))
+    T = promote_type(eltype(bmn),
                     typeof(L), typeof(p_over_q_ref)
     )
   else
-    if bm.n isa AbstractArray
-      T = promote_type(reduce(promote_type, typeof.(bm.n)), 
-                      reduce(promote_type, typeof.(bm.s)),
-                      reduce(promote_type, typeof.(bm.tilt)),
+    if bmn isa AbstractArray
+      T = promote_type(reduce(promote_type, typeof.(bmn)), 
+                      reduce(promote_type, typeof.(bms)),
+                      reduce(promote_type, typeof.(bmtilt)),
                       typeof(L), typeof(p_over_q_ref)
       )
     else
-      T = promote_type(typeof(bm.n), 
-                      typeof(bm.s),
-                      typeof(bm.tilt),
+      T = promote_type(typeof(bmn), 
+                      typeof(bms),
+                      typeof(bmtilt),
                       typeof(L), typeof(p_over_q_ref)
       )
     end
   end
-  n = T.(make_static(bm.n))
-  s = T.(make_static(bm.s))
-  tilt = T.(make_static(bm.tilt))
-  order = bm.order
-  normalized = bm.normalized
-  integrated = bm.integrated
+  n = T.(make_static(bmn))
+  s = T.(make_static(bms))
+  tilt = T.(make_static(bmtilt))
+  order = getfield(bm, :order)
+  normalized = getfield(bm, :normalized)
+  integrated = getfield(bm, :integrated)
   np = @. n*cos(order*tilt) + s*sin(order*tilt)
   sp = @. -n*sin(order*tilt) + s*cos(order*tilt)
   np = @. ifelse(!normalized, np/p_over_q_ref, np) 
@@ -169,4 +175,22 @@ function rf_phi0(rfparams)
   else
     error("RF parameter zero_phase value not set correctly.")
   end
+end
+
+#---------------------------------------------------------------------------------------------------
+
+function fringe_in(f::Fringe.T)
+  if f == Fringe.BothEnds || f == Fringe.EntranceEnd
+    return Val{true}()
+  else
+    return Val{false}()
+  end
+end
+
+function fringe_out(f::Fringe.T)
+  if f == Fringe.BothEnds || f == Fringe.ExitEnd
+    return Val{true}()
+  else
+    return Val{false}()
+  end
 end

ext/BeamTrackingBeamlinesExt/yoshida.jl

@@ -9,14 +9,16 @@
     num_steps = Int(ceil(L / ds_step))
     ds_step = L / num_steps
   end
+  fin  = fringe_in(tm.fringe_at)
+  fout = fringe_out(tm.fringe_at)
   if order == 2
-    return KernelCall(BeamTracking.order_two_integrator!, (ker, params, photon_params, ds_step, num_steps, edge_params, L))
+    return KernelCall(BeamTracking.order_two_integrator!, (ker, params, photon_params, ds_step, num_steps, edge_params, fin, fout, L))
   elseif order == 4
-    return KernelCall(BeamTracking.order_four_integrator!, (ker, params, photon_params, ds_step, num_steps, edge_params, L))
+    return KernelCall(BeamTracking.order_four_integrator!, (ker, params, photon_params, ds_step, num_steps, edge_params, fin, fout, L))
   elseif order == 6
-    return KernelCall(BeamTracking.order_six_integrator!, (ker, params, photon_params, ds_step, num_steps, edge_params, L))
+    return KernelCall(BeamTracking.order_six_integrator!, (ker, params, photon_params, ds_step, num_steps, edge_params, fin, fout, L))
   elseif order == 8
-    return KernelCall(BeamTracking.order_eight_integrator!, (ker, params, photon_params, ds_step, num_steps, edge_params, L))
+    return KernelCall(BeamTracking.order_eight_integrator!, (ker, params, photon_params, ds_step, num_steps, edge_params, fin, fout, L))
   end
 end
 
@@ -81,7 +83,7 @@ end
     q = chargeof(bunch.species)
     mc2 = massof(bunch.species)
     a = gyromagnetic_anomaly(bunch.species)
-    edge_params = ifelse(tm.fringe_on, (a, tilde_m, Ksol, 0, 0, 0), nothing)
+    edge_params = (a, tilde_m, Ksol, 0, 0, 0)
     E0 = mc2/tilde_m/beta_0
     params = (q, mc2, tm.radiation_damping_on, beta_0, gamsqr_0, tilde_m, a, Ksol, mm, kn, ks)
     photon_params = ifelse(tm.radiation_fluctuations_on, (q, mc2, E0, 0, 0, mm, kn, ks), nothing)
@@ -97,39 +99,48 @@ end
   Ksol = kn[1]
   q = chargeof(bunch.species)
   mc2 = massof(bunch.species)
+  a = gyromagnetic_anomaly(bunch.species)
+  Kn0 = ifelse(mm[2] == 1, kn[2], 0)
+  edge_params = (a, tilde_m, Ksol, Kn0, 0, 0)
   E0 = mc2/tilde_m/beta_0
-  params = (q, mc2, tm.radiation_damping_on, beta_0, gamsqr_0, tilde_m, gyromagnetic_anomaly(bunch.species), Ksol, mm, kn, ks)
+  params = (q, mc2, tm.radiation_damping_on, beta_0, gamsqr_0, tilde_m, a, Ksol, mm, kn, ks)
   photon_params = ifelse(tm.radiation_fluctuations_on, (q, mc2, E0, 0, 0, mm, kn, ks), nothing)
-  return integration_launcher(BeamTracking.sks_multipole!, params, photon_params, tm, nothing, L)
+  return integration_launcher(BeamTracking.sks_multipole!, params, photon_params, tm, edge_params, L)
 end
 
-@inline function thick_pure_bdipole(tm::DriftKick, bunch, bm, L)
+@inline function thick_pure_bdipole(tm::Union{Yoshida,DriftKick}, bunch, bm, L)
   p_over_q_ref = bunch.p_over_q_ref
   tilde_m, gamsqr_0, beta_0 = BeamTracking.drift_params(bunch.species, p_over_q_ref)
   mm = bm.order
   kn, ks = get_strengths(bm, L, p_over_q_ref)
   q = chargeof(bunch.species)
   mc2 = massof(bunch.species)
+  a = gyromagnetic_anomaly(bunch.species)
+  Kn0 = ifelse(mm == 1, kn, 0)
+  edge_params = (a, tilde_m, 0, Kn0, 0, 0)
   E0 = mc2/tilde_m/beta_0
-  params = (q, mc2, tm.radiation_damping_on, beta_0, gamsqr_0, tilde_m, gyromagnetic_anomaly(bunch.species), SA[mm], SA[kn], SA[ks])
+  params = (q, mc2, tm.radiation_damping_on, beta_0, gamsqr_0, tilde_m, a, SA[mm], SA[kn], SA[ks])
   photon_params = ifelse(tm.radiation_fluctuations_on, (q, mc2, E0, 0, 0, SA[mm], SA[kn], SA[ks]), nothing)
-  return integration_launcher(BeamTracking.dkd_multipole!, params, photon_params, tm, nothing, L)
+  return integration_launcher(BeamTracking.dkd_multipole!, params, photon_params, tm, edge_params, L)
 end
 
-@inline function thick_bdipole(tm::DriftKick, bunch, bm, L)
+@inline function thick_bdipole(tm::Union{Yoshida,DriftKick}, bunch, bm, L)
   p_over_q_ref = bunch.p_over_q_ref
   tilde_m, gamsqr_0, beta_0 = BeamTracking.drift_params(bunch.species, p_over_q_ref)
   mm = bm.order
   kn, ks = get_strengths(bm, L, p_over_q_ref)
   q = chargeof(bunch.species)
   mc2 = massof(bunch.species)
+  a = gyromagnetic_anomaly(bunch.species)
+  Kn0 = ifelse(mm[1] == 1, kn[1], 0)
+  edge_params = (a, tilde_m, 0, Kn0, 0, 0)
   E0 = mc2/tilde_m/beta_0
-  params = (q, mc2, tm.radiation_damping_on, beta_0, gamsqr_0, tilde_m, gyromagnetic_anomaly(bunch.species), mm, kn, ks)
+  params = (q, mc2, tm.radiation_damping_on, beta_0, gamsqr_0, tilde_m, a, mm, kn, ks)
   photon_params = ifelse(tm.radiation_fluctuations_on, (q, mc2, E0, 0, 0, mm, kn, ks), nothing)
-  return integration_launcher(BeamTracking.dkd_multipole!, params, photon_params, tm, nothing, L)
+  return integration_launcher(BeamTracking.dkd_multipole!, params, photon_params, tm, edge_params, L)
 end
 
-@inline function thick_pure_bdipole(tm::Union{Yoshida,BendKick}, bunch, bm1, L) 
+@inline function thick_pure_bdipole(tm::BendKick, bunch, bm1, L) 
   if isnothing(bunch.coords.q) && !(tm.radiation_damping_on || tm.radiation_fluctuations_on)
     return thick_pure_bdipole(Exact(), bunch, bm1, L)
   else
@@ -144,9 +155,11 @@ end
     q = chargeof(bunch.species)
     mc2 = massof(bunch.species)
     E0 = mc2/tilde_m/beta_0
-    params = (q, mc2, tm.radiation_damping_on, tilde_m, beta_0, gyromagnetic_anomaly(bunch.species), 0, w, w_inv, k0, SA[mm], SA[kn], SA[ks])
+    a = gyromagnetic_anomaly(bunch.species)
+    edge_params = (a, tilde_m, 0, k0, 0, 0)
+    params = (q, mc2, tm.radiation_damping_on, tilde_m, beta_0, a, 0, w, w_inv, k0, SA[mm], SA[kn], SA[ks])
     photon_params = ifelse(tm.radiation_fluctuations_on, (q, mc2, E0, 0, 0, SA[mm], SA[kn], SA[ks]), nothing)
-    return integration_launcher(BeamTracking.bkb_multipole!, params, photon_params, tm, nothing, L)
+    return integration_launcher(BeamTracking.bkb_multipole!, params, photon_params, tm, edge_params, L)
   end
 end
 
@@ -161,10 +174,12 @@ end
   w_inv = inv_rot_quaternion(0,0,tilt)
   q = chargeof(bunch.species)
   mc2 = massof(bunch.species)
+  a = gyromagnetic_anomaly(bunch.species)
+  edge_params = (a, tilde_m, 0, k0, 0, 0)
   E0 = mc2/tilde_m/beta_0
-  params = (q, mc2, tm.radiation_damping_on, tilde_m, beta_0, gyromagnetic_anomaly(bunch.species), 0, w, w_inv, k0, mm, kn, ks)
+  params = (q, mc2, tm.radiation_damping_on, tilde_m, beta_0, a, 0, w, w_inv, k0, mm, kn, ks)
   photon_params = ifelse(tm.radiation_fluctuations_on, (q, mc2, E0, 0, 0, mm, kn, ks), nothing)
-  return integration_launcher(BeamTracking.bkb_multipole!, params, photon_params, tm, nothing, L)
+  return integration_launcher(BeamTracking.bkb_multipole!, params, photon_params, tm, edge_params, L)
 end
 
 @inline function thick_bdipole(tm::MatrixKick, bunch, bm, L)
@@ -185,18 +200,12 @@ end
   w_inv = inv_rot_quaternion(0,0,tilt)
   q = chargeof(bunch.species)
   mc2 = massof(bunch.species)
+  a = gyromagnetic_anomaly(bunch.species)
+  edge_params = (a, tilde_m, 0, kn[1], 0, 0)
   E0 = mc2/tilde_m/beta_0
-  params = (q, mc2, tm.radiation_damping_on, beta_0, gamsqr_0, tilde_m, gyromagnetic_anomaly(bunch.species), w, w_inv, k1, mm, kn, ks)
+  params = (q, mc2, tm.radiation_damping_on, beta_0, gamsqr_0, tilde_m, a, w, w_inv, k1, mm, kn, ks)
   photon_params = ifelse(tm.radiation_fluctuations_on, (q, mc2, E0, 0, 0, mm, kn, ks), nothing)
-  return integration_launcher(BeamTracking.mkm_quadrupole!, params, photon_params, tm, nothing, L)
-end
-
-@inline function thick_bdipole(tm::Yoshida, bunch, bm, L)
-  if bm.order[2] == 2
-    return thick_bdipole(MatrixKick(order=tm.order, num_steps=tm.num_steps, ds_step=tm.ds_step, radiation_damping_on=tm.radiation_damping_on, radiation_fluctuations_on=tm.radiation_fluctuations_on), bunch, bm, L)
-  else
-    return thick_bdipole(BendKick(order=tm.order, num_steps=tm.num_steps, ds_step=tm.ds_step, radiation_damping_on=tm.radiation_damping_on, radiation_fluctuations_on=tm.radiation_fluctuations_on), bunch, bm, L)
-  end
+  return integration_launcher(BeamTracking.mkm_quadrupole!, params, photon_params, tm, edge_params, L)
 end
 
 @inline function thick_pure_bquadrupole(tm::Union{Yoshida,MatrixKick}, bunch, bm, L)
@@ -272,7 +281,7 @@ end
     w = rot_quaternion(0,0,tilt)
     w_inv = inv_rot_quaternion(0,0,tilt)
     a = gyromagnetic_anomaly(bunch.species)
-    edge_params = ifelse(tm.fringe_on, (a, tilde_m, 0, Kn0, e1, e2), nothing)
+    edge_params = (a, tilde_m, 0, Kn0, e1, e2)
     q = chargeof(bunch.species)
     mc2 = massof(bunch.species)
     E0 = mc2/tilde_m/beta_0
@@ -313,8 +322,24 @@ end
   p0c = BeamTracking.R_to_pc(bunch.species, p_over_q_ref)
   q = chargeof(bunch.species)
   mc2 = massof(bunch.species)
+  if mm[1] == 0
+    Ksol = kn[1]
+    if length(mm) > 1 && mm[2] == 1
+      Kn0 = kn[2]
+    else
+      Kn0 = 0
+    end
+  elseif mm[1] == 1
+    Ksol = 0
+    Kn0 = kn[1]
+  else
+    Ksol = 0
+    Kn0 = 0
+  end
+  a = gyromagnetic_anomaly(bunch.species)
+  edge_params = (a, tilde_m, Ksol, Kn0, 0, 0)
   E0 = mc2/tilde_m/beta_0
-  params = (q, mc2, tm.radiation_damping_on, beta_0, gamsqr_0, tilde_m, E_ref, p0c, gyromagnetic_anomaly(bunch.species), omega, t0, E0_over_Rref, mm, kn, ks)
+  params = (q, mc2, tm.radiation_damping_on, beta_0, gamsqr_0, tilde_m, E_ref, p0c, a, omega, t0, E0_over_Rref, mm, kn, ks)
   photon_params = ifelse(tm.radiation_fluctuations_on, (q, mc2, E0, 0, 0, mm, kn, ks), nothing)
-  return integration_launcher(BeamTracking.cavity!, params, photon_params, tm, nothing, L)
+  return integration_launcher(BeamTracking.cavity!, params, photon_params, tm, edge_params, L)
 end