Rationalise final jets (JuliaHEP#198)

* First attempt at parametrising final_jets

* Fix return values from jet selectors

The default return values from exclusive and inclusive jet selectors
were wrong, as the LorentzVectorCyl was created mistakenly with
rapidity(jet) instead of eta(jet). This went unnoticed as the FinalJets
reversed the mistake, treating eta() as if it was rapidity().

In testing, there is quite a precision loss involved in conversion to a
LorentzVectorCyl as this involves converting from (px,py,py,E) to
(pT,eta,phi,m), then re-extracting (rapidity,phi,pT). It is the
rapidity<->phi which is problematic. So, the default return type has
been changed to be a LorentzVector as this matches better the internal
representation of PseudoJet and EEJet.

Also now import eta() and pt() methods for LorentzVectorCyl and
LorentzVector.

Decrease the tolerance on the pT test from 1e-6 to 1e-7 (for Float64).

* Formatting

* Add jet selector tests

Test if different struct selectors work for selected jets, both
inclusive and exclusive.

Add mass and mass2 "imports" for LorentzVector and LorentzVectorHEP.

* Broaden test coverage

* Add new tests to runtests.jl

* Format!

* Test mass and eta as well

* Better isapprox for FinalJet

Pass kwargs properly and don't rely on the ≈ alias.

Use the default equality implementation.

Author: graeme-a-stewart

src/ClusterSequence.jl

@@ -235,18 +235,17 @@ It iterates over the clustering history and checks the transverse momentum of
 each parent jet. If the transverse momentum is greater than or equal to `ptmin`,
 the jet is added to the array of inclusive jets.
 
-Valid return types are `LorentzVectorCyl` and the jet type of the input `clusterseq`
-(`U` - either `PseudoJet` or `EEJet` depending which algorithm was used)
-(N.B. this will evolve in the future to be any subtype of `FourMomentumBase`;
-currently unrecognised types will return `LorentzVectorCyl`).
+Valid return types are `LorentzVector` `LorentzVectorCyl` or the jet type of the
+input `clusterseq` (`U` - either `PseudoJet` or `EEJet` depending which
+algorithm was used).
 
 # Example
 ```julia
 inclusive_jets(clusterseq; ptmin = 10.0)
 ```
 """
 function inclusive_jets(clusterseq::ClusterSequence{U},
-                        ::Type{T} = LorentzVectorCyl{Float64};
+                        ::Type{T} = LorentzVector{Float64};
                         ptmin = 0.0) where {T, U}
     pt2min = ptmin * ptmin
     jets_local = T[]
@@ -263,25 +262,29 @@ function inclusive_jets(clusterseq::ClusterSequence{U},
             @debug "Added inclusive jet index $iparent_jet"
             if T == U
                 push!(jets_local, jet)
+            elseif T <: LorentzVectorCyl
+                push!(jets_local, lorentzvector_cyl(jet))
+            elseif T <: LorentzVector
+                push!(jets_local, lorentzvector(jet))
             else
-                push!(jets_local,
-                      LorentzVectorCyl(pt(jet), rapidity(jet), phi(jet), mass(jet)))
+                error("Unsupported return type $T for inclusive jets")
             end
         end
     end
     jets_local
 end
 
 """
-    exclusive_jets(clusterseq::ClusterSequence{U}, ::Type{T} = LorentzVectorCyl{Float64}; dcut = nothing, njets = nothing) where {T, U}
+    exclusive_jets(clusterseq::ClusterSequence{U}, ::Type{T} = LorentzVector{Float64}; dcut = nothing, njets = nothing) where {T, U}
 
 Return all exclusive jets of a ClusterSequence, with either a specific number of
 jets or a cut on the maximum distance parameter.
 
 # Arguments
 - `clusterseq::ClusterSequence`: The `ClusterSequence` object containing the
   clustering history and jets.
-- `::Type{T} = LorentzVectorCyl{Float64}`: The return type used for the selected jets.
+- `::Type{T} = LorentzVectorCyl{Float64}`: The return type used for the selected
+  jets.
 - `dcut::Union{Nothing, Real}`: The distance parameter used to define the
   exclusive jets. If `dcut` is provided, the number of exclusive jets will be
   calculated based on this parameter.
@@ -294,10 +297,9 @@ jets or a cut on the maximum distance parameter.
 # Returns
 - An array of `T` objects representing the exclusive jets.
 
-Valid return types are `LorentzVectorCyl` and the jet type of the input `clusterseq`
-(`U` - either `PseudoJet` or `EEJet` depending which algorithm was used)
-(N.B. this will evolve in the future to be any subtype of `FourMomentumBase`;
-currently unrecognised types will return `LorentzVectorCyl`).
+Valid return types are `LorentzVector` `LorentzVectorCyl` or the jet type of the
+input `clusterseq` (`U` - either `PseudoJet` or `EEJet` depending which
+algorithm was used).
 
 # Exceptions
 - `ArgumentError`: If neither `dcut` nor `njets` is provided.
@@ -313,7 +315,7 @@ exclusive_jets(clusterseq, PseudoJet, njets = 3)
 ```
 """
 function exclusive_jets(clusterseq::ClusterSequence{U},
-                        ::Type{T} = LorentzVectorCyl{Float64};
+                        ::Type{T} = LorentzVector{Float64};
                         dcut = nothing, njets = nothing,) where {T, U}
     if isnothing(dcut) && isnothing(njets)
         throw(ArgumentError("Must pass either a dcut or an njets value"))
@@ -355,14 +357,16 @@ function exclusive_jets(clusterseq::ClusterSequence{U},
                 jet = clusterseq.jets[clusterseq.history[parent].jetp_index]
                 if T == U
                     push!(excl_jets, jet)
+                elseif T <: LorentzVectorCyl
+                    push!(excl_jets, lorentzvector_cyl(jet))
+                elseif T <: LorentzVector
+                    push!(excl_jets, lorentzvector(jet))
                 else
-                    push!(excl_jets,
-                          LorentzVectorCyl(pt(jet), rapidity(jet), phi(jet), mass(jet)))
+                    error("Unsupported return type $T for inclusive jets")
                 end
             end
         end
     end
-
     excl_jets
 end
 

src/CommonJetStructs.jl

@@ -109,6 +109,9 @@ phi(j::FourMomentum) = begin
     phi
 end
 
+# Alternative name for [0, 2π) range
+const phi02pi = phi
+
 """Used to protect against parton-level events where pt can be zero
 for some partons, giving rapidity=infinity. KtJet fails in those cases."""
 const _MaxRap = 1e5

src/JSONresults.jl

@@ -15,6 +15,28 @@ struct FinalJet
     pt::Float64
 end
 
+"""
+    FinalJet(jet::T)
+
+Convert a jet of type `T` to a `FinalJet` object. `T` must be a type that
+supports the methods `rapidity`, `phi`, and `pt`.
+"""
+FinalJet(jet::T) where {T} = FinalJet(rapidity(jet), phi(jet), pt(jet))
+
+import Base: isapprox
+
+"""
+    isapprox(fj1::FinalJet, fj2::FinalJet; kwargs...)
+
+Compare two `FinalJet` objects for approximate equality based on their rapidity,
+azimuthal angle, and transverse momentum. `kwargs` are passed to the `isapprox`
+function, e.g., `tol`.
+"""
+function isapprox(fj1::FinalJet, fj2::FinalJet; kwargs...)
+    isapprox(fj1.rap, fj2.rap; kwargs...) && isapprox(fj1.phi, fj2.phi; kwargs...) &&
+        isapprox(fj1.pt, fj2.pt; kwargs...)
+end
+
 """
     struct FinalJets
 

src/JetReconstruction.jl

@@ -21,20 +21,28 @@ using StructArrays
 
 # Import from LorentzVectorHEP methods for those 4-vector types
 pt2(p::LorentzVector) = LorentzVectorHEP.pt2(p)
+pt(p::LorentzVector) = LorentzVectorHEP.pt(p)
 phi(p::LorentzVector) = LorentzVectorHEP.phi(p)
 rapidity(p::LorentzVector) = LorentzVectorHEP.rapidity(p)
+eta(p::LorentzVector) = LorentzVectorHEP.eta(p)
 px(p::LorentzVector) = LorentzVectorHEP.px(p)
 py(p::LorentzVector) = LorentzVectorHEP.py(p)
 pz(p::LorentzVector) = LorentzVectorHEP.pz(p)
 energy(p::LorentzVector) = LorentzVectorHEP.energy(p)
+mass2(p::LorentzVector) = LorentzVectorHEP.mass2(p)
+mass(p::LorentzVector) = LorentzVectorHEP.mass(p)
 
 pt2(p::LorentzVectorCyl) = LorentzVectorHEP.pt2(p)
+pt(p::LorentzVectorCyl) = LorentzVectorHEP.pt(p)
 phi(p::LorentzVectorCyl) = LorentzVectorHEP.phi(p)
 rapidity(p::LorentzVectorCyl) = LorentzVectorHEP.rapidity(p)
+eta(p::LorentzVectorCyl) = LorentzVectorHEP.eta(p)
 px(p::LorentzVectorCyl) = LorentzVectorHEP.px(p)
 py(p::LorentzVectorCyl) = LorentzVectorHEP.py(p)
 pz(p::LorentzVectorCyl) = LorentzVectorHEP.pz(p)
 energy(p::LorentzVectorCyl) = LorentzVectorHEP.energy(p)
+mass2(p::LorentzVectorCyl) = LorentzVectorHEP.mass2(p)
+mass(p::LorentzVectorCyl) = LorentzVectorHEP.mass(p)
 
 # Some useful constants/limits that are used internally
 const max_allowable_R = 1000.0

src/Utils.jl

@@ -73,64 +73,28 @@ function read_final_state_particles(fname, ::Type{T} = PseudoJet; maxevents = -1
 end
 
 """
-    final_jets(jets::Vector{PseudoJet}, ptmin::AbstractFloat=0.0)
+    final_jets(jets::Vector{T}, ptmin::AbstractFloat=0.0)
 
-This function takes a vector of `PseudoJet` objects and a minimum transverse
-momentum `ptmin` as input. It returns a vector of `FinalJet` objects that
-satisfy the transverse momentum condition.
+This function takes a vector of `T` objects, which should be jets, and a minimum
+transverse momentum `ptmin` as input. It returns a vector of `FinalJet` objects
+that satisfy the transverse momentum condition.
 
 # Arguments
-- `jets::Vector{PseudoJet}`: A vector of `PseudoJet` objects representing the
-  input jets.
+- `jets::Vector{T}`: A vector of `T` objects representing the input jets.
 - `ptmin::AbstractFloat=0.0`: The minimum transverse momentum required for a jet
   to be included in the final jets vector.
 
 # Returns
 A vector of `FinalJet` objects that satisfy the transverse momentum condition.
 """
-function final_jets(jets::Vector{PseudoJet}, ptmin::AbstractFloat = 0.0)
+function final_jets(jets::Vector{T}, ptmin::AbstractFloat = 0.0) where {T}
     count = 0
     final_jets = Vector{FinalJet}()
-    sizehint!(final_jets, 6)
-    for jet in jets
-        dcut = ptmin^2
-        if pt2(jet) > dcut
-            count += 1
-            push!(final_jets, FinalJet(rapidity(jet), phi(jet), sqrt(pt2(jet))))
-        end
-    end
-    final_jets
-end
-
-"""Specialisation for final jets from LorentzVectors (TODO: merge into more general function)"""
-function final_jets(jets::Vector{<:LorentzVector}, ptmin::AbstractFloat = 0.0)
-    count = 0
-    final_jets = Vector{FinalJet}()
-    sizehint!(final_jets, 6)
     dcut = ptmin^2
     for jet in jets
-        if LorentzVectorHEP.pt(jet)^2 > dcut
-            count += 1
-            push!(final_jets,
-                  FinalJet(LorentzVectorHEP.rapidity(jet), LorentzVectorHEP.phi(jet),
-                           LorentzVectorHEP.pt(jet)))
-        end
-    end
-    final_jets
-end
-
-"""Specialisation for final jets from LorentzVectorCyl (TODO: merge into more general function)"""
-function final_jets(jets::Vector{<:LorentzVectorCyl}, ptmin::AbstractFloat = 0.0)
-    count = 0
-    final_jets = Vector{FinalJet}()
-    sizehint!(final_jets, 6)
-    dcut = ptmin^2
-    for jet in jets
-        if LorentzVectorHEP.pt(jet)^2 > dcut
+        if pt2(jet) > dcut
             count += 1
-            push!(final_jets,
-                  FinalJet(LorentzVectorHEP.eta(jet), LorentzVectorHEP.phi(jet),
-                           LorentzVectorHEP.pt(jet)))
+            push!(final_jets, FinalJet(jet))
         end
     end
     final_jets

test/_common.jl

@@ -139,7 +139,7 @@ function run_reco_test(test::ComparisonTest; testname = nothing)
                         normalised_phi = jet.phi < 0.0 ? jet.phi + 2π : jet.phi
                         @test jet.rap≈fastjet_jets[ievt]["jets"][ijet]["rap"] atol=1e-7
                         @test normalised_phi≈fastjet_jets[ievt]["jets"][ijet]["phi"] atol=1e-7
-                        @test jet.pt≈fastjet_jets[ievt]["jets"][ijet]["pt"] rtol=1e-6
+                        @test jet.pt≈fastjet_jets[ievt]["jets"][ijet]["pt"] rtol=1e-7
                     end
                 end
             end

test/runtests.jl

@@ -31,6 +31,9 @@ function main()
     # Utility support tests
     include("test-utils.jl")
 
+    # Check jet selectors to different jet types
+    include("test-jet-selectors.jl")
+
     # Substructure tests
     include("test-substructure.jl")
 

test/test-compare-types.jl

@@ -49,9 +49,7 @@ function do_test_compare_types(strategy::RecoStrategy.Strategy;
                 @test size(event.jets) == size(event_lv.jets)
                 # Test each jet in turn
                 for (jet, jet_lv) in zip(event.jets, event_lv.jets)
-                    @test jet.rap≈jet_lv.rap atol=1e-7
-                    @test jet.phi≈jet_lv.phi atol=1e-7
-                    @test jet.pt≈jet_lv.pt rtol=1e-6
+                    @test jet == jet_lv
                 end
             end
         end

test/test-jet-selectors.jl

@@ -0,0 +1,55 @@
+# Test that final jets can be retrieved with different
+# output types
+
+include("common.jl")
+
+"""
+    test_jet_equality(ref, test)
+
+Test two jet representations for equality.
+"""
+function test_jet_equality(ref, test)
+    for (j1, j2) in zip(ref, test)
+        @test JetReconstruction.pt(j1) ≈ JetReconstruction.pt(j2)
+        test_phi = JetReconstruction.phi02pi(j2)
+        test_phi = test_phi < 0.0 ? test_phi + 2π : test_phi
+        @test JetReconstruction.phi(j1) ≈ test_phi
+        @test JetReconstruction.rapidity(j1) ≈ JetReconstruction.rapidity(j2)
+        @test JetReconstruction.eta(j1) ≈ JetReconstruction.eta(j2)
+        @test JetReconstruction.px(j1) ≈ JetReconstruction.px(j2)
+        @test JetReconstruction.py(j1) ≈ JetReconstruction.py(j2)
+        @test JetReconstruction.pz(j1) ≈ JetReconstruction.pz(j2)
+        @test JetReconstruction.energy(j1) ≈ JetReconstruction.energy(j2)
+        @test JetReconstruction.mass2(j1) ≈ JetReconstruction.mass2(j2)
+        @test JetReconstruction.mass(j1) ≈ JetReconstruction.mass(j2)
+    end
+    nothing
+end
+
+@testset "Jet selections to types" begin
+    pp_event = first(JetReconstruction.read_final_state_particles(events_file_pp;
+                                                                  maxevents = 1))
+    ee_event = first(JetReconstruction.read_final_state_particles(events_file_ee;
+                                                                  maxevents = 1))
+
+    pp_cs = jet_reconstruct(pp_event; algorithm = JetAlgorithm.Kt)
+    ee_cs = jet_reconstruct(ee_event; algorithm = JetAlgorithm.Durham)
+
+    pp_ref_jets = inclusive_jets(pp_cs, PseudoJet; ptmin = 10.0)
+    for T in (LorentzVectorCyl, LorentzVector)
+        test_jet_equality(pp_ref_jets, inclusive_jets(pp_cs, T; ptmin = 10.0))
+    end
+
+    ee_ref_jets = exclusive_jets(ee_cs, EEJet; njets = 2)
+    for T in (LorentzVectorCyl, LorentzVector)
+        test_jet_equality(ee_ref_jets, exclusive_jets(ee_cs, T; njets = 2))
+    end
+
+    # Also check that for an unknown type we throw
+    @test_throws ErrorException begin
+        inclusive_jets(pp_cs, Float64; ptmin = 10.0)
+    end
+    @test_throws ErrorException begin
+        exclusive_jets(ee_cs, Float64; njets = 2)
+    end
+end

test/test-jet-utils.jl

@@ -41,3 +41,24 @@ eej2 = EEJet(-1.0, 3.2, -1.2, 39.0)
     @test LorentzVectorHEP.mass(lvc_pj1) ≈ JetReconstruction.mass(eej1) ≈
           JetReconstruction.mass(eej1) ≈ LorentzVectorHEP.mass(lvc_eej1)
 end
+
+@testset "Final jets extraction" begin
+    # Test final jets extraction from different jet types
+    vec_pj = [pj1, pj2]
+    finaljets_pj = JetReconstruction.final_jets(vec_pj)
+
+    vec_lorentzhep = [JetReconstruction.lorentzvector(pj1),
+        JetReconstruction.lorentzvector(pj2)]
+    finaljets_lorentzhep = JetReconstruction.final_jets(vec_lorentzhep)
+
+    vec_lorentzhepcyl = [JetReconstruction.lorentzvector_cyl(pj1),
+        JetReconstruction.lorentzvector_cyl(pj2)]
+    finaljets_lorentzhepcyl = JetReconstruction.final_jets(vec_lorentzhepcyl)
+
+    @test length(finaljets_pj) == 2
+    @test length(finaljets_lorentzhep) == 2
+    @test length(finaljets_lorentzhepcyl) == 2
+    for i in 1:2
+        @test finaljets_pj[i] ≈ finaljets_lorentzhep[i] ≈ finaljets_lorentzhepcyl[i]
+    end
+end