def test_regular():
regular = ak.Array(np.array([[i, i] for i in range(10)], dtype=np.int64))
assert str(regular.type) == "10 * 2 * int64"
assert ak.to_list(ak.to_regular(regular)) == ak.to_list(regular)
assert ak.type(ak.to_regular(regular)) == ak.type(regular)
irregular = ak.from_regular(regular)
assert ak.to_list(irregular) == ak.to_list(regular)
assert str(irregular.type) == "10 * var * int64"
assert ak.to_list(ak.from_regular(irregular)) == ak.to_list(irregular)
assert ak.type(ak.from_regular(irregular)) == ak.type(irregular)
reregular = ak.to_regular(irregular)
assert ak.to_list(reregular) == ak.to_list(regular)
assert str(reregular.type) == "10 * 2 * int64"
def test_regular_deep():
regular = ak.Array(
np.array([[[i, i, i], [i, i, i]] for i in range(10)], dtype=np.int64)
)
assert str(regular.type) == "10 * 2 * 3 * int64"
irregular = ak.from_regular(regular, axis=1)
assert ak.to_list(irregular) == ak.to_list(regular)
assert str(irregular.type) == "10 * var * 3 * int64"
reregular = ak.to_regular(irregular, axis=1)
assert ak.to_list(reregular) == ak.to_list(regular)
assert str(reregular.type) == "10 * 2 * 3 * int64"
irregular = ak.from_regular(regular, axis=2)
assert ak.to_list(irregular) == ak.to_list(regular)
assert str(irregular.type) == "10 * 2 * var * int64"
reregular = ak.to_regular(irregular, axis=2)
assert ak.to_list(reregular) == ak.to_list(regular)
assert str(reregular.type) == "10 * 2 * 3 * int64"
irregular = ak.from_regular(regular, axis=-1)
assert ak.to_list(irregular) == ak.to_list(regular)
assert str(irregular.type) == "10 * 2 * var * int64"
reregular = ak.to_regular(irregular, axis=-1)
assert ak.to_list(reregular) == ak.to_list(regular)
assert str(reregular.type) == "10 * 2 * 3 * int64"
irregular = ak.from_regular(regular, axis=-2)
assert ak.to_list(irregular) == ak.to_list(regular)
assert str(irregular.type) == "10 * var * 3 * int64"
reregular = ak.to_regular(irregular, axis=-2)
assert ak.to_list(reregular) == ak.to_list(regular)
assert str(reregular.type) == "10 * 2 * 3 * int64"
with pytest.raises(ValueError):
ak.from_regular(regular, axis=-3)
assert ak.to_list(ak.from_regular(regular, axis=0)) == ak.to_list(regular)
assert ak.type(ak.from_regular(regular, axis=0)) == ak.type(regular)
def __init__(self, ev, obj, wp, year=2018, verbose=0):
self.obj = obj
self.wp = wp
if self.wp == None:
self.selection_dict = {}
else:
self.selection_dict = obj_def[self.obj][self.wp]
self.v = verbose
self.year = year
id_level = None
if wp.lower().count('veto') or wp.lower().count('loose'):
id_level = 0
elif wp.lower().count('fake'):
id_level = 1
elif wp.lower().count('tight'):
id_level = 2
if self.obj == "Muon":
# collections are already there, so we just need to calculate missing ones
ev['Muon', 'absMiniIso'] = ev.Muon.miniPFRelIso_all * ev.Muon.pt
ev['Muon', 'ptErrRel'] = ev.Muon.ptErr / ev.Muon.pt
# this is what we are using:
# - jetRelIso if the matched jet is within deltaR<0.4, pfRelIso03_all otherwise
# - btagDeepFlavB discriminator of the matched jet if jet is within deltaR<0.4, 0 otherwise
# - pt_cone = 0.9*pt of matched jet if jet is within deltaR<0.4, pt/(pt+iso) otherwise
mask_close = (ak.fill_none(ev.Muon.delta_r(ev.Muon.matched_jet),
99) < 0.4) * 1
mask_far = ~(ak.fill_none(ev.Muon.delta_r(ev.Muon.matched_jet), 99)
< 0.4) * 1
deepJet = ak.fill_none(ev.Muon.matched_jet.btagDeepFlavB,
0) * mask_close + 0 * mask_far
jetRelIsoV2 = ev.Muon.jetRelIso * mask_close + ev.Muon.pfRelIso03_all * mask_far # default to 0 if no match
#conePt = 0.9 * ak.fill_none(ev.Muon.matched_jet.pt,0) * mask_close + ev.Muon.pt*(1 + ev.Muon.miniPFRelIso_all)*mask_far
if self.year == 2017 or self.year == 2018:
I_1 = 0.11
I_2 = 0.74
I_3 = 6.8
elif self.year == 2016:
I_1 = 0.16
I_2 = 0.76
I_3 = 7.2
PF_unflatten = ak.from_regular(
ev.Muon.miniPFRelIso_all[:, :, np.newaxis])
max_miniIso = ak.max(
ak.concatenate(
[PF_unflatten - I_1,
ak.zeros_like(PF_unflatten)], axis=2),
axis=2) #equivalent to max(0, ev.Muon.miniPFRelIso_all - I_1)
muon_pt_unflatten = ak.from_regular(ev.Muon.pt[:, :, np.newaxis])
jet_pt_unflatten = ak.from_regular(
ev.Muon.matched_jet.pt[:, :, np.newaxis])
max_pt = ak.max(
ak.concatenate([muon_pt_unflatten, jet_pt_unflatten * I_2],
axis=2),
axis=2) #max(ev.Muon.pt, ev.Muon.matched_jet.pt * I_2)
conePt = (ev.Muon.pt *
(1 + max_miniIso)) * (ev.Muon.jetPtRelv2 > I_3) + (
max_pt * ~(ev.Muon.jetPtRelv2 > I_3))
ev['Muon', 'deepJet'] = ak.copy(deepJet)
ev['Muon', 'jetRelIsoV2'] = jetRelIsoV2
ev['Muon', 'conePt'] = conePt
ev['Muon', 'id'] = ak.ones_like(conePt) * id_level
self.cand = ev.Muon
elif self.obj == "Electron":
# calculate new variables. asignment is awkward, but what can you do.
ev['Electron',
'absMiniIso'] = ev.Electron.miniPFRelIso_all * ev.Electron.pt
ev['Electron', 'etaSC'] = ev.Electron.eta + ev.Electron.deltaEtaSC
if self.year == 2017 or self.year == 2018:
I_1 = 0.07
I_2 = 0.78
I_3 = 8.0
elif self.year == 2016:
I_1 = 0.12
I_2 = 0.80
I_3 = 7.2
# the following line is only needed if we do our own matching.
# right now, we keep using the NanoAOD match, but check the deltaR distance
# jet_index, mask_match, mask_nomatch = self.matchJets(ev.Electron, ev.Jet)
# this is what we are using:
# - jetRelIso if the matched jet is within deltaR<0.4, pfRelIso03_all otherwise
# - btagDeepFlavB discriminator of the matched jet if jet is within deltaR<0.4, 0 otherwise
# - pt_cone = 0.9*pt of matched jet if jet is within deltaR<0.4, pt/(pt+iso) otherwise
mask_close = (ak.fill_none(
ev.Electron.delta_r(ev.Electron.matched_jet), 99) < 0.4) * 1
mask_far = ~(ak.fill_none(
ev.Electron.delta_r(ev.Electron.matched_jet), 99) < 0.4) * 1
deepJet = ak.fill_none(ev.Electron.matched_jet.btagDeepFlavB,
0) * mask_close
jetRelIsoV2 = ev.Electron.jetRelIso * mask_close + ev.Electron.pfRelIso03_all * mask_far # default to 0 if no match
#conePt = 0.9 * ak.fill_none(ev.Electron.matched_jet.pt,0) * mask_close + ev.Electron.pt*(1 + ev.Electron.miniPFRelIso_all)*mask_far
PF_unflatten = ak.from_regular(
ev.Electron.miniPFRelIso_all[:, :, np.newaxis])
max_miniIso = ak.max(
ak.concatenate(
[PF_unflatten - I_1,
ak.zeros_like(PF_unflatten)], axis=2),
axis=2) #equivalent to max(0, ev.Muon.miniPFRelIso_all - I_1)
electron_pt_unflatten = ak.from_regular(ev.Electron.pt[:, :,
np.newaxis])
jet_pt_unflatten = ak.from_regular(
ev.Electron.matched_jet.pt[:, :, np.newaxis])
max_pt = ak.max(
ak.concatenate([electron_pt_unflatten, jet_pt_unflatten * I_2],
axis=2),
axis=2) #max(ev.Muon.pt, ev.Muon.matched_jet.pt * I_2)
conePt = (ev.Electron.pt *
(1 + max_miniIso)) * (ev.Electron.jetPtRelv2 > I_3) + (
max_pt * ~(ev.Electron.jetPtRelv2 > I_3))
ev['Electron', 'deepJet'] = ak.copy(deepJet)
ev['Electron', 'jetRelIsoV2'] = jetRelIsoV2
ev['Electron', 'conePt'] = conePt
ev['Electron', 'id'] = ak.ones_like(conePt) * id_level
ev['Electron', 'jetRelIso'] = ev.Electron.jetRelIso
ev['Electron', 'jetPtRelv2'] = ev.Electron.jetPtRelv2
self.cand = ev.Electron
self.getSelection()
if self.obj == "Electron" and self.wp == "tight":
self.selection = self.selection & self.getElectronMVAID(
) & self.getIsolation(0.07, 0.78, 8.0) & self.isTriggerSafeNoIso()
if self.v > 0: print(" - custom ID and multi-isolation")
if self.obj == "Electron" and self.wp == "tightFCNC":
self.selection = self.selection & self.getElectronMVAID(
) & self.getFCNCIsolation(ev.Electron.jetRelIso,
ev.Electron.jetPtRelv2, I_2,
I_3) & (ev.Electron.miniPFRelIso_all <
I_1) & self.isTriggerSafeNoIso()
if self.v > 0: print(" - custom ID and multi-isolation")
if self.obj == "Muon" and self.wp == "tight":
self.selection = self.selection & self.getIsolation(
0.11, 0.74, 6.8)
if self.v > 0: print(" - custom multi-isolation")
#self.selection = self.selection & ak.fill_none(ev.Muon.matched_jet.btagDeepFlavB<0.2770, True)
#self.selection = self.selection & (ev.Muon.matched_jet.btagDeepFlavB<0.2770)
#if self.v>0: print (" - deepJet")
if self.obj == "Muon" and self.wp == "tightFCNC":
self.selection = self.selection & self.getFCNCIsolation(
ev.Muon.jetRelIso, ev.Muon.jetPtRelv2, I_2,
I_3) & (ev.Muon.miniPFRelIso_all < I_1)
if self.v > 0: print(" - custom multi-isolation")
if self.obj == "Electron" and (self.wp == "tightTTH"
or self.wp == 'fakeableTTH'
or self.wp == "tightSSTTH"
or self.wp == 'fakeableSSTTH'):
self.selection = self.selection & self.getSigmaIEtaIEta
if self.v > 0: print(" - SigmaIEtaIEta")
#self.selection = self.selection & ak.fill_none(ev.Electron.matched_jet.btagDeepFlavB<0.2770, True)
#self.selection = self.selection & (ev.Electron.matched_jet.btagDeepFlavB<0.2770)
#self.selection = self.selection & (ev.Jet[ev.Electron.jetIdx].btagDeepFlavB<0.2770)
#if self.v>0: print (" - deepJet")
if self.obj == "Electron" and self.wp == "looseFCNC":
self.selection = self.selection & (ev.Electron.miniPFRelIso_all <
0.4)
if self.obj == 'Muon' and (self.wp == 'fakeableTTH'
or self.wp == 'fakeableSSTTH'):
#self.selection = self.selection & (self.cand.deepJet < self.getThreshold(self.cand.conePt, min_pt=20, max_pt=45, low=0.2770, high=0.0494))
self.selection = self.selection & (ak.fill_none(
ev.Muon.matched_jet.btagDeepFlavB, 0) < self.getThreshold(
self.cand.conePt, min_pt=20, max_pt=45))
if self.v > 0: print(" - interpolated deepJet")
if self.obj == "Muon" and self.wp == "looseFCNC":
self.selection = self.selection & (ev.Muon.miniPFRelIso_all < 0.4)
def test_toListOffsetArray64():
assert ak.from_regular(empty).tolist() == []