From 8aaa3f1ddb1aaac61383f40e049df0287c5e931b Mon Sep 17 00:00:00 2001
From: dtabell <dtabell@users.noreply.github.com>
Date: Fri, 27 Sep 2024 15:18:38 -0600
Subject: [PATCH] Multiple changes.

z0 -> zi
export track
initial stab at tracking a quad
---
 src/BeamTracking.jl          |  2 +-
 src/symplectic/Symplectic.jl | 46 +++++++++++++++++++++++++++---------
 test/element_tests.jl        |  2 +-
 3 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/src/BeamTracking.jl b/src/BeamTracking.jl
index 32cad09..358b176 100644
--- a/src/BeamTracking.jl
+++ b/src/BeamTracking.jl
@@ -74,4 +74,4 @@ end
 include("symplectic/Symplectic.jl") 
 include("paraxial/Paraxial.jl")    
 
-end
\ No newline at end of file
+end
diff --git a/src/symplectic/Symplectic.jl b/src/symplectic/Symplectic.jl
index 6604f6a..bc73ca4 100644
--- a/src/symplectic/Symplectic.jl
+++ b/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)
diff --git a/test/element_tests.jl b/test/element_tests.jl
index b338c4d..eee6fa7 100644
--- a/test/element_tests.jl
+++ b/test/element_tests.jl
@@ -3,7 +3,7 @@ using AcceleratorLattice
 using Test
 
 # define elements
-@ele start = BeginningEle(pc_ref = 10.e6, species_ref = species("electron"));
+@ele start = BeginningEle(pc_ref = 10.e6, species_ref = Species("electron"));
 @ele d1 = Drift(L = 2.);
 
 # define beamlines