diff --git a/Project.toml b/Project.toml index 9394b00..74aa3c1 100644 --- a/Project.toml +++ b/Project.toml @@ -22,9 +22,9 @@ JetVisualisation = "Makie" Accessors = "0.1.36" EnumX = "1.0.4" JSON = "0.21.4" -Makie = "0.21" LorentzVectorHEP = "0.1.6" LorentzVectors = "0.4.3" +Makie = "0.21" julia = "1.9" [extras] diff --git a/src/AlgorithmStrategyEnums.jl b/src/AlgorithmStrategyEnums.jl index ed24065..415c020 100644 --- a/src/AlgorithmStrategyEnums.jl +++ b/src/AlgorithmStrategyEnums.jl @@ -1,23 +1,52 @@ # EnumX provides scoped enums, which are nicer using EnumX -# Valid strategy enum + +""" + enum RecoStrategy + +Scoped enumeration (using EnumX) representing the different strategies for jet reconstruction. + +## Fields +- `Best`: The best strategy. +- `N2Plain`: The plain N2 strategy. +- `N2Tiled`: The tiled N2 strategy. +""" @enumx T = Strategy RecoStrategy Best N2Plain N2Tiled const AllJetRecoStrategies = [String(Symbol(x)) for x in instances(RecoStrategy.Strategy)] -# Algorithm emun + +""" + enum T + +Scoped enumeration (using EnumX) representing different jet algorithms used in the JetReconstruction module. + +## Fields +- `AntiKt`: The Anti-Kt algorithm. +- `CA`: The Cambridge/Aachen algorithm. +- `Kt`: The Inclusive-Kt algorithm. +- `EEKt`: The Generalised e+e- kt algorithm. +- `Durham`: The e+e- kt algorithm, aka Durham. +""" @enumx T = Algorithm JetAlgorithm AntiKt CA Kt EEKt Durham const AllJetRecoAlgorithms = [String(Symbol(x)) for x in instances(JetAlgorithm.Algorithm)] -# Map from algorithms to power values used -power2algorithm = Dict(-1 => JetAlgorithm.AntiKt, + +""" + power2algorithm + +A dictionary that maps power values to corresponding jet algorithm used for pp jet reconstruction. +""" +const power2algorithm = Dict(-1 => JetAlgorithm.AntiKt, 0 => JetAlgorithm.CA, 1 => JetAlgorithm.Kt) -algorithm2power = Dict(JetAlgorithm.AntiKt => -1, - JetAlgorithm.CA => 0, - JetAlgorithm.Kt => 1) -# Map from string to an enum value (used for CLI parsing) + +""" + Base.tryparse(E::Type{<:Enum}, str::String) + +Parser that converts a string to an enum value if it exists, otherwise returns nothing. +""" Base.tryparse(E::Type{<:Enum}, str::String) = let insts = instances(E), p = findfirst(==(Symbol(str)) ∘ Symbol, insts) diff --git a/src/ClusterSequence.jl b/src/ClusterSequence.jl index adf98c4..03ce26a 100644 --- a/src/ClusterSequence.jl +++ b/src/ClusterSequence.jl @@ -1,93 +1,99 @@ # 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 using Logging -# 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" -# (TODO: this is not runtime critical, so could consider a Union{Int,Emum}?) +# Constant values for "magic numbers" that represent special states +# in the history. These are negative integers, but as this is +# not runtime critical, so could consider a Union{Int,Emum} +"Invalid child for this jet, meaning it did not recombine further" const Invalid = -3 + +"Original cluster so it has no parent" const NonexistentParent = -2 + +"Cluster recombined with beam" const BeamJet = -1 -"""A struct holding a record of jet mergers and finalisations""" +""" + struct HistoryElement + +A struct holding a record of jet mergers and finalisations + +Fields: +- parent1: Index in history where first parent of this jet was created (NonexistentParent if this jet is an original particle) +- parent2: 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) +- child: 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. +- jetp_index: 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. +- dij: The distance corresponding to the recombination at this stage of the clustering. +- max_dij_so_far: The largest recombination distance seen so far in the clustering history. +""" 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) + +Constructs a `HistoryElement` object with the given `jetp_index`, used for initialising the history with original particles. + +# Arguments +- `jetp_index`: The index of the jetp. + +# Returns +A `HistoryElement` object. + +""" 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 + +A struct holding the full history of a jet clustering sequence, including the final jets. + +# Fields +- `algorithm::JetAlgorithm.Algorithm`: The algorithm used for clustering. +- `strategy::RecoStrategy.Strategy`: The strategy used for clustering. +- `jets::Vector{PseudoJet}`: The physical PseudoJets in the cluster sequence. Each PseudoJet corresponds to a position in the history. +- `n_initial_jets::Int`: The initial number of particles used for exclusive jets. +- `history::Vector{HistoryElement}`: The branching history of the cluster sequence. Each stage in the history indicates where to look in the jets vector to get the physical PseudoJet. +- `Qtot::Any`: The total energy of the event. """ struct ClusterSequence - """Algorithm and strategy used""" algorithm::JetAlgorithm.Algorithm strategy::RecoStrategy.Strategy - - """ - 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} - - """ - Record the initial number of particlesm used for exclusive jets - """ n_initial_jets::Int - - """ - 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::Any end -"""ClusterSequence constructor, where the power value is given""" +""" + ClusterSequence(p::Int, strategy::RecoStrategy.Strategy, jets, history, Qtot) + +Constructs a `ClusterSequence` object with a specific power value mapped to a pp reconstruction algorithm. + +# Arguments +- `p::Int`: The power value for the algorithm. +- `strategy::RecoStrategy.Strategy`: The reconstruction strategy. +- `jets`: The jets to be clustered. +- `history`: The length of the jets. +- `Qtot`: The total energy of the jets. + +# Returns +A `ClusterSequence` object. +""" ClusterSequence(p::Int, strategy::RecoStrategy.Strategy, jets, history, Qtot) = begin if !haskey(power2algorithm, p) raise(ArgumentError("Unrecognised algorithm for power value p=$p")) @@ -95,11 +101,45 @@ ClusterSequence(p::Int, strategy::RecoStrategy.Strategy, jets, history, Qtot) = ClusterSequence(power2algorithm[p], strategy, jets, length(jets), history, Qtot) end -"""ClusterSequence constructor, with direct algorithm specified""" +""" + ClusterSequence(alg::JetAlgorithm.Algorithm, strategy::RecoStrategy.Strategy, jets, history, Qtot) + +Constructs a `ClusterSequence` object with a specific algorithm spcified. + +# Arguments +- `alg::JetAlgorithm.Algorithm`: The algorithm. +- `strategy::RecoStrategy.Strategy`: The reconstruction strategy. +- `jets`: The jets to be clustered. +- `history`: The length of the jets. +- `Qtot`: The total energy of the jets. + +# Returns +A `ClusterSequence` object. +""" ClusterSequence(alg::JetAlgorithm.Algorithm, strategy::RecoStrategy.Strategy, jets, history, Qtot) = ClusterSequence(alg, strategy, jets, length(jets), history, Qtot) -"""Add a new jet's history into the recombination sequence""" +""" + add_step_to_history!(clusterseq::ClusterSequence, parent1, parent2, jetp_index, dij) + +Add a new jet's history into the recombination sequence. + +Arguments: +- `clusterseq::ClusterSequence`: The cluster sequence object. +- `parent1`: The index of the first parent. +- `parent2`: The index of the second parent. +- `jetp_index`: The index of the jet. +- `dij`: The dij value. + +This function adds a new `HistoryElement` to the `history` vector of the +`clusterseq` object. The `HistoryElement` contains information about the +parents, child, jet index, dij value, and the maximum dij value so far. It also +updates the child index of the parent elements. + +If the `parent1` or `parent2` have already been recombined, an `InternalError` +is thrown. The `jetp_index` is used to update the `_cluster_hist_index` of the +corresponding `PseudoJet` object. +""" 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, @@ -137,7 +177,26 @@ add_step_to_history!(clusterseq::ClusterSequence, parent1, parent2, jetp_index, end end -"""Return all inclusive jets of a ClusterSequence with pt > ptmin""" +""" + inclusive_jets(clusterseq::ClusterSequence, ptmin = 0.0) + +Return all inclusive jets of a ClusterSequence with pt > ptmin. + +# Arguments +- `clusterseq::ClusterSequence`: The `ClusterSequence` object containing the clustering history and jets. +- `ptmin::Float64 = 0.0`: The minimum transverse momentum (pt) threshold for the inclusive jets. + +# Returns +An array of `LorentzVectorCyl` objects representing the inclusive jets. + +# Description +This function computes the inclusive jets from a given `ClusterSequence` object. It iterates over the clustering history and checks the transverse momentum of each parent jet. If the transverse momentum is greater than or equal to `ptmin`, the jet is added to the array of inclusive jets. + +# Example +```julia +inclusive_jets(clusterseq, ptmin = 10.0) +``` +""" function inclusive_jets(clusterseq::ClusterSequence, ptmin = 0.0) pt2min = ptmin * ptmin jets_local = LorentzVectorCyl[] @@ -158,15 +217,40 @@ function inclusive_jets(clusterseq::ClusterSequence, ptmin = 0.0) jets_local end +""" + exclusive_jets(clusterseq::ClusterSequence; dcut = nothing, njets = nothing) + +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. +- `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. + +**Note**: Either `dcut` or `njets` must be provided (but not both). -"""Return all exclusive jets of a ClusterSequence""" +# Returns +- `excl_jets::Array{LorentzVectorCyl}`: An array of `LorentzVectorCyl` objects representing the exclusive jets. + +# Exceptions +- `ArgumentError`: If neither `dcut` nor `njets` is provided. +- `ArgumentError`: If the algorithm used in the `ClusterSequence` object is not suitable for exclusive jets. +- `ErrorException`: If the cluster sequence is incomplete and exclusive jets are unavailable. + +# Examples +```julia +exclusive_jets(clusterseq, dcut = 20.0) +exclusive_jets(clusterseq, njets = 3) +``` +""" function exclusive_jets(clusterseq::ClusterSequence; dcut = nothing, njets = nothing) if isnothing(dcut) && isnothing(njets) throw(ArgumentError("Must pass either a dcut or an njets value")) end if !isnothing(dcut) - njets = n_exclusive_jets(clusterseq, dcut=dcut) + njets = n_exclusive_jets(clusterseq, dcut = dcut) end # Check that an algorithm was used that makes sense for exclusive jets @@ -175,7 +259,7 @@ function exclusive_jets(clusterseq::ClusterSequence; dcut = nothing, njets = not end # njets search - stop_point = 2*clusterseq.n_initial_jets - njets + 1 + stop_point = 2 * clusterseq.n_initial_jets - njets + 1 # Sanity check - never return more jets than initial particles if stop_point < clusterseq.n_initial_jets @@ -203,7 +287,23 @@ function exclusive_jets(clusterseq::ClusterSequence; dcut = nothing, njets = not end -"""Return all number of exclusive jets of a ClusterSequence that are above a certain dcut value""" +""" + n_exclusive_jets(clusterseq::ClusterSequence; dcut::AbstractFloat) + +Return the number of exclusive jets of a ClusterSequence that are above a certain dcut value. + +# Arguments +- `clusterseq::ClusterSequence`: The `ClusterSequence` object containing the clustering history. +- `dcut::AbstractFloat`: The maximum calue for the distance parameter in the reconstruction. + +# Returns +The number of exclusive jets in the `ClusterSequence` object. + +# Example +```julia +n_exclusive_jets(clusterseq, dcut = 20.0) +``` +""" function n_exclusive_jets(clusterseq::ClusterSequence; dcut::AbstractFloat) # Check that an algorithm was used that makes sense for exclusive jets if !(clusterseq.algorithm ∈ (JetAlgorithm.CA, JetAlgorithm.Kt, JetAlgorithm.EEKt, JetAlgorithm.Durham)) @@ -223,5 +323,4 @@ function n_exclusive_jets(clusterseq::ClusterSequence; dcut::AbstractFloat) # The number of jets is then given by this formula length(clusterseq.history) - i_dcut - -end \ No newline at end of file +end diff --git a/src/GenericAlgo.jl b/src/GenericAlgo.jl index d3d5395..a88e58d 100644 --- a/src/GenericAlgo.jl +++ b/src/GenericAlgo.jl @@ -1,13 +1,29 @@ -# This is the generic reconstruction algorithm that will -# switch based on the strategy, or based on the event density -# if the "Best" strategy is to be employed +""" + jet_reconstruct(particles; p = -1, R = 1.0, recombine = +, strategy = RecoStrategy.Best) +Reconstructs jets from a collection of particles using a specified algorithm and strategy + +# Arguments +- `particles`: A collection of particles used for jet reconstruction. +- `p=-1`: The power value used for the distance measure, which maps to a particular reconstruction algorithm (-1 = AntiKt, 0 = Cambridge/Aachen, 1 = Kt). +- `R=1.0`: The jet radius parameter. +- `recombine=+`: The recombination scheme used for combining particles. +- `strategy=RecoStrategy.Best`: The jet reconstruction strategy to use. `RecoStrategy.Best` makes a dynamic decision based on the number of starting particles. + +# Returns +A cluster sequence object containing the reconstructed jets and the merging history. + +# Example +```julia +jet_reconstruct(particles; p = -1, R = 0.4) +``` +""" function jet_reconstruct(particles; p = -1, R = 1.0, recombine = +, strategy = RecoStrategy.Best) # Either map to the fixed algorithm corresponding to the strategy # or to an optimal choice based on the density of initial particles if strategy == RecoStrategy.Best - # The breakpoint of ~90 is determined empirically on e+e- -> H and 0.5TeV pp -> 5GeV jets + # The breakpoint of ~80 is determined empirically on e+e- -> H and 0.5TeV pp -> 5GeV jets algorithm = length(particles) > 80 ? tiled_jet_reconstruct : plain_jet_reconstruct elseif strategy == RecoStrategy.N2Plain algorithm = plain_jet_reconstruct diff --git a/src/HepMC3.jl b/src/HepMC3.jl index 48d1c32..7c44dbe 100644 --- a/src/HepMC3.jl +++ b/src/HepMC3.jl @@ -1,35 +1,63 @@ -# Copyright (c) 2022, Philippe Gras -# -#---------------------------------------------------------------------- -# This file is part of AntiKt.jl. -# -# AntiKt.jl is free software; you can redistribute it and/or modify -# it under the terms of the GNU General Public License as published by -# the Free Software Foundation; either version 2 of the License, or -# (at your option) any later version. -# -# The algorithms that underlie FastJet have required considerable -# development. They are described in the original FastJet paper, -# hep-ph/0512210 and in the manual, arXiv:1111.6097. If you use -# FastJet as part of work towards a scientific publication, please -# quote the version you use and include a citation to the manual and -# optionally also to hep-ph/0512210. -# -# AntiKet.jl is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU General Public License -# along with FastJet. If not, see . -#---------------------------------------------------------------------- - -""" Module providing a minimal support to read HepMC3 ascii files. +""" +Module providing a minimal support to read HepMC3 ascii files. + +This module contains functions and types for reading HepMC3 ASCII files. It +provides a `Particle` struct to represent particles in the event, and a +`read_events` function to read events from a file. + +## Types +- `Particle{T}`: A struct representing a particle in the event. It contains + fields for momentum, status, PDG ID, barcode, and vertex. +- `T`: The type parameter for the `Particle` struct, specifying the precision of + the momentum components. + +## Functions +- `read_events(f, fin; maxevents=-1, skipevents=0)`: Reads events from a HepMC3 + ascii file. Each event is passed to the provided function `f` as a vector of + `Particle`s. The `maxevents` parameter specifies the maximum number of events + to read (-1 to read all available events), and the `skipevents` parameter + specifies the number of events to skip at the beginning of the file. + + Example usage: + ```julia + function process_event(particles) + # Process the event + end + + file = open("events.hepmc", "r") + read_events(process_event, file, maxevents=10, skipevents=2) + close(file) + ``` + + - `f`: The function to be called for each event. It should take a single + argument of type `Vector{Particle}`. + - `fin`: The input file stream to read from. + - `maxevents`: The maximum number of events to read. Default is -1, which + reads all available events. + - `skipevents`: The number of events to skip at the beginning of the file. + Default is 0. + +## References +- [HepMC3](https://doi.org/10.1016/j.cpc.2020.107310): The HepMC3 software + package. """ module HepMC3 using LorentzVectors +""" + struct Particle{T} + +A struct representing a particle in the HepMC3 library. + +# Fields +- `momentum::LorentzVector{T}`: The momentum of the particle. +- `status::Integer`: The status code of the particle. +- `pdgid::Integer`: The PDG ID of the particle. +- `barcode::Integer`: The barcode of the particle. +- `vertex::Integer`: The vertex ID of the particle. + +""" struct Particle{T} momentum::LorentzVector{T} status::Integer @@ -46,6 +74,46 @@ Particle{T}() where T = Particle(LorentzVector{T}(0., 0., 0., 0.), 0, 0, 0, 0) maximum number of events to read (value -1 to read all availble events) and a number of events to skip at the beginning of the file can be provided. """ + +""" + read_events(f, fin; maxevents=-1, skipevents=0) + +Reads events from a file and processes them. + +## Arguments +- `f`: A function that takes an array of `Particle{T}` as input and processes + it. +- `fin`: The input file object. +- `maxevents`: The maximum number of events to read. Default is -1, which means + all events will be read. +- `skipevents`: The number of events to skip before processing. Default is 0. + +## Description +This function reads events from a file and processes them using the provided +function `f`. Each event is represented by a series of lines in the file, where +each line corresponds to a particle. The function parses the lines and creates +`Particle{T}` objects, which are then passed to the processing function `f`. The +function `f` can perform any desired operations on the particles. + +## Example +```julia +f = open(fname) +events = Vector{PseudoJet}[] +HepMC3.read_events(f, maxevents = maxevents, skipevents = skipevents) do parts + input_particles = PseudoJet[] + for p in parts + if p.status == 1 + push!( + input_particles, + PseudoJet(p.momentum.x, p.momentum.y, p.momentum.z, p.momentum.t), + ) + end + end + push!(events, input_particles) +end +close(f) +``` +""" function read_events(f, fin; maxevents=-1, skipevents=0) T = Float64 particles = Particle{T}[] diff --git a/src/JSONresults.jl b/src/JSONresults.jl index e52788c..ccdda78 100644 --- a/src/JSONresults.jl +++ b/src/JSONresults.jl @@ -1,11 +1,30 @@ -"Simple structures to capture the state of a final jet output -from the algorithm" +""" + struct FinalJet + +A struct representing the final properties of a jet, used for +JSON serialisation. + +# Fields +- `rap::Float64`: The rapidity of the jet. +- `phi::Float64`: The azimuthal angle of the jet. +- `pt::Float64`: The transverse momentum of the jet. +""" struct FinalJet rap::Float64 phi::Float64 pt::Float64 end +""" + struct FinalJets + +A struct with the vector of all jets for a certain jet identifier, used for +JSON serialisation. + +# Fields +- `jetid::Int64`: The ID of the jet. +- `jets::Vector{FinalJet}`: A vector of `FinalJet` objects representing the jets. +""" struct FinalJets jetid::Int64 jets::Vector{FinalJet} diff --git a/src/JetReconstruction.jl b/src/JetReconstruction.jl index bb1af4e..b1fad9a 100644 --- a/src/JetReconstruction.jl +++ b/src/JetReconstruction.jl @@ -1,5 +1,17 @@ """ -Jet reconstruction (reclustering) in Julia. +Module for jet reconstruction algorithms and utilities. + +This module provides various algorithms and utilities for jet reconstruction in +high-energy physics (HEP) experiments. It includes functions for performing jet +clustering with different algorithms and strategies. + +Utility functions are also provided for reading input from HepMC3 files, saving +and loading jet objects, and visualizing jets. + +# Usage +To use this module, include it in your Julia code and call the desired functions +or use the exported types. See the documentation for each function and the +examples in this package for more information. """ module JetReconstruction @@ -59,7 +71,7 @@ export savejets, loadjets!, loadjets # utility functions, useful for different primary scripts include("Utils.jl") -export read_final_state_particles, read_final_state_particles_lv, pseudojets2vectors, final_jets +export read_final_state_particles, read_final_state_particles_lv, final_jets # Jet visualisation is an extension, see ext/JetVisualisation.jl function jetsplot() end diff --git a/src/PlainAlgo.jl b/src/PlainAlgo.jl index 22e86ea..30e3be4 100644 --- a/src/PlainAlgo.jl +++ b/src/PlainAlgo.jl @@ -1,5 +1,19 @@ using LoopVectorization +""" + dist(i, j, rapidity_array, phi_array) + +Compute the distance between points in a 2D space defined by rapidity and phi coordinates. + +# Arguments +- `i::Int`: Index of the first point to consider (indexes into `rapidity_array` and `phi_array`). +- `j::Int`: Index of the second point to consider (indexes into `rapidity_array` and `phi_array`). +- `rapidity_array::Vector{Float64}`: Array of rapidity coordinates. +- `phi_array::Vector{Float64}`: Array of phi coordinates. + +# Returns +- `distance::Float64`: The distance between the two points. +""" Base.@propagate_inbounds function dist(i, j, rapidity_array, phi_array) drapidity = rapidity_array[i] - rapidity_array[j] dphi = abs(phi_array[i] - phi_array[j]) @@ -7,14 +21,50 @@ Base.@propagate_inbounds function dist(i, j, rapidity_array, phi_array) muladd(drapidity, drapidity, dphi * dphi) end -# d_{ij} distance with i's NN (times R^2) + +""" + dij(i, kt2_array, nn, nndist) + +Compute the dij value for a given index `i` to its nearest neighbor. The nearest +neighbor is determined from `nn[i]`, and the metric distance to the nearest +neighbor is given by the distance `nndist[i]` applying the lower of the +`kt2_array` values for the two particles.ßß + +# Arguments +- `i`: The index of the element. +- `kt2_array`: An array of kt2 values. +- `nn`: An array of nearest neighbors. +- `nndist`: An array of nearest neighbor distances. + +# Returns +- The computed dij value. +""" Base.@propagate_inbounds function dij(i, kt2_array, nn, nndist) j = nn[i] d = nndist[i] d * min(kt2_array[i], kt2_array[j]) end -# finds new nn for i and checks everyone additionally +""" + upd_nn_crosscheck!(i, from, to, rapidity_array, phi_array, R2, nndist, nn) + +Update the nearest neighbor information for a given particle index `i` against +all particles in the range indexes `from` to `to`. The function updates the +`nndist` and `nn` arrays with the nearest neighbor distance and index +respectively, both for particle `i` and the checked particles `[from:to]` (hence +*crosscheck*). + +# Arguments +- `i::Int`: The index of the particle to update and check against. +- `from::Int`: The starting index of the range of particles to check against. +- `to::Int`: The ending index of the range of particles to check against. +- `rapidity_array`: An array containing the rapidity values of all particles. +- `phi_array`: An array containing the phi values of the all particles. +- `R2`: The squared jet distance threshold for considering a particle as a + neighbour. +- `nndist`: The array that stores the nearest neighbor distances. +- `nn`: The array that stores the nearest neighbor indices. +""" Base.@propagate_inbounds function upd_nn_crosscheck!(i::Int, from::Int, to::Int, rapidity_array, phi_array, R2, nndist, nn) nndist_min = R2 nn_min = i @@ -34,6 +84,26 @@ Base.@propagate_inbounds function upd_nn_crosscheck!(i::Int, from::Int, to::Int, end # finds new nn for i + +""" + upd_nn_nocross!(i, from, to, rapidity_array, phi_array, R2, nndist, nn) + +Update the nearest neighbor information for a given particle index `i` against +all particles in the range indexes `from` to `to`. The function updates the +`nndist` and `nn` arrays with the nearest neighbor distance and index +respectively, only for particle `i` (hence *nocross*). + +# Arguments +- `i::Int`: The index of the particle to update and check against. +- `from::Int`: The starting index of the range of particles to check against. +- `to::Int`: The ending index of the range of particles to check against. +- `rapidity_array`: An array containing the rapidity values of all particles. +- `phi_array`: An array containing the phi values of the all particles. +- `R2`: The squared jet distance threshold for considering a particle as a + neighbour. +- `nndist`: The array that stores the nearest neighbor distances. +- `nn`: The array that stores the nearest neighbor indices. +""" Base.@propagate_inbounds function upd_nn_nocross!(i::Int, from::Int, to::Int, rapidity_array, phi_array, R2, nndist, nn) nndist_min = R2 nn_min = i @@ -54,16 +124,40 @@ Base.@propagate_inbounds function upd_nn_nocross!(i::Int, from::Int, to::Int, ra nn[i] = nn_min end -# entire NN update step + +""" + upd_nn_step!(i, j, k, N, Nn, kt2_array, rapidity_array, phi_array, R2, nndist, nn, nndij) + +Update the nearest neighbor information after a jet merge step. + +Arguments: +- `i`: Index of the first particle in the last merge step. +- `j`: Index of the second particle in the last merge step. +- `k`: Index of the current particle for which the nearest neighbour will be updated. +- `N`: Total number of particles (currently vaild array indexes are `[1:N]`). +- `Nn`: Number of nearest neighbors to consider. +- `kt2_array`: Array of transverse momentum squared values. +- `rapidity_array`: Array of rapidity values. +- `phi_array`: Array of azimuthal angle values. +- `R2`: Distance threshold squared for nearest neighbors. +- `nndist`: Array of nearest neighbor geometric distances. +- `nn`: Array of nearest neighbor indices. +- `nndij`: Array of metric distances between particles. + +This function updates the nearest neighbor information for the current particle `k` by considering the distances to particles `i` and `j`. It checks if the distance between `k` and `i` is smaller than the current nearest neighbor distance for `k`, and updates the nearest neighbor information accordingly. It also updates the nearest neighbor information for `i` if the distance between `k` and `i` is smaller than the current nearest neighbor distance for `i`. Finally, it checks if the nearest neighbor of `k` is the total number of particles `Nn` and updates it to `j` if necessary. + +""" Base.@propagate_inbounds function upd_nn_step!(i, j, k, N, Nn, kt2_array, rapidity_array, phi_array, R2, nndist, nn, nndij) - nnk = nn[k] + nnk = nn[k] # Nearest neighbour of k if nnk == i || nnk == j - upd_nn_nocross!(k, 1, N, rapidity_array, phi_array, R2, nndist, nn) # update dist and nn + # Our old nearest neighbour is one of the merged particles + upd_nn_nocross!(k, 1, N, rapidity_array, phi_array, R2, nndist, nn) # Update dist and nn nndij[k] = dij(k, kt2_array, nn, nndist) nnk = nn[k] end if j != i && k != i + # Update particle's nearest neighbour if it's not i and the merge step was not a beam merge Δ2 = dist(i, k, rapidity_array, phi_array) if Δ2 < nndist[k] nndist[k] = Δ2 @@ -75,18 +169,37 @@ Base.@propagate_inbounds function upd_nn_step!(i, j, k, N, Nn, kt2_array, rapidi nndist[i], nn[i] = ifelse(cond, (Δ2, k), (nndist[i], nn[i])) end + # If the previous nearest neighbour was the final jet in the array before + # the merge that was just done, this jet has now been moved in the array to + # position k (to compactify the array), so we need to update the nearest + # neighbour nnk == Nn && (nn[k] = j) end """ -This is the N2Plain jet reconstruction algorithm interface, called with an arbitrary array -of particles, which supports suitable 4-vector methods, viz. - - pt2(), phi(), rapidity(), px(), py(), pz(), energy() -for each element. - -Examples of suitable types are JetReconstruction.PseudoJet and LorentzVectorHEP. -N.B. these methods need to exist in the namespace of this package, i.e. JetReconstruction.pt2(::T), -which is already done for the two types above. + plain_jet_reconstruct(particles::Vector{T}; p = -1, R = 1.0, recombine = +) where T + +Perform jet reconstruction using the plain algorithm. + +# Arguments +- `particles::Vector{T}`: A vector of particles used for jet reconstruction, any + array of particles, which supports suitable 4-vector methods, viz. pt2(), + phi(), rapidity(), px(), py(), pz(), energy(), can be used. for each element. +- `p::Int=-1`: The integer value used for jet reconstruction. +- `R::Float64=1.0`: The radius parameter used for jet reconstruction. +- `recombine::Function=+`: The recombination function used for jet + reconstruction. + +**Note** for the `particles` argument, the 4-vector methods need to exist in the +JetReconstruction package namespace. + +# Returns +- `Vector{PseudoJet}`: A vector of reconstructed jets. + +# Example +```julia +jets = plain_jet_reconstruct(particles; p = -1, R = 1.0) +``` """ function plain_jet_reconstruct(particles::Vector{T}; p = -1, R = 1.0, recombine = +) where T # Integer p if possible @@ -106,11 +219,37 @@ function plain_jet_reconstruct(particles::Vector{T}; p = -1, R = 1.0, recombine end # Now call the actual reconstruction method, tuned for our internal EDM - _plain_jet_reconstruct(particles=recombination_particles, p=p, R=R, recombine=recombine) + _plain_jet_reconstruct(particles = recombination_particles, p = p, R = R, recombine = recombine) end -function _plain_jet_reconstruct(;particles::Vector{PseudoJet}, p = -1, R = 1.0, recombine = +) +""" + _plain_jet_reconstruct(; particles::Vector{PseudoJet}, p = -1, R = 1.0, recombine = +) + +This is the internal implementation of jet reconstruction using the plain +algorithm. It takes a vector of `particles` representing the input particles and +reconstructs jets based on the specified parameters. Here the particles must be +of type `PseudoJet`. + +Users of the package should use the `plain_jet_reconstruct` function as their +entry point to this jet reconstruction. + +The power value maps to specific pp jet reconstruction algorithms: -1 = AntiKt, +0 = Cambridge/Aachen, 1 = Inclusive Kt. + +# Arguments +- `particles`: A vector of `PseudoJet` objects representing the input particles. +- `p=-1`: The power to which the transverse momentum (`pt`) of each particle is + raised. +- `R=1.0`: The jet radius parameter. +- `recombine`: The recombination function used to merge two jets. Default is `+` + (additive recombination). + +# Returns +- `clusterseq`: The resulting `ClusterSequence` object representing the + reconstructed jets. +""" +function _plain_jet_reconstruct(; particles::Vector{PseudoJet}, p = -1, R = 1.0, recombine = +) # Bounds N::Int = length(particles) # Parameters @@ -151,12 +290,13 @@ function _plain_jet_reconstruct(;particles::Vector{PseudoJet}, p = -1, R = 1.0, # Extremely odd - having these @debug statements present causes a performance # degradation of ~140μs per event on my M2 mac (20%!), even when no debugging is used # so they need to be completely commented out... - #@debug "Beginning iteration $iteration" + # Findmin and add back renormalisation to distance dij_min, i = fast_findmin(nndij, N) @fastmath dij_min /= R2 j::Int = nn[i] + #@debug "Closest compact jets are $i ($(clusterseq_index[i])) and $j ($(clusterseq_index[j]))" if i != j # Merge jets i and j @@ -170,7 +310,7 @@ function _plain_jet_reconstruct(;particles::Vector{PseudoJet}, p = -1, R = 1.0, hist_j = clusterseq.jets[clusterseq_index[j]]._cluster_hist_index # Recombine i and j into the next jet - push!(clusterseq.jets, + push!(clusterseq.jets, recombine(clusterseq.jets[clusterseq_index[i]], clusterseq.jets[clusterseq_index[j]])) # Get its index and the history index newjet_k = length(clusterseq.jets) @@ -180,7 +320,7 @@ function _plain_jet_reconstruct(;particles::Vector{PseudoJet}, p = -1, R = 1.0, add_step_to_history!(clusterseq, minmax(hist_i, hist_j)..., newjet_k, dij_min) # Update the compact arrays, reusing the i-th slot - kt2_array[i] = pt2(clusterseq.jets[newjet_k]) ^ p + kt2_array[i] = pt2(clusterseq.jets[newjet_k])^p rapidity_array[i] = rapidity(clusterseq.jets[newjet_k]) phi_array[i] = phi(clusterseq.jets[newjet_k]) clusterseq_index[i] = newjet_k diff --git a/src/Pseudojet.jl b/src/Pseudojet.jl index 1d3a739..c2d307a 100644 --- a/src/Pseudojet.jl +++ b/src/Pseudojet.jl @@ -1,31 +1,52 @@ -# Copyright (c) 2022-2023, Philippe Gras and CERN +# Adapted from PseudoJet class of c++ code of Fastjet (https://fastjet.fr, +# hep-ph/0512210, arXiv:1111.6097) # -# Adapted from PseudoJet class of c++ code from the Fastjet -# software (https://fastjet.fr, hep-ph/0512210, arXiv:1111.6097) +# Copyright (c) 2005-2020, Matteo Cacciari, Gavin P. Salam and Gregory Soyez # -# Copyright (c) 2005-2020, Matteo Cacciari, Gavin P. Salam and Gregory -# Soyez -# -# Some implementation is taken from LorentzVectorHEP.jl, (c) Jerry Ling +# Some of the implementation is taken from LorentzVectorHEP.jl, (c) Jerry Ling """Interface for composite types that includes fields px, py, py, and E that represents the components of a four-momentum vector.""" abstract type FourMomentum end -# Used to protect against parton-level events where pt can be zero -# for some partons, giving rapidity=infinity. KtJet fails in those cases. +"""Used to protect against parton-level events where pt can be zero +for some partons, giving rapidity=infinity. KtJet fails in those cases.""" const _MaxRap = 1e5 +"""Used to signal that the phi value has not yet been computed.""" const _invalid_phi = -100.0 + +"""Used to signal that the rapidity value has not yet been computed.""" const _invalid_rap = -1.e200 # @ingroup basic_classes # \class PseudoJet # Class to contain pseudojets, including minimal information of use to # jet-clustering routines. + + +""" + mutable struct PseudoJet <: FourMomentum + +The `PseudoJet` struct represents a pseudojet, a four-momentum object used in +jet reconstruction algorithms. Additonal information for the link back into the +history of the clustering is stored in the `_cluster_hist_index` field. There is +caching of the more expensive calculations for rapidity and azimuthal angle. + +# Fields +- `px::Float64`: The x-component of the momentum. +- `py::Float64`: The y-component of the momentum. +- `pz::Float64`: The z-component of the momentum. +- `E::Float64`: The energy component of the momentum. +- `_cluster_hist_index::Int`: The index of the cluster history. +- `_pt2::Float64`: The squared transverse momentum. +- `_inv_pt2::Float64`: The inverse squared transverse momentum. +- `_rap::Float64`: The rapidity. +- `_phi::Float64`: The azimuthal angle. + +""" mutable struct PseudoJet <: FourMomentum - # construct a pseudojet from explicit components px::Float64 py::Float64 pz::Float64 @@ -38,22 +59,77 @@ mutable struct PseudoJet <: FourMomentum end +""" + PseudoJet(px::Float64, py::Float64, pz::Float64, E::Float64, + _cluster_hist_index::Int, + pt2::Float64) + +Constructs a PseudoJet object with the given momentum components and energy and +history index. + +# Arguments +- `px::Float64`: The x-component of the momentum. +- `py::Float64`: The y-component of the momentum. +- `pz::Float64`: The z-component of the momentum. +- `E::Float64`: The energy. +- `_cluster_hist_index::Int`: The cluster history index. +- `pt2::Float64`: The transverse momentum squared. + +# Returns +A `PseudoJet` object. +""" PseudoJet(px::Float64, py::Float64, pz::Float64, E::Float64, _cluster_hist_index::Int, pt2::Float64) = PseudoJet(px, py, pz, E, _cluster_hist_index, pt2, 1.0 / pt2, _invalid_rap, _invalid_phi) +""" + PseudoJet(px::Float64, py::Float64, pz::Float64, E::Float64) + +Constructs a PseudoJet object with the given momentum components and energy. + +# Arguments +- `px::Float64`: The x-component of the momentum. +- `py::Float64`: The y-component of the momentum. +- `pz::Float64`: The z-component of the momentum. +- `E::Float64`: The energy. + +# Returns +A PseudoJet object. +""" PseudoJet(px::Float64, py::Float64, pz::Float64, E::Float64) = PseudoJet(px, py, pz, E, 0, px^2 + py^2) + import Base.show +""" + show(io::IO, jet::PseudoJet) + +Print a `PseudoJet` object to the specified IO stream. + +# Arguments +- `io::IO`: The IO stream to which the information will be printed. +- `jet::PseudoJet`: The `PseudoJet` object whose information will be printed. +""" show(io::IO, jet::PseudoJet) = begin print(io, "Pseudojet(px: ", jet.px, " py: ", jet.py, " pz: ", jet.pz, " E: ", jet.E, "; ", "pt: ", sqrt(jet._pt2), " eta: ", rapidity(jet), " phi: ", phi(jet), ", m: ", m(jet), ")") end +""" + set_momentum!(j::PseudoJet, px, py, pz, E) + +Set the momentum components and energy of a `PseudoJet` object. + +# Arguments +- `j::PseudoJet`: The `PseudoJet` object to set the momentum for. +- `px`: The x-component of the momentum. +- `py`: The y-component of the momentum. +- `pz`: The z-component of the momentum. +- `E`: The energy of the particle. +""" set_momentum!(j::PseudoJet, px, py, pz, E) = begin j.px = px j.py = py @@ -61,14 +137,43 @@ set_momentum!(j::PseudoJet, px, py, pz, E) = begin j.E = E j._pt2 = px^2 + py^2 j._inv_pt2 = 1.0 / j._pt2 - j._rap = _invalid_eta + j._rap = _invalid_rap j._phi = _invalid_phi end +""" + _ensure_valid_rap_phi(p::PseudoJet) + +Ensure that the rapidity and azimuthal angle of the PseudoJet `p` are valid. If +the azimuthal angle is invalid (used as a proxy for both variables), they are +set to a valid value using `_set_rap_phi!`. + +# Arguments +- `p::PseudoJet`: The PseudoJet object to ensure valid rapidity and azimuthal + angle for. +""" _ensure_valid_rap_phi(p::PseudoJet) = p._phi == _invalid_phi && _set_rap_phi!(p) -_set_rap_phi!(p::PseudoJet) = begin +""" +_set_rap_phi!(p::PseudoJet) + +Set the rapidity and azimuthal angle of the PseudoJet `p`. + +# Arguments +- `p::PseudoJet`: The PseudoJet object for which to set the rapidity and + azimuthal angle. +# Description +This function calculates and sets the rapidity and azimuthal angle of the +PseudoJet `p` based on its momentum components. The rapidity is calculated in a +way that is insensitive to roundoff errors when the momentum components are +large. If the PseudoJet represents a point with infinite rapidity, a large +number is assigned to the rapidity in order to lift the degeneracy between +different zero-pt momenta. + +Note - the ϕ angle is calculated in the range [0, 2π). +""" +_set_rap_phi!(p::PseudoJet) = begin p._phi = p._pt2 == 0.0 ? 0.0 : atan(p.py, p.px) if p._phi < 0.0 p._phi += 2π @@ -97,36 +202,135 @@ _set_rap_phi!(p::PseudoJet) = begin nothing end +""" + phi(p::PseudoJet) + +Compute the ϕ angle of a `PseudoJet` object `p`. + +Note this function is a wrapper for `phi_02pi(p)`. + +# Arguments +- `p::PseudoJet`: The `PseudoJet` object for which to compute the azimuthal angle. + +# Returns +- The azimuthal angle of `p` in the range [0, 2π). +""" phi(p::PseudoJet) = phi_02pi(p) +""" + phi_02pi(p::PseudoJet) + +Compute the azimuthal angle of a `PseudoJet` object `p` in the range [0, 2π). + +# Arguments +- `p::PseudoJet`: The `PseudoJet` object for which to compute the azimuthal angle. + +# Returns +- The azimuthal angle of `p` in the range [0, 2π). +""" phi_02pi(p::PseudoJet) = begin _ensure_valid_rap_phi(p) return p._phi end +""" + rapidity(p::PseudoJet) + +Compute the rapidity of a `PseudoJet` object. + +# Arguments +- `p::PseudoJet`: The `PseudoJet` object for which to compute the rapidity. + +# Returns +The rapidity of the `PseudoJet` object. +""" rapidity(p::PseudoJet) = begin _ensure_valid_rap_phi(p) return p._rap end +""" + pt2(p::PseudoJet) + +Get the squared transverse momentum of a PseudoJet. + +# Arguments +- `p::PseudoJet`: The PseudoJet object. + +# Returns +- The squared transverse momentum of the PseudoJet. +""" pt2(p::PseudoJet) = p._pt2 -"Returns the scalar transverse momentum" +""" + pt(p::PseudoJet) + +Compute the scalar transverse momentum (pt) of a PseudoJet. + +# Arguments +- `p::PseudoJet`: The PseudoJet object for which to compute the transverse momentum. + +# Returns +- The transverse momentum (pt) of the PseudoJet. +""" pt(p::PseudoJet) = sqrt(p._pt2) -"Returns the squared invariant mass" +""" + m2(p::PseudoJet) + +Calculate the invariant mass squared (m^2) of a PseudoJet. + +# Arguments +- `p::PseudoJet`: The PseudoJet object for which to calculate the invariant mass squared. + +# Returns +- The invariant mass squared (m^2) of the PseudoJet. +""" m2(p::PseudoJet) = (p.E + p.pz) * (p.E - p.pz) - p._pt2 -"Returns the magnitude of the momentum, |p|" +""" + mag(p::PseudoJet) + +Return the magnitude of the momentum of a `PseudoJet`, `|p|`. + +# Arguments +- `p::PseudoJet`: The `PseudoJet` object for which to compute the magnitude. + +# Returns +The magnitude of the `PseudoJet` object. +""" mag(p::PseudoJet) = sqrt(muladd(p.px, p.px, p.py^2) + p.pz^2) + +""" + CosTheta(p::PseudoJet) + +Compute the cosine of the angle between the momentum vector `p` and the z-axis. + +# Arguments +- `p::PseudoJet`: The PseudoJet object representing the momentum vector. + +# Returns +- The cosine of the angle between `p` and the z-axis. +""" @inline function CosTheta(p::PseudoJet) fZ = p.pz ptot = mag(p) return ifelse(ptot == 0.0, 1.0, fZ / ptot) end -"Returns pseudorapidity, η" + +""" + eta(p::PseudoJet) + +Compute the pseudorapidity (η) of a PseudoJet. + +# Arguments +- `p::PseudoJet`: The PseudoJet object for which to compute the pseudorapidity. + +# Returns +- The pseudorapidity (η) of the PseudoJet. +""" function eta(p::PseudoJet) cosTheta = CosTheta(p) (cosTheta^2 < 1.0) && return -0.5 * log((1.0 - cosTheta) / (1.0 + cosTheta)) @@ -136,25 +340,117 @@ function eta(p::PseudoJet) fZ > 0.0 && return 10e10 return -10e10 end + +""" + const η = eta + +Alias for the pseudorapidity function, `eta`. +""" const η = eta -# returns the invariant mass -# (If m2() is negative then -sqrt(-m2()) is returned, as in CLHEP, SciKitHEP Particle and ROOT) + +""" + m(p::PseudoJet) + +Compute the invariant mass of a `PseudoJet` object. By convention if m^2 < 0, +then -sqrt{(-m^2)} is returned. + +# Arguments +- `p::PseudoJet`: The `PseudoJet` object for which to compute the invariant + mass. + +# Returns +The invariant mass of the `PseudoJet` object. +""" m(p::PseudoJet) = begin x = m2(p) x < 0.0 ? -sqrt(-x) : sqrt(x) end -# Ensure we have accessors for jet parameters +""" + mass(p::PseudoJet) + +Compute the invariant mass (alias for `m(p)`). + +# Arguments +- `p::PseudoJet`: The PseudoJet object for which to compute the mass. + +# Returns +- The mass of the PseudoJet. +""" +mass(p::PseudoJet) = m(p) + +"""Alias for `m2` function""" +const mass2 = m2 + +""" + px(p::PseudoJet) + +Return the x-component of the momentum of a `PseudoJet`. + +# Arguments +- `p::PseudoJet`: The `PseudoJet` object. + +# Returns +- The x-component of the momentum of the `PseudoJet`. +""" px(p::PseudoJet) = p.px + +""" + py(p::PseudoJet) + +Return the y-component of the momentum of a `PseudoJet`. + +# Arguments +- `p::PseudoJet`: The `PseudoJet` object. + +# Returns +- The y-component of the momentum of the `PseudoJet`. +""" py(p::PseudoJet) = p.py + +""" + pz(p::PseudoJet) + +Return the z-component of the momentum of a `PseudoJet`. + +# Arguments +- `p::PseudoJet`: The `PseudoJet` object. + +# Returns +- The z-component of the momentum of the `PseudoJet`. +""" pz(p::PseudoJet) = p.pz -mass(p::PseudoJet) = m(p) -const mass2 = m2 + +""" + energy(p::PseudoJet) + +Return the energy of a `PseudoJet`. + +# Arguments +- `p::PseudoJet`: The `PseudoJet` object. + +# Returns +- The energy of the `PseudoJet`. +""" energy(p::PseudoJet) = p.E + + import Base.+; +""" + +(j1::PseudoJet, j2::PseudoJet) + +Addition operator for `PseudoJet` objects. + +# Arguments +- `j1::PseudoJet`: The first `PseudoJet` object. +- `j2::PseudoJet`: The second `PseudoJet` object. + +# Returns +A new `PseudoJet` object with the sum of the momenta and energy of `j1` and `j2`. +""" +(j1::PseudoJet, j2::PseudoJet) = begin PseudoJet(j1.px + j2.px, j1.py + j2.py, j1.pz + j2.pz, j1.E + j2.E) diff --git a/src/Serialize.jl b/src/Serialize.jl index 3e5dda8..4d1c51a 100644 --- a/src/Serialize.jl +++ b/src/Serialize.jl @@ -1,7 +1,31 @@ """ -`savejets(filename, jets; format="px py pz E")` + savejets(filename, jets; format="px py pz E") -Saves the given `jets` into a file with the given `filename`. Each line contains information about a single jet and is formatted according to the `format` string which defaults to `"px py pz E"` but can also contain other values in any order: `"pt2"` for pt^2, `"phi"` for azimuth, `"rapidity"` for rapidity. It is strongly NOT recommended to put something other than values and (possibly custom) separators in the `format` string. +Save jet data to a file. + +# Arguments +- `filename`: The name of the file to save the jet data to. +- `jets`: An array of jet objects to save. +- `format="px py pz E"`: (optional) A string specifying the format of the jet + data to save. The default format is "px py pz E". + +# Details +This function saves jet data to a file in a specific format. Each line in the +file represents a jet and contains the information about the jet in the +specified format. The format string can include the following placeholders: + +- "E" or "energy": Jet energy +- "px": Momentum along the x-axis +- "py": Momentum along the y-axis +- "pz": Momentum along the z-axis +- "pt2": Square of the transverse momentum +- "phi": Azimuth angle +- "rapidity": Rapidity + +Lines starting with '#' are treated as comments and are ignored. + +It is strongly NOT recommended to put something other than values and (possibly +custom) separators in the `format` string. """ function savejets(filename, jets; format="px py pz E") symbols = Dict( @@ -35,19 +59,27 @@ function savejets(filename, jets; format="px py pz E") end """ - loadjets!(filename, jets; splitby=isspace, constructor=(px,py,pz,E)->LorentzVector(E,px,py,pz), dtype=Float64) + loadjets!(filename, jets; splitby=isspace, constructor=(px,py,pz,E)->LorentzVectorHEP(E,px,py,pz), dtype=Float64) -Loads the `jets` from a file. Ignores lines that start with `'#'`. Each line gets processed in the following way: the line is split using `split(line, splitby)` or simply `split(line)` by default. Every value in this line is then converted to the `dtype` (which is `Float64` by default). These values are then used as arguments for the `constructor` function which should produce individual jets. By default, the `constructor` constructs Lorentz vectors. +Loads the `jets` from a file. Ignores lines that start with `'#'`. Each line +gets processed in the following way: the line is split using `split(line, +splitby)` or simply `split(line)` by default. Every value in this line is then +converted to the `dtype` (which is `Float64` by default). These values are then +used as arguments for the `constructor` function which should produce individual +jets. By default, the `constructor` constructs Lorentz vectors. -Everything that was already in `jets` is not affected as we only use `push!` on it. +Everything that was already in `jets` is not affected as we only use `push!` on +it. + +# Example ```julia -# example that loads jets from two files into one array +# Load jets from two files into one array jets = [] loadjets!("myjets1.dat", jets) loadjets!("myjets2.dat", jets) ``` """ -function loadjets!(filename, jets; splitby=isspace, constructor=(px,py,pz,E)->LorentzVector(E,px,py,pz), dtype=Float64) +function loadjets!(filename, jets; splitby=isspace, constructor=(px,py,pz,E)->LorentzVectorHEP(E,px,py,pz), dtype=Float64) open(filename, "r") do file for line in eachline(file) if line[1] != '#' @@ -63,15 +95,23 @@ function loadjets!(filename, jets; splitby=isspace, constructor=(px,py,pz,E)->Lo end """ - loadjets(filename; splitby=isspace, constructor=(px,py,pz,E)->LorentzVector(E,px,py,pz), VT=LorentzVector) -> jets + loadjets(filename; splitby=isspace, constructor=(px,py,pz,E)->LorentzVectorHEP(E,px,py,pz), VT=LorentzVector) -Loads the `jets` from a file, where each element of `jets` is of type `VT`. Ignores lines that start with `'#'`. Each line gets processed in the following way: the line is split using `split(line, splitby)` or simply `split(line)` by default. These values are then used as arguments for the `constructor` function which should produce individual jets of the `VT` type. By default, the `constructor` constructs Lorentz vectors. +Load jets from a file. + +# Arguments +- `filename`: The name of the file to load jets from. +- `splitby`: The delimiter used to split the data in the file. Default is + `isspace`. +- `constructor`: A function that constructs a `VT` object from the jet data. + Default is `(px,py,pz,E)->LorentzVector(E,px,py,pz)`. +- `VT`: The type of the vector used to store the jet data. Default is + `LorentzVector`. + +# Returns +- A vector of `VT` objects representing the loaded jets. -```julia -# example -jets = loadjets("myjets1.dat") -``` """ -function loadjets(filename; splitby=isspace, constructor=(px,py,pz,E)->LorentzVector(E,px,py,pz), VT=LorentzVector) +function loadjets(filename; splitby=isspace, constructor=(px,py,pz,E)->LorentzVectorHEP(E,px,py,pz), VT=LorentzVector) loadjets!(filename, VT[], splitby=splitby, constructor=constructor) end diff --git a/src/TiledAlgoLL.jl b/src/TiledAlgoLL.jl index 3d5eb45..53a6cf1 100644 --- a/src/TiledAlgoLL.jl +++ b/src/TiledAlgoLL.jl @@ -10,10 +10,18 @@ using LoopVectorization # Include struct definitions and basic operations include("TiledAlgoLLStructs.jl") + """ -Initialise the clustering history in a standard way, -Takes as input the list of stable particles as input -Returns the history and the total event energy. + initial_history(particles) + +Create an initial history for the given particles. + +# Arguments +- `particles`: The initial vector of stable particles. + +# Returns +- `history`: An array of `HistoryElement` objects. +- `Qtot`: The total energy in the event. """ function initial_history(particles) # reserve sufficient space for everything @@ -34,14 +42,40 @@ function initial_history(particles) history, Qtot end -"""Computes distance in the (eta,phi)-plane -between two jets.""" +""" + _tj_dist(jetA, jetB) + +Compute the geometric distance in the (y, ϕ)-plane between two jets in the TiledAlgoLL module. + +# Arguments +- `jetA`: The first jet. +- `jetB`: The second jet. + +# Returns +The squared distance between `jetA` and `jetB`. + +# Examples +""" _tj_dist(jetA, jetB) = begin dphi = π - abs(π - abs(jetA.phi - jetB.phi)) deta = jetA.eta - jetB.eta - return dphi * dphi + deta * deta + return muladd(dphi, dphi, deta * deta) end + +""" + _tj_diJ(jet) + +Compute the dij metric value for a given jet. + +# Arguments +- `jet`: The input jet. + +# Returns +- The dij value for the jet. + +# Example +""" _tj_diJ(jet) = begin kt2 = jet.kt2 if isvalid(jet.NN) && jet.NN.kt2 < kt2 @@ -51,10 +85,22 @@ _tj_diJ(jet) = begin end -"""Return the tile index corresponding to the given eta,phi point""" +""" + tile_index(tiling_setup, eta::Float64, phi::Float64) + +Compute the tile index for a given (eta, phi) coordinate. + +# Arguments +- `tiling_setup`: The tiling setup object containing the tile size and number of tiles. +- `eta::Float64`: The eta coordinate. +- `phi::Float64`: The phi coordinate. + +# Returns +The tile index corresponding to the (eta, phi) coordinate. +""" tile_index(tiling_setup, eta::Float64, phi::Float64) = begin # Use clamp() to restrict to the correct ranges - # - eta can be out of range by construction (open end bins) + # - eta can be out of range by construction (open ended bins) # - phi is protection against bad rounding ieta = clamp(1 + unsafe_trunc(Int, (eta - tiling_setup._tiles_eta_min) / tiling_setup._tile_size_eta), 1, tiling_setup._n_tiles_eta) iphi = clamp(unsafe_trunc(Int, phi / tiling_setup._tile_size_phi), 0, tiling_setup._n_tiles_phi) @@ -62,7 +108,24 @@ tile_index(tiling_setup, eta::Float64, phi::Float64) = begin end -"""Initialise a tiled jet from a PseudoJet (using an index into our ClusterSequence)""" +""" + tiledjet_set_jetinfo!(jet::TiledJet, clusterseq::ClusterSequence, tiling::Tiling, jets_index, R2, p) + +Initialise a tiled jet from a PseudoJet (using an index into our ClusterSequence) + +Arguments: +- `jet::TiledJet`: The TiledJet object to set the information for. +- `clusterseq::ClusterSequence`: The ClusterSequence object containing the jets. +- `tiling::Tiling`: The Tiling object containing the tile information. +- `jets_index`: The index of the jet in the ClusterSequence. +- `R2`: The jet radius parameter squared. +- `p`: The power to raise the pt2 value to. + +This function sets the eta, phi, kt2, jets_index, NN_dist, NN, tile_index, previous, and next fields of the TiledJet object. + +Returns: +- `nothing` +""" tiledjet_set_jetinfo!(jet::TiledJet, clusterseq::ClusterSequence, tiling::Tiling, jets_index, R2, p) = begin @inbounds jet.eta = rapidity(clusterseq.jets[jets_index]) @inbounds jet.phi = phi_02pi(clusterseq.jets[jets_index]) @@ -87,6 +150,31 @@ end """Full scan for nearest neighbours""" + + +""" + set_nearest_neighbours!(clusterseq::ClusterSequence, tiling::Tiling, tiledjets::Vector{TiledJet}) + +This function sets the nearest neighbor information for all jets in the +`tiledjets` vector. + +# Arguments +- `clusterseq::ClusterSequence`: The cluster sequence object. +- `tiling::Tiling`: The tiling object. +- `tiledjets::Vector{TiledJet}`: The vector of tiled jets. + +# Returns +- `NNs::Vector{TiledJet}`: The vector of nearest neighbor jets. +- `diJ::Vector{Float64}`: The vector of diJ values. + +The function iterates over each tile in the `tiling` and sets the nearest +neighbor information for each jet in the tile. It then looks for neighbor jets +in the neighboring tiles and updates the nearest neighbor information +accordingly. Finally, it creates the diJ table and returns the vectors of +nearest neighbor jets and diJ values. + +Note: The diJ values are calculated as the kt distance multiplied by R^2. +""" function set_nearest_neighbours!(clusterseq::ClusterSequence, tiling::Tiling, tiledjets::Vector{TiledJet}) # Setup the initial nearest neighbour information for tile in tiling.tiles @@ -136,7 +224,7 @@ function set_nearest_neighbours!(clusterseq::ClusterSequence, tiling::Tiling, ti for i in eachindex(diJ) jetA = tiledjets[i] diJ[i] = _tj_diJ(jetA) # kt distance * R^2 - # our compact diJ table will not be in one-to-one corresp. with non-compact jets, + # Our compact diJ table will not be in one-to-one corresp. with non-compact jets, # so set up bi-directional correspondence here. @inbounds NNs[i] = jetA jetA.dij_posn = i @@ -144,9 +232,30 @@ function set_nearest_neighbours!(clusterseq::ClusterSequence, tiling::Tiling, ti NNs, diJ end -"""Carries out the bookkeeping associated with the step of recombining -jet_i and jet_j (assuming a distance dij) and returns the index -of the recombined jet, newjet_k.""" + +""" + do_ij_recombination_step!(clusterseq::ClusterSequence, jet_i, jet_j, dij, recombine=+) + +Perform the bookkeeping associated with the step of recombining jet_i and jet_j +(assuming a distance dij). + +# Arguments +- `clusterseq::ClusterSequence`: The cluster sequence object. +- `jet_i`: The index of the first jet to be recombined. +- `jet_j`: The index of the second jet to be recombined. +- `dij`: The distance between the two jets. +- `recombine=+`: The recombination function to be used. Default is addition. + +# Returns +- `newjet_k`: The index of the newly created jet. + +# Description +This function performs the i-j recombination step in the cluster sequence. It +creates a new jet by recombining the first two jets using the specified +recombination function. The new jet is then added to the cluster sequence. The +function also updates the indices and history information of the new jet and +sorts out the history. +""" do_ij_recombination_step!(clusterseq::ClusterSequence, jet_i, jet_j, dij, recombine=+) = begin # Create the new jet by recombining the first two with # the E-scheme @@ -169,8 +278,18 @@ do_ij_recombination_step!(clusterseq::ClusterSequence, jet_i, jet_j, dij, recomb newjet_k end -"""Carries out the bookkeeping associated with the step of recombining -jet_i with the beam (i.e., finalising the jet)""" + +""" + do_iB_recombination_step!(clusterseq::ClusterSequence, jet_i, diB) + +Bookkeeping for recombining a jet with the beam (i.e., finalising the jet) by +adding a step to the history of the cluster sequence. + +# Arguments +- `clusterseq::ClusterSequence`: The cluster sequence object. +- `jet_i`: The index of the jet. +- `diB`: The diB value. +""" do_iB_recombination_step!(clusterseq::ClusterSequence, jet_i, diB) = begin # Recombine the jet with the beam add_step_to_history!(clusterseq, clusterseq.jets[jet_i]._cluster_hist_index, BeamJet, @@ -178,15 +297,24 @@ do_iB_recombination_step!(clusterseq::ClusterSequence, jet_i, diB) = begin end -"""Adds to the vector tile_union the tiles that are in the neighbourhood -of the specified tile_index, including itself and whose tagged status are -false ---start adding from position n_near_tiles-1, and increase n_near_tiles as -you go along. When a neighbour is added its tagged status is set to true. +""" + add_untagged_neighbours_to_tile_union(center_index, tile_union, n_near_tiles, tiling) -Returns the updated number of near_tiles.""" -function add_untagged_neighbours_to_tile_union(center_index, - tile_union, n_near_tiles, - tiling) +Adds to the vector tile_union the tiles that are in the neighbourhood of the +specified center_index, including itself and whose tagged status are false - +start adding from position n_near_tiles-1, and increase n_near_tiles. When a +neighbour is added its tagged status is set to true. + +# Arguments +- `center_index`: The index of the center tile. +- `tile_union`: An array to store the indices of neighbouring tiles. +- `n_near_tiles`: The number of neighbouring tiles. +- `tiling`: The tiling object containing the tile tags. + +# Returns +The updated number of near tiles. +""" +function add_untagged_neighbours_to_tile_union(center_index, tile_union, n_near_tiles, tiling) for tile_index in surrounding(center_index, tiling) @inbounds if !tiling.tags[tile_index] n_near_tiles += 1 @@ -198,13 +326,23 @@ function add_untagged_neighbours_to_tile_union(center_index, n_near_tiles end + """ -Establish the set of tiles over which we are going to -have to run searches for updated and new nearest-neighbours -- -basically a combination of vicinity of the tiles of the two old -jets and the new jet. + find_tile_neighbours!(tile_union, jetA, jetB, oldB, tiling) + +Find the union of neighbouring tiles of `jetA`, `jetB`, and `oldB` and add them +to the `tile_union`. This established the set of tiles over which searches for +updated and new nearest-neighbours must be run -Updates tile_union and returns n_near_tiles +# Arguments +- `tile_union`: The tile union to which the neighbouring tiles will be added. +- `jetA`: The first jet. +- `jetB`: The second jet. +- `oldB`: The old second jet. +- `tiling`: The tiling information. + +# Returns +The number of neighbouring tiles added to the `tile_union`. """ function find_tile_neighbours!(tile_union, jetA, jetB, oldB, tiling) n_near_tiles = add_untagged_neighbours_to_tile_union(jetA.tile_index, @@ -224,14 +362,32 @@ end """ -Main jet reconstruction algorithm entry point for generic data types + tiled_jet_reconstruct(particles::Vector{T}; p = -1, R = 1.0, recombine = +) where {T} -`particles` must support methods px, py, pz and energy (N.B. these must be in the -JetReconstruction namespace). In particular Cartesian LorentzVector structs can -be used. +Main jet reconstruction algorithm entry point for reconstructing jets using the +tiled stragegy for generic jet type T. -If a non-standard recombination is used, it must be defined for +**Note** - if a non-standard recombination is used, it must be defined for JetReconstruction.PseudoJet, as this struct is used internally. + +## Arguments +- `particles::Vector{T}`: A vector of particles used as input for jet + reconstruction. T must support methods px, py, pz and energy (defined in the + JetReconstruction namespace) +- `p::Int = -1`: The power parameter for the jet reconstruction algorithm, thus + swiching between different algorithms. +- `R::Float64 = 1.0`: The jet radius parameter for the jet reconstruction + algorithm. +- `recombine::Function = +`: The recombination function used for combining + pseudojets. + +## Returns +- `Vector{PseudoJet}`: A vector of reconstructed jets. + +## Example +```julia +tiled_jet_reconstruct(particles::Vector{LorentzVectorHEP}; p = -1, R = 0.4, recombine = +) +``` """ function tiled_jet_reconstruct(particles::Vector{T}; p = -1, R = 1.0, recombine = +) where {T} # Here we need to populate the vector of PseudoJets that are the internal @@ -248,6 +404,31 @@ end """ Main jet reconstruction algorithm, using PseudoJet objects """ + +""" + tiled_jet_reconstruct(particles::Vector{PseudoJet}; p = -1, R = 1.0, recombine = +) where {T} + +Main jet reconstruction algorithm entry point for reconstructing jets using the +tiled stragegy for the prefered internal jet type, `PseudoJet`. + +## Arguments +- `particles::Vector{PseudoJet}`: A vector of `PseudoJet` particles used as input for jet + reconstruction. +- `p::Int = -1`: The power parameter for the jet reconstruction algorithm, thus + swiching between different algorithms. +- `R::Float64 = 1.0`: The jet radius parameter for the jet reconstruction + algorithm. +- `recombine::Function = +`: The recombination function used for combining + pseudojets. + +## Returns +- `Vector{PseudoJet}`: A vector of reconstructed jets. + +## Example +```julia +tiled_jet_reconstruct(particles::Vector{PseudoJet}; p = 1, R = 1.0, recombine = +) +``` +""" function tiled_jet_reconstruct(particles::Vector{PseudoJet}; p = -1, R = 1.0, recombine = +) # Bounds N::Int = length(particles) diff --git a/src/TiledAlgoLLStructs.jl b/src/TiledAlgoLLStructs.jl index 95857b0..9a23dcb 100644 --- a/src/TiledAlgoLLStructs.jl +++ b/src/TiledAlgoLLStructs.jl @@ -1,7 +1,27 @@ # Structures used internally by the tiled algorithm -"""Structure analogous to BriefJet, but with the extra information -needed for dealing with tiles""" +"""Structure analogous to PseudoJet, but with the extra information +needed for dealing with tiles for the tiled stragegy""" + +""" + struct TiledJet + +TiledJet represents a jet in a tiled algorithm for jet reconstruction, with +additional information to track the jet's position in the tiled structures. + +# Fields +- `id::Int`: The ID of the jet. +- `eta::Float64`: The rapidity of the jet. +- `phi::Float64`: The azimuthal angle of the jet. +- `kt2::Float64`: The transverse momentum squared of the jet. +- `NN_dist::Float64`: The distance to the nearest neighbor. +- `jets_index::Int`: The index of the jet in the jet array. +- `tile_index::Int`: The index of the tile in the tile array. +- `dij_posn::Int`: The position of this jet in the dij compact array. +- `NN::TiledJet`: The nearest neighbor. +- `previous::TiledJet`: The previous jet. +- `next::TiledJet`: The next jet. +""" mutable struct TiledJet id::Int eta::Float64 @@ -13,18 +33,21 @@ mutable struct TiledJet tile_index::Int dij_posn::Int - "Nearest neighbour" NN::TiledJet previous::TiledJet next::TiledJet + TiledJet(::Type{Nothing}) = begin t = new(-1, 0., 0., 0., 0., -1, 0, 0) t.NN = t.previous = t.next = t t end + """ + TiledJet constructor with all fields. + """ TiledJet(id, eta, phi, kt2, NN_dist, jet_index, tile_index, dij_posn, NN, previous, next) = new(id, eta, phi, kt2, NN_dist, @@ -33,13 +56,52 @@ mutable struct TiledJet end +""" + const noTiledJet::TiledJet = TiledJet(Nothing) + +A constant variable representing a "blank" `TiledJet` object with invalid values. +""" const noTiledJet::TiledJet = TiledJet(Nothing) +""" + isvalid(t::TiledJet) + +Check if a `TiledJet` is valid, by seeing if it is not the `noTiledJet` object. + +# Arguments +- `t::TiledJet`: The `TiledJet` object to check. + +# Returns +- `Bool`: `true` if the `TiledJet` object is valid, `false` otherwise. +""" isvalid(t::TiledJet) = !(t === noTiledJet) """ -Move a TiledJet in front of a TiledJet list element + TiledJet(id) + +Constructs a `TiledJet` object with the given `id` and initializes its properties to zero. + +# Arguments +- `id`: The ID of the `TiledJet` object. + +# Returns +A `TiledJet` object with the specified `id` and values set to zero or noTiledJet. +""" +TiledJet(id) = TiledJet(id, 0., 0., 0., 0., + 0, 0, 0, + noTiledJet, noTiledJet, noTiledJet) + +""" + insert!(nextjet::TiledJet, jettomove::TiledJet) + +Inserts a `TiledJet` object into the linked list of `TiledJet` objects, before the `nextjet` object. The jet to move can be an isolated jet, a jet from another list or a jet from the same list + +# Arguments +- `nextjet::TiledJet`: The `TiledJet` object after which `jettomove` should be inserted. +- `jettomove::TiledJet`: The `TiledJet` object to be inserted. + +# Example """ insert!(nextjet::TiledJet, jettomove::TiledJet) = begin if !isnothing(nextjet) @@ -51,7 +113,14 @@ insert!(nextjet::TiledJet, jettomove::TiledJet) = begin nextjet = jettomove end -"""Detach a TiledJet from its list""" +""" + detach!(jet::TiledJet) + +Detach a `TiledJet` from its linked list by updating the `previous` and `next` pointers. + +# Arguments +- `jet::TiledJet`: The `TiledJet` object to detach. +""" detach!(jet::TiledJet) = begin if !isnothing(jet.previous) jet.previous.next = jet.next @@ -62,15 +131,33 @@ detach!(jet::TiledJet) = begin jet.next = jet.previous = noTiledJet end -TiledJet(id) = TiledJet(id, 0., 0., 0., 0., - 0, 0, 0, - noTiledJet, noTiledJet, noTiledJet) import Base.copy +""" + copy(j::TiledJet) + +Create a copy of a `TiledJet` object. + +# Arguments +- `j::TiledJet`: The `TiledJet` object to be copied. + +# Returns +A new `TiledJet` object with the same attributes as the input object. +""" copy(j::TiledJet) = TiledJet(j.id, j.eta, j.phi, j.kt2, j.NN_dist, j.jets_index, j.tile_index, j.dij_posn, j.NN, j.previous, j.next) # Iterator over a TiledJet walks along the chain of linked jets # until we reach a "noTiledJet" (which is !isvalid) + +""" + Base.iterate(tj::TiledJet) + +Iterate over a `TiledJet` object's linked list, walking over all jets +until the end (then the next jet is invalid). + +# Arguments +- `tj::TiledJet`: The `TiledJet` object to start to iterate over. +""" Base.iterate(tj::TiledJet) = begin isvalid(tj) ? (tj, tj) : nothing end @@ -80,8 +167,21 @@ end """ -Structure with the tiling parameters, as well as some bookkeeping -variables used during reconstruction + struct Tiling + +The `Tiling` struct represents a tiling configuration for jet reconstruction. + +# Fields +- `setup::TilingDef`: The tiling definition used for the configuration. +- `tiles::Matrix{TiledJet}`: A matrix of tiled jets, containing the first jet in + each tile (then the linked list of the first jet is followed to get access to + all jets in this tile). +- `positions::Matrix{Int}`: Used to track tiles that are on the edge of ϕ array, + where neighbours need to be wrapped around. +- `tags::Matrix{Bool}`: The matrix of tags indicating whether a tile is valid or + not (set to `false` initially, then `true` when the tile has been setup + properly). + """ struct Tiling setup::TilingDef @@ -90,7 +190,18 @@ struct Tiling tags::Matrix{Bool} end -"""Return a tiling setup with bookkeeping""" + +""" + Tiling(setup::TilingDef) + +Constructs a intial `Tiling` object based on the provided `setup` parameters. + +# Arguments +- `setup::TilingDef`: The setup parameters for the tiling. + +# Returns +A `Tiling` object. +""" Tiling(setup::TilingDef) = begin t = Tiling(setup, fill(noTiledJet, (setup._n_tiles_eta, setup._n_tiles_phi)), @@ -105,19 +216,24 @@ Tiling(setup::TilingDef) = begin t end +"Signal that a tile is on the left hand side of the tiling array in ϕ" const tile_left = -1 +"Signal that a tile is central for the tiling array in ϕ" const tile_central = 0 +"Signal that a tile is on the right hand side of the tiling array in ϕ" const tile_right = 1 - -const _n_tile_center = 1 -const _n_tile_left_neighbours = 4 -const _tile_right_neigbour_indices = 6:9 -const _n_tile_right_neighbours = 4 +"Number of neighbours for a tile in the tiling array (including itself)" const _n_tile_neighbours = 9 -const neigh_init = fill(nothing, _n_tile_neighbours) -"""Structure used for iterating over neighbour tiles""" +""" + struct Surrounding{N} + +Structure used for iterating over neighbour tiles. + +# Fields +- `indices::NTuple{N, Int}`: A tuple of `N` integers representing the indices. +""" struct Surrounding{N} indices::NTuple{N, Int} end @@ -136,6 +252,18 @@ Base.iterate(x::Surrounding{9}, state) = state > 9 ? nothing : (x.indices[state] import Base.length length(x::Surrounding{N}) where N = N +""" + surrounding(center::Int, tiling::Tiling) + +Compute the surrounding indices of a given center index in a tiling. + +# Arguments +- `center::Int`: The center index. +- `tiling::Tiling`: The tiling object. + +# Returns +- `Surrounding`: An object containing the surrounding indices. +""" surrounding(center::Int, tiling::Tiling) = begin # 4|6|9 # 3|1|8 @@ -159,6 +287,22 @@ surrounding(center::Int, tiling::Tiling) = begin end end +""" + rightneighbours(center::Int, tiling::Tiling) + +Compute the indices of the right neighbors of a given center index in a tiling. +This is used in the inital sweep to calculate the nearest neighbors, where the +search between jets for the nearest neighbour is bi-directional, thus when a +tile is considered only the right neighbours are needed to compare jet +distances as the left-hand tiles have been done from that tile already. + +# Arguments +- `center::Int`: The center index. +- `tiling::Tiling`: The tiling object. + +# Returns +- `Surrounding`: An object containing the indices of the right neighbors. +""" rightneighbours(center::Int, tiling::Tiling) = begin # |1|4 # | |3 @@ -175,7 +319,20 @@ rightneighbours(center::Int, tiling::Tiling) = begin end end -"""Remove a jet from a tiling""" + +""" + tiledjet_remove_from_tiles!(tiling, jet) + +Remove a jet from the given tiling structure. + +# Arguments +- `tiling`: The tiling structure from which the jet will be removed. +- `jet`: The jet to be removed from the tiling structure. + +# Description +This function removes a jet from the tiling structure. It adjusts the linked list +to be consistent with the removal of the jet. +""" tiledjet_remove_from_tiles!(tiling, jet) = begin if !isvalid(jet.previous) # We are at head of the tile, so reset it. diff --git a/src/TiledAlgoUtils.jl b/src/TiledAlgoUtils.jl index 96665b1..f2d5f7e 100644 --- a/src/TiledAlgoUtils.jl +++ b/src/TiledAlgoUtils.jl @@ -1,32 +1,63 @@ # Common functions and structures for Tiled reconstruction algorithms -"""Tiling definition parameters""" +""" + struct TilingDef + +A struct representing the definition of a spcific tiling scheme. + +# Fields +- `_tiles_eta_min::Float64`: The minimum rapidity of the tiles. +- `_tiles_eta_max::Float64`: The maximum rapidity of the tiles. +- `_tile_size_eta::Float64`: The size of a tile in rapidity (usually R^2). +- `_tile_size_phi::Float64`: The size of a tile in phi (usually a bit more than R^2). +- `_n_tiles_eta::Int`: The number of tiles across rapidity. +- `_n_tiles_phi::Int`: The number of tiles across phi. +- `_n_tiles::Int`: The total number of tiles. +- `_tiles_ieta_min::Int`: The minimum rapidity tile index. +- `_tiles_ieta_max::Int`: The maximum rapidity tile index. + +# Constructor + TilingDef(_tiles_eta_min, _tiles_eta_max, _tile_size_eta, _tile_size_phi, + _n_tiles_eta, _n_tiles_phi, _tiles_ieta_min, _tiles_ieta_max) + +Constructs a `TilingDef` object with the given parameters. +""" struct TilingDef - _tiles_eta_min::Float64 # Minimum rapidity - _tiles_eta_max::Float64 # Maximum rapidity - _tile_size_eta::Float64 # Size of a tile in rapidity (usually R^2) - _tile_size_phi::Float64 # Size of a tile in phi (usually a bit more than R^2) - _n_tiles_eta::Int # Number of tiles across rapidity - _n_tiles_phi::Int # Number of tiles across phi - _n_tiles::Int # Total number of tiles - _tiles_ieta_min::Int # Min_rapidity / rapidity tile size (needed?) - _tiles_ieta_max::Int # Max_rapidity / rapidity tile size (needed?) - - # Use an inner constructor as _n_tiles and _tile_linear_indexes - # are defined by the other values + _tiles_eta_min::Float64 + _tiles_eta_max::Float64 + _tile_size_eta::Float64 + _tile_size_phi::Float64 + _n_tiles_eta::Int + _n_tiles_phi::Int + _n_tiles::Int + _tiles_ieta_min::Int + _tiles_ieta_max::Int + function TilingDef(_tiles_eta_min, _tiles_eta_max, _tile_size_eta, _tile_size_phi, _n_tiles_eta, _n_tiles_phi, _tiles_ieta_min, _tiles_ieta_max) new(_tiles_eta_min, _tiles_eta_max, _tile_size_eta, _tile_size_phi, - _n_tiles_eta, _n_tiles_phi, _n_tiles_eta*_n_tiles_phi, _tiles_ieta_min, _tiles_ieta_max, - ) + _n_tiles_eta, _n_tiles_phi, _n_tiles_eta*_n_tiles_phi, _tiles_ieta_min, _tiles_ieta_max) end end + """ -Determine an extent for the rapidity tiles, by binning in -rapidity and collapsing the outer bins until they have about -1/4 the number of particles as the maximum bin. This is the -heuristic which is used by FastJet. + determine_rapidity_extent(eta::Vector{T}) where T <: AbstractFloat + +Calculate the minimum and maximum rapidities based on the input vector `eta`. +The function determines the rapidity extent by binning the multiplicities as a +function of rapidity and finding the minimum and maximum rapidities such that +the edge bins contain a certain fraction (~1/4) of the busiest bin and a minimum +number of particles. + +This is the heuristic which is used by FastJet (inline comments are from FastJet). + +## Arguments +- `eta::Vector{T}`: A vector of rapidity values. + +## Returns +- `minrap::T`: The minimum rapidity value. +- `maxrap::T`: The maximum rapidity value. """ function determine_rapidity_extent(eta::Vector{T}) where T <: AbstractFloat length(eta) == 0 && return 0.0, 0.0 @@ -37,7 +68,7 @@ function determine_rapidity_extent(eta::Vector{T}) where T <: AbstractFloat # Get the minimum and maximum rapidities and at the same time bin # the multiplicities as a function of rapidity to help decide how - # far out it's worth going + # far out it is worth going minrap = maxrap = eta[1] ibin = 0 for y in eta @@ -106,8 +137,26 @@ function determine_rapidity_extent(eta::Vector{T}) where T <: AbstractFloat minrap, maxrap end + """ -Setup the tiling parameters for this recontstruction + setup_tiling(eta::Vector{T}, Rparam::AbstractFloat) where T <: AbstractFloat + +This function sets up the tiling parameters for a reconstruction given a vector +of rapidities `eta` and a radius parameter `Rparam`. + +# Arguments +- `eta::Vector{T}`: A vector of rapidities. +- `Rparam::AbstractFloat`: The jet radius parameter. + +# Returns +- `tiling_setup`: A `TilingDef` object containing the tiling setup parameters. + +# Description +The function first decides the tile sizes based on the `Rparam` value. It then +determines the number of tiles in the phi direction (`n_tiles_phi`) based on the +tile size. Next, it determines the rapidity extent of the input `eta` vector and +adjusts the values accordingly. Finally, it creates a `TilingDef` object with +the calculated tiling parameters and returns it. """ function setup_tiling(eta::Vector{T}, Rparam::AbstractFloat) where T <: AbstractFloat # First decide tile sizes (with a lower bound to avoid huge memory use with @@ -137,8 +186,20 @@ function setup_tiling(eta::Vector{T}, Rparam::AbstractFloat) where T <: Abstract tiling_setup end + """ -Return the geometric distance between a pair of (eta,phi) coordinates + geometric_distance(eta1::AbstractFloat, phi1::AbstractFloat, eta2::AbstractFloat, phi2::AbstractFloat) + +Compute the geometric distance between two points in the rap-phi plane. + +# Arguments +- `eta1::AbstractFloat`: The eta coordinate of the first point. +- `phi1::AbstractFloat`: The phi coordinate of the first point. +- `eta2::AbstractFloat`: The eta coordinate of the second point. +- `phi2::AbstractFloat`: The phi coordinate of the second point. + +# Returns +- `distance::Float64`: The geometric distance between the two points. """ geometric_distance(eta1::AbstractFloat, phi1::AbstractFloat, eta2::AbstractFloat, phi2::AbstractFloat) = begin δeta = eta2 - eta1 @@ -146,8 +207,23 @@ geometric_distance(eta1::AbstractFloat, phi1::AbstractFloat, eta2::AbstractFloat return δeta * δeta + δphi * δphi end + """ -Return the dij metric distance between a pair of pseudojets + get_dij_dist(nn_dist, kt2_1, kt2_2, R2) + +Compute the dij metric distance between two jets. + +# Arguments +- `nn_dist`: The nearest-neighbor distance between two jets. +- `kt2_1`: The squared momentum metric value of the first jet. +- `kt2_2`: The squared momentum metric value of the second jet. +- `R2`: The jet radius parameter squared. + +# Returns +The distance between the two jets. + +If `kt2_2` is equal to 0.0, then the first jet doesn't actually have a valid +neighbour, so it's treated as a single jet adjecent to the beam. """ get_dij_dist(nn_dist, kt2_1, kt2_2, R2) = begin if kt2_2 == 0.0 @@ -156,8 +232,22 @@ get_dij_dist(nn_dist, kt2_1, kt2_2, R2) = begin return nn_dist * min(kt2_1, kt2_2) end + """ -Return the tile coordinates of an (eta, phi) pair + get_tile(tiling_setup::TilingDef, eta::AbstractFloat, phi::AbstractFloat) + +Given a `tiling_setup` object, `eta` and `phi` values, this function calculates +the tile indices for the given `eta` and `phi` values. + +# Arguments +- `tiling_setup`: A `TilingDef` object that contains the tiling setup + parameters. +- `eta`: The eta value for which to calculate the tile index. +- `phi`: The phi value for which to calculate the tile index. + +# Returns +- `ieta`: The tile index along the eta direction. +- `iphi`: The tile index along the phi direction. """ get_tile(tiling_setup::TilingDef, eta::AbstractFloat, phi::AbstractFloat) = begin # The eta clamp is necessary as the extreme bins catch overflows for high abs(eta) @@ -167,17 +257,22 @@ get_tile(tiling_setup::TilingDef, eta::AbstractFloat, phi::AbstractFloat) = begi ieta, iphi end -""" -Map an (η, ϕ) pair into a linear index, which is much faster "by hand" than using -the LinearIndices construct (like x100, which is bonkers, but there you go...) -""" -get_tile_linear_index(tiling_setup::TilingDef, i_η::Int, i_ϕ::Int) = begin - return tiling_setup._n_tiles_eta * (i_ϕ-1) + i_η -end """ -Map a linear index to a tuple of (η, ϕ) - again this is very much faster than -using CartesianIndices + get_tile_linear_index(tiling_setup::TilingDef, i_η::Int, i_ϕ::Int) + +Compute the linear index of a tile in a tiled setup. This is much faster in this +function than using the LinearIndices construct (like x100, which is bonkers, +but there you go...) + +# Arguments +- `tiling_setup::TilingDef`: The tiling setup defining the number of tiles in + each dimension. +- `i_η::Int`: The index of the tile in the η dimension. +- `i_ϕ::Int`: The index of the tile in the ϕ dimension. + +# Returns +- The linear index of the tile. """ get_tile_cartesian_indices(tiling_setup::TilingDef, index::Int) = begin return (rem(index-1, tiling_setup._n_tiles_eta)+1, div(index-1, tiling_setup._n_tiles_eta)+1) @@ -189,16 +284,43 @@ Iterator for the indexes of rightmost tiles for a given Cartesian tile index XXX O.X OOO - - η coordinate must be in range, ϕ coordinate wraps + - rapidity coordinate must be in range, ϕ coordinate wraps + +""" +""" + struct rightmost_tiles + +A struct for iterating over rightmost tiles for a given Cartesian tile index. +These are the tiles above and to the right of the given tile (X=included, O=not +included): + + XXX + O.X + OOO + +Note, rapidity coordinate must be in range, ϕ coordinate wraps + +# Fields +- `n_η::Int`: Number of η tiles +- `n_ϕ::Int`: Number of ϕ tiles +- `start_η::Int`: Centre η tile coordinate +- `start_ϕ::Int`: Centre ϕ tile coordinate """ struct rightmost_tiles - n_η::Int # Number of η tiles - n_ϕ::Int # Number of ϕ tiles - start_η::Int # Centre η tile coordinate - start_ϕ::Int # Centre ϕ tile coordinate + n_η::Int # Number of η tiles + n_ϕ::Int # Number of ϕ tiles + start_η::Int # Centre η tile coordinate + start_ϕ::Int # Centre ϕ tile coordinate end + +""" + Base.iterate(t::rightmost_tiles, state=1) + +Iterate over the `rightmost_tiles` object, returning all the rightmost tiles for +a given Cartesian tile index. +""" function Base.iterate(t::rightmost_tiles, state=1) mapping = ((-1,-1), (-1,0), (-1,1), (0,1)) if t.start_η == 1 && state == 1 @@ -217,13 +339,27 @@ function Base.iterate(t::rightmost_tiles, state=1) return nothing end + """ -Iterator for the indexes of neighbouring tiles for a given Cartesian tile index - XXX - X.X - XXX - - η coordinate must be in range, ϕ coordinate wraps + struct neighbour_tiles +A struct representing the neighbouring tiles. + +A struct for iterating over all neighbour tiles for a given Cartesian tile index. +These are the tiles above and to the right of the given tile (X=included, O=not +included): + + XXX + X.X + XXX + +Note, rapidity coordinate must be in range, ϕ coordinate wraps + +# Fields +- `n_η::Int`: Number of η tiles +- `n_ϕ::Int`: Number of ϕ tiles +- `start_η::Int`: Centre η tile coordinate +- `start_ϕ::Int`: Centre ϕ tile coordinate """ struct neighbour_tiles n_η::Int # Number of η tiles @@ -232,6 +368,12 @@ struct neighbour_tiles start_ϕ::Int # Centre ϕ tile coordinate end +""" + Base.iterate(t::neighbour_tiles, state=1) + +Iterate over the `neighbour_tiles` object, returning all the neighbour tiles for +a given Cartesian tile index. +""" function Base.iterate(t::neighbour_tiles, state=1) mapping = ((-1,-1), (-1,0), (-1,1), (0,-1), (0,1), (1,-1), (1,0), (1,1)) # Skip for top row in η diff --git a/src/Utils.jl b/src/Utils.jl index 62c71e4..1cc23d4 100644 --- a/src/Utils.jl +++ b/src/Utils.jl @@ -1,6 +1,23 @@ # Utility functions, which can be used by different top level scripts -"""Read HepMC3 event and keep final state particles (return PseudoJets)""" + +""" + read_final_state_particles(fname; maxevents = -1, skipevents = 0) + +Reads final state particles from a file and returns them as a vector of +PseudoJet objects. + +# Arguments +- `fname`: The name of the file to read particles from. +- `maxevents=-1`: The maximum number of events to read. -1 means all events will + be read. +- `skipevents`: The number of events to skip before an event is included. + Default is 0. + +# Returns +A vector of vectors of PseudoJet objects, where each inner vector represents all +the particles of a particular event. +""" function read_final_state_particles(fname; maxevents = -1, skipevents = 0) f = open(fname) @@ -27,7 +44,23 @@ function read_final_state_particles(fname; maxevents = -1, skipevents = 0) events end -"""Read HepMC3 event and keep final state particles (return LorentzVectors)""" +""" + read_final_state_particles_lv(fname; maxevents = -1, skipevents = 0) + +Reads final state particles from a file and returns them as a vector of +LorentzVector objects. + +# Arguments +- `fname`: The name of the file to read particles from. +- `maxevents=-1`: The maximum number of events to read. -1 means all events will + be read. +- `skipevents`: The number of events to skip before an event is included. + Default is 0. + +# Returns +A vector of vectors of LorentzVector objects, where each inner vector represents +all the particles of a particular event. +""" function read_final_state_particles_lv(fname; maxevents = -1, skipevents = 0) f = open(fname) @@ -54,23 +87,7 @@ function read_final_state_particles_lv(fname; maxevents = -1, skipevents = 0) events end -function pseudojets2vectors(events::Vector{Vector{PseudoJet}}) - """Convert events in PseudoJets into deep vectors""" - event_vector = Vector{Vector{Vector{Float64}}}(undef, size(events)[1]) - for (ievent, event) in enumerate(events) - jet_struct = Vector{Vector{Float64}}(undef, size(event)[1]) - for (ipart, initial_particle) in enumerate(event) - jet_struct[ipart] = [ - initial_particle.px, - initial_particle.py, - initial_particle.pz, - initial_particle.E, - ] - end - event_vector[ievent] = jet_struct - end - event_vector -end + """ Return the list of jets passing a pt cut @@ -78,21 +95,38 @@ Return the list of jets passing a pt cut The ptmin cut in these functions is slightly legacy as often the input jets were already filtered on pt """ -function final_jets(jets::Vector{Vector{Float64}}, ptmin::AbstractFloat=0.0) - count = 0 - final_jets = Vector{FinalJet}() - sizehint!(final_jets, 6) - for jet in jets - dcut = ptmin^2 - p = PseudoJet(jet[1], jet[2], jet[3], jet[4]) - if p._pt2 > dcut - count += 1 - push!(final_jets, FinalJet(rap(p), phi(p), sqrt(pt2(p)))) - end - end - final_jets -end +# function final_jets(jets::Vector{Vector{Float64}}, ptmin::AbstractFloat=0.0) +# count = 0 +# final_jets = Vector{FinalJet}() +# sizehint!(final_jets, 6) +# for jet in jets +# dcut = ptmin^2 +# p = PseudoJet(jet[1], jet[2], jet[3], jet[4]) +# if p._pt2 > dcut +# count += 1 +# push!(final_jets, FinalJet(rap(p), phi(p), sqrt(pt2(p)))) +# end +# end +# final_jets +# end + +""" + final_jets(jets::Vector{PseudoJet}, ptmin::AbstractFloat=0.0) + +This function takes a vector of `PseudoJet` objects and a minimum transverse +momentum `ptmin` as input. It returns a vector of `FinalJet` objects that +satisfy the transverse momentum condition. + +# Arguments +- `jets::Vector{PseudoJet}`: A vector of `PseudoJet` objects representing the + input jets. +- `ptmin::AbstractFloat=0.0`: The minimum transverse momentum required for a jet + to be included in the final jets vector. + +# Returns +A vector of `FinalJet` objects that satisfy the transverse momentum condition. +""" function final_jets(jets::Vector{PseudoJet}, ptmin::AbstractFloat=0.0) count = 0 final_jets = Vector{FinalJet}() @@ -107,6 +141,7 @@ function final_jets(jets::Vector{PseudoJet}, ptmin::AbstractFloat=0.0) final_jets end +"""Specialisation for final jets from LorentzVectors (TODO: merge into more general function)""" function final_jets(jets::Vector{LorentzVector}, ptmin::AbstractFloat=0.0) count = 0 final_jets = Vector{FinalJet}() @@ -121,6 +156,7 @@ function final_jets(jets::Vector{LorentzVector}, ptmin::AbstractFloat=0.0) final_jets end +"""Specialisation for final jets from LorentzVectorCyl (TODO: merge into more general function)""" function final_jets(jets::Vector{LorentzVectorCyl}, ptmin::AbstractFloat=0.0) count = 0 final_jets = Vector{FinalJet}() @@ -135,16 +171,29 @@ function final_jets(jets::Vector{LorentzVectorCyl}, ptmin::AbstractFloat=0.0) final_jets end -## Faster findmin function -"""Find the lowest value in the array, returning the value and the index""" + +""" + fast_findmin(dij, n) + +Find the minimum value and its index in the first `n` elements of the `dij` +array. The use of `@turbo` macro gives a significiant performance boost. + +# Arguments +- `dij`: An array of values. +- `n`: The number of elements to consider in the `dij` array. + +# Returns +- `dij_min`: The minimum value in the first `n` elements of the `dij` array. +- `best`: The index of the minimum value in the `dij` array. +""" fast_findmin(dij, n) = begin - # findmin(@inbounds @view dij[1:n]) - best = 1 - @inbounds dij_min = dij[1] - @turbo for here in 2:n - newmin = dij[here] < dij_min - best = newmin ? here : best - dij_min = newmin ? dij[here] : dij_min - end - dij_min, best + # findmin(@inbounds @view dij[1:n]) + best = 1 + @inbounds dij_min = dij[1] + @turbo for here in 2:n + newmin = dij[here] < dij_min + best = newmin ? here : best + dij_min = newmin ? dij[here] : dij_min + end + dij_min, best end \ No newline at end of file