Merge pull request #26 from bmad-sim/work

Work

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

@@ -9,39 +9,49 @@ using GTPSA,
 
 include("aapc.jl")
 
-export Beam,
-       Coords,
-       Particle,
+export Beam,       
        Coord, 
+       Quaternion,
+       Particle,
+
        MatrixKick,
        Linear,
+
        sr_gamma, 
        sr_gamma_m1,
        sr_beta,
        sr_pc,
        sr_ekin,
        sr_etot,
        brho,
+
        chargeof,
        massof,
        Species,
+
        sincu,
-       sinhc,
+       sinhcu,
+
        track!
 
 
 
+
+
 include("types.jl")
 
-# Empty tracking method ----------------
+# Empty tracking method to be imported by submodules 
 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")
 
 
 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}
@@ -40,12 +40,11 @@ end
 """
     track!(beam::Beam, ele::Linear.Drift) -> beam
 
-Routine to tracking through a drift using the  approximation and 
-including higher-order energy effects. 
+Track through a linear drift.
 
 ### Arguments
 - `beam` -- Input/output beam before/after tracking through
-- `ele`  -- `Drift` type element
+- `ele`  -- `Linear.Drift` type element
 """
 function track!(beam::Beam, ele::Linear.Drift)
   L = ele.L
@@ -54,65 +53,62 @@ 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
 
   return beam
 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

@@ -0,0 +1,18 @@
+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}}}
+  v::Coord{T}
+  q::U
+end
+
+function track!(beam::Beam, ele::Misc.Taylor)
+
+end
+
+
+end

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

@@ -1,13 +1,25 @@
 """
-    get_orbit_work(beam::Beam, n::Integer)
+    get_work(beam::Beam, ::Val{N}) where N -> work
 
-Returns an array of `n` `StructArray{<:Coord}` based on the 
-passed `Beam` struct which may be used as temporaries.
+Returns a tuple of `N` arrays of type `eltype(Beam.v.x)` and 
+length `length(Beam.v.x)` which may be used as temporaries.
 
 ### Arguments
-- `beam` -- Beam to extract types and number of particles from
-- `n`    -- Number of `StructArray{<:Coord}` temporaries desired
+- `beam`     -- Beam to extract types and number of particles from
+- `::Val{N}` -- Number of `N` temporary arrays desired
 """
-function get_work(beam::Beam, n::Integer)
+function get_work(beam::Beam, ::Val{N}) where {N}
+  sample = first(beam.v.x)
+  T = typeof(sample)
+  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 t
+    r = Vector{T}(undef, N_particle)
+    for idx in eachindex(r)
+      r[idx] = zero(sample)
+    end
+    r
+  end
 end