pfnano_validation.py

import awkward as ak
from coffea.nanoevents import NanoEventsFactory

fn = "root://grid-cms-xrootd.physik.rwth-aachen.de:1094//store/user/anovak/PFNano/106X_v2_17/QCD_HT1000to1500_TuneCP5_13TeV-madgraph-pythia8/RunIIFall17PFNanoAODv2-PU2017_12Apr2018_94X_mc2017_realistic_v14-v1PFNanoV2/210101_173148/0000/nano_mc2017_10.root"
events = NanoEventsFactory.from_root(fn).events()

# the Jet-PF association has valid indexes to jet collection
assert ak.all(events.JetPFCands.jetIdx < ak.num(events.Jet))
assert ak.all(events.FatJetPFCands.jetIdx < ak.num(events.FatJet))
assert ak.all(events.GenJetCands.jetIdx < ak.num(events.GenJet))
# assert ak.all(events.GenFatJetCands.jetIdx < ak.num(events.GenJetAK8))

# the Jet-PF association has valid indexes to pf collection
assert ak.all(events.JetPFCands.pFCandsIdx < ak.num(events.PFCands))
assert ak.all(events.FatJetPFCands.pFCandsIdx < ak.num(events.PFCands))
assert ak.all(events.GenJetCands.pFCandsIdx < ak.num(events.GenCands))
assert ak.all(events.GenFatJetCands.pFCandsIdx < ak.num(events.GenCands))

# the Jet-PF association is the right size
assert ak.all(ak.num(events.JetPFCands) == ak.sum(events.Jet.nConstituents, axis=1))
# assert ak.all(ak.num(events.FatJetPFCands) == ak.sum(events.FatJet.nConstituents, axis=1))
# assert ak.all(ak.num(events.GenJetCands) == ak.sum(events.GenJet.nConstituents, axis=1))
# assert ak.all(ak.num(events.GenFatJetCands) == ak.sum(events.GenJetAK8.nConstituents, axis=1))

# the Jet-PF association is filled in the same order as the jets
def diff_axis1(array):
    return array[:, 1:] - array[:, :-1]

assert ak.all(diff_axis1(events.JetPFCands.jetIdx) >= 0)
assert ak.all(diff_axis1(events.FatJetPFCands.jetIdx) >= 0)
assert ak.all(diff_axis1(events.GenJetCands.jetIdx) >= 0)
assert ak.all(diff_axis1(events.GenFatJetCands.jetIdx) >= 0)

# same checks for secondary vertices

assert ak.all(events.JetSVs.jetIdx < ak.num(events.Jet))
assert ak.all(events.FatJetSVs.jetIdx < ak.num(events.FatJet))
assert ak.all(events.GenJetSVs.jetIdx < ak.num(events.GenJet))
# assert ak.all(events.GenFatJetSVs.jetIdx < ak.num(events.GenJetAK8))

assert ak.all(events.JetSVs.sVIdx < ak.num(events.SV))
assert ak.all(events.FatJetSVs.sVIdx < ak.num(events.SV))
assert ak.all(events.GenJetSVs.sVIdx < ak.num(events.SV))
assert ak.all(events.GenFatJetSVs.sVIdx < ak.num(events.SV))

# assert ak.all(ak.num(events.JetSVs) == ak.sum(events.Jet.nSVs, axis=1))
# assert ak.all(ak.num(events.FatJetSVs) == ak.sum(events.FatJet.nSVs, axis=1))
# assert ak.all(ak.num(events.GenJetSVs) == ak.sum(events.GenJet.nSVs, axis=1))
# assert ak.all(ak.num(events.GenFatJetSVs) == ak.sum(events.GenJetAK8.nSVs, axis=1))

assert ak.all(diff_axis1(events.JetSVs.jetIdx) >= 0)
assert ak.all(diff_axis1(events.FatJetSVs.jetIdx) >= 0)
assert ak.all(diff_axis1(events.GenJetSVs.jetIdx) >= 0)
assert ak.all(diff_axis1(events.GenFatJetSVs.jetIdx) >= 0)