|
| 1 | +using JetReconstruction |
| 2 | +using CairoMakie |
| 3 | +using StatsBase |
| 4 | +using LaTeXStrings |
| 5 | + |
| 6 | +""" |
| 7 | +generate_average_lund_image(njets::Int, delta_array::Vector{Vector{Real}}, |
| 8 | + kt_array::Vector{Vector{Real}}; xrange::Tuple{Real}, |
| 9 | + yrange::Tuple{Real}, bins::Int) -> (xgrid, ygrid, avg_image) |
| 10 | +
|
| 11 | +Computes an average Lund image from a set of jets by binning the (log(1/ΔR), log(kt)) |
| 12 | +coordinates into a fixed-size 2D histogram. Each jet's histogram is normalized and |
| 13 | +then averaged across all jets. |
| 14 | +
|
| 15 | +Arguments: |
| 16 | +- `njets`: Number of jets |
| 17 | +- `delta_array`: Vector of vectors, where each inner vector contains ΔR values for a jet |
| 18 | +- `kt_array`: Vector of vectors, where each inner vector contains kt values for a jet |
| 19 | +- `xrange`: Tuple defining the x-axis (log(1/ΔR)) range |
| 20 | +- `yrange`: Tuple defining the y-axis (log(kt)) range |
| 21 | +- `bins`: Number of bins along each axis |
| 22 | +
|
| 23 | +Returns: |
| 24 | +- A tuple `(xgrid, ygrid, avg_image)` where `xgrid` and `ygrid` are coordinate axes labels and |
| 25 | + `avg_image` is the averaged 2D histogram as a matrix |
| 26 | +""" |
| 27 | +function generate_average_lund_image(njets::Int, delta_array::Vector{Vector{Real}}, |
| 28 | + kt_array::Vector{Vector{Real}}; xrange = (0.0, 4.0), |
| 29 | + yrange = (-5.0, 7.0), bins = 25) |
| 30 | + xmin, xmax = xrange |
| 31 | + ymin, ymax = yrange |
| 32 | + |
| 33 | + x_width = (xmax - xmin) / bins |
| 34 | + y_width = (ymax - ymin) / bins |
| 35 | + |
| 36 | + total_res = [] |
| 37 | + |
| 38 | + for i in 1:njets |
| 39 | + Xind = ceil.((delta_array[i] .- xmin) ./ x_width) |
| 40 | + Yind = ceil.((kt_array[i] .- ymin) ./ y_width) |
| 41 | + |
| 42 | + res = zeros(Float32, bins, bins) |
| 43 | + L1norm = 0.0 |
| 44 | + |
| 45 | + for i in 1:length(Xind) |
| 46 | + x = Int(Xind[i]) |
| 47 | + y = Int(Yind[i]) |
| 48 | + if (maximum([x, y]) < bins && minimum([x, y]) >= 1) |
| 49 | + res[x, y] += 1 |
| 50 | + L1norm += 1 |
| 51 | + end |
| 52 | + end |
| 53 | + |
| 54 | + if L1norm > 0 |
| 55 | + res[1, :] .= res[1, :] ./ L1norm |
| 56 | + push!(total_res, res) |
| 57 | + end |
| 58 | + end |
| 59 | + |
| 60 | + avg_res = mean(total_res, dims = 1)[1] |
| 61 | + x = range(xmin, xmax; length = bins) |
| 62 | + y = range(ymin, ymax; length = bins) |
| 63 | + |
| 64 | + return (x, y, avg_res) |
| 65 | +end |
| 66 | + |
| 67 | +input_file = joinpath(dirname(pathof(JetReconstruction)), |
| 68 | + "..", "test", "data", "events.pp13TeV.hepmc3.zst") |
| 69 | +events = read_final_state_particles(input_file) |
| 70 | + |
| 71 | +N = length(events) |
| 72 | +lundX = Vector{Vector{Real}}() |
| 73 | +lundY = Vector{Vector{Real}}() |
| 74 | + |
| 75 | +lund_kt(p) = p.kt |
| 76 | +lund_delta(p) = p.delta |
| 77 | + |
| 78 | +# Event to pick |
| 79 | +for event_no in 1:N |
| 80 | + cluster_seq = jet_reconstruct(events[event_no]; algorithm = JetAlgorithm.AntiKt, |
| 81 | + R = 1.0) |
| 82 | + jets = sort!(inclusive_jets(cluster_seq, PseudoJet; ptmin = 10.0), |
| 83 | + by = JetReconstruction.pt2, rev = true) |
| 84 | + |
| 85 | + @info "Generating Primary Lund Emissions for $(length(jets)) jets for Event $(event_no):" |
| 86 | + for (ijet, jet) in enumerate(jets) |
| 87 | + lundvars = generate_lund_emissions(jet, cluster_seq) |
| 88 | + println("- Jet $(ijet) has $(length(lundvars)) emissions in lund plane") |
| 89 | + kt = lund_kt.(lundvars) |
| 90 | + Δ = lund_delta.(lundvars) |
| 91 | + |
| 92 | + push!(lundX, -log.(Δ)) |
| 93 | + push!(lundY, log.(kt)) |
| 94 | + end |
| 95 | +end |
| 96 | + |
| 97 | +njets = length(lundX) |
| 98 | +x, y, avg_res = generate_average_lund_image(njets, lundX, lundY) |
| 99 | + |
| 100 | +fig = Figure() |
| 101 | +ax = Axis(fig[1, 1], xlabel = L"\ln(R/\Delta)", ylabel = L"\ln(k_T/\mathrm{GeV})", |
| 102 | + title = "Average Lund Image") |
| 103 | +hm = heatmap!(ax, x, y, avg_res; colormap = :viridis, colorrange = extrema(avg_res)) |
| 104 | +Colorbar(fig[1, 2], hm; label = "") |
| 105 | + |
| 106 | +display(fig) |
| 107 | + |
| 108 | +save("lund-gen.png", fig) |
0 commit comments