Skip to content

Commit

Permalink
Generalise to return type T
Browse files Browse the repository at this point in the history
Make the return type for read_final_state_particles a type, allowing
one function to serve both for PseudoJet and LorentzVector
  • Loading branch information
graeme-a-stewart committed Jul 12, 2024
1 parent fc239ed commit 8de4986
Showing 1 changed file with 9 additions and 31 deletions.
40 changes: 9 additions & 31 deletions src/Utils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,39 +3,38 @@
using CodecZlib

"""
read_final_state_particles(fname; maxevents = -1, skipevents = 0)
read_final_state_particles(fname; maxevents = -1, skipevents = 0, T=PseudoJet)
Reads final state particles from a file and returns them as a vector of
PseudoJet objects.
Reads final state particles from a file and returns them as a vector of type T.
# Arguments
- `fname`: The name of the HepMC3 ASCII file to read particles from. If the file
is gzipped, the function will automatically decompress it.
- `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.
- `skipevents=0`: The number of events to skip before an event is included.
- `T=PseudoJet`: The type of object to contruct and return.
# Returns
A vector of vectors of PseudoJet objects, where each inner vector represents all
A vector of vectors of T objects, where each inner vector represents all
the particles of a particular event.
"""
function read_final_state_particles(fname; maxevents = -1, skipevents = 0)
function read_final_state_particles(fname; maxevents = -1, skipevents = 0, T=PseudoJet)
f = open(fname)
if endswith(fname, ".gz")
@debug "Reading gzipped file $fname"
f = GzipDecompressorStream(f)
end

events = Vector{PseudoJet}[]
events = Vector{T}[]

ipart = 1
HepMC3.read_events(f, maxevents = maxevents, skipevents = skipevents) do parts
input_particles = PseudoJet[]
input_particles = T[]
for p in parts
if p.status == 1
push!(input_particles,
PseudoJet(p.momentum.x, p.momentum.y, p.momentum.z, p.momentum.t))
T(p.momentum.x, p.momentum.y, p.momentum.z, p.momentum.t))
end
end
push!(events, input_particles)
Expand Down Expand Up @@ -94,27 +93,6 @@ function read_final_state_particles_lv(fname; maxevents = -1, skipevents = 0)
events
end

"""
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

"""
final_jets(jets::Vector{PseudoJet}, ptmin::AbstractFloat=0.0)
Expand Down

0 comments on commit 8de4986

Please sign in to comment.