Merge pull request #37 from bmad-sim/4-mk-drift

4 mk drift

Project.toml

@@ -4,20 +4,18 @@ authors = ["mattsignorelli <[email protected]> and contributors"]
 version = "1.0.0-DEV"
 
 [deps]
-Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+AtomicAndPhysicalConstants = "5c0d271c-5419-4163-b387-496237733d8b"
 GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8"
 ReferenceFrameRotations = "74f56ac7-18b3-5285-802d-d4bd4f104033"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
-Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [compat]
-Distributions = "0.25"
-GTPSA = "1.2.0"
+AtomicAndPhysicalConstants = "0.5"
+GTPSA = "1.3.1"
 ReferenceFrameRotations = "3"
 StaticArrays = "1"
 StructArrays = "0.6.18, 0.7"
-Unitful = "1.21.0"
 julia = "1.9"
 
 [extras]

src/BeamTracking.jl

@@ -3,21 +3,19 @@ using GTPSA,
       ReferenceFrameRotations,
       StaticArrays,
       StructArrays,
-      Distributions,
-      Unitful
+      AtomicAndPhysicalConstants
+@APCdef
 
 
-include("aapc.jl")
-
 export Bunch,       
        Coord, 
        Quaternion,
        Particle,
-       
+
        MatrixKick,
        Linear,
 
-       sr_gamma, 
+       sr_gamma,
        sr_gamma_m1,
        sr_beta,
        sr_pc,
@@ -34,9 +32,7 @@ export Bunch,
 
        track!
 
-
-
-
+import GTPSA: sincu, sinhcu
 
 include("types.jl")
 
@@ -49,8 +45,8 @@ include("utils.jl")
 include("work.jl")
 
 # Modules separated:
-include("MatrixKick/MatrixKick.jl") 
-include("Linear/Linear.jl")   
+include("MatrixKick/MatrixKick.jl")
+include("Linear/Linear.jl")
 include("Misc/Misc.jl")
 
 

src/MatrixKick/MatrixKick.jl

@@ -1,5 +1,6 @@
 module MatrixKick
 using ..GTPSA: @FastGTPSA!, GTPSA
+using StructArrays
 import ..BeamTracking: track!
 using ..BeamTracking
 using ..BeamTracking: get_work
@@ -10,23 +11,40 @@ Base.@kwdef struct Drift{T}
 end
 
 Base.@kwdef struct Quadrupole{T}
-  L::T   # quadrupole length / m
+  L::T    # quadrupole length / m
   Bn1::T  # quadrupole gradient / (T·m^-1)
 end
 
 
 function track!(bunch::Bunch, ele::MatrixKick.Drift; work=get_work(bunch, Val{1}()))
+#=
+This function implements symplectic tracking through a drift,
+derived using the Hamiltonian (25.9) given in the BMad manual.
+As a consequence of using that Hamiltonian, the reference value
+of βγ must be that of a particle with the design energy.
+Should we wish to change that, we shall need to carry both
+reference and design values.
+=#
   L = ele.L
   v = bunch.v
 
-  tilde_m = 1 / bunch.beta_gamma_ref
-  beta_ref = sr_beta(bunch.beta_gamma_ref)
+  tilde_m    = 1 / bunch.beta_gamma_ref
+  gamsqr_ref = 1 + bunch.beta_gamma_ref^2
+  beta_ref   = bunch.beta_gamma_ref / sqrt(gamsqr_ref)
 
   @FastGTPSA! begin
