Add jet parents retrieval

src/ClusterSequence.jl

@@ -610,3 +610,31 @@ An vector of indices representing the original constituents of the given jet.
 function constituent_indexes(jet::T, cs::ClusterSequence{T}) where {T <: FourMomentum}
     get_all_ancestors(cs.history[jet._cluster_hist_index].jetp_index, cs)
 end
+
+"""
+    parent_jets(jet::T, cs::ClusterSequence{T})::Tuple{Union{Nothing, T}, Union{Nothing, T}} where {T <: FourMomentum}
+
+Find the parent jets of a given jet in a cluster sequence.
+
+# Arguments
+- `jet::T`: The jet for which to find the parent jets.
+- `cs::ClusterSequence`: The cluster sequence object.
+
+# Returns
+A tuple of two elements, each of which is either the parent jet object or
+`nothing` (if the jet has no parent).
+"""
+function parent_jets(jet::T,
+                     cs::ClusterSequence{T})::Tuple{Union{Nothing, T},
+                                                    Union{Nothing, T}} where {T <:
+                                                                              FourMomentum}
+    hist_idx = jet._cluster_hist_index
+    jet_history = cs.history[hist_idx]
+
+    parent1_idx, parent2_idx = jet_history.parent1, jet_history.parent2
+
+    parent1_jet = parent1_idx > 0 ? cs.jets[cs.history[parent1_idx].jetp_index] : nothing
+    parent2_jet = parent2_idx > 0 ? cs.jets[cs.history[parent2_idx].jetp_index] : nothing
+
+    return parent1_jet, parent2_jet
+end

src/JetReconstruction.jl

@@ -48,7 +48,7 @@ export RecoStrategy, JetAlgorithm
 # ClusterSequence type
 include("ClusterSequence.jl")
 export ClusterSequence, inclusive_jets, exclusive_jets, n_exclusive_jets, constituents,
-       constituent_indexes
+       constituent_indexes, parent_jets
 
 ## N2Plain algorithm
 # Algorithmic part for simple sequential implementation

src/Substructure.jl

@@ -1,25 +1,3 @@
-"""
-    has_parents(p, clusterseq) -> (boolean, Int64, Int64)
-
-Checks if the jet `p` is a child of two other jets, after clustering 
-
-# Arguments
-- `p::PseudoJet`: The jet to check.
-- `clusterseq::ClusterSequence`: The cluster sequence object.
-
-# Returns
-- `(boolean, Int64, Int64)`: true or false depending on if the jet has a parent or not. If the jet has a parent, returns the indices of the parent jets in the history element. Otherwise, returns -2 (NonexistentParent).
-"""
-function has_parents(p::PseudoJet, clusterseq::ClusterSequence)
-    history = clusterseq.history
-    N = p._cluster_hist_index
-    p1 = history[N].parent1
-    p2 = history[N].parent2
-
-    p1 == p2 == NonexistentParent ? result = false : result = true
-    return (result, p1, p2)
-end
-
 """
     deltaR(jet1, jet2) -> Float64
 
@@ -147,14 +125,10 @@ function mass_drop(jet::PseudoJet, clusterseq::ClusterSequence, tag::MassDropTag
     hist = clusterseq.history
 
     while true
-        had_parents, p1, p2 = has_parents(jet, clusterseq)
-
-        if had_parents
-            parent1 = all_jets[hist[p1].jetp_index]
-            parent2 = all_jets[hist[p2].jetp_index]
+        parent1, parent2 = parent_jets(jet, clusterseq)
 
+        if !isnothing(parent1)
             if m2(parent1) < m2(parent2)
-                p1, p2 = p2, p1
                 parent1, parent2 = parent2, parent1
             end
 
@@ -198,21 +172,17 @@ function soft_drop(jet::PseudoJet, clusterseq::ClusterSequence,
     hist = new_clusterseq.history
 
     while true
-        had_parents, p1, p2 = has_parents(new_jet, new_clusterseq)
-
-        if had_parents
-            parent1 = all_jets[hist[p1].jetp_index]
-            parent2 = all_jets[hist[p2].jetp_index]
+        parent1, parent2 = parent_jets(new_jet, new_clusterseq)
 
+        if !isnothing(parent1)
             if m2(parent1) < m2(parent2)
-                p1, p2 = p2, p1
                 parent1, parent2 = parent2, parent1
             end
 
-            pti = pt(parent1)
-            ptj = pt(parent2)
+            pt1 = pt(parent1)
+            pt2 = pt(parent2)
 
-            if min(pti, ptj) / (pti + ptj) >
+            if min(pt1, pt2) / (pt1 + pt2) >
                tag.zcut * (deltaR(parent1, parent2) / rad)^tag.b
                 return new_jet
             else

test/test-constituents.jl

@@ -1,4 +1,4 @@
-# Tests of jet constituent retrieval
+# Tests of jet constituent retrieval and parentage
 
 include("common.jl")
 
@@ -8,8 +8,9 @@ function Base.isapprox(j1::PseudoJet, j2::PseudoJet)
         isapprox(j1.py, j2.py) && isapprox(j1.pz, j2.pz)
 end
 
-# Expected constituent indexes
+# Expected constituent indexes and parent indexes
 const expected_constituent_indexes = [84, 85, 139, 86, 133, 74, 79, 124, 76, 75, 163]
+const expected_parent_indexes = [320, 335]
 
 input_file = joinpath(dirname(pathof(JetReconstruction)), "..", "test", "data",
                       "events.pp13TeV.hepmc3.gz")
@@ -25,20 +26,31 @@ pj_jets = inclusive_jets(cluster_seq; ptmin = 5.0, T = PseudoJet)
 
 @testset "Jet constituents" begin
     @testset "Constituents of jet number $(event_no)" begin
-        my_constituents = JetReconstruction.constituents(pj_jets[1], cluster_seq)
+        my_constituents = JetReconstruction.constituents(pj_jets[event_no], cluster_seq)
         @test size(my_constituents)[1] == 11
         for (i, idx) in enumerate(expected_constituent_indexes)
             @test my_constituents[i] ≈ events[1][idx]
         end
-        # @test my_constituents[1] ≈ events[1][84]
-        # @test my_constituents[1] ≈ LorentzVectorHEP.LorentzVector(0.0, 0.0, 0.0, 0.0)
-        # @test my_constituents[2] ≈ LorentzVectorHEP.LorentzVector(0.0, 0.0, 0.0, 0.0)
     end
 
     @testset "Constituent indexes for jet number $(event_no)" begin
-        my_constituent_indexes = constituent_indexes(pj_jets[1], cluster_seq)
+        my_constituent_indexes = constituent_indexes(pj_jets[event_no], cluster_seq)
         @test size(my_constituent_indexes)[1] == 11
         # Testing the index values is sufficient, the content came from the original input file!
         @test my_constituent_indexes == expected_constituent_indexes
     end
 end
+
+@testset "Jet parents" begin
+    @testset "Parent of jet number $(event_no)" begin
+        my_parents = JetReconstruction.parent_jets(pj_jets[event_no], cluster_seq)
+        @test my_parents[1] ≈
+              cluster_seq.jets[cluster_seq.history[expected_parent_indexes[1]].jetp_index]
+        @test my_parents[2] ≈
+              cluster_seq.jets[cluster_seq.history[expected_parent_indexes[2]].jetp_index]
+    end
+    @testset "Parents of input cluster" begin
+        no_parents = JetReconstruction.parent_jets(cluster_seq.jets[1], cluster_seq)
+        @test isnothing(no_parents[1]) && isnothing(no_parents[2])
+    end
+end