In [1]:

using Pkg
Pkg.status()

      Status `/home/jiling/UnROOT_RDataFrame_MiniBenchmark/Project.toml`
  [68837c9b] FHist v0.6.1
  [7073ff75] IJulia v1.23.2
  [3a55db76] LVCyl v0.1.0 `https://github.com/JuliaHEP/LVCyl.jl#master`
  [f517fe37] Polyester v0.4.2
  [3cd96dde] UnROOT v0.5.1

In [2]:

using UnROOT, FHist, Polyester, LVCyl

┌ Info: Precompiling UnROOT [3cd96dde-e98d-4713-81e9-a4a1b0235ce9]
└ @ Base loading.jl:1423
┌ Info: Precompiling FHist [68837c9b-b678-4cd5-9925-8a54edc8f695]
└ @ Base loading.jl:1423
┌ Info: Precompiling Polyester [f517fe37-dbe3-4b94-8317-1923a5111588]
└ @ Base loading.jl:1423
┌ Info: Precompiling LVCyl [3a55db76-103f-4f27-8e51-63f2a02e4e27]
└ @ Base loading.jl:1423

In [3]:

# http://opendata.web.cern.ch/record/12341
rootfname = "./Run2012BC_DoubleMuParked_Muons.root"

Out[3]:

"./Run2012BC_DoubleMuParked_Muons.root"

In [4]:

filesize(rootfname) / 1024^3 # in GB

Out[4]:

2.0903061451390386

In [5]:

const r = ROOTFile(rootfname);

In [6]:

const mytree = LazyTree(r, "Events");

In [7]:

mytree[begin:3]

Out[7]:

3 rows × 6 columns (omitted printing of 1 columns)

	nMuon	Muon_pt	Muon_eta	Muon_phi	Muon_mass
	UInt32	SubArra…	SubArra…	SubArra…	SubArra…
1	2	[10.7637, 15.7365]	[1.06683, -0.563787]	[-0.0342727, 2.54262]	[0.105658, 0.105658]
2	2	[10.5385, 16.3271]	[-0.42778, 0.349225]	[-0.274792, 2.53978]	[0.105658, 0.105658]
3	1	[3.27533]	[2.21086]	[-1.22341]	[0.105658]

In [8]:

const z_mass = 91.2
const LV32 = LorentzVectorCyl{Float32}
function reco_zz_to_4l(pts, etas, phis, masses, charges)
    idx = [Int[], Int[]]

    # Find first lepton pair with invariant mass closest to Z mass
    best_mass = -Inf
    best_i1 = best_i2 = 1
    for i1 in eachindex(pts), i2 in (i1 + 1):lastindex(pts)
        charges[i1] == charges[i2] && continue
        p1 = LV32(pts[i1], etas[i1], phis[i1], masses[i1])
        p2 = LV32(pts[i2], etas[i2], phis[i2], masses[i2])
        this_mass = (p1+p2).mass
        if (abs(z_mass - this_mass) < abs(z_mass - best_mass))
            best_mass = this_mass
            best_i1 = i1
            best_i2 = i2
        end
    end
    push!(idx[1], best_i1)
    push!(idx[1], best_i2)

    #Reconstruct second Z from remaining lepton pair
    for i in 1:4
        if (i != best_i1 && i != best_i2)
            push!(idx[2], i)
        end
    end

    # Return indices of the pairs building two Z bosons
    return idx
end

Out[8]:

reco_zz_to_4l (generic function with 1 method)

In [9]:

function filter_z_dr(idx, etas, phis)
    for pair in idx
        i1, i2 = pair
        dr = sqrt((etas[i1]-etas[i2])^2 + (phis[i1]-phis[i2])^2)
        dr < 0.02 && return false;
    end
    return true
end

Out[9]:

filter_z_dr (generic function with 1 method)

In [10]:

function filter_z_candidates(Z_mass)
    40 < Z_mass[1] < 120 || return false
    12 < Z_mass[2] < 120 || return false
    return true
end

Out[10]:

filter_z_candidates (generic function with 1 method)

In [11]:

