Fix type instability when passing type to a function, concretize default types (#187)

m-fila · web-flow · commit c4ad094f3791 · 2025-07-20T09:31:31.000+02:00
* fix type instability in inclusive and exclusive jets

* fix type instability in read_final_state_particles

* concretize default vector types

* use same convention in recombination schemes, fix docstrings and return type
diff --git a/README.md b/README.md
@@ -53,7 +53,7 @@ To obtain the final inclusive jets, use the `inclusive_jets` method:
 final_jets = inclusive_jets(cs::ClusterSequence; ptmin=0.0)
 ```
 
-Only jets passing the cut $p_T > p_{Tmin}$ will be returned. The result is returned as a `Vector{LorentzVectorHEP}`, but different return types can be specified (e.g., `T = EEJet`).
+Only jets passing the cut $p_T > p_{Tmin}$ will be returned. The result is returned as a `Vector{LorentzVectorHEP}`, but different return types can be specified (e.g., `inclusive_jets(cs::ClusterSequence, EEJet; ptmin=0.0)`).
 
 #### Sorting
 
diff --git a/docs/src/softkiller.md b/docs/src/softkiller.md
@@ -33,7 +33,7 @@ A typical workflow using SoftKiller is:
 using JetReconstruction
 
 # Read event particles (see examples/softkiller/softkiller_plots.jl)
-particles = read_final_state_particles("event.hepmc3", T=PseudoJet)[1]
+particles = read_final_state_particles("event.hepmc3", PseudoJet)[1]
 
 # Set up SoftKiller
 sk = SoftKiller(rapmax=5.0, grid_size=0.4)
diff --git a/examples/EDM4hep/EDM4hepJets.jl b/examples/EDM4hep/EDM4hepJets.jl
@@ -56,7 +56,7 @@ function main()
         if args[:printjets]
             @info begin
                 jets = "Event $(ievt)\n"
-                for jet in exclusive_jets(cs; njets = args[:njets], T = EEJet)
+                for jet in exclusive_jets(cs, EEJet; njets = args[:njets])
                     jets *= " $jet\n"
                 end
                 jets
diff --git a/examples/EDM4hep/SimpleRecoEDM4hep.jl b/examples/EDM4hep/SimpleRecoEDM4hep.jl
@@ -14,7 +14,7 @@ recps = RootIO.get(reader, evt, "ReconstructedParticles")
 
 # Reconstruct and print the jets
 cs = jet_reconstruct(recps; algorithm = JetAlgorithm.Durham)
-dijets = exclusive_jets(cs; njets = 2, T = EEJet)
+dijets = exclusive_jets(cs, EEJet; njets = 2)
 for jet in dijets
     println(jet)
 end
diff --git a/examples/constituents/jetreco-constituents-nb.jl b/examples/constituents/jetreco-constituents-nb.jl
@@ -32,7 +32,7 @@ cluster_seq = jet_reconstruct(events[event_no]; algorithm = JetAlgorithm.Kt, R =
 md"Retrieve the exclusive pj_jets, but as `PseudoJet` types"
 
 # ╔═╡ 0d8d4664-915f-4f28-9d5a-6e03cb8d7d8b
-pj_jets = inclusive_jets(cluster_seq; ptmin = 5.0, T = PseudoJet)
+pj_jets = inclusive_jets(cluster_seq, PseudoJet; ptmin = 5.0)
 
 # ╔═╡ 0bd764f9-d427-43fc-8342-603b6759ec8f
 md"Get the constituents of the first jet"
diff --git a/examples/constituents/jetreco-constituents.jl b/examples/constituents/jetreco-constituents.jl
@@ -18,7 +18,7 @@ event_no = 1
 cluster_seq = jet_reconstruct(events[event_no]; algorithm = JetAlgorithm.Kt, R = 1.0)
 
 # Retrieve the exclusive pj_jets, but as `PseudoJet` types
-pj_jets = inclusive_jets(cluster_seq; ptmin = 5.0, T = PseudoJet)
+pj_jets = inclusive_jets(cluster_seq, PseudoJet; ptmin = 5.0)
 
 # Get the constituents of the first jet
 my_constituents = JetReconstruction.constituents(pj_jets[1], cluster_seq)
diff --git a/examples/instrumented-jetreco.jl b/examples/instrumented-jetreco.jl
@@ -341,10 +341,9 @@ function main()
     else
         jet_type = PseudoJet
     end
-    events::Vector{Vector{jet_type}} = read_final_state_particles(args[:file],
+    events::Vector{Vector{jet_type}} = read_final_state_particles(args[:file], jet_type;
                                                                   maxevents = args[:maxevents],
-                                                                  skipevents = args[:skip],
-                                                                  T = jet_type)
+                                                                  skipevents = args[:skip])
 
     # Major switch between modes of running 
     if args[:alloc]
diff --git a/examples/jetreco.jl b/examples/jetreco.jl
@@ -149,10 +149,9 @@ function main()
     else
         jet_type = PseudoJet
     end
-    events::Vector{Vector{jet_type}} = read_final_state_particles(args[:file],
+    events::Vector{Vector{jet_type}} = read_final_state_particles(args[:file], jet_type;
                                                                   maxevents = args[:maxevents],
-                                                                  skipevents = args[:skip],
-                                                                  T = jet_type)
+                                                                  skipevents = args[:skip])
     if isnothing(args[:algorithm]) && isnothing(args[:power])
         @warn "Neither algorithm nor power specified, defaulting to AntiKt"
         args[:algorithm] = JetAlgorithm.AntiKt
diff --git a/examples/softkiller/softkiller_plots.jl b/examples/softkiller/softkiller_plots.jl
@@ -165,15 +165,13 @@ function main()
     args[:hard_file] = normpath(joinpath(@__DIR__, args[:hard_file]))
 
     # Reading pileup and hard event files
-    events = read_final_state_particles(args[:pileup_file],
+    events = read_final_state_particles(args[:pileup_file], jet_type;
                                         maxevents = args[:maxevents],
-                                        skipevents = args[:skip],
-                                        T = jet_type)
+                                        skipevents = args[:skip])
 
-    h_events = read_final_state_particles(args[:hard_file],
+    h_events = read_final_state_particles(args[:hard_file], jet_type;
                                           maxevents = args[:maxevents],
-                                          skipevents = args[:skip],
-                                          T = jet_type)
+                                          skipevents = args[:skip])
 
     # Set up SoftKiller grid and rapidity range
     rapmax = 5.0
diff --git a/examples/softkiller/softkiller_runtime.jl b/examples/softkiller/softkiller_runtime.jl
@@ -86,15 +86,13 @@ function main()
     args[:hard_file] = normpath(joinpath(@__DIR__, args[:hard_file]))
 
     # Reading pileup and hard event files
-    events = read_final_state_particles(args[:pileup_file],
+    events = read_final_state_particles(args[:pileup_file], jet_type;
                                         maxevents = args[:maxevents],
-                                        skipevents = args[:skip],
-                                        T = jet_type)
+                                        skipevents = args[:skip])
 
-    h_events = read_final_state_particles(args[:hard_file],
+    h_events = read_final_state_particles(args[:hard_file], jet_type;
                                           maxevents = args[:maxevents],
-                                          skipevents = args[:skip],
-                                          T = jet_type)
+                                          skipevents = args[:skip])
 
     # Set up SoftKiller grid and rapidity range
     rapmax = 5.0
diff --git a/examples/substructure/jet-grooming.jl b/examples/substructure/jet-grooming.jl
@@ -9,7 +9,7 @@ events = read_final_state_particles(input_file)
 event_no = 1
 
 cluster_seq = jet_reconstruct(events[event_no]; algorithm = JetAlgorithm.CA, R = 1.0)
-jets = inclusive_jets(cluster_seq; ptmin = 5.0, T = PseudoJet)
+jets = inclusive_jets(cluster_seq, PseudoJet; ptmin = 5.0)
 
 filter = (radius = 0.3, hardest_jets = 3)
 
diff --git a/examples/substructure/jet-tagging.jl b/examples/substructure/jet-tagging.jl
@@ -9,7 +9,7 @@ events = read_final_state_particles(input_file)
 event_no = 1
 
 cluster_seq = jet_reconstruct(events[event_no]; algorithm = JetAlgorithm.CA, R = 1.0)
-jets = inclusive_jets(cluster_seq; ptmin = 5.0, T = PseudoJet)
+jets = inclusive_jets(cluster_seq, PseudoJet; ptmin = 5.0)
 
 MDtagger = (mu = 0.67, y = 0.09)
 
diff --git a/src/ClusterSequence.jl b/src/ClusterSequence.jl
@@ -215,16 +215,16 @@ function add_step_to_history!(clusterseq::ClusterSequence, parent1, parent2, jet
 end
 
 """
-    inclusive_jets(clusterseq::ClusterSequence{U}; ptmin = 0.0, T = LorentzVectorCyl) where {U}
+    inclusive_jets(clusterseq::ClusterSequence{U}, ::Type{T} = LorentzVectorCyl{Float64}; ptmin = 0.0) where {T, U}
 
 Return all inclusive jets of a ClusterSequence with pt > ptmin.
 
 # Arguments
 - `clusterseq::ClusterSequence`: The `ClusterSequence` object containing the
   clustering history and jets.
+- `::Type{T} = LorentzVectorCyl{Float64}`: The return type used for the selected jets.
 - `ptmin::Float64 = 0.0`: The minimum transverse momentum (pt) threshold for the
   inclusive jets.
-- `T = LorentzVectorCyl`: The return type used for the selected jets.
 
 # Returns
 An array of `T` objects representing the inclusive jets.
@@ -245,8 +245,9 @@ currently unrecognised types will return `LorentzVectorCyl`).
 inclusive_jets(clusterseq; ptmin = 10.0)
 ```
 """
-function inclusive_jets(clusterseq::ClusterSequence{U}; ptmin = 0.0,
-                        T = LorentzVectorCyl) where {U}
+function inclusive_jets(clusterseq::ClusterSequence{U},
+                        ::Type{T} = LorentzVectorCyl{Float64};
+                        ptmin = 0.0) where {T, U}
     pt2min = ptmin * ptmin
     jets_local = T[]
     # sizehint!(jets_local, length(clusterseq.jets))
@@ -272,21 +273,21 @@ function inclusive_jets(clusterseq::ClusterSequence{U}; ptmin = 0.0,
 end
 
 """
-    exclusive_jets(clusterseq::ClusterSequence{U}; dcut = nothing, njets = nothing, T = LorentzVectorCyl) where {U}
+    exclusive_jets(clusterseq::ClusterSequence{U}, ::Type{T} = LorentzVectorCyl{Float64}; dcut = nothing, njets = nothing) where {T, U}
 
 Return all exclusive jets of a ClusterSequence, with either a specific number of
 jets or a cut on the maximum distance parameter.
 
 # Arguments
 - `clusterseq::ClusterSequence`: The `ClusterSequence` object containing the
   clustering history and jets.
+- `::Type{T} = LorentzVectorCyl{Float64}`: The return type used for the selected jets.
 - `dcut::Union{Nothing, Real}`: The distance parameter used to define the
   exclusive jets. If `dcut` is provided, the number of exclusive jets will be
   calculated based on this parameter.
 - `njets::Union{Nothing, Integer}`: The number of exclusive jets to be
   calculated. If `njets` is provided, the distance parameter `dcut` will be
   calculated based on this number.
-- `T = LorentzVectorCyl`: The return type used for the selected jets.
 
 **Note**: Either `dcut` or `njets` must be provided (but not both).
 
@@ -308,11 +309,12 @@ currently unrecognised types will return `LorentzVectorCyl`).
 # Examples
 ```julia
 exclusive_jets(clusterseq, dcut = 20.0)
-exclusive_jets(clusterseq, njets = 3, T = PseudoJet)
+exclusive_jets(clusterseq, PseudoJet, njets = 3)
 ```
 """
-function exclusive_jets(clusterseq::ClusterSequence{U}; dcut = nothing, njets = nothing,
-                        T = LorentzVectorCyl) where {U}
+function exclusive_jets(clusterseq::ClusterSequence{U},
+                        ::Type{T} = LorentzVectorCyl{Float64};
+                        dcut = nothing, njets = nothing,) where {T, U}
     if isnothing(dcut) && isnothing(njets)
         throw(ArgumentError("Must pass either a dcut or an njets value"))
     end
diff --git a/src/CommonJetStructs.jl b/src/CommonJetStructs.jl
@@ -236,21 +236,34 @@ function addjets_ptscheme(jet1::T, jet2::T;
 end
 
 """
