Merge pull request #25 from sophiayang627/main

Finished quad + drift testings

src/BeamTracking.jl

@@ -4,15 +4,28 @@ using GTPSA,
       StaticArrays,
       Distributions
 
+
 include("aapc.jl")
 
 export Beam,
        Coords,
+       Particle,
+       Coord, 
        Symplectic,
-       Linear
+       Linear,
+       sr_gamma, 
+       sr_gamma_m1,
+       sr_beta,
+       sr_pc,
+       sr_ekin,
+       sr_etot,
+       brho,
+       chargeof,
+       massof,
+       Species
 
 # SoA ----------------------------------
-struct Coords{T} <: FieldVector{6, T}
+Base.@kwdef struct Coords{T}
   x::Vector{T}
   px::Vector{T}
   y::Vector{T}
@@ -37,12 +50,16 @@ function Coords(
   return Coords(x, px, y, py, z, pz)
 end
 
+
+
 struct Beam{S,T}
   species::Species
   beta_gamma_0::S
   z::Coords{T}
 end
 
+
+
 function Beam(
   n::Integer; species::Species=Species("electron"), beta_gamma_0=1,
   d_x::Distribution=Normal(0,0), d_px::Distribution=Normal(0,0), 
@@ -62,12 +79,47 @@ function Beam(d::Descriptor; species::Species=Species("electron"), beta_gamma_0=
   return Beam(species, beta_gamma_0, Coords([z[1]], [z[2]], [z[3]], [z[4]], [z[5]], [z[6]]))
 end
 
+#Create a Beam with single particle for testing
+function Beam(
+  ;species::Species=Species("electron"), beta_gamma_0=1,
+  x=0.0, px=0.0, y=0.0, py=0.0, z=0.0, pz=0.0,
+  )
+
+  coords = Coords(x=[x], px=[px], y=[y], py=[py], z=[z], pz=[pz])
+
+  return Beam(species, beta_gamma_0, coords)
+end
+
+#AoS from SoA for single particle
+# Extract the phase space coord of a particle in a beam 
+Base.@kwdef struct Coord{T} <: FieldVector{6, T} # Just like Coords but only 1 Coord 
+  x::T  = 0.0
+  px::T = 0.0
+  y::T  = 0.0
+  py::T = 0.0
+  z::T  = 0.0
+  pz::T = 0.0
+end
+
+struct Particle{S,T}
+  species::Species
+  beta_gamma_0::S
+  z::Coord{T}
+end
+
+function Particle(b::Beam, n::Integer=1)
+  z = b.z
+  coord = Coord(z.x[n],z.px[n],z.y[n],z.py[n],z.z[n],z.pz[n])
+
+  return Particle(b.species, b.beta_gamma_0, coord)
+end
 
 # --------------------------------------
 include("utils.jl")
 
 # Modules separated:
 include("symplectic/Symplectic.jl") 
-include("linear/Linear.jl")    
+include("linear/Linear.jl")   
+
 
 end

src/aapc.jl

@@ -8,4 +8,6 @@ Base.@kwdef struct Species
   name::String
 end
 massof(::Species) = 510998.95069
-chargeof(::Species) = -1
+chargeof(::Species) = -1
+
+c_light = 2.99792458e8

src/linear/Linear.jl

@@ -1,31 +1,41 @@
 module Linear
-using ..BeamTracking: Coords, Beam
-using ..GTPSA: @FastGTPSA!
+using ..GTPSA: @FastGTPSA!, GTPSA
+using ..BeamTracking
 
 export track!
 
-struct Drift{T}
+Base.@kwdef struct Drift{T}
   L::T
 end
 
-struct Quadrupole{T}
+Base.@kwdef struct Quadrupole{T}
   L::T
   B1::T
 end
 
-struct SBend{T}
+Base.@kwdef struct SBend{T}
   L::T    # Arc length
   B0::T   # Field strength
   g::T    # Coordinate system curvature through element
   e1::T   # Edge 1
   e2::T   # Edge 2
 end
 
-struct Solenoid{T}
+Base.@kwdef struct Combined{T}
+  L::T 
+  B0::T 
+  B1::T
+  g::T
+  e1::T
+  e2::T
+end
+
+Base.@kwdef struct Solenoid{T}
   L::T
   Bs::T
 end
 
+
 """
     track!(ele::Linear.Drift, beamf::Beam, beami::Beam) -> beamf
 
@@ -39,20 +49,127 @@ including higher-order energy effects.
 """
 function track!(ele::Linear.Drift, beamf::Beam, beami::Beam)
   @assert !(beamf === beami) "Aliasing beamf === beami not allowed!"
+  L = ele.L
   zi = beami.z
   zf = beamf.z
 
   @FastGTPSA! begin
-  @. zf.x  = zi.x+zi.px*L
+  @. zf.x  = zi.x + zi.px * L
   @. zf.px = zi.px
-  @. zf.y  = zi.y+zi.py*L
+  @. zf.y  = zi.y + zi.py * L
   @. zf.py = zi.py
-  @. zf.z  = zi.z-L*((zi.px^2)+(zi.py^2))/(1.0+zi.pz)^2/2.0
+  @. zf.z  = zi.z + zf.pz * L
   @. zf.pz = zi.pz 
   end
 
   return beamf
 end
 
 
+"""
+Routine to linearly tracking through a quadrupole
+"""
+function track!(ele::Linear.Quadrupole,beamf::Beam,beami::Beam)
+  @assert !(beamf === beami) "Aliasing beamf === beami not allowed!"
+  zi = beami.z
+  zf = beamf.z
+  L = ele.L
+  q = chargeof(beami.species)
+
+  k1 = ele.B1 / brho(massof(beami.species),beami.beta_gamma_0,q) #quadrupole strengh
+
+  k = sqrt(abs(k1))
+
+  kl = k * L
+  greater = k1 >= 0.0 #horizontal focusing
+  smaller = k1 < 0.0 #horizontal defocusing 
+
+
+  if greater
+    #horizontal focusing
+    cx = cos(kl)
+    sx = sin(kl)
+    sxc = sincu(kl)
+    cy = cosh(kl)
+    sy = sinh(kl)
+    syc = sinhc(kl)
+  else
+     #horizontal defocusing
+     cx = cosh(kl)
+     sx = sinh(kl)
+     sxc = sinhc(kl)
+     cy = cos(kl)
+     sy = sin(kl)
+     syc = sincu(kl)
+  end
+
+  @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
+  end 
+
+  return beamf
+end 
+
+"""
+Routine to linearly tracking through a Sector Magnet
+"""
+function track!(ele::Linear.SBend,beamf::Beam,beami::Beam)
+  @assert !(beamf === beami) "Aliasing beamf === beami not allowed!"
+  zi = beami.z
+  zf = beamf.z
+  L = ele.L
+  q = chargeof(beami.species)
+
+  #curvature of B field
+  k = ele.B0 / brho(massof(beami.species),beami.beta_gamma_0,q)
+  kl = k * L
+
+
+  @FastGTPSA! begin
+  @. zf.x  = zi.x * cos(kl) +  zi.px * sin(kl) / k + zi.pz * (1 - cos(kl)) / k
+  @. zf.px = -zi.x * k * sin(kl) + zi.px * cos(kl) + zi.pz * sin(kl)
+  @. zf.y = zi.y + zi.py * L
+  @. zf.py = zi.py
+  @. zf.z = -zi.x * sin(kl) + zi.px * (cos(kl) - 1) / k + zi.z + zi.pz * (sin(kl) - kl) / k
+  @. zf.pz = zi.pz
+ end 
+end 
+
+
+"""
+Routine to linearly tracking through a Combined Magnet
+"""
+function track!(ele::Linear.Combined,beamf::Beam,beami::Beam)
+  @assert !(beamf === beami) "Aliasing beamf === beami not allowed!"
+  zi = beami.z
+  zf = beamf.z
+  L = ele.L
+
+  brho = brho(massof(beami.species),beami.beta_gamma_0,q)
+  ka = ele.B0 / brho #curvature
+  k1 =  ele.B1 / brho #quad strength
+
+  K = ka^2 + k1
+
+  sqk1 = sqrt(abs(k1)) 
+  sqK = sqrt(abs(K))
+
+  Kl = sqK * L
+  kl  = sqk1 * L
+
+  @FastGTPSA! begin
+  @. zf.x  = zi.x * cos(Kl) +  zi.px * sin(Kl) / sqK + zi.pz * (1 - cos(Kl)) * ka / K
+  @. zf.px = - zi.x * sqK * sin(Kl) + zi.px * cos(Kl) + zi.pz * sin(Kl) * ka / sqK
+  @. zf.y =  zi.y * cosh(kl) + zi.py * sinh(kl) / sqk
+  @. zf.py = zi.y * sqK * sinh(kl) + zi.py * cosh(kl)
+  @. zf.z = - zi.x * sin(Kl) * ka / sqK + zi.px * (cos(Kl)-1) * ka / K + zi.z + zi.pz * (sin(Kl) / sqK - l) *  ka^2 / K
+  @. zf.pz = zi.pz
+ end 
+end 
+
 end

src/utils.jl

@@ -16,6 +16,7 @@ function sr_gamma(beta_gamma)
 end
 
 
+
 """
     sr_gamma_m1(beta_gamma)
 
@@ -81,10 +82,8 @@ 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 * clight)
+  return sr_pc(e_rest, beta_gamma) / (ne * c_light)
 end
-
-
 ## If given ``E_\text{kin}`` instead of ``\beta\gamma``,
 ## use the following:
 #
@@ -117,3 +116,56 @@ end
 #  return sr_pc(e_rest, e_kin) / (ne * clight)
 #end
 
+
+"""
+    sincu(z)
+
+## Description
+Compute the unnormalized sinc function, ``\\operatorname{sincu}(z) = \\sin(z) / z``,
+with a correct treatment of the removable singularity at the origin.
+
+### Implementation
+The function ``\\sin(z) / z = 1 - z^2 / 3! + z^4 / 5! - z^6 / 7! + \\cdots``.
+For real values of ``z \\in (-1,1)``, one can truncate this series just before
+the ``k^\\text{th}`` term, ``z^{2k} / (2k+1)!``, and the alternating nature of
+the series guarantees the error does not exceed the value of this first truncated term.
+It follows that if ``z^2 / 6 < \\varepsilon_\\text{machine}``, simply truncating
+the series to 1 yields machine precision accuracy near the origin.
+And outside that neighborhood, one simply computes ``\\sin(z) / z``.
+On the otherhand, if we allow for complex values, one can no longer assume the
+series alternates, and one must use a more conservative threshold.
+Numerical exploration suggests that for ``|z| < 1`` the sum of terms starting
+at the ``k^\\text{th}`` term is bounded by ``|z|^{2k} / (2k)!``.
+It then follows that one should use ``|z|^2 / 2 < \\varepsilon_\\text{machine}``
+as the criterion for deciding to truncate the series to 1 near the origin.
+"""
+function sincu(z::Number)
+  threshold = sqrt(2eps())
+  if abs(z) < threshold
+    return 1.
+  else
+    return sin(z) / z
+  end
+end
+
+"""
+    sinhc(z)
+
+## Description
+Compute the hyperbolic sinc function, ``\\operatorname{sinhc}(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)
+  threshold = sqrt(2eps())
+  if abs(z) < threshold
+    return 1.
+  else
+    return sinh(z) / z
+  end
+end
+
+export sincu,
+       sinhc