diff --git a/src/low_level/int_arrays.jl b/src/low_level/int_arrays.jl new file mode 100644 index 0000000..6bc9901 --- /dev/null +++ b/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) + diff --git a/src/low_level/offset_particle.jl b/src/low_level/offset_particle.jl index 11f5618..f372741 100644 --- a/src/low_level/offset_particle.jl +++ b/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 \ No newline at end of file + return nothing +end + diff --git a/src/low_level/quad_mat2_calc.jl b/src/low_level/quad_mat2_calc.jl index 56b4234..d2aceb8 100644 --- a/src/low_level/quad_mat2_calc.jl +++ b/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 - diff --git a/src/low_level/structures.jl b/src/low_level/structures.jl index 834d653..31bbda1 100644 --- a/src/low_level/structures.jl +++ b/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,17 +68,14 @@ 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 @@ -79,7 +83,7 @@ struct Intermediate_z_Correction{T} # z energy correction 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 \ No newline at end of file +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""" + diff --git a/src/track_a_quadrupole.jl b/src/track_a_quadrupole.jl new file mode 100644 index 0000000..3b26912 --- /dev/null +++ b/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 +