-    preprocess_massless_pt(jet::T) where {T <: FourMomentum}
+    preprocess_ptscheme(jet::T, ::Type{OutputT};
+                             cluster_hist_index::Int = 0) -> OutputT where {T <: FourMomentum,
+                                                                            OutputT <: FourMomentum}
+
+Jet preprocessor for the massless ``p_T`` schemes, resetting the energy of the
+jet to be equal to the 3-momentum of the input jet.
+"""
+function preprocess_ptscheme(jet::T, ::Type{OutputT};
+                             cluster_hist_index::Int = 0) where {T <: FourMomentum,
+                                                                 OutputT <: FourMomentum}
+    OutputT(px(jet), py(jet), pz(jet), p(jet); cluster_hist_index = cluster_hist_index)
+end
+
+"""
+    preprocess_ptscheme(jet::T; cluster_hist_index::Int = 0) -> T where {T <: FourMomentum}
 
 Jet preprocessor for the massless ``p_T`` schemes, resetting the energy of the
 jet to be equal to the 3-momentum of the input jet.
 """
 function preprocess_ptscheme(jet::T;
-                             cluster_hist_index::Int = 0,
-                             jet_type = T) where {T <: FourMomentum}
-    T(px(jet), py(jet), pz(jet), p(jet); cluster_hist_index = cluster_hist_index)
+                             cluster_hist_index::Int = 0) where {T <: FourMomentum}
+    preprocess_ptscheme(jet, T; cluster_hist_index = cluster_hist_index)
 end
 
 """
