Streamlining and cleanup

Project.toml

@@ -3,8 +3,5 @@ uuid = "592a752d-6533-4762-a71b-738712ea30ec"
 authors = ["Uwe Hernandez Acosta <[email protected]>"]
 version = "0.0.1-DEV"
 
-[deps]
-
 [compat]
 julia = "1.6"
-

docs/src/95-reference.md

@@ -1,3 +1,7 @@
+```@meta
+CurrentModule=LorentzVectorBase
+```
+
 # [Reference](@id reference)
 
 ## Contents
@@ -12,6 +16,49 @@ Pages = ["95-reference.md"]
 Pages = ["95-reference.md"]
 ```
 
-```@autodocs
-Modules = [LorentzVectorBase]
+## Interface
+
+```@docs
+ coordinate_system
+coordinate_names
+```
+
+## Supported coordinate systems
+
+```@docs
+XYZT
+PxPyPzE
+PtEtaPhiM
+```
+
+## Supported getter functions
+
+```@docs
+  x
+  y
+  z
+  t
+  px
+  py
+  pz
+  E
+  pt
+  pt2
+  eta
+  phi
+  spatial_magnitude
+  spatial_magnitude2
+  mass
+  mass2
+  boost_beta
+  boost_gamma
+  mt2
+  mt
+  rapidity
+  polar_angle
+  cos_theta
+  cos_phi
+  sin_phi
+  plus_component
+  minus_component
 ```

docs/src/index.md

@@ -5,4 +5,3 @@ CurrentModule = LorentzVectorBase
 # LorentzVectorBase
 
 Documentation for [LorentzVectorBase](https://github.com/JuliaHEP/LorentzVectorBase.jl).
-

src/LorentzVectorBase.jl

@@ -2,8 +2,10 @@ module LorentzVectorBase
 
 export coordinate_system, coordinate_names
 
+include("getter.jl")
 include("interface.jl")
-include("coordinate_systems/xyze.jl")
+include("coordinate_systems/XYZT.jl")
+include("coordinate_systems/PxPyPzE.jl")
 include("coordinate_systems/PtEtaPhiM.jl")
 
 end

src/coordinate_systems/PtEtaPhiM.jl

@@ -23,7 +23,7 @@ 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)
+function E(::PtEtaPhiM, mom)
   return sqrt(px(mom)^2 + py(mom)^2 + pz(mom)^2 + invariant_mass2(mom))
 end
 

src/coordinate_systems/PxPyPzE.jl

@@ -0,0 +1,67 @@
+
+"""
+
+    PxPyPzE <: AbstractCoordinateSystem
+
+Cartesian coordinate system for four-momenta. Using this requires the implementation of the following interface functions:
+
+```julia
+LorentzVectorBase.px(::CustomFourMomentum)
+LorentzVectorBase.py(::CustomFourMomentum)
+LorentzVectorBase.pz(::CustomFourMomentum)
+LorentzVectorBase.E(::CustomFourMomentum)
+```
+
+"""
+struct PxPyPzE <: AbstractCoordinateSystem end
+coordinate_names(::PxPyPzE) = (:px, :py, :pz, :E)
+
+x(cs, p) = px(p)
+x(::XYZT, p) = x(p)
+px(::PxPyPzE, p) = px(p)
+px(::XYZT, p) = x(p)
+
+y(cs, p) = py(p)
+y(::XYZT, p) = y(p)
+py(::PxPyPzE, p) = py(p)
+py(::XYZT, p) = y(p)
+
+z(cs, p) = pz(p)
+z(::XYZT, p) = z(p)
+pz(::PxPyPzE, p) = pz(p)
+pz(::XYZT, p) = z(p)
+
+t(cs, p) = E(p)
+t(::XYZT, p) = t(p)
+E(::PxPyPzE, p) = E(p)
+E(::XYZT, p) = t(p)
+
+const DELEGATED_GETTER_FUNCTIONS = (
+  :pt,
+  :pt2,
+  :eta,
+  :phi,
+  :spatial_magnitude,
+  :spatial_magnitude2,
+  :mass,
+  :mass2,
+  :boost_beta,
+  :boost_gamma,
+  :mt2,
+  :mt,
+  :rapidity,
+  :polar_angle,
+  :cos_theta,
+  :cos_phi,
+  :sin_phi,
+  :plus_component,
+  :minus_component,
+)
+
+for func in DELEGATED_GETTER_FUNCTIONS
+  eval(
+    quote
+      ($func)(::PxPyPzE, mom) = ($func)(XYZT(), mom)
+    end,
+  )
+end

src/coordinate_systems/XYZT.jl

