def _kExtra(self, kpt, eta, nl, u, s=0, m=0): # if it is a jagged array, save the offsets then flatten everything # needed for the ternary conditions later abseta = abs(eta) kData = self._kRes[s][m][1](abseta) # type 1 is data kMC = self._kRes[s][m][0](abseta) # type 0 is MC mask = kData > kMC x = awkward.zeros_like(kpt) sigma = self._sigma(kpt, eta, nl, s, m) # Rochester cbA = beta, cbN = m, as well as cbM (always 0?) = loc and cbS = scale to transform y = (x-loc)/scale in the pdf method cbA = self._cbA[s][m](abseta, nl) cbN = self._cbN[s][m](abseta, nl) cbS = self._cbS[s][m](abseta, nl) counts = awkward.num(u) u_flat = awkward.flatten(u) loc = awkward.zeros_like(u_flat) cbA_flat = awkward.flatten(cbA) cbN_flat = awkward.flatten(cbN) cbS_flat = awkward.flatten(cbS) invcdf = awkward.unflatten( doublecrystalball.ppf(u_flat, cbA_flat, cbA_flat, cbN_flat, cbN_flat, loc, cbS_flat), counts, ) x = awkward.where( mask, (numpy.sqrt(kData * kData - kMC * kMC) * sigma * invcdf), x, ) result = awkward.where(x > -1, 1.0 / (1.0 + x), awkward.ones_like(kpt)) if isinstance(kpt, numpy.ndarray): result = numpy.array(result) return result
def _assignTruthDef(self, tree): assert self.splitIdx is not None nonSplitRecHitSimClusIdx = self._createTruthAssociation(tree) recHitTruthPID = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_pdgId",nonSplitRecHitSimClusIdx) recHitTruthEnergy = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_boundaryEnergy",nonSplitRecHitSimClusIdx) recHitDepEnergy = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_recEnergy",nonSplitRecHitSimClusIdx) if not self.use_true_muon_momentum: recHitTruthEnergy = ak1.where(np.abs(recHitTruthPID[:,:,0])==13, recHitDepEnergy, recHitTruthEnergy) recHitTruthX = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_x",nonSplitRecHitSimClusIdx) recHitTruthY = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_y",nonSplitRecHitSimClusIdx) recHitTruthZ = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_z",nonSplitRecHitSimClusIdx) recHitTruthTime = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_t",nonSplitRecHitSimClusIdx) fullyContained = ak1.where(np.abs(recHitTruthZ)[:,:,0]<323.,#somehow that seems necessary ak1.ones_like(recHitTruthZ), ak1.zeros_like(recHitTruthZ)) recHitEnergy = self._readSplitAndExpand(tree,"RecHitHGC_energy") recHitTime = self._readSplitAndExpand(tree,"RecHitHGC_time") recHitX = self._readSplitAndExpand(tree,"RecHitHGC_x") recHitY = self._readSplitAndExpand(tree,"RecHitHGC_y") recHitZ = self._readSplitAndExpand(tree,"RecHitHGC_z") # should not expand here to allow indexing as done below recHitSimClusIdx = self._splitJaggedArray(nonSplitRecHitSimClusIdx) # set noise to rec features recHitTruthEnergy = ak1.where(recHitSimClusIdx<0, recHitEnergy, recHitTruthEnergy) recHitTruthX = ak1.where(recHitSimClusIdx<0, recHitX, recHitTruthX) recHitTruthY = ak1.where(recHitSimClusIdx<0, recHitY, recHitTruthY) recHitTruthZ = ak1.where(recHitSimClusIdx<0, recHitZ, recHitTruthZ) recHitTruthTime = ak1.where(recHitSimClusIdx<0, recHitTime, recHitTruthTime) recHitDepEnergy = ak1.where(recHitSimClusIdx<0, recHitEnergy, recHitDepEnergy) recHitSpectatorFlag = self._createSpectators(tree) #remove spectator flag for noise recHitSpectatorFlag = ak1.where(recHitSimClusIdx<0 , ak1.zeros_like(recHitSpectatorFlag), recHitSpectatorFlag)#this doesn't work for some reason! #DEBUG!!! #ticlidx = self._readSplitAndExpand(tree,'RecHitHGC_TICLCandIdx') self.truth={} #DEBUG!!! self.truth['t_idx'] = self._expand(recHitSimClusIdx)# now expand to a trailing dimension self.truth['t_energy'] = recHitTruthEnergy self.truth['t_pos'] = ak1.concatenate([recHitTruthX, recHitTruthY,recHitTruthZ],axis=-1) self.truth['t_time'] = recHitTruthTime self.truth['t_pid'] = recHitTruthPID self.truth['t_spectator'] = recHitSpectatorFlag self.truth['t_fully_contained'] = fullyContained self.truth['t_rec_energy'] = recHitDepEnergy
def calculate_selection(self, syst_tag, events): """ """ electrons = events.ele electrons["label"] = -1 * awkward.ones_like(electrons.pt) if not self.is_data: electrons["label"] = awkward.where( electrons.genPartFlav == 1, awkward.ones_like(electrons.label), electrons.label) electrons["label"] = awkward.where( (electrons.genPartFlav == 3) | (electrons.genPartFlav == 4) | (electrons.genPartFlav == 5), awkward.zeros_like(electrons.label), electrons.label) fields = [x for x in events.fields if x not in ["ele", "Electron"]] for x in fields: if x == "ZCand": electrons[x] = awkward.firsts(events[x]) else: electrons[x] = events[x] electrons = awkward.flatten(electrons) dummy_cut = electrons.pt >= 0 return dummy_cut, electrons
def _assignTruthByIndexAndSplit(self, tree, label, indices, null=0): sc = label if type(label) is str: sc = ak1.from_awkward0(tree[label].array()) vals = sc[indices] vals = ak1.where(indices<0, ak1.zeros_like(vals)+null, vals) ja = self._splitJaggedArray(vals) return self._expand(ja)
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 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 _assignTruth(self, tree): nonSplitTrackSimClusIdx = self._getMatchIdxs(tree) truthEnergy = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_boundaryEnergy",nonSplitTrackSimClusIdx) truthPID = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_pdgId",nonSplitTrackSimClusIdx) truthX = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_x",nonSplitTrackSimClusIdx) truthY = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_y",nonSplitTrackSimClusIdx) truthZ = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_z",nonSplitTrackSimClusIdx) truthTime = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_t",nonSplitTrackSimClusIdx) #some manual sets zeros = ak1.zeros_like(truthEnergy) splittruthidx = self._splitJaggedArray(nonSplitTrackSimClusIdx) spectator = ak1.where(splittruthidx<0, zeros+10., zeros) trackPt = self._readSplitAndExpand(tree,"Track_pt") trackVertEta = self._readSplitAndExpand(tree,"Track_eta") trackMom = trackPt * np.cosh(trackVertEta) impactX = self._readSplitAndExpand(tree,"Track_HGCFront_x") impactY = self._readSplitAndExpand(tree,"Track_HGCFront_y") impactZ = self._readSplitAndExpand(tree,"Track_HGCFront_z") truthX = ak1.where(splittruthidx<0, impactX, truthX) truthY = ak1.where(splittruthidx<0, impactY, truthY) truthZ = ak1.where(splittruthidx<0, impactZ, truthZ) truthEnergy = ak1.where(splittruthidx<0, trackMom, truthEnergy) truthidx = self._expand(splittruthidx) self.truth={} self.truth['t_idx'] = truthidx# for now self.truth['t_energy'] = truthEnergy self.truth['t_pos'] = ak1.concatenate([truthX,truthY,truthZ],axis=-1) self.truth['t_time'] = truthTime self.truth['t_pid'] = truthPID self.truth['t_spectator'] = spectator self.truth['t_fully_contained'] = zeros+1 self.truth['t_rec_energy'] = trackMom
def _readTree(self, tree): # no truth here! Only features self._readSplits(tree, splitlabel='RecHitHGC_z') recHitEnergy = self._readSplitAndExpand(tree,"RecHitHGC_energy") recHitTime = self._readSplitAndExpand(tree,"RecHitHGC_time") recHitX = self._readSplitAndExpand(tree,"RecHitHGC_x") recHitY = self._readSplitAndExpand(tree,"RecHitHGC_y") recHitZ = self._readSplitAndExpand(tree,"RecHitHGC_z") recHitHitR = self._readSplitAndExpand(tree,"RecHitHGC_hitr") recHitR = np.sqrt(recHitX*recHitX+recHitY*recHitY+recHitZ*recHitZ) recHitTheta = np.arccos(recHitZ/recHitR) recHitEta = -np.log(np.tan(recHitTheta/2)) zeros = ak1.zeros_like(recHitEta) self.features = ak1.concatenate([ recHitEnergy, recHitEta, zeros, #indicator if it is track or not recHitTheta, recHitR, recHitX, recHitY, recHitZ, recHitTime, recHitHitR ], axis=-1) #this is just for bookkeeping self.featurenames = [ 'recHitEnergy', 'recHitEta', 'isTrack', 'recHitTheta', 'recHitR', 'recHitX', 'recHitY', 'recHitZ', 'recHitTime', 'recHitHitR' ]
def jer_smear( variation, forceStochastic, pt_gen, jetPt, etaJet, jet_energy_resolution, jet_resolution_rand_gauss, jet_energy_resolution_scale_factor, ): pt_gen = pt_gen if not forceStochastic else None if not isinstance(jetPt, awkward.highlevel.Array): raise Exception("'jetPt' must be an awkward array of some kind!") if forceStochastic: pt_gen = awkward.without_parameters(awkward.zeros_like(jetPt)) jersmear = jet_energy_resolution * jet_resolution_rand_gauss jersf = jet_energy_resolution_scale_factor[:, variation] deltaPtRel = (jetPt - pt_gen) / jetPt doHybrid = (pt_gen > 0) & (numpy.abs(deltaPtRel) < 3 * jet_energy_resolution) detSmear = 1 + (jersf - 1) * deltaPtRel stochSmear = 1 + numpy.sqrt(numpy.maximum(jersf**2 - 1, 0)) * jersmear min_jet_pt = _MIN_JET_ENERGY / numpy.cosh(etaJet) min_jet_pt_corr = min_jet_pt / jetPt smearfact = awkward.where(doHybrid, detSmear, stochSmear) smearfact = awkward.where((smearfact * jetPt) < min_jet_pt, min_jet_pt_corr, smearfact) def getfunction(layout, depth): if isinstance(layout, awkward.layout.NumpyArray) or not isinstance( layout, (awkward.layout.Content, awkward.partition.PartitionedArray)): return lambda: awkward.layout.NumpyArray(smearfact) return None smearfact = awkward._util.recursively_apply( awkward.operations.convert.to_layout(jetPt), getfunction) smearfact = awkward._util.wrap(smearfact, awkward._util.behaviorof(jetPt)) return smearfact
def _readTree(self, tree): self._readSplits(tree, splitlabel='Track_HGCFront_z') trackPt = self._readSplitAndExpand(tree,"Track_pt") trackEta = self._readSplitAndExpand(tree,"Track_HGCFront_eta") trackVertEta = self._readSplitAndExpand(tree,"Track_eta") trackMom = trackPt * np.cosh(trackVertEta) impactX = self._readSplitAndExpand(tree,"Track_HGCFront_x") impactY = self._readSplitAndExpand(tree,"Track_HGCFront_y") impactZ = self._readSplitAndExpand(tree,"Track_HGCFront_z") chi2 = self._readSplitAndExpand(tree,"Track_normChiSq") impactR = np.sqrt(impactX**2+impactY**2+impactZ**2)+1e-3 impactTheta = np.arccos(impactZ/impactR) self.features = ak1.concatenate([ trackMom, trackEta, ak1.ones_like(trackMom), #indicator if it is track or not impactTheta, impactR, impactX, impactY, impactZ, ak1.zeros_like(trackMom),#no time info (yet,could be from MTD here) chi2 #this is radius for hits, here chi2 for tracks, since it's kinda realted to the impact points resolution... ], axis=-1) #this is just for bookkeeping, keep them the same as for hits self.featurenames = [ 'recHitEnergy', 'recHitEta', 'isTrack', 'recHitTheta', 'recHitR', 'recHitX', 'recHitY', 'recHitZ', 'recHitTime', 'recHitHitR' ]
def get_charge_parent(particle): parent = find_first_parent(particle) charge = ak.zeros_like(parent) one = [[-11, -13, -15, -17, 24, 37], 1] minus_one = [[11, 13, 15, 17, -24, -37], -1] two_thirds = [[2, 4, 6, 8], 2 / 3] minus_two_thirds = [[-2, -4, -6, -8], -2 / 3] minus_one_third = [[1, 3, 5, 7], -1 / 3] one_third = [[-1, -3, -5, -7], 1 / 3] zero = [[12, 14, 16, 18, 9, 21, 22, 23, 25], 0] charge_pairs = [ one, minus_one, two_thirds, minus_two_thirds, minus_one_third, zero ] for pair in charge_pairs: for ID in pair[0]: charge = (parent == ID) * ak.ones_like(parent) * pair[1] + ( ~(parent == ID)) * charge return charge
def test(): array = ak.Array( [ [{"x": 0.0, "y": []}, {"x": 1.1, "y": [1]}, {"x": 2.2, "y": [1, 2]}], [], [ {"x": 3.3, "y": [1, 2, None, 3]}, False, False, True, {"x": 4.4, "y": [1, 2, None, 3, 4]}, ], ] ) def assert_array_type(new_array, intended_type): """helper function to check each part of the array took the intended type""" assert isinstance(new_array[0][0]["x"], intended_type) assert isinstance(new_array[0][1]["x"], intended_type) assert isinstance(new_array[0][1]["y"][0], intended_type) assert isinstance(new_array[0][2]["x"], intended_type) assert isinstance(new_array[0][2]["y"][0], intended_type) assert isinstance(new_array[0][2]["y"][1], intended_type) assert isinstance(new_array[2][0]["x"], intended_type) assert isinstance(new_array[2][0]["y"][0], intended_type) assert isinstance(new_array[2][0]["y"][1], intended_type) assert isinstance(new_array[2][0]["y"][3], intended_type) assert isinstance(new_array[2][1], intended_type) assert isinstance(new_array[2][2], intended_type) assert isinstance(new_array[2][3], intended_type) assert isinstance(new_array[2][4]["x"], intended_type) assert isinstance(new_array[2][4]["y"][0], intended_type) assert isinstance(new_array[2][4]["y"][1], intended_type) assert isinstance(new_array[2][4]["y"][3], intended_type) assert isinstance(new_array[2][4]["y"][4], intended_type) int_type = ak.full_like(array, 12, dtype=int) float_type = ak.full_like(array, 12, dtype=float) assert int_type.tolist() == float_type.tolist() assert_array_type(int_type, int) assert_array_type(float_type, float) bool_type = ak.full_like(array, 12, dtype=bool) assert_array_type(bool_type, bool) int_type = ak.full_like(array, -1.2, dtype=int) float_type = ak.full_like(array, -1.2, dtype=float) bool_type = ak.full_like(array, -1.2, dtype=bool) assert_array_type(int_type, int) assert_array_type(float_type, float) assert_array_type(bool_type, bool) int_type = ak.zeros_like(array, dtype=int) float_type = ak.zeros_like(array, dtype=float) bool_type = ak.zeros_like(array, dtype=bool) assert int_type.tolist() == float_type.tolist() assert int_type.tolist() == bool_type.tolist() assert_array_type(int_type, int) assert_array_type(float_type, float) assert_array_type(bool_type, bool) int_type = ak.ones_like(array, dtype=int) float_type = ak.ones_like(array, dtype=float) bool_type = ak.ones_like(array, dtype=bool) assert int_type.tolist() == float_type.tolist() assert int_type.tolist() == bool_type.tolist() assert_array_type(int_type, int) assert_array_type(float_type, float) assert_array_type(bool_type, bool) array = ak.Array([["one", "two", "three"], [], ["four", "five"]]) def assert_array_type(new_array, intended_type): """helper function to check each part of the array took the intended type""" assert isinstance(new_array[0][0], intended_type) assert isinstance(new_array[0][1], intended_type) assert isinstance(new_array[0][2], intended_type) assert isinstance(new_array[2][0], intended_type) assert isinstance(new_array[2][1], intended_type) int_type = ak.full_like(array, 12, dtype=int) float_type = ak.full_like(array, 12, dtype=float) assert int_type.tolist() == float_type.tolist() assert_array_type(int_type, int) assert_array_type(float_type, float) bool_type = ak.full_like(array, 12, dtype=bool) assert_array_type(bool_type, bool) int_type = ak.full_like(array, -1.2, dtype=int) float_type = ak.full_like(array, -1.2, dtype=float) bool_type = ak.full_like(array, -1.2, dtype=bool) assert_array_type(int_type, int) assert_array_type(float_type, float) assert_array_type(bool_type, bool) int_type = ak.zeros_like(array, dtype=int) float_type = ak.zeros_like(array, dtype=float) bool_type = ak.zeros_like(array, dtype=bool) assert int_type.tolist() == float_type.tolist() assert int_type.tolist() == bool_type.tolist() assert_array_type(int_type, int) assert_array_type(float_type, float) assert_array_type(bool_type, bool) int_type = ak.ones_like(array, dtype=int) float_type = ak.ones_like(array, dtype=float) bool_type = ak.ones_like(array, dtype=bool) assert int_type.tolist() == float_type.tolist() assert int_type.tolist() == bool_type.tolist() assert_array_type(int_type, int) assert_array_type(float_type, float) assert_array_type(bool_type, bool)
def test(): array = ak.Array([ [{ "x": 0.0, "y": [] }, { "x": 1.1, "y": [1] }, { "x": 2.2, "y": [1, 2] }], [], [ { "x": 3.3, "y": [1, 2, None, 3] }, False, False, True, { "x": 4.4, "y": [1, 2, None, 3, 4] }, ], ]) assert ak.full_like(array, 12.3).tolist() == [ [{ "x": 12.3, "y": [] }, { "x": 12.3, "y": [12] }, { "x": 12.3, "y": [12, 12] }], [], [ { "x": 12.3, "y": [12, 12, None, 12] }, True, True, True, { "x": 12.3, "y": [12, 12, None, 12, 12] }, ], ] assert ak.zeros_like(array).tolist() == [ [{ "x": 0.0, "y": [] }, { "x": 0.0, "y": [0] }, { "x": 0.0, "y": [0, 0] }], [], [ { "x": 0.0, "y": [0, 0, None, 0] }, False, False, False, { "x": 0.0, "y": [0, 0, None, 0, 0] }, ], ] assert ak.ones_like(array).tolist() == [ [{ "x": 1.0, "y": [] }, { "x": 1.0, "y": [1] }, { "x": 1.0, "y": [1, 1] }], [], [ { "x": 1.0, "y": [1, 1, None, 1] }, True, True, True, { "x": 1.0, "y": [1, 1, None, 1, 1] }, ], ] array = ak.Array([["one", "two", "three"], [], ["four", "five"]]) assert ak.full_like(array, "hello").tolist() == [ ["hello", "hello", "hello"], [], ["hello", "hello"], ] assert ak.full_like(array, 1).tolist() == [["1", "1", "1"], [], ["1", "1"]] assert ak.full_like(array, 0).tolist() == [["0", "0", "0"], [], ["0", "0"]] assert ak.ones_like(array).tolist() == [["1", "1", "1"], [], ["1", "1"]] assert ak.zeros_like(array).tolist() == [["", "", ""], [], ["", ""]] array = ak.Array([[b"one", b"two", b"three"], [], [b"four", b"five"]]) assert ak.full_like(array, b"hello").tolist() == [ [b"hello", b"hello", b"hello"], [], [b"hello", b"hello"], ] assert ak.full_like(array, 1).tolist() == [[b"1", b"1", b"1"], [], [b"1", b"1"]] assert ak.full_like(array, 0).tolist() == [[b"0", b"0", b"0"], [], [b"0", b"0"]] assert ak.ones_like(array).tolist() == [[b"1", b"1", b"1"], [], [b"1", b"1"]] assert ak.zeros_like(array).tolist() == [[b"", b"", b""], [], [b"", b""]]
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 mass(self): return awkward.without_parameters(awkward.zeros_like(self.pt))
def select_normal(genparts, w_decay_momid): columns = ["pt", "eta", "phi", "mass", "pdgId", "charge", "decaytype"] # get tops, defined as last copy gen_tops = genparts[(genparts.hasFlags(['isLastCopy'])) & (genparts.pdgId == 6)] gen_tops['charge'] = ak.ones_like(gen_tops.pt) * (2. / 3.) gen_tbars = genparts[(genparts.hasFlags(['isLastCopy'])) & (genparts.pdgId == -6)] gen_tbars['charge'] = ak.ones_like(gen_tbars.pt) * (-2. / 3.) # get direct top decay products (children are "isHardProcess" and "isFirstCopy") gen_bs = ak.flatten(gen_tops.children[(gen_tops.children.pdgId == 5)], axis=2) gen_bs['charge'] = ak.ones_like(gen_bs.pt) * (-1. / 3.) gen_bbars = ak.flatten( gen_tbars.children[(gen_tbars.children.pdgId == -5)], axis=2) gen_bbars['charge'] = ak.ones_like(gen_bbars.pt) * (1. / 3.) gen_wplus = ak.flatten( gen_tops.children[(gen_tops.children.pdgId == w_decay_momid)], axis=2) gen_wplus['charge'] = ak.ones_like(gen_wplus.pt) gen_wminus = ak.flatten( gen_tbars.children[(gen_tbars.children.pdgId == -1 * w_decay_momid)], axis=2) gen_wminus['charge'] = ak.ones_like(gen_wminus.pt) * (-1.) # get w decay products # get last copy of W bosons last_gen_wplus = ak.flatten(gen_tops.distinctChildren[ (gen_tops.distinctChildren.pdgId == w_decay_momid) & (gen_tops.distinctChildren.hasFlags(['isLastCopy']))], axis=2) last_gen_wminus = ak.flatten(gen_tbars.distinctChildren[ (gen_tbars.distinctChildren.pdgId == -1 * w_decay_momid) & (gen_tbars.distinctChildren.hasFlags(['isLastCopy']))], axis=2) gen_wplus_decaytype = np.zeros(len(gen_wplus)) gen_wminus_decaytype = np.zeros(len(gen_wminus)) # up/down partons from last W+ gen_wpartons_up_fromWplus = ak.flatten( last_gen_wplus.children[(np.mod(last_gen_wplus.children.pdgId, 2) == 0) & (np.abs(last_gen_wplus.children.pdgId) < 6)], axis=2) gen_wpartons_up_fromWplus['charge'] = np.sign( gen_wpartons_up_fromWplus.pdgId) * (2. / 3.) gen_wplus_decaytype[ak.num(gen_wpartons_up_fromWplus) > 0] = np.ones( ak.sum(ak.num(gen_wpartons_up_fromWplus) > 0)) * 2 gen_wpartons_dw_fromWplus = ak.flatten( last_gen_wplus.children[(np.mod(last_gen_wplus.children.pdgId, 2) == 1) & (np.abs(last_gen_wplus.children.pdgId) < 6)], axis=2) gen_wpartons_dw_fromWplus['charge'] = np.sign( gen_wpartons_up_fromWplus.pdgId) * (-1. / 3.) # up/down partons from last W- gen_wpartons_up_fromWminus = ak.flatten(last_gen_wminus.children[ (np.mod(last_gen_wminus.children.pdgId, 2) == 0) & (np.abs(last_gen_wminus.children.pdgId) < 6)], axis=2) gen_wpartons_up_fromWminus['charge'] = np.sign( gen_wpartons_up_fromWminus.pdgId) * (2. / 3.) gen_wminus_decaytype[ak.num(gen_wpartons_up_fromWminus) > 0] = np.ones( ak.sum(ak.num(gen_wpartons_up_fromWminus) > 0)) * 2 gen_wpartons_dw_fromWminus = ak.flatten(last_gen_wminus.children[ (np.mod(last_gen_wminus.children.pdgId, 2) == 1) & (np.abs(last_gen_wminus.children.pdgId) < 6)], axis=2) gen_wpartons_dw_fromWminus['charge'] = np.sign( gen_wpartons_up_fromWminus.pdgId) * (-1. / 3.) # charged leps from last W+ gen_charged_leps_fromWplus = ak.flatten( last_gen_wplus.children[(np.abs(last_gen_wplus.children.pdgId) == 11) | (np.abs(last_gen_wplus.children.pdgId) == 13) | (np.abs(last_gen_wplus.children.pdgId) == 15)], axis=2) gen_charged_leps_fromWplus['charge'] = ak.ones_like( gen_charged_leps_fromWplus.pdgId) gen_wplus_decaytype[ak.num(gen_charged_leps_fromWplus) > 0] = np.ones( ak.sum(ak.num(gen_charged_leps_fromWplus) > 0)) # add decaytype (0 is INVALID, 1 is LEPTONIC, 2 is HADRONIC) gen_wplus['decaytype'] = ak.unflatten(gen_wplus_decaytype, ak.num(gen_wplus)) gen_tops['decaytype'] = gen_wplus['decaytype'] # set decaytype for tops gen_bs['decaytype'] = gen_wplus['decaytype'] # set decaytype for bs # neutral leps from last W+ gen_neutral_leps_fromWplus = ak.flatten( last_gen_wplus.children[(np.abs(last_gen_wplus.children.pdgId) == 12) | (np.abs(last_gen_wplus.children.pdgId) == 14) | (np.abs(last_gen_wplus.children.pdgId) == 16)], axis=2) gen_neutral_leps_fromWplus['charge'] = ak.zeros_like( gen_neutral_leps_fromWplus.pt) # charged leps from last W- gen_charged_leps_fromWminus = ak.flatten(last_gen_wminus.children[ (np.abs(last_gen_wminus.children.pdgId) == 11) | (np.abs(last_gen_wminus.children.pdgId) == 13) | (np.abs(last_gen_wminus.children.pdgId) == 15)], axis=2) gen_charged_leps_fromWminus['charge'] = ak.ones_like( gen_charged_leps_fromWminus.pdgId) * (-1.) gen_wminus_decaytype[ak.num(gen_charged_leps_fromWminus) > 0] = np.ones( ak.sum(ak.num(gen_charged_leps_fromWminus) > 0)) # add decaytype (0 is INVALID, 1 is LEPTONIC, 2 is HADRONIC) gen_wminus['decaytype'] = ak.unflatten(gen_wminus_decaytype, ak.num(gen_wminus)) gen_tbars['decaytype'] = gen_wminus['decaytype'] # set decaytype for tbars gen_bbars['decaytype'] = gen_wminus['decaytype'] # set decaytype for bbars # neutral leps from last W- gen_neutral_leps_fromWminus = ak.flatten(last_gen_wminus.children[ (np.abs(last_gen_wminus.children.pdgId) == 12) | (np.abs(last_gen_wminus.children.pdgId) == 14) | (np.abs(last_gen_wminus.children.pdgId) == 16)], axis=2) gen_neutral_leps_fromWminus['charge'] = ak.zeros_like( gen_neutral_leps_fromWminus.pt) gen_wpartons_up = ak.Array({}, with_name="PtEtaPhiMLorentzVector") gen_wpartons_dw = ak.Array({}, with_name="PtEtaPhiMLorentzVector") gen_charged_leps = ak.Array({}, with_name="PtEtaPhiMLorentzVector") gen_neutral_leps = ak.Array({}, with_name="PtEtaPhiMLorentzVector") for column in columns: if column == 'decaytype': continue gen_wpartons_up[column] = ak.flatten( ak.concatenate([ gen_wpartons_up_fromWplus[column], gen_wpartons_up_fromWminus[column] ], axis=1)) # (up-type partons from W+, W-) gen_wpartons_dw[column] = ak.flatten( ak.concatenate([ gen_wpartons_dw_fromWplus[column], gen_wpartons_dw_fromWminus[column] ], axis=1)) # (dw-type partons from W+, W-) gen_charged_leps[column] = ak.flatten( ak.concatenate([ gen_charged_leps_fromWplus[column], gen_charged_leps_fromWminus[column] ], axis=1)) # (charged leps from W+, W-) gen_neutral_leps[column] = ak.flatten( ak.concatenate([ gen_neutral_leps_fromWplus[column], gen_neutral_leps_fromWminus[column] ], axis=1)) # (neutral leps from W+, W-) gen_wpartons_up = ak.unflatten( gen_wpartons_up, ak.num(gen_wpartons_up_fromWplus) + ak.num(gen_wpartons_up_fromWminus)) gen_wpartons_dw = ak.unflatten( gen_wpartons_dw, ak.num(gen_wpartons_dw_fromWplus) + ak.num(gen_wpartons_dw_fromWminus)) gen_charged_leps = ak.unflatten( gen_charged_leps, ak.num(gen_charged_leps_fromWplus) + ak.num(gen_charged_leps_fromWminus)) gen_neutral_leps = ak.unflatten( gen_neutral_leps, ak.num(gen_neutral_leps_fromWplus) + ak.num(gen_neutral_leps_fromWminus)) gen_taus = gen_charged_leps[np.abs(gen_charged_leps.pdgId) == 15] gen_nu_taus = gen_neutral_leps[np.abs(gen_neutral_leps.pdgId) == 16] # fully hadronic evts had_evts = (ak.num(gen_charged_leps) == 0) & ( ak.num(gen_neutral_leps) == 0) & (ak.num(gen_wpartons_up) == 2) & (ak.num(gen_wpartons_dw) == 2) #set_trace() # get direct tau decay products from hard processes (subset of gen_taus events above) tau_decay_prods = genparts[genparts.hasFlags( ['isDirectHardProcessTauDecayProduct'])] # only need decays to leptons (e/mu, tau nu) for event classification tau_TO_tau_nu = tau_decay_prods[np.abs(tau_decay_prods.pdgId) == 16] tau_TO_charged_lep = tau_decay_prods[(np.abs(tau_decay_prods.pdgId) == 11) | (np.abs(tau_decay_prods.pdgId) == 13)] tau_TO_neutral_lep = tau_decay_prods[(np.abs(tau_decay_prods.pdgId) == 12) | (np.abs(tau_decay_prods.pdgId) == 14)] # set decaytype for gen taus charged_lep_decaytype_array = ak.to_numpy( ak.flatten(ak.zeros_like(gen_charged_leps['pt']))) # semilep evts semilep_evts = (ak.num(gen_charged_leps) == 1) & ( ak.num(gen_neutral_leps) == 1) & (ak.num(gen_wpartons_up) == 1) & (ak.num(gen_wpartons_dw) == 1) tau_jets_evts = semilep_evts & (ak.num(gen_taus) == 1) sl_evts_mask = np.repeat(ak.to_numpy(semilep_evts), ak.to_numpy(ak.num(gen_charged_leps))) semilep_decaytype_array = np.zeros(ak.to_numpy(semilep_evts).sum(), dtype=int) # tau -> l semilep_tau_leptonic_decay = (tau_jets_evts) & ( ak.num(tau_TO_charged_lep) == 1) & ( ak.num(tau_TO_neutral_lep) == 1) & (ak.num(tau_TO_tau_nu) == 1) semilep_tau_hadronic_decay = (tau_jets_evts) & ( ~semilep_tau_leptonic_decay) semilep_decaytype_array[ semilep_tau_leptonic_decay[semilep_evts]] = np.ones( ak.to_numpy(semilep_tau_leptonic_decay).sum(), dtype=int) semilep_decaytype_array[ semilep_tau_hadronic_decay[semilep_evts]] = np.ones( ak.to_numpy(semilep_tau_hadronic_decay).sum(), dtype=int) * 2 # set charged_lep_decatype_array for semileptonic events charged_lep_decaytype_array[sl_evts_mask] = semilep_decaytype_array # dilep evts dilep_evts = (ak.num(gen_charged_leps) == 2) & ( ak.num(gen_neutral_leps) == 2) & (ak.num(gen_wpartons_up) == 0) & (ak.num(gen_wpartons_dw) == 0) lep_lep_evts = dilep_evts & (ak.num(gen_taus) == 0) lep_tau_evts = dilep_evts & (ak.num(gen_taus) == 1) tau_tau_evts = dilep_evts & (ak.num(gen_taus) == 2) dl_evts_mask = np.repeat(ak.to_numpy(dilep_evts), ak.to_numpy(ak.num(gen_charged_leps))) dilep_decaytype_array = np.zeros((ak.to_numpy(dilep_evts).sum(), 2), dtype=int) # tau + tau # tau + tau -> ll dilep_TauTau_ll_decay = (tau_tau_evts) & ( ak.num(tau_TO_charged_lep) == 2) & ( ak.num(tau_TO_neutral_lep) == 2) & (ak.num(tau_TO_tau_nu) == 2) dilep_decaytype_array[dilep_TauTau_ll_decay[dilep_evts]] = np.ones( (ak.to_numpy(dilep_TauTau_ll_decay).sum(), 2), dtype=int) # tau + tau -> hh dilep_TauTau_hh_decay = (tau_tau_evts) & ( (ak.num(tau_TO_charged_lep) + ak.num(tau_TO_neutral_lep)) == 0) & (ak.num(tau_TO_tau_nu) == 2) dilep_decaytype_array[dilep_TauTau_hh_decay[dilep_evts]] = np.ones( (ak.to_numpy(dilep_TauTau_hh_decay).sum(), 2), dtype=int) * 2 # tau + tau -> lh dilep_TauTau_lh_decay = ( tau_tau_evts) & ~(dilep_TauTau_ll_decay | dilep_TauTau_hh_decay) # set index corresponding to leptonically decaying tau to 1, default array is set to 2 dl_TauTau_to_lh_decaytype_array = np.ones( ak.to_numpy(dilep_TauTau_lh_decay).sum() * 2, dtype=int) * 2 lep_tau_mask = (np.repeat( ak.to_numpy( ak.flatten( tau_TO_charged_lep[dilep_TauTau_lh_decay].parent.pdgId)), 2) == ak.flatten(gen_charged_leps[dilep_TauTau_lh_decay].pdgId)) dl_TauTau_to_lh_decaytype_array[lep_tau_mask] = np.ones(lep_tau_mask.sum(), dtype=int) dilep_decaytype_array[dilep_TauTau_lh_decay[ dilep_evts]] = dl_TauTau_to_lh_decaytype_array.reshape( ak.to_numpy(dilep_TauTau_lh_decay).sum(), 2) # lep + tau # tau -> l dilep_LepTau_l_decay = (lep_tau_evts) & ( ak.num(tau_TO_charged_lep) == 1) & ( ak.num(tau_TO_neutral_lep) == 1) & (ak.num(tau_TO_tau_nu) == 1) # set index corresponding to tau to 1 dl_LepTau_to_Lep_decaytype_array = np.zeros( ak.to_numpy(dilep_LepTau_l_decay).sum() * 2, dtype=int) dl_LepTau_to_Lep_decaytype_array[ak.flatten( np.abs(gen_charged_leps[dilep_LepTau_l_decay].pdgId) == 15)] = np.ones( ak.sum(dilep_LepTau_l_decay), dtype=int) dilep_decaytype_array[dilep_LepTau_l_decay[ dilep_evts]] = dl_LepTau_to_Lep_decaytype_array.reshape( ak.sum(dilep_LepTau_l_decay), 2) # tau -> h dilep_LepTau_h_decay = (lep_tau_evts) & ~(dilep_LepTau_l_decay) # set index corresponding to tau to 2 dl_LepTau_to_Had_decaytype_array = np.zeros(ak.sum(dilep_LepTau_h_decay) * 2, dtype=int) dl_LepTau_to_Had_decaytype_array[ak.flatten( np.abs(gen_charged_leps[dilep_LepTau_h_decay].pdgId) == 15)] = np.ones( ak.sum(dilep_LepTau_h_decay), dtype=int) * 2 dilep_decaytype_array[dilep_LepTau_h_decay[ dilep_evts]] = dl_LepTau_to_Had_decaytype_array.reshape( ak.sum(dilep_LepTau_h_decay), 2) # set charged_lep_decatype_array for dileptonic events charged_lep_decaytype_array[dl_evts_mask] = dilep_decaytype_array.flatten() # set charged lepton decaytype (defined for taus only, e/mu are 0) (1 is LEPTONIC, 2 is HADRONIC) gen_charged_leps['decaytype'] = ak.unflatten(charged_lep_decaytype_array, ak.num(gen_charged_leps)) # make awkward arrays of (top decay prods, tbar decay prods) Gen_Top_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector") Gen_B_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector") Gen_W_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector") Gen_Wparton_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector") for column in columns: Gen_Top_Pairs[column] = ak.flatten( ak.concatenate([gen_tops[column], gen_tbars[column]], axis=1)) # (top, tbar) Gen_B_Pairs[column] = ak.flatten( ak.concatenate([gen_bs[column], gen_bbars[column]], axis=1)) # (b, bbar) Gen_W_Pairs[column] = ak.flatten( ak.concatenate([gen_wplus[column], gen_wminus[column]], axis=1)) # (W+, W-) if column is not "decaytype": Gen_Wparton_Pairs[column] = ak.flatten( ak.concatenate( [ ak.pad_none(gen_wpartons_up[column], 1, axis=1), ak.pad_none(gen_wpartons_dw[column], 1, axis=1) ], axis=1)) # (up-type wpartons, down-type wpartons) Gen_Top_Pairs = ak.unflatten(Gen_Top_Pairs, ak.num(gen_tops) + ak.num(gen_tbars)) Gen_B_Pairs = ak.unflatten(Gen_B_Pairs, ak.num(gen_bs) + ak.num(gen_bbars)) Gen_W_Pairs = ak.unflatten(Gen_W_Pairs, ak.num(gen_wplus) + ak.num(gen_wminus)) Gen_Wparton_Pairs = ak.unflatten( Gen_Wparton_Pairs, ak.num(ak.pad_none(gen_wpartons_up, 1, axis=1)) + ak.num(ak.pad_none(gen_wpartons_dw, 1, axis=1))) Gen_Wparton_Pairs = Gen_Wparton_Pairs[ak.argsort( Gen_Wparton_Pairs["pt"], ascending=False)] # sort by pt Gen_TTbar = ak.Array( { "pt": (gen_tops + gen_tbars).pt, "eta": (gen_tops + gen_tbars).eta, "phi": (gen_tops + gen_tbars).phi, "mass": (gen_tops + gen_tbars).mass, "decaytype": gen_tops["decaytype"] + gen_tbars[ "decaytype"], # 0 is for INVALID, 2 for DILEP, 3 for SEMILEP, 4 for HADRONIC }, with_name="PtEtaPhiMLorentzVector") ## make "table" of gen objects for certain decays ## DILEP DILEP_evts = ak.zip({ "TTbar": Gen_TTbar[dilep_evts] if ak.any(dilep_evts) else ak.unflatten( Gen_TTbar[dilep_evts], ak.values_astype(dilep_evts, int)), "Top": Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][dilep_evts] if ak.any(dilep_evts) else ak.unflatten( Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][dilep_evts], ak.values_astype(dilep_evts, int)), "Tbar": Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][dilep_evts] if ak.any(dilep_evts) else ak.unflatten( Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][dilep_evts], ak.values_astype(dilep_evts, int)), "B": Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][dilep_evts] if ak.any(dilep_evts) else ak.unflatten( Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][dilep_evts], ak.values_astype(dilep_evts, int)), "Bbar": Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][dilep_evts] if ak.any(dilep_evts) else ak.unflatten( Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][dilep_evts], ak.values_astype(dilep_evts, int)), "Wplus": Gen_W_Pairs[Gen_W_Pairs.charge == 1][dilep_evts] if ak.any(dilep_evts) else ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == 1][dilep_evts], ak.values_astype(dilep_evts, int)), "Wminus": Gen_W_Pairs[Gen_W_Pairs.charge == -1][dilep_evts] if ak.any(dilep_evts) else ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == -1][dilep_evts], ak.values_astype(dilep_evts, int)), "First_plus": gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts] if ak.any(dilep_evts) else ak.unflatten( gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts], ak.values_astype(dilep_evts, int)), # charged lepton always made leading "Second_plus": gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts] if ak.any(dilep_evts) else ak.unflatten( gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts], ak.values_astype(dilep_evts, int)), # neutral lepton always made subleading "First_minus": gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts] if ak.any(dilep_evts) else ak.unflatten( gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts], ak.values_astype(dilep_evts, int)), # charged lepton always made leading "Second_minus": gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts] if ak.any(dilep_evts) else ak.unflatten( gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts], ak.values_astype(dilep_evts, int)), # neutral lepton always made subleading "Up_plus": gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts] if ak.any(dilep_evts) else ak.unflatten( gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts], ak.values_astype(dilep_evts, int)), # same as second plus "Down_plus": gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts] if ak.any(dilep_evts) else ak.unflatten( gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts], ak.values_astype(dilep_evts, int)), # same as first plus "Up_minus": gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts] if ak.any(dilep_evts) else ak.unflatten( gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts], ak.values_astype(dilep_evts, int)), # same as second minus "Down_minus": gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts] if ak.any(dilep_evts) else ak.unflatten( gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts], ak.values_astype(dilep_evts, int)), # same as first minus }) ## HAD HAD_evts = ak.zip({ "TTbar": Gen_TTbar[had_evts] if ak.any(had_evts) else ak.unflatten( Gen_TTbar[had_evts], ak.values_astype(had_evts, int)), "Top": Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][had_evts] if ak.any(had_evts) else ak.unflatten( Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][had_evts], ak.values_astype(had_evts, int)), "Tbar": Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][had_evts] if ak.any(had_evts) else ak.unflatten( Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][had_evts], ak.values_astype(had_evts, int)), "B": Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][had_evts] if ak.any(had_evts) else ak.unflatten( Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][had_evts], ak.values_astype(had_evts, int)), "Bbar": Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][had_evts] if ak.any(had_evts) else ak.unflatten( Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][had_evts], ak.values_astype(had_evts, int)), "Wplus": Gen_W_Pairs[Gen_W_Pairs.charge == 1][had_evts] if ak.any(had_evts) else ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == 1][had_evts], ak.values_astype(had_evts, int)), "Wminus": Gen_W_Pairs[Gen_W_Pairs.charge == -1][had_evts] if ak.any(had_evts) else ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == -1][had_evts], ak.values_astype(had_evts, int)), "First_plus": Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0][:, 0] if ak.any(had_evts) else ak.unflatten( Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0] [:, 0], ak.values_astype( had_evts, int)), # leading positively-charged parton "Second_plus": Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0][:, 1] if ak.any(had_evts) else ak.unflatten( Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0] [:, 1], ak.values_astype( had_evts, int)), # subleading positively-charged parton "First_minus": Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0][:, 0] if ak.any(had_evts) else ak.unflatten( Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0] [:, 0], ak.values_astype( had_evts, int)), # leading negatively-charged parton "Second_minus": Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0][:, 1] if ak.any(had_evts) else ak.unflatten( Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0] [:, 1], ak.values_astype( had_evts, int)), # subleading negatively-charged parton "Up_plus": gen_wpartons_up[gen_wpartons_up.charge > 0][had_evts] if ak.any(had_evts) else ak.unflatten( gen_wpartons_up[gen_wpartons_up.charge > 0][had_evts], ak.values_astype(had_evts, int)), # positively-charged up-type parton "Down_plus": gen_wpartons_dw[gen_wpartons_dw.charge > 0][had_evts] if ak.any(had_evts) else ak.unflatten( gen_wpartons_dw[gen_wpartons_dw.charge > 0][had_evts], ak.values_astype(had_evts, int)), # positively-charged down-type parton "Up_minus": gen_wpartons_up[gen_wpartons_up.charge < 0][had_evts] if ak.any(had_evts) else ak.unflatten( gen_wpartons_up[gen_wpartons_up.charge < 0][had_evts], ak.values_astype(had_evts, int)), # negatively-charged up-type parton "Down_minus": gen_wpartons_dw[gen_wpartons_dw.charge < 0][had_evts] if ak.any(had_evts) else ak.unflatten( gen_wpartons_dw[gen_wpartons_dw.charge < 0][had_evts], ak.values_astype(had_evts, int)), # negatively-charged down-type parton }) ## SEMILEP SEMILEP_evts = ak.zip({ "TTbar": Gen_TTbar[semilep_evts] if ak.any(semilep_evts) else ak.unflatten( Gen_TTbar[semilep_evts], ak.values_astype(semilep_evts, int)), "THad": Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 2][semilep_evts] if ak.any(semilep_evts) else ak.unflatten( Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 2][semilep_evts], ak.values_astype(semilep_evts, int)), "TLep": Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 1][semilep_evts] if ak.any(semilep_evts) else ak.unflatten( Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 1][semilep_evts], ak.values_astype(semilep_evts, int)), "BHad": Gen_B_Pairs[Gen_B_Pairs.decaytype == 2][semilep_evts] if ak.any(semilep_evts) else ak.unflatten( Gen_B_Pairs[Gen_B_Pairs.decaytype == 2][semilep_evts], ak.values_astype(semilep_evts, int)), "BLep": Gen_B_Pairs[Gen_B_Pairs.decaytype == 1][semilep_evts] if ak.any(semilep_evts) else ak.unflatten( Gen_B_Pairs[Gen_B_Pairs.decaytype == 1][semilep_evts], ak.values_astype(semilep_evts, int)), "WHad": Gen_W_Pairs[Gen_W_Pairs.decaytype == 2][semilep_evts] if ak.any(semilep_evts) else ak.unflatten( Gen_W_Pairs[Gen_W_Pairs.decaytype == 2][semilep_evts], ak.values_astype(semilep_evts, int)), "WLep": Gen_W_Pairs[Gen_W_Pairs.decaytype == 1][semilep_evts] if ak.any(semilep_evts) else ak.unflatten( Gen_W_Pairs[Gen_W_Pairs.decaytype == 1][semilep_evts], ak.values_astype(semilep_evts, int)), "Lepton": gen_charged_leps[semilep_evts] if ak.any(semilep_evts) else ak.unflatten(gen_charged_leps[semilep_evts], ak.values_astype(semilep_evts, int)), "Nu": gen_neutral_leps[semilep_evts] if ak.any(semilep_evts) else ak.unflatten(gen_neutral_leps[semilep_evts], ak.values_astype(semilep_evts, int)), "WJa": Gen_Wparton_Pairs[:, 0][semilep_evts] if ak.any(semilep_evts) else ak.unflatten(Gen_Wparton_Pairs[:, 0][semilep_evts], ak.values_astype(semilep_evts, int)), "WJb": Gen_Wparton_Pairs[:, 1][semilep_evts] if ak.any(semilep_evts) else ak.unflatten(Gen_Wparton_Pairs[:, 1][semilep_evts], ak.values_astype(semilep_evts, int)), "Up_Had": gen_wpartons_up[semilep_evts] if ak.any(semilep_evts) else ak.unflatten(gen_wpartons_up[semilep_evts], ak.values_astype(semilep_evts, int)), "Down_Had": gen_wpartons_dw[semilep_evts] if ak.any(semilep_evts) else ak.unflatten(gen_wpartons_dw[semilep_evts], ak.values_astype(semilep_evts, int)), }) # make dictionary to return GenObjects = dict({ "SL": SEMILEP_evts, "DL": DILEP_evts, "Had": HAD_evts, }) return GenObjects
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 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