Fix significant performance regression (#81)

* Read correct internal EDM type for benchmarking

Use EEjet for e+e-, PseudoJet for pp

* Factorise compact array updates

Keep the main loop code cleaner by factorising
the compact array updates into functions

* Use clearer notation for jet indexes

Resolve the actual jet indexes more explicitly

* Fix ClusterSequence to be type specific

It is vital that ClusterSequence is a fully concrete type.
Change the constructor to be parametrised on the type of the jets
themselves.

* Small docs fix

* Format fixes

Author: graeme-a-stewart

docs/make.jl

@@ -18,7 +18,7 @@ makedocs(sitename = "JetReconstruction.jl",
              "Particle Inputs" => "particles.md",
              "Reconstruction Strategies" => "strategy.md",
              "Reference Docs" => Any["Public API" => "lib/public.md",
-                                "Internal API" => "lib/internal.md"],
+                                     "Internal API" => "lib/internal.md"],
              "Extras" => Any["Serialisation" => "extras/serialisation.md"]
          ])
 

examples/instrumented-jetreco.jl

@@ -67,7 +67,7 @@ happens inside the JetReconstruction package itself.
 Some other ustilities are also supported here, such as profiling and
 serialising the reconstructed jet outputs.
 """
-function jet_process(events::Vector{Vector{PseudoJet}};
+function jet_process(events::Vector{Vector{T}};
                      distance::Real = 0.4,
                      algorithm::Union{JetAlgorithm.Algorithm, Nothing} = nothing,
                      p::Union{Real, Nothing} = nothing,
@@ -80,7 +80,7 @@ function jet_process(events::Vector{Vector{PseudoJet}};
                      profile = nothing,
                      alloc::Bool = false,
                      dump::Union{String, Nothing} = nothing,
-                     dump_cs = false)
+                     dump_cs = false) where {T <: JetReconstruction.FourMomentum}
 
     # If we are dumping the results, setup the JSON structure
     if !isnothing(dump)
@@ -294,9 +294,16 @@ function main()
         logger = ConsoleLogger(stdout, Logging.Warn)
     end
     global_logger(logger)
-    events::Vector{Vector{PseudoJet}} = read_final_state_particles(args[:file],
-                                                                   maxevents = args[:maxevents],
-                                                                   skipevents = args[:skip])
+    # Try to read events into the correct type!
+    if JetReconstruction.is_ee(args[:algorithm])
+        jet_type = EEjet
+    else
+        jet_type = PseudoJet
+    end
+    events::Vector{Vector{jet_type}} = read_final_state_particles(args[:file],
+                                                                  maxevents = args[:maxevents],
+                                                                  skipevents = args[:skip],
+                                                                  T = jet_type)
     if isnothing(args[:algorithm]) && isnothing(args[:power])
         @warn "Neither algorithm nor power specified, defaulting to AntiKt"
         args[:algorithm] = JetAlgorithm.AntiKt

examples/jetreco.jl

@@ -23,15 +23,16 @@ happens inside the JetReconstruction package itself.
 
 Final jets can be serialised if the "dump" option is given
 """