-    preprocess_ptscheme(particle::Union{LorentzVector, LorentzVectorCyl};
-                             cluster_hist_index::Int = 0,
-                             jet_type = PseudoJet)
+    preprocess_ptscheme(particle::Union{LorentzVector, LorentzVectorCyl},
+                             ::Type{OutputT}= PseudoJet,;
+                             cluster_hist_index::Int = 0) -> OutputT where {OutputT <: FourMomentum}
 
 Jet preprocessor for the massless ``p_T`` schemes, resetting the energy of the
 jet to be equal to the 3-momentum of the input jet (generic particle type).
@@ -262,11 +275,11 @@ This function is used to convert a particle of type `LorentzVector` or
 `FourMomentum`. (This is a work around until `LorentzVectorBase` can be used,
 which will make the accessors uniform.)
 """
-function preprocess_ptscheme(particle::Union{LorentzVector, LorentzVectorCyl};
-                             cluster_hist_index::Int = 0,
-                             jet_type = PseudoJet)
-    T(px(particle), py(particle), pz(particle), mag(particle);
-      cluster_hist_index = cluster_hist_index)
+function preprocess_ptscheme(particle::Union{LorentzVector, LorentzVectorCyl},
+                             ::Type{OutputT} = PseudoJet;
+                             cluster_hist_index::Int = 0) where {OutputT <: FourMomentum}
+    OutputT(px(particle), py(particle), pz(particle), mag(particle);
+            cluster_hist_index = cluster_hist_index)
 end
 
 """
