starting GPU version of track a quadrupole

src/low_level/int_arrays.jl

@@ -0,0 +1,17 @@
+using CUDA
+include("structures.jl")
+
+function fill_quadrupole(elements)
+    """filling intermediate calculation arrays for track 
+a quadrupole with zeros to avoid dynamic memory allocation; 
+change element count as needed."""
+    global x_ele_int, x_ele, y_ele, px_ele, py_ele, s, c, 
+    x_lab, y_lab, px_lab, py_lab, beta, beta0, e_tot, 
+    evaluation, dz, sqrt_k, sk_l, cx, sx, a11, a12, a21,
+    a22, c1, c2, c3 = (CUDA.fill(0.0, elements) for item = 1:27)
+
+    return nothing
+end
+
+fill_quadrupole(10_000)
+

src/low_level/offset_particle.jl

@@ -1,15 +1,12 @@
 using CUDA
-import Adapt.@adapt_structure
 include("structures.jl")
 
-Adapt.@adapt_structure Particle; Adapt.@adapt_structure Intermediate_Set;
-Adapt.@adapt_structure Intermediate_Unset; Adapt.Adapt.@adapt_structure offset_and_tilt;
-
 function offset_particle_set(corrections, p_lab, int)
     """Transform from lab to element coords.
     See Bmad manual (2022-11-06) sections 5.6.1, 15.3.1 and 24.2
     **NOTE**: transverse only as of now.
     """
+
     x_ele_int, x_ele, y_ele, px_ele, py_ele, s, c = int.x_ele_int, 
     int.x_ele, int.y_ele, int.px_ele, int.py_ele, int.s, int.c
 
@@ -35,8 +32,7 @@ function offset_particle_set(corrections, p_lab, int)
 
         i += stride
     end
-
-    return
+    return nothing
 end
 
 function offset_particle_unset(corrections, p_ele, int)
@@ -68,6 +64,8 @@ function offset_particle_unset(corrections, p_ele, int)
         py_lab[i] = px[i]*s[i] + py[i]*c[i];)
 
         i += stride 
+
     end
-    return
-end
+    return nothing
+end
+

src/low_level/quad_mat2_calc.jl

@@ -1,8 +1,6 @@
-using CUDA, BenchmarkTools, Adapt
+using CUDA, Adapt
 include("structures.jl")
 
-Adapt.@adapt_structure Intermediate_Elements; Adapt.@adapt_structure quad_calc_input;
-
 function quad_mat2_calc(input, int)
     """Returns 2x2 transfer matrix elements aij and the
     coefficients to calculate the change in z position.
@@ -48,4 +46,3 @@ function quad_mat2_calc(input, int)
     end
     return
 end
-

src/low_level/structures.jl

@@ -1,24 +1,22 @@
- """Structures used for particle tracking on GPU;
-adapt type with Adapt.@adapt_structure in main program"""
+using Adapt
 
-struct Particle{T} # used in several scripts
+struct particle{T} # used in several scripts
     x::T
     px::T
     y::T
     py::T
     z::T
     pz::T
-    s::T
+    sec::Float64
     p0c::T
-    mc2::T
-
+    mc2::Float64
 end
 
-struct Drift{T} # track a drift
+struct drift{T} # track_a_drift
     L::T
 end
 
-struct offset_and_tilt{T} # offset particle
+struct offset_and_tilt{T} # offset_particle
     x_offset::T
     y_offset::T
     tilt::T
@@ -31,18 +29,27 @@ struct z_correction{T} # z energy correction
     ds::T
 end
 
-struct quad_calc_input{T}  # input vals for quad_mat2_calc
+struct quad_calc_input{T}  # quad_mat2_calc
     k1::T
-    len::T
+    len::Float64
     rel_p::T
 end
 
+struct quad_input{T}  #track_a_quadrupole
+    L::Float64
+    K1::T
+    NUM_STEPS::Number
+    X_OFFSET::T
+    Y_OFFSET::T
+    TILT::T
+end
+
 
 """Intermediate structs in order 
 to not dynamically allocate memory to
 intermediate calculations"""
 
-struct Intermediate_Drift{T} # track a drift
+struct int_drift{T} # track_a_drift
     P::T
     Px::T
     Py::T
@@ -51,7 +58,7 @@ struct Intermediate_Drift{T} # track a drift
     dz::T
 end
 
-struct Intermediate_Set{T} # offset particle
+struct int_set{T} # offset_particle
     x_ele_int::T
     x_ele::T
     y_ele::T
@@ -61,25 +68,22 @@ struct Intermediate_Set{T} # offset particle
     c::T
 end
 
-
-struct Intermediate_Unset{T} # offset particle
+struct int_unset{T}
     x_lab::T
     y_lab::T
     px_lab::T
     py_lab::T
-    s::T
-    c::T
 end
 
-struct Intermediate_z_Correction{T} # z energy correction
+struct int_z_correction{T} # z energy correction
     beta::T
     beta0::T
     e_tot::T
     evaluation::T
     dz::T
 end
 
-struct Intermediate_Elements{T} # quad_mat2_calc intermediate steps
+struct int_elements{T} # quad_mat2_calc intermediate steps
     sqrt_k::T
     sk_l::T
     cx::T
@@ -91,4 +95,14 @@ struct Intermediate_Elements{T} # quad_mat2_calc intermediate steps
     c1::T
     c2::T
     c3::T
-end
+end
+
+"""adapting structs to bitstype"""
+Adapt.@adapt_structure particle; Adapt.@adapt_structure drift; Adapt.@adapt_structure offset_and_tilt;
+Adapt.@adapt_structure z_correction; Adapt.@adapt_structure quad_calc_input; Adapt.@adapt_structure quad_input;
+Adapt.@adapt_structure int_drift; Adapt.@adapt_structure int_set; Adapt.@adapt_structure int_unset;
+Adapt.@adapt_structure int_z_correction; Adapt.@adapt_structure int_elements;
+
+"""Structures used for particle tracking on GPU;
+adapt type with Adapt.@adapt_structure in main program"""
+

src/track_a_quadrupole.jl

@@ -0,0 +1,34 @@
+using CUDA, Adapt
+include("low_level/structures.jl"); include("low_level/offset_particle.jl"); include("low_level/int_arrays.jl");
+
+fill_quadrupole(10_000)
+
+function track_a_quadrupole(p_in, quad)
+    """Tracks the incoming Particle p_in though quad element and
+    returns the outgoing particle.
+    See Bmad manual section 24.15
+    """
+    l = quad.L
+    k1 = quad.K1
+    n_step = quad.NUM_STEPS  # number of divisions
+    step_len = l / n_step  # length of division
+
+    x_off = quad.X_OFFSET
+    y_off = quad.Y_OFFSET
+    tilt = quad.TILT
+
+    corrections = offset_and_tilt(x_off, y_off, tilt)
+    int = int_set(x_ele_int, x_ele, y_ele, px_ele, py_ele, s, c)
+
+    sec = p_in.sec
+    p0c = p_in.p0c
+    mc2 = p_in.mc2
+
+    # --- TRACKING --- :
+
+    offset_particle_set(corrections, p_in, int)
+    x, px, y, py, z, pz = x_ele, px_ele, y_ele, py_ele, p_in.z, p_in.pz
+
+    return nothing
+end
+