From d2ccac7b56430c9d46698455adfd690d41fd2668 Mon Sep 17 00:00:00 2001
From: dtabell <dtabell@users.noreply.github.com>
Date: Mon, 14 Oct 2024 12:01:17 -0600
Subject: [PATCH] Begin implementing Quadrupole.

---
 src/symplectic/Symplectic.jl | 38 ++++++++++++++++++++++++++++++++++--
 test/element_tests.jl        |  5 +++--
 2 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/src/symplectic/Symplectic.jl b/src/symplectic/Symplectic.jl
index 89741ec..ba84d70 100644
--- a/src/symplectic/Symplectic.jl
+++ b/src/symplectic/Symplectic.jl
@@ -42,7 +42,7 @@ end # function track!(::Drift, ::Beam, ::Beam)
 """
 track "linear" quadrupole
 """
-function trackQL!(ele::Quadrupole, statef::Beam, statei::Beam)
+function track(ele::Quadrupole, statef::Beam, statei::Beam)
   @assert !(statef === statei) "Aliasing statef === statei not allowed!"
   L = ele.L
   zi = statei.z
@@ -62,7 +62,41 @@ function trackQL!(ele::Quadrupole, statef::Beam, statei::Beam)
     @. zf[6] = zi[6]
   end
   return statef
-end # function track(Drift)
+end # function track!(::Quadrupole)
+
+"""
+track "matrix part" of quadrupole
+"""
+function trackQuadMf!(statef::Beam, statei::Beam, s::Float64, κnum::Float64)
+  zi = statei.z
+  zf = statef.z
+
+  tilde_m = massof(ele.species_ref) / ele.pc_ref
+  β0 = ele.pc_ref / ele.E_tot_ref
+
+  begin
+    p_red = 1 + zi[6] # reduced Pz
+    κ = κnum / p_red
+    κs = κ * s
+    cx = cos(κs)
+    sx = sin(κs)
+    cy = cosh(κs)
+    sy = sinh(κs)
+    sx2 = sinc(2κs)
+    sy2 = sinch(2κs)
+    xp = px / p_red
+    yp = py / p_red
+    @. zf[1] = zi[1] * cx + xp * sx / κ
+    @. zf[2] = zi[2] * cx - p_red * κ * zi[1]
+    @. zf[3] = zi[3] * cy + yp * sy / κ
+    @. zf[4] = zi[4] * cy + p_red * κ * zi[3]
+    @. zf[5] = (zi[5] - (s / 4.) * (xp^2 * (1. + sx2) + yp^2 * (1 + sy2)
+                                    + (κ * x)^2 * (1 - sx2) + (κ * y)^2 * (1 - sy2))
+                      + (x * xp * sx^2 - y * yp * sy^2) / 2.)
+    @. zf[6] = zi[6]
+  end
+  return statef
+end # function trackQ!M::Quadrupole()
 
 
 end # module Symplectic
diff --git a/test/element_tests.jl b/test/element_tests.jl
index 7a38d4c..fb69cbe 100644
--- a/test/element_tests.jl
+++ b/test/element_tests.jl
@@ -7,7 +7,7 @@ using Test
 @ele d1 = Drift(L = 2.);
 
 # define beamlines
-bl_drift = BeamLine([ start, d1 ]);
+bl_drift = BeamLine("bl_drift", [ start, d1 ]);
 
 # expand beamlines
 lat_drift = expand([ bl_drift ])
@@ -16,7 +16,8 @@ lat_drift = expand([ bl_drift ])
 @testset "element_tests" begin
   # drift
   @test lat_drift["d1"][1].L == 2.0
-  #@test track!(d1, zf, zi) == [xz, pxf, yf, pyf, zf, pzf]
+  #@test track!(lat_drift["d1"][1].L, zf, zi);
+  #      zf == [ x_f, px_f, y_f, py_f, z_f, pz_f ]
 
   # quadrupole