Exclusive jet selections (#47)

* Exclusive jets base on njets

First implementation of exclusive jet finder, for njets

Add a record of initial jets/clusters given to the ClusterSequence.

Add an Enum for the jet algorithm (to be allowed as a CLI option
and also to record this value internally in the ClusterSequence).

* Introduce algorithm enum

Used to store the algorithm in the ClusterSequence

Can also be used to build a simple CLI to select the algorithm

* Refactor enum for strategy name

Call this RecoStrategy instead if JetRecoStrategy

(This is different from JetAlgorithm, where it's how do I define jets,
vs. RecoStrategy, how to I build the reconstructed jets
- seems to make sense to me...)

* Implement n_exclusive_jets(cs, dcut)

This returns the number of jets that would pass a cut
where the reconstruction stopped at a certain dcut
value

* Sort final jets by pt

This makes comparisons a lot easier!

* Correct normalisation of dij values

Had been multiplying by R^2, instead of dividing!

* Add first inclusive selection test

Inclusive kT with dijmax=20

* More exclusive tests

Better structuring for exclusive selection tests
using a struct for all parameters.
Tests for njets and dijcut for both Cambridge/Aachen and
Inclusive-kT.

Author: graeme-a-stewart

README.md

@@ -17,7 +17,7 @@ The algorithms include anti-${k}_\text{T}$, Cambridge/Aachen and inclusive $k_\t
 The simplest interface is to call:
 
 ```julia
-cs = jet_reconstruct(particles::Vector{T}; p = -1, R = 1.0, recombine = +, strategy = JetRecoStrategy.Best)
+cs = jet_reconstruct(particles::Vector{T}; p = -1, R = 1.0, recombine = +, strategy = RecoStrategy.Best)
 ```
 
 - `particles` - a vector of input particles for the clustering
@@ -30,7 +30,7 @@ cs = jet_reconstruct(particles::Vector{T}; p = -1, R = 1.0, recombine = +, strat
   - `1` gives inclusive $k_\text{T}$
 - `R` - the cone size parameter; no particles more geometrically distance than `R` will be merged (default 1.0)
 - `recombine` - the function used to merge two pseudojets (default is a simple 4-vector addition of $`(E, \mathbf{p})`$)
-- `strategy` - the algorithm strategy to adopt, as described below (default `JetRecoStrategy.Best`)
+- `strategy` - the algorithm strategy to adopt, as described below (default `RecoStrategy.Best`)
 
 The object returned is a `ClusterSequence`, which internally tracks all merge steps.
 
@@ -56,9 +56,9 @@ Three strategies are available for the different algorithms:
 
 | Strategy Name | Notes | Interface |
 |---|---|---|
-| `JetRecoStrategy.Best` | Dynamically switch strategy based on input particle density | `jet_reconstruct` |
-| `JetRecoStrategy.N2Plain` | Global matching of particles at each interation (works well for low $N$) | `plain_jet_reconstruct` |
-| `JetRecoStrategy.N2Tiled` | Use tiles of radius $R$ to limit search space (works well for higher $N$) | `tiled_jet_reconstruct` |
+| `RecoStrategy.Best` | Dynamically switch strategy based on input particle density | `jet_reconstruct` |
+| `RecoStrategy.N2Plain` | Global matching of particles at each interation (works well for low $N$) | `plain_jet_reconstruct` |
+| `RecoStrategy.N2Tiled` | Use tiles of radius $R$ to limit search space (works well for higher $N$) | `tiled_jet_reconstruct` |
 
 Generally one can use the `jet_reconstruct` interface, shown above, as the *Best* strategy safely as the overhead is extremely low. That interface supports a `strategy` option to switch to a different option.
 

examples/jetreco.jl

@@ -17,7 +17,7 @@ using JSON
 using LorentzVectorHEP
 using JetReconstruction
 
-function profile_code(profile, jet_reconstruction, events, niters; R=0.4, p=-1, strategy=JetRecoStrategy.N2Tiled)
+function profile_code(profile, jet_reconstruction, events, niters; R=0.4, p=-1, strategy=RecoStrategy.N2Tiled)
 	Profile.init(n = 5*10^6, delay = 0.00001)
 	profile_events(events) = begin
 		for evt in events
@@ -67,10 +67,12 @@ serialising the reconstructed jet outputs.
 """
 function jet_process(
 	events::Vector{Vector{PseudoJet}};
-	ptmin::Float64 = 5.0,
-	distance::Float64 = 0.4,
+	distance::Real = 0.4,
 	power::Integer = -1,
-	strategy::JetRecoStrategy.Strategy,
+	ptmin::Real = 5.0,
+	dcut = nothing,
+	njets = nothing,
+	strategy::RecoStrategy.Strategy,
 	nsamples::Integer = 1,
 	gcoff::Bool = false,
 	profile = nothing,
@@ -114,11 +116,18 @@ function jet_process(
 		gcoff && GC.enable(false)
 		t_start = time_ns()
 		for (ievt, event) in enumerate(events)
-			finaljets = inclusive_jets(jet_reconstruct(event, R = distance, p = power, strategy = strategy), ptmin)
+			if !isnothing(njets)
+				finaljets = exclusive_jets(jet_reconstruct(event, R = distance, p = power, strategy = strategy), njets=njets)
+			elseif !isnothing(dcut)
+				finaljets = exclusive_jets(jet_reconstruct(event, R = distance, p = power, strategy = strategy), dcut=dcut)
+			else
+				finaljets = inclusive_jets(jet_reconstruct(event, R = distance, p = power, strategy = strategy), ptmin)
+			end
 			# Only print the jet content once
 			if irun == 1
 				@info begin
 					jet_output = "Event $(ievt)\n"
+					sort!(finaljets, by = x -> pt(x), rev=true)
 					for (ijet, jet) in enumerate(finaljets)
 						jet_output *= " $(ijet) - $(jet)\n"
 					end
@@ -182,10 +191,18 @@ parse_command_line(args) = begin
 		default = 0
 
 		"--ptmin"
-		help = "Minimum p_t for final jets (GeV)"
+		help = "Minimum p_t for final inclusive jets (energy unit is the same as the input clusters, usually GeV)"
 		arg_type = Float64
 		default = 5.0
 
+		"--exclusive-dcut"
+		help = "Return all exclusive jets where further merging would have d>d_cut"
+		arg_type = Float64
+
+		"--exclusive-njets"
+		help = "Return all exclusive jets once clusterisation has produced n jets"
+		arg_type = Int
+
 		"--distance", "-R"
 		help = "Distance parameter for jet merging"
 		arg_type = Float64
@@ -198,8 +215,8 @@ parse_command_line(args) = begin
 
 		"--strategy"
 		help = """Strategy for the algorithm, valid values: $(join(JetReconstruction.AllJetRecoStrategies, ", "))"""
-		arg_type = JetRecoStrategy.Strategy
-		default = JetRecoStrategy.Best
+		arg_type = RecoStrategy.Strategy
+		default = RecoStrategy.Best
 
 		"--nsamples", "-m"
 		help = "Number of measurement points to acquire."
@@ -238,8 +255,8 @@ parse_command_line(args) = begin
 end
 
 
-function ArgParse.parse_item(::Type{JetRecoStrategy.Strategy}, x::AbstractString)
-	s = tryparse(JetRecoStrategy.Strategy, x)
+function ArgParse.parse_item(::Type{RecoStrategy.Strategy}, x::AbstractString)
+	s = tryparse(RecoStrategy.Strategy, x)
 	if s === nothing
 		throw(ErrorException("Invalid value for strategy: $(x)"))
 	end
@@ -258,8 +275,8 @@ main() = begin
 	global_logger(logger)
 	events::Vector{Vector{PseudoJet}} =
 		read_final_state_particles(args[:file], maxevents = args[:maxevents], skipevents = args[:skip])
-	jet_process(events, ptmin = args[:ptmin], distance = args[:distance], 
-        power = args[:power], strategy = args[:strategy],
+	jet_process(events, distance = args[:distance], power = args[:power], strategy = args[:strategy],
+		ptmin = args[:ptmin], dcut = args[:exclusive_dcut], njets = args[:exclusive_njets],
 		nsamples = args[:nsamples], gcoff = args[:gcoff], profile = args[:profile],
 		alloc = args[:alloc], dump = args[:dump])
 	nothing

src/AlgorithmStrategyEnums.jl

@@ -0,0 +1,23 @@
+# EnumX provides scoped enums, which are nicer
+using EnumX
+
+# Valid strategy enum
+@enumx T = Strategy RecoStrategy Best N2Plain N2Tiled
+const AllJetRecoStrategies = [String(Symbol(x)) for x in instances(RecoStrategy.Strategy)]
+
+# Algorithm emun
+@enumx T = Algorithm JetAlgorithm AntiKt Cambridge Kt EEKt Durham
+const AllJetRecoAlgorithms = [String(Symbol(x)) for x in instances(JetAlgorithm.Algorithm)]
+
+# Map from power values to algorithms
+power2algorithm = Dict(-1 => JetAlgorithm.AntiKt,
+    0 => JetAlgorithm.Cambridge,
+    1 => JetAlgorithm.Kt)
+
+# Map from string to an enum value (used for CLI parsing)
+Base.tryparse(E::Type{<:Enum}, str::String) =
+    let insts = instances(E),
+        p = findfirst(==(Symbol(str)) ∘ Symbol, insts)
+
+        p !== nothing ? insts[p] : nothing
+    end

src/ClusterSequence.jl

@@ -4,13 +4,15 @@
 # ported to this package by Graeme Stewart
 
 using Accessors
+using Logging
 
 # One cannot use an ENUM here as it's a different type
 # I don't know any good way to keep this lean and to have these necessary
 # flags other than using these "magic numbers"
-const Invalid=-3
-const NonexistentParent=-2
-const BeamJet=-1
+# (TODO: this is not runtime critical, so could consider a Union{Int,Emum}?)
+const Invalid = -3
+const NonexistentParent = -2
+const BeamJet = -1
 
 """A struct holding a record of jet mergers and finalisations"""
 struct HistoryElement
@@ -58,24 +60,44 @@ Convienence structure holding all of the relevant parameters for
 the jet reconstruction
 """
 struct ClusterSequence
+    """Algorithm and strategy used"""
+    algorithm::JetAlgorithm.Algorithm
+    strategy::RecoStrategy.Strategy
+    
     """
     This contains the physical PseudoJets; for each PseudoJet one can find
-    the corresponding position in the _history by looking at
-    _jets[i].cluster_hist_index()
+    the corresponding position in the history by looking at
+    jets[i].cluster_hist_index()
     """
     jets::Vector{PseudoJet}
 
+    """
+    Record the initial number of particlesm used for exclusive jets
+    """
+    n_initial_jets::Int
+
     """
     This vector will contain the branching history; for each stage,
-    history[i].jetp_index indicates where to look in the _jets
+    history[i].jetp_index indicates where to look in the jets
     vector to get the physical PseudoJet.
     """
     history::Vector{HistoryElement}
 
     """Total energy of the event"""
-    Qtot
+    Qtot::Any
+end
+
+"""ClusterSequence constructor, where the power value is given"""
+ClusterSequence(p::Int, strategy::RecoStrategy.Strategy, jets, history, Qtot) = begin
+    if !haskey(power2algorithm, p)
+        raise(ArgumentError("Unrecognised algorithm for power value p=$p"))
+    end
+    ClusterSequence(power2algorithm[p], strategy, jets, length(jets), history, Qtot)
 end
 
+"""ClusterSequence constructor, with direct algorithm specified"""
+ClusterSequence(alg::JetAlgorithm.Algorithm, strategy::RecoStrategy.Strategy, jets, history, Qtot) =
+    ClusterSequence(alg, strategy, jets, length(jets), history, Qtot)
 
 """Add a new jet's history into the recombination sequence"""
 add_step_to_history!(clusterseq::ClusterSequence, parent1, parent2, jetp_index, dij) = begin
@@ -117,8 +139,8 @@ end
 
 """Return all inclusive jets of a ClusterSequence with pt > ptmin"""
 function inclusive_jets(clusterseq::ClusterSequence, ptmin = 0.0)
-    dcut = ptmin * ptmin
-    jets_local = Vector{LorentzVectorCyl}(undef, 0)
+    pt2min = ptmin * ptmin
+    jets_local = LorentzVectorCyl[]
     # sizehint!(jets_local, length(clusterseq.jets))
     # For inclusive jets with a plugin algorithm, we make no
     # assumptions about anything (relation of dij to momenta,
@@ -128,10 +150,89 @@ function inclusive_jets(clusterseq::ClusterSequence, ptmin = 0.0)
         elt.parent2 == BeamJet || continue
         iparent_jet = clusterseq.history[elt.parent1].jetp_index
         jet = clusterseq.jets[iparent_jet]
-        if pt2(jet) >= dcut
+        if pt2(jet) >= pt2min
+            @debug "Added inclusive jet index $iparent_jet"
             push!(jets_local, LorentzVectorCyl(pt(jet), rapidity(jet), phi(jet), mass(jet)))
         end
     end
     jets_local
 end
 
+
+"""Return all exclusive jets of a ClusterSequence"""
+function exclusive_jets(clusterseq::ClusterSequence; dcut = nothing, njets = nothing)
+    if isnothing(dcut) && isnothing(njets)
+        throw(ArgumentError("Must pass either a dcut or an njets value"))
+    end
+
+    if !isnothing(dcut)
+        njets = n_exclusive_jets(clusterseq, dcut=dcut)
+    end
+
+    # Check that an algorithm was used that makes sense for exclusive jets
+    if !(clusterseq.algorithm ∈ (JetAlgorithm.Cambridge, JetAlgorithm.Kt, JetAlgorithm.EEKt, JetAlgorithm.Durham))
+        throw(ArgumentError("Algorithm used is not suitable for exclusive jets ($(clusterseq.algorithm))"))
+    end
+
+    # njets search
+    stop_point = 2*clusterseq.n_initial_jets - njets + 1
+
+    # Sanity check - never return more jets than initial particles
+    if stop_point < clusterseq.n_initial_jets
+        stop_point = clusterseq.n_initial_jets
+    end
+
+    # Sanity check - ensure that reconstruction proceeded to the end
+    if 2 * clusterseq.n_initial_jets != length(clusterseq.history)
+        throw(ErrorException("Cluster sequence is incomplete, exclusive jets unavailable"))
+    end
+
+    excl_jets = LorentzVectorCyl[]
+    for j in stop_point:length(clusterseq.history)
+        @info "Search $j ($(clusterseq.history[j].parent1) + $(clusterseq.history[j].parent2))"
+        for parent in (clusterseq.history[j].parent1, clusterseq.history[j].parent2)
+            if (parent < stop_point && parent > 0)
+                @info "Added exclusive jet index $(clusterseq.history[parent].jetp_index)"
+                jet = clusterseq.jets[clusterseq.history[parent].jetp_index]
+                push!(excl_jets, LorentzVectorCyl(pt(jet), rapidity(jet), phi(jet), mass(jet)))
+            end
+        end
+    end
+
+    excl_jets
+end
+
+
+"""Return all number of exclusive jets of a ClusterSequence that are above a certain dcut value"""
+function n_exclusive_jets(clusterseq::ClusterSequence; dcut::AbstractFloat)
+    # Check that an algorithm was used that makes sense for exclusive jets
+    if !(clusterseq.algorithm ∈ (JetAlgorithm.Cambridge, JetAlgorithm.Kt, JetAlgorithm.EEKt, JetAlgorithm.Durham))
+        throw(ArgumentError("Algorithm used is not suitable for exclusive jets ($(clusterseq.algorithm))"))
+    end
+
+    # Locate the point where clustering would have stopped (i.e. the
+    # first time max_dij_so_far > dcut)
+    i_dcut = length(clusterseq.history)
+    for i_history ∈ length(clusterseq.history):-1:1
+        @info "Examining $i_history, max_dij=$(clusterseq.history[i_history].max_dij_so_far)"
+        if clusterseq.history[i_history].max_dij_so_far <= dcut
+            i_dcut = i_history
+            break
+        end
+    end
+
+    # The number of jets is then given by this formula
+    length(clusterseq.history) - i_dcut
+
+    # int i = _history.size() - 1; // last jet
+    # while (i >= 0) {
+    #   if (_history[i].max_dij_so_far <= dcut) {break;}
+    #   i--;
+    # }
+    # int stop_point = i + 1;
+    # // relation between stop_point, njets assumes one extra jet disappears
+    # // at each clustering.
+    # int njets = 2*_initial_n - stop_point;
+    # return njets;
+
+end

src/GenericAlgo.jl

@@ -2,16 +2,16 @@
 # switch based on the strategy, or based on the event density
 # if the "Best" strategy is to be employed
 
-function jet_reconstruct(particles; p = -1, R = 1.0, recombine = +, strategy = JetRecoStrategy.Best)
+function jet_reconstruct(particles; p = -1, R = 1.0, recombine = +, strategy = RecoStrategy.Best)
     # Either map to the fixed algorithm corresponding to the strategy
     # or to an optimal choice based on the density of initial particles
 
-    if strategy == JetRecoStrategy.Best
+    if strategy == RecoStrategy.Best
         # The breakpoint of ~90 is determined empirically on e+e- -> H and 0.5TeV pp -> 5GeV jets
         algorithm = length(particles) > 80 ? tiled_jet_reconstruct : plain_jet_reconstruct
-    elseif strategy == JetRecoStrategy.N2Plain
+    elseif strategy == RecoStrategy.N2Plain
         algorithm = plain_jet_reconstruct
-    elseif strategy == JetRecoStrategy.N2Tiled
+    elseif strategy == RecoStrategy.N2Tiled
         algorithm = tiled_jet_reconstruct
     else
         throw(ErrorException("Invalid strategy: $(strategy)"))

src/JetRecoStrategies.jl

@@ -26,13 +26,13 @@ energy(p::LorentzVectorCyl) = LorentzVectorHEP.energy(p)
 include("Pseudojet.jl")
 export PseudoJet
 
+# Jet reconstruction strategies and algorithms (enums!)
+include("AlgorithmStrategyEnums.jl")
+export RecoStrategy, JetAlgorithm
+
 # ClusterSequence type
 include("ClusterSequence.jl")
-export ClusterSequence, inclusive_jets
-
-# Jet reconstruction strategies
-include("JetRecoStrategies.jl")
-export JetRecoStrategy
+export ClusterSequence, inclusive_jets, exclusive_jets
 
 ## N2Plain algorithm
 # Algorithmic part for simple sequential implementation

src/JetReconstruction.jl

@@ -134,7 +134,7 @@ function _plain_jet_reconstruct(;particles::Vector{PseudoJet}, p = -1, R = 1.0,
     # Current implementation mutates the particles vector, so need to copy it
     # for the cluster sequence (there is too much copying happening, so this
     # needs to be rethought and reoptimised)
-    clusterseq = ClusterSequence(particles, history, Qtot)
+    clusterseq = ClusterSequence(p, RecoStrategy.N2Plain, particles, history, Qtot)
 
     # Initialize nearest neighbours
     @simd for i in 1:N
@@ -155,7 +155,7 @@ function _plain_jet_reconstruct(;particles::Vector{PseudoJet}, p = -1, R = 1.0,
         #@debug "Beginning iteration $iteration"
         # Findmin and add back renormalisation to distance
         dij_min, i = fast_findmin(nndij, N)
-        dij_min *= R2
+        @fastmath dij_min /= R2
         j::Int = nn[i]
         #@debug "Closest compact jets are $i ($(clusterseq_index[i])) and $j ($(clusterseq_index[j]))"