track! modular now and different argument order, small updates

Project.toml

@@ -8,12 +8,14 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8"
 ReferenceFrameRotations = "74f56ac7-18b3-5285-802d-d4bd4f104033"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [compat]
 Distributions = "0.25"
 GTPSA = "1.2.0"
 ReferenceFrameRotations = "3"
 StaticArrays = "1"
+Unitful = "1.21.0"
 julia = "1.9"
 
 [extras]

src/BeamTracking.jl

@@ -2,7 +2,8 @@ module BeamTracking
 using GTPSA,
       ReferenceFrameRotations,
       StaticArrays,
-      Distributions
+      Distributions,
+      Unitful
 
 
 include("aapc.jl")
@@ -22,7 +23,10 @@ export Beam,
        brho,
        chargeof,
        massof,
-       Species
+       Species,
+       sincu,
+       sinhc,
+       track!
 
 # SoA ----------------------------------
 Base.@kwdef struct Coords{T}
@@ -58,8 +62,6 @@ struct Beam{S,T}
   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), 
@@ -114,6 +116,9 @@ function Particle(b::Beam, n::Integer=1)
   return Particle(b.species, b.beta_gamma_0, coord)
 end
 
+# Empty tracking method ----------------
+track!(beamf::Beam, ::Nothing, beami::Beam) = (beamf.z .= beami.z; return)
+
 # --------------------------------------
 include("utils.jl")
 

src/linear/Linear.jl

@@ -1,5 +1,6 @@
 module Linear
 using ..GTPSA: @FastGTPSA!, GTPSA
+import ..BeamTracking: track!
 using ..BeamTracking
 
 export track!
@@ -37,7 +38,7 @@ end
 
 
 """
-    track!(ele::Linear.Drift, beamf::Beam, beami::Beam) -> beamf
+    track!(beamf::Beam, ele::Linear.Drift, beami::Beam) -> beamf
 
 Routine to tracking through a drift using the  approximation and 
 including higher-order energy effects. 
@@ -47,7 +48,7 @@ including higher-order energy effects.
 - `beamf` -- Output beam after tracking through
 - `beami` -- Input beam before tracking through
 """
-function track!(ele::Linear.Drift, beamf::Beam, beami::Beam)
+function track!(beamf::Beam, ele::Linear.Drift, beami::Beam)
   @assert !(beamf === beami) "Aliasing beamf === beami not allowed!"
   L = ele.L
   zi = beami.z
@@ -69,7 +70,7 @@ end
 """
 Routine to linearly tracking through a quadrupole
 """
-function track!(ele::Linear.Quadrupole,beamf::Beam,beami::Beam)
+function track!(beamf::Beam, ele::Linear.Quadrupole, beami::Beam)
   @assert !(beamf === beami) "Aliasing beamf === beami not allowed!"
   zi = beami.z
   zf = beamf.z
@@ -104,12 +105,12 @@ function track!(ele::Linear.Quadrupole,beamf::Beam,beami::Beam)
   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
+  @. 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
@@ -118,7 +119,7 @@ end
 """
 Routine to linearly tracking through a Sector Magnet
 """
-function track!(ele::Linear.SBend,beamf::Beam,beami::Beam)
+function track!(beamf::Beam, ele::Linear.SBend, beami::Beam)
   @assert !(beamf === beami) "Aliasing beamf === beami not allowed!"
   zi = beami.z
   zf = beamf.z
@@ -137,14 +138,14 @@ function track!(ele::Linear.SBend,beamf::Beam,beami::Beam)
   @. 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 
 end 
 
 
 """
 Routine to linearly tracking through a Combined Magnet
 """
-function track!(ele::Linear.Combined,beamf::Beam,beami::Beam)
+function track!(beamf::Beam, ele::Linear.Combined, beami::Beam)
   @assert !(beamf === beami) "Aliasing beamf === beami not allowed!"
   zi = beami.z
   zf = beamf.z
@@ -169,7 +170,7 @@ function track!(ele::Linear.Combined,beamf::Beam,beami::Beam)
   @. 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 
 
 end

src/symplectic/Symplectic.jl

@@ -1,6 +1,7 @@
 module Symplectic
-using ..BeamTracking: Coords, Beam
-using ..GTPSA: @FastGTPSA!
+using ..GTPSA: @FastGTPSA!, GTPSA
+import ..BeamTracking: track!
+using ..BeamTracking
 
 export track!
 
@@ -14,16 +15,16 @@ struct Quadrupole{T}
 end
 
 """
-    track!(ele::Symplectic.Drift, beamf::Beam, beami::Beam) -> beamf
+    track!(beamf::Beam, ele::Drift, beami::Beam) -> beamf
 
 This function performs exact, hence symplectic, tracking of a `Beam` through a drift.
 
 ### Arguments
-  - `ele`    -- `Drift`-type element
   - `beamf` -- Output `Beam`
+  - `ele`    -- `Drift`-type element
   - `beami` -- Input `Beam`
 """
-function track!(ele::Drift, beamf::Beam, beami::Beam)
+function track!(beamf::Beam, ele::Drift, beami::Beam)
   @assert !(beamf === beami) "Aliasing beamf === beami not allowed!"
   @assert !(beamf.species === beami.species) "Input species must be equal to output!"
   L = ele.L
@@ -48,7 +49,7 @@ end # function track!(::Drift, ::Beam, ::Beam)
 """
 track quadrupole
 """
-function track!(ele::Quadrupole, beamf::Beam, beami::Beam)
+function track!(beamf::Beam, ele::Quadrupole, beami::Beam)
   @assert !(beamf === beami) "Aliasing beamf === beami not allowed!"
   L = ele.L
   zi = beami.z

src/utils.jl

@@ -167,5 +167,3 @@ function sinhc(z::Number)
   end
 end
 
-export sincu,
-       sinhc