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