diff --git a/src/EEAlgorithm.jl b/src/EEAlgorithm.jl
@@ -256,7 +256,7 @@ function ee_genkt_algorithm(particles::AbstractVector{T}; algorithm::JetAlgorith
         sizehint!(recombination_particles, length(particles) * 2)
         for (i, particle) in enumerate(particles)
             push!(recombination_particles,
-                  preprocess(particle; cluster_hist_index = i, jet_type = EEJet))
+                  preprocess(particle, EEJet; cluster_hist_index = i))
         end
     end
 
diff --git a/src/PlainAlgo.jl b/src/PlainAlgo.jl
@@ -255,7 +255,7 @@ function plain_jet_reconstruct(particles::AbstractVector{T};
         sizehint!(recombination_particles, length(particles) * 2)
         for (i, particle) in enumerate(particles)
             push!(recombination_particles,
-                  preprocess(particle; cluster_hist_index = i, jet_type = PseudoJet))
+                  preprocess(particle, PseudoJet; cluster_hist_index = i))
         end
     end
 
diff --git a/src/Substructure.jl b/src/Substructure.jl
@@ -90,7 +90,7 @@ This function reclusters the jet and iteratively checks the soft-drop condition
 function soft_drop(jet::PseudoJet, clusterseq::ClusterSequence{PseudoJet}; zcut::Real,
                    beta::Real, radius::Real = 1.0)
     new_clusterseq = recluster(jet, clusterseq; R = radius, algorithm = JetAlgorithm.CA)
