final changes

Author: HarshitNagpal29

src/ClusterSequence.jl

@@ -92,29 +92,17 @@ end
 
 Reset and initialise reusable clustering-history storage for `particles`.
 
-The vector's logical length is reset to the number of initial particles.
-Capacity is retained for the complete clustering history.
+The vector's logical length is reset to the number of initial particles while
+retaining any capacity already owned by the vector.
 
 The returned history vector is borrowed storage owned by the caller.
 """
 function initial_history!(history::Vector{HistoryElement},
                           particles)
     N = length(particles)
 
-    # Remove the previous event's logical history entries.
-    #
-    # This changes length(history) to zero, but retained vector capacity is
-    # preserved.
-    empty!(history)
-
-    # A complete hadron-collider clustering sequence contains:
-    #
-    #   N initial history entries
-    # + N recombination/finalisation entries
-    # = 2N total entries.
-    _sizehint_for_reuse!(history, 2 * N)
-
-    # Establish exactly N active initial-history slots.
+    # Establish exactly N active initial-history slots. All slots are
+    # overwritten below, and `resize!` retains any existing excess capacity.
     resize!(history, N)
 
     Qtot::Float64 = 0.0

src/JetUtils.jl

@@ -138,3 +138,68 @@ function kt_scale(jet1::T, jet2::T) where {T <: FourMomentum}
     pt2 = JetReconstruction.pt(jet2)
     return min(pt1, pt2) * deltar(jet1, jet2)
 end
+
+"""
+    construct_reco_jets(particles, ::Type{J}, preprocess) where {J <: FourMomentum}
+
+Create independently owned, reconstruction-ready jets of type `J` from
+`particles`.
+
+When `preprocess` is `nothing`, inputs that already have type `J` are copied
+directly and other inputs are converted to `J`. Otherwise `preprocess` is called
+for every input particle.
+"""
+function construct_reco_jets(particles::AbstractVector{P},
+                             ::Type{J},
+                             preprocess) where {P, J <: FourMomentum}
+    TargetNumericalType = eltype(particles[1])
+    if TargetNumericalType <: Real
+        TargetJetType = concretize_return_type(J, TargetNumericalType)
+    else
+        TargetJetType = typeof(J(particles[1]))
+    end
+
+    recombination_particles = Vector{TargetJetType}()
+    sizehint!(recombination_particles, 2 * length(particles))
+
+    return construct_reco_jets!(recombination_particles,
+                                particles,
+                                preprocess)
+end
+
+"""
+    construct_reco_jets!(recombination_particles, particles, preprocess)
+
+Reset and fill reusable reconstruction-jet storage. The destination must not
+alias the input collection.
+"""
+function construct_reco_jets!(recombination_particles::Vector{J},
+                              particles::AbstractVector{P},
+                              preprocess) where {P, J <: FourMomentum}
+    Base.mightalias(recombination_particles, particles) &&
+        throw(ArgumentError("reusable jet storage must not alias the input particles"))
+
+    N = length(particles)
+    empty!(recombination_particles)
+    _sizehint_for_reuse!(recombination_particles, 2 * N)
+
+    if isnothing(preprocess)
+        if P === J
+            append!(recombination_particles, particles)
+        else
+            for (i, particle) in enumerate(particles)
+                push!(recombination_particles,
+                      J(particle; cluster_hist_index = i))
+            end
+        end
+    else
+        for (i, particle) in enumerate(particles)
+            push!(recombination_particles,
+                  preprocess(particle,
+                             J;
+                             cluster_hist_index = i))
+        end
+    end
+
+    return recombination_particles
+end

src/TiledAlgoLL.jl

@@ -220,20 +220,6 @@ function set_nearest_neighbours!(tiling::Tiling,
     return nothing
 end
 
-function set_nearest_neighbours!(clusterseq::ClusterSequence,
-                                 tiling::Tiling,
-                                 tiledjets::Vector{TiledJet})
-    NNs = similar(clusterseq.jets, TiledJet)
-    diJ = similar(clusterseq.jets, Float64)
-
-    set_nearest_neighbours!(tiling,
-                            tiledjets,
-                            NNs,
-                            diJ)
-
-    NNs, diJ
-end
-
 """
     do_iB_recombination_step!(clusterseq::ClusterSequence, jet_i, diB)
 
