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
def add_mutriggerSF(weights, leadingmuon, year, selection): def mask(w): return np.where(selection.all('onemuon'), w, 1.) mu_pt = np.array(ak.fill_none(leadingmuon.pt, 0.)) mu_eta = np.array(ak.fill_none(abs(leadingmuon.eta), 0.)) nom = mask(compiled[f'{year}_mutrigweight_pt_abseta'](mu_pt, mu_eta)) shift = mask(compiled[f'{year}_mutrigweight_pt_abseta_mutrigweightShift']( mu_pt, mu_eta)) weights.add('mu_trigger', nom, shift, shift=True)
def test_fill_none_axis_deprecated(): ak.deprecations_as_errors = False with pytest.deprecated_call(): filled = ak.fill_none(array, 10) assert ak.to_list(filled) == [[None, 2], 10, [4, None]] with pytest.deprecated_call(): filled_twice = ak.fill_none(filled, 20) assert ak.to_list(filled_twice) == [[20, 2], 10, [4, 20]]
def isCleanJet(jets, electrons, muons, taus, drmin=0.4): ''' Returns mask to select clean jets ''' # Think we can do the jet cleaning like this? It's based on L355 of https://github.com/areinsvo/TTGamma_LongExercise/blob/FullAnalysis/ttgamma/processor.py # However, not really sure how "nearest" works, so probably should look into that more at some point # But this new version of the function does seem to be returning the same mask as the old version jetEle, jetEleDR = jets.nearest(electrons, return_metric=True) jetMu, jetMuDR = jets.nearest(muons, returyyppn_metric=True) jetTau, jetTauDR = jets.nearest(taus, returyyppn_metric=True) jetEleMask = ak.fill_none(jetEleDR > drmin, True) jetMuMask = ak.fill_none(jetMuDR > drmin, True) jetTauMask = ak.fill_none(jetTauDR > drmin, True) return (jetEleMask & jetMuMask & jetTauMask)
def add_jetTriggerSF(weights, leadingjet, year, selection): def mask(w): return np.where(selection.all('noleptons'), w, 1.) jet_pt = np.array(ak.fill_none(leadingjet.pt, 0.)) jet_msd = np.array(ak.fill_none(leadingjet.msoftdrop, 0.)) # note: uncorrected nom = mask(jet_triggerSF[f'fatjet_triggerSF{year}'].evaluate( "nominal", jet_pt, jet_msd)) up = mask(jet_triggerSF[f'fatjet_triggerSF{year}'].evaluate( "stat_up", jet_pt, jet_msd)) down = mask(jet_triggerSF[f'fatjet_triggerSF{year}'].evaluate( "stat_dn", jet_pt, jet_msd)) weights.add('jet_trigger', nom, up, down)
def test_highlevel(): array = ak.Array([[1.1, 2.2, None, 3.3], [], [4.4, None, 5.5]]) assert ak.to_list(ak.fill_none(array, 999, axis=1)) == [ [1.1, 2.2, 999, 3.3], [], [4.4, 999, 5.5], ] assert ak.to_list(ak.fill_none(array, [1, 2, 3], axis=1)) == [ [1.1, 2.2, [1, 2, 3], 3.3], [], [4.4, [1, 2, 3], 5.5], ] assert ak.to_list(ak.fill_none(array, [], axis=1)) == [ [1.1, 2.2, [], 3.3], [], [4.4, [], 5.5], ] assert ak.to_list(ak.fill_none(array, {"x": 999}, axis=1)) == [ [1.1, 2.2, { "x": 999 }, 3.3], [], [4.4, { "x": 999 }, 5.5], ] array = ak.Array([[1.1, 2.2, 3.3], None, [], None, [4.4, 5.5]]) assert ak.to_list(ak.fill_none(array, 999, axis=0)) == [ [1.1, 2.2, 3.3], 999, [], 999, [4.4, 5.5], ] assert ak.to_list(ak.fill_none(array, [1, 2, 3], axis=0)) == [ [1.1, 2.2, 3.3], [1, 2, 3], [], [1, 2, 3], [4.4, 5.5], ] assert ak.to_list(ak.fill_none(array, {"x": 999}, axis=0)) == [ [1.1, 2.2, 3.3], { "x": 999 }, [], { "x": 999 }, [4.4, 5.5], ] assert ak.to_list(ak.fill_none(array, [], axis=0)) == [ [1.1, 2.2, 3.3], [], [], [], [4.4, 5.5], ]
def find_first_parent(particle, maxgen=10): tmp_mother = particle.parent out_mother_pdg = tmp_mother.pdgId found = ak.zeros_like(particle.pdgId) for i in range(maxgen): update = ak.fill_none( (tmp_mother.pdgId != particle.pdgId) * (found == 0), False) out_mother_pdg = update * ak.fill_none(tmp_mother.pdgId, 0) + (~update) * out_mother_pdg found = found + update tmp_mother = tmp_mother.parent return out_mother_pdg
def add_jec_variables(jets, event_rho): jets["pt_raw"] = (1 - jets.rawFactor) * jets.pt jets["mass_raw"] = (1 - jets.rawFactor) * jets.mass jets["pt_gen"] = ak.values_astype(ak.fill_none(jets.matched_gen.pt, 0), np.float32) jets["event_rho"] = ak.broadcast_arrays(event_rho, jets.pt)[0] return jets
def run(): events = NanoEventsFactory.from_root( os.path.abspath("tests/samples/nano_dy.root"), persistent_cache=array_log, ).events() jets = events.Jet met = events.MET jets["pt_raw"] = (1 - jets["rawFactor"]) * jets["pt"] jets["mass_raw"] = (1 - jets["rawFactor"]) * jets["mass"] jets["pt_gen"] = ak.values_astype( ak.fill_none(jets.matched_gen.pt, 0.0), np.float32) jets["rho"] = ak.broadcast_arrays(events.fixedGridRhoFastjetAll, jets.pt)[0] jec_cache = cachetools.Cache(np.inf) weakref.finalize(jec_cache, jec_finalized.set) corrected_jets = jet_factory.build(jets, lazy_cache=jec_cache) corrected_met = met_factory.build(met, corrected_jets, lazy_cache=jec_cache) print(corrected_met.pt_orig) print(corrected_met.pt) for unc in jet_factory.uncertainties() + met_factory.uncertainties(): print(unc, corrected_met[unc].up.pt) print(unc, corrected_met[unc].down.pt) for unc in jet_factory.uncertainties(): print(unc, corrected_jets[unc].up.pt) print("Finalized:", array_log.finalized)
def set_genZ(events, genBranches, selection_options, debug): if genBranches is None: events["genZ_decayMode"] = ak.from_numpy(-1 * numpy.ones(len(events))) else: electron_idxs = abs(genBranches.pdgId) == 11 muon_idxs = abs(genBranches.pdgId) == 13 tau_idxs = abs(genBranches.pdgId) == 15 motherOfElectrons = genBranches.genPartIdxMother[electron_idxs] motherOfMuons = genBranches.genPartIdxMother[muon_idxs] motherOfTaus = genBranches.genPartIdxMother[tau_idxs] ZToEleEvents = ak.sum( (genBranches.pdgId[motherOfElectrons] == 23), axis=1) > 0 ZToMuEvents = ak.sum( (genBranches.pdgId[motherOfMuons] == 23), axis=1) > 0 ZToTauEvents = ak.sum( (genBranches.pdgId[motherOfTaus] == 23), axis=1) > 0 # W decay dudes WToEEvents = ak.sum( (abs(genBranches.pdgId[motherOfElectrons]) == 24), axis=1) > 0 WToMuEvents = ak.sum( (abs(genBranches.pdgId[motherOfMuons]) == 24), axis=1) > 0 WToTauEvents = ak.sum( (abs(genBranches.pdgId[motherOfTaus]) == 24), axis=1) > 0 events["genZ_decayMode"] = ak.from_numpy( numpy.zeros(len(events)) ) + 1 * ZToEleEvents + 2 * ZToMuEvents + 3 * ZToTauEvents + 4 * WToEEvents + 5 * WToMuEvents + 6 * WToTauEvents events["tau_motherID"] = ak.fill_none( ak.firsts(genBranches.pdgId[motherOfTaus]), 0) return events
def combine(eff, sf): # tagged SF = SF*eff / eff = SF tagged_sf = awkward.prod(sf[passbtag], axis=-1) # untagged SF = (1 - SF*eff) / (1 - eff) untagged_sf = awkward.prod(((1 - sf * eff) / (1 - eff))[~passbtag], axis=-1) return awkward.fill_none(tagged_sf * untagged_sf, 1.) # TODO: move None guard to coffea
def apply_roccor(df, rochester, is_mc): if is_mc: hasgen = ~np.isnan(ak.fill_none(df.Muon.matched_gen.pt, np.nan)) mc_rand = np.random.rand(*ak.to_numpy(ak.flatten(df.Muon.pt)).shape) mc_rand = ak.unflatten(mc_rand, ak.num(df.Muon.pt, axis=1)) corrections = np.array(ak.flatten(ak.ones_like(df.Muon.pt))) errors = np.array(ak.flatten(ak.ones_like(df.Muon.pt))) mc_kspread = rochester.kSpreadMC( df.Muon.charge[hasgen], df.Muon.pt[hasgen], df.Muon.eta[hasgen], df.Muon.phi[hasgen], df.Muon.matched_gen.pt[hasgen], ) mc_ksmear = rochester.kSmearMC( df.Muon.charge[~hasgen], df.Muon.pt[~hasgen], df.Muon.eta[~hasgen], df.Muon.phi[~hasgen], df.Muon.nTrackerLayers[~hasgen], mc_rand[~hasgen], ) errspread = rochester.kSpreadMCerror( df.Muon.charge[hasgen], df.Muon.pt[hasgen], df.Muon.eta[hasgen], df.Muon.phi[hasgen], df.Muon.matched_gen.pt[hasgen], ) errsmear = rochester.kSmearMCerror( df.Muon.charge[~hasgen], df.Muon.pt[~hasgen], df.Muon.eta[~hasgen], df.Muon.phi[~hasgen], df.Muon.nTrackerLayers[~hasgen], mc_rand[~hasgen], ) hasgen_flat = np.array(ak.flatten(hasgen)) corrections[hasgen_flat] = np.array(ak.flatten(mc_kspread)) corrections[~hasgen_flat] = np.array(ak.flatten(mc_ksmear)) errors[hasgen_flat] = np.array(ak.flatten(errspread)) errors[~hasgen_flat] = np.array(ak.flatten(errsmear)) corrections = ak.unflatten(corrections, ak.num(df.Muon.pt, axis=1)) errors = ak.unflatten(errors, ak.num(df.Muon.pt, axis=1)) else: corrections = rochester.kScaleDT(df.Muon.charge, df.Muon.pt, df.Muon.eta, df.Muon.phi) errors = rochester.kScaleDTerror(df.Muon.charge, df.Muon.pt, df.Muon.eta, df.Muon.phi) df["Muon", "pt_roch"] = df.Muon.pt * corrections df["Muon", "pt_roch_up"] = df.Muon.pt_roch + df.Muon.pt * errors df["Muon", "pt_roch_down"] = df.Muon.pt_roch - df.Muon.pt * errors
def test(): first = ak.from_numpy(np.array([1, 2, 3])) deltas = ak.Array([[1, 2], [1, 2], [1, 2, 3]]) assert np.hstack( (first[:, np.newaxis], ak.fill_none(ak.pad_none(deltas, 3, axis=-1), 999))).tolist() == [[1, 1, 2, 999], [2, 1, 2, 999], [3, 1, 2, 3]]
def prepare_jets(df, is_mc): # Initialize missing fields (needed for JEC) df["Jet", "pt_raw"] = (1 - df.Jet.rawFactor) * df.Jet.pt df["Jet", "mass_raw"] = (1 - df.Jet.rawFactor) * df.Jet.mass df["Jet", "rho"] = ak.broadcast_arrays(df.fixedGridRhoFastjetAll, df.Jet.pt)[0] if is_mc: df["Jet", "pt_gen"] = ak.values_astype(ak.fill_none(df.Jet.matched_gen.pt, 0), np.float32)
def studyAcc(ifile, title, acceptances_dict, mass): dq_events = getData(ifile, "Truth") # choosing the nevts for denominator (because all files have 10000 evts) den = len(dq_events) # choosing 500 evts for denominator den500 = False if den500: dq_events = dq_events[:500] den = 500 dq_showers = dq_events["Showers"] dq_electrons = dq_events["Electrons"] dq_sh_z = ak.fill_none(ak.pad_none(dq_showers.sz, 2, axis=1), 0) dq_sh_e = ak.fill_none(ak.pad_none(dq_showers.sedep_ecal, 2, axis=1), -1) dq_e_ge = ak.fill_none(ak.pad_none(dq_electrons.ge, 2, axis=1), -1) shower_zero_mask = ak.any(dq_sh_e == 0, axis=1) shower_zero_index = np.where(shower_zero_mask) eshower_mask = ak.all(dq_sh_e > 0.1 * dq_e_ge, axis=1) eshower_index = np.where(eshower_mask) masks = {'shower_zero': shower_zero_index, 'eshower': eshower_index} dq_evts = dq_events[ak.all(dq_e_ge > -1, axis=1)] # print(acceptances_dict) # we can print the mass, the number of events with gen e-, the number of events in the sample, and the denominator print(mass, len(dq_evts), len(dq_events), den) for maskstr, m in masks.items(): if maskstr in acceptances_dict: acceptances_dict[maskstr].append(len(m[0]) / den) else: acceptances_dict[maskstr] = [len(m[0]) / den] #print(maskstr+':', len(m[0])/den) return acceptances_dict
def local2global(stack): """Turn jagged local index into global index Signature: index,target_offsets,!local2global Outputs a content array with same shape as index content """ target_offsets = stack.pop() index = stack.pop() index = index.mask[index >= 0] + target_offsets[:-1] index = index.mask[index < target_offsets[1:]] out = numpy.array(awkward.flatten(awkward.fill_none(index, -1))) if out.dtype != numpy.int64: raise RuntimeError stack.append(out)
def get_vpt(check_offshell=False): """Only the leptonic samples have no resonance in the decay tree, and only when M is beyond the configured Breit-Wigner cutoff (usually 15*width) """ boson = ak.firsts( genpart[((genpart.pdgId == 23) | (abs(genpart.pdgId) == 24)) & genpart.hasFlags(["fromHardProcess", "isLastCopy"])]) if check_offshell: offshell = genpart[ genpart.hasFlags(["fromHardProcess", "isLastCopy"]) & ak.is_none(boson) & (abs(genpart.pdgId) >= 11) & (abs(genpart.pdgId) <= 16)].sum() return ak.where(ak.is_none(boson.pt), offshell.pt, boson.pt) return np.array(ak.fill_none(boson.pt, 0.))
def matchJets(self, obj, jet, deltaRCut=0.4): combs = ak.cartesian([obj, jet], nested=True) jet_index = ak.local_index(delta_r( combs['0'], combs['1']))[delta_r(combs['0'], combs['1']) < 0.4] jet_index_pad = ak.flatten(ak.fill_none( ak.pad_none(jet_index, target=1, clip=True, axis=2), 0), axis=2) mask = ak.num(jet_index, axis=2) > 0 # a mask for obj with a matched jet mask_match = mask * 1 + ~mask * 0 mask_nomatch = mask * 0 + ~mask * 1 return jet_index_pad, mask_match, mask_nomatch
def test(): array = ak.Array([{"x": 1}, {"x": 2}, None, {"x": 3}]) assert ak.fill_none(array, array[0]).tolist() == [ { "x": 1 }, { "x": 2 }, { "x": 1 }, { "x": 3 }, ]
def fsr_recovery(df): mask = ( (df.Muon.fsrPhotonIdx >= 0) & (df.Muon.matched_fsrPhoton.relIso03 < 1.8) & (df.Muon.matched_fsrPhoton.dROverEt2 < 0.012) & (df.Muon.matched_fsrPhoton.pt / df.Muon.pt < 0.4) & (abs(df.Muon.matched_fsrPhoton.eta) < 2.4) ) mask = ak.fill_none(mask, False) px = ak.zeros_like(df.Muon.pt) py = ak.zeros_like(df.Muon.pt) pz = ak.zeros_like(df.Muon.pt) e = ak.zeros_like(df.Muon.pt) fsr = { "pt": df.Muon.matched_fsrPhoton.pt, "eta": df.Muon.matched_fsrPhoton.eta, "phi": df.Muon.matched_fsrPhoton.phi, "mass": 0.0, } for obj in [df.Muon, fsr]: px_ = obj["pt"] * np.cos(obj["phi"]) py_ = obj["pt"] * np.sin(obj["phi"]) pz_ = obj["pt"] * np.sinh(obj["eta"]) e_ = np.sqrt(px_**2 + py_**2 + pz_**2 + obj["mass"] ** 2) px = px + px_ py = py + py_ pz = pz + pz_ e = e + e_ pt = np.sqrt(px**2 + py**2) eta = np.arcsinh(pz / pt) phi = np.arctan2(py, px) mass = np.sqrt(e**2 - px**2 - py**2 - pz**2) iso = (df.Muon.pfRelIso04_all * df.Muon.pt - df.Muon.matched_fsrPhoton.pt) / pt df["Muon", "pt_fsr"] = ak.where(mask, pt, df.Muon.pt) df["Muon", "eta_fsr"] = ak.where(mask, eta, df.Muon.eta) df["Muon", "phi_fsr"] = ak.where(mask, phi, df.Muon.phi) df["Muon", "mass_fsr"] = ak.where(mask, mass, df.Muon.mass) df["Muon", "iso_fsr"] = ak.where(mask, iso, df.Muon.pfRelIso04_all) return mask
def TTsemileptonicmatch(events): dataset = events.metadata['dataset'] if 'TTToSemiLeptonic' not in dataset: return np.zeros(len(events.FatJet.pt), dtype=bool) child = events.GenPart[ (abs(events.GenPart.pdgId) == 24) & events.GenPart.hasFlags(["isLastCopy", "fromHardProcess"])].children fatjet = ak.firsts(events.FatJet) n_matched_quarks = np.zeros(len(fatjet)) for ii in [0, 1]: for jj in [0, 1]: n_matched_quarks = n_matched_quarks + ak.fill_none( (fatjet.delta_r2(child[:, ii, jj]) < 0.8**2) & (abs(child[:, ii, jj].pdgId) < 6), 0.) print(n_matched_quarks) return n_matched_quarks
def process_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 __init__(self, ev, obj, wp, year=2018, verbose=0): self.obj = obj self.wp = wp if self.wp == None: self.selection_dict = {} else: self.selection_dict = obj_def[self.obj][self.wp] self.v = verbose self.year = year id_level = None if wp.lower().count('veto') or wp.lower().count('loose'): id_level = 0 elif wp.lower().count('fake'): id_level = 1 elif wp.lower().count('tight'): id_level = 2 if self.obj == "Muon": # collections are already there, so we just need to calculate missing ones ev['Muon', 'absMiniIso'] = ev.Muon.miniPFRelIso_all * ev.Muon.pt ev['Muon', 'ptErrRel'] = ev.Muon.ptErr / ev.Muon.pt # this is what we are using: # - jetRelIso if the matched jet is within deltaR<0.4, pfRelIso03_all otherwise # - btagDeepFlavB discriminator of the matched jet if jet is within deltaR<0.4, 0 otherwise # - pt_cone = 0.9*pt of matched jet if jet is within deltaR<0.4, pt/(pt+iso) otherwise mask_close = (ak.fill_none(ev.Muon.delta_r(ev.Muon.matched_jet), 99) < 0.4) * 1 mask_far = ~(ak.fill_none(ev.Muon.delta_r(ev.Muon.matched_jet), 99) < 0.4) * 1 deepJet = ak.fill_none(ev.Muon.matched_jet.btagDeepFlavB, 0) * mask_close + 0 * mask_far jetRelIsoV2 = ev.Muon.jetRelIso * mask_close + ev.Muon.pfRelIso03_all * mask_far # default to 0 if no match #conePt = 0.9 * ak.fill_none(ev.Muon.matched_jet.pt,0) * mask_close + ev.Muon.pt*(1 + ev.Muon.miniPFRelIso_all)*mask_far if self.year == 2017 or self.year == 2018: I_1 = 0.11 I_2 = 0.74 I_3 = 6.8 elif self.year == 2016: I_1 = 0.16 I_2 = 0.76 I_3 = 7.2 PF_unflatten = ak.from_regular( ev.Muon.miniPFRelIso_all[:, :, np.newaxis]) max_miniIso = ak.max( ak.concatenate( [PF_unflatten - I_1, ak.zeros_like(PF_unflatten)], axis=2), axis=2) #equivalent to max(0, ev.Muon.miniPFRelIso_all - I_1) muon_pt_unflatten = ak.from_regular(ev.Muon.pt[:, :, np.newaxis]) jet_pt_unflatten = ak.from_regular( ev.Muon.matched_jet.pt[:, :, np.newaxis]) max_pt = ak.max( ak.concatenate([muon_pt_unflatten, jet_pt_unflatten * I_2], axis=2), axis=2) #max(ev.Muon.pt, ev.Muon.matched_jet.pt * I_2) conePt = (ev.Muon.pt * (1 + max_miniIso)) * (ev.Muon.jetPtRelv2 > I_3) + ( max_pt * ~(ev.Muon.jetPtRelv2 > I_3)) ev['Muon', 'deepJet'] = ak.copy(deepJet) ev['Muon', 'jetRelIsoV2'] = jetRelIsoV2 ev['Muon', 'conePt'] = conePt ev['Muon', 'id'] = ak.ones_like(conePt) * id_level self.cand = ev.Muon elif self.obj == "Electron": # calculate new variables. asignment is awkward, but what can you do. ev['Electron', 'absMiniIso'] = ev.Electron.miniPFRelIso_all * ev.Electron.pt ev['Electron', 'etaSC'] = ev.Electron.eta + ev.Electron.deltaEtaSC if self.year == 2017 or self.year == 2018: I_1 = 0.07 I_2 = 0.78 I_3 = 8.0 elif self.year == 2016: I_1 = 0.12 I_2 = 0.80 I_3 = 7.2 # the following line is only needed if we do our own matching. # right now, we keep using the NanoAOD match, but check the deltaR distance # jet_index, mask_match, mask_nomatch = self.matchJets(ev.Electron, ev.Jet) # this is what we are using: # - jetRelIso if the matched jet is within deltaR<0.4, pfRelIso03_all otherwise # - btagDeepFlavB discriminator of the matched jet if jet is within deltaR<0.4, 0 otherwise # - pt_cone = 0.9*pt of matched jet if jet is within deltaR<0.4, pt/(pt+iso) otherwise mask_close = (ak.fill_none( ev.Electron.delta_r(ev.Electron.matched_jet), 99) < 0.4) * 1 mask_far = ~(ak.fill_none( ev.Electron.delta_r(ev.Electron.matched_jet), 99) < 0.4) * 1 deepJet = ak.fill_none(ev.Electron.matched_jet.btagDeepFlavB, 0) * mask_close jetRelIsoV2 = ev.Electron.jetRelIso * mask_close + ev.Electron.pfRelIso03_all * mask_far # default to 0 if no match #conePt = 0.9 * ak.fill_none(ev.Electron.matched_jet.pt,0) * mask_close + ev.Electron.pt*(1 + ev.Electron.miniPFRelIso_all)*mask_far PF_unflatten = ak.from_regular( ev.Electron.miniPFRelIso_all[:, :, np.newaxis]) max_miniIso = ak.max( ak.concatenate( [PF_unflatten - I_1, ak.zeros_like(PF_unflatten)], axis=2), axis=2) #equivalent to max(0, ev.Muon.miniPFRelIso_all - I_1) electron_pt_unflatten = ak.from_regular(ev.Electron.pt[:, :, np.newaxis]) jet_pt_unflatten = ak.from_regular( ev.Electron.matched_jet.pt[:, :, np.newaxis]) max_pt = ak.max( ak.concatenate([electron_pt_unflatten, jet_pt_unflatten * I_2], axis=2), axis=2) #max(ev.Muon.pt, ev.Muon.matched_jet.pt * I_2) conePt = (ev.Electron.pt * (1 + max_miniIso)) * (ev.Electron.jetPtRelv2 > I_3) + ( max_pt * ~(ev.Electron.jetPtRelv2 > I_3)) ev['Electron', 'deepJet'] = ak.copy(deepJet) ev['Electron', 'jetRelIsoV2'] = jetRelIsoV2 ev['Electron', 'conePt'] = conePt ev['Electron', 'id'] = ak.ones_like(conePt) * id_level ev['Electron', 'jetRelIso'] = ev.Electron.jetRelIso ev['Electron', 'jetPtRelv2'] = ev.Electron.jetPtRelv2 self.cand = ev.Electron self.getSelection() if self.obj == "Electron" and self.wp == "tight": self.selection = self.selection & self.getElectronMVAID( ) & self.getIsolation(0.07, 0.78, 8.0) & self.isTriggerSafeNoIso() if self.v > 0: print(" - custom ID and multi-isolation") if self.obj == "Electron" and self.wp == "tightFCNC": self.selection = self.selection & self.getElectronMVAID( ) & self.getFCNCIsolation(ev.Electron.jetRelIso, ev.Electron.jetPtRelv2, I_2, I_3) & (ev.Electron.miniPFRelIso_all < I_1) & self.isTriggerSafeNoIso() if self.v > 0: print(" - custom ID and multi-isolation") if self.obj == "Muon" and self.wp == "tight": self.selection = self.selection & self.getIsolation( 0.11, 0.74, 6.8) if self.v > 0: print(" - custom multi-isolation") #self.selection = self.selection & ak.fill_none(ev.Muon.matched_jet.btagDeepFlavB<0.2770, True) #self.selection = self.selection & (ev.Muon.matched_jet.btagDeepFlavB<0.2770) #if self.v>0: print (" - deepJet") if self.obj == "Muon" and self.wp == "tightFCNC": self.selection = self.selection & self.getFCNCIsolation( ev.Muon.jetRelIso, ev.Muon.jetPtRelv2, I_2, I_3) & (ev.Muon.miniPFRelIso_all < I_1) if self.v > 0: print(" - custom multi-isolation") if self.obj == "Electron" and (self.wp == "tightTTH" or self.wp == 'fakeableTTH' or self.wp == "tightSSTTH" or self.wp == 'fakeableSSTTH'): self.selection = self.selection & self.getSigmaIEtaIEta if self.v > 0: print(" - SigmaIEtaIEta") #self.selection = self.selection & ak.fill_none(ev.Electron.matched_jet.btagDeepFlavB<0.2770, True) #self.selection = self.selection & (ev.Electron.matched_jet.btagDeepFlavB<0.2770) #self.selection = self.selection & (ev.Jet[ev.Electron.jetIdx].btagDeepFlavB<0.2770) #if self.v>0: print (" - deepJet") if self.obj == "Electron" and self.wp == "looseFCNC": self.selection = self.selection & (ev.Electron.miniPFRelIso_all < 0.4) if self.obj == 'Muon' and (self.wp == 'fakeableTTH' or self.wp == 'fakeableSSTTH'): #self.selection = self.selection & (self.cand.deepJet < self.getThreshold(self.cand.conePt, min_pt=20, max_pt=45, low=0.2770, high=0.0494)) self.selection = self.selection & (ak.fill_none( ev.Muon.matched_jet.btagDeepFlavB, 0) < self.getThreshold( self.cand.conePt, min_pt=20, max_pt=45)) if self.v > 0: print(" - interpolated deepJet") if self.obj == "Muon" and self.wp == "looseFCNC": self.selection = self.selection & (ev.Muon.miniPFRelIso_all < 0.4)
def process(self, events): output = self.accumulator.identity() output['total']['all'] += len(events) # use a very loose preselection to filter the events presel = ak.num(events.Jet) > 2 ev = events[presel] dataset = ev.metadata['dataset'] # load the config - probably not needed anymore cfg = loadConfig() ## Muons muon = Collections(ev, "Muon", "vetoTTH").get() tightmuon = Collections(ev, "Muon", "tightTTH").get() dimuon = choose(muon, 2) SSmuon = ak.any((dimuon['0'].charge * dimuon['1'].charge) > 0, axis=1) leading_muon_idx = ak.singletons(ak.argmax(muon.pt, axis=1)) leading_muon = muon[leading_muon_idx] ## Electrons electron = Collections(ev, "Electron", "vetoTTH").get() tightelectron = Collections(ev, "Electron", "tightTTH").get() dielectron = choose(electron, 2) SSelectron = ak.any( (dielectron['0'].charge * dielectron['1'].charge) > 0, axis=1) leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1)) leading_electron = electron[leading_electron_idx] ## Merge electrons and muons - this should work better now in ak1 dilepton = cross(muon, electron) SSlepton = ak.any((dilepton['0'].charge * dilepton['1'].charge) > 0, axis=1) lepton = ak.concatenate([muon, electron], axis=1) leading_lepton_idx = ak.singletons(ak.argmax(lepton.pt, axis=1)) leading_lepton = lepton[leading_lepton_idx] trailing_lepton_idx = ak.singletons(ak.argmin(lepton.pt, axis=1)) trailing_lepton = lepton[trailing_lepton_idx] dilepton_mass = (leading_lepton + trailing_lepton).mass dilepton_pt = (leading_lepton + trailing_lepton).pt dilepton_dR = delta_r(leading_lepton, trailing_lepton) mt_lep_met = mt(lepton.pt, lepton.phi, ev.MET.pt, ev.MET.phi) min_mt_lep_met = ak.min(mt_lep_met, axis=1) ## Jets jet = getJets(ev, minPt=25, maxEta=4.7, pt_var='pt_nom') jet = jet[ak.argsort( jet.pt_nom, ascending=False )] # need to sort wrt smeared and recorrected jet pt jet = jet[~match(jet, muon, deltaRCut=0.4)] # remove jets that overlap with muons jet = jet[~match( jet, electron, deltaRCut=0.4)] # remove jets that overlap with electrons central = jet[(abs(jet.eta) < 2.4)] btag = getBTagsDeepFlavB( jet, year=self.year) # should study working point for DeepJet light = getBTagsDeepFlavB(jet, year=self.year, invert=True) fwd = getFwdJet(light) fwd_noPU = getFwdJet(light, puId=False) tau = getTaus(ev) track = getIsoTracks(ev) ## forward jets j_fwd = fwd[ak.singletons(ak.argmax( fwd.p, axis=1))] # highest momentum spectator high_score_btag = central[ak.argsort(central.btagDeepFlavB)][:, :2] bl = cross(lepton, high_score_btag) bl_dR = delta_r(bl['0'], bl['1']) min_bl_dR = ak.min(bl_dR, axis=1) jf = cross(j_fwd, jet) mjf = (jf['0'] + jf['1']).mass j_fwd2 = jf[ak.singletons( ak.argmax(mjf, axis=1) )]['1'] # this is the jet that forms the largest invariant mass with j_fwd delta_eta = ak.fill_none( ak.pad_none(abs(j_fwd2.eta - j_fwd.eta), 1, clip=True), 0) ## MET -> can switch to puppi MET met_pt = ev.MET.pt met_phi = ev.MET.phi ## other variables ht = ak.sum(jet.pt, axis=1) st = met_pt + ht + ak.sum(muon.pt, axis=1) + ak.sum(electron.pt, axis=1) ## event selectors filters = getFilters(ev, year=self.year, dataset=dataset) dilep = ((ak.num(tightelectron) + ak.num(tightmuon)) == 2) lep0pt = ((ak.num(electron[(electron.pt > 25)]) + ak.num(muon[(muon.pt > 25)])) > 0) lep1pt = ((ak.num(electron[(electron.pt > 20)]) + ak.num(muon[(muon.pt > 20)])) > 1) lepveto = ((ak.num(electron) + ak.num(muon)) == 2) selection = PackedSelection() selection.add('lepveto', lepveto) selection.add('dilep', dilep) selection.add('filter', (filters)) selection.add('p_T(lep0)>25', lep0pt) selection.add('p_T(lep1)>20', lep1pt) selection.add('SS', (SSlepton | SSelectron | SSmuon)) selection.add('N_jet>3', (ak.num(jet) >= 4)) selection.add('N_central>2', (ak.num(central) >= 3)) selection.add('N_btag>0', (ak.num(btag) >= 1)) selection.add('N_fwd>0', (ak.num(fwd) >= 1)) #ss_reqs = ['lepveto', 'dilep', 'filter', 'p_T(lep0)>25', 'p_T(lep1)>20', 'SS'] ss_reqs = [ 'lepveto', 'dilep', 'filter', 'p_T(lep0)>25', 'p_T(lep1)>20', 'SS' ] #bl_reqs = ss_reqs + ['N_jet>3', 'N_central>2', 'N_btag>0', 'N_fwd>0'] bl_reqs = ss_reqs + ['N_jet>3', 'N_central>2', 'N_btag>0'] ss_reqs_d = {sel: True for sel in ss_reqs} ss_selection = selection.require(**ss_reqs_d) bl_reqs_d = {sel: True for sel in bl_reqs} BL = selection.require(**bl_reqs_d) weight = Weights(len(ev)) if not dataset == 'MuonEG': # lumi weight weight.add("weight", ev.weight) # PU weight - not in the babies... weight.add("PU", ev.puWeight, weightUp=ev.puWeightUp, weightDown=ev.puWeightDown, shift=False) # b-tag SFs weight.add("btag", self.btagSF.Method1a(btag, light)) # lepton SFs weight.add("lepton", self.leptonSF.get(electron, muon)) #cutflow = Cutflow(output, ev, weight=weight) #cutflow_reqs_d = {} #for req in bl_reqs: # cutflow_reqs_d.update({req: True}) # cutflow.addRow( req, selection.require(**cutflow_reqs_d) ) labels = { 'topW_v3': 0, 'TTW': 1, 'TTZ': 2, 'TTH': 3, 'ttbar': 4, 'ttbar1l_MG': 4 } if dataset in labels: label_mult = labels[dataset] else: label_mult = 5 label = np.ones(len(ev[BL])) * label_mult output["n_lep"] += processor.column_accumulator( ak.to_numpy((ak.num(electron) + ak.num(muon))[BL])) output["n_lep_tight"] += processor.column_accumulator( ak.to_numpy((ak.num(tightelectron) + ak.num(tightmuon))[BL])) output["lead_lep_pt"] += processor.column_accumulator( ak.to_numpy(ak.flatten(leading_lepton[BL].pt, axis=1))) output["lead_lep_eta"] += processor.column_accumulator( ak.to_numpy(ak.flatten(leading_lepton[BL].eta, axis=1))) output["lead_lep_phi"] += processor.column_accumulator( ak.to_numpy(ak.flatten(leading_lepton[BL].phi, axis=1))) output["lead_lep_charge"] += processor.column_accumulator( ak.to_numpy(ak.flatten(leading_lepton[BL].charge, axis=1))) output["sublead_lep_pt"] += processor.column_accumulator( ak.to_numpy(ak.flatten(trailing_lepton[BL].pt, axis=1))) output["sublead_lep_eta"] += processor.column_accumulator( ak.to_numpy(ak.flatten(trailing_lepton[BL].eta, axis=1))) output["sublead_lep_phi"] += processor.column_accumulator( ak.to_numpy(ak.flatten(trailing_lepton[BL].phi, axis=1))) output["sublead_lep_charge"] += processor.column_accumulator( ak.to_numpy(ak.flatten(trailing_lepton[BL].charge, axis=1))) output["lead_jet_pt"] += processor.column_accumulator( ak.to_numpy(ak.flatten(jet[:, 0:1][BL].pt, axis=1))) output["lead_jet_eta"] += processor.column_accumulator( ak.to_numpy(ak.flatten(jet[:, 0:1][BL].eta, axis=1))) output["lead_jet_phi"] += processor.column_accumulator( ak.to_numpy(ak.flatten(jet[:, 0:1][BL].phi, axis=1))) output["sublead_jet_pt"] += processor.column_accumulator( ak.to_numpy(ak.flatten(jet[:, 1:2][BL].pt, axis=1))) output["sublead_jet_eta"] += processor.column_accumulator( ak.to_numpy(ak.flatten(jet[:, 1:2][BL].eta, axis=1))) output["sublead_jet_phi"] += processor.column_accumulator( ak.to_numpy(ak.flatten(jet[:, 1:2][BL].phi, axis=1))) output["lead_btag_pt"] += processor.column_accumulator( ak.to_numpy(ak.flatten(high_score_btag[:, 0:1][BL].pt, axis=1))) output["lead_btag_eta"] += processor.column_accumulator( ak.to_numpy(ak.flatten(high_score_btag[:, 0:1][BL].eta, axis=1))) output["lead_btag_phi"] += processor.column_accumulator( ak.to_numpy(ak.flatten(high_score_btag[:, 0:1][BL].phi, axis=1))) output["sublead_btag_pt"] += processor.column_accumulator( ak.to_numpy(ak.flatten(high_score_btag[:, 1:2][BL].pt, axis=1))) output["sublead_btag_eta"] += processor.column_accumulator( ak.to_numpy(ak.flatten(high_score_btag[:, 1:2][BL].eta, axis=1))) output["sublead_btag_phi"] += processor.column_accumulator( ak.to_numpy(ak.flatten(high_score_btag[:, 1:2][BL].phi, axis=1))) output["fwd_jet_p"] += processor.column_accumulator( ak.to_numpy( ak.flatten(ak.fill_none(ak.pad_none(j_fwd[BL].p, 1, clip=True), 0), axis=1))) output["fwd_jet_pt"] += processor.column_accumulator( ak.to_numpy( ak.flatten(ak.fill_none( ak.pad_none(j_fwd[BL].pt, 1, clip=True), 0), axis=1))) output["fwd_jet_eta"] += processor.column_accumulator( ak.to_numpy( ak.flatten(ak.fill_none( ak.pad_none(j_fwd[BL].eta, 1, clip=True), 0), axis=1))) output["fwd_jet_phi"] += processor.column_accumulator( ak.to_numpy( ak.flatten(ak.fill_none( ak.pad_none(j_fwd[BL].phi, 1, clip=True), 0), axis=1))) output["mjj_max"] += processor.column_accumulator( ak.to_numpy(ak.fill_none(ak.max(mjf[BL], axis=1), 0))) output["delta_eta_jj"] += processor.column_accumulator( ak.to_numpy(ak.flatten(delta_eta[BL], axis=1))) output["met"] += processor.column_accumulator(ak.to_numpy(met_pt[BL])) output["ht"] += processor.column_accumulator(ak.to_numpy(ht[BL])) output["st"] += processor.column_accumulator(ak.to_numpy(st[BL])) output["n_jet"] += processor.column_accumulator( ak.to_numpy(ak.num(jet[BL]))) output["n_btag"] += processor.column_accumulator( ak.to_numpy(ak.num(btag[BL]))) output["n_fwd"] += processor.column_accumulator( ak.to_numpy(ak.num(fwd[BL]))) output["n_central"] += processor.column_accumulator( ak.to_numpy(ak.num(central[BL]))) output["n_tau"] += processor.column_accumulator( ak.to_numpy(ak.num(tau[BL]))) output["n_track"] += processor.column_accumulator( ak.to_numpy(ak.num(track[BL]))) output["dilepton_pt"] += processor.column_accumulator( ak.to_numpy(ak.flatten(dilepton_pt[BL], axis=1))) output["dilepton_mass"] += processor.column_accumulator( ak.to_numpy(ak.flatten(dilepton_mass[BL], axis=1))) output["min_bl_dR"] += processor.column_accumulator( ak.to_numpy(min_bl_dR[BL])) output["min_mt_lep_met"] += processor.column_accumulator( ak.to_numpy(min_mt_lep_met[BL])) output["label"] += processor.column_accumulator(label) output["weight"] += processor.column_accumulator(weight.weight()[BL]) output["presel"]["all"] += len(ev[ss_selection]) output["sel"]["all"] += len(ev[BL]) return output
def test_rochester(): rochester_data = lookup_tools.txt_converters.convert_rochester_file( "tests/samples/RoccoR2018.txt.gz", loaduncs=True) rochester = lookup_tools.rochester_lookup.rochester_lookup(rochester_data) # to test 1-to-1 agreement with official Rochester requires loading C++ files # instead, preload the correct scales in the sample directory # the script tests/samples/rochester/build_rochester.py produces these official_data_k = np.load("tests/samples/nano_dimuon_rochester.npy") official_data_err = np.load("tests/samples/nano_dimuon_rochester_err.npy") official_mc_k = np.load("tests/samples/nano_dy_rochester.npy") official_mc_err = np.load("tests/samples/nano_dy_rochester_err.npy") mc_rand = np.load("tests/samples/nano_dy_rochester_rand.npy") # test against nanoaod events = NanoEventsFactory.from_root( os.path.abspath("tests/samples/nano_dimuon.root")).events() data_k = rochester.kScaleDT(events.Muon.charge, events.Muon.pt, events.Muon.eta, events.Muon.phi) data_k = np.array(ak.flatten(data_k)) assert all(np.isclose(data_k, official_data_k)) data_err = rochester.kScaleDTerror(events.Muon.charge, events.Muon.pt, events.Muon.eta, events.Muon.phi) data_err = np.array(ak.flatten(data_err), dtype=float) assert all(np.isclose(data_err, official_data_err, atol=1e-8)) # test against mc events = NanoEventsFactory.from_root( os.path.abspath("tests/samples/nano_dy.root")).events() hasgen = ~np.isnan(ak.fill_none(events.Muon.matched_gen.pt, np.nan)) mc_rand = ak.unflatten(mc_rand, ak.num(hasgen)) mc_kspread = rochester.kSpreadMC( events.Muon.charge[hasgen], events.Muon.pt[hasgen], events.Muon.eta[hasgen], events.Muon.phi[hasgen], events.Muon.matched_gen.pt[hasgen], ) mc_ksmear = rochester.kSmearMC( events.Muon.charge[~hasgen], events.Muon.pt[~hasgen], events.Muon.eta[~hasgen], events.Muon.phi[~hasgen], events.Muon.nTrackerLayers[~hasgen], mc_rand[~hasgen], ) mc_k = np.array(ak.flatten(ak.ones_like(events.Muon.pt))) hasgen_flat = np.array(ak.flatten(hasgen)) mc_k[hasgen_flat] = np.array(ak.flatten(mc_kspread)) mc_k[~hasgen_flat] = np.array(ak.flatten(mc_ksmear)) assert all(np.isclose(mc_k, official_mc_k)) mc_errspread = rochester.kSpreadMCerror( events.Muon.charge[hasgen], events.Muon.pt[hasgen], events.Muon.eta[hasgen], events.Muon.phi[hasgen], events.Muon.matched_gen.pt[hasgen], ) mc_errsmear = rochester.kSmearMCerror( events.Muon.charge[~hasgen], events.Muon.pt[~hasgen], events.Muon.eta[~hasgen], events.Muon.phi[~hasgen], events.Muon.nTrackerLayers[~hasgen], mc_rand[~hasgen], ) mc_err = np.array(ak.flatten(ak.ones_like(events.Muon.pt))) mc_err[hasgen_flat] = np.array(ak.flatten(mc_errspread)) mc_err[~hasgen_flat] = np.array(ak.flatten(mc_errsmear)) assert all(np.isclose(mc_err, official_mc_err, atol=1e-8))
def to_np_array(ak_array, maxN=100, pad=0): return ak.fill_none(ak.pad_none(ak_array, maxN, clip=True, axis=-1), pad).to_numpy()
def test(): array = ak.values_astype(ak.Array([1.1, 2.2, None, 3.3]), np.float32) assert str(ak.fill_none(array, np.float32(0)).type) == "4 * float32"
def pad_awkward_array(array, pad_length, pad_value): return awkward.fill_none(awkward.pad_none(array, pad_length, clip=True), pad_value)
def test(): a = ak.values_astype(ak.Array([1, None]), np.float32) assert ak.fill_none(a, np.float32(0)).tolist() == [1, 0] assert str(ak.fill_none(a, np.float32(0)).type) == "2 * float32" assert ak.fill_none(a, np.array(0, np.float32)).tolist() == [1, 0] assert str(ak.fill_none(a, np.array(0, np.float32)).type) == "2 * float32" assert ak.fill_none(a, np.array([0], np.float32)).tolist() == [1, [0]] assert (str(ak.fill_none(a, np.array( [0], np.float32)).type) == "2 * union[float32, 1 * float32]") assert ak.fill_none(a, np.array([[0]], np.float32)).tolist() == [1, [[0]]] assert (str(ak.fill_none(a, np.array( [[0]], np.float32)).type) == "2 * union[float32, 1 * 1 * float32]") assert ak.fill_none(a, 0).tolist() == [1, 0] assert str(ak.fill_none(a, 0).type) == "2 * float64" assert ak.fill_none(a, [0]).tolist() == [1, [0]] assert str(ak.fill_none(a, [0]).type) == "2 * union[float32, 1 * int64]" assert ak.fill_none(a, [[0]]).tolist() == [1, [[0]]] assert str(ak.fill_none( a, [[0]]).type) == "2 * union[float32, 1 * var * int64]"
def num(ar): return ak.num(ak.fill_none(ar[~ak.is_none(ar)], 0), axis=0)