Added work helper function, linear quad tracking fixes/tests

Project.toml

@@ -21,10 +21,11 @@ Unitful = "1.21.0"
 julia = "1.9"
 
 [extras]
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "Distributions", "JET", "GTPSA"]
+test = ["Test", "Distributions", "JET", "GTPSA", "BenchmarkTools"]

src/BeamTracking.jl

@@ -11,7 +11,7 @@ include("aapc.jl")
 
 export Beam,       
        Coord, 
-       Quaternion
+       Quaternion,
        Particle,
 
        MatrixKick,
@@ -30,7 +30,7 @@ export Beam,
        Species,
 
        sincu,
-       sinhc,
+       sinhcu,
 
        track!
 
@@ -46,11 +46,12 @@ track!(beam::Beam, ::Nothing) = beam
 # --------------------------------------------------
 
 include("utils.jl")
+include("work.jl")
 
 # Modules separated:
 include("MatrixKick/MatrixKick.jl") 
 include("Linear/Linear.jl")   
-#include("Misc/Misc.jl")
+include("Misc/Misc.jl")
 
 
 end

src/Linear/Linear.jl

@@ -2,16 +2,16 @@ module Linear
 using ..GTPSA: @FastGTPSA!, GTPSA
 import ..BeamTracking: track!
 using ..BeamTracking
-
+using ..BeamTracking: get_work
 export track!
 
 Base.@kwdef struct Drift{T}
   L::T
 end
 
 Base.@kwdef struct Quadrupole{T}
-  L::T
-  B1::T
+  B1::T 
+  L::T  
 end
 
 Base.@kwdef struct SBend{T}
@@ -53,12 +53,9 @@ function track!(beam::Beam, ele::Linear.Drift)
   gamma_ref = sr_gamma(beam.beta_gamma_ref)
 
   @FastGTPSA! begin
-  @. v.x  = v.x + v.px * L
-  @. v.px = v.px
-  @. v.y  = v.y + v.py * L
-  @. v.py = v.py
-  @. v.z  = v.z + v.pz*L/gamma_ref^2
-  @. v.pz = v.pz 
+  @. v.x  = v.x + L*v.px
+  @. v.y  = v.y + L*v.py
+  @. v.z  = v.z + L/gamma_ref^2*v.pz
   end
 
   # Spin unchanged
@@ -67,53 +64,51 @@ function track!(beam::Beam, ele::Linear.Drift)
 end
 
 
-"""
-Routine to linearly tracking through a quadrupole
-"""
-function track!(beamf::Beam, ele::Linear.Quadrupole, beami::Beam)
-  @assert !(beamf === beami) "Aliasing beamf === beami not allowed!"
-  zi = beami.vec
-  zf = beamf.vec
+function track!(beam::Beam, ele::Linear.Quadrupole; work=get_work(beam, Val{1}()))
+  v = beam.v
   L = ele.L
-  q = chargeof(beami.species)
 
-  k1 = ele.B1 / brho(massof(beami.species),beami.beta_gamma_ref,q) #quadrupole strengh
-
-  k = sqrt(abs(k1))
-
-  kl = k * L
-  greater = k1 >= 0.0 #horizontal focusing
-  smaller = k1 < 0.0 #horizontal defocusing 
-
+  K1n = ele.B1 / brho(massof(beam.species), beam.beta_gamma_ref, chargeof(beam.species))
+  gamma_ref = sr_gamma(beam.beta_gamma_ref)
 
-  if greater
-    #horizontal focusing
-    cx = cos(kl)
-    sx = sin(kl)
-    sxc = sincu(kl)
-    cy = cosh(kl)
-    sy = sinh(kl)
-    syc = sinhc(kl)
+  if K1n >= 0
+    k = sqrt(K1n)
+    kL = k*L
+    sx  = sin(kL)
+    cx  = cos(kL)
+    sxc = sincu(kL)
+    sy  = sinh(kL)
+    cy  = cosh(kL)
+    syc = sinhcu(kL)
+    sgn = 1
   else
-     #horizontal defocusing
-     cx = cosh(kl)
-     sx = sinh(kl)
-     sxc = sinhc(kl)
-     cy = cos(kl)
-     sy = sin(kl)
-     syc = sincu(kl)
+    k = sqrt(-K1n)
+    kL = k*L
+    sx  = sinh(kL)
+    cx  = cosh(kL)
+    sxc = sinhcu(kL)
+    sy  = sin(kL)
+    cy  = cos(kL)
+    syc = sincu(kL)
+    sgn = -1
   end
 
+
+  # Note: 0+work[1] is workaround for GTPSA @FastGTPSA! bug
   @FastGTPSA! begin
