Add rapidity methods

.gitignore

@@ -1 +1,2 @@
 /Manifest.toml
+.vscode/*

README.md

@@ -1,23 +1,32 @@
 # LorentzVectorHEP
 
 Provides two types (and the conversion between the two):
+
 - `LorentzVector` (energy, px, py, pz) (from [LorentzVectors.jl](https://github.com/JLTastet/LorentzVectors.jl))
 - `LorentzVectorCyl` (pt, eta, phi, mass)
 
 you can also use `fromPtEtaPhiE(pt, eta, phi, energy) --> LorentzVectorCyl`.
 
 and these functions for both of them:
+
 ```julia
-px, py, pz, energy, pt, eta, phi, mass
+px, py, pz, energy, pt, rapidity, eta, phi, mass
 ```
 
 as well as these utility functions:
+
 ```julia
 deltar, deltaphi, deltaeta, mt, mt2
 ```
-(some of them have alias, `ΔR, Δϕ, Δη`)
 
+(some of them have aliases, `ΔR, Δϕ, Δη`)
+
+There are some unexported methods which are useful for more specialist use cases:
+
+```julia
+mass2, pt2, mt, mt2, mag
+```
 
 ## LHC coordinate system
 
-![](https://cds.cern.ch/record/1699952/files/Figures_T_Coordinate.png)
+![LHC Coordinate System](https://cds.cern.ch/record/1699952/files/Figures_T_Coordinate.png)

src/LorentzVectorHEP.jl

@@ -4,7 +4,7 @@ using LorentzVectors # provides x, y, z, t
 
 export LorentzVectorCyl, LorentzVector
 
-export px, py, pz, energy, fast_mass, pt, eta, phi, mass
+export px, py, pz, energy, fast_mass, pt, rapidity, eta, phi, mass
 export deltaphi, deltar, deltaeta
 export ΔR, Δϕ, Δη
 export fromPtEtaPhiE

src/cartesian.jl

@@ -1,6 +1,9 @@
 Base.zero(lv::T) where T<:LorentzVector = T(0,0,0,0)
-mass(lv::LorentzVector) = sqrt(dot(lv, lv))
-pt(lv::LorentzVector) = sqrt(muladd(lv.x, lv.x, lv.y^2))
+mass2(lv::LorentzVector) = dot(lv, lv)
+"""mass value - returns a negative number for spacelike 4-vectors"""
+mass(lv::LorentzVector) = mass2(lv) < 0.0 ? -sqrt(-mass2(lv)) : sqrt(mass2(lv))
+pt2(lv::LorentzVector) = muladd(lv.x, lv.x, lv.y^2)
+pt(lv::LorentzVector) = sqrt(pt2(lv))
 mt2(lv::LorentzVector) = lv.t^2 - lv.z^2
 mt(lv::LorentzVector) = mt2(lv)<0 ? -sqrt(-mt2(lv)) : sqrt(mt2(lv))
 mag(lv::LorentzVector) = sqrt(muladd(lv.x, lv.x, lv.y^2) + lv.z^2)
@@ -15,19 +18,39 @@ pz(lv::LorentzVector) = lv.z
     return ifelse(ptot == 0.0, 1.0, fZ / ptot)
 end
 
+"""Pseudorapidity"""
 function eta(lv::LorentzVector)
     cosTheta = CosTheta(lv)
     (cosTheta^2 < 1.0) && return -0.5 * log((1.0 - cosTheta) / (1.0 + cosTheta))
     fZ = lv.z
     iszero(fZ) && return 0.0
-    # Warning("PseudoRapidity","transvers momentum = 0! return +/- 10e10");
+    # Warning("PseudoRapidity","transverse momentum = 0! return +/- 10e10");
     fZ > 0.0 && return 10e10
     return -10e10
 end
+const η = eta
+
+"""Rapidity"""
+function rapidity(lv::LorentzVector)
+    pt_squared = pt2(lv)
+    abspz = abs(pz(lv))
+    if (energy(lv) == abspz) && (pt_squared == 0.0)
+        return (-1)^(pz(lv) < 0)*(1e5 + abspz) # a very large value that depends on pz
+    end
+    m2 = max((energy(lv) + pz(lv))*(energy(lv) - pz(lv)) - pt_squared, 0.0) # mass^2
+    E_plus_z = energy(lv) + abspz
+    return (-1)^(pz(lv) > 0) * 0.5*log((pt_squared + m2)/(E_plus_z^2))
+end
+# Don't export "y" as an alias as for a normal 4-vector it's the second space coordinate
 
 function phi(lv::LorentzVector)
     return (lv.x == 0.0 && lv.y == 0.0) ? 0.0 : atan(lv.y, lv.x)
 end
+const ϕ = phi
+
+function phi02pi(lv::LorentzVector)
+    return phi(lv) < 0.0 ? phi(lv) + 2π : phi(lv)
+end
 
 function phi_mpi_pi(x)
     twopi = 2pi

src/cylindrical.jl

@@ -20,9 +20,11 @@ Base.zero(::Type{LorentzVectorCyl}) = zero(LorentzVectorCyl{Float64})
 Base.zero(lv::T) where T<:LorentzVectorCyl = T(0,0,0,0)
 
 pt(lv::LorentzVectorCyl) = lv.pt
+pt2(lv::LorentzVectorCyl) = lv.pt^2
 eta(lv::LorentzVectorCyl) = lv.eta
 phi(lv::LorentzVectorCyl) = lv.phi
 mass(lv::LorentzVectorCyl) = lv.mass
+mass2(lv::LorentzVectorCyl) = lv.mass^2
 px(v::LorentzVectorCyl) = v.pt * cos(v.phi)
 py(v::LorentzVectorCyl) = v.pt * sin(v.phi)
 pz(v::LorentzVectorCyl) = v.pt * sinh(v.eta)
@@ -74,6 +76,12 @@ function fast_mass(v1::LorentzVectorCyl, v2::LorentzVectorCyl)
         - tpt12*cos(phi1-phi2), zero(pt1)))
 end
 
+"Rapidity"
+function rapidity(lv::LorentzVectorCyl)
+    num = sqrt(lv.mass^2 + lv.pt^2 * cosh(lv.eta)^2) + lv.pt * sinh(lv.eta)
+    den = sqrt(lv.mass^2 + lv.pt^2)
+    return log(num/den)
+end
 
 # https://root.cern.ch/doc/v606/GenVector_2VectorUtil_8h_source.html#l00061
 @inline function deltaphi(v1::LorentzVectorCyl, v2::LorentzVectorCyl)

test/runtests.jl

@@ -8,6 +8,10 @@ using Test
     @test px(v1) ≈ -1424.610065192358 atol=1e-6
     @test py(v1) ≈ -1036.2899616674022 atol=1e-6
     @test pz(v1) ≈ -8725.817601790963 atol=1e-6
+    @test rapidity(v1) ≈ -2.3032199982371715 atol=1e-6
+
+    @test LorentzVectorHEP.pt2(v1) ≈ 3.1034107225e6 atol=1e-6
+    @test LorentzVectorHEP.mass2(v1) ≈ 0.011162345103999998 atol=1e-6
 
     @test isapprox((v1+v2).mass, 8.25741602000877, atol=1e-6)
     @test isapprox(fast_mass(v1,v2), 8.25741602000877, atol=1e-6)
@@ -24,6 +28,26 @@ using Test
     @test v3.mass == 5*v1.mass
     @test v3.eta == v1.eta
     @test v3.phi == v1.phi
+
+    vcart1 = LorentzVector(10.0, -2.3, 4.5, 0.23)
+    @test rapidity(vcart1) ≈ 0.02300405695442185 atol=1e-9
+    @test eta(vcart1) ≈ 0.045495409709778126 atol=1e-9
+    @test phi(vcart1) ≈ 2.0432932623119604 atol=1e-9
+
+    vcart2 = LorentzVector(10.0, 2.7, -4.1, -0.21)
+    @test rapidity(vcart2) ≈ -0.021003087817077763 atol=1e-9
+    @test eta(vcart2) ≈ -0.04276400891568771 atol=1e-9
+    @test phi(vcart2) ≈ -0.9884433806509134 atol=1e-9
+    @test LorentzVectorHEP.phi02pi(vcart2) ≈ 5.294741926528673 atol=1e-9
+
+    @test deltaeta(vcart1, vcart2) ≈ 0.08825941862546584 atol=1e-9
+    @test deltaphi(vcart1, vcart2) ≈ 3.0317366429628736 atol=1e-9
+
+    vcart3 = LorentzVector(66.0, 0.0, 0.0, 66.0)
+    @test rapidity(vcart3) ≈ 100066.0 atol=1e-9
+
+    vcart4 = LorentzVector(4.4, 8.1, 2.2, 3.3)
+    @test mass(vcart4) ≈ -7.872737770305829 atol=1e-9
 end
 
 @testset "summing" begin