@@ -361,31 +347,9 @@ function tiled_jet_reconstruct(particles::AbstractVector{T};
     # Get consistent algorithm power
     p = get_algorithm_power(p = p, algorithm = algorithm)
 
-    if isnothing(preprocess)
-        if T == PseudoJet
-            # If we don't have a preprocessor, we just need to copy to our own
-            # PseudoJet objects
-            recombination_particles = copy(particles)
-            sizehint!(recombination_particles, length(particles) * 2)
-        else
-            # We assume a constructor for PseudoJet that can ingest the appropriate
-            # type of particle
-            recombination_particles = PseudoJet[]
-            sizehint!(recombination_particles, length(particles) * 2)
-            for (i, particle) in enumerate(particles)
-                push!(recombination_particles, PseudoJet(particle; cluster_hist_index = i))
-            end
-        end
-    else
-        # We have a preprocessor function that we need to call to modify the
-        # input particles
-        recombination_particles = PseudoJet[]
-        sizehint!(recombination_particles, length(particles) * 2)
-        for (i, particle) in enumerate(particles)
-            push!(recombination_particles,
-                  preprocess(particle, PseudoJet; cluster_hist_index = i))
-        end
-    end
+    recombination_particles = construct_reco_jets(particles,
+                                                  PseudoJet,
+                                                  preprocess)
 
     _tiled_jet_reconstruct!(recombination_particles; algorithm = algorithm, p = p, R = R,
                             recombine = recombine)
@@ -454,9 +418,7 @@ function _tiled_jet_reconstruct!(particles::AbstractVector{PseudoJet};
     tiledjets = scratch.tiledjets
 
     if isnothing(history_buffer)
-        # Preserve the owning path's eager complete-history allocation. This is
-        # especially important on Julia < 1.11, where initial_history!'s
-        # no-shrink size hint is intentionally unavailable.
+        # Preserve the owning path's eager complete-history allocation.
         history_buffer = Vector{HistoryElement}(undef, N)
         sizehint!(history_buffer, 2 * N)
     end
@@ -631,9 +593,9 @@ function _n2tiled_reconstruct_with_workspace!(workspace::N2TiledWorkspace,
     resolved_power = get_algorithm_power(p = p,
                                          algorithm = algorithm)
 
-    jets = prepare_recombination_jets!(workspace,
-                                       particles,
-                                       preprocess = preprocess)
+    jets = construct_reco_jets!(workspace.jets,
+                                particles,
+                                preprocess)
 
     return _tiled_jet_reconstruct!(jets;
                                    algorithm = algorithm,

src/TiledAlgoLLStructs.jl

@@ -261,51 +261,6 @@ function TiledScratch()
                         Dict{Tuple{Int, Int}, TilingArrays}())
 end
 
-"""
-Reset and fill reusable recombination-jet storage.
-
-This mirrors the preprocessing behaviour of `tiled_jet_reconstruct`.
-
-When `isnothing(preprocess)`:
-
-- PseudoJet inputs are copied directly;
-- other supported input types are converted to PseudoJet.
-
-Otherwise, the supplied preprocessing function is called for every particle.
-"""
-function _prepare_recombination_jets!(jets::Vector{PseudoJet},
-                                      particles::AbstractVector{T};
-                                      preprocess = preprocess_escheme) where {T}
-    jets === particles &&
-        throw(ArgumentError("reusable jet storage must not alias the input particle vector"))
-
-    N = length(particles)
-
-    empty!(jets)
-    _sizehint_for_reuse!(jets, 2 * N)
-
-    if isnothing(preprocess)
-        if T == PseudoJet
-            append!(jets, particles)
-        else
-            for (i, particle) in enumerate(particles)
-                push!(jets,
-                      PseudoJet(particle;
-                                cluster_hist_index = i))
-            end
-        end
-    else
-        for (i, particle) in enumerate(particles)
-            push!(jets,
-                  preprocess(particle,
-                             PseudoJet;
-                             cluster_hist_index = i))
-        end
-    end
-
-    return jets
-end
-
 """
 Reusable state for one full-semantics N2Tiled reconstruction worker.
 
@@ -359,14 +314,6 @@ function release_n2tiled_workspace_capacity!(workspace::N2TiledWorkspace)
     return nothing
 end
 
-function prepare_recombination_jets!(workspace::N2TiledWorkspace,
-                                     particles::AbstractVector;
-                                     preprocess = preprocess_escheme)
-    return _prepare_recombination_jets!(workspace.jets,
-                                        particles;
-                                        preprocess = preprocess)
-end
-
 """
 Resize the event-sized vectors in `scratch` for an event with `N` particles.
 

src/Utils.jl

@@ -125,3 +125,42 @@ fast_findmin(dij, n) = begin
     end
     dij_min, best
 end
+
+"""
+    concretize_return_type(return_type, numerical_type)
+
+Return `return_type` unchanged when it is already a `DataType`, or apply the
+single type parameter of a `UnionAll` type using `numerical_type`.
+"""
+function concretize_return_type(return_type::Type,
+                                numerical_type::Type{<:Real})
+    return_type isa DataType && return return_type
+
+    if return_type isa UnionAll
+        n_parameters = count_typevars(return_type)
+        n_parameters == 1 ||
+            throw(ArgumentError("Type $return_type requires $n_parameters type parameters, " *
+                                "but only 1 can be inferred"))
+
+        typevar = return_type.var
+        if !(typevar.lb <: numerical_type <: typevar.ub)
+            throw(ArgumentError("Cannot parameterize $return_type with type $numerical_type"))
+        end
+
+        return return_type{numerical_type}
+    end
+
+    throw(ArgumentError("Unexpected type specified: $return_type"))
+end
+
+function count_typevars(type::UnionAll)
+    count = 0
+    current = type
+
+    while current isa UnionAll
+        count += 1
+        current = current.body
+    end
+
+    return count
+end

test/test-jet-utils.jl

@@ -95,3 +95,44 @@ end
         @test finaljets_pj[i] ≈ finaljets_lorentzhep[i] ≈ finaljets_lorentzhepcyl[i]
     end
 end
+
+@testset "Reconstruction-jet construction" begin
+    raw_jets = [PseudoJet(1.0, 2.0, 3.0, 4.0; cluster_hist_index = 11),
+        PseudoJet(2.0, 3.0, 4.0, 6.0; cluster_hist_index = 17)]
+
+    owned_jets = JetReconstruction.construct_reco_jets(raw_jets,
+                                                       PseudoJet,
+                                                       nothing)
+    @test owned_jets == raw_jets
+    @test owned_jets !== raw_jets
+    @test JetReconstruction.cluster_hist_index.(owned_jets) == [11, 17]
+
+    lorentz_particles = lorentzvector.(raw_jets)
+    converted_jets = JetReconstruction.construct_reco_jets(lorentz_particles,
+                                                           PseudoJet,
+                                                           nothing)
+    @test JetReconstruction.cluster_hist_index.(converted_jets) == [1, 2]
+    @test lorentzvector.(converted_jets) == lorentz_particles
+
+    preprocessed_jets = JetReconstruction.construct_reco_jets(raw_jets,
+                                                              PseudoJet,
+                                                              preprocess_ptscheme)
+    @test JetReconstruction.cluster_hist_index.(preprocessed_jets) == [1, 2]
+
+    reusable_jets = PseudoJet[]
+    returned_jets = JetReconstruction.construct_reco_jets!(reusable_jets,
+                                                           raw_jets,
+                                                           nothing)
+    @test returned_jets === reusable_jets
+    @test returned_jets == raw_jets
+
+    JetReconstruction.construct_reco_jets!(reusable_jets,
+                                           raw_jets[1:1],
+                                           nothing)
+    @test length(reusable_jets) == 1
+    @test reusable_jets[1] == raw_jets[1]
+
+    @test_throws ArgumentError JetReconstruction.construct_reco_jets!(raw_jets,
+                                                                      raw_jets,
+                                                                      nothing)
+end

test/test-n2tiled-kernel.jl

@@ -49,8 +49,7 @@ end
     @testset "Exact owning/workspace equivalence" begin
         workspace = N2TiledWorkspace()
 
-        cases = ((PseudoJet[], JetAlgorithm.AntiKt, nothing, 0.4),
-                 (small_event[1:1], JetAlgorithm.AntiKt, nothing, 0.8),
+        cases = ((small_event[1:1], JetAlgorithm.AntiKt, nothing, 0.8),
                  (small_event, JetAlgorithm.AntiKt, nothing, 0.4),
                  (large_event, JetAlgorithm.AntiKt, nothing, 0.2),
                  (small_event, JetAlgorithm.Kt, nothing, 0.4),
@@ -65,19 +64,20 @@ end
                                              preprocess = nothing)
             expected_inclusive = inclusive_jets(expected; ptmin = 5.0)
 
-            with_n2tiled_reconstruction(workspace,
-                                        event;
-                                        algorithm = algorithm,
-                                        p = power,
-                                        R = distance,
-                                        preprocess = nothing) do actual
+            output = with_n2tiled_reconstruction(workspace,
+                                                 event;
+                                                 algorithm = algorithm,
+                                                 p = power,
+                                                 R = distance,
+                                                 preprocess = nothing) do actual
                 @test actual.jets === workspace.jets
                 @test actual.history === workspace.history
+                @test length(actual.history) == 2 * length(event)
                 test_exact_clustersequence_equality(actual, expected)
 
-                output = inclusive_jets(actual; ptmin = 5.0)
-                @test output == expected_inclusive
+                return inclusive_jets(actual; ptmin = 5.0)
             end
+            @test output == expected_inclusive
         end
     end