Fix paramaterised types for substructure

JuliaHEP · Feb 4, 2025 · cf7018e · cf7018e
1 parent 89421bf
commit cf7018e
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 12 deletions.
diff --git a/src/ClusterSequence.jl b/src/ClusterSequence.jl
@@ -266,7 +266,7 @@ function inclusive_jets(clusterseq::ClusterSequence{U}; ptmin = 0.0,
         jet = clusterseq.jets[iparent_jet]
         if pt2(jet) >= pt2min
             @debug "Added inclusive jet index $iparent_jet"
-            if T == U
+            if T <: U
                 push!(jets_local, jet)
             else
                 push!(jets_local,

diff --git a/src/Substructure.jl b/src/Substructure.jl
@@ -142,7 +142,8 @@ The method stops at the first jet satisfying the mass and distance thresholds.
 # Returns:
 - `PseudoJet`: The jet (or subjet) satisfying the mass drop conditions, if tagging is successful, otherwise a zero-momentum PseudoJet
 """
-function mass_drop(jet::PseudoJet, clusterseq::ClusterSequence, tag::MassDropTagger)
+function mass_drop(jet::PseudoJet{T}, clusterseq::ClusterSequence,
+                   tag::MassDropTagger) where {T <: Real}
     all_jets = clusterseq.jets
     hist = clusterseq.history
 
@@ -188,11 +189,11 @@ This function reclusters the jet and iteratively checks the soft-drop condition
 # Returns:
 - `PseudoJet`: Groomed jet or zero-momentum PseudoJet if grooming fails
 """
-function soft_drop(jet::PseudoJet, clusterseq::ClusterSequence,
-                   tag::SoftDropTagger)
+function soft_drop(jet::PseudoJet{T}, clusterseq::ClusterSequence,
+                   tag::SoftDropTagger) where {T <: Real}
     rad = tag.cluster_rad
     new_clusterseq = recluster(jet, clusterseq; R = rad, algorithm = JetAlgorithm.CA)
-    new_jet = sort!(inclusive_jets(new_clusterseq; T = PseudoJet), by = pt2, rev = true)[1]
+    new_jet = sort!(inclusive_jets(new_clusterseq; T = PseudoJet{T}), by = pt2, rev = true)[1]
 
     all_jets = new_clusterseq.jets
     hist = new_clusterseq.history
@@ -238,10 +239,12 @@ Filters a jet to retain only the hardest subjets based on a specified radius and
 # Returns:
 - `PseudoJet`: Filtered jet composed of the hardest subjets
 """
-function jet_filtering(jet::PseudoJet, clusterseq::ClusterSequence, filter::JetFilter)
+function jet_filtering(jet::PseudoJet{T}, clusterseq::ClusterSequence,
+                       filter::JetFilter) where {T <: Real}
     rad = filter.filter_radius
     new_clusterseq = recluster(jet, clusterseq; R = rad, algorithm = JetAlgorithm.CA)
-    reclustered = sort!(inclusive_jets(new_clusterseq; T = PseudoJet), by = pt2, rev = true)
+    reclustered = sort!(inclusive_jets(new_clusterseq; T = PseudoJet{T}), by = pt2,
+                        rev = true)
 
     n = length(reclustered) <= filter.num_hardest_jets ? length(reclustered) :
         filter.num_hardest_jets
@@ -265,13 +268,15 @@ Trims a jet by removing subjets with transverse momentum below a specified fract
 # Returns:
 - `PseudoJet`: Trimmed jet composed of retained subjets
 """
-function jet_trimming(jet::PseudoJet, clusterseq::ClusterSequence, trim::JetTrim)
+function jet_trimming(jet::PseudoJet{T}, clusterseq::ClusterSequence,
+                      trim::JetTrim) where {T <: Real}
     rad = trim.trim_radius
     alg = trim.recluster_method
     frac2 = trim.trim_fraction^2
 
     new_clusterseq = recluster(jet, clusterseq; R = rad, algorithm = alg)
-    reclustered = sort!(inclusive_jets(new_clusterseq; T = PseudoJet), by = pt2, rev = true)
+    reclustered = sort!(inclusive_jets(new_clusterseq; T = PseudoJet{T}), by = pt2,
+                        rev = true)
 
     hard = Vector{PseudoJet}(undef, 0)
     for item in reclustered

diff --git a/test/test-constituents.jl b/test/test-constituents.jl
@@ -15,8 +15,6 @@ input_file = joinpath(dirname(pathof(JetReconstruction)), "..", "test", "data",
                       "events.pp13TeV.hepmc3.gz")
 events = read_final_state_particles(input_file)
 
-println(typeof(events))
-
 # Event to pick
 event_no = 1
 

diff --git a/test/test-substructure.jl b/test/test-substructure.jl
@@ -36,7 +36,7 @@ for m in keys(methods)
     @testset "$testname" begin
         for (ievt, evt) in enumerate(events)
             cluster_seq = jet_reconstruct(evt, p = 0, R = 1.0)
-            jets = inclusive_jets(cluster_seq; ptmin = 5.0, T = PseudoJet)
+            jets = inclusive_jets(cluster_seq; ptmin = 5.0, T = PseudoJet{Float64})
             groomed = sort!([methods[m](jet, cluster_seq, groomer)
                              for jet in jets], by = JetReconstruction.pt2, rev = true)