Make jet structs immutable and improve jet struct methods (#146)

This PR makes the jet structures used in the package immutable. This is relatively straight forward, but some care is required with the _cluster_hist_index, which is the link from the jet to the correct point in the reconstruction history. This now has to be set correctly at the time of object construction. This is addressed in a few places:

For constructing an initial vector of jets, it's necessary to pre-set the cluster history to 1...N, so the initial jets are correctly lined up
This is done in the reconstruction entry points automatically if conversion to the internal EDM is happening
When reading from HepMC3 files to a FourMomentum struct, this is done
During the reconstruction it is easy to calculate the correct point in the history. This requires the recombine function to be restructured to have the signature addjets(jet1, jet2; cluster_hist_index=index), which is not the same a simple +. However, it's not terrible and this is added to the documentation. Support is added for known schemes, currently massless Pt and Pt^2.

The bells-and-whistles reconstruction example, instrumented-jetreco.jl, can now take an --recombine option to use non-standard recombination.

Note that all of the jet struct constructors have been reworked to provide a more consistent interface and cluster_hist_index is always a named parameter. This PR also restructures the accessor/auxiliary methods for all the jet types, making these as generic as possible. To this end a new source file is added, CommonJetStructs.jl that factorises all of the generic type definitions and methods. Constructing jet types from LorentzVectors and LorentzVectorHEPs is made more consistent.

Core reconstruction times are slightly improved with these static structs, especially on x86_64.

For substructure functions, it was necessary to redo the recluster() method, to make copies of the relevant jets, with updated cluster indexes for the reclustering. This is actually a bug fix as the old way was updating the cluster indexes for the reclustering, which would break the reconstruction history for the original reconstruction.

Docstrings were improved - adding missing ones and simplifying a bit others.

Author: graeme-a-stewart

README.md

@@ -27,7 +27,7 @@ algorithm and generalised $`k_\text{T}`$ for $`e^+e^-`$.
 The simplest interface is to call:
 
 ```julia
-cs = jet_reconstruct(particles::AbstractVector{T}; algorithm = JetAlgorithm.AntiKt, R = 1.0, [p = -1,] [recombine = +,] [strategy = RecoStrategy.Best])
+cs = jet_reconstruct(particles::AbstractVector{T}; algorithm = JetAlgorithm.AntiKt, R = 1.0, [p = -1,] [strategy = RecoStrategy.Best])
 ```
 
 - `particles` - a one dimensional array (vector) of input particles for the clustering
@@ -42,7 +42,6 @@ cs = jet_reconstruct(particles::AbstractVector{T}; algorithm = JetAlgorithm.Anti
   - `JetAlgorithm.Durham` the $e^+e-$ $k_\text{T}$ algorithm, also known as the Durham algorithm
   - `JetAlgorithm.EEKt` the $e^+e-$ generalised $k_\text{T}$ algorithm
 - `R` - the cone size parameter; no particles more geometrically distance than `R` will be merged (default 1.0; note this parameter is ignored for the Durham algorithm)
-- `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 `RecoStrategy.Best`)
 
 The object returned is a `ClusterSequence`, which internally tracks all merge steps.
@@ -90,7 +89,7 @@ Another option, if one wishes to use a specific strategy, is to call that strate
 
 ```julia
 # For N2Plain strategy called directly
-plain_jet_reconstruct(particles::AbstractVector{T}; algorithm = JetAlgorithm.AntiKt, R = 1.0, recombine = +)
+plain_jet_reconstruct(particles::AbstractVector{T}; algorithm = JetAlgorithm.AntiKt, R = 1.0)
 ```
 
 Note that there is no `strategy` option in these interfaces.

docs/make.jl

@@ -18,6 +18,7 @@ makedocs(sitename = "JetReconstruction.jl",
              "Substructure" => "substructure.md",
              "Jet Helpers" => "helpers.md",
              "EDM4hep" => "EDM4hep.md",
+             "Recombination Schemes" => "recombination.md",
              "Visualisation" => "visualisation.md",
              "Contributing" => "contributing.md",
              "Reference Docs" => Any["Public API" => "lib/public.md",

docs/src/index.md

@@ -27,7 +27,7 @@ or with some of the optional arguments,
 
 ```julia
 jet_reconstruct(particles; algorithm = JetAlgorithm.GenKt, R = 0.4, 
-                p = 0.5, recombine = +, strategy = RecoStrategy.Best)
+                p = 0.5, recombine = addjets, strategy = RecoStrategy.Best)
 ```
 
 Where `particles` is a collection of 4-vector objects (see [Input Particle
@@ -40,6 +40,8 @@ algorithm (`GenKt`, `EEKt`) and `p` are needed.
 The `R` value determines the cone size; in the case of the Durham algorithm the
 `R` value is ignored.
 
+For a discussion of the `recombine` function, see [Jet Recombination](@ref).
+
 The object returned is a [`ClusterSequence`](@ref), which internally tracks all
 merge steps and is used for [Inclusive and Exclusive Selections](@ref).
 

docs/src/recombination.md

@@ -0,0 +1,134 @@
+# Jet Recombination
+
+When two jets are merged different strategies can be adopted to produce the merged jet.
+
+There are two functions to support this, which can be set by the user and passed
+as a parameters to the reconstruction algorithms. One function gives the
+necessary *preprocessing* for input particles, e.g., setting particles to be
+massless. The other controls the actual *recombination* of two particles into a
+merged jet.
+
+These functions are passed as the `preprocess` and `recombine` parameters to the
+reconstruction interfaces.
+
+## Default - Four Vector Addition
+
+The default for jet merging is simply four momentum addition, that is:
+
+``
+(\mathbf{p}_m, E_m) = (\mathbf{p_1} + \mathbf{p_2}, E_1 + E_2)
+``
+
+This is defined as the [`addjets`](@ref) function in the package, which also
+serves as an example of how the recombination functions are written.
+
+In this case, no preprocessing of particles is required and the default value of
+`preprocess = nothing` signals this.
+
+### Different Recombination Schemes
+
+Two additional recombination schemes are directly supported, the ``p_T`` and
+``p_T^2`` schemes. In these schemes the recombined jet is created to be
+*massless*, i.e., the mass is set to the 3-momentum. The transverse momentum is
+the sum of the two parent jets and the rapidity (``y``) and phi (``\phi``)
+values are weighted averages, by ``p_T`` or ``p_T^2``, of the parent jets.
+
+- `recombine =` [`addjets_ptscheme`](@ref)
+- `recombine =` [`addjets_pt2scheme`](@ref)
+
+In this case the input particles must be rescaled to be massless, setting the
+energy equal to the (three) momentum sum.
+
+- `preprocess =` [`preprocess_ptscheme`](@ref)
+- `preprocess =` [`preprocess_pt2scheme`](@ref)
+
+(In fact `preprocess_pt2scheme` is just an alias for `preprocess_ptscheme` as
+the rescaling is identical.)
+
+### Named Recombination Schemes
+
+To simplify the usage of different recombination schemes supported directly,
+there is a defined enum (scoped, using `EnumX`) for each one:
+`RecombinationScheme.SCHEME`.
+
+This enum is then used with the `RecombinationMethods` dictionary to
+obtain a named tuple in which `recombine` and `preprocess` are set, which can
+then be splatted into the [`jet_reconstruct`](@ref) interface:
+
+```julia
+myscheme = RecombinationMethods[RecombinationScheme.PtScheme]
+jet_reconstruct(event; R = distance, p = p, algorithm = algorithm,
+                                 strategy = strategy, myscheme...)
+```
+
+The supported values in the enum are:
+
+| Scheme | Implements |
+|---|---|
+| `EScheme` | Default 4-momentum addition |
+| `PtScheme` | Massless weighted average of momentum |
+| `Pt2Scheme` | Massless weighted average of momentum squared |
+
+(Should other schemes prove to be particularly desired they can be implemented
+on request.)
+
+## User Defined Recombination
+
+### Preprocessing
+
+The user must supply, if needed, a preprocessing function, which accepts an
+input particle and returns the rescaled particle. This function must accept a
+named argument `cluster_hist_index` to pass to the constructor of the resulting
+particle.
+
+```julia
+user_preprocess(jet::T; cluster_hist_index) -> T
+```
+
+An example of a preprocessing function is [`preprocess_ptscheme`](@ref).
+
+### Recombination
+
+If a different merging scheme is desired then a method must be defined
+that implements the following interface:
+
+```julia
+user_recombine(jet1::T, jet2::T; cluster_hist_index::Int) where {T <: FourMomentum} -> T
+```
+
+i.e., three arguments are needed, the two parent jets and the named argument
+`cluster_hist_index`, which is needed to identify the jet in the reconstruction
+sequence.
+
+It is recommended to use the constructor signature for the output jet of:
+
+```julia
+T(px, py, pz, E; cluster_hist_index = cluster_hist_index)
+```
+
+Where `px`, `py`, `pz` and `E` have been calculated from the inputs `jet1` and
+`jet2` as desired.
+
+However, if working in ``(p_T, y, \phi, m)`` space, use the alternative constructor
+with named parameters:
+
+```julia
+T(;pt=pt, rap=rap, phi=phi, m=m, cluster_hist_index=cluster_hist_index)
+```
+
+(Note that there is a default of `m=0.0`, which is used for massless
+recombination.)
+
+The user function should not modify the `cluster_hist_index`, but must pass in
+to the new jet's constructor to ensure that the resulting reconstruction
+[`ClusterSequence`](@ref) is valid. The recombination functions defined in the
+package serve as examples: [`addjets_ptscheme`](@ref).
+
+### Using an Custom Recombination Method
+
+To use a non-default recombination method, simply pass the recombination method
+to the [`jet_reconstruct`](@ref) entry point as the `recombine` parameter and
+the preprocessing method as `preprocess`.
+
+A very convenient way to do this is to bind these functions into a named tuple
+and splat the tuple into the arguments for the reconstruction.

examples/benchmark.sh

@@ -1,11 +1,13 @@
 #! /bin/sh
 #
 # Quick and dirty set of benchmarks for the most important cases
+trials=${1:-16}
+
 echo "pp 14TeV Tiled"
-julia --project instrumented-jetreco.jl --algorithm=AntiKt -R 0.4 ../test/data/events.pp13TeV.hepmc3.gz -S N2Tiled -m 16
+julia --project instrumented-jetreco.jl --algorithm=AntiKt -R 0.4 ../test/data/events.pp13TeV.hepmc3.gz -S N2Tiled -m $trials
 
 echo "pp 14 TeV Plain"
-julia --project instrumented-jetreco.jl --algorithm=AntiKt -R 0.4 ../test/data/events.pp13TeV.hepmc3.gz -S N2Plain -m 16
+julia --project instrumented-jetreco.jl --algorithm=AntiKt -R 0.4 ../test/data/events.pp13TeV.hepmc3.gz -S N2Plain -m $trials
 
 echo "ee H Durham"
-julia --project instrumented-jetreco.jl --algorithm=Durham ../test/data/events.eeH.hepmc3.gz -m 16
+julia --project instrumented-jetreco.jl --algorithm=Durham ../test/data/events.eeH.hepmc3.gz -m $trials

examples/instrumented-jetreco.jl

@@ -33,13 +33,14 @@ flamegraph which is saved to the `profile/profile_subdir` directory.
 """
 function profile_code(events::Vector{Vector{T}}, profile, nsamples; R = 0.4, p = -1,
                       algorithm::JetAlgorithm.Algorithm = JetAlgorithm.AntiKt,
-                      strategy = RecoStrategy.N2Tiled) where {T <:
-                                                              JetReconstruction.FourMomentum}
+                      strategy = RecoStrategy.N2Tiled,
+                      recombine = RecombinationMethods[RecombinationScheme.EScheme]) where {T <:
+                                                                                            JetReconstruction.FourMomentum}
     Profile.init(n = 5 * 10^6, delay = 0.00001)
     function profile_events(events)
         for evt in events
-            jet_reconstruct(evt, R = R, p = p, algorithm = algorithm,
-                            strategy = strategy)
+            jet_reconstruct(evt; R = R, p = p, algorithm = algorithm,
+                            strategy = strategy, recombine...)
         end
     end
     # Do a warm up run first to avoid JIT compilation costs
@@ -88,12 +89,14 @@ function allocation_stats(events::Vector{Vector{T}}; distance::Real = 0.4,
                           p::Union{Real, Nothing} = nothing,
                           algorithm::Union{JetAlgorithm.Algorithm, Nothing} = nothing,
                           strategy::RecoStrategy.Strategy,
+                          recombine = RecombinationMethods[RecombinationScheme.EScheme],
                           ptmin::Real = 5.0) where {T <: JetReconstruction.FourMomentum}
     println("Memory allocation statistics:")
     @timev for event in events
-        _ = inclusive_jets(jet_reconstruct(event, R = distance, p = p,
+        _ = inclusive_jets(jet_reconstruct(event; R = distance, p = p,
                                            algorithm = algorithm,
-                                           strategy = strategy), ptmin = ptmin)
+                                           strategy = strategy, recombine...),
+                           ptmin = ptmin)
     end
     nothing
 end
@@ -125,10 +128,11 @@ function benchmark_jet_reco(events::Vector{Vector{T}};
                             distance::Real = 0.4,
                             algorithm::Union{JetAlgorithm.Algorithm, Nothing} = nothing,
                             p::Union{Real, Nothing} = nothing,
+                            strategy::RecoStrategy.Strategy,
+                            recombine = RecombinationMethods[RecombinationScheme.EScheme],
                             ptmin::Real = 5.0,
                             dcut = nothing,
                             njets = nothing,
-                            strategy::RecoStrategy.Strategy,
                             nsamples::Integer = 1,
                             gcoff::Bool = false,
                             dump::Union{String, Nothing} = nothing,
@@ -155,8 +159,8 @@ function benchmark_jet_reco(events::Vector{Vector{T}};
         gcoff && GC.enable(false)
         t_start = time_ns()
         for (ievt, event) in enumerate(events)
-            cs = jet_reconstruct(event, R = distance, p = p, algorithm = algorithm,
-                                 strategy = strategy)
+            cs = jet_reconstruct(event; R = distance, p = p, algorithm = algorithm,
+                                 strategy = strategy, recombine...)
             if !isnothing(njets)
                 finaljets = exclusive_jets(cs; njets = njets)
             elseif !isnothing(dcut)
@@ -273,6 +277,11 @@ function parse_command_line(args)
         arg_type = RecoStrategy.Strategy
         default = RecoStrategy.Best
 
+        "--recombine"
+        help = """Recombination scheme to use for jet reconstruction: $(join(JetReconstruction.AllRecombinationSchemes, ", "))"""
+        arg_type = RecombinationScheme.Recombine
+        default = RecombinationScheme.EScheme
+
         "--nsamples", "-m"
         help = "Number of measurement points to acquire."
         arg_type = Int
@@ -344,15 +353,19 @@ function main()
     if args[:alloc]
         allocation_stats(events; distance = args[:distance],
                          p = args[:power], algorithm = args[:algorithm],
-                         strategy = args[:strategy], ptmin = args[:ptmin])
+                         strategy = args[:strategy],
+                         recombine = JetReconstruction.RecombinationMethods[args[:recombine]],
+                         ptmin = args[:ptmin])
     elseif !isnothing(args[:profile])
         profile_code(events, args[:profile], args[:nsamples];
                      R = args[:distance], p = args[:power],
-                     algorithm = args[:algorithm], strategy = args[:strategy])
+                     algorithm = args[:algorithm], strategy = args[:strategy],
+                     recombine = JetReconstruction.RecombinationMethods[args[:recombine]])
     else
         benchmark_jet_reco(events, distance = args[:distance], algorithm = args[:algorithm],
                            p = args[:power],
                            strategy = args[:strategy],
+                           recombine = JetReconstruction.RecombinationMethods[args[:recombine]],
                            ptmin = args[:ptmin], dcut = args[:exclusive_dcut],
                            njets = args[:exclusive_njets],
                            nsamples = args[:nsamples], gcoff = args[:gcoff],

examples/parse-options.jl

@@ -24,3 +24,11 @@ function ArgParse.parse_item(E::Type{RecoStrategy.Strategy}, x::AbstractString)
     end
     p
 end
+
+function ArgParse.parse_item(E::Type{RecombinationScheme.Recombine}, x::AbstractString)
+    p = do_enum_parse(E, x)
+    if p === nothing
+        throw(ErrorException("Invalid value for recombination scheme: $(x)"))
+    end
+    p
+end

examples/visualisation/visualise-jets.jl

@@ -66,8 +66,9 @@ function main()
                                                                    maxevents = args[:event],
                                                                    skipevents = args[:event])
 
-    (p, algorithm) = JetReconstruction.get_algorithm_power_consistency(p = args[:power],
-                                                                       algorithm = args[:algorithm])
+    (p,
+    algorithm) = JetReconstruction.get_algorithm_power_consistency(p = args[:power],
+                                                                   algorithm = args[:algorithm])
     cs = jet_reconstruct(events[1], R = args[:distance], p = p, algorithm = algorithm,
                          strategy = args[:strategy])
 

ext/JetVisualisation.jl

@@ -144,21 +144,22 @@ function JetReconstruction.jetsplot(cs::ClusterSequence,
     end
 
     set_theme!(jetreco_theme)
-    fig, ax, plt_obj = Module.meshscatter(jet_plot_points;
-                                          markersize = jet_plot_marker_size,
-                                          marker = jet_plot_marker,
-                                          colormap = colormap,
-                                          color = jet_plot_colours,
-                                          colorrange = (1, 256),
-                                          figure = (size = (700, 600),),
-                                          axis = (type = Axis3, perspectiveness = 0.5,
-                                                  azimuth = 2.7,
-                                                  elevation = 0.5,
-                                                  xlabel = L"\phi", ylabel = L"y",
-                                                  zlabel = L"p_T",
-                                                  limits = (0, 2π, min_rap - 0.5,
-                                                            max_rap + 0.5, 0, max_pt + 10)),
-                                          shading = NoShading)
+    fig, ax,
+    plt_obj = Module.meshscatter(jet_plot_points;
+                                 markersize = jet_plot_marker_size,
+                                 marker = jet_plot_marker,
+                                 colormap = colormap,
+                                 color = jet_plot_colours,
+                                 colorrange = (1, 256),
+                                 figure = (size = (700, 600),),
+                                 axis = (type = Axis3, perspectiveness = 0.5,
+                                         azimuth = 2.7,
+                                         elevation = 0.5,
+                                         xlabel = L"\phi", ylabel = L"y",
+                                         zlabel = L"p_T",
+                                         limits = (0, 2π, min_rap - 0.5,
+                                                   max_rap + 0.5, 0, max_pt + 10)),
+                                 shading = NoShading)
     fig, ax, plt_obj
 end
 

src/AlgorithmStrategyEnums.jl

@@ -150,3 +150,26 @@ Check if the algorithm is a e+e- reconstruction algorithm.
 function is_ee(algorithm::JetAlgorithm.Algorithm)
     return algorithm in [JetAlgorithm.EEKt, JetAlgorithm.Durham]
 end
+
+"""
+    enum RecombinationScheme
+
+An EnumX scoped enumeration representing different recombination schemes that
+are supported directly in the package.
+
+These schemes map to both a `recombine` and a `preprocess` function, which are
+used in the main reconstruction algorithm.
+"""
+@enumx T=Recombine RecombinationScheme EScheme PtScheme Pt2Scheme
+const AllRecombinationSchemes = [String(Symbol(x))
+                                 for x in instances(RecombinationScheme.Recombine)]
+
+# Note it's a bit fragile to have the dictionary and the enum built
+# separately, but it is manageable. There is a test in the CI that
+# checks that all the enums are defined in the dictionary.
+const RecombinationMethods = Dict(RecombinationScheme.EScheme => (recombine = addjets_escheme,
+                                                                  preprocess = nothing),
+                                  RecombinationScheme.PtScheme => (recombine = addjets_ptscheme,
+                                                                   preprocess = preprocess_ptscheme),
+                                  RecombinationScheme.Pt2Scheme => (recombine = addjets_pt2scheme,
+                                                                    preprocess = preprocess_pt2scheme))