Improved constituents (#98)

graeme-a-stewart · web-flow · commit 0c618d51fed2 · 2024-11-27T13:49:20.000+01:00
Improved jet constituent retrieval, with one method to retrieve jets in a collection (by reference to the reconstruction jets) and one to just retrieve the indexes.

Added tests, documentation and examples.
diff --git a/docs/README.md b/docs/README.md
@@ -2,7 +2,7 @@
 
 ## Adding Documentation
 
-`Documenter.jl` is used to generate package documentation. 
+`Documenter.jl` is used to generate package documentation.
 
 The general guidelines are to have documentation split into meaningful sections,
 none of which are too long. If a new functionality is added then create a new
diff --git a/docs/src/examples.md b/docs/src/examples.md
@@ -61,3 +61,8 @@ showing how the jets merge from their different constituents.
 The `examples/EDM4hep` folder contains examples of using EDM4hep reconstructed
 particles as input to jet reconstruction. See the specific `README.md` file in
 that directory as well as [EDM4hep Inputs](@ref).
+
+## Jet Constituents
+
+The `examples/constituents` folder shows an example of the two mechanisms to
+retrieve jet constituents.
diff --git a/docs/src/index.md b/docs/src/index.md
@@ -88,6 +88,25 @@ sorted_jets = sort!(inclusive_jets(cs::ClusterSequence; ptmin=5.0),
   by=JetReconstruction.energy, rev=true)
 ```
 
+## Jet Constituents
+
+There are two ways to retrieve jet constituents. The first way is just to
+retrieve the *indexes* of the constituent jets. These indexes refer to the
+original collection of particles passed in to the reconstruction.
+
+- [`constituent_indexes`](@ref)
+
+The alternative it to retrieve the actual jets from the reconstruction sequence.
+In this case the returned array contains references to the jet objects (of type
+`T`) used internally in the reconstruction.
+
+- [`constituents`](@ref)
+
+Note that in both these cases the cluster sequence object from the
+reconstruction is required (to avoid circular dependencies and improve memory
+management reconstructed jets do not contain a link back to their cluster
+sequence).
+
 ## References
 
 Although it has been developed further since the CHEP2023 conference, the CHEP
diff --git a/examples/EDM4hep/SimpleRecoEDM4hep.jl b/examples/EDM4hep/SimpleRecoEDM4hep.jl
@@ -6,13 +6,21 @@ using JetReconstruction
 input_file = joinpath("/", "Users", "graemes", "code", "EDM4hepJets", "data",
                       "events_196755633.root")
 reader = RootIO.Reader(input_file)
-events = RootIO.get(reader, "events")
+events = RootIO.get(reader, "events");
 
-evt = events[1]
+evt = events[1];
 
 recps = RootIO.get(reader, evt, "ReconstructedParticles")
 
+# Reconstruct and print the jets
 cs = jet_reconstruct(recps; algorithm = JetAlgorithm.Durham)
-for jet in exclusive_jets(cs; njets = 2, T = EEjet)
+dijets = exclusive_jets(cs; njets = 2, T = EEjet)
+for jet in dijets
     println(jet)
 end
+
+# Get constituents
+for (i, jet) in enumerate(dijets)
+    my_constituent_indexes = constituent_indexes(jet, cs)
+    println("Jet $i constituents: $my_constituent_indexes")
+end
diff --git a/examples/README.md b/examples/README.md
@@ -6,6 +6,9 @@ This directory has a number of example files that show how to used the
 Because of extra dependencies in these scripts, one must use the `Project.toml`
 file in this directory.
 
+Some features are demonstrated in their own subdirectories, in which case use
+the `Project.toml` file in that folder.
+
 ## `jetreco.jl`
 
 This is a basic jet reconstruction example that shows how to call the package to
diff --git a/examples/constituents/Project.toml b/examples/constituents/Project.toml
@@ -0,0 +1,4 @@
+[deps]
+JetReconstruction = "44e8cb2c-dfab-4825-9c70-d4808a591196"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+LorentzVectorHEP = "f612022c-142a-473f-8cfd-a09cf3793c6c"
diff --git a/examples/constituents/README.md b/examples/constituents/README.md
@@ -0,0 +1,7 @@
+# Jet Consitituents
+
+The `jetreco-constituents.jl` example shows how to retrieve jet constituents (by
+index and by reference to jet objects).
+
+The `jetreco-constituents-nb.jl` is the same example, just in Pluto notebook
+form.
diff --git a/examples/constituents/jetreco-constituents-nb.jl b/examples/constituents/jetreco-constituents-nb.jl
@@ -1,5 +1,5 @@
 ### A Pluto.jl notebook ###
-# v0.19.45
+# v0.20.3
 
 using Markdown
 using InteractiveUtils
@@ -8,7 +8,7 @@ using InteractiveUtils
 using Pkg
 
 # ╔═╡ cd974dcf-ab96-4ff2-b76a-e18032343581
-Pkg.activate(".")
+Pkg.activate("..")
 
 # ╔═╡ d25974b4-6531-408f-a0f7-ae7ae4a731d4
 using Revise
@@ -23,8 +23,6 @@ Perform a simple reconstruction example and show how to retrieve constituent jet
 
 # ╔═╡ b16f99a0-31ec-4e8f-99c6-7a6fcb16cbee
 md"As this is running in development, use `Pkg` to activate the local development version of the package and use `Revise` to track code changes.
-
-(Currently you must use the `jet-constituents` branch of `JetReconstruction`.)
 "
 
 # ╔═╡ 79f24ec1-a63e-4e96-bd67-49661125be66
@@ -63,9 +61,18 @@ begin
     end
 end
 
+# ╔═╡ c9ce9c76-82ef-42ff-bb2e-3b3b8085d8bc
+begin
+	my_constituent_indexes = constituent_indexes(pj_jets[1], cluster_seq)
+	println("\nConsitituent indexes for jet number $(event_no): $my_constituent_indexes")
+	for i in my_constituent_indexes
+	    println("  Constituent jet $i: $(events[1][i])")
+	end
+end
+
 # ╔═╡ Cell order:
 # ╟─dff6a188-2cbe-11ef-32d0-73c4c05efad2
-# ╟─b16f99a0-31ec-4e8f-99c6-7a6fcb16cbee
+# ╠═b16f99a0-31ec-4e8f-99c6-7a6fcb16cbee
 # ╠═f3a6edec-9d40-4044-89bc-4ff1656f634f
 # ╠═cd974dcf-ab96-4ff2-b76a-e18032343581
 # ╠═d25974b4-6531-408f-a0f7-ae7ae4a731d4
@@ -80,3 +87,4 @@ end
 # ╟─0bd764f9-d427-43fc-8342-603b6759ec8f
 # ╠═46a64c6f-51d7-4083-a953-ecc76882f21e
 # ╠═300879ca-b53d-40b3-864a-1d46f2094123
+# ╠═c9ce9c76-82ef-42ff-bb2e-3b3b8085d8bc
diff --git a/examples/constituents/jetreco-constituents.jl b/examples/constituents/jetreco-constituents.jl
@@ -1,8 +1,6 @@
 # # Jet Reconstruction Constituents Example
 #
 # Perform a simple reconstruction example and show how to retrieve constituent jets.
-#
-# N.B. currently you must use the `jet-constituents` branch of `JetReconstruction`.
 using JetReconstruction
 using LorentzVectorHEP
 using Logging
@@ -30,8 +28,9 @@ for c in my_constituents
     println(" $c")
 end
 
-# Now show how to convert to LorentzVectorCyl:
-println("\nConstituents of jet number $(event_no) as LorentzVectorCyl:")
-for c in my_constituents
-    println(" $(LorentzVectorCyl(JetReconstruction.pt(c), JetReconstruction.rapidity(c), JetReconstruction.phi(c), JetReconstruction.mass(c)))")
+# Just retrieve the indexes of the constituents
+my_constituent_indexes = constituent_indexes(pj_jets[1], cluster_seq)
+println("\nConsitituent indexes for jet number $(event_no): $my_constituent_indexes")
+for i in my_constituent_indexes
+    println("  Constituent jet $i: $(events[1][i])")
 end
diff --git a/src/ClusterSequence.jl b/src/ClusterSequence.jl
@@ -567,25 +567,42 @@ function reco_state(cs::ClusterSequence, ranks; iteration = 0, ignore_beam_merge
 end
 
 """
-    constituents(j::PseudoJet, cs::ClusterSequence)
+    constituents(jet::T, cs::ClusterSequence{T}) where T <: FourMomentum
 
 Get the constituents of a given jet in a cluster sequence.
 
 # Arguments
-- `cs::ClusterSequence`: The cluster sequence object.
-- `j::PseudoJet`: The jet for which to retrieve the constituents.
+- `cs::ClusterSequence{T}`: The cluster sequence object.
+- `jet::T`: The jet for which to retrieve the constituents.
 
 # Returns
-An array of `PseudoJet` objects representing the constituents of the given jet.
-(That is, the original clusters that were recombined to form this jet.)
+An array of jet objects (which are of the same type as the input jet)
+representing the constituents of the given jet,  
 
 """
-function constituents(j::PseudoJet, cs::ClusterSequence)
-    constituent_indexes = get_all_ancestors(cs.history[j._cluster_hist_index].jetp_index,
-                                            cs)
-    constituents = Vector{PseudoJet}()
-    for idx in constituent_indexes
+function constituents(jet::T, cs::ClusterSequence{T}) where {T <: FourMomentum}
+    constituent_idxs = constituent_indexes(jet, cs)
+    constituents = Vector{T}()
+    for idx in constituent_idxs
         push!(constituents, cs.jets[idx])
     end
     constituents
 end
+
+"""
+    constituent_indexes(jet::T, cs::ClusterSequence{T}) where T <: FourMomentum
+
+Return the indexes of the original particles which are the constituents of the
+given jet.
+
+# Arguments
+- `jet::T`: The jet for which to retrieve the constituents.
+- `cs::ClusterSequence{T}`: The cluster sequence object.
+
+# Returns
+
+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
diff --git a/src/JetReconstruction.jl b/src/JetReconstruction.jl
@@ -47,7 +47,8 @@ export RecoStrategy, JetAlgorithm
 
 # ClusterSequence type
 include("ClusterSequence.jl")
-export ClusterSequence, inclusive_jets, exclusive_jets, n_exclusive_jets, constituents
+export ClusterSequence, inclusive_jets, exclusive_jets, n_exclusive_jets, constituents,
+       constituent_indexes
 
 ## N2Plain algorithm
 # Algorithmic part for simple sequential implementation
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -45,6 +45,9 @@ function main()
     do_test_compare_types(RecoStrategy.N2Plain, algname = pp_algorithms[-1], power = -1)
     do_test_compare_types(RecoStrategy.N2Tiled, algname = pp_algorithms[-1], power = -1)
 
+    # Check jet constituents
+    include("test-constituents.jl")
+
     # Suppress these tests for now, as the examples Project.toml is rather heavy
     # because of the GLMakie dependency, plus on a CI there is no GL subsystem,
     # so things fail. The examples should be restructured to have a cleaner set
diff --git a/test/test-constituents.jl b/test/test-constituents.jl
@@ -0,0 +1,44 @@
+# Tests of jet constituent retrieval
+
+include("common.jl")
+
+# PseudoJet comparison test
+function Base.isapprox(j1::PseudoJet, j2::PseudoJet)
+    isapprox(j1.E, j2.E) && isapprox(j1.px, j2.px) &&
+        isapprox(j1.py, j2.py) && isapprox(j1.pz, j2.pz)
+end
+
+# Expected constituent indexes
+const expected_constituent_indexes = [84, 85, 139, 86, 133, 74, 79, 124, 76, 75, 163]
+
+input_file = joinpath(dirname(pathof(JetReconstruction)), "..", "test", "data",
+                      "events.pp13TeV.hepmc3.gz")
+events = read_final_state_particles(input_file)
+
+# Event to pick
+event_no = 1
+
+cluster_seq = jet_reconstruct(events[event_no], p = 1, R = 1.0)
+
+# Retrieve the exclusive pj_jets, but as `PseudoJet` types
+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)
+        @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)
+        @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
