From 7a649970f9f96a356df592308971333e40b978ca Mon Sep 17 00:00:00 2001 From: obeznosov Date: Fri, 14 Feb 2025 15:18:51 -0700 Subject: [PATCH 1/4] added MatrixKick GPU extention --- src/GPU/MatrixKick_GPUext.jl | 126 +++++++++++++++++++++++ src/GPU/main.jl | 194 +++++++++++++++++++++++++++++++++++ 2 files changed, 320 insertions(+) create mode 100644 src/GPU/MatrixKick_GPUext.jl create mode 100644 src/GPU/main.jl diff --git a/src/GPU/MatrixKick_GPUext.jl b/src/GPU/MatrixKick_GPUext.jl new file mode 100644 index 0000000..643deae --- /dev/null +++ b/src/GPU/MatrixKick_GPUext.jl @@ -0,0 +1,126 @@ +module MatrixKick_GPUext +using StructArrays +using BeamTracking +using BeamTracking: get_work +using CUDA +export track! + +Base.@kwdef struct Drift{T} + L::T # drift length / m +end + +Base.@kwdef struct Quadrupole{T} + L::T # quadrupole length / m + Bn1::T # quadrupole gradient / (T·m^-1) +end + +# GPU Kernel for Drift Tracking +function track_drift_kernel!(x, y, z, px, py, pz, L, beta_ref, tilde_m, gamsqr_ref) + i = threadIdx().x + (blockIdx().x - 1) * blockDim().x + if i <= length(x) + ps = sqrt((1.0 + pz[i])^2 - (px[i]^2 + py[i]^2)) # P_s + x[i] += px[i] * L / ps + y[i] += py[i] * L / ps + z[i] -= ( (1.0 + pz[i]) * L + * ((px[i]^2 + py[i]^2) - pz[i] * (2 + pz[i]) / gamsqr_ref) + / (beta_ref * sqrt((1.0 + pz[i])^2 + tilde_m^2) * ps + * (beta_ref * sqrt((1.0 + pz[i])^2 + tilde_m^2) + ps)) + ) + end + return nothing +end + +function track!(bunch::Bunch, ele::MatrixKick_GPUext.Drift) + L = ele.L + tilde_m = 1 / bunch.beta_gamma_ref + gamsqr_ref = 1 + bunch.beta_gamma_ref^2 + beta_ref = bunch.beta_gamma_ref / sqrt(gamsqr_ref) + + v = bunch.v # Assuming v is a StructArray with CuArrays + + # Launch GPU Kernel + CUDA.@sync @cuda threads=256 blocks=cld(length(v.x), 256) track_drift_kernel!( + v.x, v.y, v.z, v.px, v.py, v.pz, L, beta_ref, tilde_m, gamsqr_ref + ) + + return bunch +end + +# GPU Kernel for Quadrupole Matrix Kick +function track_quad_mx_kernel!(x, y, z, px, py, pz, k2_num, s) + i = threadIdx().x + (blockIdx().x - 1) * blockDim().x + if i <= length(x) + p = 1 + pz[i] + k2 = k2_num / p + ks = sqrt(abs(k2)) * s + xp = px[i] / p + yp = py[i] / p + + focus = k2_num >= 0 + defocus = k2_num < 0 + + cx = focus * cos(ks) + defocus * cosh(ks) + cy = focus * cosh(ks) + defocus * cos(ks) + sx = focus * sinc(ks) + defocus * sinh(ks) + sy = focus * sinh(ks) + defocus * sinc(ks) + + x[i] = x[i] * cx + xp * s * sx + px[i] = px[i] * cx - x[i] * p * k2 * s * sx + y[i] = y[i] * cy + yp * s * sy + py[i] = py[i] * cy + y[i] * p * k2 * s * sy + z[i] -= (s / 4) * (xp^2 + yp^2 - k2 * x[i]^2 + k2 * y[i]^2) + end + return nothing +end + +# GPU Kernel for Quadrupole Kick +function track_quad_k_kernel!(x, y, z, px, py, pz, betgam_ref, s) + i = threadIdx().x + (blockIdx().x - 1) * blockDim().x + if i <= length(x) + tilde_m = 1 / betgam_ref + beta_ref = betgam_ref / sqrt(1 + betgam_ref^2) + gamsqr_ref = 1 + betgam_ref^2 + + p = 1 + pz[i] + ptr2 = px[i]^2 + py[i]^2 + ps = sqrt(p^2 - ptr2) + + x[i] += s * px[i] / p * ptr2 / (ps * (p + ps)) + y[i] += s * py[i] / p * ptr2 / (ps * (p + ps)) + z[i] -= s * ( (1.0 + pz[i]) + * (ptr2 - pz[i] * (2 + pz[i]) / gamsqr_ref) + / ( beta_ref * sqrt((1.0 + pz[i])^2 + tilde_m^2) * ps + * (beta_ref * sqrt((1.0 + pz[i])^2 + tilde_m^2) + ps)) + - ptr2 / (2 * (1 + pz[i])^2) + ) + end + return nothing +end + +function track!(bunch::Bunch, ele::MatrixKick_GPUext.Quadrupole) + L = ele.L + k2_num = ele.Bn1 / brho(massof(bunch.species), bunch.beta_gamma_ref, chargeof(bunch.species)) + + v = bunch.v + @captured begin + # First Matrix Kick + CUDA.@sync @cuda threads=256 blocks=cld(length(v.x), 256) track_quad_mx_kernel!( + v.x, v.y, v.z, v.px, v.py, v.pz, k2_num, L / 2 + ) + + # Quadrupole Kick + CUDA.@sync @cuda threads=256 blocks=cld(length(v.x), 256) track_quad_k_kernel!( + v.x, v.y, v.z, v.px, v.py, v.pz, bunch.beta_gamma_ref, L + ) + + # Second Matrix Kick + CUDA.@sync @cuda threads=256 blocks=cld(length(v.x), 256) track_quad_mx_kernel!( + v.x, v.y, v.z, v.px, v.py, v.pz, k2_num, L / 2 + ) + end + + return bunch +end + +end # module MatrixKick_GPUext + diff --git a/src/GPU/main.jl b/src/GPU/main.jl new file mode 100644 index 0000000..e36d0e8 --- /dev/null +++ b/src/GPU/main.jl @@ -0,0 +1,194 @@ +import Pkg; Pkg.activate(".") +using CUDA +using Test + +include("MatrixKick_GPUext.jl") + +function main() + if CUDA.functional() + @info "CUDA is functional" + else + @info "CUDA is unavailible" + end + + test_ele() +end + +using BeamTracking +function test_ele() + # define elements + # -- drifts + ld1 = 0.15; # m + ld2 = 0.75; # m + ld3 = 2.00; # m + dr1 = MatrixKick_GPUext.Drift(L = ld1) + dr2 = MatrixKick_GPUext.Drift(L = ld2) + dr3 = MatrixKick_GPUext.Drift(L = ld3) + + # -- quadrupoles + lq1 = 0.05; # m + lq2 = 0.30; # m + lq3 = 0.75; # m + lq4 = 1.20; # m + gr1 = 0.20; # T/m + gr2 = 0.60; # T/m + gr3 = 1.20; # T/m + gr4 = 3.50; # T/m + qf1 = MatrixKick_GPUext.Quadrupole(L = lq1, Bn1 = +gr1) + qd1 = MatrixKick_GPUext.Quadrupole(L = lq1, Bn1 = -gr1) + qf2 = MatrixKick_GPUext.Quadrupole(L = lq2, Bn1 = +gr2) + qd2 = MatrixKick_GPUext.Quadrupole(L = lq2, Bn1 = -gr2) + qf3 = MatrixKick_GPUext.Quadrupole(L = lq3, Bn1 = +gr3) + qd3 = MatrixKick_GPUext.Quadrupole(L = lq3, Bn1 = -gr3) + qf4 = MatrixKick_GPUext.Quadrupole(L = lq4, Bn1 = +gr4) + qd4 = MatrixKick_GPUext.Quadrupole(L = lq4, Bn1 = -gr4) + + # define beams + # -- species + e_minus = Species("electron") + p_plus = Species("proton") + mec2 = massof(e_minus) # 0.51099895069 MeV + mpc2 = massof(p_plus) # 938.27208943 MeV + # -- kinetic energy + ek1 = 5.e3; # eV + ek2 = 1.e6; # eV + ek3 = 1.e9; # eV + ek4 = 250.e9; # eV + # -- βγ + bg1 = sqrt(ek1 / mec2 * (ek1 / mec2 + 2)) + bg2 = sqrt(ek2 / mec2 * (ek2 / mec2 + 2)) + bg3 = sqrt(ek3 / mec2 * (ek3 / mec2 + 2)) + bg4 = sqrt(ek4 / mpc2 * (ek4 / mpc2 + 2)) + # -- pc + pc1 = mec2 * bg1 + pc2 = mec2 * bg2 + pc3 = mec2 * bg3 + pc4 = mpc2 * bg4 + # -- beams + # beam1.initial + xi = ([ 0.000, 0.000, 0.000, 0.002, 0.00200, -0.00200 ] .* ones(1024)')[:] + pxi = ([ 0.000, 0.000, 0.000, 0.000, 0.00075, -0.00075 ] .* ones(1024)')[:] + yi = ([ 0.000, 0.000, 0.000, 0.001, 0.00100, -0.00100 ] .* ones(1024)')[:] + pyi = ([ 0.000, 0.000, 0.000, 0.000, 0.00030, -0.00030 ] .* ones(1024)')[:] + zi = ([ 0.000, 0.000, 0.000, 0.000, 0.00000, 0.00000 ] .* ones(1024)')[:] + pzi = ([ 0.000, 0.001, -0.001, 0.001, 0.00100, 0.00100 ] .* ones(1024)')[:] + + # beam1.dr1.final + xf_dr1 = ([ 0., 0., 0., 2.e-3, 2.1123876489808660e-3, -2.1123876489808660e-3 ] .* ones(1024)')[:] + pxf_dr1 = ([ 0., 0., 0., 0., 7.5e-4, -7.5e-4 ] .* ones(1024)')[:] + yf_dr1 = ([ 0., 0., 0., 1.e-3, 1.0449550595923462e-3, -1.0449550595923462e-3 ] .* ones(1024)')[:] + pyf_dr1 = ([ 0., 0., 0., 0., 3.e-4, -3.e-4 ] .* ones(1024)')[:] + zf_dr1 = ([ 0., 1.4710284465059844e-4, -1.4711135596840458e-4, 1.4710284465059844e-4, 1.4705400485512816e-4, 1.4705400485512816e-4 ] .* ones(1024)')[:] + pzf_dr1 = ([ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] .* ones(1024)')[:] + + # beam2.dr2.final + xf_dr2 = [ 0., 0., 0., 0.002, 0.0025619382449043287, -0.0025619382449043287 ] + pxf_dr2 = [ 0., 0., 0., 0., 0.00075, -0.00075 ] + yf_dr2 = [ 0., 0., 0., 0.001, 0.0012247752979617315, -0.0012247752979617315 ] + pyf_dr2 = [ 0., 0., 0., 0., 0.0003, -0.0003 ] + zf_dr2 = [ 0., 0.00008566359457101641, -0.00008589149602558208, 0.00008566359457101641, 0.00008541939559366522, 0.00008541939559366522 ] + pzf_dr2 = [ 0., 0.001, -0.001, 0.001, 0.001, 0.001 ] + + # beam3.dr3.final + xf_dr3 = [ 0., 0., 0., 0.002, 0.003498501986411544, -0.003498501986411544 ] + pxf_dr3 = [ 0., 0., 0., 0., 0.00075, -0.00075 ] + yf_dr3 = [ 0., 0., 0., 0.001, 0.0015994007945646172, -0.0015994007945646172 ] + pyf_dr3 = [ 0., 0., 0., 0., 0.0003, -0.0003 ] + zf_dr3 = [ 0., 5.209250185095532e-10, -5.224901403178192e-10, 5.209250185095532e-10, -6.506763479180273e-7, -6.506763479180273e-7 ] + pzf_dr3 = [ 0., 0.001, -0.001, 0.001, 0.001, 0.001 ] + + # beam4.dr4.final + xf_dr4 = [ 0., 0., 0., 0.002, 0.003498501986411544, -0.003498501986411544 ] + pxf_dr4 = [ 0., 0., 0., 0., 0.00075, -0.00075 ] + yf_dr4 = [ 0., 0., 0., 0.001, 0.0015994007945646172, -0.0015994007945646172 ] + pyf_dr4 = [ 0., 0., 0., 0., 0.0003, -0.0003 ] + zf_dr4 = [ 0., 2.7919184691863886e-8, -2.800306686850912e-8, 2.7919184691863886e-8, -6.232780882446728e-7, -6.232780882446728e-7 ] + pzf_dr4 = [ 0., 0.001, -0.001, 0.001, 0.001, 0.001 ] + + # beam1.qf1.final + xf_qf1 = [ 0., 0., 0., 4.4834792779340600e-3, 4.5356143287504990e-3, -4.5356143287504990e-3 ] + pxf_qf1 = [ 0., 0., 0., 1.1607821136240948e-1, 1.1776162121669208e-1, -1.1776162121669208e-1 ] + yf_qf1 = [ 0., 0., 0., 1.2400905948673489e-4, 1.3427609296030678e-4, -1.3427609296030678e-4 ] + pyf_qf1 = [ 0., 0., 0., -2.8691666098954356e-2, -2.8653744321335432e-2, 2.8653744321335432e-2 ] + zf_qf1 = [ 0., 4.903428155019947e-5, -4.903711865613486e-5, -4.8701323139842656e-5, -5.1605970700562340e-5, -5.1605970700562340e-5 ] + pzf_qf1 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] + + # beam1.qd1.final + xf_qd1 = [ 0., 0., 0., 2.4834727340278294e-4, 2.7409827330240064e-4, -2.7409827330240064e-4 ] + pxf_qd1 = [ 0., 0., 0., -5.7391734255615580e-2, -5.7299059109537503e-2, 5.7299059109537503e-2 ] + yf_qd1 = [ 0., 0., 0., 2.2414071638535435e-3, 2.2621899982980444e-3, -2.2621899982980444e-3 ] + pyf_qd1 = [ 0., 0., 0., 5.8033153149417160e-2, 5.8705242601074584e-2, -5.8705242601074584e-2 ] + zf_qd1 = [ 0., 4.9034281550199470e-5, -4.9037118656134860e-5, -1.1242623339116514e-5, -1.1118475705649755e-5, -1.1118475705649755e-5 ] + pzf_qd1 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] + + # beam2.qf2.final + xf_qf2 = [ 0., 0., 0., 2.9271671401041872e-2, 3.0251479090608963e-2, -3.0251479090608963e-2 ] + pxf_qf2 = [ 0., 0., 0., 3.2854870071095094e-1, 3.3959282727762820e-1, -3.3959282727762820e-1 ] + yf_qf2 = [ 0., 0., 0., -9.7278287676729380e-4, -9.7883210731348450e-4, 9.7883210731348450e-4 ] + pyf_qf2 = [ 0., 0., 0., 2.6415505265087600e-3, 2.3544440936341645e-3, -2.3544440936341645e-3 ] + zf_qf2 = [ 0., 3.4265437828406564e-5, -3.4356598410232830e-5, -2.3378443788359830e-3, -2.5012489357825860e-3, -2.5012489357825860e-3 ] + pzf_qf2 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] + + # beam2.qd2.final + xf_qd2 = [ 0., 0., 0., -0.0019464186216836795, -0.001961645985092862, 0.001961645985092862 ] + pxf_qd2 = [ 0., 0., 0., 0.005200326165715375, 0.004472438532736971, -0.004472438532736971 ] + yf_qd2 = [ 0., 0., 0., 0.014608704236165976, 0.014997970393040412, -0.014997970393040412 ] + pyf_qd2 = [ 0., 0., 0., 0.16398929979812307, 0.16837903611031713, -0.16837903611031713 ] + zf_qd2 = [ 0., 3.4265437828406564e-5, -3.435659841023283e-5, -5.899497183747271e-4, -6.222650472445767e-4, -6.222650472445767e-4 ] + pzf_qd2 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] + + # beam3.qf3.final + xf_qf3 = [ 0., 0., 0., 0.002205479704035602, 0.0027865339903883585, -0.0027865339903883585 ] + pxf_qf3 = [ 0., 0., 0., 0.0005576983180317967, 0.0013847532608658173, -0.0013847532608658173 ] + yf_qf3 = [ 0., 0., 0., 0.0009006624003320741, 0.0011179443238845692, -0.0011179443238845692 ] + pyf_qf3 = [ 0., 0., 0., -0.0002606852268330615, 9.513485195908057e-6, -9.513485195908057e-6 ] + zf_qf3 = [ 0., 1.9534688194108246e-10, -1.959338026191822e-10, -4.630230762164262e-8, -4.3665730716563075e-7, -4.3665730716563075e-7 ] + pzf_qf3 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] + + # beam3.qd3.final + xf_qd3 = [ 0., 0., 0., 0.0018013248008431747, 0.0023445295165179687, -0.0023445295165179687 ] + pxf_qd3 = [ 0., 0., 0., -0.0005213704536906773, 0.0001541263391654702, -0.0001541263391654702 ] + yf_qd3 = [ 0., 0., 0., 0.0011027398519220506, 0.0013351614586315588, -0.0013351614586315588 ] + pyf_qd3 = [ 0., 0., 0., 0.00027884915900320065, 0.0006096711218370137, -0.0006096711218370137 ] + zf_qd3 = [ 0., 1.9534688194108246e-10, -1.959338026191822e-10, -4.4102076263621526e-8, -1.6795543457606254e-7, -1.6795543457606254e-7 ] + pzf_qd3 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] + + # beam4.qf4.final + xf_qf4 = [ 0., 0., 0., 0.0019939877700227713, 0.002892187842249561, -0.002892187842249561 ] + pxf_qf4 = [ 0., 0., 0., -0.000010025375230046909, 0.0007377200378071899, -0.0007377200378071899 ] + yf_qf4 = [ 0., 0., 0., 0.0010030091302526826, 0.0013630102695159176, -0.0013630102695159176 ] + pyf_qf4 = [ 0., 0., 0., 5.022748546347732e-6, 0.0003059254879156335, -0.0003059254879156335 ] + zf_qf4 = [ 0., 1.675151081511833e-8, -1.680184012110547e-8, 1.6726401735327598e-8, -3.698308917351584e-7, -3.698308917351584e-7 ] + pzf_qf4 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] + + # beam4.qd4.final + xf_qd4 = [ 0., 0., 0., 0.002006018260505365, 0.00290602111426842, -0.00290602111426842 ] + pxf_qd4 = [ 0., 0., 0., 0.000010045497092695465, 0.0007623023455299651, -0.0007623023455299651 ] + yf_qd4 = [ 0., 0., 0., 0.0009969938850113856, 0.0013562739161008426, -0.0013562739161008426 ] + pyf_qd4 = [ 0., 0., 0., -5.012687615023454e-6, 0.0002940854775943521, -0.0002940854775943521 ] + zf_qd4 = [ 0., 1.675151081511833e-8, -1.680184012110547e-8, 1.6726365460219405e-8, -3.78176641902771e-7, -3.78176641902771e-7 ] + pzf_qd4 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] + + # === drifts === + # + # 5 keV electron + beam1_d = Bunch(species = e_minus, beta_gamma_ref = bg1, + x = CuArray(xi), px = CuArray(pxi), y = CuArray(yi), py = CuArray(pyi), z = CuArray(zi), pz = CuArray(pzi)) + MatrixKick_GPUext.track!(beam1_d, dr1); + + beam1_h = Bunch(species = e_minus, beta_gamma_ref = bg1, + x = Array(beam1_d.v.x), px = Array(beam1_d.v.px), y = Array(beam1_d.v.y), py = Array(beam1_d.v.py), z = Array(beam1_d.v.z), pz = Array(beam1_d.v.pz)) + + @test beam1_h.v.x ≈ xf_dr1 (rtol=5.e-13) + @test beam1_h.v.y ≈ yf_dr1 (rtol=5.e-13) + @test beam1_h.v.z ≈ zf_dr1 (rtol=5.e-13) + @test beam1_h.v.px == pxf_dr1 + @test beam1_h.v.py == pyf_dr1 + @test beam1_h.v.pz == pzf_dr1 +end + + +main() + + + From f697225178a2c3c5833a06b05169eca3c734d052 Mon Sep 17 00:00:00 2001 From: obeznosov Date: Fri, 14 Feb 2025 17:50:41 -0700 Subject: [PATCH 2/4] Implemented RF cavity test and kickdrift --- src/MatrixKick/MatrixKick.jl | 58 +++++++++++++++++++++++++++++++++++- test/element_tests.jl | 34 +++++++++++++++++++++ 2 files changed, 91 insertions(+), 1 deletion(-) diff --git a/src/MatrixKick/MatrixKick.jl b/src/MatrixKick/MatrixKick.jl index 469bf7f..5fbe451 100644 --- a/src/MatrixKick/MatrixKick.jl +++ b/src/MatrixKick/MatrixKick.jl @@ -15,6 +15,13 @@ Base.@kwdef struct Quadrupole{T} Bn1::T # quadrupole gradient / (T·m^-1) end +Base.@kwdef struct ThinLensRFCavity{T} + L::T # RF-cavity length / m + V::T # Cavity voltage / V + k::T # RF wavenumber (2π/λ) / raduans·m^-1 + phi0::T # RF phase offset / radians +end + function track!(bunch::Bunch, ele::MatrixKick.Drift; work=get_work(bunch, Val{1}())) #= @@ -67,7 +74,7 @@ function track!(bunch::Bunch, ele::MatrixKick.Quadrupole; work=get_work(bunch, V v_work = StructArray{Coord{eltype(work[1])}}((work[1], work[2], work[3], work[4], work[5], work[6])) trackQuadMx!(v_work, v, k2_num, L / 2) - trackQuadK!( v, v_work, bunch.beta_gamma_ref, L) + trackQuadK!( v, v_work, bunch.beta_gamma_ref, L) trackQuadMx!(v_work, v, k2_num, L / 2) v .= v_work @@ -150,5 +157,54 @@ function trackQuadK!(vf, vi, betgam_ref, s) end # function trackQ!M::Quadrupole() + +""" +function track!(bunch::Bunch, ele::MatrixKick.ThinLensRFCavity; work=get_work(bunch, Val{1}())) + +Apply a thin RF cavity map in-place, where momenta (p_x, p_y, p_z) +and rest mass are expressed in eV (treating c=1 for momentum units), +while the longitudinal coordinate z is in meters. + +**Algorithm**: +1) Half-drift by L/2 using the current momentum (which determines velocity). +2) Kick in pz by `Δpz = q*V (eV) × `cos(k*z + phi0)` +3) Half-drift by L/2 using the updated momentum. + +""" +function track!(bunch::Bunch, ele::MatrixKick.ThinLensRFCavity; work=get_work(bunch, Val{1}())) + v = bunch.v + q = chargeof(bunch.species) + m = massof(bunch.species) + k = ele.k + # pleccing minus according to 4.46 in manual + dE = -q * ele.V # energy gain/loss in eV + L = ele.L + + @FastGTPSA! begin + # total momentum p [eV] + #------------------------------------------------------ + # Step 1: Half-drift using old momentum + #------------------------------------------------------ + @. work[1] = sqrt(v.px^2 + v.py^2 + v.pz^2 + m^2) + @. work[1] = ifelse(work[1] > 0.0, work[1], 1.0) + @. v.z += 0.5 * L * v.pz/work[1] + + #------------------------------------------------------ + # Step 2: RF momentum kick in pz + #------------------------------------------------------ + @. v.pz += dE * cos(k*v.z + ele.phi0) + + #------------------------------------------------------ + # Step 3: Half-drift using updated momentum + #------------------------------------------------------ + @. work[1] = sqrt(v.px^2 + v.py^2 + v.pz^2 + m^2) + @. work[1] = ifelse(work[1] > 0.0, work[1], 1.0) + @. v.z += 0.5 * L * v.pz / work[1] + + end + # Spin unchanged + return bunch +end + end # module MatrixKick diff --git a/test/element_tests.jl b/test/element_tests.jl index 8026bc9..9c63e11 100644 --- a/test/element_tests.jl +++ b/test/element_tests.jl @@ -28,6 +28,16 @@ qf3 = MatrixKick.Quadrupole(L = lq3, Bn1 = +gr3) qd3 = MatrixKick.Quadrupole(L = lq3, Bn1 = -gr3) qf4 = MatrixKick.Quadrupole(L = lq4, Bn1 = +gr4) qd4 = MatrixKick.Quadrupole(L = lq4, Bn1 = -gr4) +phi0_1 = 0.0 +phi0_2 = 1.2 +rfk = 2π / 0.1 +V1 = 1.e5 +V2 = 2.e5 +lrf1 = 0.0 +lrf2 = 0.02 # m +rf1 = MatrixKick.ThinLensRFCavity(L = lrf1, V = V1, k = rfk, phi0 = phi0_1) # 100 kV, zero length +rf2 = MatrixKick.ThinLensRFCavity(L = lrf2, V = V2, k = rfk, phi0 = phi0_2) # 200 kV, non-zero length + # define beams # -- species @@ -155,6 +165,9 @@ pyf_qd4 = [ 0., 0., 0., -5.012687615023 zf_qd4 = [ 0., 1.675151081511833e-8, -1.680184012110547e-8, 1.6726365460219405e-8, -3.78176641902771e-7, -3.78176641902771e-7 ] pzf_qd4 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] +# beam6.qd4.final +pze_rf2 = [72471.55089533473, 72471.5516661298, 72471.55012453966, 72471.5516661298, 72471.5516661298, 72471.5516661298] + # test individual elements @testset "element_tests" begin # === drifts === @@ -285,6 +298,27 @@ pzf_qd4 = [ 0., 1.e-3, -1.e-3, 1.e-3, @test beam4.v.z ≈ zf_qd4 (rtol=5.e-13) @test beam4.v.pz == pzf_qd4 + beam5 = Bunch(species = e_minus, beta_gamma_ref = bg3, + x = copy(xi), px = copy(pxi), y = copy(yi), py = copy(pyi), z = copy(zi), pz = copy(pzi)) + track!(beam5, rf1); + @test beam5.v.x == xi + @test beam5.v.px == pxi + @test beam5.v.y == yi + @test beam5.v.py == pyi + @test beam5.v.z == zi + expected_dpz = -chargeof(e_minus) * V1* cos.(rfk * zi .+ phi0_1) + @test isapprox(beam5.v.pz, pzi .+ expected_dpz; atol=1e-13) + + beam6 = Bunch(species = e_minus, beta_gamma_ref = bg3, + x = copy(xi), px = copy(pxi), y = copy(yi), py = copy(pyi), z = copy(zi), pz = copy(pzi)) + track!(beam6, rf2); + @test beam6.v.x == xi + @test beam6.v.px == pxi + @test beam6.v.y == yi + @test beam6.v.py == pyi + @test beam6.v.z != zi + @test isapprox(beam6.v.pz, pze_rf2; atol=1e-13) + end From efb3fedd0ad6507d9e0b1a2bb8b38791edd6c23f Mon Sep 17 00:00:00 2001 From: obeznosov Date: Fri, 14 Feb 2025 17:57:55 -0700 Subject: [PATCH 3/4] removed GPU files --- src/GPU/MatrixKick_GPUext.jl | 126 ----------------------- src/GPU/main.jl | 194 ----------------------------------- 2 files changed, 320 deletions(-) delete mode 100644 src/GPU/MatrixKick_GPUext.jl delete mode 100644 src/GPU/main.jl diff --git a/src/GPU/MatrixKick_GPUext.jl b/src/GPU/MatrixKick_GPUext.jl deleted file mode 100644 index 643deae..0000000 --- a/src/GPU/MatrixKick_GPUext.jl +++ /dev/null @@ -1,126 +0,0 @@ -module MatrixKick_GPUext -using StructArrays -using BeamTracking -using BeamTracking: get_work -using CUDA -export track! - -Base.@kwdef struct Drift{T} - L::T # drift length / m -end - -Base.@kwdef struct Quadrupole{T} - L::T # quadrupole length / m - Bn1::T # quadrupole gradient / (T·m^-1) -end - -# GPU Kernel for Drift Tracking -function track_drift_kernel!(x, y, z, px, py, pz, L, beta_ref, tilde_m, gamsqr_ref) - i = threadIdx().x + (blockIdx().x - 1) * blockDim().x - if i <= length(x) - ps = sqrt((1.0 + pz[i])^2 - (px[i]^2 + py[i]^2)) # P_s - x[i] += px[i] * L / ps - y[i] += py[i] * L / ps - z[i] -= ( (1.0 + pz[i]) * L - * ((px[i]^2 + py[i]^2) - pz[i] * (2 + pz[i]) / gamsqr_ref) - / (beta_ref * sqrt((1.0 + pz[i])^2 + tilde_m^2) * ps - * (beta_ref * sqrt((1.0 + pz[i])^2 + tilde_m^2) + ps)) - ) - end - return nothing -end - -function track!(bunch::Bunch, ele::MatrixKick_GPUext.Drift) - L = ele.L - tilde_m = 1 / bunch.beta_gamma_ref - gamsqr_ref = 1 + bunch.beta_gamma_ref^2 - beta_ref = bunch.beta_gamma_ref / sqrt(gamsqr_ref) - - v = bunch.v # Assuming v is a StructArray with CuArrays - - # Launch GPU Kernel - CUDA.@sync @cuda threads=256 blocks=cld(length(v.x), 256) track_drift_kernel!( - v.x, v.y, v.z, v.px, v.py, v.pz, L, beta_ref, tilde_m, gamsqr_ref - ) - - return bunch -end - -# GPU Kernel for Quadrupole Matrix Kick -function track_quad_mx_kernel!(x, y, z, px, py, pz, k2_num, s) - i = threadIdx().x + (blockIdx().x - 1) * blockDim().x - if i <= length(x) - p = 1 + pz[i] - k2 = k2_num / p - ks = sqrt(abs(k2)) * s - xp = px[i] / p - yp = py[i] / p - - focus = k2_num >= 0 - defocus = k2_num < 0 - - cx = focus * cos(ks) + defocus * cosh(ks) - cy = focus * cosh(ks) + defocus * cos(ks) - sx = focus * sinc(ks) + defocus * sinh(ks) - sy = focus * sinh(ks) + defocus * sinc(ks) - - x[i] = x[i] * cx + xp * s * sx - px[i] = px[i] * cx - x[i] * p * k2 * s * sx - y[i] = y[i] * cy + yp * s * sy - py[i] = py[i] * cy + y[i] * p * k2 * s * sy - z[i] -= (s / 4) * (xp^2 + yp^2 - k2 * x[i]^2 + k2 * y[i]^2) - end - return nothing -end - -# GPU Kernel for Quadrupole Kick -function track_quad_k_kernel!(x, y, z, px, py, pz, betgam_ref, s) - i = threadIdx().x + (blockIdx().x - 1) * blockDim().x - if i <= length(x) - tilde_m = 1 / betgam_ref - beta_ref = betgam_ref / sqrt(1 + betgam_ref^2) - gamsqr_ref = 1 + betgam_ref^2 - - p = 1 + pz[i] - ptr2 = px[i]^2 + py[i]^2 - ps = sqrt(p^2 - ptr2) - - x[i] += s * px[i] / p * ptr2 / (ps * (p + ps)) - y[i] += s * py[i] / p * ptr2 / (ps * (p + ps)) - z[i] -= s * ( (1.0 + pz[i]) - * (ptr2 - pz[i] * (2 + pz[i]) / gamsqr_ref) - / ( beta_ref * sqrt((1.0 + pz[i])^2 + tilde_m^2) * ps - * (beta_ref * sqrt((1.0 + pz[i])^2 + tilde_m^2) + ps)) - - ptr2 / (2 * (1 + pz[i])^2) - ) - end - return nothing -end - -function track!(bunch::Bunch, ele::MatrixKick_GPUext.Quadrupole) - L = ele.L - k2_num = ele.Bn1 / brho(massof(bunch.species), bunch.beta_gamma_ref, chargeof(bunch.species)) - - v = bunch.v - @captured begin - # First Matrix Kick - CUDA.@sync @cuda threads=256 blocks=cld(length(v.x), 256) track_quad_mx_kernel!( - v.x, v.y, v.z, v.px, v.py, v.pz, k2_num, L / 2 - ) - - # Quadrupole Kick - CUDA.@sync @cuda threads=256 blocks=cld(length(v.x), 256) track_quad_k_kernel!( - v.x, v.y, v.z, v.px, v.py, v.pz, bunch.beta_gamma_ref, L - ) - - # Second Matrix Kick - CUDA.@sync @cuda threads=256 blocks=cld(length(v.x), 256) track_quad_mx_kernel!( - v.x, v.y, v.z, v.px, v.py, v.pz, k2_num, L / 2 - ) - end - - return bunch -end - -end # module MatrixKick_GPUext - diff --git a/src/GPU/main.jl b/src/GPU/main.jl deleted file mode 100644 index e36d0e8..0000000 --- a/src/GPU/main.jl +++ /dev/null @@ -1,194 +0,0 @@ -import Pkg; Pkg.activate(".") -using CUDA -using Test - -include("MatrixKick_GPUext.jl") - -function main() - if CUDA.functional() - @info "CUDA is functional" - else - @info "CUDA is unavailible" - end - - test_ele() -end - -using BeamTracking -function test_ele() - # define elements - # -- drifts - ld1 = 0.15; # m - ld2 = 0.75; # m - ld3 = 2.00; # m - dr1 = MatrixKick_GPUext.Drift(L = ld1) - dr2 = MatrixKick_GPUext.Drift(L = ld2) - dr3 = MatrixKick_GPUext.Drift(L = ld3) - - # -- quadrupoles - lq1 = 0.05; # m - lq2 = 0.30; # m - lq3 = 0.75; # m - lq4 = 1.20; # m - gr1 = 0.20; # T/m - gr2 = 0.60; # T/m - gr3 = 1.20; # T/m - gr4 = 3.50; # T/m - qf1 = MatrixKick_GPUext.Quadrupole(L = lq1, Bn1 = +gr1) - qd1 = MatrixKick_GPUext.Quadrupole(L = lq1, Bn1 = -gr1) - qf2 = MatrixKick_GPUext.Quadrupole(L = lq2, Bn1 = +gr2) - qd2 = MatrixKick_GPUext.Quadrupole(L = lq2, Bn1 = -gr2) - qf3 = MatrixKick_GPUext.Quadrupole(L = lq3, Bn1 = +gr3) - qd3 = MatrixKick_GPUext.Quadrupole(L = lq3, Bn1 = -gr3) - qf4 = MatrixKick_GPUext.Quadrupole(L = lq4, Bn1 = +gr4) - qd4 = MatrixKick_GPUext.Quadrupole(L = lq4, Bn1 = -gr4) - - # define beams - # -- species - e_minus = Species("electron") - p_plus = Species("proton") - mec2 = massof(e_minus) # 0.51099895069 MeV - mpc2 = massof(p_plus) # 938.27208943 MeV - # -- kinetic energy - ek1 = 5.e3; # eV - ek2 = 1.e6; # eV - ek3 = 1.e9; # eV - ek4 = 250.e9; # eV - # -- βγ - bg1 = sqrt(ek1 / mec2 * (ek1 / mec2 + 2)) - bg2 = sqrt(ek2 / mec2 * (ek2 / mec2 + 2)) - bg3 = sqrt(ek3 / mec2 * (ek3 / mec2 + 2)) - bg4 = sqrt(ek4 / mpc2 * (ek4 / mpc2 + 2)) - # -- pc - pc1 = mec2 * bg1 - pc2 = mec2 * bg2 - pc3 = mec2 * bg3 - pc4 = mpc2 * bg4 - # -- beams - # beam1.initial - xi = ([ 0.000, 0.000, 0.000, 0.002, 0.00200, -0.00200 ] .* ones(1024)')[:] - pxi = ([ 0.000, 0.000, 0.000, 0.000, 0.00075, -0.00075 ] .* ones(1024)')[:] - yi = ([ 0.000, 0.000, 0.000, 0.001, 0.00100, -0.00100 ] .* ones(1024)')[:] - pyi = ([ 0.000, 0.000, 0.000, 0.000, 0.00030, -0.00030 ] .* ones(1024)')[:] - zi = ([ 0.000, 0.000, 0.000, 0.000, 0.00000, 0.00000 ] .* ones(1024)')[:] - pzi = ([ 0.000, 0.001, -0.001, 0.001, 0.00100, 0.00100 ] .* ones(1024)')[:] - - # beam1.dr1.final - xf_dr1 = ([ 0., 0., 0., 2.e-3, 2.1123876489808660e-3, -2.1123876489808660e-3 ] .* ones(1024)')[:] - pxf_dr1 = ([ 0., 0., 0., 0., 7.5e-4, -7.5e-4 ] .* ones(1024)')[:] - yf_dr1 = ([ 0., 0., 0., 1.e-3, 1.0449550595923462e-3, -1.0449550595923462e-3 ] .* ones(1024)')[:] - pyf_dr1 = ([ 0., 0., 0., 0., 3.e-4, -3.e-4 ] .* ones(1024)')[:] - zf_dr1 = ([ 0., 1.4710284465059844e-4, -1.4711135596840458e-4, 1.4710284465059844e-4, 1.4705400485512816e-4, 1.4705400485512816e-4 ] .* ones(1024)')[:] - pzf_dr1 = ([ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] .* ones(1024)')[:] - - # beam2.dr2.final - xf_dr2 = [ 0., 0., 0., 0.002, 0.0025619382449043287, -0.0025619382449043287 ] - pxf_dr2 = [ 0., 0., 0., 0., 0.00075, -0.00075 ] - yf_dr2 = [ 0., 0., 0., 0.001, 0.0012247752979617315, -0.0012247752979617315 ] - pyf_dr2 = [ 0., 0., 0., 0., 0.0003, -0.0003 ] - zf_dr2 = [ 0., 0.00008566359457101641, -0.00008589149602558208, 0.00008566359457101641, 0.00008541939559366522, 0.00008541939559366522 ] - pzf_dr2 = [ 0., 0.001, -0.001, 0.001, 0.001, 0.001 ] - - # beam3.dr3.final - xf_dr3 = [ 0., 0., 0., 0.002, 0.003498501986411544, -0.003498501986411544 ] - pxf_dr3 = [ 0., 0., 0., 0., 0.00075, -0.00075 ] - yf_dr3 = [ 0., 0., 0., 0.001, 0.0015994007945646172, -0.0015994007945646172 ] - pyf_dr3 = [ 0., 0., 0., 0., 0.0003, -0.0003 ] - zf_dr3 = [ 0., 5.209250185095532e-10, -5.224901403178192e-10, 5.209250185095532e-10, -6.506763479180273e-7, -6.506763479180273e-7 ] - pzf_dr3 = [ 0., 0.001, -0.001, 0.001, 0.001, 0.001 ] - - # beam4.dr4.final - xf_dr4 = [ 0., 0., 0., 0.002, 0.003498501986411544, -0.003498501986411544 ] - pxf_dr4 = [ 0., 0., 0., 0., 0.00075, -0.00075 ] - yf_dr4 = [ 0., 0., 0., 0.001, 0.0015994007945646172, -0.0015994007945646172 ] - pyf_dr4 = [ 0., 0., 0., 0., 0.0003, -0.0003 ] - zf_dr4 = [ 0., 2.7919184691863886e-8, -2.800306686850912e-8, 2.7919184691863886e-8, -6.232780882446728e-7, -6.232780882446728e-7 ] - pzf_dr4 = [ 0., 0.001, -0.001, 0.001, 0.001, 0.001 ] - - # beam1.qf1.final - xf_qf1 = [ 0., 0., 0., 4.4834792779340600e-3, 4.5356143287504990e-3, -4.5356143287504990e-3 ] - pxf_qf1 = [ 0., 0., 0., 1.1607821136240948e-1, 1.1776162121669208e-1, -1.1776162121669208e-1 ] - yf_qf1 = [ 0., 0., 0., 1.2400905948673489e-4, 1.3427609296030678e-4, -1.3427609296030678e-4 ] - pyf_qf1 = [ 0., 0., 0., -2.8691666098954356e-2, -2.8653744321335432e-2, 2.8653744321335432e-2 ] - zf_qf1 = [ 0., 4.903428155019947e-5, -4.903711865613486e-5, -4.8701323139842656e-5, -5.1605970700562340e-5, -5.1605970700562340e-5 ] - pzf_qf1 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] - - # beam1.qd1.final - xf_qd1 = [ 0., 0., 0., 2.4834727340278294e-4, 2.7409827330240064e-4, -2.7409827330240064e-4 ] - pxf_qd1 = [ 0., 0., 0., -5.7391734255615580e-2, -5.7299059109537503e-2, 5.7299059109537503e-2 ] - yf_qd1 = [ 0., 0., 0., 2.2414071638535435e-3, 2.2621899982980444e-3, -2.2621899982980444e-3 ] - pyf_qd1 = [ 0., 0., 0., 5.8033153149417160e-2, 5.8705242601074584e-2, -5.8705242601074584e-2 ] - zf_qd1 = [ 0., 4.9034281550199470e-5, -4.9037118656134860e-5, -1.1242623339116514e-5, -1.1118475705649755e-5, -1.1118475705649755e-5 ] - pzf_qd1 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] - - # beam2.qf2.final - xf_qf2 = [ 0., 0., 0., 2.9271671401041872e-2, 3.0251479090608963e-2, -3.0251479090608963e-2 ] - pxf_qf2 = [ 0., 0., 0., 3.2854870071095094e-1, 3.3959282727762820e-1, -3.3959282727762820e-1 ] - yf_qf2 = [ 0., 0., 0., -9.7278287676729380e-4, -9.7883210731348450e-4, 9.7883210731348450e-4 ] - pyf_qf2 = [ 0., 0., 0., 2.6415505265087600e-3, 2.3544440936341645e-3, -2.3544440936341645e-3 ] - zf_qf2 = [ 0., 3.4265437828406564e-5, -3.4356598410232830e-5, -2.3378443788359830e-3, -2.5012489357825860e-3, -2.5012489357825860e-3 ] - pzf_qf2 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] - - # beam2.qd2.final - xf_qd2 = [ 0., 0., 0., -0.0019464186216836795, -0.001961645985092862, 0.001961645985092862 ] - pxf_qd2 = [ 0., 0., 0., 0.005200326165715375, 0.004472438532736971, -0.004472438532736971 ] - yf_qd2 = [ 0., 0., 0., 0.014608704236165976, 0.014997970393040412, -0.014997970393040412 ] - pyf_qd2 = [ 0., 0., 0., 0.16398929979812307, 0.16837903611031713, -0.16837903611031713 ] - zf_qd2 = [ 0., 3.4265437828406564e-5, -3.435659841023283e-5, -5.899497183747271e-4, -6.222650472445767e-4, -6.222650472445767e-4 ] - pzf_qd2 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] - - # beam3.qf3.final - xf_qf3 = [ 0., 0., 0., 0.002205479704035602, 0.0027865339903883585, -0.0027865339903883585 ] - pxf_qf3 = [ 0., 0., 0., 0.0005576983180317967, 0.0013847532608658173, -0.0013847532608658173 ] - yf_qf3 = [ 0., 0., 0., 0.0009006624003320741, 0.0011179443238845692, -0.0011179443238845692 ] - pyf_qf3 = [ 0., 0., 0., -0.0002606852268330615, 9.513485195908057e-6, -9.513485195908057e-6 ] - zf_qf3 = [ 0., 1.9534688194108246e-10, -1.959338026191822e-10, -4.630230762164262e-8, -4.3665730716563075e-7, -4.3665730716563075e-7 ] - pzf_qf3 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] - - # beam3.qd3.final - xf_qd3 = [ 0., 0., 0., 0.0018013248008431747, 0.0023445295165179687, -0.0023445295165179687 ] - pxf_qd3 = [ 0., 0., 0., -0.0005213704536906773, 0.0001541263391654702, -0.0001541263391654702 ] - yf_qd3 = [ 0., 0., 0., 0.0011027398519220506, 0.0013351614586315588, -0.0013351614586315588 ] - pyf_qd3 = [ 0., 0., 0., 0.00027884915900320065, 0.0006096711218370137, -0.0006096711218370137 ] - zf_qd3 = [ 0., 1.9534688194108246e-10, -1.959338026191822e-10, -4.4102076263621526e-8, -1.6795543457606254e-7, -1.6795543457606254e-7 ] - pzf_qd3 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] - - # beam4.qf4.final - xf_qf4 = [ 0., 0., 0., 0.0019939877700227713, 0.002892187842249561, -0.002892187842249561 ] - pxf_qf4 = [ 0., 0., 0., -0.000010025375230046909, 0.0007377200378071899, -0.0007377200378071899 ] - yf_qf4 = [ 0., 0., 0., 0.0010030091302526826, 0.0013630102695159176, -0.0013630102695159176 ] - pyf_qf4 = [ 0., 0., 0., 5.022748546347732e-6, 0.0003059254879156335, -0.0003059254879156335 ] - zf_qf4 = [ 0., 1.675151081511833e-8, -1.680184012110547e-8, 1.6726401735327598e-8, -3.698308917351584e-7, -3.698308917351584e-7 ] - pzf_qf4 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] - - # beam4.qd4.final - xf_qd4 = [ 0., 0., 0., 0.002006018260505365, 0.00290602111426842, -0.00290602111426842 ] - pxf_qd4 = [ 0., 0., 0., 0.000010045497092695465, 0.0007623023455299651, -0.0007623023455299651 ] - yf_qd4 = [ 0., 0., 0., 0.0009969938850113856, 0.0013562739161008426, -0.0013562739161008426 ] - pyf_qd4 = [ 0., 0., 0., -5.012687615023454e-6, 0.0002940854775943521, -0.0002940854775943521 ] - zf_qd4 = [ 0., 1.675151081511833e-8, -1.680184012110547e-8, 1.6726365460219405e-8, -3.78176641902771e-7, -3.78176641902771e-7 ] - pzf_qd4 = [ 0., 1.e-3, -1.e-3, 1.e-3, 1.e-3, 1.e-3 ] - - # === drifts === - # - # 5 keV electron - beam1_d = Bunch(species = e_minus, beta_gamma_ref = bg1, - x = CuArray(xi), px = CuArray(pxi), y = CuArray(yi), py = CuArray(pyi), z = CuArray(zi), pz = CuArray(pzi)) - MatrixKick_GPUext.track!(beam1_d, dr1); - - beam1_h = Bunch(species = e_minus, beta_gamma_ref = bg1, - x = Array(beam1_d.v.x), px = Array(beam1_d.v.px), y = Array(beam1_d.v.y), py = Array(beam1_d.v.py), z = Array(beam1_d.v.z), pz = Array(beam1_d.v.pz)) - - @test beam1_h.v.x ≈ xf_dr1 (rtol=5.e-13) - @test beam1_h.v.y ≈ yf_dr1 (rtol=5.e-13) - @test beam1_h.v.z ≈ zf_dr1 (rtol=5.e-13) - @test beam1_h.v.px == pxf_dr1 - @test beam1_h.v.py == pyf_dr1 - @test beam1_h.v.pz == pzf_dr1 -end - - -main() - - - From 98768d0c6c46415842d109ef1bd91aa717c7ce7a Mon Sep 17 00:00:00 2001 From: obeznosov Date: Thu, 20 Feb 2025 14:57:54 -0700 Subject: [PATCH 4/4] changed drift in cavity to pz/delta --- src/MatrixKick/MatrixKick.jl | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/MatrixKick/MatrixKick.jl b/src/MatrixKick/MatrixKick.jl index 5fbe451..aea0eb3 100644 --- a/src/MatrixKick/MatrixKick.jl +++ b/src/MatrixKick/MatrixKick.jl @@ -185,9 +185,10 @@ function track!(bunch::Bunch, ele::MatrixKick.ThinLensRFCavity; work=get_work(bu #------------------------------------------------------ # Step 1: Half-drift using old momentum #------------------------------------------------------ - @. work[1] = sqrt(v.px^2 + v.py^2 + v.pz^2 + m^2) + # @. work[1] = sqrt(v.px^2 + v.py^2 + v.pz^2 + m^2) + @. work[1] = sqrt((1.0 + v.pz)^2 - (v.px^2 + v.py^2)) # P_s @. work[1] = ifelse(work[1] > 0.0, work[1], 1.0) - @. v.z += 0.5 * L * v.pz/work[1] + @. v.z += 0.5 * L * v.pz / work[1] #------------------------------------------------------ # Step 2: RF momentum kick in pz @@ -197,9 +198,9 @@ function track!(bunch::Bunch, ele::MatrixKick.ThinLensRFCavity; work=get_work(bu #------------------------------------------------------ # Step 3: Half-drift using updated momentum #------------------------------------------------------ - @. work[1] = sqrt(v.px^2 + v.py^2 + v.pz^2 + m^2) + @. work[1] = sqrt((1.0 + v.pz)^2 - (v.px^2 + v.py^2)) # P_s @. work[1] = ifelse(work[1] > 0.0, work[1], 1.0) - @. v.z += 0.5 * L * v.pz / work[1] + @. v.z += 0.5 * L * v.pz / work[1] end # Spin unchanged