def test_highlevel(): array = ak.Array([[3.3, 1.1, 5.5, 1.1, 4.4], [4.4, 2.2, 1.1, 6.6], [2.2, 3.3, -1.1]]) assert ak.argmin(array) == 11 assert ak.argmax(array) == 8 assert ak.to_list(ak.argmin(array, axis=0)) == [2, 0, 2, 0, 0] assert ak.to_list(ak.argmax(array, axis=0)) == [1, 2, 0, 1, 0] assert ak.to_list(ak.argmin(array, axis=1)) == [1, 2, 2] assert ak.to_list(ak.argmax(array, axis=1)) == [2, 3, 1]
def test_argmin_argmax(): array = ak.Array( [ [ [ np.datetime64("2022"), np.datetime64("2023"), np.datetime64("2025") ], [], [np.datetime64("2027"), np.datetime64("2011")], [np.datetime64("2013")], ], [], [[np.datetime64("2017"), np.datetime64("2019")], [np.datetime64("2023")]], ], check_valid=True, ) assert ak.argmin(array) == 4 assert ak.argmax(array) == 3 assert ak.to_list(ak.argmin(array, axis=0)) == [[1, 1, 0], [1], [0, 0], [0]] assert ak.to_list(ak.argmax(array, axis=0)) == [[0, 0, 0], [1], [0, 0], [0]] assert ak.to_list(ak.argmin(array, axis=1)) == [[3, 2, 0], [], [0, 0]] assert ak.to_list(ak.argmax(array, axis=1)) == [[2, 0, 0], [], [1, 0]] array = ak.from_iter( [ [[ np.datetime64("2021-01-20"), np.datetime64("2021-01-10"), np.datetime64("2021-01-30"), ]], [[]], [None, None, None], [[ np.datetime64("2021-01-14"), np.datetime64("2021-01-15"), np.datetime64("2021-01-16"), ]], ], highlevel=False, ) assert ak.to_list(array.argmin(axis=2)) == [[1], [None], [None, None, None], [0]]
def run_deltar_matching(store, target, drname='deltaR', radius=0.4, unique=False, sort=False): """ Running a delta R matching of some object collection "store" of dimension NxS with some target collection "target" of dimension NxT, The return object will have dimension NxSxT' where objects in the T' contain all "target" objects within the delta R radius. The delta R between the store and target object will be stored in the field `deltaR`. If the unique flag is turned on, then objects in the target collection will only be associated to the closest object. If the sort flag is turned on, then the target collection will be sorted according to the computed `deltaR`. """ _, target = ak.unzip(ak.cartesian([store.eta, target], nested=True)) target[drname] = delta_r(store, target) if unique: # Additional filtering t_index = ak.argmin(target[drname], axis=-2) s_index = ak.local_index(store.eta, axis=-1) _, t_index = ak.unzip(ak.cartesian([s_index, t_index], nested=True)) target = target[s_index == t_index] # Cutting on the computed delta R target = target[target[drname] < radius] # Sorting according to the computed delta R if sort: idx = ak.argsort(target[drname], axis=-1) target = target[idx] return target
def match(first, second, **kwargs): if 'dR_cutoff' not in kwargs: raise Exception("Please specify dR cutoff for matching!") dR_cutoff = kwargs.pop('dR_cutoff', 0.3) return_match_properties = kwargs.pop('return_match_properties', False) etas = ak.cartesian({ 'first': first.eta, 'second': second.eta }, axis=1, nested=True) phis = ak.cartesian({ 'first': first.phi, 'second': second.phi }, axis=1, nested=True) dR, deta, dphi = delta_r(etas['first'], etas['second'], phis['first'], phis['second']) min_idx = ak.argmin(dR, axis=2) match_properties = { 'pt': second.pt[min_idx], 'eta': second.eta[min_idx], 'phi': second.phi[min_idx], } if return_match_properties: return ak.any(dR < dR_cutoff, axis=2), match_properties else: return ak.any(dR < dR_cutoff, axis=2)
def process(self, events): events["Electron", "pdgId"] = -11 * events.Electron.charge events["Muon", "pdgId"] = -13 * events.Muon.charge events["leptons"] = ak.concatenate( [events.Electron, events.Muon], axis=1, ) events = events[ak.num(events.leptons) >= 3] pair = ak.argcombinations(events.leptons, 2, fields=["l1", "l2"]) pair = pair[( events.leptons[pair.l1].pdgId == -events.leptons[pair.l2].pdgId)] with np.errstate(invalid="ignore"): pair = pair[ak.singletons( ak.argmin( abs((events.leptons[pair.l1] + events.leptons[pair.l2]).mass - 91.2), axis=1, ))] events = events[ak.num(pair) > 0] pair = pair[ak.num(pair) > 0][:, 0] l3 = ak.local_index(events.leptons) l3 = l3[(l3 != pair.l1) & (l3 != pair.l2)] l3 = l3[ak.argmax(events.leptons[l3].pt, axis=1, keepdims=True)] l3 = events.leptons[l3][:, 0] mt = np.sqrt(2 * l3.pt * events.MET.pt * (1 - np.cos(events.MET.delta_phi(l3)))) return (hist.Hist.new.Reg( 100, 0, 200, name="mt", label=r"$\ell$-MET transverse mass [GeV]").Double().fill(mt))
def process(self, events): jets = ak.zip( {k: getattr(events.Jet, k) for k in ["x", "y", "z", "t", "btag"]}, with_name="LorentzVector", behavior=events.Jet.behavior, ) trijet = ak.combinations(jets, 3, fields=["j1", "j2", "j3"]) trijet["p4"] = trijet.j1 + trijet.j2 + trijet.j3 trijet = ak.flatten(trijet[ak.singletons( ak.argmin(abs(trijet.p4.mass - 172.5), axis=1))]) maxBtag = np.maximum( trijet.j1.btag, np.maximum( trijet.j2.btag, trijet.j3.btag, ), ) return { "trijetpt": hist.Hist.new.Reg(100, 0, 200, name="pt3j", label="Trijet $p_{T}$ [GeV]").Double().fill( trijet.p4.pt), "maxbtag": hist.Hist.new.Reg( 100, 0, 1, name="btag", label="Max jet b-tag score").Double().fill(maxBtag), }
def test_issue434(): a = ak.Array([[0.0, 1.1, 2.2], [3.3, 4.4], [5.5]]) b = ak.Array([[9.9, 8.8, 7.7], [6.6, 5.5], [4.4]]) assert ak.to_list(b[ak.argmin(a, axis=1, keepdims=True)]) == [[9.9], [6.6], [4.4]] assert ak.to_list(b[ak.argmax(a, axis=1, keepdims=True)]) == [[7.7], [5.5], [4.4]]
def mask_nearest(ref_eta, ref_phi, obj_eta, obj_phi, threshold=10): """ This function returns the masking to the closest object in dR among the "obj" wrt the "ref", satisfying a minimum threshold """ d_eta = ref_eta - obj_eta d_phi = (ref_phi - obj_phi + PI) % (2 * PI) - PI drs = numpy.sqrt(d_eta**2 + d_phi**2) min_drs = awkward.argmin(drs, axis=1, keepdims=True) min_drs_mask = drs[min_drs] < threshold return min_drs_mask
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_minmax(): assert ak.min(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]])) == 1 + 5j assert ak.max(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]])) == 3 + 3j assert ak.min(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]]), axis=1).tolist() == [ 1 + 5j, None, 3 + 3j, ] assert ak.max(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]]), axis=1).tolist() == [ 2 + 4j, None, 3 + 3j, ] assert ak.argmin(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]]), axis=1).tolist() == [0, None, 0] assert ak.argmax(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]]), axis=1).tolist() == [1, None, 0]
def test_IndexedOptionArray(): content = ak.Array([1.1, 2.2, 3.3, 4.4, 5.5]).layout index = ak.layout.Index64(np.array([4, 2, -1, -1, 1, 0, 1])) array = ak.Array(ak.layout.IndexedOptionArray64(index, content)) assert array.tolist() == [5.5, 3.3, None, None, 2.2, 1.1, 2.2] assert ak.min(array, axis=0) == 1.1 assert ak.argmin(array, axis=0) == 5 assert ak.argmin(ak.Array([[2.2, 1.1], [None, 3.3], [2.2, 1.1]]), axis=-1).tolist() == [1, 1, 1] assert ak.argmin(ak.Array([[2.2, 1.1], [None, 3.3], [2.2, None, 1.1]]), axis=-1).tolist() == [1, 1, 2] assert ak.argmin(ak.Array([[2.2, 1.1], [3.3, None], [2.2, None, 1.1]]), axis=-1).tolist() == [1, 0, 2] assert ak.argmin(ak.Array([[2.2, 1.1, 0.0], [], [None, 0.5], [2, 1]]), axis=0).tolist() == [3, 2, 0] assert ak.argmin(ak.Array([[2.2, 1.1, 0.0], [], [0.5, None], [2, 1]]), axis=0).tolist() == [2, 3, 0] assert ak.argmin(ak.Array([[2.2, 1.1, 0.0], [0.5, None], [], [2, 1]]), axis=0).tolist() == [1, 3, 0]
def process(self, events): # Dataset parameters dataset = events.metadata['dataset'] year = self._samples[dataset]['year'] xsec = self._samples[dataset]['xsec'] sow = self._samples[dataset]['nSumOfWeights' ] isData = self._samples[dataset]['isData'] datasets = ['SingleMuon', 'SingleElectron', 'EGamma', 'MuonEG', 'DoubleMuon', 'DoubleElectron'] for d in datasets: if d in dataset: dataset = dataset.split('_')[0] # Inittialize objects met = events.GenMET e = events.GenPart[abs(events.GenPart.pdgId)==11] mu = events.GenPart[abs(events.GenPart.pdgId)==13] tau = events.GenPart[abs(events.GenPart.pdgId)==15] j = events.GenJet leading_mu = mu[ak.argmax(mu.pt,axis=-1,keepdims=True)] leading_e = e[ak.argmax(e.pt,axis=-1,keepdims=True)] nElec = ak.num(e) nMuon = ak.num(mu) nTau = ak.num(tau) twoLeps = (nElec+nMuon) == 2 threeLeps = (nElec+nMuon) == 3 twoElec = (nElec == 2) twoMuon = (nMuon == 2) e0 = e[ak.argmax(e.pt,axis=-1,keepdims=True)] m0 = mu[ak.argmax(mu.pt,axis=-1,keepdims=True)] elecs = e[ak.argsort(e.pt, ascending=False)] muons = mu[ak.argsort(mu.pt, ascending=False)] e1 = elecs e2 = elecs m1 = muons m2 = muons # Jet selection jetptname = 'pt_nom' if hasattr(j, 'pt_nom') else 'pt' njets = ak.num(j) ht = ak.sum(j.pt,axis=-1) jets = j[ak.argsort(j.pt, ascending=False)] j0 = j[ak.argmax(j.pt,axis=-1,keepdims=True)] j1 = jets j2 = jets j3 = jets nbtags = ak.num(j[abs(j.hadronFlavour)==5]) ################################################################## ### 2 same-sign leptons ################################################################## # emu singe = e [(nElec==1)&(nMuon==1)&(e .pt>-1)] singm = mu[(nElec==1)&(nMuon==1)&(mu.pt>-1)] em = ak.cartesian({"e":singe,"m":singm}) emSSmask = (em.e.pdgId*em.m.pdgId>0) emSS = em[emSSmask] nemSS = len(ak.flatten(emSS)) year = 2018 lepSF_emSS = GetLeptonSF(mu.pt, mu.eta, 'm', e.pt, e.eta, 'e', year=year) # ee and mumu # pt>-1 to preserve jagged dimensions ee = e [(nElec==2)&(nMuon==0)&(e.pt>-1)] mm = mu[(nElec==0)&(nMuon==2)&(mu.pt>-1)] eepairs = ak.combinations(ee, 2, fields=["e0","e1"]) eeSSmask = (eepairs.e0.pdgId*eepairs.e1.pdgId>0) eeonZmask = (np.abs((eepairs.e0+eepairs.e1).mass-91.2)<10) eeoffZmask = (eeonZmask==0) mmpairs = ak.combinations(mm, 2, fields=["m0","m1"]) mmSSmask = (mmpairs.m0.pdgId*mmpairs.m1.pdgId>0) mmonZmask = (np.abs((mmpairs.m0+mmpairs.m1).mass-91.2)<10) mmoffZmask = (mmonZmask==0) eeSSonZ = eepairs[eeSSmask & eeonZmask] eeSSoffZ = eepairs[eeSSmask & eeoffZmask] mmSSonZ = mmpairs[mmSSmask & mmonZmask] mmSSoffZ = mmpairs[mmSSmask & mmoffZmask] neeSS = len(ak.flatten(eeSSonZ)) + len(ak.flatten(eeSSoffZ)) nmmSS = len(ak.flatten(mmSSonZ)) + len(ak.flatten(mmSSoffZ)) lepSF_eeSS = GetLeptonSF(eepairs.e0.pt, eepairs.e0.eta, 'e', eepairs.e1.pt, eepairs.e1.eta, 'e', year=year) lepSF_mumuSS = GetLeptonSF(mmpairs.m0.pt, mmpairs.m0.eta, 'm', mmpairs.m1.pt, mmpairs.m1.eta, 'm', year=year) print('Same-sign events [ee, emu, mumu] = [%i, %i, %i]'%(neeSS, nemSS, nmmSS)) # Cuts eeSSmask = (ak.num(eeSSmask[eeSSmask])>0) mmSSmask = (ak.num(mmSSmask[mmSSmask])>0) eeonZmask = (ak.num(eeonZmask[eeonZmask])>0) eeoffZmask = (ak.num(eeoffZmask[eeoffZmask])>0) mmonZmask = (ak.num(mmonZmask[mmonZmask])>0) mmoffZmask = (ak.num(mmoffZmask[mmoffZmask])>0) emSSmask = (ak.num(emSSmask[emSSmask])>0) ################################################################## ### 3 leptons ################################################################## # eem muon_eem = mu[(nElec==2)&(nMuon==1)&(mu.pt>-1)] elec_eem = e[(nElec==2)&(nMuon==1)&( e.pt>-1)] ee_eem = ak.combinations(elec_eem, 2, fields=["e0", "e1"]) ee_eemZmask = (ee_eem.e0.pdgId*ee_eem.e1.pdgId<1)&(np.abs((ee_eem.e0+ee_eem.e1).mass-91.2)<10) ee_eemOffZmask = (ee_eem.e0.pdgId*ee_eem.e1.pdgId<1)&(np.abs((ee_eem.e0+ee_eem.e1).mass-91.2)>10) ee_eemZmask = (ak.num(ee_eemZmask[ee_eemZmask])>0) ee_eemOffZmask = (ak.num(ee_eemOffZmask[ee_eemOffZmask])>0) eepair_eem = (ee_eem.e0+ee_eem.e1) trilep_eem = eepair_eem+muon_eem #ak.cartesian({"e0":ee_eem.e0,"e1":ee_eem.e1, "m":muon_eem}) lepSF_eem = GetLeptonSF(ee_eem.e0.pt, ee_eem.e0.eta, 'e', ee_eem.e1.pt, ee_eem.e1.eta, 'e', mu.pt, mu.eta, 'm', year) # mme muon_mme = mu[(nElec==1)&(nMuon==2)&(mu.pt>-1)] elec_mme = e[(nElec==1)&(nMuon==2)&( e.pt>-1)] mm_mme = ak.combinations(muon_mme, 2, fields=["m0", "m1"]) mm_mmeZmask = (mm_mme.m0.pdgId*mm_mme.m1.pdgId<1)&(np.abs((mm_mme.m0+mm_mme.m1).mass-91.2)<10) mm_mmeOffZmask = (mm_mme.m0.pdgId*mm_mme.m1.pdgId<1)&(np.abs((mm_mme.m0+mm_mme.m1).mass-91.2)>10) mm_mmeZmask = (ak.num(mm_mmeZmask[mm_mmeZmask])>0) mm_mmeOffZmask = (ak.num(mm_mmeOffZmask[mm_mmeOffZmask])>0) mmpair_mme = (mm_mme.m0+mm_mme.m1) trilep_mme = mmpair_mme+elec_mme mZ_mme = mmpair_mme.mass mZ_eem = eepair_eem.mass m3l_eem = trilep_eem.mass m3l_mme = trilep_mme.mass lepSF_mme = GetLeptonSF(mm_mme.m0.pt, mm_mme.m0.eta, 'm', mm_mme.m1.pt, mm_mme.m1.eta, 'm', e.pt, e.eta, 'e', year) # eee and mmm eee = e[(nElec==3)&(nMuon==0)&( e.pt>-1)] mmm = mu[(nElec==0)&(nMuon==3)&(mu.pt>-1)] eee_leps = ak.combinations(eee, 3, fields=["e0", "e1", "e2"]) mmm_leps = ak.combinations(mmm, 3, fields=["m0", "m1", "m2"]) ee_pairs = ak.combinations(eee, 2, fields=["e0", "e1"]) mm_pairs = ak.combinations(mmm, 2, fields=["m0", "m1"]) ee_pairs_index = ak.argcombinations(eee, 2, fields=["e0", "e1"]) mm_pairs_index = ak.argcombinations(mmm, 2, fields=["m0", "m1"]) lepSF_eee = GetLeptonSF(eee_leps.e0.pt, eee_leps.e0.eta, 'e', eee_leps.e1.pt, eee_leps.e1.eta, 'e', eee_leps.e2.pt, eee_leps.e2.eta, 'e', year) lepSF_mmm = GetLeptonSF(mmm_leps.m0.pt, mmm_leps.m0.eta, 'm', mmm_leps.m1.pt, mmm_leps.m1.eta, 'm', mmm_leps.m2.pt, mmm_leps.m2.eta, 'm', year) mmSFOS_pairs = mm_pairs[(np.abs(mm_pairs.m0.pdgId) == np.abs(mm_pairs.m1.pdgId)) & (mm_pairs.m0.pdgId != mm_pairs.m1.pdgId)] offZmask_mm = ak.all(np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2)>10., axis=1, keepdims=True) & (ak.num(mmSFOS_pairs)>0) onZmask_mm = ak.any(np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2)<10., axis=1, keepdims=True) eeSFOS_pairs = ee_pairs[(np.abs(ee_pairs.e0.pdgId) == np.abs(ee_pairs.e1.pdgId)) & (ee_pairs.e0.pdgId != ee_pairs.e1.pdgId)] offZmask_ee = ak.all(np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2)>10, axis=1, keepdims=True) & (ak.num(eeSFOS_pairs)>0) onZmask_ee = ak.any(np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2)<10, axis=1, keepdims=True) # Create masks **for event selection** eeeOnZmask = (ak.num(onZmask_ee[onZmask_ee])>0) eeeOffZmask = (ak.num(offZmask_ee[offZmask_ee])>0) mmmOnZmask = (ak.num(onZmask_mm[onZmask_mm])>0) mmmOffZmask = (ak.num(offZmask_mm[offZmask_mm])>0) # Now we need to create masks for the leptons in order to select leptons from the Z boson candidate (in onZ categories) ZeeMask = ak.argmin(np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2),axis=1,keepdims=True) ZmmMask = ak.argmin(np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2),axis=1,keepdims=True) Zee = eeSFOS_pairs[ZeeMask] Zmm = mmSFOS_pairs[ZmmMask] eZ0= Zee.e0[ak.num(eeSFOS_pairs)>0] eZ1= Zee.e1[ak.num(eeSFOS_pairs)>0] eZ = eZ0+eZ1 mZ0= Zmm.m0[ak.num(mmSFOS_pairs)>0] mZ1= Zmm.m1[ak.num(mmSFOS_pairs)>0] mZ = mZ0+mZ1 mZ_eee = eZ.mass mZ_mmm = mZ.mass # And for the W boson ZmmIndices = mm_pairs_index[ZmmMask] ZeeIndices = ee_pairs_index[ZeeMask] eW = eee[~ZeeIndices.e0 | ~ZeeIndices.e1] mW = mmm[~ZmmIndices.m0 | ~ZmmIndices.m1] triElec = eee_leps.e0+eee_leps.e1+eee_leps.e2 triMuon = mmm_leps.m0+mmm_leps.m1+mmm_leps.m2 m3l_eee = triElec.mass m3l_mmm = triMuon.mass # Triggers trig_eeSS = passTrigger(events,'ee',isData,dataset) trig_mmSS = passTrigger(events,'mm',isData,dataset) trig_emSS = passTrigger(events,'em',isData,dataset) trig_eee = passTrigger(events,'eee',isData,dataset) trig_mmm = passTrigger(events,'mmm',isData,dataset) trig_eem = passTrigger(events,'eem',isData,dataset) trig_mme = passTrigger(events,'mme',isData,dataset) # MET filters # Weights genw = np.ones_like(events['MET_pt']) if isData else events['genWeight'] ### We need weights for: normalization, lepSF, triggerSF, pileup, btagSF... weights = {} for r in ['all', 'ee', 'mm', 'em', 'eee', 'mmm', 'eem', 'mme']: weights[r] = coffea.analysis_tools.Weights(len(events)) weights[r].add('norm',genw if isData else (xsec/sow)*genw) weights['ee'].add('lepSF_eeSS', lepSF_eeSS) weights['em'].add('lepSF_emSS', lepSF_emSS) weights['mm'].add('lepSF_mmSS', lepSF_mumuSS) weights['eee'].add('lepSF_eee', lepSF_eee) weights['mmm'].add('lepSF_mmm', lepSF_mmm) weights['mme'].add('lepSF_mme', lepSF_mme) weights['eem'].add('lepSF_eem', lepSF_eem) # Extract the EFT quadratic coefficients and optionally use them to calculate the coefficients on the w**2 quartic function # eft_coeffs is never Jagged so convert immediately to numpy for ease of use. eft_coeffs = ak.to_numpy(events['EFTfitCoefficients']) if hasattr(events, "EFTfitCoefficients") else None eft_w2_coeffs = efth.calc_w2_coeffs(eft_coeffs,self._dtype) if (self._do_errors and eft_coeffs is not None) else None # Selections and cuts selections = PackedSelection() channels2LSS = ['eeSSonZ', 'eeSSoffZ', 'mmSSonZ', 'mmSSoffZ', 'emSS'] selections.add('eeSSonZ', (eeonZmask)&(eeSSmask)&(trig_eeSS)) selections.add('eeSSoffZ', (eeoffZmask)&(eeSSmask)&(trig_eeSS)) selections.add('mmSSonZ', (mmonZmask)&(mmSSmask)&(trig_mmSS)) selections.add('mmSSoffZ', (mmoffZmask)&(mmSSmask)&(trig_mmSS)) selections.add('emSS', (emSSmask)&(trig_emSS)) channels3L = ['eemSSonZ', 'eemSSoffZ', 'mmeSSonZ', 'mmeSSoffZ'] selections.add('eemSSonZ', (ee_eemZmask)&(trig_eem)) selections.add('eemSSoffZ', (ee_eemOffZmask)&(trig_eem)) selections.add('mmeSSonZ', (mm_mmeZmask)&(trig_mme)) selections.add('mmeSSoffZ', (mm_mmeOffZmask)&(trig_mme)) channels3L += ['eeeSSonZ', 'eeeSSoffZ', 'mmmSSonZ', 'mmmSSoffZ'] selections.add('eeeSSonZ', (eeeOnZmask)&(trig_eee)) selections.add('eeeSSoffZ', (eeeOffZmask)&(trig_eee)) selections.add('mmmSSonZ', (mmmOnZmask)&(trig_mmm)) selections.add('mmmSSoffZ', (mmmOffZmask)&(trig_mmm)) levels = ['base', '2jets', '4jets', '4j1b', '4j2b'] selections.add('base', (nElec+nMuon>=2)) selections.add('2jets',(njets>=2)) selections.add('4jets',(njets>=4)) selections.add('4j1b',(njets>=4)&(nbtags>=1)) selections.add('4j2b',(njets>=4)&(nbtags>=2)) # Variables invMass_eeSSonZ = ( eeSSonZ.e0+ eeSSonZ.e1).mass invMass_eeSSoffZ = (eeSSoffZ.e0+eeSSoffZ.e1).mass invMass_mmSSonZ = ( mmSSonZ.m0+ mmSSonZ.m1).mass invMass_mmSSoffZ = (mmSSoffZ.m0+mmSSoffZ.m1).mass invMass_emSS = (emSS.e+emSS.m).mass varnames = {} varnames['met'] = met.pt varnames['ht'] = ht varnames['njets'] = njets varnames['nbtags'] = nbtags varnames['invmass'] = { 'eeSSonZ' : invMass_eeSSonZ, 'eeSSoffZ' : invMass_eeSSoffZ, 'mmSSonZ' : invMass_mmSSonZ, 'mmSSoffZ' : invMass_mmSSoffZ, 'emSS' : invMass_emSS, 'eemSSonZ' : mZ_eem, 'eemSSoffZ' : mZ_eem, 'mmeSSonZ' : mZ_mme, 'mmeSSoffZ' : mZ_mme, 'eeeSSonZ' : mZ_eee, 'eeeSSoffZ' : mZ_eee, 'mmmSSonZ' : mZ_mmm, 'mmmSSoffZ' : mZ_mmm, } varnames['m3l'] = { 'eemSSonZ' : m3l_eem, 'eemSSoffZ' : m3l_eem, 'mmeSSonZ' : m3l_mme, 'mmeSSoffZ' : m3l_mme, 'eeeSSonZ' : m3l_eee, 'eeeSSoffZ' : m3l_eee, 'mmmSSonZ' : m3l_mmm, 'mmmSSoffZ' : m3l_mmm, } varnames['e0pt' ] = e0.pt varnames['e0eta'] = e0.eta varnames['m0pt' ] = m0.pt varnames['m0eta'] = m0.eta varnames['e1pt' ] = e1 varnames['e1eta'] = e1 varnames['e2pt' ] = e2 varnames['e2eta'] = e2 varnames['m1pt' ] = m1 varnames['m1eta'] = m1 varnames['m2pt' ] = m2 varnames['m2eta'] = m2 varnames['j0pt' ] = j0.pt varnames['j0eta'] = j0.eta varnames['j1pt'] = j1 varnames['j1eta'] = j1 varnames['j2pt'] = j2 varnames['j2eta'] = j2 varnames['j3pt'] = j3 varnames['j3eta'] = j3 varnames['counts'] = np.ones_like(events.GenMET.pt) # fill Histos hout = self.accumulator.identity() normweights = weights['all'].weight().flatten() # Why does it not complain about .flatten() here? hout['SumOfEFTweights'].fill(sample=dataset, SumOfEFTweights=varnames['counts'], weight=normweights, eft_coeff=eft_coeffs, eft_err_coeff=eft_w2_coeffs) for var, v in varnames.items(): for ch in channels2LSS+channels3L: for lev in levels: weight = weights[ ch[:3] if (ch.startswith('eee') or ch.startswith('mmm') or ch.startswith('eem') or ch.startswith('mme')) else ch[:2]].weight() cuts = [ch] + [lev] cut = selections.all(*cuts) weights_flat = weight[cut].flatten() # Why does it not complain about .flatten() here? weights_ones = np.ones_like(weights_flat, dtype=np.int) eft_coeffs_cut = eft_coeffs[cut] if eft_coeffs is not None else None eft_w2_coeffs_cut = eft_w2_coeffs[cut] if eft_w2_coeffs is not None else None if var == 'invmass': if ch in ['eeeSSoffZ', 'mmmSSoffZ']: continue elif ch in ['eeeSSonZ' , 'mmmSSonZ' ]: continue #values = v[ch] else : values = ak.flatten(v[ch][cut]) hout['invmass'].fill(eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, sample=dataset, channel=ch, cut=lev, invmass=values, weight=weights_flat) elif var == 'm3l': if ch in ['eeSSonZ','eeSSoffZ', 'mmSSonZ', 'mmSSoffZ','emSS', 'eeeSSoffZ', 'mmmSSoffZ', 'eeeSSonZ' , 'mmmSSonZ']: continue values = ak.flatten(v[ch][cut]) hout['m3l'].fill(sample=dataset, channel=ch, cut=lev, m3l=values, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) else: values = v[cut] if var == 'ht' : hout[var].fill(ht=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'met' : hout[var].fill(met=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'njets' : hout[var].fill(njets=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'nbtags': hout[var].fill(nbtags=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'counts': hout[var].fill(counts=values, sample=dataset, channel=ch, cut=lev, weight=weights_ones) elif var == 'j0eta' : if lev == 'base': continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(j0eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'e0pt' : if ch in ['mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ']: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(e0pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'm0pt' : if ch in ['eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ']: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(m0pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'e0eta' : if ch in ['mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ']: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(e0eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'm0eta': if ch in ['eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ']: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(m0eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'j0pt' : if lev == 'base': continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(j0pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'j1pt': if lev == "base": continue values = values.pt[:,1] #values = ak.flatten(values) hout[var].fill(j1pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var =='j1eta': if lev == 'base': continue values = values.eta[:,1] hout[var].fill(j1eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'j2pt': if lev in ['base', "2jets"]: continue values = values.pt[:,2] hout[var].fill(j2pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'j2eta': if lev in ['base', "2jets"]: continue values = values.eta[:,2] hout[var].fill(j2eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'j3pt': if lev in ['base', "2jets"]: continue values = values.pt[:,3] hout[var].fill(j3pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'j3eta': if lev in ['base', "2jets"]: continue values = values.eta[:,3] hout[var].fill(j3eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'e1pt': if ch in ['mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ', 'mmeSSonZ', 'mmeSSoffZ', 'emSS']: continue values = values.pt[:,1] hout[var].fill(e1pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'e1eta': if ch in ['mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ', 'mmeSSonZ', 'mmeSSoffZ', 'emSS']: continue values = values.eta[:,1] hout[var].fill(e1eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'e2pt': if ch in ['eeeSSonZ', 'eeeSSoffZ']: values = values.pt[:,2] hout[var].fill(e2pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'e2eta': if ch in ['eeeSSonZ', 'eeeSSoffZ']: values = values.eta[:,2] hout[var].fill(e2eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'm1pt': if ch in ['eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ', 'eemSSonZ', 'eemSSoffZ', 'emSS']: continue values = values.pt[:,1] hout[var].fill(m1pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'm1eta': if ch in ['eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ', 'eemSSonZ', 'eemSSoffZ', 'emSS']: continue values = values.eta[:,1] hout[var].fill(m1eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'm2pt': if ch in ['mmmSSonZ', 'mmmSSoffZ']: values = values.pt[:,2] hout[var].fill(m2pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'm2eta': if ch in ['mmmSSonZ', 'mmmSSoffZ']: values = values.eta[:,2] hout[var].fill(m2eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) return hout
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 test_jagged_axis0(): assert ak.min(ak.Array([[1.1, 5.5], [4.4], [2.2, 3.3, 0.0, -10]]), axis=0).tolist() == [1.1, 3.3, 0, -10] assert ak.argmin(ak.Array([[1.1, 5.5], [4.4], [2.2, 3.3, 0.0, -10]]), axis=0).tolist() == [0, 2, 2, 2]
def test_jagged_axis1(): # first is [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [4, 3, 2]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [5, 4, 3]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[], [], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [6, 5, 4]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[1.1], [1.1, 2.2], [], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [5, 4, 3]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [5, 4, 2]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [5, 3, 2]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0], []], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [4, 3, 2]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[1.1, 999, 999], [1.1, 2.2, 999], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [4, 3, 2]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[1.1, 999, 999, 999], [1.1, 2.2, 999], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3, 999]] assert ak.argmin(array, axis=1).tolist() == [[4, 3, 2], [4, 3, 2, 0]] # first is [[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]] array = ak.Array([ [[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [4, 3, 2]] array = ak.Array([ [[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [5, 4, 3]] array = ak.Array([ [[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[], [], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [6, 5, 4]] array = ak.Array([ [[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[1.1], [1.1, 2.2], [], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [5, 4, 3]] array = ak.Array([ [[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [5, 4, 2]] array = ak.Array([ [[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [5, 3, 2]] array = ak.Array([ [[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0], []], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [4, 3, 2]] array = ak.Array([ [[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[1.1, 999, 999], [1.1, 2.2, 999], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [4, 3, 2]] array = ak.Array([ [[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], [[1.1, 999, 999, 999], [1.1, 2.2, 999], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3, 999]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 3], [4, 3, 2, 0]] # first is [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]], [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [4, 3, 2]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]], [[], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [5, 4, 3]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]], [[], [], [1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [6, 5, 4]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]], [[1.1], [1.1, 2.2], [], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [5, 4, 3]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]], [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [5, 4, 2]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]], [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [5, 3, 2]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]], [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [999, 2.0], [1.0], []], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [4, 3, 2]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]], [[1.1, 999, 999], [1.1, 2.2, 999], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [4, 3, 2]] array = ak.Array([ [[1.1], [1.1, 2.2], [1.1, 2.2, 3.3], [], [999, 2.0], [1.0]], [[1.1, 999, 999, 999], [1.1, 2.2, 999], [1.1, 2.2, 3.3], [999, 2.0], [1.0]], ]) assert ak.min(array, axis=1).tolist() == [[1, 2, 3.3], [1, 2, 3.3, 999]] assert ak.argmin(array, axis=1).tolist() == [[5, 4, 2], [4, 3, 2, 0]]
def process(self, events): output = self.accumulator.identity() # use a very loose preselection to filter the events presel = ak.num(events.Jet) > 2 ev = events[presel] dataset = ev.metadata['dataset'] # load the config - probably not needed anymore cfg = loadConfig() output['totalEvents']['all'] += len(events) output['skimmedEvents']['all'] += len(ev) ## Generated leptons '''gen_lep = ev.GenL leading_gen_lep = gen_lep[ak.singletons(ak.argmax(gen_lep.pt, axis=1))] trailing_gen_lep = gen_lep[ak.singletons(ak.argmin(gen_lep.pt, axis=1))]''' ## Muons muon = Collections(ev, "Muon", "tightTTH").get() vetomuon = Collections(ev, "Muon", "vetoTTH").get() leading_muon_idx = ak.singletons(ak.argmax(muon.pt, axis=1)) leading_muon = muon[leading_muon_idx] ## Electrons electron = Collections(ev, "Electron", "tightTTH").get() vetoelectron = Collections(ev, "Electron", "vetoTTH").get() leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1)) leading_electron = electron[leading_electron_idx] ## Merge electrons and muons - this should work better now in ak1 dilepton = cross(muon, electron) dimuon = choose(muon, 2) OS_dimuon = dimuon[(dimuon['0'].charge * dimuon['1'].charge < 0)] dielectron = choose(electron, 2) OS_dielectron = dielectron[( dielectron['0'].charge * dielectron['1'].charge < 0)] OS_dimuon_bestZmumu = OS_dimuon[ak.singletons( ak.argmin(abs(OS_dimuon.mass - 91.2), axis=1))] OS_dielectron_bestZee = OS_dielectron[ak.singletons( ak.argmin(abs(OS_dielectron.mass - 91.2), axis=1))] OS_dilepton_mass = ak.fill_none( ak.pad_none(ak.concatenate( [OS_dimuon_bestZmumu.mass, OS_dielectron_bestZee.mass], axis=1), 1, clip=True), -1) lepton = ak.concatenate([muon, electron], axis=1) leading_lepton_idx = ak.singletons(ak.argmax(lepton.pt, axis=1)) leading_lepton = lepton[leading_lepton_idx] trailing_lepton_idx = ak.singletons(ak.argmin(lepton.pt, axis=1)) trailing_lepton = lepton[trailing_lepton_idx] ## Jets jet = getJets(ev, minPt=25, maxEta=4.7, pt_var='pt_nom') jet = jet[ak.argsort( jet.pt_nom, ascending=False )] # need to sort wrt smeared and recorrected jet pt jet = jet[~match(jet, muon, deltaRCut=0.4)] # remove jets that overlap with muons jet = jet[~match( jet, electron, deltaRCut=0.4)] # remove jets that overlap with electrons central = jet[(abs(jet.eta) < 2.4)] btag = getBTagsDeepFlavB( jet, year=self.year) # should study working point for DeepJet light = getBTagsDeepFlavB(jet, year=self.year, invert=True) fwd = getFwdJet(light) fwd_noPU = getFwdJet(light, puId=False) ## forward jets j_fwd = fwd[ak.singletons(ak.argmax( fwd.p, axis=1))] # highest momentum spectator jf = cross(j_fwd, jet) mjf = (jf['0'] + jf['1']).mass # j_fwd2 = jf[ak.singletons(ak.argmax(mjf, axis=1))]['1'] # this is the jet that forms the largest invariant mass with j_fwd # delta_eta = abs(j_fwd2.eta - j_fwd.eta) ## MET -> can switch to puppi MET met_pt = ev.MET.pt met_phi = ev.MET.phi ## other variables ht = ak.sum(jet.pt, axis=1) st = met_pt + ht + ak.sum(muon.pt, axis=1) + ak.sum(electron.pt, axis=1) # define the weight weight = Weights(len(ev)) if not re.search(re.compile('MuonEG|DoubleMuon|DoubleEG|EGamma'), dataset): # lumi weight weight.add("weight", ev.weight * cfg['lumi'][self.year]) # PU weight - not in the babies... weight.add("PU", ev.puWeight, weightUp=ev.puWeightUp, weightDown=ev.puWeightDown, shift=False) # b-tag SFs weight.add("btag", self.btagSF.Method1a(btag, light)) # lepton SFs # weight.add("lepton", self.leptonSF.get(electron, muon)) cutflow = Cutflow(output, ev, weight=weight) sel = Selection( dataset=dataset, events=ev, year=self.year, ele=electron, ele_veto=vetoelectron, mu=muon, mu_veto=vetomuon, jet_all=jet, jet_central=central, jet_btag=btag, jet_fwd=fwd, met=ev.MET, ) BL = sel.trilep_baseline(cutflow=cutflow) # first, make a few super inclusive plots output['ST'].fill(dataset=dataset, ht=st[BL], weight=weight.weight()[BL]) output['PV_npvs'].fill(dataset=dataset, multiplicity=ev.PV[BL].npvs, weight=weight.weight()[BL]) output['PV_npvsGood'].fill(dataset=dataset, multiplicity=ev.PV[BL].npvsGood, weight=weight.weight()[BL]) output['N_jet'].fill(dataset=dataset, multiplicity=ak.num(jet)[BL], weight=weight.weight()[BL]) output['N_b'].fill(dataset=dataset, multiplicity=ak.num(btag)[BL], weight=weight.weight()[BL]) output['N_central'].fill(dataset=dataset, multiplicity=ak.num(central)[BL], weight=weight.weight()[BL]) output['N_ele'].fill(dataset=dataset, multiplicity=ak.num(vetoelectron)[BL], weight=weight.weight()[BL]) output['N_mu'].fill(dataset=dataset, multiplicity=ak.num(vetomuon)[BL], weight=weight.weight()[BL]) output['N_fwd'].fill(dataset=dataset, multiplicity=ak.num(fwd)[BL], weight=weight.weight()[BL]) '''output['nLepFromTop'].fill(dataset=dataset, multiplicity=ev[BL].nLepFromTop, weight=weight.weight()[BL]) output['nLepFromTau'].fill(dataset=dataset, multiplicity=ev.nLepFromTau[BL], weight=weight.weight()[BL]) output['nLepFromZ'].fill(dataset=dataset, multiplicity=ev.nLepFromZ[BL], weight=weight.weight()[BL]) output['nLepFromW'].fill(dataset=dataset, multiplicity=ev.nLepFromW[BL], weight=weight.weight()[BL]) output['nGenTau'].fill(dataset=dataset, multiplicity=ev.nGenTau[BL], weight=weight.weight()[BL]) output['nGenL'].fill(dataset=dataset, multiplicity=ak.num(ev.GenL[BL], axis=1), weight=weight.weight()[BL])''' # make a plot of the dilepton mass, but without applying the cut on the dilepton mass itself (N-1 plot) output['dilep_mass'].fill( dataset=dataset, mass=ak.flatten( OS_dilepton_mass[sel.trilep_baseline(omit=['offZ'])]), weight=weight.weight()[sel.trilep_baseline(omit=['offZ'])]) output['MET'].fill(dataset=dataset, pt=ev.MET[BL].pt, phi=ev.MET[BL].phi, weight=weight.weight()[BL]) '''output['lead_gen_lep'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(leading_gen_lep[BL].pt)), eta = ak.to_numpy(ak.flatten(leading_gen_lep[BL].eta)), phi = ak.to_numpy(ak.flatten(leading_gen_lep[BL].phi)), weight = weight.weight()[BL] ) output['trail_gen_lep'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(trailing_gen_lep[BL].pt)), eta = ak.to_numpy(ak.flatten(trailing_gen_lep[BL].eta)), phi = ak.to_numpy(ak.flatten(trailing_gen_lep[BL].phi)), weight = weight.weight()[BL] )''' output['lead_lep'].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(leading_lepton[BL].pt)), eta=ak.to_numpy(ak.flatten(leading_lepton[BL].eta)), phi=ak.to_numpy(ak.flatten(leading_lepton[BL].phi)), weight=weight.weight()[BL]) output['trail_lep'].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(trailing_lepton[BL].pt)), eta=ak.to_numpy(ak.flatten(trailing_lepton[BL].eta)), phi=ak.to_numpy(ak.flatten(trailing_lepton[BL].phi)), weight=weight.weight()[BL]) output['j1'].fill(dataset=dataset, pt=ak.flatten(jet.pt_nom[:, 0:1][BL]), eta=ak.flatten(jet.eta[:, 0:1][BL]), phi=ak.flatten(jet.phi[:, 0:1][BL]), weight=weight.weight()[BL]) output['j2'].fill(dataset=dataset, pt=ak.flatten(jet[:, 1:2][BL].pt_nom), eta=ak.flatten(jet[:, 1:2][BL].eta), phi=ak.flatten(jet[:, 1:2][BL].phi), weight=weight.weight()[BL]) #output['j3'].fill( # dataset = dataset, # pt = ak.flatten(jet[:, 2:3][BL].pt_nom), # eta = ak.flatten(jet[:, 2:3][BL].eta), # phi = ak.flatten(jet[:, 2:3][BL].phi), # weight = weight.weight()[BL] #) output['fwd_jet'].fill(dataset=dataset, pt=ak.flatten(j_fwd[BL].pt), eta=ak.flatten(j_fwd[BL].eta), phi=ak.flatten(j_fwd[BL].phi), weight=weight.weight()[BL]) output['high_p_fwd_p'].fill(dataset=dataset, p=ak.flatten(j_fwd[BL].p), weight=weight.weight()[BL]) vetolepton = ak.concatenate([vetomuon, vetoelectron], axis=1) trilep = choose3(vetolepton, 3) trilep_m = trilep.mass output['m3l'].fill(dataset=dataset, mass=ak.flatten(trilep_m[BL]), weight=weight.weight()[BL]) return output
def best_match(gen_hyp=None, jets=None, leptons=None, met=None): if gen_hyp is None: raise ValueError("Gen Objects gen_hyp needed for matching") if jets is None: raise ValueError("Reco jets needed for matching") if leptons is None: raise ValueError("Reco leptons needed for matching") if met is None: raise ValueError("Reco met needed for matching") if not ak.all(ak.num(gen_hyp) == 1): raise ValueError("Not all events for matching are semileptonic") jets_ak = ak.with_name(jets[["pt", "eta", "phi", "mass"]],"PtEtaPhiMLorentzVector") leps_ak = ak.with_name(leptons[["pt", "eta", "phi", "mass"]],"PtEtaPhiMLorentzVector") # init dict of objects matched_objects = {} # match jet closest to gen objects for genobj in ['BHad', 'BLep', 'WJa', 'WJb']: genobj_ak = ak.with_name(gen_hyp[genobj][["pt", "eta", "phi", "mass"]],"PtEtaPhiMLorentzVector") jets_akc, genobj_akc = ak.unzip(ak.cartesian([jets_ak, genobj_ak], nested=False)) deltaRs = jets_akc.delta_r(genobj_akc) # find deltaRs between jets and gen object indexOfMin = ak.unflatten(ak.argmin(deltaRs, axis=1), ak.num(genobj_ak)) passing_inds = deltaRs[indexOfMin] < 0.4 matched_jets_inds = indexOfMin[passing_inds] matched_jets = jets[matched_jets_inds] ## add matched perm objects matched_objects[genobj] = ak.Array({ 'pt' : matched_jets.pt, 'eta' : matched_jets.eta, 'phi' : matched_jets.phi, 'mass' : matched_jets.mass, 'jetIdx' : matched_jets_inds, # index of jet that the gen object is matched to in the event }, with_name="PtEtaPhiMLorentzVector") # match lepton closest to gen lepton genlep_ak = ak.with_name(gen_hyp['Lepton'][["pt", "eta", "phi", "mass"]],"PtEtaPhiMLorentzVector") lep_akc, genlep_akc = ak.unzip(ak.cartesian([leps_ak, genlep_ak], nested=False)) lepDRs = lep_akc.delta_r(genlep_akc) lepIdxOfMin = ak.unflatten(ak.argmin(lepDRs, axis=1), ak.num(genlep_ak)) passing_inds = lepDRs[lepIdxOfMin] < 0.4 matched_leps_inds = lepIdxOfMin[passing_inds] matched_leps = leptons[matched_leps_inds] ## add matched perm objects matched_objects['Lepton'] = ak.Array({key: matched_leps[key] for key in matched_leps.fields}, with_name="PtEtaPhiMLorentzVector") # solve for neutrino nu_array = np.zeros((len(ak.num(jets)), 4), dtype='float64') # convert all inputs into 2d numpy arrays of dtype=float64 (won't work if they're not float64) blep_inputs = np.stack((ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['BLep'].px, 1), -999))).astype('float64'), ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['BLep'].py, 1), -999))).astype('float64'),\ ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['BLep'].pz, 1), -999))).astype('float64'), ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['BLep'].energy, 1), -999))).astype('float64')), axis=-1) lep_inputs = np.stack((ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['Lepton'].px, 1), -999))).astype('float64'), ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['Lepton'].py, 1), -999))).astype('float64'),\ ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['Lepton'].pz, 1), -999))).astype('float64'), ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['Lepton'].energy, 1), -999))).astype('float64')), axis=-1) met_inputs = np.stack((ak.to_numpy(ak.fill_none(met.px, -999)).astype('float64'), ak.to_numpy(ak.fill_none(met.py, -999)).astype('float64')), axis=-1) nu_array = find_nu(bleps=blep_inputs, leptons=lep_inputs, met=met_inputs, nu_array=nu_array) valid_nu = ~((nu_array[:, 3] > 1e20) | (nu_array[:, 3] == 0)) # events that have a solution and matched blep # convert px, py, pz to pt, eta, phi nu_px, nu_py, nu_pz = nu_array[:, 0][valid_nu], nu_array[:, 1][valid_nu], nu_array[:, 2][valid_nu] nu_mom, nu_pt = np.sqrt(np.square(nu_px)+np.square(nu_py)+np.square(nu_pz)), np.sqrt(np.square(nu_px)+np.square(nu_py)) nu_phi = np.arctan2(nu_py, nu_px) nu_eta = np.arcsinh(nu_pz/nu_pt) matched_objects['Nu'] = ak.Array({ 'pt' : ak.unflatten(nu_pt, valid_nu.astype(int)), 'eta' : ak.unflatten(nu_eta, valid_nu.astype(int)), 'phi' : ak.unflatten(nu_phi, valid_nu.astype(int)), 'mass' : ak.zeros_like(ak.unflatten(nu_array[:, 0][valid_nu], valid_nu.astype(int))), 'chi2' : ak.unflatten(nu_array[:, 3][valid_nu], valid_nu.astype(int)), }, with_name="PtEtaPhiMLorentzVector") matched_perm = make_perm_table(bhad=matched_objects['BHad'], blep=matched_objects['BLep'], wja=matched_objects['WJa'], wjb=matched_objects['WJb'], lepton=matched_objects['Lepton'], met=met, nu=matched_objects['Nu']) return matched_perm
def process(self, events): # Dataset parameters dataset = events.metadata['dataset'] year = self._samples[dataset]['year'] xsec = self._samples[dataset]['xsec'] sow = self._samples[dataset]['nSumOfWeights'] isData = self._samples[dataset]['isData'] datasets = [ 'SingleMuon', 'SingleElectron', 'EGamma', 'MuonEG', 'DoubleMuon', 'DoubleElectron' ] for d in datasets: if d in dataset: dataset = dataset.split('_')[0] # Initialize objects met = events.MET e = events.Electron mu = events.Muon tau = events.Tau j = events.Jet # Muon selection #mu['isGood'] = isMuonMVA(mu.pt, mu.eta, mu.dxy, mu.dz, mu.miniPFRelIso_all, mu.sip3d, mu.mvaTTH, mu.mediumPromptId, mu.tightCharge, minpt=10) mu['isPres'] = isPresMuon(mu.dxy, mu.dz, mu.sip3d, mu.looseId) mu['isTight'] = isTightMuon(mu.pt, mu.eta, mu.dxy, mu.dz, mu.pfRelIso03_all, mu.sip3d, mu.mvaTTH, mu.mediumPromptId, mu.tightCharge, mu.looseId, minpt=10) mu['isGood'] = mu['isPres'] & mu['isTight'] leading_mu = mu[ak.argmax(mu.pt, axis=-1, keepdims=True)] leading_mu = leading_mu[leading_mu.isGood] mu = mu[mu.isGood] mu_pres = mu[mu.isPres] # Electron selection #e['isGood'] = isElecMVA(e.pt, e.eta, e.dxy, e.dz, e.miniPFRelIso_all, e.sip3d, e.mvaTTH, e.mvaFall17V2Iso, e.lostHits, e.convVeto, e.tightCharge, minpt=10) e['isPres'] = isPresElec(e.pt, e.eta, e.dxy, e.dz, e.miniPFRelIso_all, e.sip3d, e.lostHits, minpt=15) e['isTight'] = isTightElec(e.pt, e.eta, e.dxy, e.dz, e.miniPFRelIso_all, e.sip3d, e.mvaTTH, e.mvaFall17V2Iso, e.lostHits, e.convVeto, e.tightCharge, e.sieie, e.hoe, e.eInvMinusPInv, minpt=15) e['isClean'] = isClean(e, mu, drmin=0.05) e['isGood'] = e['isPres'] & e['isTight'] & e['isClean'] leading_e = e[ak.argmax(e.pt, axis=-1, keepdims=True)] leading_e = leading_e[leading_e.isGood] e = e[e.isGood] e_pres = e[e.isPres & e.isClean] # Tau selection tau['isPres'] = isPresTau(tau.pt, tau.eta, tau.dxy, tau.dz, tau.leadTkPtOverTauPt, tau.idAntiMu, tau.idAntiEle, tau.rawIso, tau.idDecayModeNewDMs, minpt=20) tau['isClean'] = isClean(tau, e_pres, drmin=0.4) & isClean( tau, mu_pres, drmin=0.4) tau['isGood'] = tau['isPres'] # & tau['isClean'], for the moment tau = tau[tau.isGood] nElec = ak.num(e) nMuon = ak.num(mu) nTau = ak.num(tau) twoLeps = (nElec + nMuon) == 2 threeLeps = (nElec + nMuon) == 3 twoElec = (nElec == 2) twoMuon = (nMuon == 2) e0 = e[ak.argmax(e.pt, axis=-1, keepdims=True)] m0 = mu[ak.argmax(mu.pt, axis=-1, keepdims=True)] # Jet selection jetptname = 'pt_nom' if hasattr(j, 'pt_nom') else 'pt' j['isGood'] = isTightJet(getattr(j, jetptname), j.eta, j.jetId, j.neHEF, j.neEmEF, j.chHEF, j.chEmEF, j.nConstituents) #j['isgood'] = isGoodJet(j.pt, j.eta, j.jetId) #j['isclean'] = isClean(j, e, mu) j['isClean'] = isClean(j, e, drmin=0.4) & isClean( j, mu, drmin=0.4) # & isClean(j, tau, drmin=0.4) goodJets = j[(j.isClean) & (j.isGood)] njets = ak.num(goodJets) ht = ak.sum(goodJets.pt, axis=-1) j0 = goodJets[ak.argmax(goodJets.pt, axis=-1, keepdims=True)] #nbtags = ak.num(goodJets[goodJets.btagDeepFlavB > 0.2770]) nbtags = ak.num(goodJets[goodJets.btagDeepB > 0.4941]) ################################################################## ### 2 same-sign leptons ################################################################## # emu singe = e[(nElec == 1) & (nMuon == 1) & (e.pt > -1)] singm = mu[(nElec == 1) & (nMuon == 1) & (mu.pt > -1)] em = ak.cartesian({"e": singe, "m": singm}) emSSmask = (em.e.charge * em.m.charge > 0) emSS = em[emSSmask] nemSS = len(ak.flatten(emSS)) # ee and mumu # pt>-1 to preserve jagged dimensions ee = e[(nElec == 2) & (nMuon == 0) & (e.pt > -1)] mm = mu[(nElec == 0) & (nMuon == 2) & (mu.pt > -1)] eepairs = ak.combinations(ee, 2, fields=["e0", "e1"]) eeSSmask = (eepairs.e0.charge * eepairs.e1.charge > 0) eeonZmask = (np.abs((eepairs.e0 + eepairs.e1).mass - 91.2) < 10) eeoffZmask = (eeonZmask == 0) mmpairs = ak.combinations(mm, 2, fields=["m0", "m1"]) mmSSmask = (mmpairs.m0.charge * mmpairs.m1.charge > 0) mmonZmask = (np.abs((mmpairs.m0 + mmpairs.m1).mass - 91.2) < 10) mmoffZmask = (mmonZmask == 0) eeSSonZ = eepairs[eeSSmask & eeonZmask] eeSSoffZ = eepairs[eeSSmask & eeoffZmask] mmSSonZ = mmpairs[mmSSmask & mmonZmask] mmSSoffZ = mmpairs[mmSSmask & mmoffZmask] neeSS = len(ak.flatten(eeSSonZ)) + len(ak.flatten(eeSSoffZ)) nmmSS = len(ak.flatten(mmSSonZ)) + len(ak.flatten(mmSSoffZ)) print('Same-sign events [ee, emu, mumu] = [%i, %i, %i]' % (neeSS, nemSS, nmmSS)) # Cuts eeSSmask = (ak.num(eeSSmask[eeSSmask]) > 0) mmSSmask = (ak.num(mmSSmask[mmSSmask]) > 0) eeonZmask = (ak.num(eeonZmask[eeonZmask]) > 0) eeoffZmask = (ak.num(eeoffZmask[eeoffZmask]) > 0) mmonZmask = (ak.num(mmonZmask[mmonZmask]) > 0) mmoffZmask = (ak.num(mmoffZmask[mmoffZmask]) > 0) emSSmask = (ak.num(emSSmask[emSSmask]) > 0) ################################################################## ### 3 leptons ################################################################## # eem muon_eem = mu[(nElec == 2) & (nMuon == 1) & (mu.pt > -1)] elec_eem = e[(nElec == 2) & (nMuon == 1) & (e.pt > -1)] ee_eem = ak.combinations(elec_eem, 2, fields=["e0", "e1"]) ee_eemZmask = (ee_eem.e0.charge * ee_eem.e1.charge < 1) & (np.abs( (ee_eem.e0 + ee_eem.e1).mass - 91.2) < 10) ee_eemOffZmask = (ee_eem.e0.charge * ee_eem.e1.charge < 1) & (np.abs( (ee_eem.e0 + ee_eem.e1).mass - 91.2) > 10) ee_eemZmask = (ak.num(ee_eemZmask[ee_eemZmask]) > 0) ee_eemOffZmask = (ak.num(ee_eemOffZmask[ee_eemOffZmask]) > 0) eepair_eem = (ee_eem.e0 + ee_eem.e1) trilep_eem = eepair_eem + muon_eem #ak.cartesian({"e0":ee_eem.e0,"e1":ee_eem.e1, "m":muon_eem}) # mme muon_mme = mu[(nElec == 1) & (nMuon == 2) & (mu.pt > -1)] elec_mme = e[(nElec == 1) & (nMuon == 2) & (e.pt > -1)] mm_mme = ak.combinations(muon_mme, 2, fields=["m0", "m1"]) mm_mmeZmask = (mm_mme.m0.charge * mm_mme.m1.charge < 1) & (np.abs( (mm_mme.m0 + mm_mme.m1).mass - 91.2) < 10) mm_mmeOffZmask = (mm_mme.m0.charge * mm_mme.m1.charge < 1) & (np.abs( (mm_mme.m0 + mm_mme.m1).mass - 91.2) > 10) mm_mmeZmask = (ak.num(mm_mmeZmask[mm_mmeZmask]) > 0) mm_mmeOffZmask = (ak.num(mm_mmeOffZmask[mm_mmeOffZmask]) > 0) mmpair_mme = (mm_mme.m0 + mm_mme.m1) trilep_mme = mmpair_mme + elec_mme mZ_mme = mmpair_mme.mass mZ_eem = eepair_eem.mass m3l_eem = trilep_eem.mass m3l_mme = trilep_mme.mass # eee and mmm eee = e[(nElec == 3) & (nMuon == 0) & (e.pt > -1)] mmm = mu[(nElec == 0) & (nMuon == 3) & (mu.pt > -1)] eee_leps = ak.combinations(eee, 3, fields=["e0", "e1", "e2"]) mmm_leps = ak.combinations(mmm, 3, fields=["m0", "m1", "m2"]) ee_pairs = ak.combinations(eee, 2, fields=["e0", "e1"]) mm_pairs = ak.combinations(mmm, 2, fields=["m0", "m1"]) ee_pairs_index = ak.argcombinations(eee, 2, fields=["e0", "e1"]) mm_pairs_index = ak.argcombinations(mmm, 2, fields=["m0", "m1"]) mmSFOS_pairs = mm_pairs[ (np.abs(mm_pairs.m0.pdgId) == np.abs(mm_pairs.m1.pdgId)) & (mm_pairs.m0.charge != mm_pairs.m1.charge)] offZmask_mm = ak.all( np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2) > 10., axis=1, keepdims=True) & (ak.num(mmSFOS_pairs) > 0) onZmask_mm = ak.any( np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2) < 10., axis=1, keepdims=True) eeSFOS_pairs = ee_pairs[ (np.abs(ee_pairs.e0.pdgId) == np.abs(ee_pairs.e1.pdgId)) & (ee_pairs.e0.charge != ee_pairs.e1.charge)] offZmask_ee = ak.all( np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2) > 10, axis=1, keepdims=True) & (ak.num(eeSFOS_pairs) > 0) onZmask_ee = ak.any( np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2) < 10, axis=1, keepdims=True) # Create masks **for event selection** eeeOnZmask = (ak.num(onZmask_ee[onZmask_ee]) > 0) eeeOffZmask = (ak.num(offZmask_ee[offZmask_ee]) > 0) mmmOnZmask = (ak.num(onZmask_mm[onZmask_mm]) > 0) mmmOffZmask = (ak.num(offZmask_mm[offZmask_mm]) > 0) # Now we need to create masks for the leptons in order to select leptons from the Z boson candidate (in onZ categories) ZeeMask = ak.argmin(np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2), axis=1, keepdims=True) ZmmMask = ak.argmin(np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2), axis=1, keepdims=True) Zee = eeSFOS_pairs[ZeeMask] Zmm = mmSFOS_pairs[ZmmMask] eZ0 = Zee.e0[ak.num(eeSFOS_pairs) > 0] eZ1 = Zee.e1[ak.num(eeSFOS_pairs) > 0] eZ = eZ0 + eZ1 mZ0 = Zmm.m0[ak.num(mmSFOS_pairs) > 0] mZ1 = Zmm.m1[ak.num(mmSFOS_pairs) > 0] mZ = mZ0 + mZ1 mZ_eee = eZ.mass mZ_mmm = mZ.mass # And for the W boson ZmmIndices = mm_pairs_index[ZmmMask] ZeeIndices = ee_pairs_index[ZeeMask] eW = eee[~ZeeIndices.e0 | ~ZeeIndices.e1] mW = mmm[~ZmmIndices.m0 | ~ZmmIndices.m1] triElec = eee_leps.e0 + eee_leps.e1 + eee_leps.e2 triMuon = mmm_leps.m0 + mmm_leps.m1 + mmm_leps.m2 m3l_eee = triElec.mass m3l_mmm = triMuon.mass # Triggers trig_eeSS = passTrigger(events, 'ee', isData, dataset) trig_mmSS = passTrigger(events, 'mm', isData, dataset) trig_emSS = passTrigger(events, 'em', isData, dataset) trig_eee = passTrigger(events, 'eee', isData, dataset) trig_mmm = passTrigger(events, 'mmm', isData, dataset) trig_eem = passTrigger(events, 'eem', isData, dataset) trig_mme = passTrigger(events, 'mme', isData, dataset) # MET filters # Weights genw = np.ones_like( events['MET_pt']) if isData else events['genWeight'] weights = coffea.analysis_tools.Weights(len(events)) weights.add('norm', genw if isData else (xsec / sow) * genw) eftweights = events['EFTfitCoefficients'] if hasattr( events, "EFTfitCoefficients") else [] # Selections and cuts selections = PackedSelection() channels2LSS = ['eeSSonZ', 'eeSSoffZ', 'mmSSonZ', 'mmSSoffZ', 'emSS'] selections.add('eeSSonZ', (eeonZmask) & (eeSSmask) & (trig_eeSS)) selections.add('eeSSoffZ', (eeoffZmask) & (eeSSmask) & (trig_eeSS)) selections.add('mmSSonZ', (mmonZmask) & (mmSSmask) & (trig_mmSS)) selections.add('mmSSoffZ', (mmoffZmask) & (mmSSmask) & (trig_mmSS)) selections.add('emSS', (emSSmask) & (trig_emSS)) channels3L = ['eemSSonZ', 'eemSSoffZ', 'mmeSSonZ', 'mmeSSoffZ'] selections.add('eemSSonZ', (ee_eemZmask) & (trig_eem)) selections.add('eemSSoffZ', (ee_eemOffZmask) & (trig_eem)) selections.add('mmeSSonZ', (mm_mmeZmask) & (trig_mme)) selections.add('mmeSSoffZ', (mm_mmeOffZmask) & (trig_mme)) channels3L += ['eeeSSonZ', 'eeeSSoffZ', 'mmmSSonZ', 'mmmSSoffZ'] selections.add('eeeSSonZ', (eeeOnZmask) & (trig_eee)) selections.add('eeeSSoffZ', (eeeOffZmask) & (trig_eee)) selections.add('mmmSSonZ', (mmmOnZmask) & (trig_mmm)) selections.add('mmmSSoffZ', (mmmOffZmask) & (trig_mmm)) levels = ['base', '2jets', '4jets', '4j1b', '4j2b'] selections.add('base', (nElec + nMuon >= 2)) selections.add('2jets', (njets >= 2)) selections.add('4jets', (njets >= 4)) selections.add('4j1b', (njets >= 4) & (nbtags >= 1)) selections.add('4j2b', (njets >= 4) & (nbtags >= 2)) # Variables invMass_eeSSonZ = (eeSSonZ.e0 + eeSSonZ.e1).mass invMass_eeSSoffZ = (eeSSoffZ.e0 + eeSSoffZ.e1).mass invMass_mmSSonZ = (mmSSonZ.m0 + mmSSonZ.m1).mass invMass_mmSSoffZ = (mmSSoffZ.m0 + mmSSoffZ.m1).mass invMass_emSS = (emSS.e + emSS.m).mass varnames = {} varnames['met'] = met.pt varnames['ht'] = ht varnames['njets'] = njets varnames['nbtags'] = nbtags varnames['invmass'] = { 'eeSSonZ': invMass_eeSSonZ, 'eeSSoffZ': invMass_eeSSoffZ, 'mmSSonZ': invMass_mmSSonZ, 'mmSSoffZ': invMass_mmSSoffZ, 'emSS': invMass_emSS, 'eemSSonZ': mZ_eem, 'eemSSoffZ': mZ_eem, 'mmeSSonZ': mZ_mme, 'mmeSSoffZ': mZ_mme, 'eeeSSonZ': mZ_eee, 'eeeSSoffZ': mZ_eee, 'mmmSSonZ': mZ_mmm, 'mmmSSoffZ': mZ_mmm, } varnames['m3l'] = { 'eemSSonZ': m3l_eem, 'eemSSoffZ': m3l_eem, 'mmeSSonZ': m3l_mme, 'mmeSSoffZ': m3l_mme, 'eeeSSonZ': m3l_eee, 'eeeSSoffZ': m3l_eee, 'mmmSSonZ': m3l_mmm, 'mmmSSoffZ': m3l_mmm, } varnames['e0pt'] = e0.pt varnames['e0eta'] = e0.eta varnames['m0pt'] = m0.pt varnames['m0eta'] = m0.eta varnames['j0pt'] = j0.pt varnames['j0eta'] = j0.eta varnames['counts'] = np.ones_like(events.MET.pt) # fill Histos hout = self.accumulator.identity() allweights = weights.weight().flatten( ) # Why does it not complain about .flatten() here? hout['SumOfEFTweights'].fill(eftweights, sample=dataset, SumOfEFTweights=varnames['counts'], weight=allweights) for var, v in varnames.items(): for ch in channels2LSS + channels3L: for lev in levels: weight = weights.weight() cuts = [ch] + [lev] cut = selections.all(*cuts) weights_flat = weight[cut].flatten( ) # Why does it not complain about .flatten() here? weights_ones = np.ones_like(weights_flat, dtype=np.int) eftweightsvalues = eftweights[cut] if len( eftweights) > 0 else [] if var == 'invmass': if ch in ['eeeSSoffZ', 'mmmSSoffZ']: continue elif ch in ['eeeSSonZ', 'mmmSSonZ']: continue #values = v[ch] else: values = ak.flatten(v[ch][cut]) hout['invmass'].fill(sample=dataset, channel=ch, cut=lev, invmass=values, weight=weights_flat) elif var == 'm3l': if ch in [ 'eeSSonZ', 'eeSSoffZ', 'mmSSonZ', 'mmSSoffZ', 'emSS', 'eeeSSoffZ', 'mmmSSoffZ', 'eeeSSonZ', 'mmmSSonZ' ]: continue values = ak.flatten(v[ch][cut]) hout['m3l'].fill(eftweightsvalues, sample=dataset, channel=ch, cut=lev, m3l=values, weight=weights_flat) else: values = v[cut] if var == 'ht': hout[var].fill(eftweightsvalues, ht=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) elif var == 'met': hout[var].fill(eftweightsvalues, met=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) elif var == 'njets': hout[var].fill(eftweightsvalues, njets=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) elif var == 'nbtags': hout[var].fill(eftweightsvalues, nbtags=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) elif var == 'counts': hout[var].fill(counts=values, sample=dataset, channel=ch, cut=lev, weight=weights_ones) elif var == 'j0eta': if lev == 'base': continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(eftweightsvalues, j0eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) elif var == 'e0pt': if ch in [ 'mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eftweightsvalues, e0pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat ) # Crashing here, not sure why. Related to values? elif var == 'm0pt': if ch in [ 'eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(eftweightsvalues, m0pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) elif var == 'e0eta': if ch in [ 'mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(eftweightsvalues, e0eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) elif var == 'm0eta': if ch in [ 'eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(eftweightsvalues, m0eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) elif var == 'j0pt': if lev == 'base': continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(eftweightsvalues, j0pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) return hout
def process(self, events): output = self.accumulator.identity() output['total']['all'] += len(events) # use a very loose preselection to filter the events presel = ak.num(events.Jet) > 2 ev = events[presel] dataset = ev.metadata['dataset'] # load the config - probably not needed anymore cfg = loadConfig() ## Muons muon = Collections(ev, "Muon", "vetoTTH").get() tightmuon = Collections(ev, "Muon", "tightTTH").get() dimuon = choose(muon, 2) SSmuon = ak.any((dimuon['0'].charge * dimuon['1'].charge) > 0, axis=1) leading_muon_idx = ak.singletons(ak.argmax(muon.pt, axis=1)) leading_muon = muon[leading_muon_idx] ## Electrons electron = Collections(ev, "Electron", "vetoTTH").get() tightelectron = Collections(ev, "Electron", "tightTTH").get() dielectron = choose(electron, 2) SSelectron = ak.any( (dielectron['0'].charge * dielectron['1'].charge) > 0, axis=1) leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1)) leading_electron = electron[leading_electron_idx] ## Merge electrons and muons - this should work better now in ak1 dilepton = cross(muon, electron) SSlepton = ak.any((dilepton['0'].charge * dilepton['1'].charge) > 0, axis=1) lepton = ak.concatenate([muon, electron], axis=1) leading_lepton_idx = ak.singletons(ak.argmax(lepton.pt, axis=1)) leading_lepton = lepton[leading_lepton_idx] trailing_lepton_idx = ak.singletons(ak.argmin(lepton.pt, axis=1)) trailing_lepton = lepton[trailing_lepton_idx] dilepton_mass = (leading_lepton + trailing_lepton).mass dilepton_pt = (leading_lepton + trailing_lepton).pt dilepton_dR = delta_r(leading_lepton, trailing_lepton) mt_lep_met = mt(lepton.pt, lepton.phi, ev.MET.pt, ev.MET.phi) min_mt_lep_met = ak.min(mt_lep_met, axis=1) ## Jets jet = getJets(ev, minPt=25, maxEta=4.7, pt_var='pt_nom') jet = jet[ak.argsort( jet.pt_nom, ascending=False )] # need to sort wrt smeared and recorrected jet pt jet = jet[~match(jet, muon, deltaRCut=0.4)] # remove jets that overlap with muons jet = jet[~match( jet, electron, deltaRCut=0.4)] # remove jets that overlap with electrons central = jet[(abs(jet.eta) < 2.4)] btag = getBTagsDeepFlavB( jet, year=self.year) # should study working point for DeepJet light = getBTagsDeepFlavB(jet, year=self.year, invert=True) fwd = getFwdJet(light) fwd_noPU = getFwdJet(light, puId=False) tau = getTaus(ev) track = getIsoTracks(ev) ## forward jets j_fwd = fwd[ak.singletons(ak.argmax( fwd.p, axis=1))] # highest momentum spectator high_score_btag = central[ak.argsort(central.btagDeepFlavB)][:, :2] bl = cross(lepton, high_score_btag) bl_dR = delta_r(bl['0'], bl['1']) min_bl_dR = ak.min(bl_dR, axis=1) jf = cross(j_fwd, jet) mjf = (jf['0'] + jf['1']).mass j_fwd2 = jf[ak.singletons( ak.argmax(mjf, axis=1) )]['1'] # this is the jet that forms the largest invariant mass with j_fwd delta_eta = ak.fill_none( ak.pad_none(abs(j_fwd2.eta - j_fwd.eta), 1, clip=True), 0) ## MET -> can switch to puppi MET met_pt = ev.MET.pt met_phi = ev.MET.phi ## other variables ht = ak.sum(jet.pt, axis=1) st = met_pt + ht + ak.sum(muon.pt, axis=1) + ak.sum(electron.pt, axis=1) ## event selectors filters = getFilters(ev, year=self.year, dataset=dataset) dilep = ((ak.num(tightelectron) + ak.num(tightmuon)) == 2) lep0pt = ((ak.num(electron[(electron.pt > 25)]) + ak.num(muon[(muon.pt > 25)])) > 0) lep1pt = ((ak.num(electron[(electron.pt > 20)]) + ak.num(muon[(muon.pt > 20)])) > 1) lepveto = ((ak.num(electron) + ak.num(muon)) == 2) selection = PackedSelection() selection.add('lepveto', lepveto) selection.add('dilep', dilep) selection.add('filter', (filters)) selection.add('p_T(lep0)>25', lep0pt) selection.add('p_T(lep1)>20', lep1pt) selection.add('SS', (SSlepton | SSelectron | SSmuon)) selection.add('N_jet>3', (ak.num(jet) >= 4)) selection.add('N_central>2', (ak.num(central) >= 3)) selection.add('N_btag>0', (ak.num(btag) >= 1)) selection.add('N_fwd>0', (ak.num(fwd) >= 1)) #ss_reqs = ['lepveto', 'dilep', 'filter', 'p_T(lep0)>25', 'p_T(lep1)>20', 'SS'] ss_reqs = [ 'lepveto', 'dilep', 'filter', 'p_T(lep0)>25', 'p_T(lep1)>20', 'SS' ] #bl_reqs = ss_reqs + ['N_jet>3', 'N_central>2', 'N_btag>0', 'N_fwd>0'] bl_reqs = ss_reqs + ['N_jet>3', 'N_central>2', 'N_btag>0'] ss_reqs_d = {sel: True for sel in ss_reqs} ss_selection = selection.require(**ss_reqs_d) bl_reqs_d = {sel: True for sel in bl_reqs} BL = selection.require(**bl_reqs_d) weight = Weights(len(ev)) if not dataset == 'MuonEG': # lumi weight weight.add("weight", ev.weight) # PU weight - not in the babies... weight.add("PU", ev.puWeight, weightUp=ev.puWeightUp, weightDown=ev.puWeightDown, shift=False) # b-tag SFs weight.add("btag", self.btagSF.Method1a(btag, light)) # lepton SFs weight.add("lepton", self.leptonSF.get(electron, muon)) #cutflow = Cutflow(output, ev, weight=weight) #cutflow_reqs_d = {} #for req in bl_reqs: # cutflow_reqs_d.update({req: True}) # cutflow.addRow( req, selection.require(**cutflow_reqs_d) ) labels = { 'topW_v3': 0, 'TTW': 1, 'TTZ': 2, 'TTH': 3, 'ttbar': 4, 'ttbar1l_MG': 4 } if dataset in labels: label_mult = labels[dataset] else: label_mult = 5 label = np.ones(len(ev[BL])) * label_mult output["n_lep"] += processor.column_accumulator( ak.to_numpy((ak.num(electron) + ak.num(muon))[BL])) output["n_lep_tight"] += processor.column_accumulator( ak.to_numpy((ak.num(tightelectron) + ak.num(tightmuon))[BL])) output["lead_lep_pt"] += processor.column_accumulator( ak.to_numpy(ak.flatten(leading_lepton[BL].pt, axis=1))) output["lead_lep_eta"] += processor.column_accumulator( ak.to_numpy(ak.flatten(leading_lepton[BL].eta, axis=1))) output["lead_lep_phi"] += processor.column_accumulator( ak.to_numpy(ak.flatten(leading_lepton[BL].phi, axis=1))) output["lead_lep_charge"] += processor.column_accumulator( ak.to_numpy(ak.flatten(leading_lepton[BL].charge, axis=1))) output["sublead_lep_pt"] += processor.column_accumulator( ak.to_numpy(ak.flatten(trailing_lepton[BL].pt, axis=1))) output["sublead_lep_eta"] += processor.column_accumulator( ak.to_numpy(ak.flatten(trailing_lepton[BL].eta, axis=1))) output["sublead_lep_phi"] += processor.column_accumulator( ak.to_numpy(ak.flatten(trailing_lepton[BL].phi, axis=1))) output["sublead_lep_charge"] += processor.column_accumulator( ak.to_numpy(ak.flatten(trailing_lepton[BL].charge, axis=1))) output["lead_jet_pt"] += processor.column_accumulator( ak.to_numpy(ak.flatten(jet[:, 0:1][BL].pt, axis=1))) output["lead_jet_eta"] += processor.column_accumulator( ak.to_numpy(ak.flatten(jet[:, 0:1][BL].eta, axis=1))) output["lead_jet_phi"] += processor.column_accumulator( ak.to_numpy(ak.flatten(jet[:, 0:1][BL].phi, axis=1))) output["sublead_jet_pt"] += processor.column_accumulator( ak.to_numpy(ak.flatten(jet[:, 1:2][BL].pt, axis=1))) output["sublead_jet_eta"] += processor.column_accumulator( ak.to_numpy(ak.flatten(jet[:, 1:2][BL].eta, axis=1))) output["sublead_jet_phi"] += processor.column_accumulator( ak.to_numpy(ak.flatten(jet[:, 1:2][BL].phi, axis=1))) output["lead_btag_pt"] += processor.column_accumulator( ak.to_numpy(ak.flatten(high_score_btag[:, 0:1][BL].pt, axis=1))) output["lead_btag_eta"] += processor.column_accumulator( ak.to_numpy(ak.flatten(high_score_btag[:, 0:1][BL].eta, axis=1))) output["lead_btag_phi"] += processor.column_accumulator( ak.to_numpy(ak.flatten(high_score_btag[:, 0:1][BL].phi, axis=1))) output["sublead_btag_pt"] += processor.column_accumulator( ak.to_numpy(ak.flatten(high_score_btag[:, 1:2][BL].pt, axis=1))) output["sublead_btag_eta"] += processor.column_accumulator( ak.to_numpy(ak.flatten(high_score_btag[:, 1:2][BL].eta, axis=1))) output["sublead_btag_phi"] += processor.column_accumulator( ak.to_numpy(ak.flatten(high_score_btag[:, 1:2][BL].phi, axis=1))) output["fwd_jet_p"] += processor.column_accumulator( ak.to_numpy( ak.flatten(ak.fill_none(ak.pad_none(j_fwd[BL].p, 1, clip=True), 0), axis=1))) output["fwd_jet_pt"] += processor.column_accumulator( ak.to_numpy( ak.flatten(ak.fill_none( ak.pad_none(j_fwd[BL].pt, 1, clip=True), 0), axis=1))) output["fwd_jet_eta"] += processor.column_accumulator( ak.to_numpy( ak.flatten(ak.fill_none( ak.pad_none(j_fwd[BL].eta, 1, clip=True), 0), axis=1))) output["fwd_jet_phi"] += processor.column_accumulator( ak.to_numpy( ak.flatten(ak.fill_none( ak.pad_none(j_fwd[BL].phi, 1, clip=True), 0), axis=1))) output["mjj_max"] += processor.column_accumulator( ak.to_numpy(ak.fill_none(ak.max(mjf[BL], axis=1), 0))) output["delta_eta_jj"] += processor.column_accumulator( ak.to_numpy(ak.flatten(delta_eta[BL], axis=1))) output["met"] += processor.column_accumulator(ak.to_numpy(met_pt[BL])) output["ht"] += processor.column_accumulator(ak.to_numpy(ht[BL])) output["st"] += processor.column_accumulator(ak.to_numpy(st[BL])) output["n_jet"] += processor.column_accumulator( ak.to_numpy(ak.num(jet[BL]))) output["n_btag"] += processor.column_accumulator( ak.to_numpy(ak.num(btag[BL]))) output["n_fwd"] += processor.column_accumulator( ak.to_numpy(ak.num(fwd[BL]))) output["n_central"] += processor.column_accumulator( ak.to_numpy(ak.num(central[BL]))) output["n_tau"] += processor.column_accumulator( ak.to_numpy(ak.num(tau[BL]))) output["n_track"] += processor.column_accumulator( ak.to_numpy(ak.num(track[BL]))) output["dilepton_pt"] += processor.column_accumulator( ak.to_numpy(ak.flatten(dilepton_pt[BL], axis=1))) output["dilepton_mass"] += processor.column_accumulator( ak.to_numpy(ak.flatten(dilepton_mass[BL], axis=1))) output["min_bl_dR"] += processor.column_accumulator( ak.to_numpy(min_bl_dR[BL])) output["min_mt_lep_met"] += processor.column_accumulator( ak.to_numpy(min_mt_lep_met[BL])) output["label"] += processor.column_accumulator(label) output["weight"] += processor.column_accumulator(weight.weight()[BL]) output["presel"]["all"] += len(ev[ss_selection]) output["sel"]["all"] += len(ev[BL]) return output
def 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
def process(self, events): output = self.accumulator.identity() # we can use a very loose preselection to filter the events. nothing is done with this presel, though presel = ak.num(events.Jet) > 0 ev = events[presel] dataset = ev.metadata['dataset'] # load the config - probably not needed anymore cfg = loadConfig() #output['totalEvents']['all'] += len(events) #output['skimmedEvents']['all'] += len(ev) if self.year == 2018: triggers = ev.HLT.Ele23_Ele12_CaloIdL_TrackIdL_IsoVL elif self.year == 2017: triggers = ev.HLT.Ele23_Ele12_CaloIdL_TrackIdL_IsoVL elif self.year == 2016: triggers = ev.HLT.Ele23_Ele12_CaloIdL_TrackIdL_IsoVL_DZ if self.year == 2018: lumimask = LumiMask( 'processors/Cert_314472-325175_13TeV_Legacy2018_Collisions18_JSON.txt' ) ## Electrons electron = Collections(ev, "Electron", "tight").get() electron = electron[(electron.pt > 25) & (np.abs(electron.eta) < 2.4)] loose_electron = Collections(ev, "Electron", "veto").get() loose_electron = loose_electron[(loose_electron.pt > 25) & (np.abs(loose_electron.eta) < 2.4)] SSelectron = (ak.sum(electron.charge, axis=1) != 0) & (ak.num(electron) == 2) OSelectron = (ak.sum(electron.charge, axis=1) == 0) & (ak.num(electron) == 2) dielectron = choose(electron, 2) dielectron_mass = (dielectron['0'] + dielectron['1']).mass dielectron_pt = (dielectron['0'] + dielectron['1']).pt leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1)) leading_electron = electron[(leading_electron_idx)] leading_electron = leading_electron[(leading_electron.pt > 30)] trailing_electron_idx = ak.singletons(ak.argmin(electron.pt, axis=1)) trailing_electron = electron[trailing_electron_idx] ##Muons loose_muon = Collections(ev, "Muon", "veto").get() loose_muon = loose_muon[(loose_muon.pt > 20) & (np.abs(loose_muon.eta) < 2.4)] #jets jet = getJets(ev, minPt=40, maxEta=2.4, pt_var='pt', UL=False) jet = jet[ak.argsort( jet.pt, ascending=False )] # need to sort wrt smeared and recorrected jet pt jet = jet[~match(jet, loose_muon, deltaRCut=0.4)] # remove jets that overlap with muons jet = jet[~match( jet, electron, deltaRCut=0.4)] # remove jets that overlap with electrons ## MET -> can switch to puppi MET met_pt = ev.MET.pt met_phi = ev.MET.phi #selections filters = getFilters(ev, year=self.year, dataset=dataset) mask = lumimask(ev.run, ev.luminosityBlock) ss = (SSelectron) os = (OSelectron) mass = (ak.min(np.abs(dielectron_mass - 91.2), axis=1) < 15) lead_electron = (ak.min(leading_electron.pt, axis=1) > 30) jet1 = (ak.num(jet) >= 1) jet2 = (ak.num(jet) >= 2) num_loose = ((ak.num(loose_electron) == 2) & (ak.num(loose_muon) == 0)) selection = PackedSelection() selection.add('filter', (filters)) selection.add('mask', (mask)) selection.add('ss', ss) selection.add('os', os) selection.add('mass', mass) selection.add('leading', lead_electron) selection.add('triggers', triggers) selection.add('one jet', jet1) selection.add('two jets', jet2) selection.add('num_loose', num_loose) bl_reqs = ['filter'] + ['mass'] + ['mask'] + ['triggers'] + [ 'leading' ] + ['num_loose'] #bl_reqs = ['filter'] + ['mass'] + ['triggers'] + ['leading'] + ['num_loose'] bl_reqs_d = {sel: True for sel in bl_reqs} baseline = selection.require(**bl_reqs_d) s_reqs = bl_reqs + ['ss'] s_reqs_d = {sel: True for sel in s_reqs} ss_sel = selection.require(**s_reqs_d) o_reqs = bl_reqs + ['os'] o_reqs_d = {sel: True for sel in o_reqs} os_sel = selection.require(**o_reqs_d) j1s_reqs = s_reqs + ['one jet'] j1s_reqs_d = {sel: True for sel in j1s_reqs} j1ss_sel = selection.require(**j1s_reqs_d) j1o_reqs = o_reqs + ['one jet'] j1o_reqs_d = {sel: True for sel in j1o_reqs} j1os_sel = selection.require(**j1o_reqs_d) j2s_reqs = s_reqs + ['two jets'] j2s_reqs_d = {sel: True for sel in j2s_reqs} j2ss_sel = selection.require(**j2s_reqs_d) j2o_reqs = o_reqs + ['two jets'] j2o_reqs_d = {sel: True for sel in j2o_reqs} j2os_sel = selection.require(**j2o_reqs_d) output["N_jet"].fill( dataset=dataset, multiplicity=ak.num(jet)[os_sel], ) return output
def process(self, events): # Dataset parameters dataset = events.metadata['dataset'] histAxisName = self._samples[dataset]['histAxisName'] year = self._samples[dataset]['year'] xsec = self._samples[dataset]['xsec'] sow = self._samples[dataset]['nSumOfWeights'] isData = self._samples[dataset]['isData'] datasets = [ 'SingleMuon', 'SingleElectron', 'EGamma', 'MuonEG', 'DoubleMuon', 'DoubleElectron' ] for d in datasets: if d in dataset: dataset = dataset.split('_')[0] # Initialize objects met = events.MET e = events.Electron mu = events.Muon tau = events.Tau j = events.Jet # Muon selection mu['isPres'] = isPresMuon(mu.dxy, mu.dz, mu.sip3d, mu.looseId) mu['isTight'] = isTightMuon(mu.pt, mu.eta, mu.dxy, mu.dz, mu.pfRelIso03_all, mu.sip3d, mu.mvaTTH, mu.mediumPromptId, mu.tightCharge, mu.looseId, minpt=10) mu['isGood'] = mu['isPres'] & mu['isTight'] leading_mu = mu[ak.argmax(mu.pt, axis=-1, keepdims=True)] leading_mu = leading_mu[leading_mu.isGood] mu = mu[mu.isGood] mu_pres = mu[mu.isPres] # Electron selection e['isPres'] = isPresElec(e.pt, e.eta, e.dxy, e.dz, e.miniPFRelIso_all, e.sip3d, e.lostHits, minpt=15) e['isTight'] = isTightElec(e.pt, e.eta, e.dxy, e.dz, e.miniPFRelIso_all, e.sip3d, e.mvaTTH, e.mvaFall17V2Iso, e.lostHits, e.convVeto, e.tightCharge, e.sieie, e.hoe, e.eInvMinusPInv, minpt=15) e['isClean'] = isClean(e, mu, drmin=0.05) e['isGood'] = e['isPres'] & e['isTight'] & e['isClean'] leading_e = e[ak.argmax(e.pt, axis=-1, keepdims=True)] leading_e = leading_e[leading_e.isGood] e = e[e.isGood] e_pres = e[e.isPres & e.isClean] # Tau selection tau['isPres'] = isPresTau(tau.pt, tau.eta, tau.dxy, tau.dz, tau.leadTkPtOverTauPt, tau.idAntiMu, tau.idAntiEle, tau.rawIso, tau.idDecayModeNewDMs, minpt=20) tau['isClean'] = isClean(tau, e_pres, drmin=0.4) & isClean( tau, mu_pres, drmin=0.4) tau['isGood'] = tau['isPres'] # & tau['isClean'], for the moment tau = tau[tau.isGood] nElec = ak.num(e) nMuon = ak.num(mu) nTau = ak.num(tau) twoLeps = (nElec + nMuon) == 2 threeLeps = (nElec + nMuon) == 3 twoElec = (nElec == 2) twoMuon = (nMuon == 2) e0 = e[ak.argmax(e.pt, axis=-1, keepdims=True)] m0 = mu[ak.argmax(mu.pt, axis=-1, keepdims=True)] # Attach the lepton SFs to the electron and muons collections AttachElectronSF(e, year=year) AttachMuonSF(mu, year=year) # Create a lepton (muon+electron) collection and calculate a per event lepton SF leps = ak.concatenate([e, mu], axis=-1) events['lepSF_nom'] = ak.prod(leps.sf_nom, axis=-1) events['lepSF_hi'] = ak.prod(leps.sf_hi, axis=-1) events['lepSF_lo'] = ak.prod(leps.sf_lo, axis=-1) # Jet selection jetptname = 'pt_nom' if hasattr(j, 'pt_nom') else 'pt' ### Jet energy corrections if not isData: j["pt_raw"] = (1 - j.rawFactor) * j.pt j["mass_raw"] = (1 - j.rawFactor) * j.mass j["pt_gen"] = ak.values_astype(ak.fill_none(j.matched_gen.pt, 0), np.float32) j["rho"] = ak.broadcast_arrays(events.fixedGridRhoFastjetAll, j.pt)[0] events_cache = events.caches[0] corrected_jets = jet_factory.build(j, lazy_cache=events_cache) #print('jet pt: ',j.pt) #print('cor pt: ',corrected_jets.pt) #print('jes up: ',corrected_jets.JES_jes.up.pt) #print('jes down: ',corrected_jets.JES_jes.down.pt) #print(ak.fields(corrected_jets)) ''' # SYSTEMATICS jets = corrected_jets if(self.jetSyst == 'JERUp'): jets = corrected_jets.JER.up elif(self.jetSyst == 'JERDown'): jets = corrected_jets.JER.down elif(self.jetSyst == 'JESUp'): jets = corrected_jets.JES_jes.up elif(self.jetSyst == 'JESDown'): jets = corrected_jets.JES_jes.down ''' j['isGood'] = isTightJet(getattr(j, jetptname), j.eta, j.jetId, j.neHEF, j.neEmEF, j.chHEF, j.chEmEF, j.nConstituents) #j['isgood'] = isGoodJet(j.pt, j.eta, j.jetId) #j['isclean'] = isClean(j, e, mu) j['isClean'] = isClean(j, e, drmin=0.4) & isClean( j, mu, drmin=0.4) # & isClean(j, tau, drmin=0.4) goodJets = j[(j.isClean) & (j.isGood)] njets = ak.num(goodJets) ht = ak.sum(goodJets.pt, axis=-1) j0 = goodJets[ak.argmax(goodJets.pt, axis=-1, keepdims=True)] #nbtags = ak.num(goodJets[goodJets.btagDeepFlavB > 0.2770]) # Loose DeepJet WP if year == 2017: btagwpl = 0.0532 #WP loose else: btagwpl = 0.0490 #WP loose isBtagJetsLoose = (goodJets.btagDeepB > btagwpl) isNotBtagJetsLoose = np.invert(isBtagJetsLoose) nbtagsl = ak.num(goodJets[isBtagJetsLoose]) # Medium DeepJet WP if year == 2017: btagwpm = 0.3040 #WP medium else: btagwpm = 0.2783 #WP medium isBtagJetsMedium = (goodJets.btagDeepB > btagwpm) isNotBtagJetsMedium = np.invert(isBtagJetsMedium) nbtagsm = ak.num(goodJets[isBtagJetsMedium]) # Btag SF following 1a) in https://twiki.cern.ch/twiki/bin/viewauth/CMS/BTagSFMethods btagSF = np.ones_like(ht) btagSFUp = np.ones_like(ht) btagSFDo = np.ones_like(ht) if not isData: pt = goodJets.pt abseta = np.abs(goodJets.eta) flav = goodJets.hadronFlavour bJetSF = GetBTagSF(abseta, pt, flav) bJetSFUp = GetBTagSF(abseta, pt, flav, sys=1) bJetSFDo = GetBTagSF(abseta, pt, flav, sys=-1) bJetEff = GetBtagEff(abseta, pt, flav, year) bJetEff_data = bJetEff * bJetSF bJetEff_dataUp = bJetEff * bJetSFUp bJetEff_dataDo = bJetEff * bJetSFDo pMC = ak.prod(bJetEff[isBtagJetsMedium], axis=-1) * ak.prod( (1 - bJetEff[isNotBtagJetsMedium]), axis=-1) pData = ak.prod(bJetEff_data[isBtagJetsMedium], axis=-1) * ak.prod( (1 - bJetEff_data[isNotBtagJetsMedium]), axis=-1) pDataUp = ak.prod( bJetEff_dataUp[isBtagJetsMedium], axis=-1) * ak.prod( (1 - bJetEff_dataUp[isNotBtagJetsMedium]), axis=-1) pDataDo = ak.prod( bJetEff_dataDo[isBtagJetsMedium], axis=-1) * ak.prod( (1 - bJetEff_dataDo[isNotBtagJetsMedium]), axis=-1) pMC = ak.where(pMC == 0, 1, pMC) # removeing zeroes from denominator... btagSF = pData / pMC btagSFUp = pDataUp / pMC btagSFDo = pDataUp / pMC ################################################################## ### 2 same-sign leptons ################################################################## # emu singe = e[(nElec == 1) & (nMuon == 1) & (e.pt > -1)] singm = mu[(nElec == 1) & (nMuon == 1) & (mu.pt > -1)] em = ak.cartesian({"e": singe, "m": singm}) emSSmask = (em.e.charge * em.m.charge > 0) emSS = em[emSSmask] nemSS = len(ak.flatten(emSS)) # ee and mumu # pt>-1 to preserve jagged dimensions ee = e[(nElec == 2) & (nMuon == 0) & (e.pt > -1)] mm = mu[(nElec == 0) & (nMuon == 2) & (mu.pt > -1)] sumcharge = ak.sum(e.charge, axis=-1) + ak.sum(mu.charge, axis=-1) eepairs = ak.combinations(ee, 2, fields=["e0", "e1"]) eeSSmask = (eepairs.e0.charge * eepairs.e1.charge > 0) eeonZmask = (np.abs((eepairs.e0 + eepairs.e1).mass - 91.2) < 10) eeoffZmask = (eeonZmask == 0) mmpairs = ak.combinations(mm, 2, fields=["m0", "m1"]) mmSSmask = (mmpairs.m0.charge * mmpairs.m1.charge > 0) mmonZmask = (np.abs((mmpairs.m0 + mmpairs.m1).mass - 91.2) < 10) mmoffZmask = (mmonZmask == 0) eeSSonZ = eepairs[eeSSmask & eeonZmask] eeSSoffZ = eepairs[eeSSmask & eeoffZmask] mmSSonZ = mmpairs[mmSSmask & mmonZmask] mmSSoffZ = mmpairs[mmSSmask & mmoffZmask] neeSS = len(ak.flatten(eeSSonZ)) + len(ak.flatten(eeSSoffZ)) nmmSS = len(ak.flatten(mmSSonZ)) + len(ak.flatten(mmSSoffZ)) print('Same-sign events [ee, emu, mumu] = [%i, %i, %i]' % (neeSS, nemSS, nmmSS)) # Cuts eeSSmask = (ak.num(eeSSmask[eeSSmask]) > 0) mmSSmask = (ak.num(mmSSmask[mmSSmask]) > 0) eeonZmask = (ak.num(eeonZmask[eeonZmask]) > 0) eeoffZmask = (ak.num(eeoffZmask[eeoffZmask]) > 0) mmonZmask = (ak.num(mmonZmask[mmonZmask]) > 0) mmoffZmask = (ak.num(mmoffZmask[mmoffZmask]) > 0) emSSmask = (ak.num(emSSmask[emSSmask]) > 0) ################################################################## ### 3 leptons ################################################################## # eem muon_eem = mu[(nElec == 2) & (nMuon == 1) & (mu.pt > -1)] elec_eem = e[(nElec == 2) & (nMuon == 1) & (e.pt > -1)] ee_eem = ak.combinations(elec_eem, 2, fields=["e0", "e1"]) ee_eemZmask = (ee_eem.e0.charge * ee_eem.e1.charge < 1) & (np.abs( (ee_eem.e0 + ee_eem.e1).mass - 91.2) < 10) ee_eemOffZmask = (ee_eem.e0.charge * ee_eem.e1.charge < 1) & (np.abs( (ee_eem.e0 + ee_eem.e1).mass - 91.2) > 10) ee_eemZmask = (ak.num(ee_eemZmask[ee_eemZmask]) > 0) ee_eemOffZmask = (ak.num(ee_eemOffZmask[ee_eemOffZmask]) > 0) eepair_eem = (ee_eem.e0 + ee_eem.e1) trilep_eem = eepair_eem + muon_eem #ak.cartesian({"e0":ee_eem.e0,"e1":ee_eem.e1, "m":muon_eem}) # mme muon_mme = mu[(nElec == 1) & (nMuon == 2) & (mu.pt > -1)] elec_mme = e[(nElec == 1) & (nMuon == 2) & (e.pt > -1)] mm_mme = ak.combinations(muon_mme, 2, fields=["m0", "m1"]) mm_mmeZmask = (mm_mme.m0.charge * mm_mme.m1.charge < 1) & (np.abs( (mm_mme.m0 + mm_mme.m1).mass - 91.2) < 10) mm_mmeOffZmask = (mm_mme.m0.charge * mm_mme.m1.charge < 1) & (np.abs( (mm_mme.m0 + mm_mme.m1).mass - 91.2) > 10) mm_mmeZmask = (ak.num(mm_mmeZmask[mm_mmeZmask]) > 0) mm_mmeOffZmask = (ak.num(mm_mmeOffZmask[mm_mmeOffZmask]) > 0) mmpair_mme = (mm_mme.m0 + mm_mme.m1) trilep_mme = mmpair_mme + elec_mme mZ_mme = mmpair_mme.mass mZ_eem = eepair_eem.mass m3l_eem = trilep_eem.mass m3l_mme = trilep_mme.mass # eee and mmm eee = e[(nElec == 3) & (nMuon == 0) & (e.pt > -1)] mmm = mu[(nElec == 0) & (nMuon == 3) & (mu.pt > -1)] eee_leps = ak.combinations(eee, 3, fields=["e0", "e1", "e2"]) mmm_leps = ak.combinations(mmm, 3, fields=["m0", "m1", "m2"]) ee_pairs = ak.combinations(eee, 2, fields=["e0", "e1"]) mm_pairs = ak.combinations(mmm, 2, fields=["m0", "m1"]) ee_pairs_index = ak.argcombinations(eee, 2, fields=["e0", "e1"]) mm_pairs_index = ak.argcombinations(mmm, 2, fields=["m0", "m1"]) mmSFOS_pairs = mm_pairs[ (np.abs(mm_pairs.m0.pdgId) == np.abs(mm_pairs.m1.pdgId)) & (mm_pairs.m0.charge != mm_pairs.m1.charge)] offZmask_mm = ak.all( np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2) > 10., axis=1, keepdims=True) & (ak.num(mmSFOS_pairs) > 0) onZmask_mm = ak.any( np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2) < 10., axis=1, keepdims=True) eeSFOS_pairs = ee_pairs[ (np.abs(ee_pairs.e0.pdgId) == np.abs(ee_pairs.e1.pdgId)) & (ee_pairs.e0.charge != ee_pairs.e1.charge)] offZmask_ee = ak.all( np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2) > 10, axis=1, keepdims=True) & (ak.num(eeSFOS_pairs) > 0) onZmask_ee = ak.any( np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2) < 10, axis=1, keepdims=True) # Create masks **for event selection** eeeOnZmask = (ak.num(onZmask_ee[onZmask_ee]) > 0) eeeOffZmask = (ak.num(offZmask_ee[offZmask_ee]) > 0) mmmOnZmask = (ak.num(onZmask_mm[onZmask_mm]) > 0) mmmOffZmask = (ak.num(offZmask_mm[offZmask_mm]) > 0) # Now we need to create masks for the leptons in order to select leptons from the Z boson candidate (in onZ categories) ZeeMask = ak.argmin(np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2), axis=1, keepdims=True) ZmmMask = ak.argmin(np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2), axis=1, keepdims=True) Zee = eeSFOS_pairs[ZeeMask] Zmm = mmSFOS_pairs[ZmmMask] eZ0 = Zee.e0[ak.num(eeSFOS_pairs) > 0] eZ1 = Zee.e1[ak.num(eeSFOS_pairs) > 0] eZ = eZ0 + eZ1 mZ0 = Zmm.m0[ak.num(mmSFOS_pairs) > 0] mZ1 = Zmm.m1[ak.num(mmSFOS_pairs) > 0] mZ = mZ0 + mZ1 mZ_eee = eZ.mass mZ_mmm = mZ.mass # And for the W boson ZmmIndices = mm_pairs_index[ZmmMask] ZeeIndices = ee_pairs_index[ZeeMask] eW = eee[~ZeeIndices.e0 | ~ZeeIndices.e1] mW = mmm[~ZmmIndices.m0 | ~ZmmIndices.m1] triElec = eee_leps.e0 + eee_leps.e1 + eee_leps.e2 triMuon = mmm_leps.m0 + mmm_leps.m1 + mmm_leps.m2 m3l_eee = triElec.mass m3l_mmm = triMuon.mass ################################################################## ### >=4 leptons ################################################################## # 4lep cat is4lmask = ((nElec + nMuon) >= 4) muon_4l = mu[(is4lmask) & (mu.pt > -1)] elec_4l = e[(is4lmask) & (e.pt > -1)] # selecting 4 leading leptons leptons = ak.concatenate([e, mu], axis=-1) leptons_sorted = leptons[ak.argsort(leptons.pt, axis=-1, ascending=False)] lep4l = leptons_sorted[:, 0:4] e4l = lep4l[abs(lep4l.pdgId) == 11] mu4l = lep4l[abs(lep4l.pdgId) == 13] nElec4l = ak.num(e4l) nMuon4l = ak.num(mu4l) # Triggers trig_eeSS = passTrigger(events, 'ee', isData, dataset) trig_mmSS = passTrigger(events, 'mm', isData, dataset) trig_emSS = passTrigger(events, 'em', isData, dataset) trig_eee = passTrigger(events, 'eee', isData, dataset) trig_mmm = passTrigger(events, 'mmm', isData, dataset) trig_eem = passTrigger(events, 'eem', isData, dataset) trig_mme = passTrigger(events, 'mme', isData, dataset) trig_4l = triggerFor4l(events, nMuon, nElec, isData, dataset) # MET filters # Weights genw = np.ones_like(events['event']) if ( isData or len(self._wc_names_lst) > 0) else events['genWeight'] ### We need weights for: normalization, lepSF, triggerSF, pileup, btagSF... weights = {} for r in [ 'all', 'ee', 'mm', 'em', 'eee', 'mmm', 'eem', 'mme', 'eeee', 'eeem', 'eemm', 'mmme', 'mmmm' ]: # weights[r] = coffea.analysis_tools.Weights(len(events)) weights[r] = coffea.analysis_tools.Weights(len(events), storeIndividual=True) if len(self._wc_names_lst) > 0: sow = np.ones_like( sow ) # Not valid in nanoAOD for EFT samples, MUST use SumOfEFTweights at analysis level weights[r].add('norm', genw if isData else (xsec / sow) * genw) weights[r].add('btagSF', btagSF, btagSFUp, btagSFDo) weights[r].add('lepSF', events.lepSF_nom, events.lepSF_hi, events.lepSF_lo) # Extract the EFT quadratic coefficients and optionally use them to calculate the coefficients on the w**2 quartic function # eft_coeffs is never Jagged so convert immediately to numpy for ease of use. eft_coeffs = ak.to_numpy(events['EFTfitCoefficients']) if hasattr( events, "EFTfitCoefficients") else None if eft_coeffs is not None: # Check to see if the ordering of WCs for this sample matches what want if self._samples[dataset]['WCnames'] != self._wc_names_lst: eft_coeffs = efth.remap_coeffs( self._samples[dataset]['WCnames'], self._wc_names_lst, eft_coeffs) eft_w2_coeffs = efth.calc_w2_coeffs(eft_coeffs, self._dtype) if ( self._do_errors and eft_coeffs is not None) else None # Selections and cuts selections = PackedSelection() #(dtype='uint64') channels2LSS = ['eeSSonZ', 'eeSSoffZ', 'mmSSonZ', 'mmSSoffZ', 'emSS'] selections.add('eeSSonZ', (eeonZmask) & (eeSSmask) & (trig_eeSS)) selections.add('eeSSoffZ', (eeoffZmask) & (eeSSmask) & (trig_eeSS)) selections.add('mmSSonZ', (mmonZmask) & (mmSSmask) & (trig_mmSS)) selections.add('mmSSoffZ', (mmoffZmask) & (mmSSmask) & (trig_mmSS)) selections.add('emSS', (emSSmask) & (trig_emSS)) channels3L = ['eemSSonZ', 'eemSSoffZ', 'mmeSSonZ', 'mmeSSoffZ'] selections.add('eemSSonZ', (ee_eemZmask) & (trig_eem)) selections.add('eemSSoffZ', (ee_eemOffZmask) & (trig_eem)) selections.add('mmeSSonZ', (mm_mmeZmask) & (trig_mme)) selections.add('mmeSSoffZ', (mm_mmeOffZmask) & (trig_mme)) channels3L += ['eeeSSonZ', 'eeeSSoffZ', 'mmmSSonZ', 'mmmSSoffZ'] selections.add('eeeSSonZ', (eeeOnZmask) & (trig_eee)) selections.add('eeeSSoffZ', (eeeOffZmask) & (trig_eee)) selections.add('mmmSSonZ', (mmmOnZmask) & (trig_mmm)) selections.add('mmmSSoffZ', (mmmOffZmask) & (trig_mmm)) channels4L = ['eeee', 'eeem', 'eemm', 'mmme', 'mmmm'] selections.add('eeee', ((nElec4l == 4) & (nMuon4l == 0)) & (trig_4l)) selections.add('eeem', ((nElec4l == 3) & (nMuon4l == 1)) & (trig_4l)) selections.add('eemm', ((nElec4l == 2) & (nMuon4l == 2)) & (trig_4l)) selections.add('mmme', ((nElec4l == 1) & (nMuon4l == 3)) & (trig_4l)) selections.add('mmmm', ((nElec4l == 0) & (nMuon4l == 4)) & (trig_4l)) selections.add('ch+', (sumcharge > 0)) selections.add('ch-', (sumcharge < 0)) selections.add('ch0', (sumcharge == 0)) levels = ['base', '1+bm2+bl', '1bm', '2+bm'] selections.add('base', (nElec + nMuon >= 2)) selections.add('1+bm2+bl', (nElec + nMuon >= 2) & ((nbtagsm >= 1) & (nbtagsl >= 2))) selections.add('1bm', (nElec + nMuon >= 2) & (nbtagsm == 1)) selections.add('2+bm', (nElec + nMuon >= 2) & (nbtagsm >= 2)) # Variables invMass_eeSSonZ = (eeSSonZ.e0 + eeSSonZ.e1).mass invMass_eeSSoffZ = (eeSSoffZ.e0 + eeSSoffZ.e1).mass invMass_mmSSonZ = (mmSSonZ.m0 + mmSSonZ.m1).mass invMass_mmSSoffZ = (mmSSoffZ.m0 + mmSSoffZ.m1).mass invMass_emSS = (emSS.e + emSS.m).mass varnames = {} varnames['met'] = met.pt varnames['ht'] = ht varnames['njets'] = njets varnames['invmass'] = { 'eeSSonZ': invMass_eeSSonZ, 'eeSSoffZ': invMass_eeSSoffZ, 'mmSSonZ': invMass_mmSSonZ, 'mmSSoffZ': invMass_mmSSoffZ, 'emSS': invMass_emSS, 'eemSSonZ': mZ_eem, 'eemSSoffZ': mZ_eem, 'mmeSSonZ': mZ_mme, 'mmeSSoffZ': mZ_mme, 'eeeSSonZ': mZ_eee, 'eeeSSoffZ': mZ_eee, 'mmmSSonZ': mZ_mmm, 'mmmSSoffZ': mZ_mmm, } varnames['m3l'] = { 'eemSSonZ': m3l_eem, 'eemSSoffZ': m3l_eem, 'mmeSSonZ': m3l_mme, 'mmeSSoffZ': m3l_mme, 'eeeSSonZ': m3l_eee, 'eeeSSoffZ': m3l_eee, 'mmmSSonZ': m3l_mmm, 'mmmSSoffZ': m3l_mmm, } varnames['e0pt'] = e0.pt varnames['e0eta'] = e0.eta varnames['m0pt'] = m0.pt varnames['m0eta'] = m0.eta varnames['j0pt'] = j0.pt varnames['j0eta'] = j0.eta varnames['counts'] = np.ones_like(events['event']) # systematics systList = [] if isData == False: systList = ['nominal'] if self._do_systematics: systList = systList + [ 'lepSFUp', 'lepSFDown', 'btagSFUp', 'btagSFDown' ] else: systList = ['noweight'] # fill Histos hout = self.accumulator.identity() normweights = weights['all'].weight().flatten( ) # Why does it not complain about .flatten() here? sowweights = np.ones_like(normweights) if len( self._wc_names_lst) > 0 else normweights hout['SumOfEFTweights'].fill(sample=histAxisName, SumOfEFTweights=varnames['counts'], weight=sowweights, eft_coeff=eft_coeffs, eft_err_coeff=eft_w2_coeffs) for syst in systList: for var, v in varnames.items(): for ch in channels2LSS + channels3L + channels4L: for sumcharge in ['ch+', 'ch-', 'ch0']: for lev in levels: #find the event weight to be used when filling the histograms weightSyst = syst #in the case of 'nominal', or the jet energy systematics, no weight systematic variation is used (weightSyst=None) if syst in [ 'nominal', 'JERUp', 'JERDown', 'JESUp', 'JESDown' ]: weightSyst = None # no weight systematic for these variations if syst == 'noweight': weight = np.ones(len(events)) # for data else: # call weights.weight() with the name of the systematic to be varied if ch in channels3L: ch_w = ch[:3] elif ch in channels2LSS: ch_w = ch[:2] else: ch_w = ch weight = weights['all'].weight( weightSyst ) if isData else weights[ch_w].weight( weightSyst) cuts = [ch] + [lev] + [sumcharge] cut = selections.all(*cuts) weights_flat = weight[cut].flatten( ) # Why does it not complain about .flatten() here? weights_ones = np.ones_like(weights_flat, dtype=np.int) eft_coeffs_cut = eft_coeffs[ cut] if eft_coeffs is not None else None eft_w2_coeffs_cut = eft_w2_coeffs[ cut] if eft_w2_coeffs is not None else None # filling histos if var == 'invmass': if ((ch in [ 'eeeSSoffZ', 'mmmSSoffZ', 'eeeSSonZ', 'mmmSSonZ' ]) or (ch in channels4L)): continue else: values = ak.flatten(v[ch][cut]) hout['invmass'].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, invmass=values, weight=weights_flat, systematic=syst) elif var == 'm3l': if ((ch in channels2LSS) or (ch in [ 'eeeSSoffZ', 'mmmSSoffZ', 'eeeSSonZ', 'mmmSSonZ' ]) or (ch in channels4L)): continue values = ak.flatten(v[ch][cut]) hout['m3l'].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, m3l=values, weight=weights_flat, systematic=syst) else: values = v[cut] # These all look identical, do we need if/else here? if var == 'ht': hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, ht=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'met': hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, met=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'njets': hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, njets=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'nbtags': hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, nbtags=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'counts': hout[var].fill(counts=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_ones, systematic=syst) elif var == 'j0eta': if lev == 'base': continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, j0eta=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'e0pt': if ch in [ 'mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ', 'mmmm' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, e0pt=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst ) # Crashing here, not sure why. Related to values? elif var == 'm0pt': if ch in [ 'eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ', 'eeee' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, m0pt=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'e0eta': if ch in [ 'mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ', 'mmmm' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, e0eta=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'm0eta': if ch in [ 'eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ', 'eeee' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, m0eta=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'j0pt': if lev == 'base': continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, j0pt=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) return hout
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): output = self.accumulator.identity() # we can use a very loose preselection to filter the events. nothing is done with this presel, though presel = ak.num(events.Jet) > 0 if self.year == 2016: lumimask = LumiMask( '../data/lumi/Cert_271036-284044_13TeV_Legacy2016_Collisions16_JSON.txt' ) if self.year == 2017: lumimask = LumiMask( '../data/lumi/Cert_294927-306462_13TeV_UL2017_Collisions17_GoldenJSON.txt' ) if self.year == 2018: lumimask = LumiMask( '../data/lumi/Cert_314472-325175_13TeV_Legacy2018_Collisions18_JSON.txt' ) ev = events[presel] dataset = ev.metadata['dataset'] # load the config - probably not needed anymore cfg = loadConfig() output['totalEvents']['all'] += len(events) output['skimmedEvents']['all'] += len(ev) if self.year == 2018: triggers = ev.HLT.Ele23_Ele12_CaloIdL_TrackIdL_IsoVL elif self.year == 2017: triggers = ev.HLT.Ele23_Ele12_CaloIdL_TrackIdL_IsoVL elif self.year == 2016: triggers = ev.HLT.Ele23_Ele12_CaloIdL_TrackIdL_IsoVL_DZ ## Electrons electron = Collections(ev, "Electron", "tightFCNC", 0, self.year).get() electron = electron[(electron.pt > 25) & (np.abs(electron.eta) < 2.4)] loose_electron = Collections(ev, "Electron", "looseFCNC", 0, self.year).get() loose_electron = loose_electron[(loose_electron.pt > 25) & (np.abs(loose_electron.eta) < 2.4)] SSelectron = (ak.sum(electron.charge, axis=1) != 0) & (ak.num(electron) == 2) OSelectron = (ak.sum(electron.charge, axis=1) == 0) & (ak.num(electron) == 2) dielectron = choose(electron, 2) dielectron_mass = (dielectron['0'] + dielectron['1']).mass dielectron_pt = (dielectron['0'] + dielectron['1']).pt leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1)) leading_electron = electron[(leading_electron_idx)] leading_electron = leading_electron[(leading_electron.pt > 30)] trailing_electron_idx = ak.singletons(ak.argmin(electron.pt, axis=1)) trailing_electron = electron[trailing_electron_idx] ##Muons loose_muon = Collections(ev, "Muon", "looseFCNC", 0, self.year).get() loose_muon = loose_muon[(loose_muon.pt > 20) & (np.abs(loose_muon.eta) < 2.4)] #jets jet = getJets(ev, minPt=40, maxEta=2.4, pt_var='pt') jet = jet[~match(jet, loose_muon, deltaRCut=0.4)] # remove jets that overlap with muons jet = jet[~match( jet, electron, deltaRCut=0.4)] # remove jets that overlap with electrons ## MET -> can switch to puppi MET met_pt = ev.MET.pt met_phi = ev.MET.phi #weights weight = Weights(len(ev)) weight2 = Weights(len(ev)) weight2.add("charge flip", self.charge_flip_ratio.flip_weight(electron)) #selections filters = getFilters(ev, year=self.year, dataset=dataset, UL=False) mask = lumimask(ev.run, ev.luminosityBlock) ss = (SSelectron) os = (OSelectron) mass = (ak.min(np.abs(dielectron_mass - 91.2), axis=1) < 15) lead_electron = (ak.min(leading_electron.pt, axis=1) > 30) jet1 = (ak.num(jet) >= 1) jet2 = (ak.num(jet) >= 2) num_loose = ((ak.num(loose_electron) == 2) & (ak.num(loose_muon) == 0)) selection = PackedSelection() selection.add('filter', (filters)) selection.add('mask', (mask)) selection.add('ss', ss) selection.add('os', os) selection.add('mass', mass) selection.add('leading', lead_electron) selection.add('triggers', triggers) selection.add('one jet', jet1) selection.add('two jets', jet2) selection.add('num_loose', num_loose) bl_reqs = ['filter'] + ['triggers'] + ['mask'] bl_reqs_d = {sel: True for sel in bl_reqs} baseline = selection.require(**bl_reqs_d) s_reqs = bl_reqs + ['ss'] + ['mass'] + ['num_loose'] + ['leading'] s_reqs_d = {sel: True for sel in s_reqs} ss_sel = selection.require(**s_reqs_d) o_reqs = bl_reqs + ['os'] + ['mass'] + ['num_loose'] + ['leading'] o_reqs_d = {sel: True for sel in o_reqs} os_sel = selection.require(**o_reqs_d) j1s_reqs = s_reqs + ['one jet'] j1s_reqs_d = {sel: True for sel in j1s_reqs} j1ss_sel = selection.require(**j1s_reqs_d) j1o_reqs = o_reqs + ['one jet'] j1o_reqs_d = {sel: True for sel in j1o_reqs} j1os_sel = selection.require(**j1o_reqs_d) j2s_reqs = s_reqs + ['two jets'] j2s_reqs_d = {sel: True for sel in j2s_reqs} j2ss_sel = selection.require(**j2s_reqs_d) j2o_reqs = o_reqs + ['two jets'] j2o_reqs_d = {sel: True for sel in j2o_reqs} j2os_sel = selection.require(**j2o_reqs_d) #outputs output["electron_data1"].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(leading_electron[os_sel].pt)), eta=ak.to_numpy(ak.flatten(leading_electron[os_sel].eta)), phi=ak.to_numpy(ak.flatten(leading_electron[os_sel].phi)), weight=weight2.weight()[os_sel]) output["electron_data2"].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(trailing_electron[os_sel].pt)), eta=ak.to_numpy(ak.flatten(trailing_electron[os_sel].eta)), phi=ak.to_numpy(ak.flatten(trailing_electron[os_sel].phi)), weight=weight2.weight()[os_sel]) output["electron_data3"].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(leading_electron[j1os_sel].pt)), eta=ak.to_numpy(ak.flatten(leading_electron[j1os_sel].eta)), phi=ak.to_numpy(ak.flatten(leading_electron[j1os_sel].phi)), weight=weight2.weight()[j1os_sel]) output["electron_data4"].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(trailing_electron[j1os_sel].pt)), eta=ak.to_numpy(ak.flatten(trailing_electron[j1os_sel].eta)), phi=ak.to_numpy(ak.flatten(trailing_electron[j1os_sel].phi)), weight=weight2.weight()[j1os_sel]) output["electron_data5"].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(leading_electron[j2os_sel].pt)), eta=ak.to_numpy(ak.flatten(leading_electron[j2os_sel].eta)), phi=ak.to_numpy(ak.flatten(leading_electron[j2os_sel].phi)), weight=weight2.weight()[j2os_sel]) output["electron_data6"].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(trailing_electron[j2os_sel].pt)), eta=ak.to_numpy(ak.flatten(trailing_electron[j2os_sel].eta)), phi=ak.to_numpy(ak.flatten(trailing_electron[j2os_sel].phi)), weight=weight2.weight()[j2os_sel]) output["electron_data7"].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(leading_electron[ss_sel].pt)), eta=ak.to_numpy(ak.flatten(leading_electron[ss_sel].eta)), phi=ak.to_numpy(ak.flatten(leading_electron[ss_sel].phi)), weight=weight.weight()[ss_sel]) output["electron_data8"].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(trailing_electron[ss_sel].pt)), eta=ak.to_numpy(ak.flatten(trailing_electron[ss_sel].eta)), phi=ak.to_numpy(ak.flatten(trailing_electron[ss_sel].phi)), weight=weight.weight()[ss_sel]) output["electron_data9"].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(leading_electron[j1ss_sel].pt)), eta=ak.to_numpy(ak.flatten(leading_electron[j1ss_sel].eta)), phi=ak.to_numpy(ak.flatten(leading_electron[j1ss_sel].phi)), weight=weight.weight()[j1ss_sel]) output["electron_data10"].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(trailing_electron[j1ss_sel].pt)), eta=ak.to_numpy(ak.flatten(trailing_electron[j1ss_sel].eta)), phi=ak.to_numpy(ak.flatten(trailing_electron[j1ss_sel].phi)), weight=weight.weight()[j1ss_sel]) output["electron_data11"].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(leading_electron[j2ss_sel].pt)), eta=ak.to_numpy(ak.flatten(leading_electron[j2ss_sel].eta)), phi=ak.to_numpy(ak.flatten(leading_electron[j2ss_sel].phi)), weight=weight.weight()[j2ss_sel]) output["electron_data12"].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(trailing_electron[j2ss_sel].pt)), eta=ak.to_numpy(ak.flatten(trailing_electron[j2ss_sel].eta)), phi=ak.to_numpy(ak.flatten(trailing_electron[j2ss_sel].phi)), weight=weight.weight()[j2ss_sel]) output["dilep_mass1"].fill( dataset=dataset, mass=ak.to_numpy(ak.flatten(dielectron_mass[os_sel])), pt=ak.to_numpy(ak.flatten(dielectron_pt[os_sel])), weight=weight2.weight()[os_sel]) output["dilep_mass2"].fill( dataset=dataset, mass=ak.to_numpy(ak.flatten(dielectron_mass[j1os_sel])), pt=ak.to_numpy(ak.flatten(dielectron_pt[j1os_sel])), weight=weight2.weight()[j1os_sel]) output["dilep_mass3"].fill( dataset=dataset, mass=ak.to_numpy(ak.flatten(dielectron_mass[j2os_sel])), pt=ak.to_numpy(ak.flatten(dielectron_pt[j2os_sel])), weight=weight2.weight()[j2os_sel]) output["dilep_mass4"].fill( dataset=dataset, mass=ak.to_numpy(ak.flatten(dielectron_mass[ss_sel])), pt=ak.to_numpy(ak.flatten(dielectron_pt[ss_sel])), weight=weight.weight()[ss_sel]) output["dilep_mass5"].fill( dataset=dataset, mass=ak.to_numpy(ak.flatten(dielectron_mass[j1ss_sel])), pt=ak.to_numpy(ak.flatten(dielectron_pt[j1ss_sel])), weight=weight.weight()[j1ss_sel]) output["dilep_mass6"].fill( dataset=dataset, mass=ak.to_numpy(ak.flatten(dielectron_mass[j2ss_sel])), pt=ak.to_numpy(ak.flatten(dielectron_pt[j2ss_sel])), weight=weight.weight()[j2ss_sel]) output["MET"].fill(dataset=dataset, pt=met_pt[os_sel], weight=weight2.weight()[os_sel]) output["MET2"].fill(dataset=dataset, pt=met_pt[j1os_sel], weight=weight2.weight()[j1os_sel]) output["MET3"].fill(dataset=dataset, pt=met_pt[j2os_sel], weight=weight2.weight()[j2os_sel]) output["MET4"].fill(dataset=dataset, pt=met_pt[ss_sel], weight=weight.weight()[ss_sel]) output["MET5"].fill(dataset=dataset, pt=met_pt[j1ss_sel], weight=weight.weight()[j1ss_sel]) output["MET6"].fill(dataset=dataset, pt=met_pt[j2ss_sel], weight=weight.weight()[j2ss_sel]) output["N_jet"].fill(dataset=dataset, multiplicity=ak.num(jet)[os_sel], weight=weight2.weight()[os_sel]) output["N_jet2"].fill(dataset=dataset, multiplicity=ak.num(jet)[j1os_sel], weight=weight2.weight()[j1os_sel]) output["N_jet3"].fill(dataset=dataset, multiplicity=ak.num(jet)[j2os_sel], weight=weight2.weight()[j2os_sel]) output["N_jet4"].fill(dataset=dataset, multiplicity=ak.num(jet)[ss_sel], weight=weight.weight()[ss_sel]) output["N_jet5"].fill(dataset=dataset, multiplicity=ak.num(jet)[j1ss_sel], weight=weight.weight()[j1ss_sel]) output["N_jet6"].fill(dataset=dataset, multiplicity=ak.num(jet)[j2ss_sel], weight=weight.weight()[j2ss_sel]) output["PV_npvsGood"].fill(dataset=dataset, multiplicity=ev.PV[os_sel].npvsGood, weight=weight2.weight()[os_sel]) output["PV_npvsGood2"].fill(dataset=dataset, multiplicity=ev.PV[j1os_sel].npvsGood, weight=weight2.weight()[j1os_sel]) output["PV_npvsGood3"].fill(dataset=dataset, multiplicity=ev.PV[j2os_sel].npvsGood, weight=weight2.weight()[j2os_sel]) output["PV_npvsGood4"].fill(dataset=dataset, multiplicity=ev.PV[ss_sel].npvsGood, weight=weight.weight()[ss_sel]) output["PV_npvsGood5"].fill(dataset=dataset, multiplicity=ev.PV[j1ss_sel].npvsGood, weight=weight.weight()[j1ss_sel]) output["PV_npvsGood6"].fill(dataset=dataset, multiplicity=ev.PV[j2ss_sel].npvsGood, weight=weight.weight()[j2ss_sel]) return output