-  @. work[1] = 1 / sqrt((1.0 + v.pz)^2 - (v.px^2 + v.py^2))
-  @. v.x  .= v.x + v.px * L * work[1]
-  @. v.y  .= v.y + v.py * L * work[1]
-  @. v.z  .= v.z - (1.0 + v.pz) * (L * work[1] - L / (beta_ref * sqrt((1.0 + v.pz)^2 + tilde_m^2)))
+  @. work[1] = sqrt((1.0 + v.pz)^2 - (v.px^2 + v.py^2))  # P_s
+  @. v.x  .= v.x + v.px * L / work[1]
+  @. v.y  .= v.y + v.py * L / work[1]
+  #@. v.z  .= v.z - ( (1.0 + v.pz) * L
+  #                   * (1. / work[1] - 1. / (beta_ref * sqrt((1.0 + v.pz)^2 + tilde_m^2))) )
+  # high-precision computation of z-final
+  @. v.z  .= v.z - ( (1.0 + v.pz) * L
+                     * ((v.px^2 + v.py^2) - v.pz * (2 + v.pz) / gamsqr_ref)
+                     / ( beta_ref * sqrt((1.0 + v.pz)^2 + tilde_m^2) * work[1]
+                         * (beta_ref * sqrt((1.0 + v.pz)^2 + tilde_m^2) + work[1])
+                       )
+                   )
   end
 
   # Spin unchanged
@@ -39,58 +57,55 @@ end # function track!(::Bunch, ::Drift)
 # This integrator uses the so-called Matrix-Kick-Matrix method to implement
 # an integrator accurate though second-order in the integration step-size.
 function track!(bunch::Bunch, ele::MatrixKick.Quadrupole; work=get_work(bunch, Val{6}()))
-  @assert !(beamf === beami) "Aliasing beamf === beami not allowed!"
   L = ele.L
 
   # κ^2 (kappa-squared) := (q g / P0) / (1 + δ)
   # numerator of κ^2
-  k2_num = Bn1 / brho(massof(bunch.species), bunch.beta_gamma_ref, chargeof(bunch.species))
+  k2_num = ele.Bn1 / brho(massof(bunch.species), bunch.beta_gamma_ref, chargeof(bunch.species))
 
   v = bunch.v
   v_work = StructArray{Coord{eltype(work[1])}}((work[1], work[2], work[3], work[4], work[5], work[6]))
 
   trackQuadMx!(v_work, v, k2_num, L / 2)
-  trackQuadK!( v, v_work, L)
+  trackQuadK!( v, v_work, bunch.beta_gamma_ref, L)
   trackQuadMx!(v_work, v, k2_num, L / 2)
 
-
   v .= v_work
-
   return bunch
-end # function track!(::Quadrupole)
+end # function track!(::Bunch, ::Quadrupole)
 
 
 """
-    trackQuadMx!(beamf::Bunch, beami::Bunch, k2_num::Float64, s::Float64)
+    trackQuadMx!(vf::StructArray, vi::StructArray, k2_num::Float64, s::Float64)
 
 track "matrix part" of quadrupole
 """
 function trackQuadMx!(vf, vi, k2_num, s)
   focus   = k2_num >= 0  # horizontally focusing
   defocus = k2_num <  0  # horizontally defocusing
 
-  p = @. 1 + vi.pz    # reduced momentum, P/P0 = 1 + δ
+  p =  @. 1 + vi.pz    # reduced momentum, P/P0 = 1 + δ
   k2 = @. k2_num / p  # κ^2 for each particle
   ks = @. sqrt(abs(k2)) * s  # |κ|s