@@ -0,0 +1,174 @@
+"""
+
+    XYZT <: AbstractCoordinateSystem
+
+Cartesian coordinate system for four-vectors. Using this requires the implementation of the following interface functions:
+
+```julia
+LorentzVectorBase.x(::CustomFourVector)
+LorentzVectorBase.y(::CustomFourVector)
+LorentzVectorBase.z(::CustomFourVector)
+LorentzVectorBase.t(::CustomFourVector)
+```
+
+"""
+struct XYZT <: AbstractCoordinateSystem end
+coordinate_names(::XYZT) = (:x, :y, :z, :t)
+
+####
+# Interface functions
+####
+
+function t end
+
+function x end
+
+function y end
+
+function z end
+
+####
+# derived components
+####
+
+@inline function spatial_magnitude2(::XYZT, mom)
+  return x(mom)^2 + y(mom)^2 + z(mom)^2
+end
+
+@inline function spatial_magnitude(cs::XYZT, mom)
+  return sqrt(spatial_magnitude2(mom))
+end
+
+@inline function mass2(::XYZT, mom)
+  return t(mom)^2 - x(mom)^2 - y(mom)^2 - z(mom)^2
+end
+
+@inline function mass(::XYZT, mom)
+  m2 = mass2(mom)
+  if m2 < zero(m2)
+    # Think about including this warning, maybe optional with a global PRINT_WARNINGS switch.
+    #@warn("The square of the invariant mass (m2=P*P) is negative. The value -sqrt(-m2) is returned.")
+    return -sqrt(-m2)
+  else
+    return sqrt(m2)
+  end
+end
+
+@inline function boost_beta(::XYZT, mom)
+  en = t(mom)
+  rho = spatial_magnitude(mom)
+  if !iszero(en)
+    return rho / en
+  elseif iszero(rho)
+    return zero(rho)
+  else
+    throw(
+      ArgumentError(
+        "there is no beta for a four-vector with vanishing time/energy component"
+      ),
+    )
+  end
+end
+
+@inline function boost_gamma(::XYZT, mom)
+  beta = boost_beta(mom)
+  return inv(sqrt(one(beta) - beta^2))
+end
+
+########################
+# transverse coordinates
+########################
+@inline function pt2(::XYZT, mom)
+  return x(mom)^2 + y(mom)^2
+end
+
+@inline function pt(::XYZT, mom)
+  return sqrt(pt2(mom))
+end
+
+@inline function mt2(::XYZT, mom)
+  return t(mom)^2 - z(mom)^2
+end
+
+function mt(::XYZT, 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(::XYZT, mom)
+  en = t(mom)
+  zcomp = z(mom)
+  return 0.5 * log((en + zcomp) / (en - zcomp))
+end
+
+function eta(::XYZT, mom)
+  cth = cos_theta(mom)
+
+  if cth^2 < one(cth)
+    return -0.5 * log((1 - cth) / (1 + cth))
+  end
+
+  zcomp = z(mom)
+  if iszero(zcomp)
+    return zero(zcomp)
+  end
+
+  @warn "Pseudorapidity (η): transverse momentum is zero! return +/- 10e10"
+  if zcomp > zero(zcomp)
+    return 10e10
+  end
+
+  return 10e-10
+end
+
+#######################
+# spherical coordinates
+#######################
+@inline function polar_angle(::XYZT, mom)
+  xcomp = x(mom)
+  ycomp = y(mom)
+  zcomp = z(mom)
+
+  return if iszero(xcomp) && iszero(ycomp) && iszero(zcomp)
+    zero(xcomp)
+  else
+    atan(transverse_momentum(mom), zcomp)
+  end
+end
+
+@inline function cos_theta(::XYZT, mom)
+  r = spatial_magnitude(mom)
+  return iszero(r) ? one(x(mom)) : z(mom) / r
+end
+
+function phi(::XYZT, mom)
+  xcomp = x(mom)
+  ycomp = y(mom)
+  return iszero(xcomp) && iszero(ycomp) ? zero(xcomp) : atan(ycomp, xcomp)
+end
+
+function cos_phi(::XYZT, mom)
+  pT = transverse_momentum(mom)
+  return iszero(pT) ? one(pT) : x(mom) / pT
+end
+
+function sin_phi(::XYZT, mom)
+  pT = transverse_momentum(mom)
+  return iszero(pT) ? zero(pT) : y(mom) / pT
+end
+
+########################
+# light cone coordinates
+########################
+@inline function plus_component(::XYZT, mom)
+  return 0.5 * (t(mom) + z(mom))
+end
+
+@inline function minus_component(::XYZT, mom)
+  return 0.5 * (t(mom) - z(mom))
+end