-  @. zf.x  = zi.x * cx + zi.px * sxc * L
-  @. zf.px = (-1 * (greater) + 1 * (smaller)) * zi.x * k * sx + zi.px * cx
-  @. zf.y = zi.y * cy + zi.py * syc * L 
-  @. zf.py = (1 * (greater) - 1 * (smaller)) * zi.y * k * sy + zi.py * cy
-  @. zf.z = zi.z + zi.pz * L 
-  @. zf.pz = zi.pz
+  @. work[1]  = cx*v.x + sxc*L*v.px
+  @. v.px     = -sgn*k*sx*v.x + cx*v.px
+  @. v.x      = 0+work[1] 
+
+  @. work[1]  = cy*v.y + syc*L*v.py
+  @. v.py     = sgn*k*sy*v.y + cy*v.py
+  @. v.y      = 0+work[1] 
+
+  @. v.z      = v.z + L/gamma_ref^2*v.pz
   end 
+
 
-  return beamf
+  return beam
 end 
 
 """

src/MatrixKick/MatrixKick.jl

@@ -2,7 +2,7 @@ module MatrixKick
 using ..GTPSA: @FastGTPSA!, GTPSA
 import ..BeamTracking: track!
 using ..BeamTracking
-
+using ..BeamTracking: get_work
 export track!
 
 struct Drift{T}

src/Misc/Misc.jl

@@ -2,6 +2,7 @@ module Misc
 using ..GTPSA: @FastGTPSA!, GTPSA, evaluate!, TPS
 import ..BeamTracking: track!
 using ..BeamTracking
+using ..BeamTracking: get_work
 export track!
 
 Base.@kwdef struct Taylor{T<:TPS,U<:Union{Nothing,Quaternion{<:T}}}

src/aapc.jl

@@ -7,7 +7,7 @@ Temp code for AAPC
 Base.@kwdef struct Species
   name::String
 end
-massof(::Species) = 510998.95069*u"eV/c^2"
-chargeof(::Species) = -1*u"q"
+massof(::Species) = 510998.95069
+chargeof(::Species) = -1
 
-c_light = 2.99792458e8*u"m/s"
+const c_light = 2.99792458e8

src/utils.jl

@@ -82,7 +82,7 @@ For a particle with a given rest energy and relativistic parameter
 DTA: Need to handle energy units other than ``\\mathrm{eV}``..
 """
 function brho(e_rest, beta_gamma, ne = 1)
-  return sr_pc(e_rest, beta_gamma) / (ne * c_light)
+  return (sr_pc(e_rest, beta_gamma) / (ne * c_light))
 end
 ## If given ``E_\text{kin}`` instead of ``\beta\gamma``,
 ## use the following:
@@ -148,17 +148,18 @@ function sincu(z::Number)
   end
 end
 
+
 """
-    sinhc(z)
+    sinhcu(z)
 
 ## Description
-Compute the hyperbolic sinc function, ``\\operatorname{sinhc}(z) = \\operatorname{sinh}(z) / z``,
+Compute the hyperbolic sinc function, ``\\operatorname{sinhcu}(z) = \\operatorname{sinh}(z) / z``,
 with a correct treatment of the removable singularity at the origin.
 
 ### Implementation
 See sincu for notes about implementation.
 """
-function sinhc(z::Number)
+function sinhcu(z::Number)
   threshold = sqrt(2eps())
   if abs(z) < threshold
     return 1.
@@ -167,3 +168,10 @@ function sinhc(z::Number)
   end
 end
 
+# These will go in GTPSA 
+sincu(z::TPS) = sinc(z/pi)
+sinhcu(z::TPS) = sinhc(z/pi)
+sincu(z::GTPSA.TempTPS) = (GTPSA.div!(z,z,pi); return GTPSA.__t_sinc(z);)
+sinhcu(z::GTPSA.TempTPS) = (GTPSA.div!(z,z,pi); return GTPSA.__t_sinhc(z);)
+
+

src/work.jl