-  xp = @. v.px / p  # x'
-  yp = @. v.py / p  # y'
-
-  cx =  @. focus * cos(ks)    + defocus * cosh(ks)
-  cy =  @. focus * cosh(ks)   + defocus * cos(ks)
-  sx =  @. focus * sincu(ks)  + defocus * sinhc(ks)
-  sy =  @. focus * sinhc(ks)  + defocus * sincu(ks)
-  sx2 = @. focus * sincu(2ks) + defocus * sinhc(2ks)
-  sy2 = @. focus * sinhc(2ks) + defocus * sincu(2ks)
-  sxz = @. focus * sin(ks)^2  + defocus * sinh(ks)^2
-  syz = @. focus * sinh(ks)^2 + defocus * sin(ks)^2
+  xp = @. vi.px / p  # x'
+  yp = @. vi.py / p  # y'
+
+  cx =  @. focus * cos(ks)     + defocus * cosh(ks)
+  cy =  @. focus * cosh(ks)    + defocus * cos(ks)
+  sx =  @. focus * sincu(ks)   + defocus * sinhcu(ks)
+  sy =  @. focus * sinhcu(ks)  + defocus * sincu(ks)
+  sx2 = @. focus * sincu(2ks)  + defocus * sinhcu(2ks)
+  sy2 = @. focus * sinhcu(2ks) + defocus * sincu(2ks)
+  sxz = @. focus * sin(ks)^2   - defocus * sinh(ks)^2
+  syz = @. focus * sinh(ks)^2  - defocus * sin(ks)^2
 
   @. vf.x  = vi.x  * cx + xp * s * sx
-  @. vf.px = vi.px * cx - k2 * p * vi.x * s * sx
+  @. vf.px = vi.px * cx - vi.x * p * k2 * s * sx
   @. vf.y  = vi.y  * cy + yp * s * sy
-  @. vf.py = vi.py * cy + k2 * p * vi.y * s * sy
+  @. vf.py = vi.py * cy + vi.y * p * k2 * s * sy
   @. vf.z  = (vi.z - (s / 4) * (xp^2 * (1 + sx2) + yp^2 * (1 + sy2)
                                 + k2 * vi.x^2 * (1 - sx2) - k2 * vi.y^2 * (1 - sy2))
-                   + (vi.x * xp * sxz - vi.y * yp * syz) / 2.)
+                   + (vi.x * xp * sxz - vi.y * yp * syz) / 2)
   @. vf.pz = vi.pz
 
   return vf
@@ -109,18 +124,24 @@ which suffers a loss of precision when ``|A| \\ll 1``. To combat that
 problem, we rewrite it in the form ``A / (1 + \\sqrt{1-A})``---more
 complicated, yes, but far more accurate.
 """
-function trackQuadK!(vf, vi, s)
-  tilde_m = 1 / beami.beta_gamma_ref  # mc^2 / p0·c
-  beta_ref = sr_beta(beami.beta_gamma_ref)
+function trackQuadK!(vf, vi, betgam_ref, s)
+  tilde_m = 1 / betgam_ref  # mc^2 / p0·c
+  beta_ref = sr_beta(betgam_ref)
+  beta_ref = sr_beta(betgam_ref)
+  gamsqr_ref = 1 + betgam_ref^2
 
   p    = @. 1 + vi.pz  # reduced momentum, P/P0 = 1 + δ
-  ptr2 = @. px^r + py^2
+  ptr2 = @. vi.px^2 + vi.py^2
   ps   = @. sqrt(p^2 - ptr2)
 
   @. vf.x  = vi.x + s * vi.px / p * ptr2 / (ps * (p + ps))
   @. vf.y  = vi.y + s * vi.py / p * ptr2 / (ps * (p + ps))
-  @. vf.z  = vi.z - s * (p / ps - (1/2) * ptr2 / p^2
-                         - p / (beta_ref * sqrt(p^2 + tilde_m^2)))
+  @. vf.z  = vi.z - s * ( (1.0 + vi.pz)
+                         * (ptr2 - vi.pz * (2 + vi.pz) / gamsqr_ref)
+                         / ( beta_ref * sqrt((1.0 + vi.pz)^2 + tilde_m^2) * ps
+                             * (beta_ref * sqrt((1.0 + vi.pz)^2 + tilde_m^2) + ps)
+                           ) - ptr2 / (2 * (1 + vi.pz)^2)
+                        )
   @. vf.px = vi.px
   @. vf.py = vi.py
   @. vf.pz = vi.pz

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:
@@ -166,11 +166,3 @@ function sinhcu(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);)
-
-
-