-function jet_process(events::Vector{Vector{PseudoJet}};
+function jet_process(events::Vector{Vector{T}};
                      distance::Real = 0.4,
                      algorithm::Union{JetAlgorithm.Algorithm, Nothing} = nothing,
                      p::Union{Real, Nothing} = nothing,
                      ptmin::Real = 5.0,
                      dcut = nothing,
                      njets = nothing,
                      strategy::RecoStrategy.Strategy,
-                     dump::Union{String, Nothing} = nothing)
+                     dump::Union{String, Nothing} = nothing) where {T <:
+                                                                    JetReconstruction.FourMomentum}
 
     # Set consistent algorithm and power
     (p, algorithm) = JetReconstruction.get_algorithm_power_consistency(p = p,
@@ -138,9 +139,16 @@ function main()
     args = parse_command_line(ARGS)
     logger = ConsoleLogger(stdout, Logging.Info)
     global_logger(logger)
-    events::Vector{Vector{PseudoJet}} = read_final_state_particles(args[:file],
-                                                                   maxevents = args[:maxevents],
-                                                                   skipevents = args[:skip])
+    # Try to read events into the correct type!
+    if JetReconstruction.is_ee(args[:algorithm])
+        jet_type = EEjet
+    else
+        jet_type = PseudoJet
+    end
+    events::Vector{Vector{jet_type}} = read_final_state_particles(args[:file],
+                                                                  maxevents = args[:maxevents],
+                                                                  skipevents = args[:skip],
+                                                                  T = jet_type)
     if isnothing(args[:algorithm]) && isnothing(args[:power])
         @warn "Neither algorithm nor power specified, defaulting to AntiKt"
         args[:algorithm] = JetAlgorithm.AntiKt

src/ClusterSequence.jl

@@ -108,28 +108,28 @@ final jets.
 - `power::Float64`: The power value used for the clustering algorithm (not that
   this value is always stored as a Float64 to be type stable)
 - `R::Float64`: The R parameter used for the clustering algorithm.
-- `jets::Vector{PseudoJet}`: The physical PseudoJets in the cluster sequence.
-  Each PseudoJet corresponds to a position in the history.
+- `jets::Vector{T}`: The actual jets in the cluster sequence, which are of type
+  `T <: FourMomentum`.
 - `n_initial_jets::Int`: The initial number of particles used for exclusive
   jets.
 - `history::Vector{HistoryElement}`: The branching history of the cluster
   sequence. Each stage in the history indicates where to look in the jets vector
   to get the physical PseudoJet.
 - `Qtot::Any`: The total energy of the event.
 """
-struct ClusterSequence
+struct ClusterSequence{T}
     algorithm::JetAlgorithm.Algorithm
     power::Float64
     R::Float64
     strategy::RecoStrategy.Strategy
-    jets::Vector{FourMomentum}
+    jets::Vector{T}
     n_initial_jets::Int
     history::Vector{HistoryElement}
     Qtot::Any
 end
 
 """
-    ClusterSequence(algorithm::JetAlgorithm.Algorithm, p::Real, R::Float64, strategy::RecoStrategy.Strategy, jets, history, Qtot)
+    ClusterSequence(algorithm::JetAlgorithm.Algorithm, p::Real, R::Float64, strategy::RecoStrategy.Strategy, jets::T, history, Qtot) where T <: FourMomentum
 
 Construct a `ClusterSequence` object.
 
@@ -138,13 +138,14 @@ Construct a `ClusterSequence` object.
 - `p::Real`: The power value used for the clustering algorithm.
 - `R::Float64`: The R parameter used for the clustering algorithm.
 - `strategy::RecoStrategy.Strategy`: The strategy used for clustering.
-- `jets::Vector{PseudoJet}`: The physical PseudoJets in the cluster sequence.
+- `jets::Vector{T}`: The jets in the cluster sequence, which are of T <: FourMomentum
 - `history::Vector{HistoryElement}`: The branching history of the cluster
   sequence.
 - `Qtot::Any`: The total energy of the event.
 """
-ClusterSequence(algorithm::JetAlgorithm.Algorithm, p::Real, R::Float64, strategy::RecoStrategy.Strategy, jets, history, Qtot) = begin
-    ClusterSequence(algorithm, Float64(p), R, strategy, jets, length(jets), history, Qtot)
+ClusterSequence(algorithm::JetAlgorithm.Algorithm, p::Real, R::Float64, strategy::RecoStrategy.Strategy, jets::Vector{T}, history, Qtot) where {T <: FourMomentum} = begin
+    ClusterSequence{T}(algorithm, Float64(p), R, strategy, jets, length(jets), history,
+                       Qtot)
 end
 
 """

src/EEAlgorithm.jl

@@ -151,6 +151,45 @@ function ee_check_consistency(clusterseq, eereco, N)
     @debug "Consistency check passed at $msg"
 end
 
+function fill_reco_array!(eereco, particles, R2, p)
+    for i in eachindex(particles)
+        eereco.index[i] = i
+        eereco.nni[i] = 0
+        eereco.nndist[i] = R2
+        # eereco.dijdist[i] = UNDEF # Does not need to be initialised
+        eereco.nx[i] = nx(particles[i])
+        eereco.ny[i] = ny(particles[i])
+        eereco.nz[i] = nz(particles[i])
+        eereco.E2p[i] = energy(particles[i])^(2p)
+    end
+end
+
+@inline function insert_new_jet!(eereco, i, newjet_k, R2, merged_jet, p)
+    eereco.index[i] = newjet_k
+    eereco.nni[i] = 0
+    eereco.nndist[i] = R2
+    eereco.nx[i] = nx(merged_jet)
+    eereco.ny[i] = ny(merged_jet)
+    eereco.nz[i] = nz(merged_jet)
+    eereco.E2p[i] = energy(merged_jet)^(2p)
+end
+
+"""
+    copy_to_slot!(eereco, i, j)
+
+Copy the contents of slot `i` in the `eereco` array to slot `j`.
+"""
+@inline function copy_to_slot!(eereco, i, j)
+    eereco.index[j] = eereco.index[i]
+    eereco.nni[j] = eereco.nni[i]
+    eereco.nndist[j] = eereco.nndist[i]
+    eereco.dijdist[j] = eereco.dijdist[i]
+    eereco.nx[j] = eereco.nx[i]
+    eereco.ny[j] = eereco.ny[i]
+    eereco.nz[j] = eereco.nz[i]
+    eereco.E2p[j] = eereco.E2p[i]
+end
+
 """
     ee_genkt_algorithm(particles::Vector{T}; p = -1, R = 4.0,
                        algorithm::JetAlgorithm.Algorithm = JetAlgorithm.Durham,
@@ -251,16 +290,7 @@ function _ee_genkt_algorithm(; particles::Vector{EEjet}, p = 1, R = 4.0,
     # jet information and populate it accordingly
     # We need N slots for this array
     eereco = StructArray{EERecoJet}(undef, N)
-    for i in eachindex(particles)
-        eereco.index[i] = i
-        eereco.nni[i] = 0
-        eereco.nndist[i] = R2
-        # eereco[i].dijdist = UNDEF # Not needed
-        eereco.nx[i] = nx(particles[i])
-        eereco.ny[i] = ny(particles[i])
-        eereco.nz[i] = nz(particles[i])
-        eereco.E2p[i] = energy(particles[i])^(2p)
-    end
+    fill_reco_array!(eereco, particles, R2, p)
 
     # Setup the initial history and get the total energy
     history, Qtot = initial_history(particles)
@@ -301,13 +331,17 @@ function _ee_genkt_algorithm(; particles::Vector{EEjet}, p = 1, R = 4.0,
                 ijetA, ijetB = ijetB, ijetA
             end
 
+            # Resolve the jet indexes to access the actual jets
+            jetA_idx = eereco[ijetA].index
+            jetB_idx = eereco[ijetB].index
+
             # Source "history" for merge
-            hist_jetA = clusterseq.jets[eereco[ijetA].index]._cluster_hist_index
-            hist_jetB = clusterseq.jets[eereco[ijetB].index]._cluster_hist_index
+            hist_jetA = clusterseq.jets[jetA_idx]._cluster_hist_index
+            hist_jetB = clusterseq.jets[jetB_idx]._cluster_hist_index
 
             # Recombine jetA and jetB into the next jet
-            merged_jet = recombine(clusterseq.jets[eereco[ijetA].index],
-                                   clusterseq.jets[eereco[ijetB].index])
+            merged_jet = recombine(clusterseq.jets[jetA_idx],
+                                   clusterseq.jets[jetB_idx])
             merged_jet._cluster_hist_index = length(clusterseq.history) + 1
 
             # Now add the jet to the sequence, and update the history
@@ -317,26 +351,13 @@ function _ee_genkt_algorithm(; particles::Vector{EEjet}, p = 1, R = 4.0,
                                  newjet_k, dij_min)
 
             # Update the compact arrays, reusing the JetA slot
-            eereco.index[ijetA] = newjet_k
-            eereco.nni[ijetA] = 0
-            eereco.nndist[ijetA] = R2
-            eereco.nx[ijetA] = nx(merged_jet)
-            eereco.ny[ijetA] = ny(merged_jet)
-            eereco.nz[ijetA] = nz(merged_jet)
-            eereco.E2p[ijetA] = energy(merged_jet)^(2p)
+            insert_new_jet!(eereco, ijetA, newjet_k, R2, merged_jet, p)
         end
 
         # Squash step - copy the final jet's compact data into the jetB slot
         # unless we are at the end of the array, in which case do nothing
         if ijetB != N
-            eereco.index[ijetB] = eereco.index[N]
-            eereco.nni[ijetB] = eereco.nni[N]
-            eereco.nndist[ijetB] = eereco.nndist[N]
-            eereco.dijdist[ijetB] = eereco.dijdist[N]
-            eereco.nx[ijetB] = eereco.nx[N]
-            eereco.ny[ijetB] = eereco.ny[N]
-            eereco.nz[ijetB] = eereco.nz[N]
-            eereco.E2p[ijetB] = eereco.E2p[N]
+            copy_to_slot!(eereco, N, ijetB)
         end
 
         # Now number of active jets is decreased by one

src/Utils.jl

@@ -37,7 +37,7 @@ function read_final_state_particles(fname; maxevents = -1, skipevents = 0, T = P
             if p.status == 1
                 # Annoyingly PseudoJet and LorentzVector constructors
                 # disagree on the order of arguments...
-                if T == PseudoJet
+                if T <: FourMomentum
                     particle = T(p.momentum.x, p.momentum.y, p.momentum.z, p.momentum.t)
                 else
                     particle = T(p.momentum.t, p.momentum.x, p.momentum.y, p.momentum.z)