function compute_z_masses_4l(idx, pts, etas, phis, masses)
    z_masses = zeros(2)
    for (i, pair) in enumerate(idx)
        i1, i2 = pair
        p1 = LV32(pts[i1], etas[i1], phis[i1], masses[i1])
        p2 = LV32(pts[i2], etas[i2], phis[i2], masses[i2])
        z_masses[i] = (p1+p2).mass
    end
   return abs(z_masses[1] - z_mass) < abs(z_masses[2] - z_mass) ?
      z_masses : reverse!(z_masses);
end

Out[11]:

compute_z_masses_4l (generic function with 1 method)

In [12]:

function compute_higgs_mass_4l(idx, pts, etas, phis, masses)
    ((i1, i2), (i3, i4)) = idx
    p1 = LV32(pts[i1], etas[i1], phis[i1], masses[i1]);
    p2 = LV32(pts[i2], etas[i2], phis[i2], masses[i2]);
    p3 = LV32(pts[i3], etas[i3], phis[i3], masses[i3]);
    p4 = LV32(pts[i4], etas[i4], phis[i4], masses[i4]);
    return (p1 + p2 + p3 + p4).mass
end

Out[12]:

compute_higgs_mass_4l (generic function with 1 method)

In [13]:

let H = Hist1D(Float64; bins = range(70, 180; length=36), overflow=true)
    #@time @batch for evt in mytree
    @time for evt in mytree
        evt.nMuon != 4 && continue
        
        pts, etas  = evt.Muon_pt, evt.Muon_eta
        
        (all(pts .> 5) && all(abs.(etas) .< 2.4)) || continue
        
        phis, masses, charges = evt.Muon_phi, evt.Muon_mass, evt.Muon_charge
        
        sum(charges) != 0 && continue
        
        Z_idx = reco_zz_to_4l(pts, etas, phis, masses, charges)
        
        filter_z_dr(Z_idx, etas, phis) || continue

        Z_mass = compute_z_masses_4l(Z_idx, pts, etas, phis, masses)

        filter_z_candidates(Z_mass) || continue
        h_mass = compute_higgs_mass_4l(Z_idx, pts, etas, phis, masses)
        push!(H, h_mass)
    end
    H
end

 23.127858 seconds (18.04 M allocations: 24.054 GiB, 19.15% gc time)

Out[13]:

edges: 70.0:3.142857142857143:180.0
bin counts: [13923.0, 4302.0, 4300.0, 4252.0, 3949.0, 3550.0, 3362.0, 3125.0, 2734.0, 2657.0 … 852.0, 814.0, 745.0, 651.0, 592.0, 537.0, 542.0, 501.0, 445.0, 5319.0]
maximum count: 13923.0
total count: 80487.0

In [14]:

Threads.nthreads()

Out[14]:

In [15]:

let H = Hist1D(Float64; bins = range(70, 180; length=36), overflow=true)
    @time @batch for evt in mytree
        evt.nMuon != 4 && continue
        
        pts, etas  = evt.Muon_pt, evt.Muon_eta
        
        (all(pts .> 5) && all(abs.(etas) .< 2.4)) || continue
        
        phis, masses, charges = evt.Muon_phi, evt.Muon_mass, evt.Muon_charge
        
        sum(charges) != 0 && continue
        
        Z_idx = reco_zz_to_4l(pts, etas, phis, masses, charges)
        
        filter_z_dr(Z_idx, etas, phis) || continue

        Z_mass = compute_z_masses_4l(Z_idx, pts, etas, phis, masses)

        filter_z_candidates(Z_mass) || continue

        h_mass = compute_higgs_mass_4l(Z_idx, pts, etas, phis, masses)
        push!(H, h_mass)
    end
    H
end

  6.702241 seconds (21.68 M allocations: 25.134 GiB, 5.03% gc time, 0.35% compilation time)

Out[15]:

edges: 70.0:3.142857142857143:180.0
bin counts: [13923.0, 4302.0, 4300.0, 4252.0, 3949.0, 3550.0, 3362.0, 3125.0, 2734.0, 2657.0 … 852.0, 814.0, 745.0, 651.0, 592.0, 537.0, 542.0, 501.0, 445.0, 5319.0]
maximum count: 13923.0
total count: 80487.0