LorenzVectorBase interfaces (#223)

Delegate to 4-vector and 4-momentum calculations to LorentzVectorBase 
package. Support generic conversion from LorentzVectorBase. Rename
CommonJetStructs to CommonJet.

Add LorentzVectorBase package and define ::FourMomentum (EEJet and
PseudoJet) as PxPyPzE coordinate systems.

Update getter functions to call the delegated function from
LorentzVectorBase. For some functions where edge case behaviour
differs from previous behaviour the output is modified to match
what happened before:
- phi should be in the range [0, 2pi)
- rapidity and pseudorapidity should have a maximum cut off for
  stability

Small documentation improvements for PseudoJet accessors which use
internal cached values for efficiency.

For PseudoJet and EEJet update the the constructor from ::Any to test
that the LorentzVectorBase interface is supported and construct from
those accessors.

Tests are added (test-jet-constructors.jl) to check the conversion works and
that it also fails with ArgumentError when a type does not support the
interface. Test explicitly construction from LorentzVector and
LorentzVectorCyl.

Documentation on input particles is updated.

The CommonJetStructs source had a lot more in it than just struct definitions, so the
name was misleading, renamed to CommonJet,

Fix a minor inconsistency in the cross constructors for PseudoJet and
EEJet.

Use the feature that bare identifiers imply the keyword argument
name.

Author: graeme-a-stewart

Project.toml

@@ -1,7 +1,7 @@
 name = "JetReconstruction"
 uuid = "44e8cb2c-dfab-4825-9c70-d4808a591196"
-authors = ["Atell Krasnopolski <delta_atell@protonmail.com>", "Graeme A Stewart <graeme.andrew.stewart@cern.ch", "Philippe Gras <philippe.gras@cern.ch>"]
 version = "1.0.0-dev"
+authors = ["Atell Krasnopolski <delta_atell@protonmail.com>", "Graeme A Stewart <graeme.andrew.stewart@cern.ch", "Philippe Gras <philippe.gras@cern.ch>"]
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
@@ -10,6 +10,7 @@ CodecZstd = "6b39b394-51ab-5f42-8807-6242bab2b4c2"
 EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
+LorentzVectorBase = "592a752d-6533-4762-a71b-738712ea30ec"
 LorentzVectorHEP = "f612022c-142a-473f-8cfd-a09cf3793c6c"
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
@@ -34,6 +35,7 @@ EnumX = "1.0.4"
 JSON = "0.21, 1.2"
 Logging = "1.10"
 LoopVectorization = "0.12.170"
+LorentzVectorBase = "0.1.0"
 LorentzVectorHEP = "0.1.6"
 Makie = "0.20, 0.21, 0.22, 0.23, 0.24"
 MuladdMacro = "0.2.4"

docs/src/particles.md

@@ -1,24 +1,18 @@
 # Input Particle Types
 
 For the `particles` input to the reconstruction any one dimensional
-`AbstractArray{T, 1}` can be used, where the type `T` has to implement methods
-to extract the 4-vector components, viz, the following are required:
+`AbstractArray{T, 1}` can be used, in general the type `T` should implement the
+[LorentzVectorBase](https://github.com/JuliaHEP/LorentzVectorBase.jl)
+interface. That is the `T` should be a valid coordinate system and the
+following methods are required:
 
-- `JetReconstuction.px(particle::T)`
-- `JetReconstuction.py(particle::T)`
-- `JetReconstuction.pz(particle::T)`
-- `JetReconstuction.energy(particle::T)`
+- `LorentzVectorBase.px(particle::T)`
+- `LorentzVectorBase.py(particle::T)`
+- `LorentzVectorBase.pz(particle::T)`
+- `LorentzVectorBase.energy(particle::T)`
 
-Currently built-in supported types are
-[`LorentzVectorHEP`](https://github.com/JuliaHEP/LorentzVectorHEP.jl), the
-`PseudoJet` and `EEJet`s from this package, and `ReconstructedParticles` from
-[EDM4hep Inputs](@ref).
+In addition, direct support for the construction of particles from
+`LorentzVector` and `LorentzVectorCyl` types is provided.
 
-If you require support for a different input collection type then ensure you
-define the `px()`, etc. methods *for your specific type* and *in the
-`JetReconstruction` package*. This use of what might be considered type piracy
-is blessed as long as you are en *end user* of the jet reconstruction package.
-
-If your type is used in several places or by different users, please consider
-writing a package extension that will support your type, following the model for
-EDM4hep in `ext/EDM4hepJets.jl`.
+Direct use of the `PseudoJet` and `EEJet`s from this package, and
+`ReconstructedParticles` from [EDM4hep Inputs](@ref) is supported.

src/CommonJetStructs.jl src/CommonJet.jlsrc/CommonJetStructs.jl renamed to src/CommonJet.jl

@@ -4,9 +4,22 @@
 Interface for composite types that includes fields px, py, py, and E
 that represents the components of a four-momentum vector. All concrete
 jets that are used in the package are subtypes of this type.
+
+These types are also added as LorentzVectorBase `PxPyPzE` coordinate systems.
 """
 abstract type FourMomentum end
 
+"""
+    LorentzVectorBase.coordinate_system(::FourMomentum)
+
+All of this package's FourMomentum types are `PxPyPzE` coordinate systems.
+"""
+LorentzVectorBase.coordinate_system(::FourMomentum) = LorentzVectorBase.PxPyPzE()
+LorentzVectorBase.px(v::FourMomentum) = v.px
+LorentzVectorBase.py(v::FourMomentum) = v.py
+LorentzVectorBase.pz(v::FourMomentum) = v.pz
+LorentzVectorBase.E(v::FourMomentum) = v.E
+
 # Define here all common functions that can be used for all jet types <: FourMomentum
 """
     px(j::FourMomentum)
@@ -50,21 +63,28 @@ cluster_hist_index(j::FourMomentum) = j._cluster_hist_index
 
 Return the squared momentum of the four-momentum vector of `j`.
 """
-p2(j::FourMomentum) = j.px^2 + j.py^2 + j.pz^2
+p2(j::FourMomentum) = LorentzVectorBase.spatial_magnitude2(j)
 
 """
     p(j::FourMomentum)
 
 Return the momentum of the four-momentum vector of `j`.
 """
-p(j::FourMomentum) = sqrt(p2(j))
+p(j::FourMomentum) = LorentzVectorBase.spatial_magnitude(j)
+
+"""
+    mag(j::FourMomentum)
+
+Return the spatial magnitude (alias for `p()`).
+"""
+const mag = p
 
 """
     pt2(j::FourMomentum)
 
 Return the squared transverse momentum of the four-momentum vector of `j`.
 """
-pt2(j::FourMomentum) = j.px^2 + j.py^2
+pt2(j::FourMomentum) = LorentzVectorBase.transverse_momentum2(j)
 
 """Alias for `pt2` function"""
 const kt2 = pt2
@@ -73,14 +93,14 @@ const kt2 = pt2
     pt(j::FourMomentum)
 Return the momentum of the four-momentum vector of `j`.
 """
-pt(j::FourMomentum) = sqrt(pt2(j))
+pt(j::FourMomentum) = LorentzVectorBase.transverse_momentum(j)
 
 """
     mass2(j::FourMomentum)
 
 Return the invariant mass squared of the four-momentum vector `j`.
 """
-mass2(j::FourMomentum) = energy(j)^2 - p2(j)
+mass2(j::FourMomentum) = LorentzVectorBase.mass2(j)
 
 """Alias for `mass2` function"""
 const m2 = mass2
@@ -91,7 +111,7 @@ const m2 = mass2
 Return the invariant mass of a four momentum `j`. By convention if ``m^2 < 0``,
 then ``-sqrt{(-m^2)}`` is returned.
 """
-mass(j::FourMomentum) = m2(j) < 0.0 ? -sqrt(-m2(j)) : sqrt(m2(j))
+mass(j::FourMomentum) = LorentzVectorBase.mass(j)
 
 """Alias for `mass` function"""
 const m = mass
@@ -101,12 +121,9 @@ const m = mass
 
 Return the azimuthal angle, ϕ, of the four momentum `j` in the range [0, 2π).
 """
-phi(j::FourMomentum) = begin
-    phi = pt2(j) == 0.0 ? 0.0 : atan(j.py, j.px)
-    if phi < 0.0
-        phi += 2π
-    end
-    phi
+function phi(j::FourMomentum)
+    phi = LorentzVectorBase.phi(j)
+    return phi < 0.0 ? phi + 2π : phi
 end
 
 # Alternative name for [0, 2π) range
@@ -120,35 +137,16 @@ const _MaxRap = 1e5
     rapidity(j::FourMomentum)
 
 Return the rapidity of the four momentum `j`.
+
+In this package, for numerical stability, we clamp the rapidity, but in a way
+that allows for predictability, based on pz.
 """
 function rapidity(j::FourMomentum)
-    if energy(j) == abs(pz(j)) && iszero(pt2(j))
-        MaxRapHere = _MaxRap + abs(pz(j))
-        rapidity = (pz(j) >= 0.0) ? MaxRapHere : -MaxRapHere
-    else
-        # get the rapidity in a way that's modestly insensitive to roundoff
-        # error when things pz,E are large (actually the best we can do without
-        # explicit knowledge of mass)
-        effective_m2 = max(0.0, m2(j)) # force non tachyonic mass
-        E_plus_pz = energy(j) + abs(pz(j)) # the safer of p+, p-
-        rapidity = 0.5 * log((pt2(j) + effective_m2) / (E_plus_pz * E_plus_pz))
-        if pz(j) > 0
-            rapidity = -rapidity
-        end
-    end
-    rapidity
+    rap = LorentzVectorBase.rapidity(j)
+    MaxRapHere = _MaxRap + abs(pz(j))
+    return clamp(rap, -MaxRapHere, MaxRapHere)
 end
 
-"""
-    mag(jet::T) where {T <: FourMomentum}
-
-Return the magnitude of the momentum of a jet, `|p|`.
-
-# Returns
-The magnitude of the jet.
-"""
-mag(jet::T) where {T <: FourMomentum} = sqrt(muladd(jet.px, jet.px, jet.py^2) + jet.pz^2)
-
 """
     CosTheta(jet::T) where {T <: FourMomentum}
 
@@ -157,29 +155,18 @@ Compute the cosine of the angle between the momentum vector of `jet` and the z-a
 # Returns
 - The cosine of the angle between the jet and the z-axis.
 """
-@inline function CosTheta(jet::T) where {T <: FourMomentum}
-    fZ = jet.pz
-    ptot = mag(jet)
-    return ifelse(ptot == 0.0, 1.0, fZ / ptot)
-end
+CosTheta(j::FourMomentum) = LorentzVectorBase.cos_theta(j)
 
 """
     eta(jet::T) where {T <: FourMomentum}
 
-Compute the pseudorapidity (η) of a jet.
-
-# Returns
-- The pseudorapidity (η) of the jet.
-"""
-function eta(jet::T) where {T <: FourMomentum}
-    cosTheta = CosTheta(jet)
-    (cosTheta^2 < 1.0) && return -0.5 * log((1.0 - cosTheta) / (1.0 + cosTheta))
-    fZ = jet.pz
-    iszero(fZ) && return 0.0
-    # Warning("PseudoRapidity","transverse momentum = 0! return +/- 10e10");
-    fZ > 0.0 ? _MaxRap : -_MaxRap
+Returns the pseudorapidity (η or eta) of a jet. For this package we
+clamp the return value to our package's `_MaxRap`.
+"""
+function eta(j::FourMomentum)
+    eta = LorentzVectorBase.eta(j)
+    clamp(eta, -_MaxRap, _MaxRap)
 end
-
 """
     const η = eta
 

src/EEJet.jl

@@ -83,36 +83,49 @@ end
     EEJet(jet::LorentzVector; cluster_hist_index::Int = 0)
 
 Construct a EEJet from a `LorentzVector` object with optional cluster index.
+
+The `cluster_hist_index` is optional, but needed if the `jet` is part of a
+reconstruction sequence. If not provided, it defaults to `0` as an "invalid"
+value.
 """
 function EEJet(jet::LorentzVector; cluster_hist_index::Int = 0)
-    EEJet(jet.x, jet.y, jet.z, jet.t; cluster_hist_index = cluster_hist_index)
+    EEJet(jet.x, jet.y, jet.z, jet.t; cluster_hist_index)
 end
 
 """
     EEJet(jet::Any; cluster_hist_index::Int = 0)
 
-Construct a EEJet from a generic object `jet` with the given cluster index.
-This functions also for `LorentzVectorCyl` objects.
+Construct an EEJet from a generic object `jet` with the given cluster index.
+These generic jets must implement the LorentzVectorBase interface from
+`LorentzVectorBase.jl`.
 
 # Details
 
-This function is used to convert a generic object `jet` into an `EEJet`, for
-this to work the object must have the methods `px`, `py`, `pz`, and `energy`
-defined, which are used to extract the four-momentum components of the object.
+This function is used to convert a generic object `jet` into a `EEJet`.
+These generic jets should implement the LorentzVectorBase interface: the
+`LorentzVectorBase.{px,py,pz,energy}()` methods will be called to retrieve
+the four momentum components.
 
 The `cluster_hist_index` is optional, but needed if the `jet` is part of a
 reconstruction sequence. If not provided, it defaults to `0` as an "invalid"
 value.
 """
 function EEJet(jet::Any; cluster_hist_index::Int = 0)
-    EEJet(px(jet), py(jet), pz(jet), energy(jet);
-          cluster_hist_index = cluster_hist_index)
+    # Check that the interface is implemented
+    if hasmethod(LorentzVectorBase.coordinate_system, (typeof(jet),))
+        return EEJet(LorentzVectorBase.px(jet), LorentzVectorBase.py(jet),
+                     LorentzVectorBase.pz(jet), LorentzVectorBase.energy(jet);
+                     cluster_hist_index)
+    else
+        throw(ArgumentError("EEJet cannot be constructed from object of type '$(typeof(jet))'"))
+    end
 end
 
 """
     p2(eej::EEJet)
 
-Return the squared momentum of the `EEJet` object `eej`.
+Return the squared momentum of the `EEJet` object `eej`. This accessor uses the
+pre-calculated value that the struct has.
 """
 p2(eej::EEJet) = eej._p2
 

src/JetReconstruction.jl

@@ -15,6 +15,7 @@ examples in this package for more information.
 """
 module JetReconstruction
 
+using LorentzVectorBase
 using LorentzVectorHEP
 using MuladdMacro
 using StructArrays
@@ -48,7 +49,7 @@ mass(p::LorentzVectorCyl) = LorentzVectorHEP.mass(p)
 const max_allowable_R = 1000.0
 
 # Pseudojet and EEJet types
-include("CommonJetStructs.jl")
+include("CommonJet.jl")
 include("PseudoJet.jl")
 include("EEJet.jl")
 include("JetUtils.jl")

src/JetUtils.jl

@@ -3,14 +3,20 @@
 # Functions that create each jet type from the other need to be defined here,
 # after the structs are defined.
 """
-    PseudoJet(jet::EEJet)
+    PseudoJet(jet::EEJet; cluster_hist_index::Int=0)
 
-Constructs a `PseudoJet` object from an `EEJet` object, with the same four
-momentum and cluster history index.
+Constructs a `PseudoJet` from an `EEJet`.
+
+# Details
+
+The `cluster_hist_index` is set to the value of the `cluster_hist_index` of the
+EEJet if `0` is passed. Otherwise it is set to the value, `>0`, passed in.
 """
-function PseudoJet(eej::EEJet)
+function PseudoJet(eej::EEJet; cluster_hist_index::Int = 0)
+    cluster_hist_index = cluster_hist_index == 0 ? eej._cluster_hist_index :
+                         cluster_hist_index
     PseudoJet(px(eej), py(eej), pz(eej), energy(eej);
-              cluster_hist_index = cluster_hist_index(eej))
+              cluster_hist_index)
 end
 
 """
@@ -24,10 +30,10 @@ The `cluster_hist_index` is set to the value of the `cluster_hist_index` of the
 PseudoJet if `0` is passed. Otherwise it is set to the value, `>0`, passed in.
 """
 function EEJet(jet::PseudoJet; cluster_hist_index::Int = 0)
-    new_cluster_hist_index = cluster_hist_index == 0 ? jet._cluster_hist_index :
-                             cluster_hist_index
+    cluster_hist_index = cluster_hist_index == 0 ? jet._cluster_hist_index :
+                         cluster_hist_index
     EEJet(px(jet), py(jet), pz(jet), energy(jet);
-          cluster_hist_index = new_cluster_hist_index)
+          cluster_hist_index)
 end
 
 # Functions to convert jets to types from other packages