Skip to content

Commit

Permalink
Merge pull request #27 from bmad-sim/work
Browse files Browse the repository at this point in the history 
tests update
  • Loading branch information
@mattsignorelli
mattsignorelli authored Nov 11, 2024
2 parents e7bd67c + 8c7253f commit 11646ed
Show file tree
Hide file tree
Showing 2 changed files with 66 additions and 76 deletions.
2 changes: 1 addition & 1 deletion src/Linear/Linear.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ Track through a linear drift.
- `beam` -- Input/output beam before/after tracking through
- `ele` -- `Linear.Drift` type element
"""
function track!(beam::Beam, ele::Linear.Drift)
function track!(beam::Beam, ele::Linear.Drift; work=get_work(beam, Val{0}()))
L = ele.L
v = beam.v

Expand Down
140 changes: 65 additions & 75 deletions test/runtests.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,89 +5,79 @@ using Test,
BenchmarkTools,
GTPSA

@testset "BeamTracking.jl" begin
# Test setup
BenchmarkTools.DEFAULT_PARAMETERS.gctrial = false
BenchmarkTools.DEFAULT_PARAMETERS.evals = 2

# Linear tests -----------------------
d = Descriptor(6, 1) # 6 variables, order 1
beta_gamma_ref = 1.0
gamma_ref = sr_gamma(beta_gamma_ref)
species = Species("electron")
brho_ref = brho(massof(species), beta_gamma_ref, chargeof(species))
tol = 1e-14

# Linear drift test:
beam = Beam(d, beta_gamma_ref=beta_gamma_ref) # Creates Beam as identity map using GTPSA Descriptor
L_d = 5.0
drift = Linear.Drift(L=L_d)
track!(beam, drift)

M_drift_expected = [1.0 L_d 0.0 0.0 0.0 0.0;
0.0 1.0 0.0 0.0 0.0 0.0;
0.0 0.0 1.0 L_d 0.0 0.0;
0.0 0.0 0.0 1.0 0.0 0.0;
0.0 0.0 0.0 0.0 1.0 L_d/gamma_ref^2;
0.0 0.0 0.0 0.0 0.0 1.0]
BenchmarkTools.DEFAULT_PARAMETERS.gctrial = false
BenchmarkTools.DEFAULT_PARAMETERS.evals = 2

# Soon we generalize this to test_map...
function test_matrix(ele, n_work, M_expected; type_stable=true, no_allocs=true, tol=1e-14,
beta_gamma_ref=1.0, species=Species("electron") )

p = Particle(beam, 1) # Converts SoA to 1 struct
d = Descriptor(6, 1) # 6 variables, order 1
beam = Beam(d, beta_gamma_ref=beta_gamma_ref, species=species)
work = BeamTracking.get_work(beam, Val(n_work))

# 1) Correctness:
@test norm(GTPSA.jacobian(p.v) - M_drift_expected) < tol
# 2) Type-stability:
@test_opt track!(beam, drift)
# 3) Allocations:
@test @ballocated(track!($beam, $drift)) == 0
track!(beam, ele, work=work)
p = Particle(beam, 1)

# Linear Quadrupole test:
# Focusing:
beam = Beam(d, beta_gamma_ref=beta_gamma_ref)
L_q = 1.2
K1n = 0.36
B1 = K1n*brho_ref
qf = Linear.Quadrupole(B1=B1,L=L_q)
work = BeamTracking.get_work(beam, Val{1}())
track!(beam, qf, work=work)

p = Particle(beam, 1) # Converts SoA to 1 struct

M_qf_x = [cos(sqrt(K1n)*L_q) sincu(sqrt(K1n)*L_q)*L_q;
-sqrt(K1n)*sin(sqrt(K1n)*L_q) cos(sqrt(K1n)*L_q) ;]
M_qf_y = [cosh(sqrt(K1n)*L_q) sinhcu(sqrt(K1n)*L_q)*L_q;
sqrt(K1n)*sinh(sqrt(K1n)*L_q) cosh(sqrt(K1n)*L_q) ;]
M_qf_z = [1.0 L_q/gamma_ref^2;
0.0 1.0 ;]
# 1) Correctness
@test norm(GTPSA.jacobian(p.v) - M_expected) < tol
# 2) Type stability
if type_stable
@test_opt track!(beam, ele, work=work)
end
# 3) No Allocations
if no_allocs
@test @ballocated(track!($beam, $ele, work=$work)) == 0
end
end

M_qf_expected = zeros(6,6)
M_qf_expected[1:2,1:2] = M_qf_x
M_qf_expected[3:4,3:4] = M_qf_y
M_qf_expected[5:6,5:6] = M_qf_z
# 1) Correctness:
@test norm(GTPSA.jacobian(p.v) - M_qf_expected) < tol
# 2) Type-stability:
@test_opt track!(beam, qf, work=work)
# 3) Allocations:
@test @ballocated(track!($beam, $qf, work=$work)) == 0
@testset "BeamTracking.jl" begin
# Linear tests -----------------------
beta_gamma_ref = 1.0
gamma_ref = sr_gamma(beta_gamma_ref)
species = Species("electron")
brho_ref = brho(massof(species), beta_gamma_ref, chargeof(species))
tol = 1e-14

# Defocusing:
beam = Beam(d, beta_gamma_ref=beta_gamma_ref)
qd = Linear.Quadrupole(B1=-B1,L=L_q)
track!(beam, qd, work=work)
# Linear drift test:
L_d = 5.0
drift = Linear.Drift(L=L_d)
M_drift_expected = [1.0 L_d 0.0 0.0 0.0 0.0;
0.0 1.0 0.0 0.0 0.0 0.0;
0.0 0.0 1.0 L_d 0.0 0.0;
0.0 0.0 0.0 1.0 0.0 0.0;
0.0 0.0 0.0 0.0 1.0 L_d/gamma_ref^2;
0.0 0.0 0.0 0.0 0.0 1.0]

test_matrix(drift, 0, M_drift_expected, beta_gamma_ref=beta_gamma_ref, species=species)

# Linear Quadrupole test:
# Focusing:
L_q = 1.2
K1n = 0.36
B1 = K1n*brho_ref
qf = Linear.Quadrupole(B1=B1,L=L_q)
M_qf_x = [cos(sqrt(K1n)*L_q) sincu(sqrt(K1n)*L_q)*L_q;
-sqrt(K1n)*sin(sqrt(K1n)*L_q) cos(sqrt(K1n)*L_q) ;]
M_qf_y = [cosh(sqrt(K1n)*L_q) sinhcu(sqrt(K1n)*L_q)*L_q;
sqrt(K1n)*sinh(sqrt(K1n)*L_q) cosh(sqrt(K1n)*L_q) ;]
M_qf_z = [1.0 L_q/gamma_ref^2;
0.0 1.0 ;]
M_qf_expected = zeros(6,6)
M_qf_expected[1:2,1:2] = M_qf_x
M_qf_expected[3:4,3:4] = M_qf_y
M_qf_expected[5:6,5:6] = M_qf_z

p = Particle(beam, 1)
test_matrix(qf, 1, M_qf_expected, beta_gamma_ref=beta_gamma_ref, species=species)

M_qd_expected = zeros(6,6)
M_qd_expected[1:2,1:2] = M_qf_y
M_qd_expected[3:4,3:4] = M_qf_x
M_qd_expected[5:6,5:6] = M_qf_z
# Defocusing:
qd = Linear.Quadrupole(B1=-B1,L=L_q)
M_qd_expected = zeros(6,6)
M_qd_expected[1:2,1:2] = M_qf_y
M_qd_expected[3:4,3:4] = M_qf_x
M_qd_expected[5:6,5:6] = M_qf_z

# 1) Correctness:
@test norm(GTPSA.jacobian(p.v) - M_qd_expected) < tol
# 2) Type-stability:
@test_opt track!(beam, qd, work=work)
# 3) Allocations:
@test @ballocated(track!($beam, $qd, work=$work)) == 0
test_matrix(qd, 1, M_qd_expected, beta_gamma_ref=beta_gamma_ref, species=species)
end

#include("linear.jl")

0 comments on commit 11646ed

Please sign in to comment.