Cluster sequence return type

.gitignore

@@ -1,9 +1,38 @@
-/Manifest.toml
-/jetsavetest*.dat
-/demo.jl
-.vscode
+# Files generated by invoking Julia with --code-coverage
+*.jl.cov
+*.jl.*.cov
+
+# Files generated by invoking Julia with --track-allocation
+*.jl.mem
+
+# System-specific files and directories generated by the BinaryProvider and BinDeps packages
+# They contain absolute paths specific to the host computer, and so should not be committed
+deps/deps.jl
+deps/build.log
+deps/downloads/
+deps/usr/
+deps/src/
+
+# Build artifacts for creating documentation generated by the Documenter package
+docs/build/
+docs/site/
+
+# File generated by Pkg, the package manager, based on a corresponding Project.toml
+# It records a fixed state of all packages used by the project. As such, it should not be
+# committed for packages, but should be committed for applications that require a static
+# environment.
+Manifest.toml
+
+# Editor files
+.vscode/*
+
+# Test data files and benchmarking
+jetsavetest*.dat
+benchmark/*
+
+# Misc files
+.DS_Store
 /notebooks/*
+/profile/*
 /statprof/*
 /debug/*
-*.mem
-/benchmark/*

examples/jetreco.jl

@@ -17,18 +17,20 @@ using JSON
 using LorentzVectorHEP
 using JetReconstruction
 
-function profile_code(jet_reconstruction, events, niters)
+function profile_code(profile, jet_reconstruction, events, niters; R=0.4, p=-1, strategy=JetRecoStrategy.N2Tiled)
 	Profile.init(n = 5*10^6, delay = 0.00001)
 	profile_events(events) = begin
 		for evt in events
-			jet_reconstruction(evt, R = 0.4)
+			jet_reconstruction(evt, R=R, p=p, strategy=strategy)
 		end
 	end
 	profile_events(events[1:1])
 	@profile for i ∈ 1:niters
 		profile_events(events)
 	end
-	statprofilehtml()
+	profile_path = joinpath("profile", profile, "profsvg.svg")
+	mkpath(dirname(profile_path))
+	statprofilehtml(path=dirname(profile_path))
 	fcolor = FlameGraphs.FlameColors(
 		reverse(colormap("Blues", 15))[1:5],
 		colorant"slategray4",
@@ -38,14 +40,14 @@ function profile_code(jet_reconstruction, events, niters)
 	)
 	ProfileSVG.save(
 		fcolor,
-		joinpath("statprof", "profsvg.svg");
+		profile_path,
 		combine = true,
 		timeunit = :ms,
 		font = "Arial, Helvetica, sans-serif",
 	)
 	println(
 		"Flame graph from ProfileSVG.jl at file://",
-		abspath("statprof/profsvg.svg"),
+		abspath(profile_path),
 		"\n",
 		"""
 \tRed tint:          Runtime dispatch
@@ -71,7 +73,7 @@ function jet_process(
 	strategy::JetRecoStrategy.Strategy,
 	nsamples::Integer = 1,
 	gcoff::Bool = false,
-	profile::Bool = false,
+	profile = nothing,
     alloc::Bool = false,
 	dump::Union{String, Nothing} = nothing,
 )
@@ -83,24 +85,22 @@ function jet_process(
 	end
 
 	# Warmup code if we are doing a multi-sample timing run
-	if nsamples > 1 || profile
+	if nsamples > 1 || !isnothing(profile)
 		@info "Doing initial warm-up run"
 		for event in events
-			finaljets, _ = generic_jet_reconstruct(event, R = distance, p = power, strategy = strategy)
-			final_jets(finaljets, ptmin)
+			_ = inclusive_jets(generic_jet_reconstruct(event, R = distance, p = power, strategy = strategy), ptmin)
 		end
 	end
 
-	if profile
-		profile_code(generic_jet_reconstruct, events, nsamples)
+	if !isnothing(profile)
+		profile_code(profile, generic_jet_reconstruct, events, nsamples, R = distance, p = power, strategy = strategy)
 		return nothing
 	end
 
     if alloc
         println("Memory allocation statistics:")
         @timev for event in events
-            finaljets, _ = generic_jet_reconstruct(event, R = distance, p = power, strategy = strategy)
-			final_jets(finaljets, ptmin)
+			_ = inclusive_jets(generic_jet_reconstruct(event, R = distance, p = power, strategy = strategy), ptmin)
         end
         return nothing
     end
@@ -114,19 +114,18 @@ function jet_process(
 		gcoff && GC.enable(false)
 		t_start = time_ns()
 		for (ievt, event) in enumerate(events)
-			finaljets, _ = generic_jet_reconstruct(event, R = distance, p = power, ptmin=ptmin, strategy = strategy)
-			fj = final_jets(finaljets, ptmin)
+			finaljets = inclusive_jets(generic_jet_reconstruct(event, R = distance, p = power, strategy = strategy), ptmin)
 			# Only print the jet content once
 			if irun == 1
 				@info begin
 					jet_output = "Event $(ievt)\n"
-					for (ijet, jet) in enumerate(fj)
+					for (ijet, jet) in enumerate(finaljets)
 						jet_output *= " $(ijet) - $(jet)\n"
 					end
 					"$(jet_output)"
 				end
 				if !isnothing(dump)
-					push!(jet_collection, FinalJets(ievt, fj))
+					push!(jet_collection, FinalJets(ievt, finaljets))
 				end
 			end
 		end
@@ -212,8 +211,9 @@ parse_command_line(args) = begin
 		action = :store_true
 
 		"--profile"
-		help = "Profile code and generate a flame graph."
-		action = :store_true
+		help = """Profile code and generate a flame graph; the results 
+		will be stored in 'profile/PROFILE' subdirectory, where PROFILE is the value
+		given to this option)"""
 
         "--alloc"
         help = "Provide memory allocation statistics."

src/ClusterSequence.jl

@@ -0,0 +1,135 @@
+# Structure definitions for the Tiled algorithm, with linked list
+# and TiledJets (a la FastJet)
+# Original Julia implementation by Philippe Gras,
+# ported to this package by Graeme Stewart
+
+using Accessors
+
+# One cannot use an ENUM here as it's a different type
+# I don't know any good way to keep this lean and to have these necessary
+# flags other than using these "magic numbers"
+const Invalid=-3
+const NonexistentParent=-2
+const BeamJet=-1
+
+"""A struct holding a record of jet mergers and finalisations"""
+struct HistoryElement
+    """Index in _history where first parent of this jet
+    was created (NonexistentParent if this jet is an
+    original particle)"""
+    parent1::Int
+
+    """Index in _history where second parent of this jet
+    was created (NonexistentParent if this jet is an
+    original particle); BeamJet if this history entry
+    just labels the fact that the jet has recombined
+    with the beam)"""
+    parent2::Int
+
+    """Index in _history where the current jet is
+    recombined with another jet to form its child. It
+    is Invalid if this jet does not further
+    recombine."""
+    child::Int
+
+    """Index in the _jets vector where we will find the
+    PseudoJet object corresponding to this jet
+    (i.e. the jet created at this entry of the
+    history). NB: if this element of the history
+    corresponds to a beam recombination, then
+    jetp_index=Invalid."""
+    jetp_index::Int
+
+    """The distance corresponding to the recombination
+       at this stage of the clustering."""
+    dij::Float64
+
+    """The largest recombination distance seen
+       so far in the clustering history."""
+    max_dij_so_far::Float64
+end
+
+"""Used for initial particles"""
+HistoryElement(jetp_index) = HistoryElement(NonexistentParent, NonexistentParent, Invalid, jetp_index, 0.0, 0.0)
+
+
+"""
+Convienence structure holding all of the relevant parameters for
+the jet reconstruction
+"""
+struct ClusterSequence
+    """
+    This contains the physical PseudoJets; for each PseudoJet one can find
+    the corresponding position in the _history by looking at
+    _jets[i].cluster_hist_index()
+    """
+    jets::Vector{PseudoJet}
+
+    """
+    This vector will contain the branching history; for each stage,
+    history[i].jetp_index indicates where to look in the _jets
+    vector to get the physical PseudoJet.
+    """
+    history::Vector{HistoryElement}
+
+    """Total energy of the event"""
+    Qtot
+end
+
+"""Return all inclusive jets of a ClusterSequence with pt > ptmin"""
+function inclusive_jets(clusterseq::ClusterSequence, ptmin = 0.0)
+    dcut = ptmin * ptmin
+    jets_local = Vector{LorentzVectorCyl}(undef, 0)
+    # sizehint!(jets_local, length(clusterseq.jets))
+    # For inclusive jets with a plugin algorithm, we make no
+    # assumptions about anything (relation of dij to momenta,
+    # ordering of the dij, etc.)
+    # for elt in Iterators.reverse(clusterseq.history)
+    for elt in clusterseq.history
+        elt.parent2 == BeamJet || continue
+        iparent_jet = clusterseq.history[elt.parent1].jetp_index
+        jet = clusterseq.jets[iparent_jet]
+        if pt2(jet) >= dcut
+            push!(jets_local, LorentzVectorCyl(pt(jet), rapidity(jet), phi(jet), mass(jet)))
+        end
+    end
+    jets_local
+end
+
+"""Add a new jet's history into the recombination sequence"""
+add_step_to_history!(clusterseq::ClusterSequence, parent1, parent2, jetp_index, dij) = begin
+    max_dij_so_far = max(dij, clusterseq.history[end].max_dij_so_far)
+    push!(clusterseq.history, HistoryElement(parent1, parent2, Invalid,
+        jetp_index, dij, max_dij_so_far))
+
+    local_step = length(clusterseq.history)
+
+    # Sanity check: make sure the particles have not already been recombined
+    #
+    # Note that good practice would make this an assert (since this is
+    # a serious internal issue). However, we decided to throw an
+    # InternalError so that the end user can decide to catch it and
+    # retry the clustering with a different strategy.
+    @assert parent1 >= 1
+    if clusterseq.history[parent1].child != Invalid
+        error("Internal error. Trying to recombine an object that has previsously been recombined.")
+    end
+
+    hist_elem = clusterseq.history[parent1]
+    clusterseq.history[parent1] = @set hist_elem.child = local_step
+
+    if parent2 >= 1
+        clusterseq.history[parent2].child == Invalid || error(
+            "Internal error. Trying to recombine an object that has previsously been recombined.  Parent " * string(parent2) * "'s child index " * string(clusterseq.history[parent1].child) * ". Parent jet index: " *
+            string(clusterseq.history[parent2].jetp_index) * ".",
+        )
+        hist_elem = clusterseq.history[parent2]
+        clusterseq.history[parent2] = @set hist_elem.child = local_step
+    end
+
+    # Get cross-referencing right from PseudoJets
+    if jetp_index != Invalid
+        @assert jetp_index >= 1
+        clusterseq.jets[jetp_index]._cluster_hist_index = local_step
+    end
+end

src/GenericAlgo.jl

@@ -2,7 +2,7 @@
 # switch based on the strategy, or based on the event density
 # if the "Best" strategy is to be employed
 
-function generic_jet_reconstruct(particles; p = -1, R = 1.0, recombine = +, ptmin = 0.0, strategy = JetRecoStrategy.Best)
+function generic_jet_reconstruct(particles; p = -1, R = 1.0, recombine = +, strategy = JetRecoStrategy.Best)
     # Either map to the fixed algorithm corresponding to the strategy
     # or to an optimal choice based on the density of initial particles
 
@@ -18,5 +18,5 @@ function generic_jet_reconstruct(particles; p = -1, R = 1.0, recombine = +, ptmi
     end
 
     # Now call the chosen algorithm, passing through the other parameters
-    algorithm(particles; p = p, R = R, recombine = recombine, ptmin = ptmin)
+    algorithm(particles; p = p, R = R, recombine = recombine)
 end

src/JetReconstruction.jl

@@ -22,12 +22,13 @@ py(p::LorentzVectorCyl) = LorentzVectorHEP.py(p)
 pz(p::LorentzVectorCyl) = LorentzVectorHEP.pz(p)
 energy(p::LorentzVectorCyl) = LorentzVectorHEP.energy(p)
 
-# Philipp's pseudojet type
+# Pseudojet type
 include("Pseudojet.jl")
 export PseudoJet
 
-# Simple HepMC3 reader
-include("HepMC3.jl")
+# ClusterSequence type
+include("ClusterSequence.jl")
+export ClusterSequence, inclusive_jets
 
 # Jet reconstruction strategies
 include("JetRecoStrategies.jl")
@@ -49,6 +50,9 @@ export tiled_jet_reconstruct
 include("GenericAlgo.jl")
 export generic_jet_reconstruct
 
+# Simple HepMC3 reader
+include("HepMC3.jl")
+
 # jet serialisation (saving to file)
 include("Serialize.jl")
 export savejets, loadjets!, loadjets