@@ -9,13 +9,13 @@ length `length(Beam.v.x)` which may be used as temporaries.
 - `::Val{N}` -- Number of `N` temporary arrays desired
 """
 function get_work(beam::Beam, ::Val{N}) where {N}
-  sample = one(beam.v.x)
+  sample = first(beam.v.x)
   T = typeof(sample)
-  N_particle = length(Beam.v.x)
+  N_particle = length(beam.v.x)
 
   # Instead of using zeros, we do this to ensure 
   # same GTPSA descriptor if T isa TPS.
-  return ntuple(Val{N}()) do
+  return ntuple(Val{N}()) do t
     r = Vector{T}(undef, N_particle)
     for idx in eachindex(r)
       r[idx] = zero(sample)

test/runtests.jl

@@ -2,19 +2,27 @@ using Test,
       BeamTracking,
       Distributions,
       JET,
+      BenchmarkTools,
       GTPSA
 
 @testset "BeamTracking.jl" begin
-    # Linear drift test:
-    tol = 1e-14
+    # Test setup
+    BenchmarkTools.DEFAULT_PARAMETERS.gctrial = false
+    BenchmarkTools.DEFAULT_PARAMETERS.evals = 2
+
+    # Linear tests -----------------------
     d = Descriptor(6, 1) # 6 variables, order 1
-    beam = Beam(d, beta_gamma_ref=1.0) # Creates Beam as identity map using GTPSA Descriptor
-    gamma_ref = sr_gamma(beam.beta_gamma_ref)
+    beta_gamma_ref = 1.0
+    gamma_ref = sr_gamma(beta_gamma_ref)
+    species = Species("electron")
+    brho_ref = brho(massof(species), beta_gamma_ref, chargeof(species)) 
+    tol = 1e-14
+
+    # Linear drift test:
+    beam = Beam(d, beta_gamma_ref=beta_gamma_ref) # Creates Beam as identity map using GTPSA Descriptor
     L_d = 5.0
     drift = Linear.Drift(L=L_d)
     track!(beam, drift)
-
-    p = Particle(beam, 1) # Converts SoA to 1 struct
 
     M_drift_expected = [1.0  L_d  0.0  0.0  0.0  0.0;
                         0.0  1.0  0.0  0.0  0.0  0.0;
@@ -23,12 +31,63 @@ using Test,
                         0.0  0.0  0.0  0.0  1.0  L_d/gamma_ref^2;
                         0.0  0.0  0.0  0.0  0.0  1.0]
 
+    p = Particle(beam, 1) # Converts SoA to 1 struct
+
     # 1) Correctness:
     @test norm(GTPSA.jacobian(p.v) - M_drift_expected) < tol
     # 2) Type-stability:
     @test_opt track!(beam, drift) 
     # 3) Allocations:
-    @test @allocations(track!(beam, drift)) == 0
+    @test @ballocated(track!($beam, $drift)) == 0
+
+    # Linear Quadrupole test:
+    # Focusing:
+    beam = Beam(d, beta_gamma_ref=beta_gamma_ref)
+    L_q = 1.2
+    K1n = 0.36
+    B1 = K1n*brho_ref
+    qf = Linear.Quadrupole(B1=B1,L=L_q)
+    work = BeamTracking.get_work(beam, Val{1}())
+    track!(beam, qf, work=work)
+
+    p = Particle(beam, 1) # Converts SoA to 1 struct
+
+    M_qf_x = [cos(sqrt(K1n)*L_q)            sincu(sqrt(K1n)*L_q)*L_q;  
+              -sqrt(K1n)*sin(sqrt(K1n)*L_q) cos(sqrt(K1n)*L_q)      ;]
+    M_qf_y = [cosh(sqrt(K1n)*L_q)           sinhcu(sqrt(K1n)*L_q)*L_q; 
+              sqrt(K1n)*sinh(sqrt(K1n)*L_q) cosh(sqrt(K1n)*L_q)      ;]
+    M_qf_z = [1.0 L_q/gamma_ref^2;
+              0.0 1.0            ;]
+
+    M_qf_expected = zeros(6,6)
+    M_qf_expected[1:2,1:2] = M_qf_x
+    M_qf_expected[3:4,3:4] = M_qf_y
+    M_qf_expected[5:6,5:6] = M_qf_z
+    # 1) Correctness:
+    @test norm(GTPSA.jacobian(p.v) - M_qf_expected) < tol
+    # 2) Type-stability:
+    @test_opt track!(beam, qf, work=work) 
+    # 3) Allocations:
+    @test @ballocated(track!($beam, $qf, work=$work)) == 0
+
+    # Defocusing:
+    beam = Beam(d, beta_gamma_ref=beta_gamma_ref)
+    qd = Linear.Quadrupole(B1=-B1,L=L_q)
+    track!(beam, qd, work=work)
+
+    p = Particle(beam, 1)
+
+    M_qd_expected = zeros(6,6)
+    M_qd_expected[1:2,1:2] = M_qf_y
+    M_qd_expected[3:4,3:4] = M_qf_x
+    M_qd_expected[5:6,5:6] = M_qf_z
+
+    # 1) Correctness:
+    @test norm(GTPSA.jacobian(p.v) - M_qd_expected) < tol
+    # 2) Type-stability:
+    @test_opt track!(beam, qd, work=work) 
+    # 3) Allocations:
+    @test @ballocated(track!($beam, $qd, work=$work)) == 0
 end
 
 #include("linear.jl")