-    new_jet = sort!(inclusive_jets(new_clusterseq; T = PseudoJet), by = pt2, rev = true)[1]
+    new_jet = sort!(inclusive_jets(new_clusterseq, PseudoJet), by = pt2, rev = true)[1]
 
     all_jets = new_clusterseq.jets
     hist = new_clusterseq.history
@@ -137,7 +137,7 @@ Filters a jet to retain only the hardest subjets based on a specified radius and
 function jet_filtering(jet::PseudoJet, clusterseq::ClusterSequence{PseudoJet}; radius::Real,
                        hardest_jets::Integer)
     new_clusterseq = recluster(jet, clusterseq; R = radius, algorithm = JetAlgorithm.CA)
-    reclustered = sort!(inclusive_jets(new_clusterseq; T = PseudoJet), by = pt2, rev = true)
+    reclustered = sort!(inclusive_jets(new_clusterseq, PseudoJet), by = pt2, rev = true)
 
     n = length(reclustered) <= hardest_jets ? length(reclustered) : hardest_jets
     hard = reclustered[1:n]
@@ -167,7 +167,7 @@ function jet_trimming(jet::PseudoJet, clusterseq::ClusterSequence{PseudoJet}; ra
     frac2 = fraction^2
 
     new_clusterseq = recluster(jet, clusterseq; R = radius, algorithm = recluster_method)
-    reclustered = sort!(inclusive_jets(new_clusterseq; T = PseudoJet), by = pt2, rev = true)
+    reclustered = sort!(inclusive_jets(new_clusterseq, PseudoJet), by = pt2, rev = true)
 
     hard = Vector{PseudoJet}(undef, 0)
     for item in reclustered
diff --git a/src/TiledAlgoLL.jl b/src/TiledAlgoLL.jl
@@ -350,7 +350,7 @@ function tiled_jet_reconstruct(particles::AbstractVector{T};
         sizehint!(recombination_particles, length(particles) * 2)
         for (i, particle) in enumerate(particles)
             push!(recombination_particles,
-                  preprocess(particle; cluster_hist_index = i, jet_type = PseudoJet))
+                  preprocess(particle, PseudoJet; cluster_hist_index = i))
         end
     end
 
diff --git a/src/Utils.jl b/src/Utils.jl
@@ -21,25 +21,26 @@ open_with_stream(fname::AbstractString) = begin
 end
 
 """
-    read_final_state_particles(fname; maxevents = -1, skipevents = 0, T=PseudoJet)
+    read_final_state_particles(fname, ::Type{T} = PseudoJet; maxevents = -1, skipevents = 0) where {T}
 
 Reads final state particles from a file and returns them as a vector of type T.
 
 # Arguments
 - `fname`: The name of the HepMC3 ASCII file to read particles from. If the file
   is gzipped, the function will automatically decompress it.
+- `::Type{T}=PseudoJet`: The type of object to construct and return.
 - `maxevents=-1`: The maximum number of events to read. -1 means all events will
   be read.
 - `skipevents=0`: The number of events to skip before an event is included.
-- `T=PseudoJet`: The type of object to construct and return.
 
 # Returns
 A vector of vectors of T objects, where each inner vector represents all
 the particles of a particular event. In particular T can be `PseudoJet` or
 a `LorentzVector` type. Note, if T is not `PseudoJet`, the order of the
 arguments in the constructor must be `(t, x, y, z)`.
 """
-function read_final_state_particles(fname; maxevents = -1, skipevents = 0, T = PseudoJet)
+function read_final_state_particles(fname, ::Type{T} = PseudoJet; maxevents = -1,
+                                    skipevents = 0) where {T}
     f = open_with_stream(fname)
     events = Vector{T}[]
 
@@ -102,7 +103,7 @@ function final_jets(jets::Vector{PseudoJet}, ptmin::AbstractFloat = 0.0)
 end
 
 """Specialisation for final jets from LorentzVectors (TODO: merge into more general function)"""
-function final_jets(jets::Vector{LorentzVector}, ptmin::AbstractFloat = 0.0)
+function final_jets(jets::Vector{<:LorentzVector}, ptmin::AbstractFloat = 0.0)
     count = 0
     final_jets = Vector{FinalJet}()
     sizehint!(final_jets, 6)
@@ -119,7 +120,7 @@ function final_jets(jets::Vector{LorentzVector}, ptmin::AbstractFloat = 0.0)
 end
 
 """Specialisation for final jets from LorentzVectorCyl (TODO: merge into more general function)"""
-function final_jets(jets::Vector{LorentzVectorCyl}, ptmin::AbstractFloat = 0.0)
+function final_jets(jets::Vector{<:LorentzVectorCyl}, ptmin::AbstractFloat = 0.0)
     count = 0
     final_jets = Vector{FinalJet}()
     sizehint!(final_jets, 6)
diff --git a/test/_common.jl b/test/_common.jl
@@ -91,7 +91,7 @@ function sort_jets!(jet_array::Vector{FinalJet})
     sort!(jet_array, by = jet_pt, rev = true)
 end
 
-function sort_jets!(jet_array::Vector{LorentzVectorCyl})
+function sort_jets!(jet_array::Vector{<:LorentzVectorCyl})
     jet_pt(jet) = jet.pt
     sort!(jet_array, by = jet_pt, rev = true)
 end
diff --git a/test/data/hepmc32summary.jl b/test/data/hepmc32summary.jl
@@ -35,8 +35,9 @@ function main()
     args = parse_command_line(ARGS)
 
     for file in args[:files]
-        events = read_final_state_particles(file; maxevents = args[:maxevents],
-                                            skipevents = args[:skip], T = LorentzVector)
+        events = read_final_state_particles(file, LorentzVector{Float64};
+                                            maxevents = args[:maxevents],
+                                            skipevents = args[:skip])
         n_events = length(events)
         n_particles = Int[]
         for e in events
diff --git a/test/test-compare-types.jl b/test/test-compare-types.jl
diff --git a/test/test-constituents.jl b/test/test-constituents.jl
diff --git a/test/test-substructure.jl b/test/test-substructure.jl
