jpata / hepaccelerate-cms Goto Github PK

import uproot

tt = uproot.open("data/myNanoProdMc2016_NANO.root").get("Events")
mu_pt = tt.array("Muon_pt")
print(mu_pt.content.shape)
print(mu_pt.offsets)

for iev in range(len(mu_pt.offsets) - 1):
    imu1 = mu_pt.offsets[iev]
    imu2 = mu_pt.offsets[iev+1]
    print(f"Event {iev} has {imu2-imu1} muons (offsets {imu1}:{imu2})")
    for imu in range(imu1, imu2):
        print(imu, mu_pt.content[imu])

mu_pt.content[mu_pt.offsets[0:-1]]

#Corrects the muon momentum and isolation, if a matched FSR photon with dR<0.4 is found
@numba.njit
def correct_muon_with_fsr(
        muons_offsets, fsr_offsets,
        muons_pt, muons_eta, muons_phi, muons_mass, muons_iso, muons_fsrIndex
        fsr_pt, fsr_eta, fsr_phi, fsr_mass,
        out_muons_pt, out_muons_eta, out_muons_phi, out_muons_mass, out_muons_iso
    ):
    for iev in numba.prange(len(muons_offsets) - 1):

        #loop over muons in event
        mu1 = muons_offsets[iev]
        mu2 = muons_offsets[iev + 1]

        for imu in range(mu1, mu2):

            #relative FSR index in the event 
            fsr_idx_relative = muons_fsrIndex[imu]
            if fsr_idx_relative >= 0:
                #absolute index in the full FSR vector for all events
                ifsr = fsr_offsets[iev] + fsr_idx_relative

                this_mu_pt = muons_pt[imu]
                ... (and others)

                this_fsr_pt = fsr_pt[ifsr]
                ... (and others)

                #compute dR between muon and FSR, check <0.4
                #compute and set corrected isolation
                #compute and set corrected momentum