def test_firsts(): assert ak.to_list( ak.firsts( ak.singletons(ak.Array([1.1, 2.2, None, 3.3, None, None, 4.4, 5.5])), axis=1 ) ) == [1.1, 2.2, None, 3.3, None, None, 4.4, 5.5] assert ( ak.to_list( ak.firsts( ak.singletons( ak.Array([[1.1, 2.2, None], [3.3, None], [None], [4.4, 5.5]]) ), axis=2, ) ) == [[1.1, 2.2, None], [3.3, None], [None], [4.4, 5.5]] ) assert ( ak.to_list( ak.firsts( ak.singletons( ak.Array( [[[1.1, 2.2, None]], [[3.3, None]], [[None]], [[4.4, 5.5]]] ) ), axis=3, ) ) == [[[1.1, 2.2, None]], [[3.3, None]], [[None]], [[4.4, 5.5]]] )
def get_leading_pair(taus, Pt_thr): tau_1 = taus[0] tau_2 = taus[1] # select only true taus that pass generator preselection num_tau_mask_1 = num_mask_eff(tau_1, Pt_thr=Pt_thr) num_tau_mask_2 = num_mask_eff(tau_2, Pt_thr=Pt_thr) pair_mask = num_tau_mask_1 & num_tau_mask_2 ev_mask = ditau_selection(num_tau_mask_1, num_tau_mask_2) tau_1_selected = (tau_1[pair_mask])[ev_mask] tau_2_selected = (tau_2[pair_mask])[ev_mask] # get arrays of leading and subleading taus tau_leading = ak.firsts(tau_1_selected, axis=-1) tau_subleading = ak.firsts(tau_2_selected, axis=-1) return tau_leading, tau_subleading
def calculate_selection(self, syst_tag, events): """ """ electrons = events.ele electrons["label"] = -1 * awkward.ones_like(electrons.pt) if not self.is_data: electrons["label"] = awkward.where( electrons.genPartFlav == 1, awkward.ones_like(electrons.label), electrons.label) electrons["label"] = awkward.where( (electrons.genPartFlav == 3) | (electrons.genPartFlav == 4) | (electrons.genPartFlav == 5), awkward.zeros_like(electrons.label), electrons.label) fields = [x for x in events.fields if x not in ["ele", "Electron"]] for x in fields: if x == "ZCand": electrons[x] = awkward.firsts(events[x]) else: electrons[x] = events[x] electrons = awkward.flatten(electrons) dummy_cut = electrons.pt >= 0 return dummy_cut, electrons
def set_genZ(events, genBranches, selection_options, debug): if genBranches is None: events["genZ_decayMode"] = ak.from_numpy(-1 * numpy.ones(len(events))) else: electron_idxs = abs(genBranches.pdgId) == 11 muon_idxs = abs(genBranches.pdgId) == 13 tau_idxs = abs(genBranches.pdgId) == 15 motherOfElectrons = genBranches.genPartIdxMother[electron_idxs] motherOfMuons = genBranches.genPartIdxMother[muon_idxs] motherOfTaus = genBranches.genPartIdxMother[tau_idxs] ZToEleEvents = ak.sum( (genBranches.pdgId[motherOfElectrons] == 23), axis=1) > 0 ZToMuEvents = ak.sum( (genBranches.pdgId[motherOfMuons] == 23), axis=1) > 0 ZToTauEvents = ak.sum( (genBranches.pdgId[motherOfTaus] == 23), axis=1) > 0 # W decay dudes WToEEvents = ak.sum( (abs(genBranches.pdgId[motherOfElectrons]) == 24), axis=1) > 0 WToMuEvents = ak.sum( (abs(genBranches.pdgId[motherOfMuons]) == 24), axis=1) > 0 WToTauEvents = ak.sum( (abs(genBranches.pdgId[motherOfTaus]) == 24), axis=1) > 0 events["genZ_decayMode"] = ak.from_numpy( numpy.zeros(len(events)) ) + 1 * ZToEleEvents + 2 * ZToMuEvents + 3 * ZToTauEvents + 4 * WToEEvents + 5 * WToMuEvents + 6 * WToTauEvents events["tau_motherID"] = ak.fill_none( ak.firsts(genBranches.pdgId[motherOfTaus]), 0) return events
def VQQgenmatch(events): dataset = events.metadata['dataset'] if 'WJetsToQQ' in dataset: motherId = 24 elif 'ZJetsToQQ' in dataset: motherId = 23 elif 'VectorDiJet' in dataset: motherId = 55 else: return np.ones(len(events.FatJet.pt), dtype=bool) #events.FatJet.pt.ones_like() mother = ak.flatten(events.GenPart[ (abs(events.GenPart.pdgId) == motherId) & events.GenPart.hasFlags(["isLastCopy", "fromHardProcess"])]) print(mother) print(mother.children.pdgId) print(mother.children) try: q0 = mother.children[:, 0] q1 = mother.children[:, 1] except: q0 = mother.children[:, 0] q1 = mother.children[:, 1] print(q0.pdgId, q1.pdgId) leading_jet = ak.firsts(events.FatJet) return (leading_jet.delta_r2(q0) < 0.8 * 0.8) & (leading_jet.delta_r2(q1) < 0.8 * 0.8)
def test_cartesian(): muon = ak.Array([[{"pt": 1.0}], []], with_name="muon") electron = ak.Array([[], [{"pt": 1.0}]], with_name="electron") muon = muon[muon.pt > 5] electron = electron[electron.pt > 5] leptons = ak.concatenate([muon, electron], axis=1) candidate = ak.firsts(leptons) assert ak.to_list(ak.Array(candidate)) == [None, None] result = ak.cartesian([candidate, candidate], axis=0) assert ak.to_list(result) == [ (None, None), (None, None), (None, None), (None, None), ] result = ak.cartesian([candidate, ak.Array([[1, 2, 3], []])], axis=1) assert ak.to_list(result) == [None, None] one, two = ak.broadcast_arrays(candidate, ak.Array([[1, 2, 3], []])) assert ak.to_list(one) == [None, None] assert ak.to_list(two) == [None, None]
def test_firsts_singletons(): array = ak.Array([None, 1.1, 2.2, None, 3.3, None, None, 4.4, 5.5, None]) one = ak.singletons(array) assert ak.to_list(one) == [[], [1.1], [2.2], [], [3.3], [], [], [4.4], [5.5], []] two = ak.firsts(one) assert ak.to_list(two) == [None, 1.1, 2.2, None, 3.3, None, None, 4.4, 5.5, None] array = ak.repartition(array, 3) assert isinstance(array.layout, ak.partition.PartitionedArray) one = ak.singletons(array) assert isinstance(one.layout, ak.partition.PartitionedArray) assert ak.to_list(one) == [[], [1.1], [2.2], [], [3.3], [], [], [4.4], [5.5], []] two = ak.firsts(one) assert isinstance(two.layout, ak.partition.PartitionedArray) assert ak.to_list(two) == [None, 1.1, 2.2, None, 3.3, None, None, 4.4, 5.5, None]
def nearest( self, other, axis=1, metric=lambda a, b: a.delta_r(b), return_metric=False, threshold=None, ): """Return nearest object to this one Finds item in ``other`` satisfying ``min(metric(self, other))``. The two arrays should be broadcast-compatible on all axes other than the specified axis, which will be used to form a cartesian product. If axis=None, broadcast arrays directly. The return shape will be that of ``self``. Parameters ---------- other : awkward.Array Another array with same shape in all but ``axis`` axis : int, optional The axis to form the cartesian product (default 1). If None, the metric is directly evaluated on the input arrays (i.e. they should broadcast) metric : callable A function of two arguments, returning a scalar. The default metric is `delta_r`. return_metric : bool, optional If true, return both the closest object and its metric (default false) threshold : Number, optional If set, any objects with ``metric > threshold`` will be masked from the result """ mval, (a, b) = self.metric_table(other, axis, metric, return_combinations=True) if axis is None: # NotImplementedError: awkward.firsts with axis=-1 axis = other.layout.purelist_depth - 2 mmin = awkward.argmin(mval, axis=axis + 1, keepdims=True) out = awkward.firsts(b[mmin], axis=axis + 1) metric = awkward.firsts(mval[mmin], axis=axis + 1) if threshold is not None: out = out.mask[metric <= threshold] if return_metric: return out, metric return out
def test_firsts(): array = ak.Array([[[0, 1, 2], []], [[3, 4]], [], [[5], [6, 7, 8, 9]]]) assert ak.to_list(ak.firsts(array, axis=0)) == [[0, 1, 2], []] assert ak.to_list(ak.firsts(array, axis=1)) == [[0, 1, 2], [3, 4], None, [5]] assert ak.to_list(ak.firsts(array, axis=2)) == [[0, None], [3], [], [5, 6]] assert ak.to_list(ak.firsts(array, axis=-1)) == [[0, None], [3], [], [5, 6]] assert ak.to_list(ak.firsts(array, axis=-2)) == [[0, 1, 2], [3, 4], None, [5]] assert ak.to_list(ak.firsts(array, axis=-3)) == [[0, 1, 2], []] with pytest.raises(ValueError): ak.firsts(array, axis=-4)
def get_vpt(check_offshell=False): """Only the leptonic samples have no resonance in the decay tree, and only when M is beyond the configured Breit-Wigner cutoff (usually 15*width) """ boson = ak.firsts( genpart[((genpart.pdgId == 23) | (abs(genpart.pdgId) == 24)) & genpart.hasFlags(["fromHardProcess", "isLastCopy"])]) if check_offshell: offshell = genpart[ genpart.hasFlags(["fromHardProcess", "isLastCopy"]) & ak.is_none(boson) & (abs(genpart.pdgId) >= 11) & (abs(genpart.pdgId) <= 16)].sum() return ak.where(ak.is_none(boson.pt), offshell.pt, boson.pt) return np.array(ak.fill_none(boson.pt, 0.))
def TTsemileptonicmatch(events): dataset = events.metadata['dataset'] if 'TTToSemiLeptonic' not in dataset: return np.zeros(len(events.FatJet.pt), dtype=bool) child = events.GenPart[ (abs(events.GenPart.pdgId) == 24) & events.GenPart.hasFlags(["isLastCopy", "fromHardProcess"])].children fatjet = ak.firsts(events.FatJet) n_matched_quarks = np.zeros(len(fatjet)) for ii in [0, 1]: for jj in [0, 1]: n_matched_quarks = n_matched_quarks + ak.fill_none( (fatjet.delta_r2(child[:, ii, jj]) < 0.8**2) & (abs(child[:, ii, jj].pdgId) < 6), 0.) print(n_matched_quarks) return n_matched_quarks
def process(self, events): # Initialize accumulator out = self.accumulator.identity() dataset = sample_name # events.metadata['dataset'] # Stop processing if there is no event remain if len(events) == 0: return out # Cut flow cut0 = np.zeros(len(events)) # <----- Helper functions ------># # Sort by PT helper function def sort_by_pt(ele, pho, jet): ele = ele[ak.argsort(ele.pt, ascending=False, axis=1)] pho = pho[ak.argsort(pho.pt, ascending=False, axis=1)] jet = jet[ak.argsort(jet.pt, ascending=False, axis=1)] return ele, pho, jet # Lorentz vectors from coffea.nanoevents.methods import vector ak.behavior.update(vector.behavior) def TLorentz_vector(vec): vec = ak.zip( { "x": vec.x, "y": vec.y, "z": vec.z, "t": vec.t }, with_name="LorentzVector", ) return vec def TLorentz_vector_cylinder(vec): vec = ak.zip( { "pt": vec.pt, "eta": vec.eta, "phi": vec.phi, "mass": vec.mass, }, with_name="PtEtaPhiMLorentzVector", ) return vec # Cut-based ID modification @numba.njit def PhotonVID(vid, idBit): rBit = 0 for x in range(0, 7): rBit |= (1 << x) if ((vid >> (x * 2)) & 0b11 >= idBit) else 0 return rBit # Inverse Sieie and upper limit @numba.njit def make_fake_obj_mask(Pho, builder): # for eventIdx,pho in enumerate(tqdm(Pho)): # --Event Loop for eventIdx, pho in enumerate(Pho): builder.begin_list() if len(pho) < 1: continue for phoIdx, _ in enumerate(pho): # --Photon Loop vid = Pho[eventIdx][phoIdx].vidNestedWPBitmap vid_cuts1 = PhotonVID(vid, 1) # Loose photon vid_cuts2 = PhotonVID(vid, 2) # Medium photon vid_cuts3 = PhotonVID(vid, 3) # Tight photon # Field name # |0|0|0|0|0|0|0| # |IsoPho|IsoNeu|IsoChg|Sieie|hoe|scEta|PT| # 1. Turn off cut (ex turn off Sieie # |1|1|1|0|1|1|1| = |1|1|1|0|1|1|1| # 2. Inverse cut (ex inverse Sieie) # |1|1|1|1|1|1|1| = |1|1|1|0|1|1|1| # if (vid_cuts2 & 0b1111111 == 0b1111111): # Cut applied # if vid_cuts2 & 0b1111111 == 0b1110111: # Inverse Sieie if (vid_cuts2 & 0b1110111 == 0b1110111): # Without Sieie builder.boolean(True) else: builder.boolean(False) builder.end_list() return builder # Golden Json file if self._year == "2018": injson = "/x5/cms/jwkim/gitdir/JWCorp/JW_analysis/Coffea_WZG/Corrections/Cert_314472-325175_13TeV_Legacy2018_Collisions18_JSON.txt.RunABCD" if self._year == "2017": injson = "/x5/cms/jwkim/gitdir/JWCorp/JW_analysis/Coffea_WZG/Corrections/Cert_294927-306462_13TeV_UL2017_Collisions17_GoldenJSON.txt" # --- Selection Initial_events = events # Good Run ( Golden Json files ) from coffea import lumi_tools lumi_mask_builder = lumi_tools.LumiMask(injson) lumimask = ak.Array( lumi_mask_builder.__call__(events.run, events.luminosityBlock)) events = events[lumimask] # print("{0}% of files pass good-run conditions".format(len(events)/ len(Initial_events))) # Stop processing if there is no event remain if len(events) == 0: return out ##----------- Cut flow1: Passing Triggers # 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] events.Electron, events.Photon, events.Jet = sort_by_pt( events.Electron, events.Photon, events.Jet) # Apply cut1 Initial_events = events # events = events[single_ele_triggers_arr | double_ele_triggers_arr] events = events[double_ele_triggers_arr] cut1 = np.ones(len(events)) # Set Particles Electron = events.Electron Muon = events.Muon Photon = events.Photon MET = events.MET Jet = events.Jet # Stop processing if there is no event remain if len(Electron) == 0: return out # --Muon ( only used to calculate dR ) MuSelmask = ((Muon.pt >= 10) & (abs(Muon.eta) <= 2.5) & (Muon.tightId) & (Muon.pfRelIso04_all < 0.15)) # Muon = ak.mask(Muon,MuSelmask) Muon = Muon[MuSelmask] ##----------- Cut flow2: Electron Selection EleSelmask = ((Electron.pt >= 20) & (np.abs(Electron.eta + Electron.deltaEtaSC) < 1.479) & (Electron.cutBased > 2) & (abs(Electron.dxy) < 0.05) & (abs(Electron.dz) < 0.1)) | ( (Electron.pt >= 20) & (np.abs(Electron.eta + Electron.deltaEtaSC) > 1.479) & (np.abs(Electron.eta + Electron.deltaEtaSC) <= 2.5) & (Electron.cutBased > 2) & (abs(Electron.dxy) < 0.1) & (abs(Electron.dz) < 0.2)) Electron = Electron[EleSelmask] # apply cut 2 Tri_electron_mask = ak.num(Electron) >= 2 Electron = Electron[Tri_electron_mask] Photon = Photon[Tri_electron_mask] Jet = Jet[Tri_electron_mask] MET = MET[Tri_electron_mask] Muon = Muon[Tri_electron_mask] events = events[Tri_electron_mask] # Stop processing if there is no event remain if len(Electron) == 0: return out cut2 = np.ones(len(Photon)) * 2 ##----------- Cut flow3: Photon Selection # Basic photon selection isgap_mask = (abs(Photon.eta) < 1.442) | ((abs(Photon.eta) > 1.566) & (abs(Photon.eta) < 2.5)) Pixel_seed_mask = ~Photon.pixelSeed PT_mask = Photon.pt >= 20 # dR cut with selected Muon and Electrons dr_pho_ele_mask = ak.all(Photon.metric_table(Electron) >= 0.5, axis=-1) # default metric table: delta_r dr_pho_mu_mask = ak.all(Photon.metric_table(Muon) >= 0.5, axis=-1) # Photon basic mask PhoSelmask = (PT_mask & isgap_mask & Pixel_seed_mask & dr_pho_ele_mask & dr_pho_mu_mask) Photon = Photon[PhoSelmask] # Apply cut 3 A_photon_mask = ak.num(Photon) > 0 Electron = Electron[A_photon_mask] Photon = Photon[A_photon_mask] Jet = Jet[A_photon_mask] Muon = Muon[A_photon_mask] MET = MET[A_photon_mask] events = events[A_photon_mask] # ID for fake photon Photon_template_mask = make_fake_obj_mask( Photon, ak.ArrayBuilder()).snapshot() Photon = Photon[Photon_template_mask] # Apply cut -Fake Photon - A_photon_mask = ak.num(Photon) > 0 Electron = Electron[A_photon_mask] Photon = Photon[A_photon_mask] Jet = Jet[A_photon_mask] Muon = Muon[A_photon_mask] MET = MET[A_photon_mask] events = events[A_photon_mask] # Stop processing if there is no event remain if len(Electron) == 0: return out cut3 = np.ones(len(Photon)) * 3 ##----------- Cut flow4: Select 2 OSSF electrons from Z @numba.njit def find_2lep(events_leptons, builder): for leptons in events_leptons: builder.begin_list() nlep = len(leptons) for i0 in range(nlep): for i1 in range(i0 + 1, nlep): if leptons[i0].charge + leptons[i1].charge != 0: continue if nlep == 2: builder.begin_tuple(2) builder.index(0).integer(i0) builder.index(1).integer(i1) builder.end_tuple() else: for i2 in range(nlep): if len({i0, i1, i2}) < 3: continue builder.begin_tuple(3) builder.index(0).integer(i0) builder.index(1).integer(i1) builder.index(2).integer(i2) builder.end_tuple() builder.end_list() return builder ossf_idx = find_2lep(Electron, ak.ArrayBuilder()).snapshot() # OSSF cut ossf_mask = ak.num(ossf_idx) >= 1 ossf_idx = ossf_idx[ossf_mask] Electron = Electron[ossf_mask] Photon = Photon[ossf_mask] Jet = Jet[ossf_mask] MET = MET[ossf_mask] Double_electron = [Electron[ossf_idx[idx]] for idx in "01"] from coffea.nanoevents.methods import vector ak.behavior.update(vector.behavior) Diele = ak.zip({ "lep1": Double_electron[0], "lep2": Double_electron[1], "p4": TLorentz_vector(Double_electron[0] + Double_electron[1]), }) bestZ_idx = ak.singletons( ak.argmin(abs(Diele.p4.mass - 91.1876), axis=1)) Diele = Diele[bestZ_idx] cut4 = np.ones(len(Electron)) * 4 ##----------- Cut flow 5: Event Selection def make_leading_pair(target, base): return target[ak.argmax(base.pt, axis=1, keepdims=True)] leading_pho = make_leading_pair(Photon, Photon) # Mee cut Mee_cut_mask = ak.firsts(Diele.p4.mass) > 4 # Electron PT cuts Elept_mask = ak.firsts((Diele.lep1.pt >= 25) & (Diele.lep2.pt >= 20)) # MET cuts MET_mask = MET.pt > 20 # --------Mask -------# Event_sel_mask = Mee_cut_mask & Elept_mask & MET_mask Diele_sel = Diele[Event_sel_mask] leading_pho_sel = leading_pho[Event_sel_mask] Jet_sel = Jet[Event_sel_mask] MET_sel = MET[Event_sel_mask] cut5 = np.ones(len(Diele)) * 5 # Photon EE and EB isEE_mask = leading_pho.isScEtaEE isEB_mask = leading_pho.isScEtaEB Pho_EE = leading_pho[isEE_mask & Event_sel_mask] Pho_EB = leading_pho[isEB_mask & Event_sel_mask] # -------------------- Flatten variables ---------------------------# # -- Ele1 --# Ele1_PT = ak.flatten(Diele_sel.lep1.pt) Ele1_Eta = ak.flatten(Diele_sel.lep1.eta) Ele1_Phi = ak.flatten(Diele_sel.lep1.phi) # -- Ele2 --# Ele2_PT = ak.flatten(Diele_sel.lep2.pt) Ele2_Eta = ak.flatten(Diele_sel.lep2.eta) Ele2_Phi = ak.flatten(Diele_sel.lep2.phi) # -- Pho -- # Pho_PT = ak.flatten(leading_pho_sel.pt) Pho_Eta = ak.flatten(leading_pho_sel.eta) Pho_Phi = ak.flatten(leading_pho_sel.phi) # -- Pho EB --# Pho_EB_PT = ak.flatten(Pho_EB.pt) Pho_EB_Eta = ak.flatten(Pho_EB.eta) Pho_EB_Phi = ak.flatten(Pho_EB.phi) Pho_EB_Isochg = ak.flatten(Pho_EE.pfRelIso03_chg) Pho_EB_Sieie = ak.flatten(Pho_EE.sieie) # -- Pho EE --# Pho_EE_PT = ak.flatten(Pho_EE.pt) Pho_EE_Eta = ak.flatten(Pho_EE.eta) Pho_EE_Phi = ak.flatten(Pho_EE.phi) Pho_EE_Isochg = ak.flatten(Pho_EE.pfRelIso03_chg) Pho_EE_Sieie = ak.flatten(Pho_EE.sieie) # --Kinematics --# Diele_mass = ak.flatten(Diele_sel.p4.mass) eeg_vec = Diele_sel.p4 + leading_pho_sel eeg_mass = ak.flatten(eeg_vec.mass) leading_ele, subleading_ele = ak.flatten( TLorentz_vector_cylinder(Diele_sel.lep1)), ak.flatten( TLorentz_vector_cylinder(Diele_sel.lep2)) dR_e1pho = ak.flatten( leading_ele.delta_r(leading_pho_sel)) # dR pho,ele1 dR_e2pho = ak.flatten( subleading_ele.delta_r(leading_pho_sel)) # dR pho,ele2 dR_jpho = ak.flatten(Jet_sel[:, 0].delta_r(leading_pho_sel)) MET_PT = ak.to_numpy(MET_sel.pt) # -------------------- Sieie bins---------------------------# def make_bins(pt, eta, sieie, bin_range_str): bin_dict = { "PT_1_eta_1": (pt > 20) & (pt < 30) & (eta < 1), "PT_1_eta_2": (pt > 20) & (pt < 30) & (eta > 1) & (eta < 1.5), "PT_1_eta_3": (pt > 20) & (pt < 30) & (eta > 1.5) & (eta < 2), "PT_1_eta_4": (pt > 20) & (pt < 30) & (eta > 2) & (eta < 2.5), "PT_2_eta_1": (pt > 30) & (pt < 40) & (eta < 1), "PT_2_eta_2": (pt > 30) & (pt < 40) & (eta > 1) & (eta < 1.5), "PT_2_eta_3": (pt > 30) & (pt < 40) & (eta > 1.5) & (eta < 2), "PT_2_eta_4": (pt > 30) & (pt < 40) & (eta > 2) & (eta < 2.5), "PT_3_eta_1": (pt > 40) & (pt < 50) & (eta < 1), "PT_3_eta_2": (pt > 40) & (pt < 50) & (eta > 1) & (eta < 1.5), "PT_3_eta_3": (pt > 40) & (pt < 50) & (eta > 1.5) & (eta < 2), "PT_3_eta_4": (pt > 40) & (pt < 50) & (eta > 2) & (eta < 2.5), "PT_4_eta_1": (pt > 50) & (eta < 1), "PT_4_eta_2": (pt > 50) & (eta > 1) & (eta < 1.5), "PT_4_eta_3": (pt > 50) & (eta > 1.5) & (eta < 2), "PT_4_eta_4": (pt > 50) & (eta > 2) & (eta < 2.5), } binmask = bin_dict[bin_range_str] return ak.to_numpy(sieie[binmask]) bin_name_list = [ "PT_1_eta_1", "PT_1_eta_2", "PT_1_eta_3", "PT_1_eta_4", "PT_2_eta_1", "PT_2_eta_2", "PT_2_eta_3", "PT_2_eta_4", "PT_3_eta_1", "PT_3_eta_2", "PT_3_eta_3", "PT_3_eta_4", "PT_4_eta_1", "PT_4_eta_2", "PT_4_eta_3", "PT_4_eta_4", ] binned_sieie_hist = {} for name in bin_name_list: binned_sieie_hist[name] = make_bins( ak.flatten(leading_pho_sel.pt), ak.flatten(abs(leading_pho_sel.eta)), ak.flatten(leading_pho_sel.sieie), name, ) # -------------------- Fill hist ---------------------------# # Initial events out["sumw"][dataset] += len(Initial_events) # print("cut5: ",len(cut5)) # Cut flow loop for cut in [cut0, cut1, cut2, cut3, cut4, cut5]: out["cutflow"].fill(dataset=dataset, cutflow=cut) # --Ele1 -- # out["ele1pt"].fill(dataset=dataset, ele1pt=Ele1_PT) out["ele1eta"].fill(dataset=dataset, ele1eta=Ele1_Eta) out["ele1phi"].fill(dataset=dataset, ele1phi=Ele1_Phi) # --Ele2 -- # out["ele2pt"].fill(dataset=dataset, ele2pt=Ele2_PT) out["ele2eta"].fill(dataset=dataset, ele2eta=Ele2_Eta) out["ele2phi"].fill(dataset=dataset, ele2phi=Ele2_Phi) # --Photon-- # out["phopt"].fill(dataset=dataset, phopt=Pho_PT) out["phoeta"].fill(dataset=dataset, phoeta=Pho_Eta) out["phophi"].fill(dataset=dataset, phophi=Pho_Phi) # --Photon EB --# out["pho_EB_pt"].fill( dataset=dataset, pho_EB_pt=Pho_EB_PT, ) out["pho_EB_eta"].fill( dataset=dataset, pho_EB_eta=Pho_EB_Eta, ) out["pho_EB_phi"].fill( dataset=dataset, pho_EB_phi=Pho_EB_Phi, ) out["pho_EB_sieie"].fill( dataset=dataset, pho_EB_sieie=Pho_EB_Sieie, ) out["pho_EB_Iso_chg"].fill(dataset=dataset, pho_EB_Iso_chg=Pho_EB_Isochg) # --Photon EE --# out["pho_EE_pt"].fill( dataset=dataset, pho_EE_pt=Pho_EE_PT, ) out["pho_EE_eta"].fill( dataset=dataset, pho_EE_eta=Pho_EE_Eta, ) out["pho_EE_phi"].fill( dataset=dataset, pho_EE_phi=Pho_EE_Phi, ) out["pho_EE_sieie"].fill( dataset=dataset, pho_EE_sieie=Pho_EE_Sieie, ) out["pho_EE_Iso_chg"].fill(dataset=dataset, pho_EE_Iso_chg=Pho_EE_Isochg) # -- Kinematic variables -- # out["mass"].fill(dataset=dataset, Mee=Diele_mass) out["mass_eea"].fill(dataset=dataset, mass_eea=eeg_mass) out["met"].fill(dataset=dataset, met=MET_PT) out["dR_ae1"].fill(dataset=dataset, dR_ae1=dR_e1pho) out["dR_ae2"].fill(dataset=dataset, dR_ae2=dR_e2pho) out["dR_aj"].fill(dataset=dataset, dR_aj=dR_jpho) # test_target = binned_sieie_hist['PT_1_eta_1'] # print("CheckAAA: ",test_target[test_target > 0.05]) # print("CheckBBB: ",test_target[test_target > 0]) # -- Binned sieie hist -- # if len(binned_sieie_hist["PT_1_eta_1"] > 0): out["PT_1_eta_1"].fill(dataset=dataset, PT_1_eta_1=binned_sieie_hist["PT_1_eta_1"]) if len(binned_sieie_hist["PT_1_eta_2"] > 0): out["PT_1_eta_2"].fill(dataset=dataset, PT_1_eta_2=binned_sieie_hist["PT_1_eta_2"]) if len(binned_sieie_hist["PT_1_eta_3"] > 0): out["PT_1_eta_3"].fill(dataset=dataset, PT_1_eta_3=binned_sieie_hist["PT_1_eta_3"]) if len(binned_sieie_hist["PT_1_eta_4"] > 0): out["PT_1_eta_4"].fill(dataset=dataset, PT_1_eta_4=binned_sieie_hist["PT_1_eta_4"]) if len(binned_sieie_hist["PT_2_eta_1"] > 0): out["PT_2_eta_1"].fill(dataset=dataset, PT_2_eta_1=binned_sieie_hist["PT_2_eta_1"]) if len(binned_sieie_hist["PT_2_eta_2"] > 0): out["PT_2_eta_2"].fill(dataset=dataset, PT_2_eta_2=binned_sieie_hist["PT_2_eta_2"]) if len(binned_sieie_hist["PT_2_eta_3"] > 0): out["PT_2_eta_3"].fill(dataset=dataset, PT_2_eta_3=binned_sieie_hist["PT_2_eta_3"]) if len(binned_sieie_hist["PT_2_eta_4"] > 0): out["PT_2_eta_4"].fill(dataset=dataset, PT_2_eta_4=binned_sieie_hist["PT_2_eta_4"]) if len(binned_sieie_hist["PT_3_eta_1"] > 0): out["PT_3_eta_1"].fill(dataset=dataset, PT_3_eta_1=binned_sieie_hist["PT_3_eta_1"]) if len(binned_sieie_hist["PT_3_eta_2"] > 0): out["PT_3_eta_2"].fill(dataset=dataset, PT_3_eta_2=binned_sieie_hist["PT_3_eta_2"]) if len(binned_sieie_hist["PT_3_eta_3"] > 0): out["PT_3_eta_3"].fill(dataset=dataset, PT_3_eta_3=binned_sieie_hist["PT_3_eta_3"]) if len(binned_sieie_hist["PT_3_eta_4"] > 0): out["PT_3_eta_4"].fill(dataset=dataset, PT_3_eta_4=binned_sieie_hist["PT_3_eta_4"]) if len(binned_sieie_hist["PT_4_eta_1"] > 0): out["PT_4_eta_1"].fill(dataset=dataset, PT_4_eta_1=binned_sieie_hist["PT_4_eta_1"]) if len(binned_sieie_hist["PT_4_eta_2"] > 0): out["PT_4_eta_2"].fill(dataset=dataset, PT_4_eta_2=binned_sieie_hist["PT_4_eta_2"]) if len(binned_sieie_hist["PT_4_eta_3"] > 0): out["PT_4_eta_3"].fill(dataset=dataset, PT_4_eta_3=binned_sieie_hist["PT_4_eta_3"]) if len(binned_sieie_hist["PT_4_eta_4"] > 0): print("## show me the last bin: ", binned_sieie_hist["PT_4_eta_4"]) out["PT_4_eta_4"].fill(dataset=dataset, PT_4_eta_4=binned_sieie_hist["PT_4_eta_4"]) return out
def process(self, events): # Initialize accumulator out = self.accumulator.identity() dataset = sample_name # events.metadata['dataset'] # Data or MC isData = "genWeight" not in events.fields isFake = self._isFake # Stop processing if there is no event remain if len(events) == 0: return out # Golden Json file if (self._year == "2018") and isData: injson = "/x5/cms/jwkim/gitdir/JWCorp/JW_analysis/Coffea_WZG/Corrections/Cert_314472-325175_13TeV_Legacy2018_Collisions18_JSON.txt.RunABD" if (self._year == "2017") and isData: injson = "/x5/cms/jwkim/gitdir/JWCorp/JW_analysis/Coffea_WZG/Corrections/Cert_294927-306462_13TeV_UL2017_Collisions17_GoldenJSON.txt" # <----- Get Scale factors ------># if not isData: # Egamma reco ID get_ele_reco_above20_sf = self._corrections[ "get_ele_reco_above20_sf"][self._year] get_ele_medium_id_sf = self._corrections["get_ele_medium_id_sf"][ self._year] get_pho_medium_id_sf = self._corrections["get_pho_medium_id_sf"][ self._year] # DoubleEG trigger # 2016, 2017 are not applied yet if self._year == "2018": 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] # Muon ID, Iso get_mu_tight_id_sf = self._corrections["get_mu_tight_id_sf"][ self._year] get_mu_tight_iso_sf = self._corrections["get_mu_tight_iso_sf"][ self._year] # PU weight with custom made npy and multi-indexing pu_weight_idx = ak.values_astype(events.Pileup.nTrueInt, "int64") pu = self._puweight_arr[pu_weight_idx] # <----- Helper functions ------># # Sort by PT helper function def sort_by_pt(ele, pho, jet): ele = ele[ak.argsort(ele.pt, ascending=False, axis=1)] pho = pho[ak.argsort(pho.pt, ascending=False, axis=1)] jet = jet[ak.argsort(jet.pt, ascending=False, axis=1)] return ele, pho, jet # Lorentz vectors from coffea.nanoevents.methods import vector ak.behavior.update(vector.behavior) def TLorentz_vector(vec): vec = ak.zip( { "x": vec.x, "y": vec.y, "z": vec.z, "t": vec.t }, with_name="LorentzVector", ) return vec def TLorentz_vector_cylinder(vec): vec = ak.zip( { "pt": vec.pt, "eta": vec.eta, "phi": vec.phi, "mass": vec.mass, }, with_name="PtEtaPhiMLorentzVector", ) return vec # <----- Selection ------># Initial_events = events # Good Run ( Golden Json files ) from coffea import lumi_tools if isData: lumi_mask_builder = lumi_tools.LumiMask(injson) lumimask = ak.Array( lumi_mask_builder.__call__(events.run, events.luminosityBlock)) events = events[lumimask] # print("{0}% of files pass good-run conditions".format(len(events)/ len(Initial_events))) # Stop processing if there is no event remain if len(events) == 0: return out # Cut flow cut0 = np.zeros(len(events)) ##----------- Cut flow1: Passing Triggers # 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] events.Electron, events.Photon, events.Jet = sort_by_pt( events.Electron, events.Photon, events.Jet) # Good Primary vertex nPV = events.PV.npvsGood nPV_nw = events.PV.npvsGood if not isData: nPV = nPV * pu print(pu) # Apply cut1 events = events[double_ele_triggers_arr] if not isData: pu = pu[double_ele_triggers_arr] # Stop processing if there is no event remain if len(events) == 0: return out cut1 = np.ones(len(events)) # Set Particles Electron = events.Electron Muon = events.Muon Photon = events.Photon MET = events.MET Jet = events.Jet ##----------- Cut flow2: Muon Selection MuSelmask = ((Muon.pt >= 10) & (abs(Muon.eta) <= 2.5) & (Muon.tightId) & (Muon.pfRelIso04_all < 0.15)) Muon = Muon[MuSelmask] # Exatly one muon Muon_sel_mask = ak.num(Muon) == 1 Electron = Electron[Muon_sel_mask] Photon = Photon[Muon_sel_mask] Jet = Jet[Muon_sel_mask] MET = MET[Muon_sel_mask] Muon = Muon[Muon_sel_mask] events = events[Muon_sel_mask] if not isData: pu = pu[Muon_sel_mask] # Stop processing if there is no event remain if len(Electron) == 0: return out cut2 = np.ones(len(Photon)) * 2 ##----------- Cut flow3: Electron Selection EleSelmask = ((Electron.pt >= 10) & (np.abs(Electron.eta + Electron.deltaEtaSC) < 1.479) & (Electron.cutBased > 2) & (abs(Electron.dxy) < 0.05) & (abs(Electron.dz) < 0.1)) | ( (Electron.pt >= 10) & (np.abs(Electron.eta + Electron.deltaEtaSC) > 1.479) & (np.abs(Electron.eta + Electron.deltaEtaSC) <= 2.5) & (Electron.cutBased > 2) & (abs(Electron.dxy) < 0.1) & (abs(Electron.dz) < 0.2)) Electron = Electron[EleSelmask] # Exactly two electrons ee_mask = ak.num(Electron) == 2 Electron = Electron[ee_mask] Photon = Photon[ee_mask] Jet = Jet[ee_mask] MET = MET[ee_mask] Muon = Muon[ee_mask] if not isData: pu = pu[ee_mask] events = events[ee_mask] # Stop processing if there is no event remain if len(Electron) == 0: return out cut3 = np.ones(len(Photon)) * 3 ##----------- Cut flow4: Photon Selection # Basic photon selection isgap_mask = (abs(Photon.eta) < 1.442) | ((abs(Photon.eta) > 1.566) & (abs(Photon.eta) < 2.5)) Pixel_seed_mask = ~Photon.pixelSeed if (dataset == "ZZ") and (self._year == "2017"): PT_ID_mask = (Photon.pt >= 20) & ( Photon.cutBasedBitmap >= 3 ) # 2^0(Loose) + 2^1(Medium) + 2^2(Tights) else: PT_ID_mask = (Photon.pt >= 20) & (Photon.cutBased > 1) # dR cut with selected Muon and Electrons dr_pho_ele_mask = ak.all(Photon.metric_table(Electron) >= 0.5, axis=-1) # default metric table: delta_r dr_pho_mu_mask = ak.all(Photon.metric_table(Muon) >= 0.5, axis=-1) # genPartFlav cut """ if dataset == "WZG": isPrompt = (Photon.genPartFlav == 1) | (Photon.genPartFlav == 11) PhoSelmask = PT_ID_mask & isgap_mask & Pixel_seed_mask & isPrompt & dr_pho_ele_mask & dr_pho_mu_mask elif dataset == "WZ": isPrompt = (Photon.genPartFlav == 1) PhoSelmask = PT_ID_mask & isgap_mask & Pixel_seed_mask & ~isPrompt & dr_pho_ele_mask & dr_pho_mu_mask else: PhoSelmask = PT_ID_mask & isgap_mask & Pixel_seed_mask & dr_pho_ele_mask & dr_pho_mu_mask """ # Add genPartFlav to remove Fake Photon in MC samples ( They are already considered by data driven method ) if not isData: genPartFlav_mask = (Photon.genPartFlav == 1) PhoSelmask = (genPartFlav_mask & PT_ID_mask & isgap_mask & Pixel_seed_mask & dr_pho_ele_mask & dr_pho_mu_mask) else: PhoSelmask = (PT_ID_mask & isgap_mask & Pixel_seed_mask & dr_pho_ele_mask & dr_pho_mu_mask) Photon = Photon[PhoSelmask] # Apply cut 4 A_photon_mask = ak.num(Photon) > 0 Electron = Electron[A_photon_mask] Photon = Photon[A_photon_mask] Jet = Jet[A_photon_mask] Muon = Muon[A_photon_mask] MET = MET[A_photon_mask] if not isData: pu = pu[A_photon_mask] events = events[A_photon_mask] # Stop processing if there is no event remain if len(Electron) == 0: return out def make_leading_pair(target, base): return target[ak.argmax(base.pt, axis=1, keepdims=True)] leading_pho = make_leading_pair(Photon, Photon) # -------------------- Make Fake Photon BKGs---------------------------# def make_bins(pt, eta, bin_range_str): bin_dict = { "PT_1_eta_1": (pt > 20) & (pt < 30) & (eta < 1), "PT_1_eta_2": (pt > 20) & (pt < 30) & (eta > 1) & (eta < 1.5), "PT_1_eta_3": (pt > 20) & (pt < 30) & (eta > 1.5) & (eta < 2), "PT_1_eta_4": (pt > 20) & (pt < 30) & (eta > 2) & (eta < 2.5), "PT_2_eta_1": (pt > 30) & (pt < 40) & (eta < 1), "PT_2_eta_2": (pt > 30) & (pt < 40) & (eta > 1) & (eta < 1.5), "PT_2_eta_3": (pt > 30) & (pt < 40) & (eta > 1.5) & (eta < 2), "PT_2_eta_4": (pt > 30) & (pt < 40) & (eta > 2) & (eta < 2.5), "PT_3_eta_1": (pt > 40) & (pt < 50) & (eta < 1), "PT_3_eta_2": (pt > 40) & (pt < 50) & (eta > 1) & (eta < 1.5), "PT_3_eta_3": (pt > 40) & (pt < 50) & (eta > 1.5) & (eta < 2), "PT_3_eta_4": (pt > 40) & (pt < 50) & (eta > 2) & (eta < 2.5), "PT_4_eta_1": (pt > 50) & (eta < 1), "PT_4_eta_2": (pt > 50) & (eta > 1) & (eta < 1.5), "PT_4_eta_3": (pt > 50) & (eta > 1.5) & (eta < 2), "PT_4_eta_4": (pt > 50) & (eta > 2) & (eta < 2.5), } binmask = bin_dict[bin_range_str] return binmask bin_name_list = [ "PT_1_eta_1", "PT_1_eta_2", "PT_1_eta_3", "PT_1_eta_4", "PT_2_eta_1", "PT_2_eta_2", "PT_2_eta_3", "PT_2_eta_4", "PT_3_eta_1", "PT_3_eta_2", "PT_3_eta_3", "PT_3_eta_4", "PT_4_eta_1", "PT_4_eta_2", "PT_4_eta_3", "PT_4_eta_4", ] ## -- Fake-fraction Lookup table --## if isFake: # Make Bin-range mask binned_pteta_mask = {} for name in bin_name_list: binned_pteta_mask[name] = make_bins( ak.flatten(leading_pho.pt), ak.flatten(abs(leading_pho.eta)), name, ) # Read Fake fraction --> Mapping bin name to int() if self._year == "2018": in_dict = np.load("Fitting_2018/Fit_results.npy", allow_pickle="True")[()] if self._year == "2017": in_dict = np.load("Fitting_2017/Fit_results.npy", allow_pickle="True")[()] idx = 0 fake_dict = {} for i, j in in_dict.items(): fake_dict[idx] = j idx += 1 # Reconstruct Fake_weight fw = 0 for i, j in binned_pteta_mask.items(): fw = fw + j * fake_dict[bin_name_list.index(i)] # Process 0 weight to 1 @numba.njit def zero_one(x): if x == 0: x = 1 return x vec_zero_one = np.vectorize(zero_one) fw = vec_zero_one(fw) else: fw = np.ones(len(events)) cut4 = np.ones(len(Photon)) * 4 print("Fake fraction weight: ", len(fw), len(cut4), fw) ##----------- Cut flow5: OSSF ossf_mask = Electron.charge[:, 0] + Electron.charge[:, 1] == 0 # Apply cut 5 Electron = Electron[ossf_mask] Photon = Photon[ossf_mask] fw = fw[ossf_mask] Jet = Jet[ossf_mask] MET = MET[ossf_mask] Muon = Muon[ossf_mask] if not isData: pu = pu[ossf_mask] events = events[ossf_mask] # Stop processing if there is no event remain if len(Electron) == 0: return out cut5 = np.ones(ak.sum(ak.num(Electron) > 0)) * 5 # Define Electron Triplet Diele = ak.zip({ "lep1": Electron[:, 0], "lep2": Electron[:, 1], "p4": TLorentz_vector(Electron[:, 0] + Electron[:, 1]), }) leading_ele = Diele.lep1 subleading_ele = Diele.lep2 def make_leading_pair(target, base): return target[ak.argmax(base.pt, axis=1, keepdims=True)] leading_pho = make_leading_pair(Photon, Photon) # -- Scale Factor for each electron # Trigger weight helper function 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: ## -------------< Egamma ID and Reco Scale factor > -----------------## get_pho_medium_id_sf = get_pho_medium_id_sf( ak.flatten(leading_pho.eta), ak.flatten(leading_pho.pt)) ele_reco_sf = get_ele_reco_above20_sf( leading_ele.deltaEtaSC + leading_ele.eta, leading_ele.pt, ) * get_ele_reco_above20_sf( subleading_ele.deltaEtaSC + subleading_ele.eta, subleading_ele.pt, ) ele_medium_id_sf = get_ele_medium_id_sf( leading_ele.deltaEtaSC + leading_ele.eta, leading_ele.pt, ) * get_ele_medium_id_sf( subleading_ele.deltaEtaSC + subleading_ele.eta, subleading_ele.pt, ) ## -------------< Muon ID and Iso Scale factor > -----------------## get_mu_tight_id_sf = get_mu_tight_id_sf(ak.flatten(abs(Muon.eta)), ak.flatten(Muon.pt)) get_mu_tight_iso_sf = get_mu_tight_iso_sf( ak.flatten(abs(Muon.eta)), ak.flatten(Muon.pt)) ## -------------< Double Electron Trigger Scale factor > -----------------## eta1 = leading_ele.deltaEtaSC + leading_ele.eta eta2 = subleading_ele.deltaEtaSC + subleading_ele.eta pt1 = leading_ele.pt pt2 = subleading_ele.pt # -- 2017,2016 are not applied yet if self._year == "2018": ele_trig_weight = Trigger_Weight(eta1, pt1, eta2, pt2) ##----------- Cut flow6: Baseline selection # Mee cut Mee_cut_mask = Diele.p4.mass > 4 # Lepton PT cuts Leppt_mask = ak.firsts((Diele.lep1.pt >= 25) & (Diele.lep2.pt >= 20) & (Muon.pt >= 25)) # MET cuts MET_mask = MET.pt > 20 # Baseline # Assemble!! Baseline_mask = Leppt_mask & MET_mask & Mee_cut_mask # SR,CR # Apply cut6 Diele_base = Diele[Baseline_mask] leading_pho_base = leading_pho[Baseline_mask] Jet_base = Jet[Baseline_mask] MET_base = MET[Baseline_mask] Muon_base = Muon[Baseline_mask] events_base = events[Baseline_mask] # Photon EE and EB isEE_mask = leading_pho.isScEtaEE isEB_mask = leading_pho.isScEtaEB Pho_EE_base = leading_pho[isEE_mask & Baseline_mask] Pho_EB_base = leading_pho[isEB_mask & Baseline_mask] # Stop processing if there is no event remain if len(leading_pho_base) == 0: return out cut6 = np.ones(ak.sum(ak.num(leading_pho_base) > 0)) * 6 base_arr_dict = { "Diele_sel": Diele_base, "leading_pho_sel": leading_pho_base, "Jet_sel": Jet_base, "MET_sel": MET_base, "Muon_sel": Muon_base, "Pho_EE_sel": Pho_EE_base, "Pho_EB_sel": Pho_EB_base, } ##----------- << SR >> Zmass_window_mask = abs(Diele.p4.mass - 91.1876) < 15 MET_mask = MET.pt > 30 bjet_veto = ak.firsts(Jet.btagDeepB > 0.7665) == 0 Mlll_mask = ((Diele.p4 + Muon[:, 0]).mass) > 100 SR_mask = Zmass_window_mask & MET_mask & bjet_veto & Mlll_mask SR_mask = Baseline_mask & SR_mask Diele_SR = Diele[SR_mask] leading_pho_SR = leading_pho[SR_mask] Muon_SR = Muon[SR_mask] MET_SR = MET[SR_mask] Jet_SR = Jet[SR_mask] events_SR = events[SR_mask] Pho_EE_SR = leading_pho[isEE_mask & SR_mask] Pho_EB_SR = leading_pho[isEB_mask & SR_mask] SR_arr_dict = { "Diele_sel": Diele_SR, "leading_pho_sel": leading_pho_SR, "Jet_sel": Jet_SR, "MET_sel": MET_SR, "Muon_sel": Muon_SR, "Pho_EE_sel": Pho_EE_SR, "Pho_EB_sel": Pho_EB_SR, } ##----------- << CR-Z+Jets >> Zmass_window_mask = abs(Diele.p4.mass - 91.1876) < 15 MET_mask = MET.pt <= 30 bjet_veto = ak.firsts(Jet.btagDeepB > 0.7665) == 0 Mlll_mask = ((Diele.p4 + Muon[:, 0]).mass) > 100 CR_ZJets_mask = Zmass_window_mask & MET_mask & bjet_veto & Mlll_mask CR_ZJets_mask = Baseline_mask & CR_ZJets_mask Diele_CR_ZJets = Diele[CR_ZJets_mask] leading_pho_CR_ZJets = leading_pho[CR_ZJets_mask] Muon_CR_ZJets = Muon[CR_ZJets_mask] MET_CR_ZJets = MET[CR_ZJets_mask] Jet_CR_ZJets = Jet[CR_ZJets_mask] events_CR_ZJets = events[CR_ZJets_mask] Pho_EE_CR_ZJets = leading_pho[isEE_mask & CR_ZJets_mask] Pho_EB_CR_ZJets = leading_pho[isEB_mask & CR_ZJets_mask] CR_ZJets_arr_dict = { "Diele_sel": Diele_CR_ZJets, "leading_pho_sel": leading_pho_CR_ZJets, "Jet_sel": Jet_CR_ZJets, "MET_sel": MET_CR_ZJets, "Muon_sel": Muon_CR_ZJets, "Pho_EE_sel": Pho_EE_CR_ZJets, "Pho_EB_sel": Pho_EB_CR_ZJets, } ##----------- << CR-T-enriched >> Zmass_window_mask = abs(Diele.p4.mass - 91.1876) > 5 MET_mask = MET.pt > 30 bjet_veto = ak.firsts(Jet.btagDeepB > 0.7665) > 0 Mlll_mask = ((Diele.p4 + Muon[:, 0]).mass) > 100 CR_Tenri_mask = Zmass_window_mask & MET_mask & bjet_veto & Mlll_mask CR_Tenri_mask = Baseline_mask & CR_Tenri_mask Diele_CR_t = Diele[CR_Tenri_mask] leading_pho_CR_t = leading_pho[CR_Tenri_mask] Muon_CR_t = Muon[CR_Tenri_mask] MET_CR_t = MET[CR_Tenri_mask] Jet_CR_t = Jet[CR_Tenri_mask] events_CR_t = events[CR_Tenri_mask] Pho_EE_CR_t = leading_pho[isEE_mask & CR_Tenri_mask] Pho_EB_CR_t = leading_pho[isEB_mask & CR_Tenri_mask] CR_tEnriched_arr_dict = { "Diele_sel": Diele_CR_t, "leading_pho_sel": leading_pho_CR_t, "Jet_sel": Jet_CR_t, "MET_sel": MET_CR_t, "Muon_sel": Muon_CR_t, "Pho_EE_sel": Pho_EE_CR_t, "Pho_EB_sel": Pho_EB_CR_t, } ##----------- << CR-Conversion >> Zmass_window_mask = abs(Diele.p4.mass - 91.1876) > 15 MET_mask = MET.pt <= 30 bjet_veto = ak.firsts(Jet.btagDeepB > 0.7665) == 0 Mlll_mask = ((Diele.p4 + Muon[:, 0]).mass) <= 100 CR_conv_mask = Zmass_window_mask & MET_mask & bjet_veto & Mlll_mask CR_conv_mask = Baseline_mask & CR_conv_mask Diele_CR_conv = Diele[CR_conv_mask] leading_pho_CR_conv = leading_pho[CR_conv_mask] Muon_CR_conv = Muon[CR_conv_mask] MET_CR_conv = MET[CR_conv_mask] Jet_CR_conv = Jet[CR_conv_mask] events_CR_conv = events[CR_conv_mask] Pho_EE_CR_conv = leading_pho[isEE_mask & CR_conv_mask] Pho_EB_CR_conv = leading_pho[isEB_mask & CR_conv_mask] CR_Conversion_dict = { "Diele_sel": Diele_CR_conv, "leading_pho_sel": leading_pho_CR_conv, "Jet_sel": Jet_CR_conv, "MET_sel": MET_CR_conv, "Muon_sel": Muon_CR_conv, "Pho_EE_sel": Pho_EE_CR_conv, "Pho_EB_sel": Pho_EB_CR_conv, } ## -------------------- Prepare making hist --------------# regions = { "Baseline": base_arr_dict, "Signal": SR_arr_dict, "CR_ZJets": CR_ZJets_arr_dict, "CR_tEnriched": CR_tEnriched_arr_dict, "CR_conversion": CR_Conversion_dict, } mask_dict = { "Baseline": Baseline_mask, "Signal": SR_mask, "CR_ZJets": CR_ZJets_mask, "CR_tEnriched": CR_Tenri_mask, "CR_conversion": CR_conv_mask, } for region, arr_dict in regions.items(): # Photon phoPT = ak.flatten(arr_dict["leading_pho_sel"].pt) phoEta = ak.flatten(arr_dict["leading_pho_sel"].eta) phoPhi = ak.flatten(arr_dict["leading_pho_sel"].phi) # Photon EE if len(arr_dict["Pho_EE_sel"].pt) != 0: Pho_EE_PT = ak.flatten(arr_dict["Pho_EE_sel"].pt) Pho_EE_Eta = ak.flatten(arr_dict["Pho_EE_sel"].eta) Pho_EE_Phi = ak.flatten(arr_dict["Pho_EE_sel"].phi) Pho_EE_sieie = ak.flatten(arr_dict["Pho_EE_sel"].sieie) Pho_EE_Iso_charge = ak.flatten( arr_dict["Pho_EE_sel"].pfRelIso03_chg) # Photon EB if len(arr_dict["Pho_EB_sel"].pt) != 0: Pho_EB_PT = ak.flatten(arr_dict["Pho_EB_sel"].pt) Pho_EB_Eta = ak.flatten(arr_dict["Pho_EB_sel"].eta) Pho_EB_Phi = ak.flatten(arr_dict["Pho_EB_sel"].phi) Pho_EB_sieie = ak.flatten(arr_dict["Pho_EB_sel"].sieie) Pho_EB_Iso_charge = ak.flatten( arr_dict["Pho_EB_sel"].pfRelIso03_chg) # Electrons ele1PT = arr_dict["Diele_sel"].lep1.pt ele1Eta = arr_dict["Diele_sel"].lep1.eta ele1Phi = arr_dict["Diele_sel"].lep1.phi ele2PT = arr_dict["Diele_sel"].lep2.pt ele2Eta = arr_dict["Diele_sel"].lep2.eta ele2Phi = arr_dict["Diele_sel"].lep2.phi # Muon muPT = ak.flatten(arr_dict["Muon_sel"].pt) muEta = ak.flatten(arr_dict["Muon_sel"].eta) muPhi = ak.flatten(arr_dict["Muon_sel"].phi) # MET met = ak.to_numpy(arr_dict["MET_sel"].pt) # M(eea) M(ee) diele = arr_dict["Diele_sel"].p4 lll_vec = diele + arr_dict["Muon_sel"][:, 0] Mlll = lll_vec.mass Mee = diele.mass # W MT (--> beta) MT = np.sqrt( 2 * arr_dict["Muon_sel"].pt * arr_dict["MET_sel"].pt * (1 - np.cos( abs(arr_dict["MET_sel"].delta_phi(arr_dict["Muon_sel"]))))) MT = np.array(ak.firsts(MT)) # --- Apply weight and hist weights = processor.Weights(len(cut5)) # --- 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 = mask_dict[region] cuts_pho_EE = ak.flatten(isEE_mask) cuts_pho_EB = ak.flatten(isEB_mask) if isFake: weights.add("fake_fraction", fw) # Weight and SF here if not (isData | isFake): weights.add("pileup", pu) weights.add("ele_id", ele_medium_id_sf) weights.add("ele_reco", ele_reco_sf) weights.add("pho_id", get_pho_medium_id_sf) weights.add("mu_id", get_mu_tight_id_sf) weights.add("mu_iso", get_mu_tight_id_sf) # 2016,2017 are not applied yet if self._year == "2018": weights.add("ele_trigger", ele_trig_weight) # ---------------------------- Fill hist --------------------------------------# # Initial events out["sumw"][dataset] += len(Initial_events) print( "region: {0} ### cut0: {1},cut1: {2}, cut2: {3},cut3: {4},cut4: {5},cut5: {6},cut6: {7}, cut7: {8}" .format(region, len(cut0), len(cut1), len(cut2), len(cut3), len(cut4), len(cut5), len(cut6), len(met))) # Fill hist # -- met -- # out["met"].fill( dataset=dataset, region=region, met=met, weight=skim_weight(weights.weight() * cuts), ) # --mass -- # out["MT"].fill( dataset=dataset, region=region, MT=MT, weight=skim_weight(weights.weight() * cuts), ) out["mass"].fill( dataset=dataset, region=region, mass=Mee, weight=skim_weight(weights.weight() * cuts), ) out["mass_lll"].fill( dataset=dataset, region=region, mass_lll=Mlll, weight=skim_weight(weights.weight() * cuts), ) # -- Muon -- # out["mupt"].fill( dataset=dataset, region=region, mupt=muPT, weight=skim_weight(weights.weight() * cuts), ) out["mueta"].fill( dataset=dataset, region=region, mueta=muEta, weight=skim_weight(weights.weight() * cuts), ) out["muphi"].fill( dataset=dataset, region=region, muphi=muPhi, weight=skim_weight(weights.weight() * cuts), ) # -- Electron -- # out["ele1pt"].fill( dataset=dataset, region=region, ele1pt=ele1PT, weight=skim_weight(weights.weight() * cuts), ) out["ele1eta"].fill( dataset=dataset, region=region, ele1eta=ele1Eta, weight=skim_weight(weights.weight() * cuts), ) out["ele1phi"].fill( dataset=dataset, region=region, ele1phi=ele1Phi, weight=skim_weight(weights.weight() * cuts), ) out["ele2pt"].fill( dataset=dataset, region=region, ele2pt=ele2PT, weight=skim_weight(weights.weight() * cuts), ) out["ele2eta"].fill( dataset=dataset, region=region, ele2eta=ele2Eta, weight=skim_weight(weights.weight() * cuts), ) out["ele2phi"].fill( dataset=dataset, region=region, ele2phi=ele2Phi, weight=skim_weight(weights.weight() * cuts), ) # -- Photon -- # out["phopt"].fill( dataset=dataset, region=region, phopt=phoPT, weight=skim_weight(weights.weight() * cuts), ) out["phoeta"].fill( dataset=dataset, region=region, phoeta=phoEta, weight=skim_weight(weights.weight() * cuts), ) out["phophi"].fill( dataset=dataset, region=region, phophi=phoPhi, weight=skim_weight(weights.weight() * cuts), ) if len(arr_dict["Pho_EE_sel"].pt) != 0: out["pho_EE_pt"].fill( dataset=dataset, region=region, pho_EE_pt=Pho_EE_PT, weight=skim_weight(weights.weight() * cuts * cuts_pho_EE), ) out["pho_EE_eta"].fill( dataset=dataset, region=region, pho_EE_eta=Pho_EE_Eta, weight=skim_weight(weights.weight() * cuts * cuts_pho_EE), ) out["pho_EE_phi"].fill( dataset=dataset, region=region, pho_EE_phi=Pho_EE_Phi, weight=skim_weight(weights.weight() * cuts * cuts_pho_EE), ) out["pho_EE_sieie"].fill( dataset=dataset, region=region, pho_EE_sieie=Pho_EE_sieie, weight=skim_weight(weights.weight() * cuts * cuts_pho_EE), ) out["pho_EE_Iso_chg"].fill( dataset=dataset, region=region, pho_EE_Iso_chg=Pho_EE_Iso_charge, weight=skim_weight(weights.weight() * cuts * cuts_pho_EE), ) if len(arr_dict["Pho_EB_sel"].pt) != 0: out["pho_EB_pt"].fill( dataset=dataset, region=region, pho_EB_pt=Pho_EB_PT, weight=skim_weight(weights.weight() * cuts * cuts_pho_EB), ) out["pho_EB_eta"].fill( dataset=dataset, region=region, pho_EB_eta=Pho_EB_Eta, weight=skim_weight(weights.weight() * cuts * cuts_pho_EB), ) out["pho_EB_phi"].fill( dataset=dataset, region=region, pho_EB_phi=Pho_EB_Phi, weight=skim_weight(weights.weight() * cuts * cuts_pho_EB), ) out["pho_EB_sieie"].fill( dataset=dataset, region=region, pho_EB_sieie=Pho_EB_sieie, weight=skim_weight(weights.weight() * cuts * cuts_pho_EB), ) out["pho_EB_Iso_chg"].fill( dataset=dataset, region=region, pho_EB_Iso_chg=Pho_EB_Iso_charge, weight=skim_weight(weights.weight() * cuts * cuts_pho_EB), ) return out
def get_hpt(): boson = ak.firsts( genpart[(genpart.pdgId == 25) & genpart.hasFlags(["fromHardProcess", "isLastCopy"])]) return np.array(ak.fill_none(boson.pt, 0.))
def process(self, events): # Initialize accumulator out = self.accumulator.identity() dataset = sample_name #events.metadata['dataset'] # Data or MC isData = 'genWeight' not in events.fields #Stop processing if there is no event remain if len(events) == 0: return out # Golden Json file if (self._year == "2018") and isData: injson = "/x5/cms/jwkim/gitdir/JWCorp/JW_analysis/Coffea_WZG/Corrections/Cert_314472-325175_13TeV_Legacy2018_Collisions18_JSON.txt.RunABCD" if (self._year == "2017") and isData: injson = "/x5/cms/jwkim/gitdir/JWCorp/JW_analysis/Coffea_WZG/Corrections/Cert_294927-306462_13TeV_UL2017_Collisions17_GoldenJSON.txt" # <----- Get Scale factors ------># if not isData: # Egamma reco ID get_ele_reco_above20_sf = self._corrections[ 'get_ele_reco_above20_sf'][self._year] get_ele_medium_id_sf = self._corrections['get_ele_medium_id_sf'][ self._year] get_pho_medium_id_sf = self._corrections['get_pho_medium_id_sf'][ self._year] # DoubleEG trigger # 2016, 2017 are not applied yet if self._year == "2018": 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(events.Pileup.nTrueInt, "int64") pu = self._puweight_arr[pu_weight_idx] selection = processor.PackedSelection() # Cut flow cut0 = np.zeros(len(events)) # <----- Helper functions ------># # Sort by PT helper function def sort_by_pt(ele, pho, jet): ele = ele[ak.argsort(ele.pt, ascending=False, axis=1)] pho = pho[ak.argsort(pho.pt, ascending=False, axis=1)] jet = jet[ak.argsort(jet.pt, ascending=False, axis=1)] return ele, pho, jet # Lorentz vectors from coffea.nanoevents.methods import vector ak.behavior.update(vector.behavior) def TLorentz_vector(vec): vec = ak.zip({ "x": vec.x, "y": vec.y, "z": vec.z, "t": vec.t }, with_name="LorentzVector") return vec def TLorentz_vector_cylinder(vec): vec = ak.zip( { "pt": vec.pt, "eta": vec.eta, "phi": vec.phi, "mass": vec.mass, }, with_name="PtEtaPhiMLorentzVector", ) return vec # Cut-based ID modification @numba.njit def PhotonVID(vid, idBit): rBit = 0 for x in range(0, 7): rBit |= (1 << x) if ((vid >> (x * 2)) & 0b11 >= idBit) else 0 return rBit # Inverse Sieie and upper limit @numba.njit def make_fake_obj_mask(Pho, builder): #for eventIdx,pho in enumerate(tqdm(Pho)): # --Event Loop for eventIdx, pho in enumerate(Pho): builder.begin_list() if len(pho) < 1: continue for phoIdx, _ in enumerate(pho): # --Photon Loop vid = Pho[eventIdx][phoIdx].vidNestedWPBitmap vid_cuts1 = PhotonVID(vid, 1) # Loose photon vid_cuts2 = PhotonVID(vid, 2) # Medium photon vid_cuts3 = PhotonVID(vid, 3) # Tight photon # Field name # |0|0|0|0|0|0|0| # |IsoPho|IsoNeu|IsoChg|Sieie|hoe|scEta|PT| # 1. Turn off cut (ex turn off Sieie # |1|1|1|0|1|1|1| = |1|1|1|0|1|1|1| # 2. Inverse cut (ex inverse Sieie) # |1|1|1|1|1|1|1| = |1|1|1|0|1|1|1| #if (vid_cuts2 & 0b1111111 == 0b1111111): # Cut applied #if (vid_cuts2 & 0b1111111 == 0b1110111): # Inverse Sieie if (vid_cuts2 & 0b1110111 == 0b1110111): # Without Sieie builder.boolean(True) else: builder.boolean(False) builder.end_list() return builder # <----- Selection ------># Initial_events = events # Good Run ( Golden Json files ) from coffea import lumi_tools if isData: lumi_mask_builder = lumi_tools.LumiMask(injson) lumimask = ak.Array( lumi_mask_builder.__call__(events.run, events.luminosityBlock)) events = events[lumimask] #print("{0}% of files pass good-run conditions".format(len(events)/ len(Initial_events))) # Stop processing if there is no event remain if len(events) == 0: return out ##----------- Cut flow1: Passing Triggers # 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] events.Electron, events.Photon, events.Jet = sort_by_pt( events.Electron, events.Photon, events.Jet) # Good Primary vertex nPV = events.PV.npvsGood if not isData: nPV = nPV * pu nPV_nw = nPV # Apply cut1 events = events[double_ele_triggers_arr] if not isData: pu = pu[double_ele_triggers_arr] cut1 = np.ones(len(events)) # Set Particles Electron = events.Electron Muon = events.Muon Photon = events.Photon MET = events.MET Jet = events.Jet # Stop processing if there is no event remain if len(Electron) == 0: return out # --Gen Photon for dR genparts = events.GenPart pdgID_mask = (genparts.pdgId == 22) # mask2: isPrompt | fromHardProcess | isLastCopy mask2 = (1 << 0) | (1 << 8) | (1 << 13) # https://github.com/PKUHEPEWK/WGamma/blob/master/2018/wgRealPhotonTemplateModule.py status_mask = ((genparts.statusFlags & mask2) == mask2) gen_photons = genparts[pdgID_mask & status_mask] assert (ak.all(ak.num(gen_photons) == 1) ) # Raise error if len(gen_photon) != 1 # --Muon ( only used to calculate dR ) MuSelmask = (Muon.pt >= 10) & (abs( Muon.eta) <= 2.5) & (Muon.tightId) & (Muon.pfRelIso04_all < 0.15) Muon = Muon[MuSelmask] ##----------- Cut flow2: Electron Selection EleSelmask = ((Electron.pt >= 20) & (np.abs(Electron.eta + Electron.deltaEtaSC) < 1.479) & (Electron.cutBased > 2) & (abs(Electron.dxy) < 0.05) & (abs(Electron.dz) < 0.1)) | \ ((Electron.pt >= 20) & (np.abs(Electron.eta + Electron.deltaEtaSC) > 1.479) & (np.abs(Electron.eta + Electron.deltaEtaSC) <= 2.5) & (Electron.cutBased > 2) & (abs(Electron.dxy) < 0.1) & (abs(Electron.dz) < 0.2)) Electron = Electron[EleSelmask] # apply cut 2 Tri_electron_mask = ak.num(Electron) >= 2 Electron = Electron[Tri_electron_mask] Photon = Photon[Tri_electron_mask] Jet = Jet[Tri_electron_mask] MET = MET[Tri_electron_mask] Muon = Muon[Tri_electron_mask] if not isData: pu = pu[Tri_electron_mask] events = events[Tri_electron_mask] gen_photons = gen_photons[Tri_electron_mask] # Stop processing if there is no event remain if len(Electron) == 0: return out cut2 = np.ones(len(Photon)) * 2 ##----------- Cut flow3: Photon Selection # Basic photon selection isgap_mask = (abs(Photon.eta) < 1.442) | ((abs(Photon.eta) > 1.566) & (abs(Photon.eta) < 2.5)) Pixel_seed_mask = ~Photon.pixelSeed PT_mask = Photon.pt >= 20 # dR cut with selected Muon and Electrons dr_pho_ele_mask = ak.all(Photon.metric_table(Electron) >= 0.5, axis=-1) # default metric table: delta_r dr_pho_mu_mask = ak.all(Photon.metric_table(Muon) >= 0.5, axis=-1) PhoSelmask = PT_mask & isgap_mask & Pixel_seed_mask & dr_pho_ele_mask & dr_pho_mu_mask Photon = Photon[PhoSelmask] # Apply cut 3 A_photon_mask = ak.num(Photon) > 0 Electron = Electron[A_photon_mask] Photon = Photon[A_photon_mask] Jet = Jet[A_photon_mask] Muon = Muon[A_photon_mask] MET = MET[A_photon_mask] if not isData: pu = pu[A_photon_mask] events = events[A_photon_mask] gen_photons = gen_photons[A_photon_mask] Photon_template_mask = make_fake_obj_mask( Photon, ak.ArrayBuilder()).snapshot() Photon = Photon[Photon_template_mask] # Apply cut 3 A_photon_mask = ak.num(Photon) > 0 Electron = Electron[A_photon_mask] Photon = Photon[A_photon_mask] Jet = Jet[A_photon_mask] Muon = Muon[A_photon_mask] MET = MET[A_photon_mask] if not isData: pu = pu[A_photon_mask] events = events[A_photon_mask] gen_photons = gen_photons[A_photon_mask] # Stop processing if there is no event remain if len(Electron) == 0: return out cut3 = np.ones(len(Photon)) * 3 ## -- Additional photon selection: Photon gen-matching # Choose Photons that dR(genPhoton,Photon) <= 0.1 gen_match_photon_mask = ak.all(Photon.metric_table(gen_photons) <= 0.1, axis=-1) # Apply cut Photon = Photon[gen_match_photon_mask] gen_match_photon_evt_mask = ak.num(Photon) >= 1 Electron = Electron[gen_match_photon_evt_mask] Photon = Photon[gen_match_photon_evt_mask] Jet = Jet[gen_match_photon_evt_mask] MET = MET[gen_match_photon_evt_mask] gen_photons = gen_photons[gen_match_photon_evt_mask] if not isData: pu = pu[gen_match_photon_evt_mask] events = events[gen_match_photon_evt_mask] ##----------- Cut flow4: Select 2 OSSF electrons from Z @numba.njit def find_2lep(events_leptons, builder): for leptons in events_leptons: builder.begin_list() nlep = len(leptons) for i0 in range(nlep): for i1 in range(i0 + 1, nlep): if leptons[i0].charge + leptons[i1].charge != 0: continue if nlep == 2: builder.begin_tuple(2) builder.index(0).integer(i0) builder.index(1).integer(i1) builder.end_tuple() else: for i2 in range(nlep): if len({i0, i1, i2}) < 3: continue builder.begin_tuple(3) builder.index(0).integer(i0) builder.index(1).integer(i1) builder.index(2).integer(i2) builder.end_tuple() builder.end_list() return builder ossf_idx = find_2lep(Electron, ak.ArrayBuilder()).snapshot() # OSSF cut ossf_mask = ak.num(ossf_idx) >= 1 ossf_idx = ossf_idx[ossf_mask] Electron = Electron[ossf_mask] Photon = Photon[ossf_mask] Jet = Jet[ossf_mask] MET = MET[ossf_mask] events = events[ossf_mask] if not isData: pu = pu[ossf_mask] Double_electron = [Electron[ossf_idx[idx]] for idx in "01"] from coffea.nanoevents.methods import vector ak.behavior.update(vector.behavior) Diele = ak.zip({ "lep1": Double_electron[0], "lep2": Double_electron[1], "p4": TLorentz_vector(Double_electron[0] + Double_electron[1]) }) bestZ_idx = ak.singletons( ak.argmin(abs(Diele.p4.mass - 91.1876), axis=1)) Diele = Diele[bestZ_idx] # Stop processing if there is no event remain if len(Electron) == 0: return out cut4 = np.ones(len(Electron)) * 4 leading_ele = Diele.lep1 subleading_ele = Diele.lep2 def make_leading_pair(target, base): return target[ak.argmax(base.pt, axis=1, keepdims=True)] leading_pho = make_leading_pair(Photon, Photon) # -- Scale Factor for each electron # Trigger weight helper function 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: ## -------------< Egamma ID and Reco Scale factor > -----------------## get_pho_medium_id_sf = get_pho_medium_id_sf( ak.flatten(leading_pho.eta), ak.flatten(leading_pho.pt)) ele_reco_sf = get_ele_reco_above20_sf( ak.flatten(leading_ele.deltaEtaSC + leading_ele.eta), ak.flatten(leading_ele.pt)) * get_ele_reco_above20_sf( ak.flatten(subleading_ele.deltaEtaSC + subleading_ele.eta), ak.flatten(subleading_ele.pt)) ele_medium_id_sf = get_ele_medium_id_sf( ak.flatten(leading_ele.deltaEtaSC + leading_ele.eta), ak.flatten(leading_ele.pt)) * get_ele_medium_id_sf( ak.flatten(subleading_ele.deltaEtaSC + subleading_ele.eta), ak.flatten(subleading_ele.pt)) ## -------------< Double Electron Trigger Scale factor > -----------------## 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) # -- 2017,2016 are not applied yet if self._year == '2018': ele_trig_weight = Trigger_Weight(eta1, pt1, eta2, pt2) ##----------- Cut flow5: Event selection # Mee cut Mee_cut_mask = ak.firsts(Diele.p4.mass) > 4 # Electron PT cuts Elept_mask = ak.firsts((Diele.lep1.pt >= 25) & (Diele.lep2.pt >= 20)) # MET cuts MET_mask = MET.pt > 20 # --------Mask -------# Event_sel_mask = Mee_cut_mask & Elept_mask & MET_mask Diele_sel = Diele[Event_sel_mask] leading_pho_sel = leading_pho[Event_sel_mask] Jet_sel = Jet[Event_sel_mask] MET_sel = MET[Event_sel_mask] # Photon EE and EB isEE_mask = leading_pho.isScEtaEE isEB_mask = leading_pho.isScEtaEB Pho_EE = leading_pho[isEE_mask & Event_sel_mask] Pho_EB = leading_pho[isEB_mask & Event_sel_mask] #Stop processing if there is no event remain if len(leading_pho_sel) == 0: return out cut5 = np.ones(len(Diele)) * 5 # -------------------- Flatten variables ---------------------------# # -- Ele1 --# Ele1_PT = ak.flatten(Diele_sel.lep1.pt) Ele1_Eta = ak.flatten(Diele_sel.lep1.eta) Ele1_Phi = ak.flatten(Diele_sel.lep1.phi) # -- Ele2 --# Ele2_PT = ak.flatten(Diele_sel.lep2.pt) Ele2_Eta = ak.flatten(Diele_sel.lep2.eta) Ele2_Phi = ak.flatten(Diele_sel.lep2.phi) # -- Pho -- # Pho_PT = ak.flatten(leading_pho_sel.pt) Pho_Eta = ak.flatten(leading_pho_sel.eta) Pho_Phi = ak.flatten(leading_pho_sel.phi) # -- Pho EB --# Pho_EB_PT = ak.flatten(Pho_EB.pt) Pho_EB_Eta = ak.flatten(Pho_EB.eta) Pho_EB_Phi = ak.flatten(Pho_EB.phi) Pho_EB_Isochg = ak.flatten(Pho_EE.pfRelIso03_chg) Pho_EB_Sieie = ak.flatten(Pho_EE.sieie) # -- Pho EE --# Pho_EE_PT = ak.flatten(Pho_EE.pt) Pho_EE_Eta = ak.flatten(Pho_EE.eta) Pho_EE_Phi = ak.flatten(Pho_EE.phi) Pho_EE_Isochg = ak.flatten(Pho_EE.pfRelIso03_chg) Pho_EE_Sieie = ak.flatten(Pho_EE.sieie) # --Kinematics --# Diele_mass = ak.flatten(Diele_sel.p4.mass) leading_ele, subleading_ele = ak.flatten( TLorentz_vector_cylinder(Diele_sel.lep1)), ak.flatten( TLorentz_vector_cylinder(Diele_sel.lep2)) dR_e1pho = ak.flatten( leading_ele.delta_r(leading_pho_sel)) # dR pho,ele1 dR_e2pho = ak.flatten( subleading_ele.delta_r(leading_pho_sel)) # dR pho,ele2 dR_jpho = ak.flatten(Jet_sel[:, 0].delta_r(leading_pho_sel)) MET_PT = ak.to_numpy(MET_sel.pt) # -------------------- Sieie bins---------------------------# def make_bins(pt, eta, sieie, bin_range_str): bin_dict = { 'PT_1_eta_1': (pt > 20) & (pt < 30) & (eta < 1), 'PT_1_eta_2': (pt > 20) & (pt < 30) & (eta > 1) & (eta < 1.5), 'PT_1_eta_3': (pt > 20) & (pt < 30) & (eta > 1.5) & (eta < 2), 'PT_1_eta_4': (pt > 20) & (pt < 30) & (eta > 2) & (eta < 2.5), 'PT_2_eta_1': (pt > 30) & (pt < 40) & (eta < 1), 'PT_2_eta_2': (pt > 30) & (pt < 40) & (eta > 1) & (eta < 1.5), 'PT_2_eta_3': (pt > 30) & (pt < 40) & (eta > 1.5) & (eta < 2), 'PT_2_eta_4': (pt > 30) & (pt < 40) & (eta > 2) & (eta < 2.5), 'PT_3_eta_1': (pt > 40) & (pt < 50) & (eta < 1), 'PT_3_eta_2': (pt > 40) & (pt < 50) & (eta > 1) & (eta < 1.5), 'PT_3_eta_3': (pt > 40) & (pt < 50) & (eta > 1.5) & (eta < 2), 'PT_3_eta_4': (pt > 40) & (pt < 50) & (eta > 2) & (eta < 2.5), 'PT_4_eta_1': (pt > 50) & (eta < 1), 'PT_4_eta_2': (pt > 50) & (eta > 1) & (eta < 1.5), 'PT_4_eta_3': (pt > 50) & (eta > 1.5) & (eta < 2), 'PT_4_eta_4': (pt > 50) & (eta > 2) & (eta < 2.5) } binmask = bin_dict[bin_range_str] return ak.to_numpy(sieie[binmask]), binmask bin_name_list = [ 'PT_1_eta_1', 'PT_1_eta_2', 'PT_1_eta_3', 'PT_1_eta_4', 'PT_2_eta_1', 'PT_2_eta_2', 'PT_2_eta_3', 'PT_2_eta_4', 'PT_3_eta_1', 'PT_3_eta_2', 'PT_3_eta_3', 'PT_3_eta_4', 'PT_4_eta_1', 'PT_4_eta_2', 'PT_4_eta_3', 'PT_4_eta_4' ] binned_sieie_hist = {} binmask_dict = {} for name in bin_name_list: binned_sieie_hist[name], _ = make_bins( ak.flatten(leading_pho_sel.pt), ak.flatten(abs(leading_pho_sel.eta)), ak.flatten(leading_pho_sel.sieie), name) _, binmask_dict[name] = make_bins(ak.flatten(leading_pho.pt), ak.flatten(abs(leading_pho.eta)), ak.flatten(leading_pho.sieie), name) print("Show me the last bin: ", binned_sieie_hist['PT_4_eta_4']) # --- Apply weight and hist weights = processor.Weights(len(cut4)) # --- 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 = Event_sel_mask cuts_pho_EE = ak.flatten(isEE_mask) cuts_pho_EB = ak.flatten(isEB_mask) print( "cut0: {0}, cut1: {1}, cut2: {2}, cut3: {3}, cut4: {4} ,cut5 {5} ". format(len(Initial_events), len(cut1), len(cut2), len(cut3), len(cut4), len(cut5))) # Weight and SF here if not isData: weights.add('pileup', pu) weights.add('ele_id', ele_medium_id_sf) weights.add('pho_id', get_pho_medium_id_sf) weights.add('ele_reco', ele_reco_sf) # 2016,2017 are not applied yet if self._year == "2018": weights.add('ele_trigger', ele_trig_weight) # ---------------------------- Fill hist --------------------------------------# # Initial events out["sumw"][dataset] += len(Initial_events) # Cut flow loop for cut in [cut0, cut1, cut2, cut3, cut4, cut5]: out["cutflow"].fill(dataset=dataset, cutflow=cut) # Primary vertex out['nPV'].fill( dataset=dataset, nPV=nPV, ) out['nPV_nw'].fill(dataset=dataset, nPV_nw=nPV_nw) # Fill hist # -- met -- # out["met"].fill(dataset=dataset, met=MET_PT, weight=skim_weight(weights.weight() * cuts)) # --mass -- # out["mass"].fill(dataset=dataset, mass=Diele_mass, weight=skim_weight(weights.weight() * cuts)) # -- Ele1 -- # out["ele1pt"].fill(dataset=dataset, ele1pt=Ele1_PT, weight=skim_weight(weights.weight() * cuts)) out["ele1eta"].fill(dataset=dataset, ele1eta=Ele1_Eta, weight=skim_weight(weights.weight() * cuts)) out["ele1phi"].fill(dataset=dataset, ele1phi=Ele1_Phi, weight=skim_weight(weights.weight() * cuts)) # --Ele2 --# out["ele2pt"].fill(dataset=dataset, ele2pt=Ele2_PT, weight=skim_weight(weights.weight() * cuts)) out["ele2eta"].fill(dataset=dataset, ele2eta=Ele2_Eta, weight=skim_weight(weights.weight() * cuts)) out["ele2phi"].fill(dataset=dataset, ele2phi=Ele2_Phi, weight=skim_weight(weights.weight() * cuts)) # -- Photon -- # out["phopt"].fill(dataset=dataset, phopt=Pho_PT, weight=skim_weight(weights.weight() * cuts)) out["phoeta"].fill(dataset=dataset, phoeta=Pho_Eta, weight=skim_weight(weights.weight() * cuts)) out["phophi"].fill(dataset=dataset, phophi=Pho_Phi, weight=skim_weight(weights.weight() * cuts)) # -- Binned sieie hist -- # if len(binned_sieie_hist['PT_1_eta_1'] > 0): out['PT_1_eta_1'].fill(dataset=dataset, PT_1_eta_1=binned_sieie_hist['PT_1_eta_1']) if len(binned_sieie_hist['PT_1_eta_2'] > 0): out['PT_1_eta_2'].fill(dataset=dataset, PT_1_eta_2=binned_sieie_hist['PT_1_eta_2']) if len(binned_sieie_hist['PT_1_eta_3'] > 0): out['PT_1_eta_3'].fill(dataset=dataset, PT_1_eta_3=binned_sieie_hist['PT_1_eta_3']) if len(binned_sieie_hist['PT_1_eta_4'] > 0): out['PT_1_eta_4'].fill(dataset=dataset, PT_1_eta_4=binned_sieie_hist['PT_1_eta_4']) if len(binned_sieie_hist['PT_2_eta_1'] > 0): out['PT_2_eta_1'].fill(dataset=dataset, PT_2_eta_1=binned_sieie_hist['PT_2_eta_1']) if len(binned_sieie_hist['PT_2_eta_2'] > 0): out['PT_2_eta_2'].fill(dataset=dataset, PT_2_eta_2=binned_sieie_hist['PT_2_eta_2']) if len(binned_sieie_hist['PT_2_eta_3'] > 0): out['PT_2_eta_3'].fill(dataset=dataset, PT_2_eta_3=binned_sieie_hist['PT_2_eta_3']) if len(binned_sieie_hist['PT_2_eta_4'] > 0): out['PT_2_eta_4'].fill(dataset=dataset, PT_2_eta_4=binned_sieie_hist['PT_2_eta_4']) if len(binned_sieie_hist['PT_3_eta_1'] > 0): out['PT_3_eta_1'].fill(dataset=dataset, PT_3_eta_1=binned_sieie_hist['PT_3_eta_1']) if len(binned_sieie_hist['PT_3_eta_2'] > 0): out['PT_3_eta_2'].fill(dataset=dataset, PT_3_eta_2=binned_sieie_hist['PT_3_eta_2']) if len(binned_sieie_hist['PT_3_eta_3'] > 0): out['PT_3_eta_3'].fill(dataset=dataset, PT_3_eta_3=binned_sieie_hist['PT_3_eta_3']) if len(binned_sieie_hist['PT_3_eta_4'] > 0): out['PT_3_eta_4'].fill(dataset=dataset, PT_3_eta_4=binned_sieie_hist['PT_3_eta_4']) if len(binned_sieie_hist['PT_4_eta_1'] > 0): out['PT_4_eta_1'].fill(dataset=dataset, PT_4_eta_1=binned_sieie_hist['PT_4_eta_1']) if len(binned_sieie_hist['PT_4_eta_2'] > 0): out['PT_4_eta_2'].fill(dataset=dataset, PT_4_eta_2=binned_sieie_hist['PT_4_eta_2']) if len(binned_sieie_hist['PT_4_eta_3'] > 0): out['PT_4_eta_3'].fill(dataset=dataset, PT_4_eta_3=binned_sieie_hist['PT_4_eta_3']) if len(binned_sieie_hist['PT_4_eta_4'] > 0): out['PT_4_eta_4'].fill(dataset=dataset, PT_4_eta_4=binned_sieie_hist['PT_4_eta_4']) return out
def process(self, events): # Initialize accumulator out = self.accumulator.identity() dataset = sample_name # events.metadata['dataset'] # Data or MC isData = "genWeight" not in events.fields isFake = self._isFake # Stop processing if there is no event remain if len(events) == 0: return out # Golden Json file if (self._year == "2018") and isData: injson = "/x5/cms/jwkim/gitdir/JWCorp/JW_analysis/Coffea_WZG/Corrections/Cert_314472-325175_13TeV_Legacy2018_Collisions18_JSON.txt.RunABD" if (self._year == "2017") and isData: injson = "/x5/cms/jwkim/gitdir/JWCorp/JW_analysis/Coffea_WZG/Corrections/Cert_294927-306462_13TeV_UL2017_Collisions17_GoldenJSON.txt" # <----- Get Scale factors ------># if not isData: # Egamma reco ID get_ele_reco_above20_sf = self._corrections["get_ele_reco_above20_sf"][ self._year ] get_ele_medium_id_sf = self._corrections["get_ele_medium_id_sf"][self._year] get_pho_medium_id_sf = self._corrections["get_pho_medium_id_sf"][self._year] # DoubleEG trigger # 2016, 2017 are not applied yet if self._year == "2018": 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(events.Pileup.nTrueInt, "int64") pu = self._puweight_arr[pu_weight_idx] print("## pu_idx: ",len(pu_weight_idx),pu_weight_idx) print("## pu_arr: ",len(self._puweight_arr),self._puweight_arr) print("## pu:",len(pu),pu) selection = processor.PackedSelection() # Cut flow cut0 = np.zeros(len(events)) out["cutflow"].fill(dataset=dataset, cutflow=cut0) # <----- Helper functions ------># # Sort by PT helper function def sort_by_pt(ele, pho, jet): ele = ele[ak.argsort(ele.pt, ascending=False, axis=1)] pho = pho[ak.argsort(pho.pt, ascending=False, axis=1)] jet = jet[ak.argsort(jet.pt, ascending=False, axis=1)] return ele, pho, jet # Lorentz vectors from coffea.nanoevents.methods import vector ak.behavior.update(vector.behavior) def TLorentz_vector(vec): vec = ak.zip( {"x": vec.x, "y": vec.y, "z": vec.z, "t": vec.t}, with_name="LorentzVector", ) return vec def TLorentz_vector_cylinder(vec): vec = ak.zip( { "pt": vec.pt, "eta": vec.eta, "phi": vec.phi, "mass": vec.mass, }, with_name="PtEtaPhiMLorentzVector", ) return vec # <----- Selection ------># Initial_events = events # Good Run ( Golden Json files ) from coffea import lumi_tools if isData: lumi_mask_builder = lumi_tools.LumiMask(injson) lumimask = ak.Array( lumi_mask_builder.__call__(events.run, events.luminosityBlock) ) events = events[lumimask] # print("{0}% of files pass good-run conditions".format(len(events)/ len(Initial_events))) # Stop processing if there is no event remain if len(events) == 0: return out ##----------- Cut flow1: Passing Triggers # 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] events.Electron, events.Photon, events.Jet = sort_by_pt( events.Electron, events.Photon, events.Jet ) # Good Primary vertex nPV = events.PV.npvsGood nPV_nw = events.PV.npvsGood if not isData: nPV = nPV * pu print(pu) # Apply cut1 events = events[double_ele_triggers_arr] if not isData: pu = pu[double_ele_triggers_arr] # Stop processing if there is no event remain if len(events) == 0: return out cut1 = np.ones(len(events)) out["cutflow"].fill(dataset=dataset, cutflow=cut1) # Set Particles Electron = events.Electron Muon = events.Muon Photon = events.Photon MET = events.MET Jet = events.Jet # --Muon ( only used to calculate dR ) MuSelmask = ( (Muon.pt >= 10) & (abs(Muon.eta) <= 2.5) & (Muon.tightId) & (Muon.pfRelIso04_all < 0.15) ) Muon = Muon[MuSelmask] # --Loose Muon ( For Loose Muon veto ) LoooseMuSelmask = ( (Muon.pt > 20) & (abs(Muon.eta) < 2.4) & (Muon.isPFcand) & (Muon.isGlobal | Muon.isTracker) & (Muon.pfRelIso03_all < 0.25) ) # Reference: VBS Zgamma+2jets VetoMuon = Muon[LoooseMuSelmask] ##----------- Cut flow2: Electron Selection EleSelmask = ( (Electron.pt >= 10) & (np.abs(Electron.eta + Electron.deltaEtaSC) < 1.479) & (Electron.cutBased > 2) & (abs(Electron.dxy) < 0.05) & (abs(Electron.dz) < 0.1) ) | ( (Electron.pt >= 10) & (np.abs(Electron.eta + Electron.deltaEtaSC) > 1.479) & (np.abs(Electron.eta + Electron.deltaEtaSC) <= 2.5) & (Electron.cutBased > 2) & (abs(Electron.dxy) < 0.1) & (abs(Electron.dz) < 0.2) ) Electron = Electron[EleSelmask] # Event with 3 Electrons # apply cut 2 Tri_electron_mask = ak.num(Electron) == 3 Electron = Electron[Tri_electron_mask] Photon = Photon[Tri_electron_mask] Jet = Jet[Tri_electron_mask] MET = MET[Tri_electron_mask] Muon = Muon[Tri_electron_mask] VetoMuon = VetoMuon[Tri_electron_mask] if not isData: pu = pu[Tri_electron_mask] events = events[Tri_electron_mask] # Stop processing if there is no event remain if len(Electron) == 0: return out cut2 = np.ones(len(Photon)) * 2 out["cutflow"].fill(dataset=dataset, cutflow=cut2) ##----------- Cut flow3: 4th lepton veto (Loose Muon) # Veto 4th Loose muon # apply cut 3 fourth_lepton_veto = ak.num(VetoMuon) < 1 Electron = Electron[fourth_lepton_veto] Photon = Photon[fourth_lepton_veto] Jet = Jet[fourth_lepton_veto] MET = MET[fourth_lepton_veto] Muon = Muon[fourth_lepton_veto] if not isData: pu = pu[fourth_lepton_veto] events = events[fourth_lepton_veto] # Stop processing if there is no event remain if len(Electron) == 0: return out cut3 = np.ones(len(Photon)) * 3 out["cutflow"].fill(dataset=dataset, cutflow=cut3) ##----------- Cut flow4: Photon Selection # Basic photon selection isgap_mask = (abs(Photon.eta) < 1.442) | ( (abs(Photon.eta) > 1.566) & (abs(Photon.eta) < 2.5) ) Pixel_seed_mask = ~Photon.pixelSeed if (dataset == "ZZ") and (self._year == "2017"): PT_ID_mask = (Photon.pt >= 20) & ( Photon.cutBasedBitmap >= 3 ) # 2^0(Loose) + 2^1(Medium) + 2^2(Tights) else: PT_ID_mask = (Photon.pt >= 20) & (Photon.cutBased > 1) # dR cut with selected Muon and Electrons dr_pho_ele_mask = ak.all( Photon.metric_table(Electron) >= 0.5, axis=-1 ) # default metric table: delta_r dr_pho_mu_mask = ak.all(Photon.metric_table(Muon) >= 0.5, axis=-1) # genPartFlav cut """ if dataset == "WZG": isPrompt = (Photon.genPartFlav == 1) | (Photon.genPartFlav == 11) PhoSelmask = PT_ID_mask & isgap_mask & Pixel_seed_mask & isPrompt & dr_pho_ele_mask & dr_pho_mu_mask elif dataset == "WZ": isPrompt = (Photon.genPartFlav == 1) PhoSelmask = PT_ID_mask & isgap_mask & Pixel_seed_mask & ~isPrompt & dr_pho_ele_mask & dr_pho_mu_mask else: PhoSelmask = PT_ID_mask & isgap_mask & Pixel_seed_mask & dr_pho_ele_mask & dr_pho_mu_mask """ PhoSelmask = ( PT_ID_mask & isgap_mask & Pixel_seed_mask & dr_pho_ele_mask & dr_pho_mu_mask ) Photon = Photon[PhoSelmask] # Apply cut 4 A_photon_mask = ak.num(Photon) > 0 Electron = Electron[A_photon_mask] Photon = Photon[A_photon_mask] Jet = Jet[A_photon_mask] Muon = Muon[A_photon_mask] MET = MET[A_photon_mask] if not isData: pu = pu[A_photon_mask] events = events[A_photon_mask] # Stop processing if there is no event remain if len(Electron) == 0: return out cut4 = np.ones(len(Photon)) * 4 out["cutflow"].fill(dataset=dataset, cutflow=cut4) ##----------- Cut flow5: OSSF # OSSF index maker @numba.njit def find_3lep(events_leptons, builder): for leptons in events_leptons: builder.begin_list() nlep = len(leptons) for i0 in range(nlep): for i1 in range(i0 + 1, nlep): if leptons[i0].charge + leptons[i1].charge != 0: continue for i2 in range(nlep): if len({i0, i1, i2}) < 3: continue builder.begin_tuple(3) builder.index(0).integer(i0) builder.index(1).integer(i1) builder.index(2).integer(i2) builder.end_tuple() builder.end_list() return builder eee_triplet_idx = find_3lep(Electron, ak.ArrayBuilder()).snapshot() ossf_mask = ak.num(eee_triplet_idx) == 2 # Apply cut 5 eee_triplet_idx = eee_triplet_idx[ossf_mask] Electron = Electron[ossf_mask] Photon = Photon[ossf_mask] Jet = Jet[ossf_mask] MET = MET[ossf_mask] if not isData: pu = pu[ossf_mask] events = events[ossf_mask] # Stop processing if there is no event remain if len(Electron) == 0: return out cut5 = np.ones(ak.sum(ak.num(Electron) > 0)) * 5 out["cutflow"].fill(dataset=dataset, cutflow=cut5) # Define Electron Triplet Triple_electron = [Electron[eee_triplet_idx[idx]] for idx in "012"] Triple_eee = ak.zip( { "lep1": Triple_electron[0], "lep2": Triple_electron[1], "lep3": Triple_electron[2], "p4": TLorentz_vector(Triple_electron[0] + Triple_electron[1]), } ) # Ele pair selector --> Close to Z mass bestZ_idx = ak.singletons(ak.argmin(abs(Triple_eee.p4.mass - 91.1876), axis=1)) Triple_eee = Triple_eee[bestZ_idx] leading_ele = Triple_eee.lep1 subleading_ele = Triple_eee.lep2 third_ele = Triple_eee.lep3 def make_leading_pair(target, base): return target[ak.argmax(base.pt, axis=1, keepdims=True)] leading_pho = make_leading_pair(Photon, Photon) # -- Scale Factor for each electron # Trigger weight helper function 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: ## -------------< Egamma ID and Reco Scale factor > -----------------## get_pho_medium_id_sf = get_pho_medium_id_sf( ak.flatten(leading_pho.eta), ak.flatten(leading_pho.pt) ) ele_reco_sf = ( get_ele_reco_above20_sf( ak.flatten(leading_ele.deltaEtaSC + leading_ele.eta), ak.flatten(leading_ele.pt), ) * get_ele_reco_above20_sf( ak.flatten(subleading_ele.deltaEtaSC + subleading_ele.eta), ak.flatten(subleading_ele.pt), ) * get_ele_reco_above20_sf( ak.flatten(third_ele.deltaEtaSC + third_ele.eta), ak.flatten(third_ele.pt), ) ) ele_medium_id_sf = ( get_ele_medium_id_sf( ak.flatten(leading_ele.deltaEtaSC + leading_ele.eta), ak.flatten(leading_ele.pt), ) * get_ele_medium_id_sf( ak.flatten(subleading_ele.deltaEtaSC + subleading_ele.eta), ak.flatten(subleading_ele.pt), ) * get_ele_medium_id_sf( ak.flatten(third_ele.deltaEtaSC + third_ele.eta), ak.flatten(third_ele.pt), ) ) ## -------------< Double Electron Trigger Scale factor > -----------------## 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) # -- 2017,2016 are not applied yet if self._year == "2018": ele_trig_weight = Trigger_Weight(eta1, pt1, eta2, pt2) ##----------- Cut flow6: Event selection # Mee cut diele = Triple_eee.p4 Mee_cut_mask = ak.firsts(diele.mass) > 4 # Z mass window # zmass_window_mask = ak.firsts(abs(diele.mass - 91.1876)) < 15 # SR, CR_ZZA, CR_Z+jets, CR_Conversion # zmass_window_mask = ak.firsts(abs(diele.mass - 91.1876)) > 5 # CR_t-enriched # zmass_window_mask = ak.firsts(abs(diele.mass - 91.1876)) > 15 # CR_Conversion # M(eee) cut SR, CR_ZZA, CR_Z+jets, CR_t enriched # eee = Triple_eee.lep1 + Triple_eee.lep2 + Triple_eee.lep3 # Meee_cut_mask = ak.firsts(eee.mass > 100) # Meee_cut_mask = ak.firsts(eee.mass <= 100) # b-Jet veto cut #SR, CR_ZZA, CR_Z+jets, CR_Conversion # bjet_mask = (Jet.btagCSVV2 > 0.4184) & (Jet.pt > 30) # bjet_veto_mask = ak.num(Jet[bjet_mask]) == 0 # bjet_veto_mask = ak.num(Jet[bjet_mask]) > 0 # CR_t-enriched # Electron PT cuts Elept_mask = ak.firsts( (leading_ele.pt >= 25) & (subleading_ele.pt >= 10) & (third_ele.pt >= 25) ) # MET cuts MET_mask = MET > 20 # Baseline # MET_mask = MET.pt > 30 # SR, CR-ZZE, CR-t-entirched # MET_mask = MET.pt <= 30 # CR-Z+jets. CR-Conversion # Mask # Event_sel_mask = Elept_mask & MET_mask & bjet_veto_mask & Mee_cut_mask & zmass_window_mask & Meee_cut_mask # SR,CR Event_sel_mask = Elept_mask & MET_mask & Mee_cut_mask # SR,CR # Apply cut6 Triple_eee_sel = Triple_eee[Event_sel_mask] leading_pho_sel = leading_pho[Event_sel_mask] MET_sel = MET[Event_sel_mask] events = events[Event_sel_mask] # Photon EE and EB isEE_mask = leading_pho.isScEtaEE isEB_mask = leading_pho.isScEtaEB Pho_EE = leading_pho[isEE_mask & Event_sel_mask] Pho_EB = leading_pho[isEB_mask & Event_sel_mask] # Stop processing if there is no event remain if len(leading_pho_sel) == 0: return out cut6 = np.ones(ak.sum(ak.num(leading_pho_sel) > 0)) * 6 out["cutflow"].fill(dataset=dataset, cutflow=cut6) ## -------------------- Prepare making hist --------------# # Photon phoPT = ak.flatten(leading_pho_sel.pt) phoEta = ak.flatten(leading_pho_sel.eta) phoPhi = ak.flatten(leading_pho_sel.phi) # Photon EE if len(Pho_EE.pt) != 0: Pho_EE_PT = ak.flatten(Pho_EE.pt) Pho_EE_Eta = ak.flatten(Pho_EE.eta) Pho_EE_Phi = ak.flatten(Pho_EE.phi) Pho_EE_sieie = ak.flatten(Pho_EE.sieie) Pho_EE_hoe = ak.flatten(Pho_EE.hoe) Pho_EE_Iso_charge = ak.flatten(Pho_EE.pfRelIso03_chg) # Photon EB if len(Pho_EB.pt) != 0: Pho_EB_PT = ak.flatten(Pho_EB.pt) Pho_EB_Eta = ak.flatten(Pho_EB.eta) Pho_EB_Phi = ak.flatten(Pho_EB.phi) Pho_EB_sieie = ak.flatten(Pho_EB.sieie) Pho_EB_hoe = ak.flatten(Pho_EB.hoe) Pho_EB_Iso_charge = ak.flatten(Pho_EB.pfRelIso03_chg) # Electrons ele1PT = ak.flatten(Triple_eee_sel.lep1.pt) ele1Eta = ak.flatten(Triple_eee_sel.lep1.eta) ele1Phi = ak.flatten(Triple_eee_sel.lep1.phi) ele2PT = ak.flatten(Triple_eee_sel.lep2.pt) ele2Eta = ak.flatten(Triple_eee_sel.lep2.eta) ele2Phi = ak.flatten(Triple_eee_sel.lep2.phi) ele3PT = ak.flatten(Triple_eee_sel.lep3.pt) ele3Eta = ak.flatten(Triple_eee_sel.lep3.eta) ele3Phi = ak.flatten(Triple_eee_sel.lep3.phi) charge = ak.flatten(Triple_eee.lep1.charge + Triple_eee.lep2.charge) # MET met = ak.to_numpy(MET_sel) # M(eea) M(ee) diele = Triple_eee_sel.p4 eeg_vec = diele + leading_pho_sel Meea = ak.flatten(eeg_vec.mass) Mee = ak.flatten(Triple_eee_sel.p4.mass) # --- Apply weight and hist if isFake: weights = processor.Weights(len(cut6)) else: weights = processor.Weights(len(cut5)) # -------------------- Sieie bins---------------------------# def make_bins(pt, eta, bin_range_str): bin_dict = { "PT_1_eta_1": (pt > 20) & (pt < 30) & (eta < 1), "PT_1_eta_2": (pt > 20) & (pt < 30) & (eta > 1) & (eta < 1.5), "PT_1_eta_3": (pt > 20) & (pt < 30) & (eta > 1.5) & (eta < 2), "PT_1_eta_4": (pt > 20) & (pt < 30) & (eta > 2) & (eta < 2.5), "PT_2_eta_1": (pt > 30) & (pt < 40) & (eta < 1), "PT_2_eta_2": (pt > 30) & (pt < 40) & (eta > 1) & (eta < 1.5), "PT_2_eta_3": (pt > 30) & (pt < 40) & (eta > 1.5) & (eta < 2), "PT_2_eta_4": (pt > 30) & (pt < 40) & (eta > 2) & (eta < 2.5), "PT_3_eta_1": (pt > 40) & (pt < 50) & (eta < 1), "PT_3_eta_2": (pt > 40) & (pt < 50) & (eta > 1) & (eta < 1.5), "PT_3_eta_3": (pt > 40) & (pt < 50) & (eta > 1.5) & (eta < 2), "PT_3_eta_4": (pt > 40) & (pt < 50) & (eta > 2) & (eta < 2.5), "PT_4_eta_1": (pt > 50) & (eta < 1), "PT_4_eta_2": (pt > 50) & (eta > 1) & (eta < 1.5), "PT_4_eta_3": (pt > 50) & (eta > 1.5) & (eta < 2), "PT_4_eta_4": (pt > 50) & (eta > 2) & (eta < 2.5), } binmask = bin_dict[bin_range_str] return binmask bin_name_list = [ "PT_1_eta_1", "PT_1_eta_2", "PT_1_eta_3", "PT_1_eta_4", "PT_2_eta_1", "PT_2_eta_2", "PT_2_eta_3", "PT_2_eta_4", "PT_3_eta_1", "PT_3_eta_2", "PT_3_eta_3", "PT_3_eta_4", "PT_4_eta_1", "PT_4_eta_2", "PT_4_eta_3", "PT_4_eta_4", ] ## -- Fake-fraction Lookup table --## if isFake: # Make Bin-range mask binned_pteta_mask = {} for name in bin_name_list: binned_pteta_mask[name] = make_bins( ak.flatten(leading_pho_sel.pt), ak.flatten(abs(leading_pho_sel.eta)), name, ) # Read Fake fraction --> Mapping bin name to int() in_dict = np.load('Fitting_v2/results_210517.npy',allow_pickle="True")[()] idx=0 fake_dict ={} for i,j in in_dict.items(): fake_dict[idx] = j idx+=1 # Reconstruct Fake_weight fw= 0 for i,j in binned_pteta_mask.items(): fw = fw + j*fake_dict[bin_name_list.index(i)] # Process 0 weight to 1 @numba.njit def zero_one(x): if x == 0: x = 1 return x vec_zero_one = np.vectorize(zero_one) fw = vec_zero_one(fw) # --- skim cut-weight if not isFake: def skim_weight(arr): mask1 = ~ak.is_none(arr) subarr = arr[mask1] mask2 = subarr != 0 return ak.to_numpy(subarr[mask2]) else: def skim_weight(arr): return arr if not isFake: cuts = Event_sel_mask cuts_pho_EE = ak.flatten(isEE_mask) cuts_pho_EB = ak.flatten(isEB_mask) if isFake: cuts = np.ones(len(Event_sel_mask)) cuts_pho_EE = ak.flatten(isEE_mask & Event_sel_mask) cuts_pho_EB = ak.flatten(isEB_mask & Event_sel_mask) if isFake: weights.add("fake_fraction", fw) # Weight and SF here if not (isData | isFake): weights.add("pileup", pu) weights.add("ele_id", ele_medium_id_sf) weights.add("pho_id", get_pho_medium_id_sf) weights.add("ele_reco", ele_reco_sf) # 2016,2017 are not applied yet if self._year == "2018": weights.add("ele_trigger", ele_trig_weight) # ---------------------------- Fill hist --------------------------------------# # Initial events out["sumw"][dataset] += len(Initial_events) print("cut1: {0},cut2: {1},cut3: {2},cut4: {3},cut5: {4},cut6: {5},cut7: {6}".format(len(cut0), len(cut1), len(cut2), len(cut3), len(cut4), len(cut5),len(cut6))) ## Cut flow loop #for cut in [cut0, cut1, cut2, cut3, cut4, cut5,cut6]: # out["cutflow"].fill(dataset=dataset, cutflow=cut) # Primary vertex out["nPV"].fill( dataset=dataset, nPV=nPV, ) out["nPV_nw"].fill(dataset=dataset, nPV_nw=nPV_nw) # Fill hist # -- met -- # out["met"].fill( dataset=dataset, met=met, weight=skim_weight(weights.weight() * cuts) ) # --mass -- # out["mass"].fill( dataset=dataset, mass=Mee, weight=skim_weight(weights.weight() * cuts) ) out["mass_eea"].fill( dataset=dataset, mass_eea=Meea, weight=skim_weight(weights.weight() * cuts) ) # -- Electron -- # 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), ) out["ele3pt"].fill( dataset=dataset, ele3pt=ele3PT, weight=skim_weight(weights.weight() * cuts) ) # -- Photon -- # out["phopt"].fill( dataset=dataset, phopt=phoPT, weight=skim_weight(weights.weight() * cuts) ) out["phoeta"].fill( dataset=dataset, phoeta=phoEta, weight=skim_weight(weights.weight() * cuts) ) out["phophi"].fill( dataset=dataset, phophi=phoPhi, weight=skim_weight(weights.weight() * cuts) ) if len(Pho_EE.pt) != 0: out["pho_EE_pt"].fill( dataset=dataset, pho_EE_pt=Pho_EE_PT, weight=skim_weight(weights.weight() * cuts * cuts_pho_EE), ) out["pho_EE_eta"].fill( dataset=dataset, pho_EE_eta=Pho_EE_Eta, weight=skim_weight(weights.weight() * cuts * cuts_pho_EE), ) out["pho_EE_phi"].fill( dataset=dataset, pho_EE_phi=Pho_EE_Phi, weight=skim_weight(weights.weight() * cuts * cuts_pho_EE), ) out["pho_EE_hoe"].fill( dataset=dataset, pho_EE_hoe=Pho_EE_hoe, weight=skim_weight(weights.weight() * cuts * cuts_pho_EE), ) out["pho_EE_sieie"].fill( dataset=dataset, pho_EE_sieie=Pho_EE_sieie, weight=skim_weight(weights.weight() * cuts * cuts_pho_EE), ) out["pho_EE_Iso_chg"].fill( dataset=dataset, pho_EE_Iso_chg=Pho_EE_Iso_charge, weight=skim_weight(weights.weight() * cuts * cuts_pho_EE), ) if len(Pho_EB.pt) != 0: out["pho_EB_pt"].fill( dataset=dataset, pho_EB_pt=Pho_EB_PT, weight=skim_weight(weights.weight() * cuts * cuts_pho_EB), ) out["pho_EB_eta"].fill( dataset=dataset, pho_EB_eta=Pho_EB_Eta, weight=skim_weight(weights.weight() * cuts * cuts_pho_EB), ) out["pho_EB_phi"].fill( dataset=dataset, pho_EB_phi=Pho_EB_Phi, weight=skim_weight(weights.weight() * cuts * cuts_pho_EB), ) out["pho_EB_hoe"].fill( dataset=dataset, pho_EB_hoe=Pho_EB_hoe, weight=skim_weight(weights.weight() * cuts * cuts_pho_EB), ) out["pho_EB_sieie"].fill( dataset=dataset, pho_EB_sieie=Pho_EB_sieie, weight=skim_weight(weights.weight() * cuts * cuts_pho_EB), ) out["pho_EB_Iso_chg"].fill( dataset=dataset, pho_EB_Iso_chg=Pho_EB_Iso_charge, weight=skim_weight(weights.weight() * cuts * cuts_pho_EB), ) return out
inspected_event_weight = weight/abs(weight) inspected_event_weight = inspected_event_weight.sum() # 4. Particle Definition and Selections print(".... Particle and Event selection") # make Jet selection mask Jet_sel_mask = (Jet.PT > 35 ) & (abs(Jet.Eta) < 2.4) Jet = Jet[Jet_sel_mask] # selected jet # Baseline selection NJet_mask = ak.num(Jet) >= 9 # selected jet >=9 mask HT_mask = ak.firsts(HT >= 700) # HT cut Nbtag_mask = ak.sum(Jet.BTag,axis=1) >= 3 # at leat 3 b-tag base_line_mask = NJet_mask & HT_mask & Nbtag_mask # Apply cut using boolean indexing Jet = Jet[base_line_mask] HT = HT[base_line_mask] inspected_event_weight = inspected_event_weight[base_line_mask] print("After N jet cut: ",len(Jet),len(HT),len(inspected_event_weight)) # 5. Flatten the arrays ( 2D -> 1D ) print(".... Write file") # Fltten and chage type Awkward array to numpy array
def process(self, events): output = self.accumulator.identity() dataset = events.metadata['dataset'] isRealData = 'genWeight' not in events.fields if not isRealData: output['sumw'][dataset] += sum(events.genWeight) JECversion = JECversions[str(self.year)]['MC'] else: output['nbtagmu'][dataset] += ak.count(events.event) JECversion = JECversions[str( self.year)]['Data'][dataset.split('BTagMu')[1]] ############ # Some corrections weights = processor.Weights(len(events)) corrections = {} if not isRealData: weights.add('genWeight', events.genWeight) weights.add( 'pileup_weight', self.puReweight(self.puFile, self.nTrueFile, dataset)(events.Pileup.nPU)) events.FatJet = self.applyJEC(events.FatJet, events.fixedGridRhoFastjetAll, events.caches[0], 'AK8PFPuppi', isRealData, JECversion) cuts = processor.PackedSelection() ############ # Trigger selection if self.year == 2016: if 'BTagMu_AK4Jet300_Mu5' not in events.HLT.fields: self.triggers = [ trigger.replace('AK4', '') for trigger in self.triggers ] elif self.year == 2018: for (i, trigger) in enumerate(self.triggers): if trigger.strip("HLT_") not in events.HLT.fields: self.triggers[i] = trigger + "_noalgo" trig_arrs = [ events.HLT[_trig.strip("HLT_")] for _trig in self.triggers ] req_trig = np.zeros(len(events), dtype='bool') for t in trig_arrs: req_trig = req_trig | t cuts.add('trigger', ak.to_numpy(req_trig)) ############ # Basic cuts ## Muon cuts # muon twiki: https://twiki.cern.ch/twiki/bin/view/CMS/SWGuideMuonIdRun2 events.Muon = events.Muon[(events.Muon.pt > 5) & (abs(events.Muon.eta < 2.4)) & (events.Muon.tightId != 1) & (events.Muon.pfRelIso04_all > 0.15)] events.Muon = ak.pad_none(events.Muon, 2, axis=1) ## Jet cuts (not used) events.Jet = events.Jet[(events.Jet.pt > 25) & (abs(events.Jet.eta) <= 2.5)] #req_jets = (ak.count(events.Jet.pt, axis=1) >= 2) ## FatJet cuts events.FatJet = events.FatJet[ (events.FatJet.pt > self._mask_fatjets['basic']['pt_cut']) & (abs(events.FatJet.eta) <= self._mask_fatjets['basic']['eta_cut']) & (events.FatJet.jetId > self._mask_fatjets['basic']['jetId_cut']) & (ak.count(events.FatJet.subjets.pt, axis=2) >= 2)] ## subjet sel to crosscheck #print(events['FatJetSVs']) ## Event level variables eventVariables = {} eventVariables['nfatjet'] = ak.num(events.FatJet) ## Leading jet variables leadfatjet = ak.firsts(events.FatJet) leadfatjet['tau21'] = leadfatjet.tau2 / leadfatjet.tau1 subjet1 = ak.pad_none(leadfatjet.subjets, 2)[:, 0] subjet2 = ak.pad_none(leadfatjet.subjets, 2)[:, 1] leadfatjet['nsv1'] = get_nsv(subjet1, events.SV) leadfatjet['nsv2'] = get_nsv(subjet2, events.SV) leadfatjet['nmusj1'] = ak.num(subjet1.delta_r(events.Muon) < 0.4) leadfatjet['nmusj2'] = ak.num(subjet2.delta_r(events.Muon) < 0.4) fatjet_mutag = (leadfatjet.nmusj1 >= 1) & (leadfatjet.nmusj2 >= 1) cuts.add('fatjet_mutag', ak.to_numpy(fatjet_mutag)) for DDX in self._mask_DDX.keys(): for wp, cut in self._mask_DDX[DDX].items(): DDX_pass = (leadfatjet[f'btag{DDX}vLV2'] > cut) DDX_fail = (leadfatjet[f'btag{DDX}vLV2'] < cut) cuts.add(f'{DDX}_pass{wp}wp', ak.to_numpy(DDX_pass)) cuts.add(f'{DDX}_fail{wp}wp', ak.to_numpy(DDX_fail)) flavors = {} if not isRealData: flavors['b'] = (leadfatjet.hadronFlavour == 5) flavors['c'] = (leadfatjet.hadronFlavour == 4) flavors['l'] = (leadfatjet.hadronFlavour < 4) flavors['bb'] = abs(leadfatjet.hadronFlavour == 5) & ( leadfatjet.nBHadrons >= 2) #& (leadfatjet.nCHadrons == 0) flavors['cc'] = abs(leadfatjet.hadronFlavour == 4) & ( leadfatjet.nBHadrons == 0) & (leadfatjet.nCHadrons >= 2) #flavors['ll'] = abs(leadfatjet.hadronFlavour < 4) & (leadfatjet.nBHadrons == 0) & (leadfatjet.nCHadrons == 0) flavors['b'] = flavors['b'] & ~flavors['bb'] flavors['c'] = flavors['c'] & ~flavors['cc'] flavors['l'] = flavors['l'] & ~flavors['bb'] & ~flavors[ 'cc'] & ~flavors['b'] & ~flavors['c'] #flavors['others'] = ~flavors['l'] & ~flavors['bb'] & ~flavors['cc'] & ~flavors['b'] & ~flavors['c'] else: flavors['Data'] = np.ones(len(events), dtype='bool') for selname, cut in self._mask_fatjets.items(): sel = (leadfatjet.pt > cut['pt_cut']) & \ (leadfatjet.msoftdrop > cut['mass_cut']) & \ (abs(leadfatjet.eta) < cut['eta_cut']) & \ (leadfatjet.jetId >= cut['jetId_cut']) & \ (leadfatjet.tau21 < cut['tau21_cut']) #(leadfatjet.Hbb > cut['Hbb']) cuts.add(selname, ak.to_numpy(sel)) selection = {} selection['basic'] = {'trigger', 'basic'} selection['pt350msd50'] = {'trigger', 'fatjet_mutag', 'pt350msd50'} selection['msd100tau06'] = {'trigger', 'fatjet_mutag', 'msd100tau06'} selection['pt400msd100tau06'] = { 'trigger', 'fatjet_mutag', 'pt400msd100tau06' } for mask_f in self._final_mask: for DDX in self._mask_DDX.keys(): for wp, cut in self._mask_DDX[DDX].items(): selection[f'{mask_f}{DDX}pass{wp}wp'] = selection[ mask_f].copy() selection[f'{mask_f}{DDX}pass{wp}wp'].add( f'{DDX}_pass{wp}wp') selection[f'{mask_f}{DDX}fail{wp}wp'] = selection[ mask_f].copy() selection[f'{mask_f}{DDX}fail{wp}wp'].add( f'{DDX}_fail{wp}wp') for histname, h in output.items(): sel = [r for r in selection.keys() if r in histname.split('_')] if ((histname in self.fatjet_hists) | ('hist2d_fatjet' in histname)): for flav, mask in flavors.items(): weight = weights.weight() * cuts.all( *selection[sel[0]]) * ak.to_numpy(mask) fields = { k: ak.fill_none(leadfatjet[k], -9999) for k in h.fields if k in dir(leadfatjet) } h.fill(dataset=dataset, flavor=flav, **fields, weight=weight) if histname in self.event_hists: for flav, mask in flavors.items(): weight = weights.weight() * cuts.all( *selection[sel[0]]) * ak.to_numpy(mask) fields = { k: ak.fill_none(eventVariables[k], -9999) for k in h.fields if k in eventVariables.keys() } h.fill(dataset=dataset, flavor=flav, **fields, weight=weight) return output
def process_shift(self, events, shift_name): dataset = events.metadata['dataset'] isRealData = not hasattr(events, "genWeight") selection = PackedSelection() weights = Weights(len(events), storeIndividual=True) output = self.make_output() if shift_name is None and not isRealData: output['sumw'] = ak.sum(events.genWeight) if isRealData or self._newTrigger: trigger = np.zeros(len(events), dtype='bool') for t in self._triggers[self._year]: if t in events.HLT.fields: trigger = trigger | events.HLT[t] selection.add('trigger', trigger) del trigger else: selection.add('trigger', np.ones(len(events), dtype='bool')) if isRealData: selection.add( 'lumimask', lumiMasks[self._year](events.run, events.luminosityBlock)) else: selection.add('lumimask', np.ones(len(events), dtype='bool')) if isRealData and self._skipRunB and self._year == '2017': selection.add('dropB', events.run > 299329) else: selection.add('dropB', np.ones(len(events), dtype='bool')) if isRealData: trigger = np.zeros(len(events), dtype='bool') for t in self._muontriggers[self._year]: if t in events.HLT.fields: trigger |= np.array(events.HLT[t]) selection.add('muontrigger', trigger) del trigger else: selection.add('muontrigger', np.ones(len(events), dtype='bool')) metfilter = np.ones(len(events), dtype='bool') for flag in self._met_filters[ self._year]['data' if isRealData else 'mc']: metfilter &= np.array(events.Flag[flag]) selection.add('metfilter', metfilter) del metfilter fatjets = events.FatJet fatjets['msdcorr'] = corrected_msoftdrop(fatjets) fatjets['qcdrho'] = 2 * np.log(fatjets.msdcorr / fatjets.pt) fatjets['n2ddt'] = fatjets.n2b1 - n2ddt_shift(fatjets, year=self._year) fatjets['msdcorr_full'] = fatjets['msdcorr'] * self._msdSF[self._year] candidatejet = fatjets[ # https://github.com/DAZSLE/BaconAnalyzer/blob/master/Analyzer/src/VJetLoader.cc#L269 (fatjets.pt > 200) & (abs(fatjets.eta) < 2.5) & fatjets.isTight # this is loose in sampleContainer ] candidatejet = candidatejet[:, : 2] # Only consider first two to match generators if self._jet_arbitration == 'pt': candidatejet = ak.firsts(candidatejet) elif self._jet_arbitration == 'mass': candidatejet = ak.firsts(candidatejet[ak.argmax( candidatejet.msdcorr, axis=1, keepdims=True)]) elif self._jet_arbitration == 'n2': candidatejet = ak.firsts(candidatejet[ak.argmin(candidatejet.n2ddt, axis=1, keepdims=True)]) elif self._jet_arbitration == 'ddb': candidatejet = ak.firsts(candidatejet[ak.argmax( candidatejet.btagDDBvLV2, axis=1, keepdims=True)]) elif self._jet_arbitration == 'ddc': candidatejet = ak.firsts(candidatejet[ak.argmax( candidatejet.btagDDCvLV2, axis=1, keepdims=True)]) else: raise RuntimeError("Unknown candidate jet arbitration") if self._tagger == 'v1': bvl = candidatejet.btagDDBvL cvl = candidatejet.btagDDCvL cvb = candidatejet.btagDDCvB elif self._tagger == 'v2': bvl = candidatejet.btagDDBvLV2 cvl = candidatejet.btagDDCvLV2 cvb = candidatejet.btagDDCvBV2 elif self._tagger == 'v3': bvl = candidatejet.particleNetMD_Xbb cvl = candidatejet.particleNetMD_Xcc / ( 1 - candidatejet.particleNetMD_Xbb) cvb = candidatejet.particleNetMD_Xcc / ( candidatejet.particleNetMD_Xcc + candidatejet.particleNetMD_Xbb) elif self._tagger == 'v4': bvl = candidatejet.particleNetMD_Xbb cvl = candidatejet.btagDDCvLV2 cvb = candidatejet.particleNetMD_Xcc / ( candidatejet.particleNetMD_Xcc + candidatejet.particleNetMD_Xbb) else: raise ValueError("Not an option") selection.add('minjetkin', (candidatejet.pt >= 450) & (candidatejet.pt < 1200) & (candidatejet.msdcorr >= 40.) & (candidatejet.msdcorr < 201.) & (abs(candidatejet.eta) < 2.5)) selection.add('_strict_mass', (candidatejet.msdcorr > 85) & (candidatejet.msdcorr < 130)) selection.add('_high_score', cvl > 0.8) selection.add('minjetkinmu', (candidatejet.pt >= 400) & (candidatejet.pt < 1200) & (candidatejet.msdcorr >= 40.) & (candidatejet.msdcorr < 201.) & (abs(candidatejet.eta) < 2.5)) selection.add('minjetkinw', (candidatejet.pt >= 200) & (candidatejet.pt < 1200) & (candidatejet.msdcorr >= 40.) & (candidatejet.msdcorr < 201.) & (abs(candidatejet.eta) < 2.5)) selection.add('jetid', candidatejet.isTight) selection.add('n2ddt', (candidatejet.n2ddt < 0.)) if not self._tagger == 'v2': selection.add('ddbpass', (bvl >= 0.89)) selection.add('ddcpass', (cvl >= 0.83)) selection.add('ddcvbpass', (cvb >= 0.2)) else: selection.add('ddbpass', (bvl >= 0.7)) selection.add('ddcpass', (cvl >= 0.45)) selection.add('ddcvbpass', (cvb >= 0.03)) jets = events.Jet jets = jets[(jets.pt > 30.) & (abs(jets.eta) < 2.5) & jets.isTight] # only consider first 4 jets to be consistent with old framework jets = jets[:, :4] dphi = abs(jets.delta_phi(candidatejet)) selection.add( 'antiak4btagMediumOppHem', ak.max(jets[dphi > np.pi / 2][self._ak4tagBranch], axis=1, mask_identity=False) < BTagEfficiency.btagWPs[self._ak4tagger][self._year]['medium']) ak4_away = jets[dphi > 0.8] selection.add( 'ak4btagMedium08', ak.max(ak4_away[self._ak4tagBranch], axis=1, mask_identity=False) > BTagEfficiency.btagWPs[self._ak4tagger][self._year]['medium']) met = events.MET selection.add('met', met.pt < 140.) goodmuon = ((events.Muon.pt > 10) & (abs(events.Muon.eta) < 2.4) & (events.Muon.pfRelIso04_all < 0.25) & events.Muon.looseId) nmuons = ak.sum(goodmuon, axis=1) leadingmuon = ak.firsts(events.Muon[goodmuon]) if self._looseTau: goodelectron = ((events.Electron.pt > 10) & (abs(events.Electron.eta) < 2.5) & (events.Electron.cutBased >= events.Electron.VETO)) nelectrons = ak.sum(goodelectron, axis=1) ntaus = ak.sum( ((events.Tau.pt > 20) & (abs(events.Tau.eta) < 2.3) & events.Tau.idDecayMode & ((events.Tau.idMVAoldDM2017v2 & 2) != 0) & ak.all(events.Tau.metric_table(events.Muon[goodmuon]) > 0.4, axis=2) & ak.all(events.Tau.metric_table( events.Electron[goodelectron]) > 0.4, axis=2)), axis=1, ) else: goodelectron = ( (events.Electron.pt > 10) & (abs(events.Electron.eta) < 2.5) & (events.Electron.cutBased >= events.Electron.LOOSE)) nelectrons = ak.sum(goodelectron, axis=1) ntaus = ak.sum( (events.Tau.pt > 20) & events.Tau.idDecayMode # bacon iso looser than Nano selection & ak.all(events.Tau.metric_table(events.Muon[goodmuon]) > 0.4, axis=2) & ak.all(events.Tau.metric_table(events.Electron[goodelectron]) > 0.4, axis=2), axis=1, ) selection.add('noleptons', (nmuons == 0) & (nelectrons == 0) & (ntaus == 0)) selection.add('onemuon', (nmuons == 1) & (nelectrons == 0) & (ntaus == 0)) selection.add('muonkin', (leadingmuon.pt > 55.) & (abs(leadingmuon.eta) < 2.1)) selection.add('muonDphiAK8', abs(leadingmuon.delta_phi(candidatejet)) > 2 * np.pi / 3) # W-Tag (Tag and Probe) # tag side selection.add( 'ak4btagMediumOppHem', ak.max(jets[dphi > np.pi / 2][self._ak4tagBranch], axis=1, mask_identity=False) > BTagEfficiency.btagWPs[self._ak4tagger][self._year]['medium']) selection.add('met40p', met.pt > 40.) selection.add('tightMuon', (leadingmuon.tightId) & (leadingmuon.pt > 53.)) # selection.add('ptrecoW', (leadingmuon + met).pt > 250.) selection.add('ptrecoW200', (leadingmuon + met).pt > 200.) selection.add( 'ak4btagNearMu', leadingmuon.delta_r(leadingmuon.nearest(ak4_away, axis=None)) < 2.0) _bjets = jets[self._ak4tagBranch] > BTagEfficiency.btagWPs[ self._ak4tagger][self._year]['medium'] # _nearAK8 = jets.delta_r(candidatejet) < 0.8 # _nearMu = jets.delta_r(ak.firsts(events.Muon)) < 0.3 # selection.add('ak4btagOld', ak.sum(_bjets & ~_nearAK8 & ~_nearMu, axis=1) >= 1) _nearAK8 = jets.delta_r(candidatejet) < 0.8 _nearMu = jets.delta_r(leadingmuon) < 0.3 selection.add('ak4btagOld', ak.sum(_bjets & ~_nearAK8 & ~_nearMu, axis=1) >= 1) # _nearAK8 = jets.delta_r(candidatejet) < 0.8 # _nearMu = jets.delta_r(candidatejet.nearest(events.Muon[goodmuon], axis=None)) < 0.3 # selection.add('ak4btagNew', ak.sum(_bjets & ~_nearAK8 & ~_nearMu, axis=1) >= 1) # probe side selection.add('minWjetpteta', (candidatejet.pt >= 200) & (abs(candidatejet.eta) < 2.4)) # selection.add('noNearMuon', candidatejet.delta_r(candidatejet.nearest(events.Muon[goodmuon], axis=None)) > 1.0) selection.add('noNearMuon', candidatejet.delta_r(leadingmuon) > 1.0) ##### if isRealData: genflavor = ak.zeros_like(candidatejet.pt) else: if 'HToCC' in dataset or 'HToBB' in dataset: if self._ewkHcorr: add_HiggsEW_kFactors(weights, events.GenPart, dataset) weights.add('genweight', events.genWeight) if "PSWeight" in events.fields: add_ps_weight(weights, events.PSWeight) else: add_ps_weight(weights, None) if "LHEPdfWeight" in events.fields: add_pdf_weight(weights, events.LHEPdfWeight) else: add_pdf_weight(weights, None) if "LHEScaleWeight" in events.fields: add_scalevar_7pt(weights, events.LHEScaleWeight) add_scalevar_3pt(weights, events.LHEScaleWeight) else: add_scalevar_7pt(weights, []) add_scalevar_3pt(weights, []) add_pileup_weight(weights, events.Pileup.nPU, self._year, dataset) bosons = getBosons(events.GenPart) matchedBoson = candidatejet.nearest(bosons, axis=None, threshold=0.8) if self._tightMatch: match_mask = ( (candidatejet.pt - matchedBoson.pt) / matchedBoson.pt < 0.5) & ((candidatejet.msdcorr - matchedBoson.mass) / matchedBoson.mass < 0.3) selmatchedBoson = ak.mask(matchedBoson, match_mask) genflavor = bosonFlavor(selmatchedBoson) else: genflavor = bosonFlavor(matchedBoson) genBosonPt = ak.fill_none(ak.firsts(bosons.pt), 0) if self._newVjetsKfactor: add_VJets_kFactors(weights, events.GenPart, dataset) else: add_VJets_NLOkFactor(weights, genBosonPt, self._year, dataset) if shift_name is None: output['btagWeight'].fill(val=self._btagSF.addBtagWeight( weights, ak4_away, self._ak4tagBranch)) if self._nnlops_rew and dataset in [ 'GluGluHToCC_M125_13TeV_powheg_pythia8' ]: weights.add('minlo_rew', powheg_to_nnlops(ak.to_numpy(genBosonPt))) if self._newTrigger: add_jetTriggerSF( weights, ak.firsts(fatjets), self._year if not self._skipRunB else f'{self._year}CDEF', selection) else: add_jetTriggerWeight(weights, candidatejet.msdcorr, candidatejet.pt, self._year) add_mutriggerSF(weights, leadingmuon, self._year, selection) add_mucorrectionsSF(weights, leadingmuon, self._year, selection) if self._year in ("2016", "2017"): weights.add("L1Prefiring", events.L1PreFiringWeight.Nom, events.L1PreFiringWeight.Up, events.L1PreFiringWeight.Dn) logger.debug("Weight statistics: %r" % weights.weightStatistics) msd_matched = candidatejet.msdcorr * self._msdSF[self._year] * ( genflavor > 0) + candidatejet.msdcorr * (genflavor == 0) regions = { 'signal': [ 'noleptons', 'minjetkin', 'met', 'metfilter', 'jetid', 'antiak4btagMediumOppHem', 'n2ddt', 'trigger', 'lumimask' ], 'signal_noddt': [ 'noleptons', 'minjetkin', 'met', 'jetid', 'antiak4btagMediumOppHem', 'trigger', 'lumimask', 'metfilter' ], # 'muoncontrol': ['minjetkinmu', 'jetid', 'n2ddt', 'ak4btagMedium08', 'onemuon', 'muonkin', 'muonDphiAK8', 'muontrigger', 'lumimask', 'metfilter'], 'muoncontrol': [ 'onemuon', 'muonkin', 'muonDphiAK8', 'metfilter', 'minjetkinmu', 'jetid', 'ak4btagMedium08', 'n2ddt', 'muontrigger', 'lumimask' ], 'muoncontrol_noddt': [ 'onemuon', 'muonkin', 'muonDphiAK8', 'jetid', 'metfilter', 'minjetkinmu', 'jetid', 'ak4btagMedium08', 'muontrigger', 'lumimask' ], 'wtag': [ 'onemuon', 'tightMuon', 'minjetkinw', 'jetid', 'met40p', 'metfilter', 'ptrecoW200', 'ak4btagOld', 'muontrigger', 'lumimask' ], 'wtag0': [ 'onemuon', 'tightMuon', 'met40p', 'metfilter', 'ptrecoW200', 'ak4btagOld', 'muontrigger', 'lumimask' ], 'wtag2': [ 'onemuon', 'tightMuon', 'minjetkinw', 'jetid', 'ak4btagMediumOppHem', 'met40p', 'metfilter', 'ptrecoW200', 'ak4btagOld', 'muontrigger', 'lumimask' ], 'noselection': [], } def normalize(val, cut): if cut is None: ar = ak.to_numpy(ak.fill_none(val, np.nan)) return ar else: ar = ak.to_numpy(ak.fill_none(val[cut], np.nan)) return ar import time tic = time.time() if shift_name is None: for region, cuts in regions.items(): allcuts = set([]) cut = selection.all(*allcuts) output['cutflow_msd'].fill(region=region, genflavor=normalize( genflavor, None), cut=0, weight=weights.weight(), msd=normalize(msd_matched, None)) output['cutflow_eta'].fill(region=region, genflavor=normalize(genflavor, cut), cut=0, weight=weights.weight()[cut], eta=normalize( candidatejet.eta, cut)) output['cutflow_pt'].fill(region=region, genflavor=normalize(genflavor, cut), cut=0, weight=weights.weight()[cut], pt=normalize(candidatejet.pt, cut)) for i, cut in enumerate(cuts + ['ddcvbpass', 'ddcpass']): allcuts.add(cut) cut = selection.all(*allcuts) output['cutflow_msd'].fill(region=region, genflavor=normalize( genflavor, cut), cut=i + 1, weight=weights.weight()[cut], msd=normalize(msd_matched, cut)) output['cutflow_eta'].fill( region=region, genflavor=normalize(genflavor, cut), cut=i + 1, weight=weights.weight()[cut], eta=normalize(candidatejet.eta, cut)) output['cutflow_pt'].fill( region=region, genflavor=normalize(genflavor, cut), cut=i + 1, weight=weights.weight()[cut], pt=normalize(candidatejet.pt, cut)) if self._evtVizInfo and 'ddcpass' in allcuts and isRealData and region == 'signal': if 'event' not in events.fields: continue _cut = selection.all(*allcuts, '_strict_mass', '_high_score') # _cut = selection.all('_strict_mass'') output['to_check'][ 'mass'] += processor.column_accumulator( normalize(msd_matched, _cut)) nfatjet = ak.sum( ((fatjets.pt > 200) & (abs(fatjets.eta) < 2.5) & fatjets.isTight), axis=1) output['to_check'][ 'njet'] += processor.column_accumulator( normalize(nfatjet, _cut)) output['to_check'][ 'fname'] += processor.column_accumulator( np.array([events.metadata['filename']] * len(normalize(msd_matched, _cut)))) output['to_check'][ 'event'] += processor.column_accumulator( normalize(events.event, _cut)) output['to_check'][ 'luminosityBlock'] += processor.column_accumulator( normalize(events.luminosityBlock, _cut)) output['to_check'][ 'run'] += processor.column_accumulator( normalize(events.run, _cut)) if shift_name is None: systematics = [None] + list(weights.variations) else: systematics = [shift_name] def fill(region, systematic, wmod=None): selections = regions[region] cut = selection.all(*selections) sname = 'nominal' if systematic is None else systematic if wmod is None: if systematic in weights.variations: weight = weights.weight(modifier=systematic)[cut] else: weight = weights.weight()[cut] else: weight = weights.weight()[cut] * wmod[cut] output['templates'].fill( region=region, systematic=sname, runid=runmap(events.run)[cut], genflavor=normalize(genflavor, cut), pt=normalize(candidatejet.pt, cut), msd=normalize(msd_matched, cut), ddb=normalize(bvl, cut), ddc=normalize(cvl, cut), ddcvb=normalize(cvb, cut), weight=weight, ) if region in [ 'wtag', 'wtag0', 'wtag2', 'wtag3', 'wtag4', 'wtag5', 'wtag6', 'wtag7', 'noselection' ]: # and sname in ['nominal', 'pileup_weightDown', 'pileup_weightUp', 'jet_triggerDown', 'jet_triggerUp']: output['wtag'].fill( region=region, systematic=sname, genflavor=normalize(genflavor, cut), pt=normalize(candidatejet.pt, cut), msd=normalize(msd_matched, cut), n2ddt=normalize(candidatejet.n2ddt, cut), ddc=normalize(cvl, cut), ddcvb=normalize(cvb, cut), weight=weight, ) # if region in ['signal', 'noselection']: # output['etaphi'].fill( # region=region, # systematic=sname, # runid=runmap(events.run)[cut], # genflavor=normalize(genflavor, cut), # pt=normalize(candidatejet.pt, cut), # eta=normalize(candidatejet.eta, cut), # phi=normalize(candidatejet.phi, cut), # ddc=normalize(cvl, cut), # ddcvb=normalize(cvb, cut), # ), if not isRealData: if wmod is not None: _custom_weight = events.genWeight[cut] * wmod[cut] else: _custom_weight = np.ones_like(weight) output['genresponse_noweight'].fill( region=region, systematic=sname, pt=normalize(candidatejet.pt, cut), genpt=normalize(genBosonPt, cut), weight=_custom_weight, ) output['genresponse'].fill( region=region, systematic=sname, pt=normalize(candidatejet.pt, cut), genpt=normalize(genBosonPt, cut), weight=weight, ) if systematic is None: output['signal_opt'].fill( region=region, genflavor=normalize(genflavor, cut), ddc=normalize(cvl, cut), ddcvb=normalize(cvb, cut), msd=normalize(msd_matched, cut), weight=weight, ) output['signal_optb'].fill( region=region, genflavor=normalize(genflavor, cut), ddb=normalize(bvl, cut), msd=normalize(msd_matched, cut), weight=weight, ) for region in regions: cut = selection.all(*(set(regions[region]) - {'n2ddt'})) if shift_name is None: output['nminus1_n2ddt'].fill( region=region, n2ddt=normalize(candidatejet.n2ddt, cut), weight=weights.weight()[cut], ) for systematic in systematics: if isRealData and systematic is not None: continue fill(region, systematic) if shift_name is None and 'GluGluH' in dataset and 'LHEWeight' in events.fields: for i in range(9): fill(region, 'LHEScale_%d' % i, events.LHEScaleWeight[:, i]) for c in events.LHEWeight.fields[1:]: fill(region, 'LHEWeight_%s' % c, events.LHEWeight[c]) toc = time.time() output["filltime"] = toc - tic if shift_name is None: output["weightStats"] = weights.weightStatistics return {dataset: output}
def process(self, events): def normalize(val, cut): return ak.to_numpy(ak.fill_none( val[cut], np.nan)) #val[cut].pad(1, clip=True).fillna(0).flatten() def fill(region, cuts, systematic=None, wmod=None): print('filling %s' % region) selections = cuts cut = selection.all(*selections) if 'signal' in region: weight = weights_signal.weight()[cut] elif 'muonCR' in region: weight = weights_muonCR.weight()[cut] elif 'VtaggingCR' in region: weight = weights_VtaggingCR.weight()[cut] output['templates'].fill( dataset=dataset, region=region, pt=normalize(candidatejet.pt, cut), msd=normalize(candidatejet.msdcorr, cut), n2ddt=normalize(candidatejet.n2ddt, cut), #gruddt=normalize(candidatejet.gruddt, cut), in_v3_ddt=normalize(candidatejet.in_v3_ddt, cut), hadW=normalize(candidatejet.nmatcheddau, cut), weight=weight, ), output['event'].fill( dataset=dataset, region=region, MET=events.MET.pt[cut], #nJet=fatjets.counts[cut], nPFConstituents=normalize(candidatejet.nPFConstituents, cut), weight=weight, ), output['deepAK8'].fill( dataset=dataset, region=region, deepTagMDWqq=normalize(candidatejet.deepTagMDWqq, cut), deepTagMDZqq=normalize(candidatejet.deepTagMDZqq, cut), msd=normalize(candidatejet.msdcorr, cut), #genflavor=genflavor[cut], weight=weight, ), output['in_v3'].fill( dataset=dataset, region=region, #genflavor=genflavor[cut], in_v3=normalize(candidatejet.in_v3, cut), n2=normalize(candidatejet.n2b1, cut), gru=normalize(candidatejet.gru, cut), weight=weight, ), if 'muonCR' in dataset or 'VtaggingCR' in dataset: output['muon'].fill( dataset=dataset, region=region, mu_pt=normalize(candidatemuon.pt, cut), mu_eta=normalize(candidatemuon.eta, cut), mu_pfRelIso04_all=normalize(candidatemuon.pfRelIso04_all, cut), weight=weight, ), #common jet kinematics gru = events.GRU IN = events.IN fatjets = events.FatJet fatjets['msdcorr'] = corrected_msoftdrop(fatjets) fatjets['qcdrho'] = 2 * np.log(fatjets.msdcorr / fatjets.pt) fatjets['gruddt'] = gru.v25 - shift( fatjets, algo='gruddt', year='2017') fatjets['gru'] = gru.v25 fatjets['in_v3'] = IN.v3 fatjets['in_v3_ddt'] = IN.v3 - shift( fatjets, algo='inddt', year='2017') fatjets['in_v3_ddt_90pctl'] = IN.v3 - shift( fatjets, algo='inddt90pctl', year='2017') fatjets['n2ddt'] = fatjets.n2b1 - n2ddt_shift(fatjets, year='2017') fatjets['nmatcheddau'] = TTsemileptonicmatch(events) dataset = events.metadata['dataset'] print('process dataset', dataset) isRealData = not hasattr(events, 'genWeight') output = self.accumulator.identity() if (len(events) == 0): return output selection = PackedSelection('uint64') weights_signal = Weights(len(events)) weights_muonCR = Weights(len(events)) weights_VtaggingCR = Weights(len(events)) if not isRealData: output['sumw'][dataset] += ak.sum(events.genWeight) ####################### if 'signal' in self._region: if isRealData: trigger_fatjet = np.zeros(len(events), dtype='bool') for t in self._triggers[self._year]: try: trigger_fatjet = trigger_fatjet | events.HLT[t] except: print('trigger %s not available' % t) continue else: trigger_fatjet = np.ones(len(events), dtype='bool') fatjets["genMatchFull"] = VQQgenmatch(events) candidatejet = ak.firsts(fatjets) candidatejet["genMatchFull"] = VQQgenmatch(events) nelectrons = ak.sum( (events.Electron.pt > 10.) & (abs(events.Electron.eta) < 2.5) & (events.Electron.cutBased >= events.Electron.VETO), axis=1, ) nmuons = ak.sum( (events.Muon.pt > 10) & (abs(events.Muon.eta) < 2.1) & (events.Muon.pfRelIso04_all < 0.4) & (events.Muon.looseId), axis=1, ) ntaus = ak.sum( (events.Tau.pt > 20.) & (events.Tau.idDecayMode) & (events.Tau.rawIso < 5) & (abs(events.Tau.eta) < 2.3), axis=1, ) cuts = { "S_fatjet_trigger": trigger_fatjet, "S_pt": candidatejet.pt > 525, "S_eta": (abs(candidatejet.eta) < 2.5), "S_msdcorr": (candidatejet.msdcorr > 40), "S_rho": ((candidatejet.qcdrho > -5.5) & (candidatejet.qcdrho < -2.)), "S_jetid": (candidatejet.isTight), "S_VQQgenmatch": (candidatejet.genMatchFull), "S_noelectron": (nelectrons == 0), "S_nomuon": (nmuons == 0), "S_notau": (ntaus == 0), } for name, cut in cuts.items(): print(name, cut) selection.add(name, cut) if isRealData: genflavor = 0 #candidatejet.pt.zeros_like().pad(1, clip=True).fillna(-1).flatten() if not isRealData: weights_signal.add('genweight', events.genWeight) #add_pileup_weight(weights_signal, events.Pileup.nPU, self._year, dataset) add_jetTriggerWeight(weights_signal, candidatejet.msdcorr, candidatejet.pt, self._year) bosons = getBosons(events.GenPart) genBosonPt = ak.fill_none(ak.firsts(bosons.pt), 0) add_VJets_NLOkFactor(weights_signal, genBosonPt, self._year, dataset) #genflavor = matchedBosonFlavor(candidatejet, bosons).pad(1, clip=True).fillna(-1).flatten() allcuts_signal = set() output['cutflow_signal'][dataset]['none'] += float( weights_signal.weight().sum()) for cut in cuts: allcuts_signal.add(cut) output['cutflow_signal'][dataset][cut] += float( weights_signal.weight()[selection.all( *allcuts_signal)].sum()) fill('signal', cuts.keys()) ####################### if 'muonCR' in self._region: if isRealData: trigger_muon = np.zeros(len(events), dtype='bool') for t in self._muontriggers[self._year]: trigger_muon = trigger_muon | events.HLT[t] else: trigger_muon = np.ones(len(events), dtype='bool') candidatejet = ak.firsts(fatjets) candidatemuon = events.Muon[:, :5] jets = events.Jet[((events.Jet.pt > 50.) & (abs(events.Jet.eta) < 2.5) & (events.Jet.isTight))][:, :4] dphi = abs(jets.delta_phi(candidatejet)) ak4_away = jets[(dphi > 0.8)] nelectrons = ak.sum( (events.Electron.pt > 10.) & (abs(events.Electron.eta) < 2.5) & (events.Electron.cutBased >= events.Electron.VETO), axis=1, ) nmuons = ak.sum( (events.Muon.pt > 10) & (abs(events.Muon.eta) < 2.4) & (events.Muon.pfRelIso04_all < 0.25) & (events.Muon.looseId), axis=1, ) ntaus = ak.sum( (events.Tau.pt > 20.) & (events.Tau.idDecayMode) & (events.Tau.rawIso < 5) & (abs(events.Tau.eta) < 2.3) & (events.Tau.idMVAoldDM2017v1 >= 16), axis=1, ) cuts = { "CR1_muon_trigger": trigger_muon, "CR1_jet_pt": (candidatejet.pt > 525), "CR1_jet_eta": (abs(candidatejet.eta) < 2.5), "CR1_jet_msd": (candidatejet.msdcorr > 40), "CR1_jet_rho": ((candidatejet.qcdrho > -5.5) & (candidatejet.qcdrho < -2.)), "CR1_mu_pt": ak.any(candidatemuon.pt > 55, axis=1), "CR1_mu_eta": ak.any(abs(candidatemuon.eta) < 2.1, axis=1), "CR1_mu_IDLoose": ak.any(candidatemuon.looseId, axis=1), "CR1_mu_isolationTight": ak.any(candidatemuon.pfRelIso04_all < 0.15, axis=1), "CR1_muonDphiAK8": ak.any( abs(candidatemuon.delta_phi(candidatejet)) > 2 * np.pi / 3, axis=1), "CR1_ak4btagMedium08": (ak.max(ak4_away.btagCSVV2, axis=1, mask_identity=False) > BTagEfficiency.btagWPs[self._year]['medium'] ), #(ak4_away.btagCSVV2.max() > 0.8838), "CR1_noelectron": (nelectrons == 0), "CR1_onemuon": (nmuons == 1), "CR1_notau": (ntaus == 0), } for name, cut in cuts.items(): selection.add(name, cut) if isRealData: genflavor = 0 #candidatejet.pt.zeros_like().pad(1, clip=True).fillna(-1).flatten() if not isRealData: weights_muonCR.add('genweight', events.genWeight) #add_pileup_weight(weights_muonCR, events.Pileup.nPU, self._year, dataset) #add_singleMuTriggerWeight(weights, candidatejet.msdcorr, candidatejet.pt, self._year) bosons = getBosons(events.GenPart) genBosonPt = ak.fill_none(ak.firsts(bosons.pt), 0) #add_VJets_NLOkFactor(weights, genBosonPt, self._year, dataset) #genflavor = matchedBosonFlavor(candidatejet, bosons).pad(1, clip=True).fillna(-1).flatten() allcuts_ttbar_muoncontrol = set() output['cutflow_muonCR'][dataset]['none'] += float( weights_muonCR.weight().sum()) for cut in cuts: allcuts_ttbar_muoncontrol.add(cut) output['cutflow_muonCR'][dataset][cut] += float( weights_muonCR.weight()[selection.all( *allcuts_ttbar_muoncontrol)].sum()) fill('muonCR', cuts.keys()) ####################### if 'VtaggingCR' in self._region: if isRealData: trigger_muon = np.zeros(len(events), dtype='bool') for t in self._muontriggers[self._year]: trigger_muon = trigger_muon | events.HLT[t] else: trigger_muon = np.ones(len(events), dtype='bool') candidatejet = ak.firsts(fatjets) candidatemuon = ak.firsts(events.Muon) jets = events.Jet[((events.Jet.pt > 30.) & (abs(events.Jet.eta) < 2.4))][:, :4] dr_ak4_ak8 = jets.delta_r(candidatejet) dr_ak4_muon = jets.delta_r(candidatemuon) ak4_away = jets[(dr_ak4_ak8 > 0.8)] # & (dr_ak4_muon > 0.4)] mu_p4 = ak.zip( { "pt": ak.fill_none(candidatemuon.pt, 0), "eta": ak.fill_none(candidatemuon.eta, 0), "phi": ak.fill_none(candidatemuon.phi, 0), "mass": ak.fill_none(candidatemuon.mass, 0), }, with_name="PtEtaPhiMLorentzVector") met_p4 = ak.zip( { "pt": ak.from_iter([[v] for v in events.MET.pt]), "eta": ak.from_iter([[v] for v in np.zeros(len(events))]), "phi": ak.from_iter([[v] for v in events.MET.phi]), "mass": ak.from_iter([[v] for v in np.zeros(len(events))]), }, with_name="PtEtaPhiMLorentzVector") Wleptoniccandidate = mu_p4 + met_p4 nelectrons = ak.sum( ((events.Electron.pt > 10.) & (abs(events.Electron.eta) < 2.5) & (events.Electron.cutBased >= events.Electron.VETO)), axis=1, ) n_tight_muon = ak.sum( ((events.Muon.pt > 53) & (abs(events.Muon.eta) < 2.1) & (events.Muon.tightId)), axis=1, ) n_loose_muon = ak.sum( ((events.Muon.pt > 20) & (events.Muon.looseId) & (abs(events.Muon.eta) < 2.4)), axis=1, ) ntaus = ak.sum( ((events.Tau.pt > 20.) & (events.Tau.idDecayMode) & (events.Tau.rawIso < 5) & (abs(events.Tau.eta) < 2.3) & (events.Tau.idMVAoldDM2017v1 >= 16)), axis=1, ) cuts = { "CR2_muon_trigger": trigger_muon, "CR2_jet_pt": (candidatejet.pt > 200), "CR2_jet_eta": (abs(candidatejet.eta) < 2.5), "CR2_jet_msd": (candidatejet.msdcorr > 40), "CR2_mu_pt": candidatemuon.pt > 53, "CR2_mu_eta": (abs(candidatemuon.eta) < 2.1), "CR2_mu_IDTight": candidatemuon.tightId, "CR2_mu_isolationTight": (candidatemuon.pfRelIso04_all < 0.15), "CR2_muonDphiAK8": abs(candidatemuon.delta_phi(candidatejet)) > 2 * np.pi / 3, "CR2_ak4btagMedium08": (ak.max(ak4_away.btagCSVV2, axis=1, mask_identity=False) > BTagEfficiency.btagWPs[self._year]['medium']), "CR2_leptonicW": ak.flatten(Wleptoniccandidate.pt > 200), "CR2_MET": (events.MET.pt > 40.), "CR2_noelectron": (nelectrons == 0), "CR2_one_tightMuon": (n_tight_muon == 1), "CR2_one_looseMuon": (n_loose_muon == 1), #"CR2_notau" : (ntaus==0), } for name, cut in cuts.items(): print(name, cut) selection.add(name, cut) #weights.add('metfilter', events.Flag.METFilters) if isRealData: genflavor = 0 #candidatejet.pt.zeros_like().pad(1, clip=True).fillna(-1).flatten() if not isRealData: weights_VtaggingCR.add('genweight', events.genWeight) #add_pileup_weight(weights_VtaggingCR, events.Pileup.nPU, self._year, dataset) #add_singleMuTriggerWeight(weights, abs(candidatemuon.eta), candidatemuon.pt, self._year) bosons = getBosons(events.GenPart) genBosonPt = ak.fill_none(ak.firsts(bosons.pt), 0) #add_VJets_NLOkFactor(weights, genBosonPt, self._year, dataset) #genflavor = matchedBosonFlavor(candidatejet, bosons).pad(1, clip=True).fillna(-1).flatten() #b-tag weights allcuts_vselection = set() output['cutflow_VtaggingCR'][dataset]['none'] += float( weights_VtaggingCR.weight().sum()) for cut in cuts: allcuts_vselection.add(cut) output['cutflow_VtaggingCR'][dataset][cut] += float( weights_VtaggingCR.weight()[selection.all( *allcuts_vselection)].sum()) fill('VtaggingCR', cuts.keys()) return output