Merge branch 'main' into compathelper/new_version/2024-09-11-00-59-26…

…-586-03444303641

Project.toml

@@ -13,8 +13,7 @@ ReferenceFrameRotations = "74f56ac7-18b3-5285-802d-d4bd4f104033"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
-Distributions = "0.25"
-GTPSA = "1"
+GTPSA = "1.2.0"
 ReferenceFrameRotations = "3"
 StaticArrays = "1"
 julia = "1.9"

src/BeamTracking.jl

@@ -5,11 +5,6 @@ using AcceleratorLattice,
       StaticArrays,
       Distributions
 
-# Temporary until AtomicAndPhysicalConstants is cleaned up ----
-import AtomicAndPhysicalConstants: AtomicAndPhysicalConstants
-const Species = AtomicAndPhysicalConstants.Particle
-# -------------------------------------------------------------
-
 export Beam,
        Coords,
        Symplectic,
@@ -59,11 +54,19 @@ function Beam(
   return Beam(species, coords)
 end
 
+# Creates a Beam as identity GTPSA
+function Beam(d::Descriptor; species::Species=Species("electron"))
+  GTPSA.numvars(d) == 6 || error("Invalid GTPSA Descriptor! Number of variables must be equal to 6.")
+  z = vars(d)
+  return Beam(species, Coords([z[1]], [z[2]], [z[3]], [z[4]], [z[5]], [z[6]]))
+end
+
+
 # --------------------------------------
 
 
 # Modules separated:
 include("symplectic/Symplectic.jl") 
 include("paraxial/Paraxial.jl")    
 
-end
+end

src/paraxial/Paraxial.jl

@@ -1,10 +1,15 @@
 module Paraxial
 using ..BeamTracking: Coords, Beam
 using ..GTPSA
-using ..AcceleratorLattice
+#using ..AcceleratorLattice
+
+struct ParaxialDrift{T}
+  L::T
+end
+
 
 """
-    track!(L, beamf::Beam, beam0::Beam) -> beamf
+    track!(ele::ParaxialDrift, beamf::Beam, beam0::Beam) -> beamf
 
 Routine to tracking through a drift using the paraxial approximation and 
 including higher-order energy effects. 
@@ -14,21 +19,22 @@ including higher-order energy effects.
 - `beamf` -- Output beam after tracking through
 - `beam0` -- Input beam before tracking through
 """
-function track!(ele::Drift, beamf::Beam, beam0::Beam)
+function track!(L, beamf::Beam, beam0::Beam)
   @assert !(beamf === beam0) "Aliasing beamf === beam0 not allowed!"
-  L = ele.L
   z0 = beam0.z
   zf = beamf.z
-
-  @FastGTPSA begin
+  
+  @FastGTPSA! begin
   @. zf[1] = z0[1]+z0[2]*L/(1.0+z0[6])
   @. zf[2] = z0[2]
   @. zf[3] = z0[3]+z0[4]*L/(1.0+z0[6])
   @. zf[4] = z0[4]
   @. zf[5] = z0[5]-L*((z0[2]^2)+(z0[4]^2))/(1.0+z0[6])^2/2.0
   @. zf[6] = z0[6] 
   end
+
   return beamf
 end
 
+
 end

src/symplectic/Symplectic.jl

@@ -3,6 +3,8 @@ using ..BeamTracking: Coords, Beam
 using ..GTPSA
 using ..AcceleratorLattice
 
+export track!
+
 """
     track!(ele::Drift, statef::Beam, state0::Beam) -> statef
 
@@ -14,21 +16,43 @@ This function performs exact, hence symplectic, tracking of a `Beam` through a d
   - `state0` -- Input `Beam`
 """ track!
 
-function track!(ele::Drift, statef::Beam, state0::Beam)
-  @assert !(statef === state0) "Aliasing statef === state0 not allowed!"
+function track!(ele::Drift, statef::Beam, statei::Beam)
+  @assert !(statef === statei) "Aliasing statef === statei not allowed!"
+  L = ele.L
+  zi = statei.z
+  zf = statef.z
+
+  begin
+    @. ps = sqrt((1.0 + zi[6])^2 - zi[2]^2 - zi[4]^2)
+    @. et = sqrt((1.0 + zi[6])^2 + tilde_m^2)
+    @. zf[1] = zi[1] + zi[2] * L / ps
+    @. zf[2] = zi[2]
+    @. zf[3] = zi[3] + zi[4] * L / ps
+    @. zf[4] = zi[4]
+    @. zf[5] = zi[5] - (1.0 + zi[6]) * (L / ps - L / (β0 * et))
+    @. zf[6] = zi[6]
+  end
+  return statef
+end # function track!(Drift)
+
+"""
+track "linear" quadrupole
+"""
+function trackQL!(ele::Quadrupole, statef::Beam, statei::Beam)
+  @assert !(statef === statei) "Aliasing statef === statei not allowed!"
   L = ele.L
-  z0 = state0.z
+  zi = statei.z
   zf = statef.z
 
   begin
-    @. ps = sqrt((1.0 + z0[6])^2 - z0[2]^2 - z0[4]^2)
-    @. et = sqrt((1.0 + z0[6])^2 + tilde_m^2)
-    @. zf[1] = z0[1] + z0[2] * L / ps
-    @. zf[2] = z0[2]
-    @. zf[3] = z0[3] + z0[4] * L / ps
-    @. zf[4] = z0[4]
-    @. zf[5] = z0[5] - (1.0 + z0[6]) * (L / ps - L / (β0 * et))
-    @. zf[6] = z0[6]
+    @. ps = sqrt((1.0 + zi[6])^2 - zi[2]^2 - zi[4]^2)
+    @. et = sqrt((1.0 + zi[6])^2 + tilde_m^2)
+    @. zf[1] = zi[1] + zi[2] * L / ps
+    @. zf[2] = zi[2]
+    @. zf[3] = zi[3] + zi[4] * L / ps
+    @. zf[4] = zi[4]
+    @. zf[5] = zi[5] - (1.0 + zi[6]) * (L / ps - L / (β0 * et))
+    @. zf[6] = zi[6]
   end
   return statef
 end # function track(Drift)

test/element_tests.jl

@@ -0,0 +1,24 @@
+using BeamTracking
+using AcceleratorLattice
+using Test
+
+# define elements
+@ele start = BeginningEle(pc_ref = 10.e6, species_ref = Species("electron"));
+@ele d1 = Drift(L = 2.);
+
+# define beamlines
+bl_drift = BeamLine([ start, d1 ]);
+
+# expand beamlines
+lat_drift = expand([ bl_drift ])
+
+# test individual elements
+@testset "element_tests" begin
+  # drift
+  @test lat_drift["d1"][1].L == 2.0
+  #@test track!(d1, zf, zi) == [xz, pxf, yf, pyf, zf, pzf]
+
+  # quadrupole
+
+end
+

test/runtests.jl

@@ -1,7 +1,13 @@
 using BeamTracking
-using Test, Distributions
+using Distributions
+using Test
 
 @testset "BeamTracking.jl" begin
     beam0 = Beam(10000, d_x=Normal(0, 1e-3), d_px=Normal(0, 1e-4), d_pz=Normal(0, 1e-2))
     beamf = Beam(10000)
+
+  @testset "element_tests" begin
+    include("element_tests.jl")
+  end
+
 end