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 #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] # 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 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