def test_IndexedOptionArray_combinations(): content = awkward1.from_iter([[[0, 1, 2], [], [3, 4]], [], [[5]], [[6, 7, 8, 9]], [[], [10, 11, 12]]], highlevel=False) index = awkward1.layout.Index64(numpy.array([0, 1, -1, -1, 4], dtype=numpy.int64)) array = awkward1.Array(awkward1.layout.IndexedOptionArray64(index, content)) assert awkward1.to_list(array) == [[[0, 1, 2], [], [3, 4]], [], None, None, [[], [10, 11, 12]]] assert awkward1.to_list(awkward1.combinations(array, 2, axis=0)) == [([[0, 1, 2], [], [3, 4]], []), ([[0, 1, 2], [], [3, 4]], None), ([[0, 1, 2], [], [3, 4]], None), ([[0, 1, 2], [], [3, 4]], [[], [10, 11, 12]]), ([], None), ([], None), ([], [[], [10, 11, 12]]), (None, None), (None, [[], [10, 11, 12]]), (None, [[], [10, 11, 12]])] assert awkward1.to_list(awkward1.combinations(array, 2, axis=1)) == [[([0, 1, 2], []), ([0, 1, 2], [3, 4]), ([], [3, 4])], [], None, None, [([], [10, 11, 12])]] assert awkward1.to_list(awkward1.combinations(array, 2, axis=2)) == [[[(0, 1), (0, 2), (1, 2)], [], [(3, 4)]], [], None, None, [[], [(10, 11), (10, 12), (11, 12)]]]
def test_ByteMaskedArray_combinations(): content = awkward1.from_iter([[[0, 1, 2], [], [3, 4]], [], [[5]], [[6, 7, 8, 9]], [[], [10, 11, 12]]], highlevel=False) mask = awkward1.layout.Index8(numpy.array([0, 0, 1, 1, 0], dtype=numpy.int8)) array = awkward1.Array(awkward1.layout.ByteMaskedArray(mask, content, valid_when=False)) assert awkward1.to_list(array) == [[[0, 1, 2], [], [3, 4]], [], None, None, [[], [10, 11, 12]]] assert awkward1.to_list(awkward1.combinations(array, 2, axis=0)) == [([[0, 1, 2], [], [3, 4]], []), ([[0, 1, 2], [], [3, 4]], None), ([[0, 1, 2], [], [3, 4]], None), ([[0, 1, 2], [], [3, 4]], [[], [10, 11, 12]]), ([], None), ([], None), ([], [[], [10, 11, 12]]), (None, None), (None, [[], [10, 11, 12]]), (None, [[], [10, 11, 12]])] assert awkward1.to_list(awkward1.combinations(array, 2, axis=1)) == [[([0, 1, 2], []), ([0, 1, 2], [3, 4]), ([], [3, 4])], [], None, None, [([], [10, 11, 12])]] assert awkward1.to_list(awkward1.combinations(array, 2, axis=2)) == [[[(0, 1), (0, 2), (1, 2)], [], [(3, 4)]], [], None, None, [[], [(10, 11), (10, 12), (11, 12)]]]
def test_axis0(): array = awkward1.Array([0.0, 1.1, 2.2, 3.3]) assert awkward1.to_list(awkward1.combinations(array, 2, replacement=False, axis=0)) == [(0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 2.2), (1.1, 3.3), (2.2, 3.3)] assert awkward1.to_list(awkward1.combinations(array, 2, replacement=False, axis=0, fields=["x", "y"])) == [{"x": 0.0, "y": 1.1}, {"x": 0.0, "y": 2.2}, {"x": 0.0, "y": 3.3}, {"x": 1.1, "y": 2.2}, {"x": 1.1, "y": 3.3}, {"x": 2.2, "y": 3.3}] assert awkward1.combinations(array, 2, replacement=False, axis=0, parameters={"some": "param"}).layout.parameters["some"] == "param" assert awkward1.to_list(awkward1.combinations(array, 3, replacement=False, axis=0)) == [(0.0, 1.1, 2.2), (0.0, 1.1, 3.3), (0.0, 2.2, 3.3), (1.1, 2.2, 3.3)]
def process(self, events): # Initialize accumulator out = self.accumulator.identity() # Event selection: opposite charged same flavor Electron = events.Electron Electron_mask = (Electron.pt > 20) & (np.abs(Electron.eta) < 2.5) & (Electron.cutBased > 1) Ele_channel_mask = ak.num(Electron[Electron_mask]) > 1 Ele_channel_events = events[Ele_channel_mask] Ele = Ele_channel_events.Electron # All possible pairs of Electron in each event ele_pairs = ak.combinations(Ele, 2, axis=1) # TLorentz vector sum of ele_pairs ele_left, ele_right = ak.unzip(ele_pairs) diele = ele_left + ele_right diffsign_diele = diele[diele.charge == 0] leading_diffsign_diele = diffsign_diele[ak.argmax(diffsign_diele.pt, axis=1, keepdims=True)] #Mee = ak.flatten(leading_diffsign_diele.mass) # This makes type error ( primitive expected but ?float given ) Mee = ak.to_numpy(leading_diffsign_diele.mass) Mee = Mee.flatten() out.fill(dataset=events.metadata["dataset"], mass=Mee) return out
def choose(first, n=2): tmp = ak.combinations(first, n) combs = tmp['0'] for i in range(1, n): combs = combs.__add__(tmp[str(i)]) for i in range(n): combs[str(i)] = tmp[str(i)] return combs
def choose3(first, n=3): from warnings import warn warn( "Deprecation Warning: The choose3 function will be removed. Use choose(first, n=3) instead." ) tmp = ak.combinations(first, n) combs = (tmp['0'] + tmp['1'] + tmp['2']) combs['0'] = tmp['0'] combs['1'] = tmp['1'] combs['2'] = tmp['2'] return combs
def process(self, events): output = self.accumulator.identity() dataset = events.metadata["dataset"] dimuon = ak.combinations(events.Muon, 2) dimuon = dimuon["0"] + dimuon["1"] output["pt"].fill(dataset=dataset, pt=ak.flatten(events.Muon.pt)) output["mass"].fill(dataset=dataset, mass=ak.flatten(dimuon.mass)) output["cutflow"]["%s_pt" % dataset] += sum(ak.num(events.Muon)) output["cutflow"]["%s_mass" % dataset] += sum(ak.num(dimuon)) return output
def test_ListOffsetArray(): array = awkward1.Array([[0.0, 1.1, 2.2, 3.3], [], [4.4, 5.5, 6.6], [7.7], [8.8, 9.9, 10.0, 11.1, 12.2]]) assert awkward1.to_list(awkward1.combinations(array, 2, replacement=False)) == [[(0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 2.2), (1.1, 3.3), (2.2, 3.3)], [], [(4.4, 5.5), (4.4, 6.6), (5.5, 6.6)], [], [(8.8, 9.9), (8.8, 10.0), (8.8, 11.1), (8.8, 12.2), (9.9, 10.0), (9.9, 11.1), (9.9, 12.2), (10.0, 11.1), (10.0, 12.2), (11.1, 12.2)]] assert awkward1.to_list(awkward1.combinations(array, 2, replacement=False, fields=["x", "y"])) == [[{"x": 0.0, "y": 1.1}, {"x": 0.0, "y": 2.2}, {"x": 0.0, "y": 3.3}, {"x": 1.1, "y": 2.2}, {"x": 1.1, "y": 3.3}, {"x": 2.2, "y": 3.3}], [], [{"x": 4.4, "y": 5.5}, {"x": 4.4, "y": 6.6}, {"x": 5.5, "y": 6.6}], [], [{"x": 8.8, "y": 9.9}, {"x": 8.8, "y": 10.0}, {"x": 8.8, "y": 11.1}, {"x": 8.8, "y": 12.2}, {"x": 9.9, "y": 10.0}, {"x": 9.9, "y": 11.1}, {"x": 9.9, "y": 12.2}, {"x": 10.0, "y": 11.1}, {"x": 10.0, "y": 12.2}, {"x": 11.1, "y": 12.2}]] tmp = awkward1.combinations(array, 2, replacement=False, parameters={"some": "param"}).layout if isinstance(tmp, awkward1.partition.PartitionedArray): assert awkward1.partition.first(tmp).content.parameters["some"] == "param" else: assert tmp.content.parameters["some"] == "param" assert awkward1.to_list(awkward1.combinations(array, 2, replacement=True)) == [[(0.0, 0.0), (0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 1.1), (1.1, 2.2), (1.1, 3.3), (2.2, 2.2), (2.2, 3.3), (3.3, 3.3)], [], [(4.4, 4.4), (4.4, 5.5), (4.4, 6.6), (5.5, 5.5), (5.5, 6.6), (6.6, 6.6)], [(7.7, 7.7)], [(8.8, 8.8), (8.8, 9.9), (8.8, 10.0), (8.8, 11.1), (8.8, 12.2), (9.9, 9.9), (9.9, 10.0), (9.9, 11.1), (9.9, 12.2), (10.0, 10.0), (10.0, 11.1), (10.0, 12.2), (11.1, 11.1), (11.1, 12.2), (12.2, 12.2)]] assert awkward1.to_list(awkward1.combinations(array, 3, replacement=False)) == [[(0.0, 1.1, 2.2), (0.0, 1.1, 3.3), (0.0, 2.2, 3.3), (1.1, 2.2, 3.3)], [], [(4.4, 5.5, 6.6)], [], [(8.8, 9.9, 10.0), (8.8, 9.9, 11.1), (8.8, 9.9, 12.2), (8.8, 10.0, 11.1), (8.8, 10.0, 12.2), (8.8, 11.1, 12.2), (9.9, 10.0, 11.1), (9.9, 10.0, 12.2), (9.9, 11.1, 12.2), (10.0, 11.1, 12.2)]] assert awkward1.to_list(awkward1.combinations(array, 3, replacement=True)) == [[(0.0, 0.0, 0.0), (0.0, 0.0, 1.1), (0.0, 0.0, 2.2), (0.0, 0.0, 3.3), (0.0, 1.1, 1.1), (0.0, 1.1, 2.2), (0.0, 1.1, 3.3), (0.0, 2.2, 2.2), (0.0, 2.2, 3.3), (0.0, 3.3, 3.3), (1.1, 1.1, 1.1), (1.1, 1.1, 2.2), (1.1, 1.1, 3.3), (1.1, 2.2, 2.2), (1.1, 2.2, 3.3), (1.1, 3.3, 3.3), (2.2, 2.2, 2.2), (2.2, 2.2, 3.3), (2.2, 3.3, 3.3), (3.3, 3.3, 3.3)], [], [(4.4, 4.4, 4.4), (4.4, 4.4, 5.5), (4.4, 4.4, 6.6), (4.4, 5.5, 5.5), (4.4, 5.5, 6.6), (4.4, 6.6, 6.6), (5.5, 5.5, 5.5), (5.5, 5.5, 6.6), (5.5, 6.6, 6.6), (6.6, 6.6, 6.6)], [(7.7, 7.7, 7.7)], [(8.8, 8.8, 8.8), (8.8, 8.8, 9.9), (8.8, 8.8, 10.0), (8.8, 8.8, 11.1), (8.8, 8.8, 12.2), (8.8, 9.9, 9.9), (8.8, 9.9, 10.0), (8.8, 9.9, 11.1), (8.8, 9.9, 12.2), (8.8, 10.0, 10.0), (8.8, 10.0, 11.1), (8.8, 10.0, 12.2), (8.8, 11.1, 11.1), (8.8, 11.1, 12.2), (8.8, 12.2, 12.2), (9.9, 9.9, 9.9), (9.9, 9.9, 10.0), (9.9, 9.9, 11.1), (9.9, 9.9, 12.2), (9.9, 10.0, 10.0), (9.9, 10.0, 11.1), (9.9, 10.0, 12.2), (9.9, 11.1, 11.1), (9.9, 11.1, 12.2), (9.9, 12.2, 12.2), (10.0, 10.0, 10.0), (10.0, 10.0, 11.1), (10.0, 10.0, 12.2), (10.0, 11.1, 11.1), (10.0, 11.1, 12.2), (10.0, 12.2, 12.2), (11.1, 11.1, 11.1), (11.1, 11.1, 12.2), (11.1, 12.2, 12.2), (12.2, 12.2, 12.2)]]
def test_RegularArray(): array = awkward1.Array(numpy.array([[0.0, 1.1, 2.2, 3.3], [4.4, 5.5, 6.6, 7.7]])) assert awkward1.to_list(awkward1.combinations(array, 2, replacement=False)) == [[(0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 2.2), (1.1, 3.3), (2.2, 3.3)], [(4.4, 5.5), (4.4, 6.6), (4.4, 7.7), (5.5, 6.6), (5.5, 7.7), (6.6, 7.7)]] assert awkward1.to_list(awkward1.combinations(array, 2, replacement=False, fields=["x", "y"])) == [[{"x": 0.0, "y": 1.1}, {"x": 0.0, "y": 2.2}, {"x": 0.0, "y": 3.3}, {"x": 1.1, "y": 2.2}, {"x": 1.1, "y": 3.3}, {"x": 2.2, "y": 3.3}], [{"x": 4.4, "y": 5.5}, {"x": 4.4, "y": 6.6}, {"x": 4.4, "y": 7.7}, {"x": 5.5, "y": 6.6}, {"x": 5.5, "y": 7.7}, {"x": 6.6, "y": 7.7}]] tmp = awkward1.combinations(array, 2, replacement=False, parameters={"some": "param"}).layout if isinstance(tmp, awkward1.partition.PartitionedArray): assert awkward1.partition.first(tmp).content.parameters["some"] == "param" else: assert tmp.content.parameters["some"] == "param" assert awkward1.to_list(awkward1.combinations(array, 2, replacement=True)) == [[(0.0, 0.0), (0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 1.1), (1.1, 2.2), (1.1, 3.3), (2.2, 2.2), (2.2, 3.3), (3.3, 3.3)], [(4.4, 4.4), (4.4, 5.5), (4.4, 6.6), (4.4, 7.7), (5.5, 5.5), (5.5, 6.6), (5.5, 7.7), (6.6, 6.6), (6.6, 7.7), (7.7, 7.7)]] assert awkward1.to_list(awkward1.combinations(array, 3, replacement=False)) == [[(0.0, 1.1, 2.2), (0.0, 1.1, 3.3), (0.0, 2.2, 3.3), (1.1, 2.2, 3.3)], [(4.4, 5.5, 6.6), (4.4, 5.5, 7.7), (4.4, 6.6, 7.7), (5.5, 6.6, 7.7)]] assert awkward1.to_list(awkward1.combinations(array, 3, replacement=True)) == [[(0.0, 0.0, 0.0), (0.0, 0.0, 1.1), (0.0, 0.0, 2.2), (0.0, 0.0, 3.3), (0.0, 1.1, 1.1), (0.0, 1.1, 2.2), (0.0, 1.1, 3.3), (0.0, 2.2, 2.2), (0.0, 2.2, 3.3), (0.0, 3.3, 3.3), (1.1, 1.1, 1.1), (1.1, 1.1, 2.2), (1.1, 1.1, 3.3), (1.1, 2.2, 2.2), (1.1, 2.2, 3.3), (1.1, 3.3, 3.3), (2.2, 2.2, 2.2), (2.2, 2.2, 3.3), (2.2, 3.3, 3.3), (3.3, 3.3, 3.3)], [(4.4, 4.4, 4.4), (4.4, 4.4, 5.5), (4.4, 4.4, 6.6), (4.4, 4.4, 7.7), (4.4, 5.5, 5.5), (4.4, 5.5, 6.6), (4.4, 5.5, 7.7), (4.4, 6.6, 6.6), (4.4, 6.6, 7.7), (4.4, 7.7, 7.7), (5.5, 5.5, 5.5), (5.5, 5.5, 6.6), (5.5, 5.5, 7.7), (5.5, 6.6, 6.6), (5.5, 6.6, 7.7), (5.5, 7.7, 7.7), (6.6, 6.6, 6.6), (6.6, 6.6, 7.7), (6.6, 7.7, 7.7), (7.7, 7.7, 7.7)]]
final = final[(ak.count(qqb.pdgId, axis=1) == 3)] finaljets = final.Jet qqb = qqb[(ak.count(qqb.pdgId, axis=1) == 3)] #Implementing Tight Jet Cuts on Training Data finaljetSel = (abs(finaljets.eta) < 2.4) & (finaljets.pt > 30) finalJets = finaljets[finaljetSel] #Use nearest to match Jets matchedGenJets = qqb.nearest(final.GenJet) matchedJets = matchedGenJets.nearest(finalJets) print("matched genpart to genjet and finally to reco jets") #Assigning True false to sets of 3 jets test = matchedJets.genJetIdx combs = ak.combinations(finalJets, 3, replacement=False) t1 = (combs['0'].genJetIdx == test[:, 0]) | ( combs['0'].genJetIdx == test[:, 1]) | (combs['0'].genJetIdx == test[:, 2]) t2 = (combs['1'].genJetIdx == test[:, 0]) | ( combs['1'].genJetIdx == test[:, 1]) | (combs['1'].genJetIdx == test[:, 2]) t3 = (combs['2'].genJetIdx == test[:, 0]) | ( combs['2'].genJetIdx == test[:, 1]) | (combs['2'].genJetIdx == test[:, 2]) t = t1 & t2 & t3 trutharray = ak.flatten(t) print("matching a validation array for every combo of 3 jets") #Zipping into CSV for training jetcombos = ak.flatten(combs) j1, j2, j3 = ak.unzip(jetcombos) dR1_2 = j1.delta_r(j2)
def test_axis2(): array = awkward1.Array([[[0.0, 1.1, 2.2, 3.3], [], [4.4, 5.5, 6.6]], [], [[7.7], [8.8, 9.9, 10.0, 11.1, 12.2]]]) assert awkward1.to_list(awkward1.combinations(array, 2, axis=1, replacement=False)) == [[([0.0, 1.1, 2.2, 3.3], []), ([0.0, 1.1, 2.2, 3.3], [4.4, 5.5, 6.6]), ([], [4.4, 5.5, 6.6])], [], [([7.7], [8.8, 9.9, 10.0, 11.1, 12.2])]] assert awkward1.to_list(awkward1.combinations(array, 2, axis=2, replacement=False)) == [[[(0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 2.2), (1.1, 3.3), (2.2, 3.3)], [], [(4.4, 5.5), (4.4, 6.6), (5.5, 6.6)]], [], [[], [(8.8, 9.9), (8.8, 10.0), (8.8, 11.1), (8.8, 12.2), (9.9, 10.0), (9.9, 11.1), (9.9, 12.2), (10.0, 11.1), (10.0, 12.2), (11.1, 12.2)]]]
def test_IndexedArray(): array = awkward1.Array([[0.0, 1.1, 2.2, 3.3], [], [4.4, 5.5, 6.6], None, [7.7], None, [8.8, 9.9, 10.0, 11.1, 12.2]]) assert awkward1.to_list(awkward1.combinations(array, 2, replacement=False)) == [[(0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 2.2), (1.1, 3.3), (2.2, 3.3)], [], [(4.4, 5.5), (4.4, 6.6), (5.5, 6.6)], None, [], None, [(8.8, 9.9), (8.8, 10.0), (8.8, 11.1), (8.8, 12.2), (9.9, 10.0), (9.9, 11.1), (9.9, 12.2), (10.0, 11.1), (10.0, 12.2), (11.1, 12.2)]]
def choose(first, n=2): tmp = ak.combinations(first, n) combs = (tmp['0'] + tmp['1']) combs['0'] = tmp['0'] combs['1'] = tmp['1'] return combs
def process(self, events): # Initialize accumulator out = self.accumulator.identity() dataset = setname #events.metadata['dataset'] isData = 'genWeight' not in events.fields selection = processor.PackedSelection() # Cut flow cut0 = np.zeros(len(events)) # --- Selection # << flat dim helper function >> def flat_dim(arr): sub_arr = ak.flatten(arr) mask = ~ak.is_none(sub_arr) return ak.to_numpy(sub_arr[mask]) # << drop na helper function >> def drop_na(arr): mask = ~ak.is_none(arr) return arr[mask] # << drop na helper function >> def drop_na_np(arr): mask = ~np.isnan(arr) return arr[mask] # double lepton trigger is_double_ele_trigger=True if not is_double_ele_trigger: double_ele_triggers_arr=np.ones(len(events), dtype=np.bool) else: double_ele_triggers_arr = np.zeros(len(events), dtype=np.bool) for path in self._doubleelectron_triggers[self._year]: if path not in events.HLT.fields: continue double_ele_triggers_arr = double_ele_triggers_arr | events.HLT[path] # single lepton trigger is_single_ele_trigger=True if not is_single_ele_trigger: single_ele_triggers_arr=np.ones(len(events), dtype=np.bool) else: single_ele_triggers_arr = np.zeros(len(events), dtype=np.bool) for path in self._singleelectron_triggers[self._year]: if path not in events.HLT.fields: continue single_ele_triggers_arr = single_ele_triggers_arr | events.HLT[path] Initial_events = events print("#### Initial events: ",Initial_events) #events = events[single_ele_triggers_arr | double_ele_triggers_arr] events = events[double_ele_triggers_arr] ##----------- Cut flow1: Passing Triggers cut1 = np.ones(len(events)) print("#### cut1: ",len(cut1)) # Particle Identification Electron = events.Electron def Electron_selection(ele): return(ele.pt > 25) & (np.abs(ele.eta) < 2.5) & (ele.cutBased > 2) # Electron channel Electron_mask = Electron_selection(Electron) Ele_channel_mask = ak.num(Electron[Electron_mask]) > 1 Ele_channel_events = events[Ele_channel_mask] ##----------- Cut flow2: Electron channel cut2 = np.ones(len(Ele_channel_events)) * 2 print("#### cut2: ",len(cut2)) # --- Calculate Scale factor weight if not isData: # PU weight with lookup table <-- On developing --> #get_pu_weight = self._corrections['get_pu_weight'][self._year] #pu = get_pu_weight(events.Pileup.nTrueInt) get_ele_reco_sf = self._corrections['get_ele_reco_sf'][self._year] get_ele_loose_id_sf = self._corrections['get_ele_loose_id_sf'][self._year] get_ele_trig_leg1_SF = self._corrections['get_ele_trig_leg1_SF'][self._year] get_ele_trig_leg1_data_Eff = self._corrections['get_ele_trig_leg1_data_Eff'][self._year] get_ele_trig_leg1_mc_Eff = self._corrections['get_ele_trig_leg1_mc_Eff'][self._year] get_ele_trig_leg2_SF = self._corrections['get_ele_trig_leg2_SF'][self._year] get_ele_trig_leg2_data_Eff = self._corrections['get_ele_trig_leg2_data_Eff'][self._year] get_ele_trig_leg2_mc_Eff = self._corrections['get_ele_trig_leg2_mc_Eff'][self._year] # PU weight with custom made npy and multi-indexing pu_weight_idx = ak.values_astype(Ele_channel_events.Pileup.nTrueInt,"int64") pu = self._puweight_arr[pu_weight_idx] nPV = Ele_channel_events.PV.npvsGood else: nPV = Ele_channel_events.PV.npvsGood # Electron array Ele = Ele_channel_events.Electron Electron_mask = Electron_selection(Ele) Ele_sel = Ele[Electron_mask] # Electron pair ele_pairs = ak.combinations(Ele_sel,2,axis=1) ele_left, ele_right = ak.unzip(ele_pairs) diele = ele_left + ele_right # OS os_mask = diele.charge == 0 os_diele = diele[os_mask] os_ele_left = ele_left[os_mask] os_ele_right = ele_right[os_mask] os_event_mask = ak.num(os_diele) > 0 Ele_os_channel_events = Ele_channel_events[os_event_mask] #selection.add('ossf',os_event_mask) # Helper function: High PT argmax def make_leading_pair(target,base): return target[ak.argmax(base.pt,axis=1,keepdims=True)] # -- Only Leading pair -- leading_diele = make_leading_pair(diele,diele) leading_ele = make_leading_pair(ele_left,diele) subleading_ele= make_leading_pair(ele_right,diele) # -- Scale Factor for each electron def Trigger_Weight(eta1,pt1,eta2,pt2): per_ev_MC =\ get_ele_trig_leg1_mc_Eff(eta1,pt1) * get_ele_trig_leg2_mc_Eff(eta2,pt2) +\ get_ele_trig_leg1_mc_Eff(eta2,pt2) * get_ele_trig_leg2_mc_Eff(eta1,pt1) -\ get_ele_trig_leg1_mc_Eff(eta1,pt1) * get_ele_trig_leg1_mc_Eff(eta2,pt2) per_ev_data =\ get_ele_trig_leg1_data_Eff(eta1,pt1) * get_ele_trig_leg1_SF(eta1,pt1) * get_ele_trig_leg2_data_Eff(eta2,pt2) * get_ele_trig_leg2_SF(eta2,pt2) +\ get_ele_trig_leg1_data_Eff(eta2,pt2) * get_ele_trig_leg1_SF(eta2,pt2) * get_ele_trig_leg2_data_Eff(eta1,pt1) * get_ele_trig_leg2_SF(eta1,pt1) -\ get_ele_trig_leg1_data_Eff(eta1,pt1) * get_ele_trig_leg1_SF(eta1,pt1) * get_ele_trig_leg1_data_Eff(eta2,pt2) * get_ele_trig_leg1_SF(eta2,pt2) return per_ev_data/per_ev_MC if not isData: ele_loose_id_sf = get_ele_loose_id_sf(ak.flatten(leading_ele.deltaEtaSC + leading_ele.eta),ak.flatten(leading_ele.pt))* get_ele_loose_id_sf(ak.flatten(subleading_ele.deltaEtaSC + subleading_ele.eta),ak.flatten(subleading_ele.pt)) #print("Ele ID SC---->",ele_loose_id_sf) ele_reco_sf = get_ele_reco_sf(ak.flatten(leading_ele.deltaEtaSC + leading_ele.eta),ak.flatten(leading_ele.pt))* get_ele_reco_sf(ak.flatten(subleading_ele.deltaEtaSC + subleading_ele.eta),ak.flatten(subleading_ele.pt)) #print("Ele RECO SC---->",ele_reco_sf) eta1 = ak.flatten(leading_ele.deltaEtaSC + leading_ele.eta) eta2 = ak.flatten(subleading_ele.deltaEtaSC + subleading_ele.eta) pt1 = ak.flatten(leading_ele.pt) pt2 = ak.flatten(subleading_ele.pt) ele_trig_weight = Trigger_Weight(eta1,pt1,eta2,pt2) print("#### Test print trigger weight ####") print(ele_trig_weight) # --OS and Leading pair -- leading_os_diele = make_leading_pair(os_diele,os_diele) leading_os_ele = make_leading_pair(os_ele_left,os_diele) subleading_os_ele= make_leading_pair(os_ele_right,os_diele) ##----------- Cut flow3: OSSF cut3 = np.ones(len(flat_dim(leading_os_diele))) * 3 print("#### cut3: ",len(cut3)) # Helper function: Zmass window def makeZmass_window_mask(dielecs,start=60,end=120): mask = (dielecs.mass >= start) & (dielecs.mass <= end) return mask # -- OS and Leading pair -- Zmass_mask_os = makeZmass_window_mask(leading_os_diele) leading_os_Zwindow_ele = leading_os_ele[Zmass_mask_os] subleading_os_Zwindow_ele = subleading_os_ele[Zmass_mask_os] leading_os_Zwindow_diele = leading_os_diele[Zmass_mask_os] # for masking Zmass_event_mask = makeZmass_window_mask(leading_diele) Zmass_os_event_mask= ak.flatten(os_event_mask * Zmass_event_mask) Ele_Zmass_os_events = Ele_channel_events[Zmass_os_event_mask] ##----------- Cut flow4: Zmass cut4 = np.ones(len(flat_dim(leading_os_Zwindow_diele))) * 4 print("#### cut4: ",len(cut4)) ## << Selection method -- Need validation >> #print("a--->",len(Ele_channel_events)) #print("b--->",len(Ele_os_channel_events)) #print("b2--->",len(cut3)) #print("c--->",len(Ele_Zmass_os_events)) #print("c2--->",len(cut4)) ele1PT = flat_dim(leading_os_Zwindow_ele.pt) ele1Eta = flat_dim(leading_os_Zwindow_ele.eta) ele1Phi = flat_dim(leading_os_Zwindow_ele.phi) ele2PT = flat_dim(subleading_os_Zwindow_ele.pt) ele2Eta = flat_dim(subleading_os_Zwindow_ele.eta) ele2Phi = flat_dim(subleading_os_Zwindow_ele.phi) Mee = flat_dim(leading_os_Zwindow_diele.mass) charge = flat_dim(leading_os_Zwindow_diele.charge) # --- Apply weight and hist weights = processor.Weights(len(cut2)) # --- skim cut-weight def skim_weight(arr): mask1 = ~ak.is_none(arr) subarr = arr[mask1] mask2 = subarr !=0 return ak.to_numpy(subarr[mask2]) cuts = ak.flatten(Zmass_mask_os) if not isData: weights.add('pileup',pu) weights.add('ele_id',ele_loose_id_sf) weights.add('ele_reco',ele_reco_sf) #weights.add('ele_trigger',ele_trig_weight) # Initial events out["sumw"][dataset] += len(Initial_events) # Cut flow loop for cut in [cut0,cut1,cut2,cut3,cut4]: out["cutflow"].fill( dataset = dataset, cutflow=cut ) # Primary vertex out['nPV'].fill( dataset=dataset, nPV = nPV, weight = weights.weight() ) out['nPV_nw'].fill( dataset=dataset, nPV_nw = nPV ) # Physics varibles passing Zwindow out["mass"].fill( dataset=dataset, mass=Mee, weight = skim_weight(weights.weight() * cuts) ) out["ele1pt"].fill( dataset=dataset, ele1pt=ele1PT, weight = skim_weight(weights.weight() * cuts) ) out["ele1eta"].fill( dataset=dataset, ele1eta=ele1Eta, weight = skim_weight(weights.weight() * cuts) ) out["ele1phi"].fill( dataset=dataset, ele1phi=ele1Phi, weight = skim_weight(weights.weight() * cuts) ) out["ele2pt"].fill( dataset=dataset, ele2pt=ele2PT, weight = skim_weight(weights.weight() * cuts) ) out["ele2eta"].fill( dataset=dataset, ele2eta=ele2Eta, weight = skim_weight(weights.weight() * cuts) ) out["ele2phi"].fill( dataset=dataset, ele2phi=ele2Phi, weight = skim_weight(weights.weight() * cuts) ) return out
#Open ROOT file from Pythia Delphes production with FCCSW file = uproot.open(f"{path}/output/rootfiles/FCCDelphesOutput_100events.root") tree = file['events'] #Get TTree from file #Get the charged tracks tr = tree.arrays(filter_name="pfcharged.core*") tr["p"] = np.sqrt(tr["pfcharged.core.p4.px"]**2 + tr["pfcharged.core.p4.py"]**2 + tr["pfcharged.core.p4.pz"]**2) tr["pt"] = np.sqrt(tr["pfcharged.core.p4.px"]**2 + tr["pfcharged.core.p4.py"]**2) tr['eta'] = np.log((tr['p'] + tr['pfcharged.core.p4.pz'])/(tr['p'] - tr['pfcharged.core.p4.pz']))/2 tr['phi'] = np.arctan2(tr['pfcharged.core.p4.py'], tr['pfcharged.core.p4.px']) #Pions pi_cut = abs(tr["pfcharged.core.pdgId"]) == 211 pi = tr[pi_cut] #Number of pions in each event pi_sum = ak.num(pi_cut) #Keep events with 2 or more pions pi = pi[pi_sum >= 2] #Make pion pairs per event pi_pairs = ak.combinations(pi, 2) #pt of first pion pair from first event #print(pi_pairs[0]["pt"][0]) pi1, pi2 = ak.unzip(pi_pairs) m_pipi = calc_invariant_mass(pi1, pi2) plt.hist(ak.flatten(m_pipi),bins=400) plt.show()
def process(self, events): # Initialize accumulator out = self.accumulator.identity() dataset = setname #events.metadata['dataset'] isData = 'genWeight' not in events.fields selection = processor.PackedSelection() # --- Calculate Scale factor weight if not isData: # PU weight get_pu_weight = self._corrections['get_pu_weight'][self._year] pu = get_pu_weight(events.Pileup.nTrueInt) # --- Selection # flat dim helper function def flat_dim(arr): sub_arr = ak.flatten(arr) mask = ~ak.is_none(sub_arr) return ak.to_numpy(sub_arr[mask]) # double lepton trigger is_double_ele_trigger = True if not is_double_ele_trigger: double_ele_triggers_arr = np.ones(len(events), dtype=np.bool) else: double_ele_triggers_arr = np.zeros(len(events), dtype=np.bool) for path in self._doubleelectron_triggers[self._year]: if path not in events.HLT.fields: continue double_ele_triggers_arr = double_ele_triggers_arr | events.HLT[ path] #print("Double Electron triggers") #print(double_ele_triggers_arr,len(double_ele_triggers_arr)) # single lepton trigger is_single_ele_trigger = True if not is_single_ele_trigger: single_ele_triggers_arr = np.ones(len(events), dtype=np.bool) else: single_ele_triggers_arr = np.zeros(len(events), dtype=np.bool) for path in self._singleelectron_triggers[self._year]: if path not in events.HLT.fields: continue single_ele_triggers_arr = single_ele_triggers_arr | events.HLT[ path] #print("Single Electron triggers") #print(single_ele_triggers_arr,len(single_ele_triggers_arr)) Initial_events = events print("{0} number of events are detected".format(len(Initial_events))) events = events[single_ele_triggers_arr | double_ele_triggers_arr] print( "Total {0} number of events are remain after triggger | Eff: {1}". format(len(events), len(events) / len(Initial_events) * 100)) # Particle Identification Electron = events.Electron def Electron_selection(ele): return (ele.pt > 20) & (np.abs(ele.eta) < 2.5) & (ele.cutBased > 1) # Electron channel Electron_mask = Electron_selection(Electron) Ele_channel_mask = ak.num(Electron[Electron_mask]) > 1 Ele_channel_events = events[Ele_channel_mask] N_Ele_Channel_events = len(Ele_channel_events) print("#1 Ele channel evts: {0} --> {1}".format( len(events), N_Ele_Channel_events)) # Electron array Ele = Ele_channel_events.Electron Electron_mask = Electron_selection(Ele) Ele_sel = Ele[Electron_mask] # Electron pair ele_pairs = ak.combinations(Ele_sel, 2, axis=1) ele_left, ele_right = ak.unzip(ele_pairs) diele = ele_left + ele_right # OS os_mask = diele.charge == 0 os_diele = diele[os_mask] os_ele_left = ele_left[os_mask] os_ele_right = ele_right[os_mask] # Helper function: High PT argmax def make_leading_pair(target, base): return target[ak.argmax(base.pt, axis=1, keepdims=True)] # -- Only Leading pair -- #leading_diele = make_leading_pair(diele,diele) #leading_ele = make_leading_pair(ele_left,diele) #subleading_ele= make_leading_pair(ele_right,diele) # --OS and Leading pair -- leading_os_diele = make_leading_pair(os_diele, os_diele) leading_os_ele = make_leading_pair(os_ele_left, os_diele) subleading_os_ele = make_leading_pair(os_ele_right, os_diele) # Helper function: Zmass window def makeZmass_window_mask(dielecs, start=60, end=120): mask = (dielecs.mass >= start) & (dielecs.mass <= end) return mask # -- Only Leading pair -- #Zmass_mask = makeZmass_window_mask(leading_diele) #leading_Zwindow_ele = leading_ele[Zmass_mask] #subleading_Zwindow_ele = subleading_ele[Zmass_mask] #leading_Zwindow_diele = leading_diele[Zmass_mask] # -- OS and Leading pair -- Zmass_mask_os = makeZmass_window_mask(leading_os_diele) leading_os_Zwindow_ele = leading_os_ele[Zmass_mask_os] subleading_os_Zwindow_ele = subleading_os_ele[Zmass_mask_os] leading_os_Zwindow_diele = leading_os_diele[Zmass_mask_os] #ele1PT = flat_dim(leading_Zwindow_ele.pt) #ele1Eta = flat_dim(leading_Zwindow_ele.eta) #ele1Phi = flat_dim(leading_Zwindow_ele.phi) #ele2PT = flat_dim(subleading_Zwindow_ele.pt) #ele2Eta = flat_dim(subleading_Zwindow_ele.eta) #ele2Phi = flat_dim(subleading_Zwindow_ele.phi) #Mee = flat_dim(leading_Zwindow_diele.mass) #charge = flat_dim(leading_Zwindow_diele.charge) os_ele1PT = flat_dim(leading_os_Zwindow_ele.pt) os_ele1Eta = flat_dim(leading_os_Zwindow_ele.eta) os_ele1Phi = flat_dim(leading_os_Zwindow_ele.phi) os_ele2PT = flat_dim(subleading_os_Zwindow_ele.pt) os_ele2Eta = flat_dim(subleading_os_Zwindow_ele.eta) os_ele2Phi = flat_dim(subleading_os_Zwindow_ele.phi) os_Mee = flat_dim(leading_os_Zwindow_diele.mass) os_charge = flat_dim(leading_os_Zwindow_diele.charge) #print("#5 Leading PT : {0}".format(len(ele1PT))) #print("#5 Leading os PT {0}".format(len(os_ele1PT))) # --- Apply weight --- on progress # #if not isData: # weights = processor.Weights(len(events)) # weights.add('pileup',pu) out["sumw"][dataset] += len(events) out["os_mass"].fill(dataset=dataset, os_mass=os_Mee) out["os_ele1pt"].fill(dataset=dataset, os_ele1pt=os_ele1PT) out["os_ele1eta"].fill(dataset=dataset, os_ele1eta=os_ele1Eta) out["os_ele1phi"].fill(dataset=dataset, os_ele1phi=os_ele1Phi) out["os_ele2pt"].fill(dataset=dataset, os_ele2pt=os_ele2PT) out["os_ele2eta"].fill(dataset=dataset, os_ele2eta=os_ele2Eta) out["os_ele2phi"].fill(dataset=dataset, os_ele2phi=os_ele2Phi) return out