Fixes for pp reconstruction to use typed PseudoJets

It is very important to pass the fully typed PseudoJet{T} around, not a generic PseudoJet

Author: graeme-a-stewart

src/PlainAlgo.jl

@@ -230,13 +230,14 @@ function plain_jet_reconstruct(particles::AbstractArray{T, 1}; p::Union{Real, No
     # Integer p if possible
     p = (round(p) == p) ? Int(p) : p
 
-    if T == PseudoJet
+    if T isa PseudoJet
         # recombination_particles will become part of the cluster sequence, so size it for
         # the starting particles and all N recombinations
         recombination_particles = copy(particles)
         sizehint!(recombination_particles, length(particles) * 2)
     else
-        recombination_particles = PseudoJet[]
+        ParticleType = typeof(px(particles[1]))
+        recombination_particles = PseudoJet{ParticleType}[]
         sizehint!(recombination_particles, length(particles) * 2)
         for i in eachindex(particles)
             push!(recombination_particles,
@@ -278,23 +279,26 @@ generalised k_t algorithm.
 - `clusterseq`: The resulting `ClusterSequence` object representing the
   reconstructed jets.
 """
-function _plain_jet_reconstruct(; particles::Vector{PseudoJet}, p = -1, R = 1.0,
+function _plain_jet_reconstruct(; particles::Vector{PseudoJet{T}}, p = -1, R = 1.0,
                                 algorithm::JetAlgorithm.Algorithm = JetAlgorithm.AntiKt,
-                                recombine = +)
+                                recombine = +) where T <: Real
     # Bounds
     N::Int = length(particles)
     # Parameters
     R2 = R^2
 
+    # Numerical type?
+    ParticleType = typeof(particles[1].E)
+
     # Optimised compact arrays for determining the next merge step
     # We make sure these arrays are type stable - have seen issues where, depending on the values
     # returned by the methods, they can become unstable and performance degrades
-    kt2_array::Vector{Float64} = pt2.(particles) .^ p
-    phi_array::Vector{Float64} = phi.(particles)
-    rapidity_array::Vector{Float64} = rapidity.(particles)
+    kt2_array::Vector{ParticleType} = pt2.(particles) .^ p
+    phi_array::Vector{ParticleType} = phi.(particles)
+    rapidity_array::Vector{ParticleType} = rapidity.(particles)
     nn::Vector{Int} = Vector(1:N) # nearest neighbours
-    nndist::Vector{Float64} = fill(float(R2), N) # geometric distances to the nearest neighbour
-    nndij::Vector{Float64} = zeros(N) # dij metric distance
+    nndist::Vector{ParticleType} = fill(float(R2), N) # geometric distances to the nearest neighbour
+    nndij::Vector{ParticleType} = zeros(N) # dij metric distance
 
     # Maps index from the compact array to the clusterseq jet vector
     clusterseq_index::Vector{Int} = collect(1:N)
@@ -304,7 +308,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(algorithm, p, R, RecoStrategy.N2Plain, particles, history,
+    clusterseq = ClusterSequence{PseudoJet{ParticleType}}(algorithm, p, R, RecoStrategy.N2Plain, particles, history,
                                  Qtot)
 
     # Initialize nearest neighbours

src/Pseudojet.jl

@@ -69,10 +69,16 @@ history index.
 A `PseudoJet` object.
 """
 PseudoJet(px::T, py::T, pz::T, E::T,
-_cluster_hist_index::Int,
-pt2::T) where {T <: Real} = PseudoJet{T}(px,
-                                         py, pz, E, _cluster_hist_index,
-                                         pt2, 1.0 / pt2, _invalid_rap, _invalid_phi)
+    _cluster_hist_index::Int,
+    pt2::T) where {T <: Real} = PseudoJet{T}(px,
+                                            py, pz, E, _cluster_hist_index,
+                                             pt2, 1.0 / pt2, _invalid_rap, _invalid_phi)
+
+PseudoJet{T}(px::T, py::T, pz::T, E::T,
+            _cluster_hist_index::Int,
+            pt2::T) where {T <: Real} = PseudoJet{T}(px,
+                                                    py, pz, E, _cluster_hist_index,
+                                                    pt2, 1.0 / pt2, _invalid_rap, _invalid_phi)
 
 """
     PseudoJet(px::T, py::T, pz::T, E::T) where T <: Real
@@ -90,6 +96,9 @@ A PseudoJet object.
 """
 PseudoJet(px::T, py::T,
 pz::T, E::T) where {T <: Real} = PseudoJet(px, py, pz, E, 0, px^2 + py^2)
+PseudoJet{T}(px::T, py::T,
+pz::T, E::T) where {T <: Real} = PseudoJet{T}(px, py, pz, E, 0, px^2 + py^2)
+
 
 import Base.show
 """

src/TiledAlgoLL.jl

@@ -374,17 +374,18 @@ function tiled_jet_reconstruct(particles::AbstractArray{T, 1}; p::Union{Real, No
     (p, algorithm) = get_algorithm_power_consistency(p = p, algorithm = algorithm)
 
     # If we have PseudoJets, we can just call the main algorithm...
-    if T == PseudoJet
+    if T isa PseudoJet
         # recombination_particles will become part of the cluster sequence, so size it for
         # the starting particles and all N recombinations
         recombination_particles = copy(particles)
         sizehint!(recombination_particles, length(particles) * 2)
     else
-        recombination_particles = PseudoJet[]
+        ParticleType = typeof(px(particles[1]))
+        recombination_particles = PseudoJet{ParticleType}[]
         sizehint!(recombination_particles, length(particles) * 2)
         for i in eachindex(particles)
             push!(recombination_particles,
-                  PseudoJet(px(particles[i]), py(particles[i]), pz(particles[i]),
+                  PseudoJet{ParticleType}(px(particles[i]), py(particles[i]), pz(particles[i]),
                             energy(particles[i])))
         end
     end
@@ -421,9 +422,9 @@ of data types are done.
 tiled_jet_reconstruct(particles::Vector{PseudoJet}; p = 1, R = 1.0, recombine = +)
 ```
 """
-function _tiled_jet_reconstruct(particles::Vector{PseudoJet}; p::Real = -1, R = 1.0,
+function _tiled_jet_reconstruct(particles::Vector{PseudoJet{T}}; p::Real = -1, R = 1.0,
                                 algorithm::JetAlgorithm.Algorithm = JetAlgorithm.AntiKt,
-                                recombine = +)
+                                recombine = +) where T <: Real
     # Bounds
     N::Int = length(particles)
 
@@ -439,13 +440,16 @@ function _tiled_jet_reconstruct(particles::Vector{PseudoJet}; p::Real = -1, R =
     R2::Float64 = R * R
     p = (round(p) == p) ? Int(p) : p # integer p if possible
 
+    # Numerical type?
+    ParticleType = typeof(particles[1].E)
+
     # This will be used quite deep inside loops, so declare it here so that
     # memory (de)allocation gets done only once
     tile_union = Vector{Int}(undef, 3 * _n_tile_neighbours)
 
     # Container for pseudojets, sized for all initial particles, plus all of the
     # merged jets that can be created during reconstruction
-    jets = PseudoJet[]
+    jets = PseudoJet{ParticleType}[]
     sizehint!(jets, N * 2)
     resize!(jets, N)
 
@@ -456,15 +460,15 @@ function _tiled_jet_reconstruct(particles::Vector{PseudoJet}; p::Real = -1, R =
     history, Qtot = initial_history(jets)
 
     # Now get the tiling setup
-    _eta = Vector{Float64}(undef, length(particles))
+    _eta = Vector{ParticleType}(undef, length(particles))
     for ijet in 1:length(particles)
         _eta[ijet] = rapidity(particles[ijet])
     end
 
     tiling = Tiling(setup_tiling(_eta, R))
 
     # ClusterSequence is the struct that holds the state of the reconstruction
-    clusterseq = ClusterSequence(algorithm, p, R, RecoStrategy.N2Tiled, jets, history, Qtot)
+    clusterseq = ClusterSequence{PseudoJet{ParticleType}}(algorithm, p, R, RecoStrategy.N2Tiled, jets, history, Qtot)
 
     # Tiled jets is a structure that has additional variables for tracking which tile a jet is in
     tiledjets = similar(clusterseq.jets, TiledJet)