Add rapidity methods

Add methods that return the true rapidity of LorentzVectors (E, p)
these are defined for both cartesian and cylindrical cases as rap()

Add methods for the square of mass and the square of momentum
as these can be useful for more specialist calculations (but not
exported)

.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, rap, 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, rap, eta, phi, mass
 export deltaphi, deltar, deltaeta
 export ΔR, Δϕ, Δη
 export fromPtEtaPhiE

src/cartesian.jl

@@ -1,6 +1,8 @@
 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(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 +17,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 rap(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 rap(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,7 @@ 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 rap(v1) ≈ -2.3032199982371715 atol=1e-6
 
     @test isapprox((v1+v2).mass, 8.25741602000877, atol=1e-6)
     @test isapprox(fast_mass(v1,v2), 8.25741602000877, atol=1e-6)