Skip to content

Commit 46d7df1

Browse files
szabo137Uwe Hernandez AcostaMoelfAntonReinhard
authored
Streamlining and cleanup (#24)
* fixed local docs building * cleanup * formatting * added PxPyPzE * added getter docs * removed old docstrings * cleanup * formatting * finished docstrings for getter functions * fixed tests for XYZT * fixed doc building * Update src/coordinate_systems/XYZT.jl Co-authored-by: Anton Reinhard <[email protected]> * Update src/coordinate_systems/XYZT.jl Co-authored-by: Anton Reinhard <[email protected]> * Apply suggestions from code review Co-authored-by: Anton Reinhard <[email protected]> * updated tests, fixed bug in XYZT eta * increase coverage in XYZT tests * formatting * fixed getter for cyl coords * Apply suggestions from code review Co-authored-by: Anton Reinhard <[email protected]> * fix the nits * fixed eta in XYZT --------- Co-authored-by: Uwe Hernandez Acosta <[email protected]> Co-authored-by: Jerry Ling <[email protected]> Co-authored-by: Anton Reinhard <[email protected]>
1 parent 4430fee commit 46d7df1

13 files changed

+959
-511
lines changed

Project.toml

-3
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,5 @@ uuid = "592a752d-6533-4762-a71b-738712ea30ec"
33
authors = ["Uwe Hernandez Acosta <[email protected]>"]
44
version = "0.0.1-DEV"
55

6-
[deps]
7-
86
[compat]
97
julia = "1.6"
10-

docs/src/95-reference.md

+49-2
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,7 @@
1+
```@meta
2+
CurrentModule=LorentzVectorBase
3+
```
4+
15
# [Reference](@id reference)
26

37
## Contents
@@ -12,6 +16,49 @@ Pages = ["95-reference.md"]
1216
Pages = ["95-reference.md"]
1317
```
1418

15-
```@autodocs
16-
Modules = [LorentzVectorBase]
19+
## Interface
20+
21+
```@docs
22+
coordinate_system
23+
coordinate_names
24+
```
25+
26+
## Supported coordinate systems
27+
28+
```@docs
29+
XYZT
30+
PxPyPzE
31+
PtEtaPhiM
32+
```
33+
34+
## Supported getter functions
35+
36+
```@docs
37+
x
38+
y
39+
z
40+
t
41+
px
42+
py
43+
pz
44+
E
45+
pt
46+
pt2
47+
eta
48+
phi
49+
spatial_magnitude
50+
spatial_magnitude2
51+
mass
52+
mass2
53+
boost_beta
54+
boost_gamma
55+
mt2
56+
mt
57+
rapidity
58+
polar_angle
59+
cos_theta
60+
cos_phi
61+
sin_phi
62+
plus_component
63+
minus_component
1764
```

docs/src/index.md

-1
Original file line numberDiff line numberDiff line change
@@ -5,4 +5,3 @@ CurrentModule = LorentzVectorBase
55
# LorentzVectorBase
66

77
Documentation for [LorentzVectorBase](https://github.com/JuliaHEP/LorentzVectorBase.jl).
8-

src/LorentzVectorBase.jl

+3-1
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,10 @@ module LorentzVectorBase
22

33
export coordinate_system, coordinate_names
44

5+
include("getter.jl")
56
include("interface.jl")
6-
include("coordinate_systems/xyze.jl")
7+
include("coordinate_systems/XYZT.jl")
8+
include("coordinate_systems/PxPyPzE.jl")
79
include("coordinate_systems/PtEtaPhiM.jl")
810

911
end

src/coordinate_systems/PtEtaPhiM.jl

+1-1
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ px(::PtEtaPhiM, v) = pt(v) * cos(phi(v))
2323
py(::PtEtaPhiM, v) = pt(v) * sin(phi(v))
2424
pz(::PtEtaPhiM, v) = pt(v) * sinh(eta(v))
2525

26-
function energy(::PtEtaPhiM, mom)
26+
function E(::PtEtaPhiM, mom)
2727
return sqrt(px(mom)^2 + py(mom)^2 + pz(mom)^2 + invariant_mass2(mom))
2828
end
2929

src/coordinate_systems/PxPyPzE.jl

+67
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
2+
"""
3+
4+
PxPyPzE <: AbstractCoordinateSystem
5+
6+
Cartesian coordinate system for four-momenta. Using this requires the implementation of the following interface functions:
7+
8+
```julia
9+
LorentzVectorBase.px(::CustomFourMomentum)
10+
LorentzVectorBase.py(::CustomFourMomentum)
11+
LorentzVectorBase.pz(::CustomFourMomentum)
12+
LorentzVectorBase.E(::CustomFourMomentum)
13+
```
14+
15+
"""
16+
struct PxPyPzE <: AbstractCoordinateSystem end
17+
coordinate_names(::PxPyPzE) = (:px, :py, :pz, :E)
18+
19+
x(cs, p) = px(p)
20+
x(::XYZT, p) = x(p)
21+
px(::PxPyPzE, p) = px(p)
22+
px(::XYZT, p) = x(p)
23+
24+
y(cs, p) = py(p)
25+
y(::XYZT, p) = y(p)
26+
py(::PxPyPzE, p) = py(p)
27+
py(::XYZT, p) = y(p)
28+
29+
z(cs, p) = pz(p)
30+
z(::XYZT, p) = z(p)
31+
pz(::PxPyPzE, p) = pz(p)
32+
pz(::XYZT, p) = z(p)
33+
34+
t(cs, p) = E(p)
35+
t(::XYZT, p) = t(p)
36+
E(::PxPyPzE, p) = E(p)
37+
E(::XYZT, p) = t(p)
38+
39+
const DELEGATED_GETTER_FUNCTIONS = (
40+
:pt,
41+
:pt2,
42+
:eta,
43+
:phi,
44+
:spatial_magnitude,
45+
:spatial_magnitude2,
46+
:mass,
47+
:mass2,
48+
:boost_beta,
49+
:boost_gamma,
50+
:mt2,
51+
:mt,
52+
:rapidity,
53+
:polar_angle,
54+
:cos_theta,
55+
:cos_phi,
56+
:sin_phi,
57+
:plus_component,
58+
:minus_component,
59+
)
60+
61+
for func in DELEGATED_GETTER_FUNCTIONS
62+
eval(
63+
quote
64+
($func)(::PxPyPzE, mom) = ($func)(XYZT(), mom)
65+
end,
66+
)
67+
end

src/coordinate_systems/XYZT.jl

+174
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
1+
"""
2+
3+
XYZT <: AbstractCoordinateSystem
4+
5+
Cartesian coordinate system for four-vectors. Using this requires the implementation of the following interface functions:
6+
7+
```julia
8+
LorentzVectorBase.x(::CustomFourVector)
9+
LorentzVectorBase.y(::CustomFourVector)
10+
LorentzVectorBase.z(::CustomFourVector)
11+
LorentzVectorBase.t(::CustomFourVector)
12+
```
13+
14+
"""
15+
struct XYZT <: AbstractCoordinateSystem end
16+
coordinate_names(::XYZT) = (:x, :y, :z, :t)
17+
18+
####
19+
# Interface functions
20+
####
21+
22+
function t end
23+
24+
function x end
25+
26+
function y end
27+
28+
function z end
29+
30+
####
31+
# derived components
32+
####
33+
34+
@inline function spatial_magnitude2(::XYZT, mom)
35+
return x(mom)^2 + y(mom)^2 + z(mom)^2
36+
end
37+
38+
@inline function spatial_magnitude(cs::XYZT, mom)
39+
return sqrt(spatial_magnitude2(mom))
40+
end
41+
42+
@inline function mass2(::XYZT, mom)
43+
return t(mom)^2 - x(mom)^2 - y(mom)^2 - z(mom)^2
44+
end
45+
46+
@inline function mass(::XYZT, mom)
47+
m2 = mass2(mom)
48+
if m2 < zero(m2)
49+
# Think about including this warning, maybe optional with a global PRINT_WARNINGS switch.
50+
#@warn("The square of the invariant mass (m2=P*P) is negative. The value -sqrt(-m2) is returned.")
51+
return -sqrt(-m2)
52+
else
53+
return sqrt(m2)
54+
end
55+
end
56+
57+
@inline function boost_beta(::XYZT, mom)
58+
en = t(mom)
59+
rho = spatial_magnitude(mom)
60+
if !iszero(en)
61+
return rho / en
62+
elseif iszero(rho)
63+
return zero(rho)
64+
else
65+
throw(
66+
ArgumentError(
67+
"there is no beta for a four-vector with vanishing time/energy component"
68+
),
69+
)
70+
end
71+
end
72+
73+
@inline function boost_gamma(::XYZT, mom)
74+
beta = boost_beta(mom)
75+
return inv(sqrt(one(beta) - beta^2))
76+
end
77+
78+
########################
79+
# transverse coordinates
80+
########################
81+
@inline function pt2(::XYZT, mom)
82+
return x(mom)^2 + y(mom)^2
83+
end
84+
85+
@inline function pt(::XYZT, mom)
86+
return sqrt(pt2(mom))
87+
end
88+
89+
@inline function mt2(::XYZT, mom)
90+
return t(mom)^2 - z(mom)^2
91+
end
92+
93+
function mt(::XYZT, mom)
94+
mT2 = mt2(mom)
95+
if mT2 < zero(mT2)
96+
# add optional waring: negative transverse mass -> -sqrt(-mT2) is returned.
97+
-sqrt(-mT2)
98+
else
99+
sqrt(mT2)
100+
end
101+
end
102+
103+
@inline function rapidity(::XYZT, mom)
104+
en = t(mom)
105+
zcomp = z(mom)
106+
return 0.5 * log((en + zcomp) / (en - zcomp))
107+
end
108+
109+
function eta(::XYZT, mom)
110+
cth = cos_theta(mom)
111+
112+
if cth^2 < one(cth)
113+
return -0.5 * log((1 - cth) / (1 + cth))
114+
end
115+
116+
zcomp = z(mom)
117+
if iszero(zcomp)
118+
return zero(zcomp)
119+
end
120+
121+
@warn "Pseudorapidity (η): transverse momentum is zero! return +/- 10e10"
122+
if zcomp > zero(zcomp)
123+
return 10e10
124+
end
125+
126+
return -10e10
127+
end
128+
129+
#######################
130+
# spherical coordinates
131+
#######################
132+
@inline function polar_angle(::XYZT, mom)
133+
xcomp = x(mom)
134+
ycomp = y(mom)
135+
zcomp = z(mom)
136+
137+
return if iszero(xcomp) && iszero(ycomp) && iszero(zcomp)
138+
zero(xcomp)
139+
else
140+
atan(transverse_momentum(mom), zcomp)
141+
end
142+
end
143+
144+
@inline function cos_theta(::XYZT, mom)
145+
r = spatial_magnitude(mom)
146+
return iszero(r) ? one(x(mom)) : z(mom) / r
147+
end
148+
149+
function phi(::XYZT, mom)
150+
xcomp = x(mom)
151+
ycomp = y(mom)
152+
return iszero(xcomp) && iszero(ycomp) ? zero(xcomp) : atan(ycomp, xcomp)
153+
end
154+
155+
function cos_phi(::XYZT, mom)
156+
pT = transverse_momentum(mom)
157+
return iszero(pT) ? one(pT) : x(mom) / pT
158+
end
159+
160+
function sin_phi(::XYZT, mom)
161+
pT = transverse_momentum(mom)
162+
return iszero(pT) ? zero(pT) : y(mom) / pT
163+
end
164+
165+
########################
166+
# light cone coordinates
167+
########################
168+
@inline function plus_component(::XYZT, mom)
169+
return 0.5 * (t(mom) + z(mom))
170+
end
171+
172+
@inline function minus_component(::XYZT, mom)
173+
return 0.5 * (t(mom) - z(mom))
174+
end

0 commit comments

Comments
 (0)