Skip to content

Commit 4430fee

Browse files
mfarrington1mfarringUwe Hernandez Acostaszabo137
authored
Added tests for cylindrical coordinates. (#16)
* Fixed typo in coordinate_names(::PtEtaPhiM) * Modified polar_angle(::PtEtaPhiM, mom) to make it more concise and efficient. * Corrected the definition of polar_angle(::PtEtaPhiM, mom) to match that used in ROOT. * Added tests for cylindrical coordinates. * formatting * Update test/runtests.jl Co-authored-by: Uwe Hernandez Acosta <[email protected]> --------- Co-authored-by: mfarring <[email protected]> Co-authored-by: Uwe Hernandez Acosta <[email protected]> Co-authored-by: Uwe Hernandez Acosta <[email protected]>
1 parent ba9d0eb commit 4430fee

File tree

3 files changed

+142
-11
lines changed

3 files changed

+142
-11
lines changed

‎src/coordinate_systems/PtEtaPhiM.jl

+5-11
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
"""
22
3-
PtEtaPhiM <: AbstractCoordinateSystem
3+
PtEtaPhiM <: AbstractCoordinateSystem
44
55
Cylindrical coordinate system for four-momenta. Using this requires the implementation of the following interface functions:
66
@@ -13,7 +13,7 @@ mass(::CustomFourMomentum)
1313
1414
"""
1515
struct PtEtaPhiM <: AbstractCoordinateSystem end
16-
coordinate_names(::PtEtaPhiM) = (:pt, :eta, :phi, :energy)
16+
coordinate_names(::PtEtaPhiM) = (:pt, :eta, :phi, :mass)
1717

1818
####
1919
# derived components
@@ -92,15 +92,9 @@ end
9292
#######################
9393

9494
@inline function polar_angle(::PtEtaPhiM, mom)
95-
xcomp = px(mom)
96-
ycomp = py(mom)
97-
zcomp = pz(mom)
98-
99-
return if iszero(xcomp) && iszero(ycomp) && iszero(zcomp)
100-
zero(xcomp)
101-
else
102-
atan(transverse_momentum(mom), zcomp)
103-
end
95+
η = eta(mom)
96+
r = spatial_magnitude(mom)
97+
return iszero(r) && iszero(pz(mom)) ? zero(pz(mom)) : 2 * atan(exp(-η))
10498
end
10599

106100
@inline function cos_theta(::PtEtaPhiM, mom)

‎test/PtEtaPhiM.jl

+133
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
using Random
2+
using LorentzVectorBase
3+
4+
const RNG = MersenneTwister(137137)
5+
6+
struct CustomCylMom
7+
pt
8+
eta
9+
phi
10+
m
11+
end
12+
13+
LorentzVectorBase.coordinate_system(::CustomCylMom) = LorentzVectorBase.PtEtaPhiM()
14+
LorentzVectorBase.pt(mom::CustomCylMom) = mom.pt
15+
LorentzVectorBase.eta(mom::CustomCylMom) = mom.eta
16+
LorentzVectorBase.phi(mom::CustomCylMom) = mom.phi
17+
LorentzVectorBase.mass(mom::CustomCylMom) = mom.m
18+
19+
x, y, z = rand(RNG, 3)
20+
m = abs(rand(RNG))
21+
pt = sqrt(x^2 + y^2)
22+
η = atanh(z / sqrt(x^2 + y^2 + z^2))
23+
ϕ = (x == 0.0 && y == 0.0) ? 0.0 : atan(y, x)
24+
25+
mom_onshell = CustomCylMom(pt, η, ϕ, m)
26+
mom_zero = CustomCylMom(0.0, 0.0, 0.0, 0.0)
27+
28+
@testset "spatial_magnitude consistency" begin
29+
for mom in [mom_onshell, mom_zero]
30+
@test isapprox(
31+
LorentzVectorBase.spatial_magnitude(mom),
32+
sqrt(LorentzVectorBase.spatial_magnitude2(mom)),
33+
)
34+
end
35+
end
36+
37+
@testset "spatial_magnitude values" begin
38+
@test isapprox(LorentzVectorBase.spatial_magnitude2(mom_onshell), x^2 + y^2 + z^2)
39+
@test isapprox(LorentzVectorBase.spatial_magnitude(mom_onshell), sqrt(x^2 + y^2 + z^2))
40+
end
41+
42+
@testset "mass2 consistency" begin
43+
for mom in [mom_onshell, mom_zero]
44+
@test isapprox(LorentzVectorBase.mass(mom), sqrt(LorentzVectorBase.mass2(mom)))
45+
@test LorentzVectorBase.invariant_mass2(mom) == LorentzVectorBase.mass2(mom)
46+
@test LorentzVectorBase.invariant_mass(mom) == LorentzVectorBase.mass(mom)
47+
@test isapprox(
48+
LorentzVectorBase.invariant_mass(mom), sqrt(LorentzVectorBase.invariant_mass2(mom))
49+
)
50+
end
51+
end
52+
53+
@testset "mass value" begin
54+
@test LorentzVectorBase.mass2(mom_onshell) == m^2
55+
@test LorentzVectorBase.mass(mom_onshell) == m
56+
@test LorentzVectorBase.mass(mom_zero) == 0.0
57+
end
58+
59+
@testset "energy" begin
60+
for mom in [mom_onshell, mom_zero]
61+
@test isapprox(
62+
LorentzVectorBase.energy(mom)^2,
63+
LorentzVectorBase.spatial_magnitude2(mom) + LorentzVectorBase.invariant_mass2(mom),
64+
)
65+
end
66+
end
67+
68+
@testset "momentum components" begin
69+
@test LorentzVectorBase.pt(mom_onshell) == pt
70+
@test LorentzVectorBase.eta(mom_onshell) == η
71+
@test LorentzVectorBase.phi(mom_onshell) == ϕ
72+
@test LorentzVectorBase.mass(mom_onshell) == m
73+
@test isapprox(
74+
LorentzVectorBase.boost_beta(mom_onshell),
75+
LorentzVectorBase.spatial_magnitude(mom_onshell) /
76+
LorentzVectorBase.energy(mom_onshell),
77+
)
78+
79+
@test LorentzVectorBase.pt(mom_zero) == 0.0
80+
@test LorentzVectorBase.eta(mom_zero) == 0.0
81+
@test LorentzVectorBase.phi(mom_zero) == 0.0
82+
@test LorentzVectorBase.mass(mom_zero) == 0.0
83+
@test LorentzVectorBase.boost_beta(mom_zero) == 0.0
84+
end
85+
86+
@testset "transverse coordinates values" begin
87+
@test isapprox(LorentzVectorBase.pt2(mom_onshell), pt^2)
88+
@test isapprox(
89+
LorentzVectorBase.mt2(mom_onshell), LorentzVectorBase.energy(mom_onshell)^2 - z^2
90+
)
91+
@test isapprox(
92+
LorentzVectorBase.mt(mom_onshell), sqrt(LorentzVectorBase.mt2(mom_onshell))
93+
)
94+
@test isapprox(
95+
LorentzVectorBase.rapidity(mom_onshell),
96+
0.5 * log(
97+
(LorentzVectorBase.energy(mom_onshell) + z) /
98+
(LorentzVectorBase.energy(mom_onshell) - z),
99+
),
100+
)
101+
102+
@test isapprox(LorentzVectorBase.pt2(mom_zero), 0.0)
103+
@test isapprox(LorentzVectorBase.mt2(mom_zero), 0.0)
104+
@test isapprox(LorentzVectorBase.mt(mom_zero), 0.0)
105+
end
106+
107+
@testset "cartesian coordinate values" begin
108+
@test isapprox(LorentzVectorBase.px(mom_onshell), x)
109+
@test isapprox(LorentzVectorBase.py(mom_onshell), y)
110+
@test isapprox(LorentzVectorBase.pz(mom_onshell), z)
111+
112+
@test isapprox(LorentzVectorBase.px(mom_zero), 0.0)
113+
@test isapprox(LorentzVectorBase.py(mom_zero), 0.0)
114+
@test isapprox(LorentzVectorBase.pz(mom_zero), 0.0)
115+
end
116+
117+
@testset "spherical coordinate values" begin
118+
if pt != 0 && z != 0
119+
@test isapprox(LorentzVectorBase.polar_angle(mom_onshell), atan(pt, z))
120+
else
121+
@test isapprox(LorentzVectorBase.polar_angle(mom_onshell), 0)
122+
end
123+
124+
r = LorentzVectorBase.spatial_magnitude(mom_onshell)
125+
@test isapprox(LorentzVectorBase.cos_theta(mom_onshell), iszero(r) ? 1.0 : z / r)
126+
@test isapprox(LorentzVectorBase.cos_phi(mom_onshell), cos(ϕ))
127+
@test isapprox(LorentzVectorBase.sin_phi(mom_onshell), sin(ϕ))
128+
129+
@test isapprox(LorentzVectorBase.polar_angle(mom_zero), 0.0)
130+
@test isapprox(LorentzVectorBase.cos_theta(mom_zero), 1.0)
131+
@test isapprox(LorentzVectorBase.cos_phi(mom_zero), 1.0)
132+
@test isapprox(LorentzVectorBase.sin_phi(mom_zero), 0.0)
133+
end

‎test/runtests.jl

+4
Original file line numberDiff line numberDiff line change
@@ -10,4 +10,8 @@ begin
1010
@time @safetestset "xyze" begin
1111
include("xyze.jl")
1212
end
13+
14+
@time @safetestset "PtEtaPhiM" begin
15+
include("PtEtaPhiM.jl")
16+
end
1317
end

0 commit comments

Comments
 (0)