add PtEtaPhiM coordinates used by LorentzVectorHEP.jl (#4)

* add PtEtaPhiM coordinates used by LorentzVectorHEP.jl

Author: Moelf

src/LorentzVectorBase.jl

@@ -4,5 +4,6 @@ export coordinate_system, coordinate_names
 
 include("interface.jl")
 include("coordinate_systems/exyz.jl")
+include("coordinate_systems/PtEtaPhiM.jl")
 
 end

src/coordinate_systems/PtEtaPhiM.jl

@@ -0,0 +1,131 @@
+"""
+
+    PtEtaPhiM <: AbstractCoordinateSystem  
+
+Cylindrical coordinate system for four-momenta. Using this requires the implementation of the following interface functions:
+
+```julia
+pt(::CustomFourMomentum)
+eta(::CustomFourMomentum)
+phi(::CustomFourMomentum)
+mass(::CustomFourMomentum)
+```
+
+"""
+struct PtEtaPhiM <: AbstractCoordinateSystem end
+coordinate_names(::PtEtaPhiM) = (:pt, :eta, :phi, :energy)
+
+
+####
+# derived components
+####
+
+px(::PtEtaPhiM, v) = pt(v) * cos(phi(v))
+py(::PtEtaPhiM, v) = pt(v) * sin(phi(v))
+pz(::PtEtaPhiM, v) = pt(v) * sinh(eta(v))
+
+function energy(::PtEtaPhiM, mom)
+    sqrt(px(mom)^2 + py(mom)^2 + pz(mom)^2 + invariant_mass2(mom))
+end
+
+@inline function spatial_magnitude2(::PtEtaPhiM, mom)
+    return pt(mom)^2 * (1+sinh(eta(mom))^2)
+end
+
+@inline function spatial_magnitude(::PtEtaPhiM, mom)
+    return sqrt(spatial_magnitude2(mom))
+end
+
+@inline function mass2(::PtEtaPhiM, mom)
+    mass(mom)^2
+end
+
+@inline function boost_beta(::PtEtaPhiM, mom)
+    en = energy(mom)
+    rho = spatial_magnitude(mom)
+    if !iszero(en)
+        rho / en
+    elseif iszero(rho)
+        return zero(rho)
+    else
+        throw(
+            ArgumentError(
+                "There is no beta for a four-momentum with vanishing time/energy component"
+            ),
+        )
+    end
+end
+
+@inline function boost_gamma(::PtEtaPhiM, mom)
+    return inv(sqrt(one(energy(mom)) - boost_beta(mom)^2))
+end
+
+########################
+# transverse coordinates
+########################
+
+@inline function pt2(::PtEtaPhiM, mom)
+    return pt(mom)^2
+end
+
+@inline function mt2(::PtEtaPhiM, mom)
+    return energy(mom)^2 - pz(mom)^2
+end
+
+function mt(::PtEtaPhiM, mom)
+    mT2 = mt2(mom)
+    if mT2 < zero(mT2)
+        # add optional waring: negative transverse mass -> -sqrt(-mT2) is returned.
+        -sqrt(-mT2)
+    else
+        sqrt(mT2)
+    end
+end
+
+@inline function rapidity(::PtEtaPhiM, mom)
+    E = energy(mom)
+    pz1 = pz(mom)
+    return 0.5*log((E+pz1)/(E-pz1))
+end
+
+#######################
+# spherical coordinates
+#######################
+
+@inline function polar_angle(::PtEtaPhiM, mom)
+    xcomp = px(mom)
+    ycomp = py(mom)
+    zcomp = pz(mom)
+
+    return if iszero(xcomp) && iszero(ycomp) && iszero(zcomp)
+        zero(xcomp)
+    else
+        atan(transverse_momentum(mom), zcomp)
+    end
+end
+
+@inline function cos_theta(::PtEtaPhiM, mom)
+    r = spatial_magnitude(mom)
+    return iszero(r) ? one(px(mom)) : pz(mom) / r
+end
+
+function cos_phi(::PtEtaPhiM, mom)
+    pT = pt(mom)
+    return iszero(pT) ? one(pT) : px(mom) / pT
+end
+
+function sin_phi(::PtEtaPhiM, mom)
+    pT = pt(mom)
+    return iszero(pT) ? zero(pT) : py(mom) / pT
+end
+
+########################
+# light cone coordinates
+########################
+@inline function plus_component(::PtEtaPhiM, mom)
+    return 0.5 * (energy(mom) + pz(mom))
+end
+
+@inline function minus_component(::PtEtaPhiM, mom)
+    return 0.5 * (energy(mom) - pz(mom))
+end

src/coordinate_systems/exyz.jl

@@ -13,7 +13,7 @@ energy(::CustomFourMomentum)
 
 """
 struct EXYZ <: AbstractCoordinateSystem end
-coordinate_names(::EXYZ) = (:energy, :px, :py, :pz)
+coordinate_names(::EXYZ) = (:px, :py, :pz, :energy)
 
 ####
 # Interface functions
@@ -93,7 +93,7 @@ end
 
 """
 
-    invariant_mass2(::EXYZ,mom)
+    mass2(::EXYZ,mom)
 
 Return the squared invariant mass of a given four-momentum, i.e. the minkowski dot with itself. 
 
@@ -103,13 +103,13 @@ Return the squared invariant mass of a given four-momentum, i.e. the minkowski d
 
 
 """
-@inline function invariant_mass2(::EXYZ, mom)
+@inline function mass2(::EXYZ, mom)
     return energy(mom)^2 - px(mom)^2 - py(mom)^2 - pz(mom)^2
 end
 
 """
 
-    invariant_mass(::EXYZ,mom)
+    mass(::EXYZ,mom)
 
 Return the the invariant mass of a given four-momentum, i.e. the square root of the minkowski dot with itself. 
 
@@ -123,8 +123,8 @@ Return the the invariant mass of a given four-momentum, i.e. the square root of
     If the squared invariant mass `m2` is negative, `-sqrt(-m2)` is returned. 
 
 """
-@inline function invariant_mass(cs::EXYZ, mom)
-    m2 = invariant_mass2(mom)
+@inline function mass(::EXYZ, mom)
+    m2 = mass2(mom)
     if m2 < zero(m2)
         # Think about including this waring, maybe optional with a global PRINT_WARINGS switch.
         #@warn("The square of the invariant mass (m2=P*P) is negative. The value -sqrt(-m2) is returned.")
@@ -181,7 +181,7 @@ end
 ########################
 """
 
-    transverse_momentum2(::EXYZ,mom)
+    pt2(::EXYZ,mom)
 
 Return the squared transverse momentum for a given four-momentum, i.e. the sum of its squared 1- and 2-component.
 
@@ -194,13 +194,13 @@ Return the squared transverse momentum for a given four-momentum, i.e. the sum o
     The transverse components are defined w.r.t. to the 3-axis. 
 
 """
-@inline function transverse_momentum2(::EXYZ, mom)
+@inline function pt2(::EXYZ, mom)
     return px(mom)^2 + py(mom)^2
 end
 
 """
 
-    transverse_momentum(::EXYZ, mom)
+    pt(::EXYZ, mom)
 
 Return the transverse momentum for a given four-momentum, i.e. the spatial_magnitude of its transverse components.
 
@@ -214,13 +214,13 @@ Return the transverse momentum for a given four-momentum, i.e. the spatial_magni
     The transverse components are defined w.r.t. to the 3-axis. 
 
 """
-@inline function transverse_momentum(::EXYZ, mom)
-    return sqrt(transverse_momentum2(mom))
+@inline function pt(::EXYZ, mom)
+    return sqrt(pt2(mom))
 end
 
 """
 
-    transverse_mass2(::EXYZ, mom)
+    mt2(::EXYZ, mom)
 
 Return the squared transverse mass for a given four-momentum, i.e. the difference of its squared 0- and 3-component.
 
@@ -235,13 +235,13 @@ Return the squared transverse mass for a given four-momentum, i.e. the differenc
 
 
 """
-@inline function transverse_mass2(::EXYZ, mom)
+@inline function mt2(::EXYZ, mom)
     return energy(mom)^2 - pz(mom)^2
 end
 
 """
 
-    transverse_mass(::EXYZ,mom)
+    mt(::EXYZ,mom)
 
 Return the transverse momentum for a given four-momentum, i.e. the square root of its squared transverse mass.
 
@@ -259,8 +259,8 @@ Return the transverse momentum for a given four-momentum, i.e. the square root o
     If the squared transverse mass `mt2` is negative, `-sqrt(-mt2)` is returned.
 
 """
-function transverse_mass(::EXYZ, mom)
-    mT2 = transverse_mass2(mom)
+function mt(::EXYZ, mom)
+    mT2 = mt2(mom)
     if mT2 < zero(mT2)
         # add optional waring: negative transverse mass -> -sqrt(-mT2) is returned.
         -sqrt(-mT2)
@@ -292,11 +292,11 @@ Return the [rapidity](https://en.wikipedia.org/wiki/Rapidity) for a given four-m
 end
 
 """
-    pseudorapidity(::EXYZ, mom)
+    eta(::EXYZ, mom)
 
 TBW
 """
-function pseudorapidity(::EXYZ, mom)
+function eta(::EXYZ, mom)
     cth = cos_theta(mom)
 
     if cth^2 < one(cth)
@@ -308,7 +308,7 @@ function pseudorapidity(::EXYZ, mom)
         return zero(z)
     end
 
-    @warn "Pseudorapidity: transverse momentum is zero! return +/- 10e10"
+    @warn "Pseudorapidity (η): transverse momentum is zero! return +/- 10e10"
     if z > zero(z)
         return 10e10
     end
@@ -364,7 +364,7 @@ end
 
 """
 
-    azimuthal_angle(::EXYZ, mom)
+    phi(::EXYZ, mom)
 
 Return the phi angle for a given four-momentum, i.e. the azimuthal angle of its spatial components in [spherical coordinates](https://en.wikipedia.org/wiki/Spherical_coordinate_system).
 
@@ -377,7 +377,7 @@ Return the phi angle for a given four-momentum, i.e. the azimuthal angle of its
     The [spherical coordinates](https://en.wikipedia.org/wiki/Spherical_coordinate_system) are defined w.r.t. to the 3-axis. 
 
 """
-function azimuthal_angle(::EXYZ, mom)
+function phi(::EXYZ, mom)
     xcomp = px(mom)
     ycomp = py(mom)
     return iszero(xcomp) && iszero(ycomp) ? zero(xcomp) : atan(ycomp, xcomp)

src/interface.jl

@@ -15,25 +15,25 @@ function coordinate_names end
 function coordinate_system end
 
 const FOURMOMENTUM_GETTER_FUNCTIONS = (
-    :energy,
     :px,
     :py,
     :pz,
-    :spatial_magnitude2,
+    :energy,
+    :pt,
+    :pt2,
+    :eta,
+    :phi,
     :spatial_magnitude,
-    :invariant_mass2,
-    :invariant_mass,
+    :spatial_magnitude2,
+    :mass,
+    :mass2,
     :boost_beta,
     :boost_gamma,
-    :transverse_momentum2,
-    :transverse_momentum,
-    :transverse_mass2,
-    :transverse_mass,
+    :mt2,
+    :mt,
     :rapidity,
-    :pseudorapidity,
     :polar_angle,
     :cos_theta,
-    :azimuthal_angle,
     :cos_phi,
     :sin_phi,
     :plus_component,
@@ -49,16 +49,16 @@ for func in FOURMOMENTUM_GETTER_FUNCTIONS
 end
 
 const FOURMOMENTUM_GETTER_ALIASSES = Dict(
-    :mag2 => :spatial_magnitude2,
-    :mag => :spatial_magnitude,
-    :mass2 => :invariant_mass2,
-    :mass => :invariant_mass,
-    :pt2 => :transverse_momentum2,
-    :perp2 => :transverse_momentum2,
-    :pt => :transverse_momentum,
-    :perp => :transverse_momentum,
-    :mt2 => :transverse_mass2,
-    :mt => :transverse_mass,
+    :invariant_mass => :mass,
+    :invariant_mass2 => :mass2,
+    :transverse_momentum => :pt,
+    :transverse_momentum2 => :pt2,
+    :perp => :pt,
+    :perp2 => :pt2,
+    :transverse_mass => :mt,
+    :transverse_mass2 => :mt2,
+    :azimuthal_angle => :phi,
+    :pseudorapidity => :eta,
 )
 
 for (alias, func) in FOURMOMENTUM_GETTER_ALIASSES