def test():
data = ak.Array([[1, 3, 5, 4, 2], [], [2, 3, 1], [5]])
assert ak.min(data, axis=1, initial=4).tolist() == [1, None, 1, 4]
assert ak.max(data, axis=1, initial=4).tolist() == [5, None, 4, 5]
data = ak.Array([[1.1, 3.3, 5.5, 4.4, 2.2], [], [2.2, 3.3, 1.1], [5.5]])
assert ak.min(data, axis=1, initial=4).tolist() == [1.1, None, 1.1, 4]
assert ak.max(data, axis=1, initial=4).tolist() == [5.5, None, 4, 5.5]
def test():
array = ak.Array([
[0, 1, None],
[None, 3],
[],
None,
[4, 5, None, 6],
])
assert ak.sum(array, axis=-1).tolist() == [
1, # 0 + 1
3, # 3
0, # list is empty
None, # list is missing
15, # 4 + 5 + 6
]
assert ak.sum(array, axis=-2).tolist() == [
4,
9,
0,
6,
# 0+4
# 1+3+5
# no data
# 6
]
assert ak.min(array, axis=-1).tolist() == [
0, # min([0, 1])
3, # min([3])
None, # list is empty
None, # list is missing
4, # min([4, 5, 6])
]
assert ak.min(array, axis=-2).tolist() == [
0,
1,
None,
6,
# 0,4
# 1,3,5
# no data
# 6
]
# first bug-fix: single '?'
assert str(ak.min(array, axis=-1).type) == "5 * ?int64"
# second bug-fix: correct mask_identity=False behavior
assert ak.ptp(array, axis=-1).tolist() == [1, 0, None, None, 2]
assert ak.ptp(array, axis=-1,
mask_identity=False).tolist() == [1, 0, 0, None, 2]
assert ak.ptp(array, axis=-2).tolist() == [4, 4, None, 0]
assert ak.ptp(array, axis=-2, mask_identity=False).tolist() == [4, 4, 0, 0]
def test_date_time():
numpy_array = np.array(
["2020-07-27T10:41:11", "2019-01-01", "2020-01-01"], "datetime64[s]"
)
array = ak.Array(numpy_array)
assert str(array.type) == "3 * datetime64"
assert array.tolist() == [
np.datetime64("2020-07-27T10:41:11"),
np.datetime64("2019-01-01T00:00:00"),
np.datetime64("2020-01-01T00:00:00"),
]
for i in range(len(array)):
assert ak.to_numpy(array)[i] == numpy_array[i]
date_time = np.datetime64("2020-07-27T10:41:11.200000011", "us")
array1 = ak.Array(np.array(["2020-07-27T10:41:11.200000011"], "datetime64[us]"))
assert np.datetime64(array1[0], "us") == date_time
assert ak.to_list(ak.from_iter(array1)) == [
np.datetime64("2020-07-27T10:41:11.200000")
]
assert ak.max(array) == numpy_array[0]
assert ak.min(array) == numpy_array[1]
def test_minmax():
assert ak.min(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]])) == 1 + 5j
assert ak.max(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]])) == 3 + 3j
assert ak.min(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]]),
axis=1).tolist() == [
1 + 5j,
None,
3 + 3j,
]
assert ak.max(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]]),
axis=1).tolist() == [
2 + 4j,
None,
3 + 3j,
]
assert ak.argmin(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]]),
axis=1).tolist() == [0, None, 0]
assert ak.argmax(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]]),
axis=1).tolist() == [1, None, 0]
def test():
layout = ak.layout.ListOffsetArray64(
ak.layout.Index64(np.array([0, 1], dtype=np.int64)),
ak.layout.IndexedArray64(
ak.layout.Index64(np.array([0, 1, 2, 3], dtype=np.int64)),
ak.layout.RegularArray(
ak.layout.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4])), 4
),
),
)
array = ak.Array(layout)
assert str(ak.min(array, axis=-1, mask_identity=False).type) == "1 * var * float64"
def test_IndexedOptionArray():
content = ak.Array([1.1, 2.2, 3.3, 4.4, 5.5]).layout
index = ak.layout.Index64(np.array([4, 2, -1, -1, 1, 0, 1]))
array = ak.Array(ak.layout.IndexedOptionArray64(index, content))
assert array.tolist() == [5.5, 3.3, None, None, 2.2, 1.1, 2.2]
assert ak.min(array, axis=0) == 1.1
assert ak.argmin(array, axis=0) == 5
assert ak.argmin(ak.Array([[2.2, 1.1], [None, 3.3], [2.2, 1.1]]),
axis=-1).tolist() == [1, 1, 1]
assert ak.argmin(ak.Array([[2.2, 1.1], [None, 3.3], [2.2, None, 1.1]]),
axis=-1).tolist() == [1, 1, 2]
assert ak.argmin(ak.Array([[2.2, 1.1], [3.3, None], [2.2, None, 1.1]]),
axis=-1).tolist() == [1, 0, 2]
assert ak.argmin(ak.Array([[2.2, 1.1, 0.0], [], [None, 0.5], [2, 1]]),
axis=0).tolist() == [3, 2, 0]
assert ak.argmin(ak.Array([[2.2, 1.1, 0.0], [], [0.5, None], [2, 1]]),
axis=0).tolist() == [2, 3, 0]
assert ak.argmin(ak.Array([[2.2, 1.1, 0.0], [0.5, None], [], [2, 1]]),
axis=0).tolist() == [1, 3, 0]
def process(self, events):
events = events[
ak.num(events.Jet) >
0] #corrects for rare case where there isn't a single jet in event
output = self.accumulator.identity()
# we can use a very loose preselection to filter the events. nothing is done with this presel, though
presel = ak.num(events.Jet) >= 0
ev = events[presel]
dataset = ev.metadata['dataset']
# load the config - probably not needed anymore
# cfg = loadConfig()
output['totalEvents']['all'] += len(events)
output['skimmedEvents']['all'] += len(ev)
### For FCNC, we want electron -> tightTTH
electron = Collections(ev, "Electron", "tightFCNC").get()
fakeableelectron = Collections(ev, "Electron", "fakeableFCNC").get()
muon = Collections(ev, "Muon", "tightFCNC").get()
fakeablemuon = Collections(ev, "Muon", "fakeableFCNC").get()
##Jets
Jets = events.Jet
## MET -> can switch to puppi MET
met_pt = ev.MET.pt
met_phi = ev.MET.phi
lepton = fakeablemuon #ak.concatenate([fakeablemuon, fakeableelectron], axis=1)
mt_lep_met = mt(lepton.pt, lepton.phi, ev.MET.pt, ev.MET.phi)
min_mt_lep_met = ak.min(mt_lep_met, axis=1)
selection = PackedSelection()
selection.add('MET<20', (ev.MET.pt < 20))
selection.add('mt<20', min_mt_lep_met < 20)
#selection.add('MET<19', (ev.MET.pt<19) )
selection_reqs = ['MET<20', 'mt<20'] #, 'MET<19']
fcnc_reqs_d = {sel: True for sel in selection_reqs}
fcnc_selection = selection.require(**fcnc_reqs_d)
# define the weight
weight = Weights(len(ev))
if not dataset == 'MuonEG':
# generator weight
weight.add("weight", ev.genWeight)
jets = getJets(
ev, maxEta=2.4, minPt=25, pt_var='pt'
) #& (ak.num(jets[~match(jets, fakeablemuon, deltaRCut=1.0)])>=1)
single_muon_sel = (ak.num(muon) == 1) & (ak.num(fakeablemuon) == 1) | (
ak.num(muon) == 0) & (ak.num(fakeablemuon) == 1)
single_electron_sel = (ak.num(electron) == 1) & (
ak.num(fakeableelectron)
== 1) | (ak.num(electron) == 0) & (ak.num(fakeableelectron) == 1)
fcnc_muon_sel = (ak.num(
jets[~match(jets, fakeablemuon, deltaRCut=1.0)]) >=
1) & fcnc_selection & single_muon_sel
fcnc_electron_sel = (ak.num(
jets[~match(jets, fakeableelectron, deltaRCut=1.0)]) >=
1) & fcnc_selection & single_electron_sel
tight_muon_sel = (ak.num(muon) == 1) & fcnc_muon_sel
loose_muon_sel = (ak.num(fakeablemuon) == 1) & fcnc_muon_sel
tight_electron_sel = (ak.num(electron) == 1) & fcnc_electron_sel
loose_electron_sel = (ak.num(fakeableelectron)
== 1) & fcnc_electron_sel
output['single_mu_fakeable'].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(fakeablemuon[loose_muon_sel].conePt)),
eta=np.abs(
ak.to_numpy(ak.flatten(fakeablemuon[loose_muon_sel].eta))))
output['single_mu'].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(muon[tight_muon_sel].conePt)),
eta=np.abs(ak.to_numpy(ak.flatten(muon[tight_muon_sel].eta))))
output['single_e_fakeable'].fill(
dataset=dataset,
pt=ak.to_numpy(
ak.flatten(fakeableelectron[loose_electron_sel].conePt)),
eta=np.abs(
ak.to_numpy(
ak.flatten(fakeableelectron[loose_electron_sel].eta))))
output['single_e'].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(electron[tight_electron_sel].conePt)),
eta=np.abs(
ak.to_numpy(ak.flatten(electron[tight_electron_sel].eta))))
if self.debug:
#create pandas dataframe for debugging
passed_events = ev[tight_muon_sel]
passed_muons = muon[tight_muon_sel]
event_p = ak.to_pandas(passed_events[["event"]])
event_p["MET_PT"] = passed_events["MET"]["pt"]
event_p["mt"] = min_mt_lep_met[tight_muon_sel]
event_p["num_tight_mu"] = ak.to_numpy(ak.num(muon)[tight_muon_sel])
event_p["num_loose_mu"] = ak.num(fakeablemuon)[tight_muon_sel]
muon_p = ak.to_pandas(
ak.flatten(passed_muons)[[
"pt", "conePt", "eta", "dz", "dxy", "ptErrRel",
"miniPFRelIso_all", "jetRelIsoV2", "jetRelIso",
"jetPtRelv2"
]])
#convert to numpy array for the output
events_array = pd.concat([muon_p, event_p], axis=1)
events_to_add = [6886009]
for e in events_to_add:
tmp_event = ev[ev.event == e]
added_event = ak.to_pandas(tmp_event[["event"]])
added_event["MET_PT"] = tmp_event["MET"]["pt"]
added_event["mt"] = min_mt_lep_met[ev.event == e]
added_event["num_tight_mu"] = ak.to_numpy(
ak.num(muon)[ev.event == e])
added_event["num_loose_mu"] = ak.to_numpy(
ak.num(fakeablemuon)[ev.event == e])
add_muon = ak.to_pandas(
ak.flatten(muon[ev.event == e])[[
"pt", "conePt", "eta", "dz", "dxy", "ptErrRel",
"miniPFRelIso_all", "jetRelIsoV2", "jetRelIso",
"jetPtRelv2"
]])
add_concat = pd.concat([add_muon, added_event], axis=1)
events_array = pd.concat([events_array, add_concat], axis=0)
output['muons_df'] += processor.column_accumulator(
events_array.to_numpy())
return output
def process(self, events):
output = self.accumulator.identity()
output['total']['all'] += len(events)
# use a very loose preselection to filter the events
presel = ak.num(events.Jet) > 2
ev = events[presel]
dataset = ev.metadata['dataset']
# load the config - probably not needed anymore
cfg = loadConfig()
## Muons
muon = Collections(ev, "Muon", "vetoTTH").get()
tightmuon = Collections(ev, "Muon", "tightTTH").get()
dimuon = choose(muon, 2)
SSmuon = ak.any((dimuon['0'].charge * dimuon['1'].charge) > 0, axis=1)
leading_muon_idx = ak.singletons(ak.argmax(muon.pt, axis=1))
leading_muon = muon[leading_muon_idx]
## Electrons
electron = Collections(ev, "Electron", "vetoTTH").get()
tightelectron = Collections(ev, "Electron", "tightTTH").get()
dielectron = choose(electron, 2)
SSelectron = ak.any(
(dielectron['0'].charge * dielectron['1'].charge) > 0, axis=1)
leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1))
leading_electron = electron[leading_electron_idx]
## Merge electrons and muons - this should work better now in ak1
dilepton = cross(muon, electron)
SSlepton = ak.any((dilepton['0'].charge * dilepton['1'].charge) > 0,
axis=1)
lepton = ak.concatenate([muon, electron], axis=1)
leading_lepton_idx = ak.singletons(ak.argmax(lepton.pt, axis=1))
leading_lepton = lepton[leading_lepton_idx]
trailing_lepton_idx = ak.singletons(ak.argmin(lepton.pt, axis=1))
trailing_lepton = lepton[trailing_lepton_idx]
dilepton_mass = (leading_lepton + trailing_lepton).mass
dilepton_pt = (leading_lepton + trailing_lepton).pt
dilepton_dR = delta_r(leading_lepton, trailing_lepton)
mt_lep_met = mt(lepton.pt, lepton.phi, ev.MET.pt, ev.MET.phi)
min_mt_lep_met = ak.min(mt_lep_met, axis=1)
## Jets
jet = getJets(ev, minPt=25, maxEta=4.7, pt_var='pt_nom')
jet = jet[ak.argsort(
jet.pt_nom, ascending=False
)] # need to sort wrt smeared and recorrected jet pt
jet = jet[~match(jet, muon,
deltaRCut=0.4)] # remove jets that overlap with muons
jet = jet[~match(
jet, electron,
deltaRCut=0.4)] # remove jets that overlap with electrons
central = jet[(abs(jet.eta) < 2.4)]
btag = getBTagsDeepFlavB(
jet, year=self.year) # should study working point for DeepJet
light = getBTagsDeepFlavB(jet, year=self.year, invert=True)
fwd = getFwdJet(light)
fwd_noPU = getFwdJet(light, puId=False)
tau = getTaus(ev)
track = getIsoTracks(ev)
## forward jets
j_fwd = fwd[ak.singletons(ak.argmax(
fwd.p, axis=1))] # highest momentum spectator
high_score_btag = central[ak.argsort(central.btagDeepFlavB)][:, :2]
bl = cross(lepton, high_score_btag)
bl_dR = delta_r(bl['0'], bl['1'])
min_bl_dR = ak.min(bl_dR, axis=1)
jf = cross(j_fwd, jet)
mjf = (jf['0'] + jf['1']).mass
j_fwd2 = jf[ak.singletons(
ak.argmax(mjf, axis=1)
)]['1'] # this is the jet that forms the largest invariant mass with j_fwd
delta_eta = ak.fill_none(
ak.pad_none(abs(j_fwd2.eta - j_fwd.eta), 1, clip=True), 0)
## MET -> can switch to puppi MET
met_pt = ev.MET.pt
met_phi = ev.MET.phi
## other variables
ht = ak.sum(jet.pt, axis=1)
st = met_pt + ht + ak.sum(muon.pt, axis=1) + ak.sum(electron.pt,
axis=1)
## event selectors
filters = getFilters(ev, year=self.year, dataset=dataset)
dilep = ((ak.num(tightelectron) + ak.num(tightmuon)) == 2)
lep0pt = ((ak.num(electron[(electron.pt > 25)]) +
ak.num(muon[(muon.pt > 25)])) > 0)
lep1pt = ((ak.num(electron[(electron.pt > 20)]) +
ak.num(muon[(muon.pt > 20)])) > 1)
lepveto = ((ak.num(electron) + ak.num(muon)) == 2)
selection = PackedSelection()
selection.add('lepveto', lepveto)
selection.add('dilep', dilep)
selection.add('filter', (filters))
selection.add('p_T(lep0)>25', lep0pt)
selection.add('p_T(lep1)>20', lep1pt)
selection.add('SS', (SSlepton | SSelectron | SSmuon))
selection.add('N_jet>3', (ak.num(jet) >= 4))
selection.add('N_central>2', (ak.num(central) >= 3))
selection.add('N_btag>0', (ak.num(btag) >= 1))
selection.add('N_fwd>0', (ak.num(fwd) >= 1))
#ss_reqs = ['lepveto', 'dilep', 'filter', 'p_T(lep0)>25', 'p_T(lep1)>20', 'SS']
ss_reqs = [
'lepveto', 'dilep', 'filter', 'p_T(lep0)>25', 'p_T(lep1)>20', 'SS'
]
#bl_reqs = ss_reqs + ['N_jet>3', 'N_central>2', 'N_btag>0', 'N_fwd>0']
bl_reqs = ss_reqs + ['N_jet>3', 'N_central>2', 'N_btag>0']
ss_reqs_d = {sel: True for sel in ss_reqs}
ss_selection = selection.require(**ss_reqs_d)
bl_reqs_d = {sel: True for sel in bl_reqs}
BL = selection.require(**bl_reqs_d)
weight = Weights(len(ev))
if not dataset == 'MuonEG':
# lumi weight
weight.add("weight", ev.weight)
# PU weight - not in the babies...
weight.add("PU",
ev.puWeight,
weightUp=ev.puWeightUp,
weightDown=ev.puWeightDown,
shift=False)
# b-tag SFs
weight.add("btag", self.btagSF.Method1a(btag, light))
# lepton SFs
weight.add("lepton", self.leptonSF.get(electron, muon))
#cutflow = Cutflow(output, ev, weight=weight)
#cutflow_reqs_d = {}
#for req in bl_reqs:
# cutflow_reqs_d.update({req: True})
# cutflow.addRow( req, selection.require(**cutflow_reqs_d) )
labels = {
'topW_v3': 0,
'TTW': 1,
'TTZ': 2,
'TTH': 3,
'ttbar': 4,
'ttbar1l_MG': 4
}
if dataset in labels:
label_mult = labels[dataset]
else:
label_mult = 5
label = np.ones(len(ev[BL])) * label_mult
output["n_lep"] += processor.column_accumulator(
ak.to_numpy((ak.num(electron) + ak.num(muon))[BL]))
output["n_lep_tight"] += processor.column_accumulator(
ak.to_numpy((ak.num(tightelectron) + ak.num(tightmuon))[BL]))
output["lead_lep_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(leading_lepton[BL].pt, axis=1)))
output["lead_lep_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(leading_lepton[BL].eta, axis=1)))
output["lead_lep_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(leading_lepton[BL].phi, axis=1)))
output["lead_lep_charge"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(leading_lepton[BL].charge, axis=1)))
output["sublead_lep_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(trailing_lepton[BL].pt, axis=1)))
output["sublead_lep_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(trailing_lepton[BL].eta, axis=1)))
output["sublead_lep_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(trailing_lepton[BL].phi, axis=1)))
output["sublead_lep_charge"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(trailing_lepton[BL].charge, axis=1)))
output["lead_jet_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 0:1][BL].pt, axis=1)))
output["lead_jet_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 0:1][BL].eta, axis=1)))
output["lead_jet_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 0:1][BL].phi, axis=1)))
output["sublead_jet_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 1:2][BL].pt, axis=1)))
output["sublead_jet_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 1:2][BL].eta, axis=1)))
output["sublead_jet_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 1:2][BL].phi, axis=1)))
output["lead_btag_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 0:1][BL].pt, axis=1)))
output["lead_btag_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 0:1][BL].eta, axis=1)))
output["lead_btag_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 0:1][BL].phi, axis=1)))
output["sublead_btag_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 1:2][BL].pt, axis=1)))
output["sublead_btag_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 1:2][BL].eta, axis=1)))
output["sublead_btag_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 1:2][BL].phi, axis=1)))
output["fwd_jet_p"] += processor.column_accumulator(
ak.to_numpy(
ak.flatten(ak.fill_none(ak.pad_none(j_fwd[BL].p, 1, clip=True),
0),
axis=1)))
output["fwd_jet_pt"] += processor.column_accumulator(
ak.to_numpy(
ak.flatten(ak.fill_none(
ak.pad_none(j_fwd[BL].pt, 1, clip=True), 0),
axis=1)))
output["fwd_jet_eta"] += processor.column_accumulator(
ak.to_numpy(
ak.flatten(ak.fill_none(
ak.pad_none(j_fwd[BL].eta, 1, clip=True), 0),
axis=1)))
output["fwd_jet_phi"] += processor.column_accumulator(
ak.to_numpy(
ak.flatten(ak.fill_none(
ak.pad_none(j_fwd[BL].phi, 1, clip=True), 0),
axis=1)))
output["mjj_max"] += processor.column_accumulator(
ak.to_numpy(ak.fill_none(ak.max(mjf[BL], axis=1), 0)))
output["delta_eta_jj"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(delta_eta[BL], axis=1)))
output["met"] += processor.column_accumulator(ak.to_numpy(met_pt[BL]))
output["ht"] += processor.column_accumulator(ak.to_numpy(ht[BL]))
output["st"] += processor.column_accumulator(ak.to_numpy(st[BL]))
output["n_jet"] += processor.column_accumulator(
ak.to_numpy(ak.num(jet[BL])))
output["n_btag"] += processor.column_accumulator(
ak.to_numpy(ak.num(btag[BL])))
output["n_fwd"] += processor.column_accumulator(
ak.to_numpy(ak.num(fwd[BL])))
output["n_central"] += processor.column_accumulator(
ak.to_numpy(ak.num(central[BL])))
output["n_tau"] += processor.column_accumulator(
ak.to_numpy(ak.num(tau[BL])))
output["n_track"] += processor.column_accumulator(
ak.to_numpy(ak.num(track[BL])))
output["dilepton_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(dilepton_pt[BL], axis=1)))
output["dilepton_mass"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(dilepton_mass[BL], axis=1)))
output["min_bl_dR"] += processor.column_accumulator(
ak.to_numpy(min_bl_dR[BL]))
output["min_mt_lep_met"] += processor.column_accumulator(
ak.to_numpy(min_mt_lep_met[BL]))
output["label"] += processor.column_accumulator(label)
output["weight"] += processor.column_accumulator(weight.weight()[BL])
output["presel"]["all"] += len(ev[ss_selection])
output["sel"]["all"] += len(ev[BL])
return output
def test_min_max():
array = ak.Array(
[
[
np.datetime64("2020-03-27T10:41:11"),
np.datetime64("2020-01-27T10:41:11"),
np.datetime64("2020-05"),
np.datetime64("2020-01-27T10:41:11"),
np.datetime64("2020-04-27T10:41:11"),
],
[
np.datetime64("2020-04-27"),
np.datetime64("2020-02-27T10:41:11"),
np.datetime64("2020-01-27T10:41:11"),
np.datetime64("2020-06-27T10:41:11"),
],
[
np.datetime64("2020-02-27T10:41:11"),
np.datetime64("2020-03-27T10:41:11"),
np.datetime64("2020-01-27T10:41:11"),
],
]
)
assert ak.to_list(array) == [
[
np.datetime64("2020-03-27T10:41:11"),
np.datetime64("2020-01-27T10:41:11"),
np.datetime64("2020-05"),
np.datetime64("2020-01-27T10:41:11"),
np.datetime64("2020-04-27T10:41:11"),
],
[
np.datetime64("2020-04-27"),
np.datetime64("2020-02-27T10:41:11"),
np.datetime64("2020-01-27T10:41:11"),
np.datetime64("2020-06-27T10:41:11"),
],
[
np.datetime64("2020-02-27T10:41:11"),
np.datetime64("2020-03-27T10:41:11"),
np.datetime64("2020-01-27T10:41:11"),
],
]
assert ak.min(array) == np.datetime64("2020-01-27T10:41:11")
assert ak.max(array) == np.datetime64("2020-06-27T10:41:11")
assert ak.to_list(ak.min(array, axis=0)) == [
np.datetime64("2020-02-27T10:41:11"),
np.datetime64("2020-01-27T10:41:11"),
np.datetime64("2020-01-27T10:41:11"),
np.datetime64("2020-01-27T10:41:11"),
np.datetime64("2020-04-27T10:41:11"),
]
assert ak.to_list(ak.max(array, axis=0)) == [
np.datetime64("2020-04-27T00:00:00"),
np.datetime64("2020-03-27T10:41:11"),
np.datetime64("2020-05-01T20:56:24"),
np.datetime64("2020-06-27T10:41:11"),
np.datetime64("2020-04-27T10:41:11"),
]
assert ak.to_list(ak.min(array, axis=1)) == [
np.datetime64("2020-01-27T10:41:11"),
np.datetime64("2020-01-27T10:41:11"),
np.datetime64("2020-01-27T10:41:11"),
]
assert ak.to_list(ak.max(array, axis=1)) == [
np.datetime64("2020-05-01T20:56:24"),
np.datetime64("2020-06-27T10:41:11"),
np.datetime64("2020-03-27T10:41:11"),
]
def varGetter(dataset, events, scaleFactor):
varVal = {}
luminosity = 21071.0 + 38654.0
evtw = luminosity * events.Weight * scaleFactor
eCounter = np.where(evtw >= 0, 1, -1)
obj = ob.Objects(events)
jets = obj.goodJets()
bjets = obj.goodBJets(dataset, jets)
fjets = obj.goodFatJets()
gfjets = obj.goodGenFatJets()
electrons = obj.goodElectrons()
muons = obj.goodMuons()
met = events.MET
metPhi = events.METPhi
mtAK8 = events.MT_AK8
madHT = events.madHT
jetAK8Eta = fjets.eta
jetAK8Phi = fjets.phi
j1_etaAK8 = jetVar_i(jetAK8Eta, 0)
j2_etaAK8 = jetVar_i(jetAK8Eta, 1)
j1_phiAK8 = jetVar_i(jetAK8Phi, 0)
j2_phiAK8 = jetVar_i(jetAK8Phi, 1)
## GenJetsAK8_hvCategory is only present in the signal samples, not the V17 background
jetCats = []
bkgKeys = ["QCD", "TTJets", "WJets", "ZJets"]
isSignal = True
for k in bkgKeys:
if k in dataset:
isSignal = False
break
if isSignal:
jetsAK8GenInd = fjets.genIndex
for gji in range(len(jetsAK8GenInd)):
genInd = jetsAK8GenInd[gji]
genCat = gfjets.hvCategory[gji]
if (len(genCat) > 0) and (len(genInd) > 0):
if np.max(genInd) < len(genCat):
jetCats.append(list(genCat[genInd]))
else:
jetCats.append([-1] * len(genInd))
else:
jetCats.append([-1] * len(genInd))
jetCats = ak.Array(jetCats)
varVal['JetsAK8_hvCategory'] = [jetCats, 'fjw']
else:
varVal['JetsAK8_hvCategory'] = [
awkwardReshape(fjets,
np.ones(len(evtw)) * -1), 'fjw'
]
ew = awkwardReshape(electrons, evtw)
mw = awkwardReshape(muons, evtw)
jw = awkwardReshape(jets, evtw)
fjw = awkwardReshape(fjets, evtw)
ht = ak.sum(jets.pt, axis=1)
st = ht + met
metrht = met / ht
metrst = met / st
# AK4 Jet Variables
jetPhi = jets.phi
jetEta = jets.eta
j1_eta = jetVar_i(jetEta, 0)
j2_eta = jetVar_i(jetEta, 1)
j1_phi = jetVar_i(jetPhi, 0)
j2_phi = jetVar_i(jetPhi, 1)
dPhij1 = deltaPhi(j1_phi, metPhi)
dPhij2 = deltaPhi(j2_phi, metPhi)
dPhij1rdPhij2 = dPhij1 / dPhij2
dPhiMinj = ak.min(deltaPhi(jetPhi, metPhi), axis=1, mask_identity=False)
dEtaj12 = deltaEta(j1_eta, j2_eta)
deltaR12j = delta_R(j1_eta, j2_eta, j1_phi, j2_phi)
# AK8 Jet Variables
jetAK8pT = fjets.pt
jetAK8Phi = fjets.phi
jetAK8Eta = fjets.eta
jetAK8M = fjets.mass
j1_etaAK8 = jetVar_i(jetAK8Eta, 0)
j2_etaAK8 = jetVar_i(jetAK8Eta, 1)
j1_phiAK8 = jetVar_i(jetAK8Phi, 0)
j2_phiAK8 = jetVar_i(jetAK8Phi, 1)
dPhij1AK8 = deltaPhi(j1_phiAK8, metPhi)
dPhij2AK8 = deltaPhi(j2_phiAK8, metPhi)
dPhij1rdPhij2AK8 = dPhij1AK8 / dPhij2AK8
dPhijAK8 = deltaPhi(jetAK8Phi, metPhi)
dPhiMinjAK8 = ak.min(dPhijAK8, axis=1, mask_identity=False)
dEtaj12AK8 = deltaEta(j1_etaAK8, j2_etaAK8)
deltaR12jAK8 = delta_R(j1_etaAK8, j2_etaAK8, j1_phiAK8, j2_phiAK8)
tau1 = fjets.NsubjettinessTau1
tau2 = fjets.NsubjettinessTau2
tau3 = fjets.NsubjettinessTau3
J_tau21 = tau2 / tau1
J_tau32 = tau3 / tau2
J1_tau21 = tauRatio(tau2, tau1, 0)
J1_tau32 = tauRatio(tau3, tau2, 0)
J2_tau21 = tauRatio(tau2, tau1, 1)
J2_tau32 = tauRatio(tau3, tau2, 1)
if len(bjets) > 0:
nBJets = ak.num(bjets)
else:
nBJets = np.zeros(len(evtw))
varVal['jets'] = jets
varVal['fjets'] = fjets
varVal['eCounter'] = [eCounter, 'w1']
varVal['evtw'] = [evtw, 'evtw']
varVal['jw'] = [jw, 'jw']
varVal['fjw'] = [fjw, 'fjw']
varVal['njets'] = [ak.num(jets), 'evtw']
varVal['njetsAK8'] = [ak.num(fjets), 'evtw']
varVal['nb'] = [nBJets, 'evtw']
varVal['nl'] = [(ak.num(electrons) + ak.num(muons)), 'evtw']
varVal['ht'] = [ht, 'evtw']
varVal['st'] = [st, 'evtw']
varVal['met'] = [met, 'evtw']
varVal['metPhi'] = [metPhi, 'evtw']
varVal['madHT'] = [madHT, 'evtw']
varVal['jPt'] = [jets.pt, 'jw']
varVal['jEta'] = [jetEta, 'jw']
varVal['jPhi'] = [jetPhi, 'jw']
varVal['jAxismajor'] = [jets.axismajor, 'jw']
varVal['jAxisminor'] = [jets.axisminor, 'jw']
varVal['jPtD'] = [jets.ptD, 'jw']
varVal['dPhiMinjMET'] = [dPhiMinj, 'evtw']
varVal['dEtaj12'] = [dEtaj12, 'evtw']
varVal['dRJ12'] = [deltaR12j, 'evtw']
varVal['jPtAK8'] = [fjets.pt, 'fjw']
varVal['jEtaAK8'] = [jetAK8Eta, 'fjw']
varVal['jPhiAK8'] = [jetAK8Phi, 'fjw']
varVal['jAxismajorAK8'] = [fjets.axismajor, 'fjw']
varVal['jAxisminorAK8'] = [fjets.axisminor, 'fjw']
varVal['jChEMEFractAK8'] = [fjets.chargedEmEnergyFraction, 'fjw']
varVal['jChHadEFractAK8'] = [fjets.chargedHadronEnergyFraction, 'fjw']
varVal['jChHadMultAK8'] = [fjets.chargedHadronMultiplicity, 'fjw']
varVal['jChMultAK8'] = [fjets.chargedMultiplicity, 'fjw']
varVal['jdoubleBDiscriminatorAK8'] = [fjets.doubleBDiscriminator, 'fjw']
varVal['jecfN2b1AK8'] = [fjets.ecfN2b1, 'fjw']
varVal['jecfN2b2AK8'] = [fjets.ecfN2b2, 'fjw']
varVal['jecfN3b1AK8'] = [fjets.ecfN3b1, 'fjw']
varVal['jecfN3b2AK8'] = [fjets.ecfN3b2, 'fjw']
varVal['jEleEFractAK8'] = [fjets.electronEnergyFraction, 'fjw']
varVal['jEleMultAK8'] = [fjets.electronMultiplicity, 'fjw']
varVal['jGirthAK8'] = [fjets.girth, 'fjw']
varVal['jHfEMEFractAK8'] = [fjets.hfEMEnergyFraction, 'fjw']
varVal['jHfHadEFractAK8'] = [fjets.hfHadronEnergyFraction, 'fjw']
varVal['jMultAK8'] = [fjets.multiplicity, 'fjw']
varVal['jMuEFractAK8'] = [fjets.muonEnergyFraction, 'fjw']
varVal['jMuMultAK8'] = [fjets.muonMultiplicity, 'fjw']
varVal['jNeuEmEFractAK8'] = [fjets.neutralEmEnergyFraction, 'fjw']
varVal['jNeuHadEFractAK8'] = [fjets.neutralHadronEnergyFraction, 'fjw']
varVal['jNeuHadMultAK8'] = [fjets.neutralHadronMultiplicity, 'fjw']
varVal['jNeuMultAK8'] = [fjets.neutralMultiplicity, 'fjw']
varVal['jTau1AK8'] = [tau1, 'fjw']
varVal['jTau2AK8'] = [tau2, 'fjw']
varVal['jTau3AK8'] = [tau3, 'fjw']
varVal['jTau21AK8'] = [J_tau21, 'fjw']
varVal['jTau32AK8'] = [J_tau32, 'fjw']
varVal['jNumBhadronsAK8'] = [fjets.NumBhadrons, 'fjw']
varVal['jNumChadronsAK8'] = [fjets.NumChadrons, 'fjw']
varVal['jPhoEFractAK8'] = [fjets.photonEnergyFraction, 'fjw']
varVal['jPhoMultAK8'] = [fjets.photonMultiplicity, 'fjw']
varVal['jPtDAK8'] = [fjets.ptD, 'fjw']
varVal['jSoftDropMassAK8'] = [fjets.softDropMass, 'fjw']
varVal['dPhijMETAK8'] = [dPhijAK8, 'fjw']
varVal['dPhiMinjMETAK8'] = [dPhiMinjAK8, 'evtw']
varVal['dEtaj12AK8'] = [dEtaj12AK8, 'evtw']
varVal['dRJ12AK8'] = [deltaR12jAK8, 'evtw']
varVal['mT'] = [mtAK8, 'evtw']
varVal['METrHT_pt30'] = [metrht, 'evtw']
varVal['METrST_pt30'] = [metrst, 'evtw']
varVal['j1Pt'] = [jetVar_i(jets.pt, 0), 'evtw']
varVal['j1Eta'] = [j1_eta, 'evtw']
varVal['j1Phi'] = [j1_phi, 'evtw']
varVal['j1Axismajor'] = [jetVar_i(jets.axismajor, 0), 'evtw']
varVal['j1Axisminor'] = [jetVar_i(jets.axisminor, 0), 'evtw']
varVal['j1PtD'] = [jetVar_i(jets.ptD, 0), 'evtw']
varVal['dPhij1MET'] = [dPhij1, 'evtw']
varVal['j2Pt'] = [jetVar_i(jets.pt, 1), 'evtw']
varVal['j2Eta'] = [j2_eta, 'evtw']
varVal['j2Phi'] = [j2_phi, 'evtw']
varVal['j2Axismajor'] = [jetVar_i(jets.axismajor, 1), 'evtw']
varVal['j2Axisminor'] = [jetVar_i(jets.axisminor, 1), 'evtw']
varVal['j2PtD'] = [jetVar_i(jets.ptD, 1), 'evtw']
varVal['dPhij2MET'] = [dPhij2, 'evtw']
varVal['dPhij1rdPhij2'] = [dPhij1rdPhij2, 'evtw']
varVal['j1PtAK8'] = [jetVar_i(fjets.pt, 0), 'evtw']
varVal['j1EtaAK8'] = [j1_etaAK8, 'evtw']
varVal['j1PhiAK8'] = [j1_phiAK8, 'evtw']
varVal['j1AxismajorAK8'] = [jetVar_i(fjets.axismajor, 0), 'evtw']
varVal['j1AxisminorAK8'] = [jetVar_i(fjets.axisminor, 0), 'evtw']
varVal['j1GirthAK8'] = [jetVar_i(fjets.girth, 0), 'evtw']
varVal['j1PtDAK8'] = [jetVar_i(fjets.ptD, 0), 'evtw']
varVal['j1Tau1AK8'] = [jetVar_i(tau1, 0), 'evtw']
varVal['j1Tau2AK8'] = [jetVar_i(tau2, 0), 'evtw']
varVal['j1Tau3AK8'] = [jetVar_i(tau3, 0), 'evtw']
varVal['j1Tau21AK8'] = [J1_tau21, 'evtw']
varVal['j1Tau32AK8'] = [J1_tau32, 'evtw']
varVal['j1SoftDropMassAK8'] = [jetVar_i(fjets.softDropMass, 0), 'evtw']
varVal['dPhij1METAK8'] = [dPhij1AK8, 'evtw']
varVal['j2PtAK8'] = [jetVar_i(fjets.pt, 1), 'evtw']
varVal['j2EtaAK8'] = [j2_etaAK8, 'evtw']
varVal['j2PhiAK8'] = [j2_phiAK8, 'evtw']
varVal['j2AxismajorAK8'] = [jetVar_i(fjets.axismajor, 1), 'evtw']
varVal['j2AxisminorAK8'] = [jetVar_i(fjets.axisminor, 1), 'evtw']
varVal['j2GirthAK8'] = [jetVar_i(fjets.girth, 1), 'evtw']
varVal['j2PtDAK8'] = [jetVar_i(fjets.ptD, 1), 'evtw']
varVal['j2Tau1AK8'] = [jetVar_i(tau1, 1), 'evtw']
varVal['j2Tau2AK8'] = [jetVar_i(tau2, 1), 'evtw']
varVal['j2Tau3AK8'] = [jetVar_i(tau3, 1), 'evtw']
varVal['j2Tau21AK8'] = [J2_tau21, 'evtw']
varVal['j2Tau32AK8'] = [J2_tau32, 'evtw']
varVal['j2SoftDropMassAK8'] = [jetVar_i(fjets.softDropMass, 1), 'evtw']
varVal['dPhij2METAK8'] = [dPhij2AK8, 'evtw']
varVal['dPhij1rdPhij2AK8'] = [dPhij1rdPhij2AK8, 'evtw']
# varVal['GenJetsAK8_hvCategory'] = [GenJetsAK8_hvCategory.flatten(),gfjweight]
# varVal['mT2_f4_msm'] = [f4msmCom_vec(jetAK8pT,jetAK8Eta,jetAK8Phi,jetAK8M,met,metPhi,""),'evtw']
# varVal['mT2_f4_msm_dEta'] = [f4msmCom_vec(jetAK8pT,jetAK8Eta,jetAK8Phi,jetAK8M,met,metPhi,"dEta"),'evtw']
# varVal['mT2_f4_msm_dPhi'] = [f4msmCom_vec(jetAK8pT,jetAK8Eta,jetAK8Phi,jetAK8M,met,metPhi,"dPhi"),'evtw']
# varVal['mT2_f4_msm_dR'] = [f4msmCom_vec(jetAK8pT,jetAK8Eta,jetAK8Phi,jetAK8M,met,metPhi,"dR"),'evtw']
# varVal['GenJetsAK8_darkPtFrac'] = [GenJetsAK8_darkPtFrac.flatten(),gfjweight]
# varVal['GenMT2_AK8'] = [GenMT2_AK8,'evtw']
return varVal
def test_jagged_axis0():
assert ak.min(ak.Array([[1.1, 5.5], [4.4], [2.2, 3.3, 0.0, -10]]),
axis=0).tolist() == [1.1, 3.3, 0, -10]
assert ak.argmin(ak.Array([[1.1, 5.5], [4.4], [2.2, 3.3, 0.0, -10]]),
axis=0).tolist() == [0, 2, 2, 2]
def test_jagged_axis1():
# first is [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [4, 3, 2]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [5, 4, 3]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[], [], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [6, 5, 4]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[1.1], [1.1, 2.2], [], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [5, 4, 3]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [5, 4, 2]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [5, 3, 2]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0], []],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [4, 3, 2]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[1.1, 999, 999], [1.1, 2.2, 999], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [4, 3, 2]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[1.1, 999, 999, 999], [1.1, 2.2, 999], [1.1, 2.2, 3.3], [999, 2.0],
[1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3, 999]]
assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [4, 3, 2, 0]]
# first is [[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]]
array = ak.Array([
[[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [4, 3, 2]]
array = ak.Array([
[[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [5, 4, 3]]
array = ak.Array([
[[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[], [], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [6, 5, 4]]
array = ak.Array([
[[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[1.1], [1.1, 2.2], [], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [5, 4, 3]]
array = ak.Array([
[[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [5, 4, 2]]
array = ak.Array([
[[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [5, 3, 2]]
array = ak.Array([
[[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0], []],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [4, 3, 2]]
array = ak.Array([
[[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[1.1, 999, 999], [1.1, 2.2, 999], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [4, 3, 2]]
array = ak.Array([
[[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
[[1.1, 999, 999, 999], [1.1, 2.2, 999], [1.1, 2.2, 3.3], [999, 2.0],
[1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3, 999]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [4, 3, 2, 0]]
# first is [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]],
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [4, 3, 2]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]],
[[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [5, 4, 3]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]],
[[], [], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [6, 5, 4]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]],
[[1.1], [1.1, 2.2], [], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [5, 4, 3]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]],
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [5, 4, 2]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]],
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [5, 3, 2]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]],
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0], []],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [4, 3, 2]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]],
[[1.1, 999, 999], [1.1, 2.2, 999], [1.1, 2.2, 3.3], [999, 2.0], [1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [4, 3, 2]]
array = ak.Array([
[[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]],
[[1.1, 999, 999, 999], [1.1, 2.2, 999], [1.1, 2.2, 3.3], [999, 2.0],
[1.0]],
])
assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3, 999]]
assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [4, 3, 2, 0]]
def test_actual_issue():
ak_array = ak.Array([[[1, 2, 3], [], [4, 3, 2]], [[4, 5, 6], [], [2, 3,
4]]])
assert ak.min(ak_array, axis=0).tolist() == [[1, 2, 3], [], [2, 3, 2]]
def make_radius_compatibility_distributions():
global tree
all_obj_arrays = tree.arrays(filter_name="pT3*",
entry_start=0,
entry_stop=5,
library="ak")
matchedMask = all_obj_arrays.pT3_isFake == 0
layers = np.array(
list(map(process_layers, ak.flatten(all_obj_arrays.pT3_layer_binary))))
# layerTypes = np.array(list(map(process_layerType, layers)))
layerTypes = np.array(list(map(process_numbers, layers)))
unique_layerTypes = np.unique(layerTypes, axis=0)
unique_layerTypes = np.append(unique_layerTypes, "")
print(unique_layerTypes)
for layerType in unique_layerTypes:
print("layerType = {}".format(layerType))
pixelRadius = ak.to_numpy(ak.flatten(all_obj_arrays.pT3_pixelRadius))
pixelRadiusResMin = ak.to_numpy(
ak.flatten(all_obj_arrays.pT3_pixelRadiusMin))
pixelRadiusResMax = ak.to_numpy(
ak.flatten(all_obj_arrays.pT3_pixelRadiusMax))
tripletRadius = ak.to_numpy(
ak.flatten(all_obj_arrays.pT3_tripletRadius))
tripletRadiusResMin = ak.to_numpy(
ak.flatten(all_obj_arrays.pT3_tripletRadiusMin))
tripletRadiusResMax = ak.to_numpy(
ak.flatten(all_obj_arrays.pT3_tripletRadiusMax))
simRadius = ak.flatten(all_obj_arrays.pT3_matched_pt /
(2.99792458e-3 * 3.8))
simRadius = ak.flatten(simRadius)
pixelRadiusMin = ak.to_numpy(
ak.min([pixelRadiusResMin, pixelRadius2SMin], axis=0))
pixelRadiusMax = ak.to_numpy(
ak.max([pixelRadiusResMax, pixelRadius2SMax], axis=0))
tripletRadiusMin = ak.to_numpy(
ak.min([tripletRadiusResMin, tripletRadius2SMin], axis=0))
tripletRadiusMax = ak.to_numpy(
ak.max([tripletRadiusResMax, tripletRadius2SMax], axis=0))
qArrayInnerOuter = compute_interval_overlap(1.0 / pixelRadiusMax,
1.0 / pixelRadiusMin,
1.0 / tripletRadiusMax,
1.0 / tripletRadiusMin)
for name, qArray in {"innerOuter": qArrayInnerOuter}.items():
print("qName = ", name)
if layerType == "":
qArraySimTrackMatched = qArray[ak.to_numpy(
ak.flatten(matchedMask))]
else:
qArray = qArray[layerTypes == layerType]
qArraySimTrackMatched = qArray[ak.to_numpy(
ak.flatten(matchedMask)[layerTypes == layerType])]
print(
"{} total integral = {}, {} integral below zero = {}, sim-matched {} total integral = {}, sim-matched {} integral above zero = {}"
.format(name, len(qArray), name, sum(qArray < 0), name,
len(qArraySimTrackMatched), name,
sum(qArraySimTrackMatched > 0)))
make_plots(
qArray, qArraySimTrackMatched,
"overlap between 1/{} and 1/{}".format("Inner",
name[5:]), layerType)
def test_highlevel():
array = ak.Array(
[[[2, 3, 5], [], [7, 11], [13]], [], [[17, 19], [23]]], check_valid=True
)
assert ak.count(array) == 9
assert ak.to_list(ak.count(array, axis=-1)) == [[3, 0, 2, 1], [], [2, 1]]
assert ak.to_list(ak.count(array, axis=2)) == [[3, 0, 2, 1], [], [2, 1]]
assert ak.to_list(ak.count(array, axis=-1, keepdims=True)) == [
[[3], [0], [2], [1]],
[],
[[2], [1]],
]
assert ak.to_list(ak.count(array, axis=-2)) == [[3, 2, 1], [], [2, 1]]
assert ak.to_list(ak.count(array, axis=1)) == [[3, 2, 1], [], [2, 1]]
assert ak.to_list(ak.count(array, axis=-2, keepdims=True)) == [
[[3, 2, 1]],
[[]],
[[2, 1]],
]
assert ak.count_nonzero(array) == 9
assert ak.to_list(ak.count_nonzero(array, axis=-1)) == [[3, 0, 2, 1], [], [2, 1]]
assert ak.to_list(ak.count_nonzero(array, axis=-2)) == [[3, 2, 1], [], [2, 1]]
assert ak.sum(array) == 2 + 3 + 5 + 7 + 11 + 13 + 17 + 19 + 23
assert ak.to_list(ak.sum(array, axis=-1)) == [
[2 + 3 + 5, 0, 7 + 11, 13],
[],
[17 + 19, 23],
]
assert ak.to_list(ak.sum(array, axis=-2)) == [
[2 + 7 + 13, 3 + 11, 5],
[],
[17 + 23, 19],
]
assert ak.prod(array) == 2 * 3 * 5 * 7 * 11 * 13 * 17 * 19 * 23
assert ak.to_list(ak.prod(array, axis=-1)) == [
[2 * 3 * 5, 1, 7 * 11, 13],
[],
[17 * 19, 23],
]
assert ak.to_list(ak.prod(array, axis=-2)) == [
[2 * 7 * 13, 3 * 11, 5],
[],
[17 * 23, 19],
]
assert ak.min(array) == 2
assert ak.to_list(ak.min(array, axis=-1)) == [[2, None, 7, 13], [], [17, 23]]
assert ak.to_list(ak.min(array, axis=-2)) == [[2, 3, 5], [], [17, 19]]
assert ak.max(array) == 23
assert ak.to_list(ak.max(array, axis=-1)) == [[5, None, 11, 13], [], [19, 23]]
assert ak.to_list(ak.max(array, axis=-2)) == [[13, 11, 5], [], [23, 19]]
array = ak.Array(
[
[[True, False, True], [], [False, False], [True]],
[],
[[False, True], [True]],
],
check_valid=True,
)
assert ak.any(array) == True
assert ak.to_list(ak.any(array, axis=-1)) == [
[True, False, False, True],
[],
[True, True],
]
assert ak.to_list(ak.any(array, axis=-2)) == [[True, False, True], [], [True, True]]
assert ak.all(array) == False
assert ak.to_list(ak.all(array, axis=-1)) == [
[False, True, False, True],
[],
[False, True],
]
assert ak.to_list(ak.all(array, axis=-2)) == [
[False, False, True],
[],
[False, True],
]
def uproot_tree_to_numpy(fname,
MeanNormTuple,
inbranches_listlist,
nMaxslist,
nevents,
treename="ttree",
stop=None,
branches=None):
# array = uproot_root2array(fname, treename, stop=stop, branches=branches)
# Read in total number of events
totallengthperjet = 0
for i in range(len(nMaxslist)):
if nMaxslist[i] >= 0:
totallengthperjet += len(inbranches_listlist[i]) * nMaxslist[i]
else:
totallengthperjet += len(inbranches_listlist[i]) #flat branch
# branches = [ak.fill_none(ak.pad_none(tree[barr, target=feature_length), 0.) for feature_length, arr in zip( nMaxslist, inbranches_listlist)]
tree = u3.open(fname)[treename]
branches = [
ak.fill_none(
ak.pad_none(tree[branch_name].array(),
target=feature_length,
axis=-1,
clip=True if feature_length > 1 else False), 0.)
for feature_length, branch_list in zip(nMaxslist, inbranches_listlist)
for branch_name in branch_list
]
branchnames = [n for names in inbranches_listlist for n in names]
feature_lenghts = [
f for branches, f in zip(inbranches_listlist, nMaxslist)
for _ in branches
]
means = [
m[0] for branches, m in zip(inbranches_listlist, MeanNormTuple)
for _ in branches
]
norms = [
m[1] for branches, m in zip(inbranches_listlist, MeanNormTuple)
for _ in branches
]
print("Debugigng means and norms")
print(means)
print(norms)
print(branchnames)
branches_numpy = []
for br, brname, fl, mean, norm in zip(branches, branchnames,
feature_lenghts, means, norms):
print("DBG {}".format(brname))
print(br)
print("Length: {}".format(len(br)))
if brname == "TagVarCSV_trackJetDistVal":
print("BONUS DEBUG!")
print("Min: {}, Max: {}".format(ak.min(ak.count(br, axis=-1)),
ak.max(ak.count(br, axis=-1))))
if fl > 1:
# branches_numpy.append( (ak.to_numpy( br ) - mean) / norm)
branches_numpy.append((ak.to_numpy(br) - 0.) / 1.)
elif fl == 1:
# branches_numpy.append( (np.expand_dims( ak.to_numpy( br ), axis=-1) - mean)/norm )
branches_numpy.append(
(np.expand_dims(ak.to_numpy(br), axis=-1) - 0.) / 1.)
print("FINISHED THIS LOOP, YOU ARE PERFECT! :) ")
numpyarray = np.concatenate(branches_numpy, axis=-1)
print("\n" * 5)
print("Some metrics about this numpy array")
print(np.mean(numpyarray, axis=0))
print(np.std(numpyarray, axis=0))
print("Normalize array")
numpyarray = (numpyarray - np.mean(numpyarray, axis=0)) / np.std(
numpyarray, axis=0)
print("Some metrics about this numpy array")
print(np.mean(numpyarray, axis=0))
print(np.std(numpyarray, axis=0))
return numpyarray
def process(self, events):
output = self.accumulator.identity()
# we can use a very loose preselection to filter the events. nothing is done with this presel, though
presel = ak.num(events.Jet) > 0
ev = events[presel]
dataset = ev.metadata['dataset']
# load the config - probably not needed anymore
cfg = loadConfig()
#output['totalEvents']['all'] += len(events)
#output['skimmedEvents']['all'] += len(ev)
if self.year == 2018:
triggers = ev.HLT.Ele23_Ele12_CaloIdL_TrackIdL_IsoVL
elif self.year == 2017:
triggers = ev.HLT.Ele23_Ele12_CaloIdL_TrackIdL_IsoVL
elif self.year == 2016:
triggers = ev.HLT.Ele23_Ele12_CaloIdL_TrackIdL_IsoVL_DZ
if self.year == 2018:
lumimask = LumiMask(
'processors/Cert_314472-325175_13TeV_Legacy2018_Collisions18_JSON.txt'
)
## Electrons
electron = Collections(ev, "Electron", "tight").get()
electron = electron[(electron.pt > 25) & (np.abs(electron.eta) < 2.4)]
loose_electron = Collections(ev, "Electron", "veto").get()
loose_electron = loose_electron[(loose_electron.pt > 25)
& (np.abs(loose_electron.eta) < 2.4)]
SSelectron = (ak.sum(electron.charge, axis=1) != 0) & (ak.num(electron)
== 2)
OSelectron = (ak.sum(electron.charge, axis=1) == 0) & (ak.num(electron)
== 2)
dielectron = choose(electron, 2)
dielectron_mass = (dielectron['0'] + dielectron['1']).mass
dielectron_pt = (dielectron['0'] + dielectron['1']).pt
leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1))
leading_electron = electron[(leading_electron_idx)]
leading_electron = leading_electron[(leading_electron.pt > 30)]
trailing_electron_idx = ak.singletons(ak.argmin(electron.pt, axis=1))
trailing_electron = electron[trailing_electron_idx]
##Muons
loose_muon = Collections(ev, "Muon", "veto").get()
loose_muon = loose_muon[(loose_muon.pt > 20)
& (np.abs(loose_muon.eta) < 2.4)]
#jets
jet = getJets(ev, minPt=40, maxEta=2.4, pt_var='pt', UL=False)
jet = jet[ak.argsort(
jet.pt, ascending=False
)] # need to sort wrt smeared and recorrected jet pt
jet = jet[~match(jet, loose_muon,
deltaRCut=0.4)] # remove jets that overlap with muons
jet = jet[~match(
jet, electron,
deltaRCut=0.4)] # remove jets that overlap with electrons
## MET -> can switch to puppi MET
met_pt = ev.MET.pt
met_phi = ev.MET.phi
#selections
filters = getFilters(ev, year=self.year, dataset=dataset)
mask = lumimask(ev.run, ev.luminosityBlock)
ss = (SSelectron)
os = (OSelectron)
mass = (ak.min(np.abs(dielectron_mass - 91.2), axis=1) < 15)
lead_electron = (ak.min(leading_electron.pt, axis=1) > 30)
jet1 = (ak.num(jet) >= 1)
jet2 = (ak.num(jet) >= 2)
num_loose = ((ak.num(loose_electron) == 2) & (ak.num(loose_muon) == 0))
selection = PackedSelection()
selection.add('filter', (filters))
selection.add('mask', (mask))
selection.add('ss', ss)
selection.add('os', os)
selection.add('mass', mass)
selection.add('leading', lead_electron)
selection.add('triggers', triggers)
selection.add('one jet', jet1)
selection.add('two jets', jet2)
selection.add('num_loose', num_loose)
bl_reqs = ['filter'] + ['mass'] + ['mask'] + ['triggers'] + [
'leading'
] + ['num_loose']
#bl_reqs = ['filter'] + ['mass'] + ['triggers'] + ['leading'] + ['num_loose']
bl_reqs_d = {sel: True for sel in bl_reqs}
baseline = selection.require(**bl_reqs_d)
s_reqs = bl_reqs + ['ss']
s_reqs_d = {sel: True for sel in s_reqs}
ss_sel = selection.require(**s_reqs_d)
o_reqs = bl_reqs + ['os']
o_reqs_d = {sel: True for sel in o_reqs}
os_sel = selection.require(**o_reqs_d)
j1s_reqs = s_reqs + ['one jet']
j1s_reqs_d = {sel: True for sel in j1s_reqs}
j1ss_sel = selection.require(**j1s_reqs_d)
j1o_reqs = o_reqs + ['one jet']
j1o_reqs_d = {sel: True for sel in j1o_reqs}
j1os_sel = selection.require(**j1o_reqs_d)
j2s_reqs = s_reqs + ['two jets']
j2s_reqs_d = {sel: True for sel in j2s_reqs}
j2ss_sel = selection.require(**j2s_reqs_d)
j2o_reqs = o_reqs + ['two jets']
j2o_reqs_d = {sel: True for sel in j2o_reqs}
j2os_sel = selection.require(**j2o_reqs_d)
output["N_jet"].fill(
dataset=dataset,
multiplicity=ak.num(jet)[os_sel],
)
return output
def process(self, events):
output = self.accumulator.identity()
# we can use a very loose preselection to filter the events. nothing is done with this presel, though
presel = ak.num(events.Jet) > 0
if self.year == 2016:
lumimask = LumiMask(
'../data/lumi/Cert_271036-284044_13TeV_Legacy2016_Collisions16_JSON.txt'
)
if self.year == 2017:
lumimask = LumiMask(
'../data/lumi/Cert_294927-306462_13TeV_UL2017_Collisions17_GoldenJSON.txt'
)
if self.year == 2018:
lumimask = LumiMask(
'../data/lumi/Cert_314472-325175_13TeV_Legacy2018_Collisions18_JSON.txt'
)
ev = events[presel]
dataset = ev.metadata['dataset']
# load the config - probably not needed anymore
cfg = loadConfig()
output['totalEvents']['all'] += len(events)
output['skimmedEvents']['all'] += len(ev)
if self.year == 2018:
triggers = ev.HLT.Ele23_Ele12_CaloIdL_TrackIdL_IsoVL
elif self.year == 2017:
triggers = ev.HLT.Ele23_Ele12_CaloIdL_TrackIdL_IsoVL
elif self.year == 2016:
triggers = ev.HLT.Ele23_Ele12_CaloIdL_TrackIdL_IsoVL_DZ
## Electrons
electron = Collections(ev, "Electron", "tightFCNC", 0, self.year).get()
electron = electron[(electron.pt > 25) & (np.abs(electron.eta) < 2.4)]
loose_electron = Collections(ev, "Electron", "looseFCNC", 0,
self.year).get()
loose_electron = loose_electron[(loose_electron.pt > 25)
& (np.abs(loose_electron.eta) < 2.4)]
SSelectron = (ak.sum(electron.charge, axis=1) != 0) & (ak.num(electron)
== 2)
OSelectron = (ak.sum(electron.charge, axis=1) == 0) & (ak.num(electron)
== 2)
dielectron = choose(electron, 2)
dielectron_mass = (dielectron['0'] + dielectron['1']).mass
dielectron_pt = (dielectron['0'] + dielectron['1']).pt
leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1))
leading_electron = electron[(leading_electron_idx)]
leading_electron = leading_electron[(leading_electron.pt > 30)]
trailing_electron_idx = ak.singletons(ak.argmin(electron.pt, axis=1))
trailing_electron = electron[trailing_electron_idx]
##Muons
loose_muon = Collections(ev, "Muon", "looseFCNC", 0, self.year).get()
loose_muon = loose_muon[(loose_muon.pt > 20)
& (np.abs(loose_muon.eta) < 2.4)]
#jets
jet = getJets(ev, minPt=40, maxEta=2.4, pt_var='pt')
jet = jet[~match(jet, loose_muon,
deltaRCut=0.4)] # remove jets that overlap with muons
jet = jet[~match(
jet, electron,
deltaRCut=0.4)] # remove jets that overlap with electrons
## MET -> can switch to puppi MET
met_pt = ev.MET.pt
met_phi = ev.MET.phi
#weights
weight = Weights(len(ev))
weight2 = Weights(len(ev))
weight2.add("charge flip",
self.charge_flip_ratio.flip_weight(electron))
#selections
filters = getFilters(ev, year=self.year, dataset=dataset, UL=False)
mask = lumimask(ev.run, ev.luminosityBlock)
ss = (SSelectron)
os = (OSelectron)
mass = (ak.min(np.abs(dielectron_mass - 91.2), axis=1) < 15)
lead_electron = (ak.min(leading_electron.pt, axis=1) > 30)
jet1 = (ak.num(jet) >= 1)
jet2 = (ak.num(jet) >= 2)
num_loose = ((ak.num(loose_electron) == 2) & (ak.num(loose_muon) == 0))
selection = PackedSelection()
selection.add('filter', (filters))
selection.add('mask', (mask))
selection.add('ss', ss)
selection.add('os', os)
selection.add('mass', mass)
selection.add('leading', lead_electron)
selection.add('triggers', triggers)
selection.add('one jet', jet1)
selection.add('two jets', jet2)
selection.add('num_loose', num_loose)
bl_reqs = ['filter'] + ['triggers'] + ['mask']
bl_reqs_d = {sel: True for sel in bl_reqs}
baseline = selection.require(**bl_reqs_d)
s_reqs = bl_reqs + ['ss'] + ['mass'] + ['num_loose'] + ['leading']
s_reqs_d = {sel: True for sel in s_reqs}
ss_sel = selection.require(**s_reqs_d)
o_reqs = bl_reqs + ['os'] + ['mass'] + ['num_loose'] + ['leading']
o_reqs_d = {sel: True for sel in o_reqs}
os_sel = selection.require(**o_reqs_d)
j1s_reqs = s_reqs + ['one jet']
j1s_reqs_d = {sel: True for sel in j1s_reqs}
j1ss_sel = selection.require(**j1s_reqs_d)
j1o_reqs = o_reqs + ['one jet']
j1o_reqs_d = {sel: True for sel in j1o_reqs}
j1os_sel = selection.require(**j1o_reqs_d)
j2s_reqs = s_reqs + ['two jets']
j2s_reqs_d = {sel: True for sel in j2s_reqs}
j2ss_sel = selection.require(**j2s_reqs_d)
j2o_reqs = o_reqs + ['two jets']
j2o_reqs_d = {sel: True for sel in j2o_reqs}
j2os_sel = selection.require(**j2o_reqs_d)
#outputs
output["electron_data1"].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(leading_electron[os_sel].pt)),
eta=ak.to_numpy(ak.flatten(leading_electron[os_sel].eta)),
phi=ak.to_numpy(ak.flatten(leading_electron[os_sel].phi)),
weight=weight2.weight()[os_sel])
output["electron_data2"].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(trailing_electron[os_sel].pt)),
eta=ak.to_numpy(ak.flatten(trailing_electron[os_sel].eta)),
phi=ak.to_numpy(ak.flatten(trailing_electron[os_sel].phi)),
weight=weight2.weight()[os_sel])
output["electron_data3"].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(leading_electron[j1os_sel].pt)),
eta=ak.to_numpy(ak.flatten(leading_electron[j1os_sel].eta)),
phi=ak.to_numpy(ak.flatten(leading_electron[j1os_sel].phi)),
weight=weight2.weight()[j1os_sel])
output["electron_data4"].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(trailing_electron[j1os_sel].pt)),
eta=ak.to_numpy(ak.flatten(trailing_electron[j1os_sel].eta)),
phi=ak.to_numpy(ak.flatten(trailing_electron[j1os_sel].phi)),
weight=weight2.weight()[j1os_sel])
output["electron_data5"].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(leading_electron[j2os_sel].pt)),
eta=ak.to_numpy(ak.flatten(leading_electron[j2os_sel].eta)),
phi=ak.to_numpy(ak.flatten(leading_electron[j2os_sel].phi)),
weight=weight2.weight()[j2os_sel])
output["electron_data6"].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(trailing_electron[j2os_sel].pt)),
eta=ak.to_numpy(ak.flatten(trailing_electron[j2os_sel].eta)),
phi=ak.to_numpy(ak.flatten(trailing_electron[j2os_sel].phi)),
weight=weight2.weight()[j2os_sel])
output["electron_data7"].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(leading_electron[ss_sel].pt)),
eta=ak.to_numpy(ak.flatten(leading_electron[ss_sel].eta)),
phi=ak.to_numpy(ak.flatten(leading_electron[ss_sel].phi)),
weight=weight.weight()[ss_sel])
output["electron_data8"].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(trailing_electron[ss_sel].pt)),
eta=ak.to_numpy(ak.flatten(trailing_electron[ss_sel].eta)),
phi=ak.to_numpy(ak.flatten(trailing_electron[ss_sel].phi)),
weight=weight.weight()[ss_sel])
output["electron_data9"].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(leading_electron[j1ss_sel].pt)),
eta=ak.to_numpy(ak.flatten(leading_electron[j1ss_sel].eta)),
phi=ak.to_numpy(ak.flatten(leading_electron[j1ss_sel].phi)),
weight=weight.weight()[j1ss_sel])
output["electron_data10"].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(trailing_electron[j1ss_sel].pt)),
eta=ak.to_numpy(ak.flatten(trailing_electron[j1ss_sel].eta)),
phi=ak.to_numpy(ak.flatten(trailing_electron[j1ss_sel].phi)),
weight=weight.weight()[j1ss_sel])
output["electron_data11"].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(leading_electron[j2ss_sel].pt)),
eta=ak.to_numpy(ak.flatten(leading_electron[j2ss_sel].eta)),
phi=ak.to_numpy(ak.flatten(leading_electron[j2ss_sel].phi)),
weight=weight.weight()[j2ss_sel])
output["electron_data12"].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(trailing_electron[j2ss_sel].pt)),
eta=ak.to_numpy(ak.flatten(trailing_electron[j2ss_sel].eta)),
phi=ak.to_numpy(ak.flatten(trailing_electron[j2ss_sel].phi)),
weight=weight.weight()[j2ss_sel])
output["dilep_mass1"].fill(
dataset=dataset,
mass=ak.to_numpy(ak.flatten(dielectron_mass[os_sel])),
pt=ak.to_numpy(ak.flatten(dielectron_pt[os_sel])),
weight=weight2.weight()[os_sel])
output["dilep_mass2"].fill(
dataset=dataset,
mass=ak.to_numpy(ak.flatten(dielectron_mass[j1os_sel])),
pt=ak.to_numpy(ak.flatten(dielectron_pt[j1os_sel])),
weight=weight2.weight()[j1os_sel])
output["dilep_mass3"].fill(
dataset=dataset,
mass=ak.to_numpy(ak.flatten(dielectron_mass[j2os_sel])),
pt=ak.to_numpy(ak.flatten(dielectron_pt[j2os_sel])),
weight=weight2.weight()[j2os_sel])
output["dilep_mass4"].fill(
dataset=dataset,
mass=ak.to_numpy(ak.flatten(dielectron_mass[ss_sel])),
pt=ak.to_numpy(ak.flatten(dielectron_pt[ss_sel])),
weight=weight.weight()[ss_sel])
output["dilep_mass5"].fill(
dataset=dataset,
mass=ak.to_numpy(ak.flatten(dielectron_mass[j1ss_sel])),
pt=ak.to_numpy(ak.flatten(dielectron_pt[j1ss_sel])),
weight=weight.weight()[j1ss_sel])
output["dilep_mass6"].fill(
dataset=dataset,
mass=ak.to_numpy(ak.flatten(dielectron_mass[j2ss_sel])),
pt=ak.to_numpy(ak.flatten(dielectron_pt[j2ss_sel])),
weight=weight.weight()[j2ss_sel])
output["MET"].fill(dataset=dataset,
pt=met_pt[os_sel],
weight=weight2.weight()[os_sel])
output["MET2"].fill(dataset=dataset,
pt=met_pt[j1os_sel],
weight=weight2.weight()[j1os_sel])
output["MET3"].fill(dataset=dataset,
pt=met_pt[j2os_sel],
weight=weight2.weight()[j2os_sel])
output["MET4"].fill(dataset=dataset,
pt=met_pt[ss_sel],
weight=weight.weight()[ss_sel])
output["MET5"].fill(dataset=dataset,
pt=met_pt[j1ss_sel],
weight=weight.weight()[j1ss_sel])
output["MET6"].fill(dataset=dataset,
pt=met_pt[j2ss_sel],
weight=weight.weight()[j2ss_sel])
output["N_jet"].fill(dataset=dataset,
multiplicity=ak.num(jet)[os_sel],
weight=weight2.weight()[os_sel])
output["N_jet2"].fill(dataset=dataset,
multiplicity=ak.num(jet)[j1os_sel],
weight=weight2.weight()[j1os_sel])
output["N_jet3"].fill(dataset=dataset,
multiplicity=ak.num(jet)[j2os_sel],
weight=weight2.weight()[j2os_sel])
output["N_jet4"].fill(dataset=dataset,
multiplicity=ak.num(jet)[ss_sel],
weight=weight.weight()[ss_sel])
output["N_jet5"].fill(dataset=dataset,
multiplicity=ak.num(jet)[j1ss_sel],
weight=weight.weight()[j1ss_sel])
output["N_jet6"].fill(dataset=dataset,
multiplicity=ak.num(jet)[j2ss_sel],
weight=weight.weight()[j2ss_sel])
output["PV_npvsGood"].fill(dataset=dataset,
multiplicity=ev.PV[os_sel].npvsGood,
weight=weight2.weight()[os_sel])
output["PV_npvsGood2"].fill(dataset=dataset,
multiplicity=ev.PV[j1os_sel].npvsGood,
weight=weight2.weight()[j1os_sel])
output["PV_npvsGood3"].fill(dataset=dataset,
multiplicity=ev.PV[j2os_sel].npvsGood,
weight=weight2.weight()[j2os_sel])
output["PV_npvsGood4"].fill(dataset=dataset,
multiplicity=ev.PV[ss_sel].npvsGood,
weight=weight.weight()[ss_sel])
output["PV_npvsGood5"].fill(dataset=dataset,
multiplicity=ev.PV[j1ss_sel].npvsGood,
weight=weight.weight()[j1ss_sel])
output["PV_npvsGood6"].fill(dataset=dataset,
multiplicity=ev.PV[j2ss_sel].npvsGood,
weight=weight.weight()[j2ss_sel])
return output
