def test_cartesian(): muon = ak.Array([[{"pt": 1.0}], []], with_name="muon") electron = ak.Array([[], [{"pt": 1.0}]], with_name="electron") muon = muon[muon.pt > 5] electron = electron[electron.pt > 5] leptons = ak.concatenate([muon, electron], axis=1) candidate = ak.firsts(leptons) assert ak.to_list(ak.Array(candidate)) == [None, None] result = ak.cartesian([candidate, candidate], axis=0) assert ak.to_list(result) == [ (None, None), (None, None), (None, None), (None, None), ] result = ak.cartesian([candidate, ak.Array([[1, 2, 3], []])], axis=1) assert ak.to_list(result) == [None, None] one, two = ak.broadcast_arrays(candidate, ak.Array([[1, 2, 3], []])) assert ak.to_list(one) == [None, None] assert ak.to_list(two) == [None, None]
def run_deltar_matching(store, target, drname='deltaR', radius=0.4, unique=False, sort=False): """ Running a delta R matching of some object collection "store" of dimension NxS with some target collection "target" of dimension NxT, The return object will have dimension NxSxT' where objects in the T' contain all "target" objects within the delta R radius. The delta R between the store and target object will be stored in the field `deltaR`. If the unique flag is turned on, then objects in the target collection will only be associated to the closest object. If the sort flag is turned on, then the target collection will be sorted according to the computed `deltaR`. """ _, target = ak.unzip(ak.cartesian([store.eta, target], nested=True)) target[drname] = delta_r(store, target) if unique: # Additional filtering t_index = ak.argmin(target[drname], axis=-2) s_index = ak.local_index(store.eta, axis=-1) _, t_index = ak.unzip(ak.cartesian([s_index, t_index], nested=True)) target = target[s_index == t_index] # Cutting on the computed delta R target = target[target[drname] < radius] # Sorting according to the computed delta R if sort: idx = ak.argsort(target[drname], axis=-1) target = target[idx] return target
def match(first, second, **kwargs): if 'dR_cutoff' not in kwargs: raise Exception("Please specify dR cutoff for matching!") dR_cutoff = kwargs.pop('dR_cutoff', 0.3) return_match_properties = kwargs.pop('return_match_properties', False) etas = ak.cartesian({ 'first': first.eta, 'second': second.eta }, axis=1, nested=True) phis = ak.cartesian({ 'first': first.phi, 'second': second.phi }, axis=1, nested=True) dR, deta, dphi = delta_r(etas['first'], etas['second'], phis['first'], phis['second']) min_idx = ak.argmin(dR, axis=2) match_properties = { 'pt': second.pt[min_idx], 'eta': second.eta[min_idx], 'phi': second.phi[min_idx], } if return_match_properties: return ak.any(dR < dR_cutoff, axis=2), match_properties else: return ak.any(dR < dR_cutoff, axis=2)
def association(cand1, cand2): ''' Function for association of the particles. The cuts that operates on all of them and computation of quantities can go here. individual cuts can go on the main processing''' #cut_dstar_back = ((cand2.deltamr > 0.143) & (cand2.deltamr < 0.148)) #cand2 = cand2[cut_dstar_back] asso = ak.cartesian([cand1, cand2]) asso = asso[asso.slot0.vtxIdx == asso.slot1.vtxIdx] asso = asso[ak.num(asso) > 0] cand1 = ak.zip({ 'pt': asso.slot0.pt, 'eta': asso.slot0.eta, 'phi': asso.slot0.phi, 'mass': asso.slot0.mass, 'charge': asso.slot0.charge}, with_name="PtEtaPhiMCandidate") cand2 = ak.zip({ 'pt': asso.slot1.pt, 'eta': asso.slot1.eta, 'phi': asso.slot1.phi, 'mass': asso.slot1.mass, 'charge': asso.slot1.charge}, with_name="PtEtaPhiMCandidate") asso['deltarap'] = asso.slot0.rap - asso.slot1.rap asso['cand'] = cand1 + cand2 return asso
def fill_gen_jets(df, output): gjets = df.GenJet gleptons = df.GenPart[(abs(df.GenPart.pdgId) == 13) | (abs(df.GenPart.pdgId) == 11) | (abs(df.GenPart.pdgId) == 15)] gl_pair = ak.cartesian({ "jet": gjets, "lepton": gleptons }, axis=1, nested=True) _, _, dr_gl = delta_r( gl_pair["jet"].eta, gl_pair["lepton"].eta, gl_pair["jet"].phi, gl_pair["lepton"].phi, ) isolated = ak.all((dr_gl > 0.3), axis=-1) gjet1 = ak.to_pandas(gjets[isolated]).loc[pd.IndexSlice[:, 0], ["pt", "eta", "phi", "mass"]] gjet2 = ak.to_pandas(gjets[isolated]).loc[pd.IndexSlice[:, 1], ["pt", "eta", "phi", "mass"]] gjet1.index = gjet1.index.droplevel("subentry") gjet2.index = gjet2.index.droplevel("subentry") gjsum = p4_sum(gjet1, gjet2) for var in ["pt", "eta", "phi", "mass"]: output[f"gjet1_{var}"] = gjet1[var] output[f"gjet2_{var}"] = gjet2[var] output[f"gjj_{var}"] = gjsum[var] output["gjj_dEta"], output["gjj_dPhi"], output["gjj_dR"] = delta_r( output.gjet1_eta, output.gjet2_eta, output.gjet1_phi, output.gjet2_phi) return output
def gen_jet_pair_mass(df): gjmass = None gjets = df.GenJet gleptons = df.GenPart[(abs(df.GenPart.pdgId) == 13) | (abs(df.GenPart.pdgId) == 11) | (abs(df.GenPart.pdgId) == 15)] gl_pair = ak.cartesian({ "jet": gjets, "lepton": gleptons }, axis=1, nested=True) _, _, dr_gl = delta_r( gl_pair["jet"].eta, gl_pair["lepton"].eta, gl_pair["jet"].phi, gl_pair["lepton"].phi, ) isolated = ak.all((dr_gl > 0.3), axis=-1) if ak.count(gjets[isolated], axis=None) > 0: # TODO: convert only relevant fields! gjet1 = ak.to_pandas(gjets[isolated]).loc[pd.IndexSlice[:, 0], ["pt", "eta", "phi", "mass"]] gjet2 = ak.to_pandas(gjets[isolated]).loc[pd.IndexSlice[:, 1], ["pt", "eta", "phi", "mass"]] gjet1.index = gjet1.index.droplevel("subentry") gjet2.index = gjet2.index.droplevel("subentry") gjsum = p4_sum(gjet1, gjet2) gjmass = gjsum.mass return gjmass
def test_broken(): ex = ak.Array([[1, 2, 3], [], [4, 5]]) p4 = ak.zip({"x": ex}) p4c = ak.cartesian({"a": p4, "b": p4}) df = ak.to_pandas(p4c) assert df["a", "x"].values.tolist() == [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 5, 5] assert df["b", "x"].values.tolist() == [1, 2, 3, 1, 2, 3, 1, 2, 3, 4, 5, 4, 5]
def metric_table(self, other, axis=1, metric=lambda a, b: a.delta_r(b), return_combinations=False): """Return a list of a metric evaluated between this object and another. The two arrays should be broadcast-compatible on all axes other than the specified axis, which will be used to form a cartesian product. If axis=None, broadcast arrays directly. The return shape will be that of ``self`` with a new axis with shape of ``other`` appended at the specified axis depths. Parameters ---------- other : awkward.Array Another array with same shape in all but ``axis`` axis : int, optional The axis to form the cartesian product (default 1). If None, the metric is directly evaluated on the input arrays (i.e. they should broadcast) metric : callable A function of two arguments, returning a scalar. The default metric is `delta_r`. return_combinations : bool If True return the combinations of inputs as well as an unzipped tuple """ if axis is None: a, b = self, other else: a, b = awkward.unzip( awkward.cartesian([self, other], axis=axis, nested=True)) mval = metric(a, b) if return_combinations: return mval, (a, b) return mval
def test(): simple = {"what": "ever"} one = ak.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]], behavior=simple) two = ak.Array([["one", "two"], ["three"], ["four", "five"]], behavior=simple) three = ak.cartesian({"one": one, "two": two}) assert three.behavior == {"what": "ever"}
def test(): nparray = ak.layout.NumpyArray(np.array([1, 2, 3], dtype=np.int64)) listarray = ak.layout.ListArray64( ak.layout.Index64(np.array([0], dtype=np.int64)), ak.layout.Index64(np.array([3], dtype=np.int64)), nparray, ) indexedarray = ak.layout.IndexedArray64(ak.layout.Index64([]), listarray) cart = ak.cartesian([indexedarray, indexedarray], nested=True) assert str(cart.type) == "0 * var * var * (int64, int64)" assert ak.to_list(cart) == [] cart = ak.cartesian([indexedarray, indexedarray], nested=False) assert str(cart.type) == "0 * var * (int64, int64)" assert ak.to_list(cart) == []
def test_empty_arrays_cartesian(): one = ak.Array([]) two = one = ak.Array([]) with pytest.raises(ValueError) as err: ak.to_list(ak.cartesian([one, two])) assert isinstance(err.value, ValueError) ak.to_list(ak.concatenate([one, two], axis=0))
def test(): one = ak.Array([1, 2, 3, 4]) two = ak.Array(["aa", "bb", "cc", "dd"]) with pytest.raises(ValueError): ak.cartesian([one, two], axis=1) with pytest.raises(ValueError): ak.combinations(two, 2, axis=1) two = ak.Array([["aa", "bb"], ["cc"], [], ["dd"]]) assert ak.to_list(ak.cartesian([one, two], axis=1)) == [ [(1, "aa"), (1, "bb")], [(2, "cc")], [], [(4, "dd")], ] assert ak.to_list(ak.combinations(two, 2, axis=1)) == [[("aa", "bb")], [], [], []] with pytest.raises(ValueError): ak.cartesian([one, two], axis=2) with pytest.raises(ValueError): ak.combinations(two, 2, axis=2)
def L1THLTTauMatching(L1taus, taus): dR_matching = 0.5 tau_inpair, L1_inpair = ak.unzip(ak.cartesian([taus, L1taus], nested=True)) # dR = delta_r(L1_inpair, tau_inpair) dR = delta_r(tau_inpair, L1_inpair) # # print(dR[range(2)]) # # consider only L1taus for which there is at least 1 matched tau # tau_inpair = tau_inpair[dR<dR_matching] mask = ak.sum(dR < dR_matching, axis=-1) > 0 # tau_inpair = tau_inpair[mask] # # take first matched tau for each L1tau # L2taus = tau_inpair[:,:,0] # # print(L2taus[range(2)]) L2taus = taus[mask] return L2taus
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_axis2(): one = ak.Array([[[0, 1, 2], [], [3, 4]], [[0, 1, 2], [], [3, 4]]]) two = ak.Array([[[100, 200], [300], [400, 500]], [[100, 200], [300], [400, 500]]]) assert ak.to_list(ak.cartesian([one, two], axis=2)) == [ [ [(0, 100), (0, 200), (1, 100), (1, 200), (2, 100), (2, 200)], [], [(3, 400), (3, 500), (4, 400), (4, 500)], ], [ [(0, 100), (0, 200), (1, 100), (1, 200), (2, 100), (2, 200)], [], [(3, 400), (3, 500), (4, 400), (4, 500)], ], ]
def fill_gen_jets(df, output): features = ["PT", "Eta", "Phi", "Mass"] gjets = df.GenJet[features] print(df.GenJet.fields) gleptons = df.MuonMedium # gleptons = df.GenPart[ # (abs(df.GenPart.pdgId) == 13) # | (abs(df.GenPart.pdgId) == 11) # | (abs(df.GenPart.pdgId) == 15) # ] gl_pair = ak.cartesian({"jet": gjets, "lepton": gleptons}, axis=1, nested=True) _, _, dr_gl = delta_r( gl_pair["jet"].Eta, gl_pair["lepton"].Eta, gl_pair["jet"].Phi, gl_pair["lepton"].Phi, ) isolated = ak.all((dr_gl > 0.3), axis=-1) gjet1 = ak.to_pandas(gjets[isolated]).loc[pd.IndexSlice[:, 0], features] gjet2 = ak.to_pandas(gjets[isolated]).loc[pd.IndexSlice[:, 1], features] gjet1.index = gjet1.index.droplevel("subentry") gjet2.index = gjet2.index.droplevel("subentry") feat_map = {"pt": "PT", "eta": "Eta", "phi": "Phi", "mass": "Mass"} for var in ["pt", "eta", "phi", "mass"]: gjet1[var] = gjet1[feat_map[var]] gjet2[var] = gjet2[feat_map[var]] gjsum = p4_sum(gjet1, gjet2) for var in ["pt", "eta", "phi", "mass"]: output[f"gjet1_{var}"] = gjet1[var] output[f"gjet2_{var}"] = gjet2[var] output[f"gjj_{var}"] = gjsum[var] output["gjj_dEta"], output["gjj_dPhi"], output["gjj_dR"] = delta_r( output.gjet1_eta, output.gjet2_eta, output.gjet1_phi, output.gjet2_phi ) return output
def match2(first, second, deltaRCut=0.4): drCut2 = deltaRCut**2 combs = ak.cartesian([first, second], nested=True) return ak.any((combs['0'].delta_r2(combs['1']) < drCut2), axis=2)
def delta_r_v2(first, second): combs = ak.cartesian([first, second], nested=True) return np.sqrt(delta_r2(combs['0'], combs['1']))
def test_axis0(): one = ak.Array([1.1, 2.2, 3.3]) two = ak.Array([100, 200, 300, 400, 500]) three = ak.Array(["a", "b"]) assert ak.to_list(ak.cartesian([one], axis=0)) == [(1.1, ), (2.2, ), (3.3, )] assert ak.to_list(ak.cartesian({"x": one}, axis=0)) == [ { "x": 1.1 }, { "x": 2.2 }, { "x": 3.3 }, ] assert ak.to_list(ak.cartesian([one, two], axis=0)) == [ (1.1, 100), (1.1, 200), (1.1, 300), (1.1, 400), (1.1, 500), (2.2, 100), (2.2, 200), (2.2, 300), (2.2, 400), (2.2, 500), (3.3, 100), (3.3, 200), (3.3, 300), (3.3, 400), (3.3, 500), ] if not ak._util.py27 and not ak._util.py35: assert ak.to_list(ak.cartesian({ "x": one, "y": two }, axis=0)) == [ { "x": 1.1, "y": 100 }, { "x": 1.1, "y": 200 }, { "x": 1.1, "y": 300 }, { "x": 1.1, "y": 400 }, { "x": 1.1, "y": 500 }, { "x": 2.2, "y": 100 }, { "x": 2.2, "y": 200 }, { "x": 2.2, "y": 300 }, { "x": 2.2, "y": 400 }, { "x": 2.2, "y": 500 }, { "x": 3.3, "y": 100 }, { "x": 3.3, "y": 200 }, { "x": 3.3, "y": 300 }, { "x": 3.3, "y": 400 }, { "x": 3.3, "y": 500 }, ] assert ak.to_list(ak.cartesian([one, two, three], axis=0)) == [ (1.1, 100, "a"), (1.1, 100, "b"), (1.1, 200, "a"), (1.1, 200, "b"), (1.1, 300, "a"), (1.1, 300, "b"), (1.1, 400, "a"), (1.1, 400, "b"), (1.1, 500, "a"), (1.1, 500, "b"), (2.2, 100, "a"), (2.2, 100, "b"), (2.2, 200, "a"), (2.2, 200, "b"), (2.2, 300, "a"), (2.2, 300, "b"), (2.2, 400, "a"), (2.2, 400, "b"), (2.2, 500, "a"), (2.2, 500, "b"), (3.3, 100, "a"), (3.3, 100, "b"), (3.3, 200, "a"), (3.3, 200, "b"), (3.3, 300, "a"), (3.3, 300, "b"), (3.3, 400, "a"), (3.3, 400, "b"), (3.3, 500, "a"), (3.3, 500, "b"), ] assert ak.to_list(ak.cartesian([one, two, three], axis=0, nested=[0])) == [ [ (1.1, 100, "a"), (1.1, 100, "b"), (1.1, 200, "a"), (1.1, 200, "b"), (1.1, 300, "a"), ], [ (1.1, 300, "b"), (1.1, 400, "a"), (1.1, 400, "b"), (1.1, 500, "a"), (1.1, 500, "b"), ], [ (2.2, 100, "a"), (2.2, 100, "b"), (2.2, 200, "a"), (2.2, 200, "b"), (2.2, 300, "a"), ], [ (2.2, 300, "b"), (2.2, 400, "a"), (2.2, 400, "b"), (2.2, 500, "a"), (2.2, 500, "b"), ], [ (3.3, 100, "a"), (3.3, 100, "b"), (3.3, 200, "a"), (3.3, 200, "b"), (3.3, 300, "a"), ], [ (3.3, 300, "b"), (3.3, 400, "a"), (3.3, 400, "b"), (3.3, 500, "a"), (3.3, 500, "b"), ], ] assert ak.to_list(ak.cartesian([one, two, three], axis=0, nested=[1])) == [ [(1.1, 100, "a"), (1.1, 100, "b")], [(1.1, 200, "a"), (1.1, 200, "b")], [(1.1, 300, "a"), (1.1, 300, "b")], [(1.1, 400, "a"), (1.1, 400, "b")], [(1.1, 500, "a"), (1.1, 500, "b")], [(2.2, 100, "a"), (2.2, 100, "b")], [(2.2, 200, "a"), (2.2, 200, "b")], [(2.2, 300, "a"), (2.2, 300, "b")], [(2.2, 400, "a"), (2.2, 400, "b")], [(2.2, 500, "a"), (2.2, 500, "b")], [(3.3, 100, "a"), (3.3, 100, "b")], [(3.3, 200, "a"), (3.3, 200, "b")], [(3.3, 300, "a"), (3.3, 300, "b")], [(3.3, 400, "a"), (3.3, 400, "b")], [(3.3, 500, "a"), (3.3, 500, "b")], ] assert ak.to_list(ak.cartesian([one, two, three], axis=0, nested=[0, 1])) == [ [ [(1.1, 100, "a"), (1.1, 100, "b")], [(1.1, 200, "a"), (1.1, 200, "b")], [(1.1, 300, "a"), (1.1, 300, "b")], [(1.1, 400, "a"), (1.1, 400, "b")], [(1.1, 500, "a"), (1.1, 500, "b")], ], [ [(2.2, 100, "a"), (2.2, 100, "b")], [(2.2, 200, "a"), (2.2, 200, "b")], [(2.2, 300, "a"), (2.2, 300, "b")], [(2.2, 400, "a"), (2.2, 400, "b")], [(2.2, 500, "a"), (2.2, 500, "b")], ], [ [(3.3, 100, "a"), (3.3, 100, "b")], [(3.3, 200, "a"), (3.3, 200, "b")], [(3.3, 300, "a"), (3.3, 300, "b")], [(3.3, 400, "a"), (3.3, 400, "b")], [(3.3, 500, "a"), (3.3, 500, "b")], ], ] assert ak.to_list(ak.cartesian([one, two, three], axis=0, nested=[])) == ak.to_list( ak.cartesian([one, two, three], axis=0, nested=False)) assert ak.to_list(ak.cartesian([one, two, three], axis=0, nested=[])) == ak.to_list( ak.cartesian([one, two, three], axis=0, nested=None)) assert ak.to_list(ak.cartesian([one, two, three], axis=0, nested=[0, 1])) == ak.to_list( ak.cartesian([one, two, three], axis=0, nested=True))
def process(self, events): # Dataset parameters dataset = events.metadata['dataset'] histAxisName = self._samples[dataset]['histAxisName'] year = self._samples[dataset]['year'] xsec = self._samples[dataset]['xsec'] sow = self._samples[dataset]['nSumOfWeights'] isData = self._samples[dataset]['isData'] datasets = [ 'SingleMuon', 'SingleElectron', 'EGamma', 'MuonEG', 'DoubleMuon', 'DoubleElectron' ] for d in datasets: if d in dataset: dataset = dataset.split('_')[0] # Initialize objects met = events.MET e = events.Electron mu = events.Muon tau = events.Tau j = events.Jet # Muon selection mu['isPres'] = isPresMuon(mu.dxy, mu.dz, mu.sip3d, mu.looseId) mu['isTight'] = isTightMuon(mu.pt, mu.eta, mu.dxy, mu.dz, mu.pfRelIso03_all, mu.sip3d, mu.mvaTTH, mu.mediumPromptId, mu.tightCharge, mu.looseId, minpt=10) mu['isGood'] = mu['isPres'] & mu['isTight'] leading_mu = mu[ak.argmax(mu.pt, axis=-1, keepdims=True)] leading_mu = leading_mu[leading_mu.isGood] mu = mu[mu.isGood] mu_pres = mu[mu.isPres] # Electron selection e['isPres'] = isPresElec(e.pt, e.eta, e.dxy, e.dz, e.miniPFRelIso_all, e.sip3d, e.lostHits, minpt=15) e['isTight'] = isTightElec(e.pt, e.eta, e.dxy, e.dz, e.miniPFRelIso_all, e.sip3d, e.mvaTTH, e.mvaFall17V2Iso, e.lostHits, e.convVeto, e.tightCharge, e.sieie, e.hoe, e.eInvMinusPInv, minpt=15) e['isClean'] = isClean(e, mu, drmin=0.05) e['isGood'] = e['isPres'] & e['isTight'] & e['isClean'] leading_e = e[ak.argmax(e.pt, axis=-1, keepdims=True)] leading_e = leading_e[leading_e.isGood] e = e[e.isGood] e_pres = e[e.isPres & e.isClean] # Tau selection tau['isPres'] = isPresTau(tau.pt, tau.eta, tau.dxy, tau.dz, tau.leadTkPtOverTauPt, tau.idAntiMu, tau.idAntiEle, tau.rawIso, tau.idDecayModeNewDMs, minpt=20) tau['isClean'] = isClean(tau, e_pres, drmin=0.4) & isClean( tau, mu_pres, drmin=0.4) tau['isGood'] = tau['isPres'] # & tau['isClean'], for the moment tau = tau[tau.isGood] nElec = ak.num(e) nMuon = ak.num(mu) nTau = ak.num(tau) twoLeps = (nElec + nMuon) == 2 threeLeps = (nElec + nMuon) == 3 twoElec = (nElec == 2) twoMuon = (nMuon == 2) e0 = e[ak.argmax(e.pt, axis=-1, keepdims=True)] m0 = mu[ak.argmax(mu.pt, axis=-1, keepdims=True)] # Attach the lepton SFs to the electron and muons collections AttachElectronSF(e, year=year) AttachMuonSF(mu, year=year) # Create a lepton (muon+electron) collection and calculate a per event lepton SF leps = ak.concatenate([e, mu], axis=-1) events['lepSF_nom'] = ak.prod(leps.sf_nom, axis=-1) events['lepSF_hi'] = ak.prod(leps.sf_hi, axis=-1) events['lepSF_lo'] = ak.prod(leps.sf_lo, axis=-1) # Jet selection jetptname = 'pt_nom' if hasattr(j, 'pt_nom') else 'pt' ### Jet energy corrections if not isData: j["pt_raw"] = (1 - j.rawFactor) * j.pt j["mass_raw"] = (1 - j.rawFactor) * j.mass j["pt_gen"] = ak.values_astype(ak.fill_none(j.matched_gen.pt, 0), np.float32) j["rho"] = ak.broadcast_arrays(events.fixedGridRhoFastjetAll, j.pt)[0] events_cache = events.caches[0] corrected_jets = jet_factory.build(j, lazy_cache=events_cache) #print('jet pt: ',j.pt) #print('cor pt: ',corrected_jets.pt) #print('jes up: ',corrected_jets.JES_jes.up.pt) #print('jes down: ',corrected_jets.JES_jes.down.pt) #print(ak.fields(corrected_jets)) ''' # SYSTEMATICS jets = corrected_jets if(self.jetSyst == 'JERUp'): jets = corrected_jets.JER.up elif(self.jetSyst == 'JERDown'): jets = corrected_jets.JER.down elif(self.jetSyst == 'JESUp'): jets = corrected_jets.JES_jes.up elif(self.jetSyst == 'JESDown'): jets = corrected_jets.JES_jes.down ''' j['isGood'] = isTightJet(getattr(j, jetptname), j.eta, j.jetId, j.neHEF, j.neEmEF, j.chHEF, j.chEmEF, j.nConstituents) #j['isgood'] = isGoodJet(j.pt, j.eta, j.jetId) #j['isclean'] = isClean(j, e, mu) j['isClean'] = isClean(j, e, drmin=0.4) & isClean( j, mu, drmin=0.4) # & isClean(j, tau, drmin=0.4) goodJets = j[(j.isClean) & (j.isGood)] njets = ak.num(goodJets) ht = ak.sum(goodJets.pt, axis=-1) j0 = goodJets[ak.argmax(goodJets.pt, axis=-1, keepdims=True)] #nbtags = ak.num(goodJets[goodJets.btagDeepFlavB > 0.2770]) # Loose DeepJet WP if year == 2017: btagwpl = 0.0532 #WP loose else: btagwpl = 0.0490 #WP loose isBtagJetsLoose = (goodJets.btagDeepB > btagwpl) isNotBtagJetsLoose = np.invert(isBtagJetsLoose) nbtagsl = ak.num(goodJets[isBtagJetsLoose]) # Medium DeepJet WP if year == 2017: btagwpm = 0.3040 #WP medium else: btagwpm = 0.2783 #WP medium isBtagJetsMedium = (goodJets.btagDeepB > btagwpm) isNotBtagJetsMedium = np.invert(isBtagJetsMedium) nbtagsm = ak.num(goodJets[isBtagJetsMedium]) # Btag SF following 1a) in https://twiki.cern.ch/twiki/bin/viewauth/CMS/BTagSFMethods btagSF = np.ones_like(ht) btagSFUp = np.ones_like(ht) btagSFDo = np.ones_like(ht) if not isData: pt = goodJets.pt abseta = np.abs(goodJets.eta) flav = goodJets.hadronFlavour bJetSF = GetBTagSF(abseta, pt, flav) bJetSFUp = GetBTagSF(abseta, pt, flav, sys=1) bJetSFDo = GetBTagSF(abseta, pt, flav, sys=-1) bJetEff = GetBtagEff(abseta, pt, flav, year) bJetEff_data = bJetEff * bJetSF bJetEff_dataUp = bJetEff * bJetSFUp bJetEff_dataDo = bJetEff * bJetSFDo pMC = ak.prod(bJetEff[isBtagJetsMedium], axis=-1) * ak.prod( (1 - bJetEff[isNotBtagJetsMedium]), axis=-1) pData = ak.prod(bJetEff_data[isBtagJetsMedium], axis=-1) * ak.prod( (1 - bJetEff_data[isNotBtagJetsMedium]), axis=-1) pDataUp = ak.prod( bJetEff_dataUp[isBtagJetsMedium], axis=-1) * ak.prod( (1 - bJetEff_dataUp[isNotBtagJetsMedium]), axis=-1) pDataDo = ak.prod( bJetEff_dataDo[isBtagJetsMedium], axis=-1) * ak.prod( (1 - bJetEff_dataDo[isNotBtagJetsMedium]), axis=-1) pMC = ak.where(pMC == 0, 1, pMC) # removeing zeroes from denominator... btagSF = pData / pMC btagSFUp = pDataUp / pMC btagSFDo = pDataUp / pMC ################################################################## ### 2 same-sign leptons ################################################################## # emu singe = e[(nElec == 1) & (nMuon == 1) & (e.pt > -1)] singm = mu[(nElec == 1) & (nMuon == 1) & (mu.pt > -1)] em = ak.cartesian({"e": singe, "m": singm}) emSSmask = (em.e.charge * em.m.charge > 0) emSS = em[emSSmask] nemSS = len(ak.flatten(emSS)) # ee and mumu # pt>-1 to preserve jagged dimensions ee = e[(nElec == 2) & (nMuon == 0) & (e.pt > -1)] mm = mu[(nElec == 0) & (nMuon == 2) & (mu.pt > -1)] sumcharge = ak.sum(e.charge, axis=-1) + ak.sum(mu.charge, axis=-1) eepairs = ak.combinations(ee, 2, fields=["e0", "e1"]) eeSSmask = (eepairs.e0.charge * eepairs.e1.charge > 0) eeonZmask = (np.abs((eepairs.e0 + eepairs.e1).mass - 91.2) < 10) eeoffZmask = (eeonZmask == 0) mmpairs = ak.combinations(mm, 2, fields=["m0", "m1"]) mmSSmask = (mmpairs.m0.charge * mmpairs.m1.charge > 0) mmonZmask = (np.abs((mmpairs.m0 + mmpairs.m1).mass - 91.2) < 10) mmoffZmask = (mmonZmask == 0) eeSSonZ = eepairs[eeSSmask & eeonZmask] eeSSoffZ = eepairs[eeSSmask & eeoffZmask] mmSSonZ = mmpairs[mmSSmask & mmonZmask] mmSSoffZ = mmpairs[mmSSmask & mmoffZmask] neeSS = len(ak.flatten(eeSSonZ)) + len(ak.flatten(eeSSoffZ)) nmmSS = len(ak.flatten(mmSSonZ)) + len(ak.flatten(mmSSoffZ)) print('Same-sign events [ee, emu, mumu] = [%i, %i, %i]' % (neeSS, nemSS, nmmSS)) # Cuts eeSSmask = (ak.num(eeSSmask[eeSSmask]) > 0) mmSSmask = (ak.num(mmSSmask[mmSSmask]) > 0) eeonZmask = (ak.num(eeonZmask[eeonZmask]) > 0) eeoffZmask = (ak.num(eeoffZmask[eeoffZmask]) > 0) mmonZmask = (ak.num(mmonZmask[mmonZmask]) > 0) mmoffZmask = (ak.num(mmoffZmask[mmoffZmask]) > 0) emSSmask = (ak.num(emSSmask[emSSmask]) > 0) ################################################################## ### 3 leptons ################################################################## # eem muon_eem = mu[(nElec == 2) & (nMuon == 1) & (mu.pt > -1)] elec_eem = e[(nElec == 2) & (nMuon == 1) & (e.pt > -1)] ee_eem = ak.combinations(elec_eem, 2, fields=["e0", "e1"]) ee_eemZmask = (ee_eem.e0.charge * ee_eem.e1.charge < 1) & (np.abs( (ee_eem.e0 + ee_eem.e1).mass - 91.2) < 10) ee_eemOffZmask = (ee_eem.e0.charge * ee_eem.e1.charge < 1) & (np.abs( (ee_eem.e0 + ee_eem.e1).mass - 91.2) > 10) ee_eemZmask = (ak.num(ee_eemZmask[ee_eemZmask]) > 0) ee_eemOffZmask = (ak.num(ee_eemOffZmask[ee_eemOffZmask]) > 0) eepair_eem = (ee_eem.e0 + ee_eem.e1) trilep_eem = eepair_eem + muon_eem #ak.cartesian({"e0":ee_eem.e0,"e1":ee_eem.e1, "m":muon_eem}) # mme muon_mme = mu[(nElec == 1) & (nMuon == 2) & (mu.pt > -1)] elec_mme = e[(nElec == 1) & (nMuon == 2) & (e.pt > -1)] mm_mme = ak.combinations(muon_mme, 2, fields=["m0", "m1"]) mm_mmeZmask = (mm_mme.m0.charge * mm_mme.m1.charge < 1) & (np.abs( (mm_mme.m0 + mm_mme.m1).mass - 91.2) < 10) mm_mmeOffZmask = (mm_mme.m0.charge * mm_mme.m1.charge < 1) & (np.abs( (mm_mme.m0 + mm_mme.m1).mass - 91.2) > 10) mm_mmeZmask = (ak.num(mm_mmeZmask[mm_mmeZmask]) > 0) mm_mmeOffZmask = (ak.num(mm_mmeOffZmask[mm_mmeOffZmask]) > 0) mmpair_mme = (mm_mme.m0 + mm_mme.m1) trilep_mme = mmpair_mme + elec_mme mZ_mme = mmpair_mme.mass mZ_eem = eepair_eem.mass m3l_eem = trilep_eem.mass m3l_mme = trilep_mme.mass # eee and mmm eee = e[(nElec == 3) & (nMuon == 0) & (e.pt > -1)] mmm = mu[(nElec == 0) & (nMuon == 3) & (mu.pt > -1)] eee_leps = ak.combinations(eee, 3, fields=["e0", "e1", "e2"]) mmm_leps = ak.combinations(mmm, 3, fields=["m0", "m1", "m2"]) ee_pairs = ak.combinations(eee, 2, fields=["e0", "e1"]) mm_pairs = ak.combinations(mmm, 2, fields=["m0", "m1"]) ee_pairs_index = ak.argcombinations(eee, 2, fields=["e0", "e1"]) mm_pairs_index = ak.argcombinations(mmm, 2, fields=["m0", "m1"]) mmSFOS_pairs = mm_pairs[ (np.abs(mm_pairs.m0.pdgId) == np.abs(mm_pairs.m1.pdgId)) & (mm_pairs.m0.charge != mm_pairs.m1.charge)] offZmask_mm = ak.all( np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2) > 10., axis=1, keepdims=True) & (ak.num(mmSFOS_pairs) > 0) onZmask_mm = ak.any( np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2) < 10., axis=1, keepdims=True) eeSFOS_pairs = ee_pairs[ (np.abs(ee_pairs.e0.pdgId) == np.abs(ee_pairs.e1.pdgId)) & (ee_pairs.e0.charge != ee_pairs.e1.charge)] offZmask_ee = ak.all( np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2) > 10, axis=1, keepdims=True) & (ak.num(eeSFOS_pairs) > 0) onZmask_ee = ak.any( np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2) < 10, axis=1, keepdims=True) # Create masks **for event selection** eeeOnZmask = (ak.num(onZmask_ee[onZmask_ee]) > 0) eeeOffZmask = (ak.num(offZmask_ee[offZmask_ee]) > 0) mmmOnZmask = (ak.num(onZmask_mm[onZmask_mm]) > 0) mmmOffZmask = (ak.num(offZmask_mm[offZmask_mm]) > 0) # Now we need to create masks for the leptons in order to select leptons from the Z boson candidate (in onZ categories) ZeeMask = ak.argmin(np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2), axis=1, keepdims=True) ZmmMask = ak.argmin(np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2), axis=1, keepdims=True) Zee = eeSFOS_pairs[ZeeMask] Zmm = mmSFOS_pairs[ZmmMask] eZ0 = Zee.e0[ak.num(eeSFOS_pairs) > 0] eZ1 = Zee.e1[ak.num(eeSFOS_pairs) > 0] eZ = eZ0 + eZ1 mZ0 = Zmm.m0[ak.num(mmSFOS_pairs) > 0] mZ1 = Zmm.m1[ak.num(mmSFOS_pairs) > 0] mZ = mZ0 + mZ1 mZ_eee = eZ.mass mZ_mmm = mZ.mass # And for the W boson ZmmIndices = mm_pairs_index[ZmmMask] ZeeIndices = ee_pairs_index[ZeeMask] eW = eee[~ZeeIndices.e0 | ~ZeeIndices.e1] mW = mmm[~ZmmIndices.m0 | ~ZmmIndices.m1] triElec = eee_leps.e0 + eee_leps.e1 + eee_leps.e2 triMuon = mmm_leps.m0 + mmm_leps.m1 + mmm_leps.m2 m3l_eee = triElec.mass m3l_mmm = triMuon.mass ################################################################## ### >=4 leptons ################################################################## # 4lep cat is4lmask = ((nElec + nMuon) >= 4) muon_4l = mu[(is4lmask) & (mu.pt > -1)] elec_4l = e[(is4lmask) & (e.pt > -1)] # selecting 4 leading leptons leptons = ak.concatenate([e, mu], axis=-1) leptons_sorted = leptons[ak.argsort(leptons.pt, axis=-1, ascending=False)] lep4l = leptons_sorted[:, 0:4] e4l = lep4l[abs(lep4l.pdgId) == 11] mu4l = lep4l[abs(lep4l.pdgId) == 13] nElec4l = ak.num(e4l) nMuon4l = ak.num(mu4l) # Triggers trig_eeSS = passTrigger(events, 'ee', isData, dataset) trig_mmSS = passTrigger(events, 'mm', isData, dataset) trig_emSS = passTrigger(events, 'em', isData, dataset) trig_eee = passTrigger(events, 'eee', isData, dataset) trig_mmm = passTrigger(events, 'mmm', isData, dataset) trig_eem = passTrigger(events, 'eem', isData, dataset) trig_mme = passTrigger(events, 'mme', isData, dataset) trig_4l = triggerFor4l(events, nMuon, nElec, isData, dataset) # MET filters # Weights genw = np.ones_like(events['event']) if ( isData or len(self._wc_names_lst) > 0) else events['genWeight'] ### We need weights for: normalization, lepSF, triggerSF, pileup, btagSF... weights = {} for r in [ 'all', 'ee', 'mm', 'em', 'eee', 'mmm', 'eem', 'mme', 'eeee', 'eeem', 'eemm', 'mmme', 'mmmm' ]: # weights[r] = coffea.analysis_tools.Weights(len(events)) weights[r] = coffea.analysis_tools.Weights(len(events), storeIndividual=True) if len(self._wc_names_lst) > 0: sow = np.ones_like( sow ) # Not valid in nanoAOD for EFT samples, MUST use SumOfEFTweights at analysis level weights[r].add('norm', genw if isData else (xsec / sow) * genw) weights[r].add('btagSF', btagSF, btagSFUp, btagSFDo) weights[r].add('lepSF', events.lepSF_nom, events.lepSF_hi, events.lepSF_lo) # Extract the EFT quadratic coefficients and optionally use them to calculate the coefficients on the w**2 quartic function # eft_coeffs is never Jagged so convert immediately to numpy for ease of use. eft_coeffs = ak.to_numpy(events['EFTfitCoefficients']) if hasattr( events, "EFTfitCoefficients") else None if eft_coeffs is not None: # Check to see if the ordering of WCs for this sample matches what want if self._samples[dataset]['WCnames'] != self._wc_names_lst: eft_coeffs = efth.remap_coeffs( self._samples[dataset]['WCnames'], self._wc_names_lst, eft_coeffs) eft_w2_coeffs = efth.calc_w2_coeffs(eft_coeffs, self._dtype) if ( self._do_errors and eft_coeffs is not None) else None # Selections and cuts selections = PackedSelection() #(dtype='uint64') channels2LSS = ['eeSSonZ', 'eeSSoffZ', 'mmSSonZ', 'mmSSoffZ', 'emSS'] selections.add('eeSSonZ', (eeonZmask) & (eeSSmask) & (trig_eeSS)) selections.add('eeSSoffZ', (eeoffZmask) & (eeSSmask) & (trig_eeSS)) selections.add('mmSSonZ', (mmonZmask) & (mmSSmask) & (trig_mmSS)) selections.add('mmSSoffZ', (mmoffZmask) & (mmSSmask) & (trig_mmSS)) selections.add('emSS', (emSSmask) & (trig_emSS)) channels3L = ['eemSSonZ', 'eemSSoffZ', 'mmeSSonZ', 'mmeSSoffZ'] selections.add('eemSSonZ', (ee_eemZmask) & (trig_eem)) selections.add('eemSSoffZ', (ee_eemOffZmask) & (trig_eem)) selections.add('mmeSSonZ', (mm_mmeZmask) & (trig_mme)) selections.add('mmeSSoffZ', (mm_mmeOffZmask) & (trig_mme)) channels3L += ['eeeSSonZ', 'eeeSSoffZ', 'mmmSSonZ', 'mmmSSoffZ'] selections.add('eeeSSonZ', (eeeOnZmask) & (trig_eee)) selections.add('eeeSSoffZ', (eeeOffZmask) & (trig_eee)) selections.add('mmmSSonZ', (mmmOnZmask) & (trig_mmm)) selections.add('mmmSSoffZ', (mmmOffZmask) & (trig_mmm)) channels4L = ['eeee', 'eeem', 'eemm', 'mmme', 'mmmm'] selections.add('eeee', ((nElec4l == 4) & (nMuon4l == 0)) & (trig_4l)) selections.add('eeem', ((nElec4l == 3) & (nMuon4l == 1)) & (trig_4l)) selections.add('eemm', ((nElec4l == 2) & (nMuon4l == 2)) & (trig_4l)) selections.add('mmme', ((nElec4l == 1) & (nMuon4l == 3)) & (trig_4l)) selections.add('mmmm', ((nElec4l == 0) & (nMuon4l == 4)) & (trig_4l)) selections.add('ch+', (sumcharge > 0)) selections.add('ch-', (sumcharge < 0)) selections.add('ch0', (sumcharge == 0)) levels = ['base', '1+bm2+bl', '1bm', '2+bm'] selections.add('base', (nElec + nMuon >= 2)) selections.add('1+bm2+bl', (nElec + nMuon >= 2) & ((nbtagsm >= 1) & (nbtagsl >= 2))) selections.add('1bm', (nElec + nMuon >= 2) & (nbtagsm == 1)) selections.add('2+bm', (nElec + nMuon >= 2) & (nbtagsm >= 2)) # Variables invMass_eeSSonZ = (eeSSonZ.e0 + eeSSonZ.e1).mass invMass_eeSSoffZ = (eeSSoffZ.e0 + eeSSoffZ.e1).mass invMass_mmSSonZ = (mmSSonZ.m0 + mmSSonZ.m1).mass invMass_mmSSoffZ = (mmSSoffZ.m0 + mmSSoffZ.m1).mass invMass_emSS = (emSS.e + emSS.m).mass varnames = {} varnames['met'] = met.pt varnames['ht'] = ht varnames['njets'] = njets varnames['invmass'] = { 'eeSSonZ': invMass_eeSSonZ, 'eeSSoffZ': invMass_eeSSoffZ, 'mmSSonZ': invMass_mmSSonZ, 'mmSSoffZ': invMass_mmSSoffZ, 'emSS': invMass_emSS, 'eemSSonZ': mZ_eem, 'eemSSoffZ': mZ_eem, 'mmeSSonZ': mZ_mme, 'mmeSSoffZ': mZ_mme, 'eeeSSonZ': mZ_eee, 'eeeSSoffZ': mZ_eee, 'mmmSSonZ': mZ_mmm, 'mmmSSoffZ': mZ_mmm, } varnames['m3l'] = { 'eemSSonZ': m3l_eem, 'eemSSoffZ': m3l_eem, 'mmeSSonZ': m3l_mme, 'mmeSSoffZ': m3l_mme, 'eeeSSonZ': m3l_eee, 'eeeSSoffZ': m3l_eee, 'mmmSSonZ': m3l_mmm, 'mmmSSoffZ': m3l_mmm, } varnames['e0pt'] = e0.pt varnames['e0eta'] = e0.eta varnames['m0pt'] = m0.pt varnames['m0eta'] = m0.eta varnames['j0pt'] = j0.pt varnames['j0eta'] = j0.eta varnames['counts'] = np.ones_like(events['event']) # systematics systList = [] if isData == False: systList = ['nominal'] if self._do_systematics: systList = systList + [ 'lepSFUp', 'lepSFDown', 'btagSFUp', 'btagSFDown' ] else: systList = ['noweight'] # fill Histos hout = self.accumulator.identity() normweights = weights['all'].weight().flatten( ) # Why does it not complain about .flatten() here? sowweights = np.ones_like(normweights) if len( self._wc_names_lst) > 0 else normweights hout['SumOfEFTweights'].fill(sample=histAxisName, SumOfEFTweights=varnames['counts'], weight=sowweights, eft_coeff=eft_coeffs, eft_err_coeff=eft_w2_coeffs) for syst in systList: for var, v in varnames.items(): for ch in channels2LSS + channels3L + channels4L: for sumcharge in ['ch+', 'ch-', 'ch0']: for lev in levels: #find the event weight to be used when filling the histograms weightSyst = syst #in the case of 'nominal', or the jet energy systematics, no weight systematic variation is used (weightSyst=None) if syst in [ 'nominal', 'JERUp', 'JERDown', 'JESUp', 'JESDown' ]: weightSyst = None # no weight systematic for these variations if syst == 'noweight': weight = np.ones(len(events)) # for data else: # call weights.weight() with the name of the systematic to be varied if ch in channels3L: ch_w = ch[:3] elif ch in channels2LSS: ch_w = ch[:2] else: ch_w = ch weight = weights['all'].weight( weightSyst ) if isData else weights[ch_w].weight( weightSyst) cuts = [ch] + [lev] + [sumcharge] cut = selections.all(*cuts) weights_flat = weight[cut].flatten( ) # Why does it not complain about .flatten() here? weights_ones = np.ones_like(weights_flat, dtype=np.int) eft_coeffs_cut = eft_coeffs[ cut] if eft_coeffs is not None else None eft_w2_coeffs_cut = eft_w2_coeffs[ cut] if eft_w2_coeffs is not None else None # filling histos if var == 'invmass': if ((ch in [ 'eeeSSoffZ', 'mmmSSoffZ', 'eeeSSonZ', 'mmmSSonZ' ]) or (ch in channels4L)): continue else: values = ak.flatten(v[ch][cut]) hout['invmass'].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, invmass=values, weight=weights_flat, systematic=syst) elif var == 'm3l': if ((ch in channels2LSS) or (ch in [ 'eeeSSoffZ', 'mmmSSoffZ', 'eeeSSonZ', 'mmmSSonZ' ]) or (ch in channels4L)): continue values = ak.flatten(v[ch][cut]) hout['m3l'].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, m3l=values, weight=weights_flat, systematic=syst) else: values = v[cut] # These all look identical, do we need if/else here? if var == 'ht': hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, ht=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'met': hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, met=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'njets': hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, njets=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'nbtags': hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, nbtags=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'counts': hout[var].fill(counts=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_ones, systematic=syst) elif var == 'j0eta': if lev == 'base': continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, j0eta=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'e0pt': if ch in [ 'mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ', 'mmmm' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, e0pt=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst ) # Crashing here, not sure why. Related to values? elif var == 'm0pt': if ch in [ 'eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ', 'eeee' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, m0pt=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'e0eta': if ch in [ 'mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ', 'mmmm' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, e0eta=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'm0eta': if ch in [ 'eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ', 'eeee' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, m0eta=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'j0pt': if lev == 'base': continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, j0pt=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) return hout
def cross(first, second): tmp = ak.cartesian([first, second]) combs = (tmp['0'] + tmp['1']) combs['0'] = tmp['0'] combs['1'] = tmp['1'] return combs
def eval_model(model, dataset, config, outdir): ibatch = 0 jetdef = fastjet.JetDefinition(fastjet.antikt_algorithm, 0.4) for elem in tqdm(dataset, desc="Evaluating model"): y_pred = model.predict(elem["X"], verbose=False) np_outfile = "{}/pred_batch{}.npz".format(outdir, ibatch) ygen = unpack_target(elem["ygen"], config["dataset"]["num_output_classes"], config) ycand = unpack_target(elem["ycand"], config["dataset"]["num_output_classes"], config) outs = {} for key in y_pred.keys(): outs["gen_{}".format(key)] = ygen[key].numpy() outs["cand_{}".format(key)] = ycand[key].numpy() outs["pred_{}".format(key)] = y_pred[key] jets_coll = {} jets_const = {} for typ in ["gen", "cand", "pred"]: cls_id = np.argmax(outs["{}_cls".format(typ)], axis=-1) valid = cls_id != 0 pt = awkward.from_iter( [y[m][:, 0] for y, m in zip(outs["{}_pt".format(typ)], valid)]) eta = awkward.from_iter([ y[m][:, 0] for y, m in zip(outs["{}_eta".format(typ)], valid) ]) phi = np.arctan2(outs["{}_sin_phi".format(typ)], outs["{}_cos_phi".format(typ)]) phi = awkward.from_iter([y[m][:, 0] for y, m in zip(phi, valid)]) e = awkward.from_iter([ y[m][:, 0] for y, m in zip(outs["{}_energy".format(typ)], valid) ]) vec = vector.arr({"pt": pt, "eta": eta, "phi": phi, "e": e}) cluster = fastjet.ClusterSequence(vec.to_xyzt(), jetdef) jets = cluster.inclusive_jets() jet_constituents = cluster.constituent_index() jets_coll[typ] = jets[jets.pt > 5.0] jets_const[typ] = jet_constituents[jets.pt > 5.0] for key in ["pt", "eta", "phi", "energy"]: outs["jets_gen_{}".format(key)] = awkward.to_numpy( awkward.flatten(getattr(jets_coll["gen"], key))) outs["jets_cand_{}".format(key)] = awkward.to_numpy( awkward.flatten(getattr(jets_coll["cand"], key))) outs["jets_pred_{}".format(key)] = awkward.to_numpy( awkward.flatten(getattr(jets_coll["pred"], key))) # DeltaR match between genjets and PF/MLPF jets cart = awkward.cartesian([jets_coll["gen"], jets_coll["pred"]], nested=True) jets_a, jets_b = awkward.unzip(cart) drs = deltar(jets_a, jets_b) match_gen_to_pred = [awkward.where(d < 0.1) for d in drs] m0 = awkward.from_iter([m[0] for m in match_gen_to_pred]) m1 = awkward.from_iter([m[1] for m in match_gen_to_pred]) j1s = jets_coll["gen"][m0] j2s = jets_coll["pred"][m1] outs["jets_pt_gen_to_pred"] = np.stack([ awkward.to_numpy(awkward.flatten(j1s.pt)), awkward.to_numpy(awkward.flatten(j2s.pt)) ], axis=-1) cart = awkward.cartesian([jets_coll["gen"], jets_coll["cand"]], nested=True) jets_a, jets_b = awkward.unzip(cart) drs = deltar(jets_a, jets_b) match_gen_to_pred = [awkward.where(d < 0.1) for d in drs] m0 = awkward.from_iter([m[0] for m in match_gen_to_pred]) m1 = awkward.from_iter([m[1] for m in match_gen_to_pred]) j1s = jets_coll["gen"][m0] j2s = jets_coll["cand"][m1] outs["jets_pt_gen_to_cand"] = np.stack([ awkward.to_numpy(awkward.flatten(j1s.pt)), awkward.to_numpy(awkward.flatten(j2s.pt)) ], axis=-1) np.savez(np_outfile, X=elem["X"], **outs) ibatch += 1
def best_match(gen_hyp=None, jets=None, leptons=None, met=None): if gen_hyp is None: raise ValueError("Gen Objects gen_hyp needed for matching") if jets is None: raise ValueError("Reco jets needed for matching") if leptons is None: raise ValueError("Reco leptons needed for matching") if met is None: raise ValueError("Reco met needed for matching") if not ak.all(ak.num(gen_hyp) == 1): raise ValueError("Not all events for matching are semileptonic") jets_ak = ak.with_name(jets[["pt", "eta", "phi", "mass"]],"PtEtaPhiMLorentzVector") leps_ak = ak.with_name(leptons[["pt", "eta", "phi", "mass"]],"PtEtaPhiMLorentzVector") # init dict of objects matched_objects = {} # match jet closest to gen objects for genobj in ['BHad', 'BLep', 'WJa', 'WJb']: genobj_ak = ak.with_name(gen_hyp[genobj][["pt", "eta", "phi", "mass"]],"PtEtaPhiMLorentzVector") jets_akc, genobj_akc = ak.unzip(ak.cartesian([jets_ak, genobj_ak], nested=False)) deltaRs = jets_akc.delta_r(genobj_akc) # find deltaRs between jets and gen object indexOfMin = ak.unflatten(ak.argmin(deltaRs, axis=1), ak.num(genobj_ak)) passing_inds = deltaRs[indexOfMin] < 0.4 matched_jets_inds = indexOfMin[passing_inds] matched_jets = jets[matched_jets_inds] ## add matched perm objects matched_objects[genobj] = ak.Array({ 'pt' : matched_jets.pt, 'eta' : matched_jets.eta, 'phi' : matched_jets.phi, 'mass' : matched_jets.mass, 'jetIdx' : matched_jets_inds, # index of jet that the gen object is matched to in the event }, with_name="PtEtaPhiMLorentzVector") # match lepton closest to gen lepton genlep_ak = ak.with_name(gen_hyp['Lepton'][["pt", "eta", "phi", "mass"]],"PtEtaPhiMLorentzVector") lep_akc, genlep_akc = ak.unzip(ak.cartesian([leps_ak, genlep_ak], nested=False)) lepDRs = lep_akc.delta_r(genlep_akc) lepIdxOfMin = ak.unflatten(ak.argmin(lepDRs, axis=1), ak.num(genlep_ak)) passing_inds = lepDRs[lepIdxOfMin] < 0.4 matched_leps_inds = lepIdxOfMin[passing_inds] matched_leps = leptons[matched_leps_inds] ## add matched perm objects matched_objects['Lepton'] = ak.Array({key: matched_leps[key] for key in matched_leps.fields}, with_name="PtEtaPhiMLorentzVector") # solve for neutrino nu_array = np.zeros((len(ak.num(jets)), 4), dtype='float64') # convert all inputs into 2d numpy arrays of dtype=float64 (won't work if they're not float64) blep_inputs = np.stack((ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['BLep'].px, 1), -999))).astype('float64'), ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['BLep'].py, 1), -999))).astype('float64'),\ ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['BLep'].pz, 1), -999))).astype('float64'), ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['BLep'].energy, 1), -999))).astype('float64')), axis=-1) lep_inputs = np.stack((ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['Lepton'].px, 1), -999))).astype('float64'), ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['Lepton'].py, 1), -999))).astype('float64'),\ ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['Lepton'].pz, 1), -999))).astype('float64'), ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['Lepton'].energy, 1), -999))).astype('float64')), axis=-1) met_inputs = np.stack((ak.to_numpy(ak.fill_none(met.px, -999)).astype('float64'), ak.to_numpy(ak.fill_none(met.py, -999)).astype('float64')), axis=-1) nu_array = find_nu(bleps=blep_inputs, leptons=lep_inputs, met=met_inputs, nu_array=nu_array) valid_nu = ~((nu_array[:, 3] > 1e20) | (nu_array[:, 3] == 0)) # events that have a solution and matched blep # convert px, py, pz to pt, eta, phi nu_px, nu_py, nu_pz = nu_array[:, 0][valid_nu], nu_array[:, 1][valid_nu], nu_array[:, 2][valid_nu] nu_mom, nu_pt = np.sqrt(np.square(nu_px)+np.square(nu_py)+np.square(nu_pz)), np.sqrt(np.square(nu_px)+np.square(nu_py)) nu_phi = np.arctan2(nu_py, nu_px) nu_eta = np.arcsinh(nu_pz/nu_pt) matched_objects['Nu'] = ak.Array({ 'pt' : ak.unflatten(nu_pt, valid_nu.astype(int)), 'eta' : ak.unflatten(nu_eta, valid_nu.astype(int)), 'phi' : ak.unflatten(nu_phi, valid_nu.astype(int)), 'mass' : ak.zeros_like(ak.unflatten(nu_array[:, 0][valid_nu], valid_nu.astype(int))), 'chi2' : ak.unflatten(nu_array[:, 3][valid_nu], valid_nu.astype(int)), }, with_name="PtEtaPhiMLorentzVector") matched_perm = make_perm_table(bhad=matched_objects['BHad'], blep=matched_objects['BLep'], wja=matched_objects['WJa'], wjb=matched_objects['WJb'], lepton=matched_objects['Lepton'], met=met, nu=matched_objects['Nu']) return matched_perm
def process(self, events): # Dataset parameters dataset = events.metadata['dataset'] year = self._samples[dataset]['year'] xsec = self._samples[dataset]['xsec'] sow = self._samples[dataset]['nSumOfWeights' ] isData = self._samples[dataset]['isData'] datasets = ['SingleMuon', 'SingleElectron', 'EGamma', 'MuonEG', 'DoubleMuon', 'DoubleElectron'] for d in datasets: if d in dataset: dataset = dataset.split('_')[0] # Inittialize objects met = events.GenMET e = events.GenPart[abs(events.GenPart.pdgId)==11] mu = events.GenPart[abs(events.GenPart.pdgId)==13] tau = events.GenPart[abs(events.GenPart.pdgId)==15] j = events.GenJet leading_mu = mu[ak.argmax(mu.pt,axis=-1,keepdims=True)] leading_e = e[ak.argmax(e.pt,axis=-1,keepdims=True)] nElec = ak.num(e) nMuon = ak.num(mu) nTau = ak.num(tau) twoLeps = (nElec+nMuon) == 2 threeLeps = (nElec+nMuon) == 3 twoElec = (nElec == 2) twoMuon = (nMuon == 2) e0 = e[ak.argmax(e.pt,axis=-1,keepdims=True)] m0 = mu[ak.argmax(mu.pt,axis=-1,keepdims=True)] elecs = e[ak.argsort(e.pt, ascending=False)] muons = mu[ak.argsort(mu.pt, ascending=False)] e1 = elecs e2 = elecs m1 = muons m2 = muons # Jet selection jetptname = 'pt_nom' if hasattr(j, 'pt_nom') else 'pt' njets = ak.num(j) ht = ak.sum(j.pt,axis=-1) jets = j[ak.argsort(j.pt, ascending=False)] j0 = j[ak.argmax(j.pt,axis=-1,keepdims=True)] j1 = jets j2 = jets j3 = jets nbtags = ak.num(j[abs(j.hadronFlavour)==5]) ################################################################## ### 2 same-sign leptons ################################################################## # emu singe = e [(nElec==1)&(nMuon==1)&(e .pt>-1)] singm = mu[(nElec==1)&(nMuon==1)&(mu.pt>-1)] em = ak.cartesian({"e":singe,"m":singm}) emSSmask = (em.e.pdgId*em.m.pdgId>0) emSS = em[emSSmask] nemSS = len(ak.flatten(emSS)) year = 2018 lepSF_emSS = GetLeptonSF(mu.pt, mu.eta, 'm', e.pt, e.eta, 'e', year=year) # ee and mumu # pt>-1 to preserve jagged dimensions ee = e [(nElec==2)&(nMuon==0)&(e.pt>-1)] mm = mu[(nElec==0)&(nMuon==2)&(mu.pt>-1)] eepairs = ak.combinations(ee, 2, fields=["e0","e1"]) eeSSmask = (eepairs.e0.pdgId*eepairs.e1.pdgId>0) eeonZmask = (np.abs((eepairs.e0+eepairs.e1).mass-91.2)<10) eeoffZmask = (eeonZmask==0) mmpairs = ak.combinations(mm, 2, fields=["m0","m1"]) mmSSmask = (mmpairs.m0.pdgId*mmpairs.m1.pdgId>0) mmonZmask = (np.abs((mmpairs.m0+mmpairs.m1).mass-91.2)<10) mmoffZmask = (mmonZmask==0) eeSSonZ = eepairs[eeSSmask & eeonZmask] eeSSoffZ = eepairs[eeSSmask & eeoffZmask] mmSSonZ = mmpairs[mmSSmask & mmonZmask] mmSSoffZ = mmpairs[mmSSmask & mmoffZmask] neeSS = len(ak.flatten(eeSSonZ)) + len(ak.flatten(eeSSoffZ)) nmmSS = len(ak.flatten(mmSSonZ)) + len(ak.flatten(mmSSoffZ)) lepSF_eeSS = GetLeptonSF(eepairs.e0.pt, eepairs.e0.eta, 'e', eepairs.e1.pt, eepairs.e1.eta, 'e', year=year) lepSF_mumuSS = GetLeptonSF(mmpairs.m0.pt, mmpairs.m0.eta, 'm', mmpairs.m1.pt, mmpairs.m1.eta, 'm', year=year) print('Same-sign events [ee, emu, mumu] = [%i, %i, %i]'%(neeSS, nemSS, nmmSS)) # Cuts eeSSmask = (ak.num(eeSSmask[eeSSmask])>0) mmSSmask = (ak.num(mmSSmask[mmSSmask])>0) eeonZmask = (ak.num(eeonZmask[eeonZmask])>0) eeoffZmask = (ak.num(eeoffZmask[eeoffZmask])>0) mmonZmask = (ak.num(mmonZmask[mmonZmask])>0) mmoffZmask = (ak.num(mmoffZmask[mmoffZmask])>0) emSSmask = (ak.num(emSSmask[emSSmask])>0) ################################################################## ### 3 leptons ################################################################## # eem muon_eem = mu[(nElec==2)&(nMuon==1)&(mu.pt>-1)] elec_eem = e[(nElec==2)&(nMuon==1)&( e.pt>-1)] ee_eem = ak.combinations(elec_eem, 2, fields=["e0", "e1"]) ee_eemZmask = (ee_eem.e0.pdgId*ee_eem.e1.pdgId<1)&(np.abs((ee_eem.e0+ee_eem.e1).mass-91.2)<10) ee_eemOffZmask = (ee_eem.e0.pdgId*ee_eem.e1.pdgId<1)&(np.abs((ee_eem.e0+ee_eem.e1).mass-91.2)>10) ee_eemZmask = (ak.num(ee_eemZmask[ee_eemZmask])>0) ee_eemOffZmask = (ak.num(ee_eemOffZmask[ee_eemOffZmask])>0) eepair_eem = (ee_eem.e0+ee_eem.e1) trilep_eem = eepair_eem+muon_eem #ak.cartesian({"e0":ee_eem.e0,"e1":ee_eem.e1, "m":muon_eem}) lepSF_eem = GetLeptonSF(ee_eem.e0.pt, ee_eem.e0.eta, 'e', ee_eem.e1.pt, ee_eem.e1.eta, 'e', mu.pt, mu.eta, 'm', year) # mme muon_mme = mu[(nElec==1)&(nMuon==2)&(mu.pt>-1)] elec_mme = e[(nElec==1)&(nMuon==2)&( e.pt>-1)] mm_mme = ak.combinations(muon_mme, 2, fields=["m0", "m1"]) mm_mmeZmask = (mm_mme.m0.pdgId*mm_mme.m1.pdgId<1)&(np.abs((mm_mme.m0+mm_mme.m1).mass-91.2)<10) mm_mmeOffZmask = (mm_mme.m0.pdgId*mm_mme.m1.pdgId<1)&(np.abs((mm_mme.m0+mm_mme.m1).mass-91.2)>10) mm_mmeZmask = (ak.num(mm_mmeZmask[mm_mmeZmask])>0) mm_mmeOffZmask = (ak.num(mm_mmeOffZmask[mm_mmeOffZmask])>0) mmpair_mme = (mm_mme.m0+mm_mme.m1) trilep_mme = mmpair_mme+elec_mme mZ_mme = mmpair_mme.mass mZ_eem = eepair_eem.mass m3l_eem = trilep_eem.mass m3l_mme = trilep_mme.mass lepSF_mme = GetLeptonSF(mm_mme.m0.pt, mm_mme.m0.eta, 'm', mm_mme.m1.pt, mm_mme.m1.eta, 'm', e.pt, e.eta, 'e', year) # eee and mmm eee = e[(nElec==3)&(nMuon==0)&( e.pt>-1)] mmm = mu[(nElec==0)&(nMuon==3)&(mu.pt>-1)] eee_leps = ak.combinations(eee, 3, fields=["e0", "e1", "e2"]) mmm_leps = ak.combinations(mmm, 3, fields=["m0", "m1", "m2"]) ee_pairs = ak.combinations(eee, 2, fields=["e0", "e1"]) mm_pairs = ak.combinations(mmm, 2, fields=["m0", "m1"]) ee_pairs_index = ak.argcombinations(eee, 2, fields=["e0", "e1"]) mm_pairs_index = ak.argcombinations(mmm, 2, fields=["m0", "m1"]) lepSF_eee = GetLeptonSF(eee_leps.e0.pt, eee_leps.e0.eta, 'e', eee_leps.e1.pt, eee_leps.e1.eta, 'e', eee_leps.e2.pt, eee_leps.e2.eta, 'e', year) lepSF_mmm = GetLeptonSF(mmm_leps.m0.pt, mmm_leps.m0.eta, 'm', mmm_leps.m1.pt, mmm_leps.m1.eta, 'm', mmm_leps.m2.pt, mmm_leps.m2.eta, 'm', year) mmSFOS_pairs = mm_pairs[(np.abs(mm_pairs.m0.pdgId) == np.abs(mm_pairs.m1.pdgId)) & (mm_pairs.m0.pdgId != mm_pairs.m1.pdgId)] offZmask_mm = ak.all(np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2)>10., axis=1, keepdims=True) & (ak.num(mmSFOS_pairs)>0) onZmask_mm = ak.any(np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2)<10., axis=1, keepdims=True) eeSFOS_pairs = ee_pairs[(np.abs(ee_pairs.e0.pdgId) == np.abs(ee_pairs.e1.pdgId)) & (ee_pairs.e0.pdgId != ee_pairs.e1.pdgId)] offZmask_ee = ak.all(np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2)>10, axis=1, keepdims=True) & (ak.num(eeSFOS_pairs)>0) onZmask_ee = ak.any(np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2)<10, axis=1, keepdims=True) # Create masks **for event selection** eeeOnZmask = (ak.num(onZmask_ee[onZmask_ee])>0) eeeOffZmask = (ak.num(offZmask_ee[offZmask_ee])>0) mmmOnZmask = (ak.num(onZmask_mm[onZmask_mm])>0) mmmOffZmask = (ak.num(offZmask_mm[offZmask_mm])>0) # Now we need to create masks for the leptons in order to select leptons from the Z boson candidate (in onZ categories) ZeeMask = ak.argmin(np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2),axis=1,keepdims=True) ZmmMask = ak.argmin(np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2),axis=1,keepdims=True) Zee = eeSFOS_pairs[ZeeMask] Zmm = mmSFOS_pairs[ZmmMask] eZ0= Zee.e0[ak.num(eeSFOS_pairs)>0] eZ1= Zee.e1[ak.num(eeSFOS_pairs)>0] eZ = eZ0+eZ1 mZ0= Zmm.m0[ak.num(mmSFOS_pairs)>0] mZ1= Zmm.m1[ak.num(mmSFOS_pairs)>0] mZ = mZ0+mZ1 mZ_eee = eZ.mass mZ_mmm = mZ.mass # And for the W boson ZmmIndices = mm_pairs_index[ZmmMask] ZeeIndices = ee_pairs_index[ZeeMask] eW = eee[~ZeeIndices.e0 | ~ZeeIndices.e1] mW = mmm[~ZmmIndices.m0 | ~ZmmIndices.m1] triElec = eee_leps.e0+eee_leps.e1+eee_leps.e2 triMuon = mmm_leps.m0+mmm_leps.m1+mmm_leps.m2 m3l_eee = triElec.mass m3l_mmm = triMuon.mass # Triggers trig_eeSS = passTrigger(events,'ee',isData,dataset) trig_mmSS = passTrigger(events,'mm',isData,dataset) trig_emSS = passTrigger(events,'em',isData,dataset) trig_eee = passTrigger(events,'eee',isData,dataset) trig_mmm = passTrigger(events,'mmm',isData,dataset) trig_eem = passTrigger(events,'eem',isData,dataset) trig_mme = passTrigger(events,'mme',isData,dataset) # MET filters # Weights genw = np.ones_like(events['MET_pt']) if isData else events['genWeight'] ### We need weights for: normalization, lepSF, triggerSF, pileup, btagSF... weights = {} for r in ['all', 'ee', 'mm', 'em', 'eee', 'mmm', 'eem', 'mme']: weights[r] = coffea.analysis_tools.Weights(len(events)) weights[r].add('norm',genw if isData else (xsec/sow)*genw) weights['ee'].add('lepSF_eeSS', lepSF_eeSS) weights['em'].add('lepSF_emSS', lepSF_emSS) weights['mm'].add('lepSF_mmSS', lepSF_mumuSS) weights['eee'].add('lepSF_eee', lepSF_eee) weights['mmm'].add('lepSF_mmm', lepSF_mmm) weights['mme'].add('lepSF_mme', lepSF_mme) weights['eem'].add('lepSF_eem', lepSF_eem) # Extract the EFT quadratic coefficients and optionally use them to calculate the coefficients on the w**2 quartic function # eft_coeffs is never Jagged so convert immediately to numpy for ease of use. eft_coeffs = ak.to_numpy(events['EFTfitCoefficients']) if hasattr(events, "EFTfitCoefficients") else None eft_w2_coeffs = efth.calc_w2_coeffs(eft_coeffs,self._dtype) if (self._do_errors and eft_coeffs is not None) else None # Selections and cuts selections = PackedSelection() channels2LSS = ['eeSSonZ', 'eeSSoffZ', 'mmSSonZ', 'mmSSoffZ', 'emSS'] selections.add('eeSSonZ', (eeonZmask)&(eeSSmask)&(trig_eeSS)) selections.add('eeSSoffZ', (eeoffZmask)&(eeSSmask)&(trig_eeSS)) selections.add('mmSSonZ', (mmonZmask)&(mmSSmask)&(trig_mmSS)) selections.add('mmSSoffZ', (mmoffZmask)&(mmSSmask)&(trig_mmSS)) selections.add('emSS', (emSSmask)&(trig_emSS)) channels3L = ['eemSSonZ', 'eemSSoffZ', 'mmeSSonZ', 'mmeSSoffZ'] selections.add('eemSSonZ', (ee_eemZmask)&(trig_eem)) selections.add('eemSSoffZ', (ee_eemOffZmask)&(trig_eem)) selections.add('mmeSSonZ', (mm_mmeZmask)&(trig_mme)) selections.add('mmeSSoffZ', (mm_mmeOffZmask)&(trig_mme)) channels3L += ['eeeSSonZ', 'eeeSSoffZ', 'mmmSSonZ', 'mmmSSoffZ'] selections.add('eeeSSonZ', (eeeOnZmask)&(trig_eee)) selections.add('eeeSSoffZ', (eeeOffZmask)&(trig_eee)) selections.add('mmmSSonZ', (mmmOnZmask)&(trig_mmm)) selections.add('mmmSSoffZ', (mmmOffZmask)&(trig_mmm)) levels = ['base', '2jets', '4jets', '4j1b', '4j2b'] selections.add('base', (nElec+nMuon>=2)) selections.add('2jets',(njets>=2)) selections.add('4jets',(njets>=4)) selections.add('4j1b',(njets>=4)&(nbtags>=1)) selections.add('4j2b',(njets>=4)&(nbtags>=2)) # Variables invMass_eeSSonZ = ( eeSSonZ.e0+ eeSSonZ.e1).mass invMass_eeSSoffZ = (eeSSoffZ.e0+eeSSoffZ.e1).mass invMass_mmSSonZ = ( mmSSonZ.m0+ mmSSonZ.m1).mass invMass_mmSSoffZ = (mmSSoffZ.m0+mmSSoffZ.m1).mass invMass_emSS = (emSS.e+emSS.m).mass varnames = {} varnames['met'] = met.pt varnames['ht'] = ht varnames['njets'] = njets varnames['nbtags'] = nbtags varnames['invmass'] = { 'eeSSonZ' : invMass_eeSSonZ, 'eeSSoffZ' : invMass_eeSSoffZ, 'mmSSonZ' : invMass_mmSSonZ, 'mmSSoffZ' : invMass_mmSSoffZ, 'emSS' : invMass_emSS, 'eemSSonZ' : mZ_eem, 'eemSSoffZ' : mZ_eem, 'mmeSSonZ' : mZ_mme, 'mmeSSoffZ' : mZ_mme, 'eeeSSonZ' : mZ_eee, 'eeeSSoffZ' : mZ_eee, 'mmmSSonZ' : mZ_mmm, 'mmmSSoffZ' : mZ_mmm, } varnames['m3l'] = { 'eemSSonZ' : m3l_eem, 'eemSSoffZ' : m3l_eem, 'mmeSSonZ' : m3l_mme, 'mmeSSoffZ' : m3l_mme, 'eeeSSonZ' : m3l_eee, 'eeeSSoffZ' : m3l_eee, 'mmmSSonZ' : m3l_mmm, 'mmmSSoffZ' : m3l_mmm, } varnames['e0pt' ] = e0.pt varnames['e0eta'] = e0.eta varnames['m0pt' ] = m0.pt varnames['m0eta'] = m0.eta varnames['e1pt' ] = e1 varnames['e1eta'] = e1 varnames['e2pt' ] = e2 varnames['e2eta'] = e2 varnames['m1pt' ] = m1 varnames['m1eta'] = m1 varnames['m2pt' ] = m2 varnames['m2eta'] = m2 varnames['j0pt' ] = j0.pt varnames['j0eta'] = j0.eta varnames['j1pt'] = j1 varnames['j1eta'] = j1 varnames['j2pt'] = j2 varnames['j2eta'] = j2 varnames['j3pt'] = j3 varnames['j3eta'] = j3 varnames['counts'] = np.ones_like(events.GenMET.pt) # fill Histos hout = self.accumulator.identity() normweights = weights['all'].weight().flatten() # Why does it not complain about .flatten() here? hout['SumOfEFTweights'].fill(sample=dataset, SumOfEFTweights=varnames['counts'], weight=normweights, eft_coeff=eft_coeffs, eft_err_coeff=eft_w2_coeffs) for var, v in varnames.items(): for ch in channels2LSS+channels3L: for lev in levels: weight = weights[ ch[:3] if (ch.startswith('eee') or ch.startswith('mmm') or ch.startswith('eem') or ch.startswith('mme')) else ch[:2]].weight() cuts = [ch] + [lev] cut = selections.all(*cuts) weights_flat = weight[cut].flatten() # Why does it not complain about .flatten() here? weights_ones = np.ones_like(weights_flat, dtype=np.int) eft_coeffs_cut = eft_coeffs[cut] if eft_coeffs is not None else None eft_w2_coeffs_cut = eft_w2_coeffs[cut] if eft_w2_coeffs is not None else None if var == 'invmass': if ch in ['eeeSSoffZ', 'mmmSSoffZ']: continue elif ch in ['eeeSSonZ' , 'mmmSSonZ' ]: continue #values = v[ch] else : values = ak.flatten(v[ch][cut]) hout['invmass'].fill(eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, sample=dataset, channel=ch, cut=lev, invmass=values, weight=weights_flat) elif var == 'm3l': if ch in ['eeSSonZ','eeSSoffZ', 'mmSSonZ', 'mmSSoffZ','emSS', 'eeeSSoffZ', 'mmmSSoffZ', 'eeeSSonZ' , 'mmmSSonZ']: continue values = ak.flatten(v[ch][cut]) hout['m3l'].fill(sample=dataset, channel=ch, cut=lev, m3l=values, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) else: values = v[cut] if var == 'ht' : hout[var].fill(ht=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'met' : hout[var].fill(met=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'njets' : hout[var].fill(njets=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'nbtags': hout[var].fill(nbtags=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'counts': hout[var].fill(counts=values, sample=dataset, channel=ch, cut=lev, weight=weights_ones) elif var == 'j0eta' : if lev == 'base': continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(j0eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'e0pt' : if ch in ['mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ']: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(e0pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'm0pt' : if ch in ['eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ']: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(m0pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'e0eta' : if ch in ['mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ']: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(e0eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'm0eta': if ch in ['eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ']: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(m0eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'j0pt' : if lev == 'base': continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(j0pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'j1pt': if lev == "base": continue values = values.pt[:,1] #values = ak.flatten(values) hout[var].fill(j1pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var =='j1eta': if lev == 'base': continue values = values.eta[:,1] hout[var].fill(j1eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'j2pt': if lev in ['base', "2jets"]: continue values = values.pt[:,2] hout[var].fill(j2pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'j2eta': if lev in ['base', "2jets"]: continue values = values.eta[:,2] hout[var].fill(j2eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'j3pt': if lev in ['base', "2jets"]: continue values = values.pt[:,3] hout[var].fill(j3pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'j3eta': if lev in ['base', "2jets"]: continue values = values.eta[:,3] hout[var].fill(j3eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'e1pt': if ch in ['mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ', 'mmeSSonZ', 'mmeSSoffZ', 'emSS']: continue values = values.pt[:,1] hout[var].fill(e1pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'e1eta': if ch in ['mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ', 'mmeSSonZ', 'mmeSSoffZ', 'emSS']: continue values = values.eta[:,1] hout[var].fill(e1eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'e2pt': if ch in ['eeeSSonZ', 'eeeSSoffZ']: values = values.pt[:,2] hout[var].fill(e2pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'e2eta': if ch in ['eeeSSonZ', 'eeeSSoffZ']: values = values.eta[:,2] hout[var].fill(e2eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'm1pt': if ch in ['eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ', 'eemSSonZ', 'eemSSoffZ', 'emSS']: continue values = values.pt[:,1] hout[var].fill(m1pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'm1eta': if ch in ['eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ', 'eemSSonZ', 'eemSSoffZ', 'emSS']: continue values = values.eta[:,1] hout[var].fill(m1eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'm2pt': if ch in ['mmmSSonZ', 'mmmSSoffZ']: values = values.pt[:,2] hout[var].fill(m2pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) elif var == 'm2eta': if ch in ['mmmSSonZ', 'mmmSSoffZ']: values = values.eta[:,2] hout[var].fill(m2eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut) return hout
def test_cartesian(): one = ak.Array(np.arange(2 * 3 * 5 * 7).reshape(2, 3, 5, 7).tolist()) two = ak.Array(np.arange(2 * 3 * 5 * 7).reshape(2, 3, 5, 7).tolist()) assert (str(ak.cartesian([one, two], axis=0, nested=True).type) == "2 * 2 * (var * var * var * int64, var * var * var * int64)") assert (str(ak.cartesian([one, two], axis=1, nested=True).type) == "2 * var * var * (var * var * int64, var * var * int64)") assert (str(ak.cartesian([one, two], axis=2, nested=True).type) == "2 * var * var * var * (var * int64, var * int64)") assert (str(ak.cartesian( [one, two], axis=3, nested=True).type) == "2 * var * var * var * var * (int64, int64)") assert (str(ak.cartesian( [one, two], axis=-1, nested=True).type) == "2 * var * var * var * var * (int64, int64)") assert (str(ak.cartesian([one, two], axis=-2, nested=True).type) == "2 * var * var * var * (var * int64, var * int64)") assert (str(ak.cartesian([one, two], axis=-3, nested=True).type) == "2 * var * var * (var * var * int64, var * var * int64)") assert (str(ak.cartesian([one, two], axis=-4, nested=True).type) == "2 * 2 * (var * var * var * int64, var * var * var * int64)") with pytest.raises(ValueError): ak.cartesian([one, two], axis=-5, nested=True) assert (str(ak.cartesian([one, two], axis=0).type) == "4 * (var * var * var * int64, var * var * var * int64)") assert (str(ak.cartesian( [one, two], axis=1).type) == "2 * var * (var * var * int64, var * var * int64)") assert (str(ak.cartesian( [one, two], axis=2).type) == "2 * var * var * (var * int64, var * int64)") assert (str(ak.cartesian( [one, two], axis=3).type) == "2 * var * var * var * (int64, int64)") assert (str(ak.cartesian( [one, two], axis=-1).type) == "2 * var * var * var * (int64, int64)") assert (str(ak.cartesian( [one, two], axis=-2).type) == "2 * var * var * (var * int64, var * int64)") assert (str(ak.cartesian( [one, two], axis=-3).type) == "2 * var * (var * var * int64, var * var * int64)") assert (str(ak.cartesian([one, two], axis=-4).type) == "4 * (var * var * var * int64, var * var * var * int64)") with pytest.raises(ValueError): ak.cartesian([one, two], axis=-5)
def test_axis1(): one = ak.Array([[0, 1, 2], [], [3, 4]]) two = ak.Array([[100, 200], [300], [400, 500]]) three = ak.Array([["a", "b"], ["c", "d"], ["e"]]) assert ak.to_list(ak.cartesian([one])) == [[(0, ), (1, ), (2, )], [], [(3, ), (4, )]] assert ak.to_list(ak.cartesian({"x": one})) == [ [{ "x": 0 }, { "x": 1 }, { "x": 2 }], [], [{ "x": 3 }, { "x": 4 }], ] assert ak.to_list(ak.cartesian([one, two])) == [ [(0, 100), (0, 200), (1, 100), (1, 200), (2, 100), (2, 200)], [], [(3, 400), (3, 500), (4, 400), (4, 500)], ] if not ak._util.py27 and not ak._util.py35: assert ak.to_list(ak.cartesian({ "x": one, "y": two })) == [ [ { "x": 0, "y": 100 }, { "x": 0, "y": 200 }, { "x": 1, "y": 100 }, { "x": 1, "y": 200 }, { "x": 2, "y": 100 }, { "x": 2, "y": 200 }, ], [], [ { "x": 3, "y": 400 }, { "x": 3, "y": 500 }, { "x": 4, "y": 400 }, { "x": 4, "y": 500 }, ], ] assert ak.to_list(ak.cartesian([one, two, three])) == [ [ (0, 100, "a"), (0, 100, "b"), (0, 200, "a"), (0, 200, "b"), (1, 100, "a"), (1, 100, "b"), (1, 200, "a"), (1, 200, "b"), (2, 100, "a"), (2, 100, "b"), (2, 200, "a"), (2, 200, "b"), ], [], [(3, 400, "e"), (3, 500, "e"), (4, 400, "e"), (4, 500, "e")], ] assert ak.to_list(ak.cartesian([one, two, three], nested=[0])) == [ [ [(0, 100, "a"), (0, 100, "b"), (0, 200, "a"), (0, 200, "b")], [(1, 100, "a"), (1, 100, "b"), (1, 200, "a"), (1, 200, "b")], [(2, 100, "a"), (2, 100, "b"), (2, 200, "a"), (2, 200, "b")], ], [], [[(3, 400, "e"), (3, 500, "e")], [(4, 400, "e"), (4, 500, "e")]], ] assert ak.to_list(ak.cartesian([one, two, three], nested=[1])) == [ [ [(0, 100, "a"), (0, 100, "b")], [(0, 200, "a"), (0, 200, "b")], [(1, 100, "a"), (1, 100, "b")], [(1, 200, "a"), (1, 200, "b")], [(2, 100, "a"), (2, 100, "b")], [(2, 200, "a"), (2, 200, "b")], ], [], [[(3, 400, "e")], [(3, 500, "e")], [(4, 400, "e")], [(4, 500, "e")]], ] assert ak.to_list(ak.cartesian([one, two, three], nested=[0, 1])) == [ [ [[(0, 100, "a"), (0, 100, "b")], [(0, 200, "a"), (0, 200, "b")]], [[(1, 100, "a"), (1, 100, "b")], [(1, 200, "a"), (1, 200, "b")]], [[(2, 100, "a"), (2, 100, "b")], [(2, 200, "a"), (2, 200, "b")]], ], [], [[[(3, 400, "e")], [(3, 500, "e")]], [[(4, 400, "e")], [(4, 500, "e")]]], ] assert ak.to_list(ak.cartesian([one, two, three], nested=[])) == ak.to_list( ak.cartesian([one, two, three], nested=False)) assert ak.to_list(ak.cartesian([one, two, three], nested=[])) == ak.to_list( ak.cartesian([one, two, three], nested=None)) assert ak.to_list(ak.cartesian([one, two, three], nested=[0, 1])) == ak.to_list( ak.cartesian([one, two, three], nested=True))
def process(self, events): # Dataset parameters dataset = events.metadata['dataset'] year = self._samples[dataset]['year'] xsec = self._samples[dataset]['xsec'] sow = self._samples[dataset]['nSumOfWeights'] isData = self._samples[dataset]['isData'] datasets = [ 'SingleMuon', 'SingleElectron', 'EGamma', 'MuonEG', 'DoubleMuon', 'DoubleElectron' ] for d in datasets: if d in dataset: dataset = dataset.split('_')[0] # Initialize objects met = events.MET e = events.Electron mu = events.Muon tau = events.Tau j = events.Jet # Muon selection #mu['isGood'] = isMuonMVA(mu.pt, mu.eta, mu.dxy, mu.dz, mu.miniPFRelIso_all, mu.sip3d, mu.mvaTTH, mu.mediumPromptId, mu.tightCharge, minpt=10) mu['isPres'] = isPresMuon(mu.dxy, mu.dz, mu.sip3d, mu.looseId) mu['isTight'] = isTightMuon(mu.pt, mu.eta, mu.dxy, mu.dz, mu.pfRelIso03_all, mu.sip3d, mu.mvaTTH, mu.mediumPromptId, mu.tightCharge, mu.looseId, minpt=10) mu['isGood'] = mu['isPres'] & mu['isTight'] leading_mu = mu[ak.argmax(mu.pt, axis=-1, keepdims=True)] leading_mu = leading_mu[leading_mu.isGood] mu = mu[mu.isGood] mu_pres = mu[mu.isPres] # Electron selection #e['isGood'] = isElecMVA(e.pt, e.eta, e.dxy, e.dz, e.miniPFRelIso_all, e.sip3d, e.mvaTTH, e.mvaFall17V2Iso, e.lostHits, e.convVeto, e.tightCharge, minpt=10) e['isPres'] = isPresElec(e.pt, e.eta, e.dxy, e.dz, e.miniPFRelIso_all, e.sip3d, e.lostHits, minpt=15) e['isTight'] = isTightElec(e.pt, e.eta, e.dxy, e.dz, e.miniPFRelIso_all, e.sip3d, e.mvaTTH, e.mvaFall17V2Iso, e.lostHits, e.convVeto, e.tightCharge, e.sieie, e.hoe, e.eInvMinusPInv, minpt=15) e['isClean'] = isClean(e, mu, drmin=0.05) e['isGood'] = e['isPres'] & e['isTight'] & e['isClean'] leading_e = e[ak.argmax(e.pt, axis=-1, keepdims=True)] leading_e = leading_e[leading_e.isGood] e = e[e.isGood] e_pres = e[e.isPres & e.isClean] # Tau selection tau['isPres'] = isPresTau(tau.pt, tau.eta, tau.dxy, tau.dz, tau.leadTkPtOverTauPt, tau.idAntiMu, tau.idAntiEle, tau.rawIso, tau.idDecayModeNewDMs, minpt=20) tau['isClean'] = isClean(tau, e_pres, drmin=0.4) & isClean( tau, mu_pres, drmin=0.4) tau['isGood'] = tau['isPres'] # & tau['isClean'], for the moment tau = tau[tau.isGood] nElec = ak.num(e) nMuon = ak.num(mu) nTau = ak.num(tau) twoLeps = (nElec + nMuon) == 2 threeLeps = (nElec + nMuon) == 3 twoElec = (nElec == 2) twoMuon = (nMuon == 2) e0 = e[ak.argmax(e.pt, axis=-1, keepdims=True)] m0 = mu[ak.argmax(mu.pt, axis=-1, keepdims=True)] # Jet selection jetptname = 'pt_nom' if hasattr(j, 'pt_nom') else 'pt' j['isGood'] = isTightJet(getattr(j, jetptname), j.eta, j.jetId, j.neHEF, j.neEmEF, j.chHEF, j.chEmEF, j.nConstituents) #j['isgood'] = isGoodJet(j.pt, j.eta, j.jetId) #j['isclean'] = isClean(j, e, mu) j['isClean'] = isClean(j, e, drmin=0.4) & isClean( j, mu, drmin=0.4) # & isClean(j, tau, drmin=0.4) goodJets = j[(j.isClean) & (j.isGood)] njets = ak.num(goodJets) ht = ak.sum(goodJets.pt, axis=-1) j0 = goodJets[ak.argmax(goodJets.pt, axis=-1, keepdims=True)] #nbtags = ak.num(goodJets[goodJets.btagDeepFlavB > 0.2770]) nbtags = ak.num(goodJets[goodJets.btagDeepB > 0.4941]) ################################################################## ### 2 same-sign leptons ################################################################## # emu singe = e[(nElec == 1) & (nMuon == 1) & (e.pt > -1)] singm = mu[(nElec == 1) & (nMuon == 1) & (mu.pt > -1)] em = ak.cartesian({"e": singe, "m": singm}) emSSmask = (em.e.charge * em.m.charge > 0) emSS = em[emSSmask] nemSS = len(ak.flatten(emSS)) # ee and mumu # pt>-1 to preserve jagged dimensions ee = e[(nElec == 2) & (nMuon == 0) & (e.pt > -1)] mm = mu[(nElec == 0) & (nMuon == 2) & (mu.pt > -1)] eepairs = ak.combinations(ee, 2, fields=["e0", "e1"]) eeSSmask = (eepairs.e0.charge * eepairs.e1.charge > 0) eeonZmask = (np.abs((eepairs.e0 + eepairs.e1).mass - 91.2) < 10) eeoffZmask = (eeonZmask == 0) mmpairs = ak.combinations(mm, 2, fields=["m0", "m1"]) mmSSmask = (mmpairs.m0.charge * mmpairs.m1.charge > 0) mmonZmask = (np.abs((mmpairs.m0 + mmpairs.m1).mass - 91.2) < 10) mmoffZmask = (mmonZmask == 0) eeSSonZ = eepairs[eeSSmask & eeonZmask] eeSSoffZ = eepairs[eeSSmask & eeoffZmask] mmSSonZ = mmpairs[mmSSmask & mmonZmask] mmSSoffZ = mmpairs[mmSSmask & mmoffZmask] neeSS = len(ak.flatten(eeSSonZ)) + len(ak.flatten(eeSSoffZ)) nmmSS = len(ak.flatten(mmSSonZ)) + len(ak.flatten(mmSSoffZ)) print('Same-sign events [ee, emu, mumu] = [%i, %i, %i]' % (neeSS, nemSS, nmmSS)) # Cuts eeSSmask = (ak.num(eeSSmask[eeSSmask]) > 0) mmSSmask = (ak.num(mmSSmask[mmSSmask]) > 0) eeonZmask = (ak.num(eeonZmask[eeonZmask]) > 0) eeoffZmask = (ak.num(eeoffZmask[eeoffZmask]) > 0) mmonZmask = (ak.num(mmonZmask[mmonZmask]) > 0) mmoffZmask = (ak.num(mmoffZmask[mmoffZmask]) > 0) emSSmask = (ak.num(emSSmask[emSSmask]) > 0) ################################################################## ### 3 leptons ################################################################## # eem muon_eem = mu[(nElec == 2) & (nMuon == 1) & (mu.pt > -1)] elec_eem = e[(nElec == 2) & (nMuon == 1) & (e.pt > -1)] ee_eem = ak.combinations(elec_eem, 2, fields=["e0", "e1"]) ee_eemZmask = (ee_eem.e0.charge * ee_eem.e1.charge < 1) & (np.abs( (ee_eem.e0 + ee_eem.e1).mass - 91.2) < 10) ee_eemOffZmask = (ee_eem.e0.charge * ee_eem.e1.charge < 1) & (np.abs( (ee_eem.e0 + ee_eem.e1).mass - 91.2) > 10) ee_eemZmask = (ak.num(ee_eemZmask[ee_eemZmask]) > 0) ee_eemOffZmask = (ak.num(ee_eemOffZmask[ee_eemOffZmask]) > 0) eepair_eem = (ee_eem.e0 + ee_eem.e1) trilep_eem = eepair_eem + muon_eem #ak.cartesian({"e0":ee_eem.e0,"e1":ee_eem.e1, "m":muon_eem}) # mme muon_mme = mu[(nElec == 1) & (nMuon == 2) & (mu.pt > -1)] elec_mme = e[(nElec == 1) & (nMuon == 2) & (e.pt > -1)] mm_mme = ak.combinations(muon_mme, 2, fields=["m0", "m1"]) mm_mmeZmask = (mm_mme.m0.charge * mm_mme.m1.charge < 1) & (np.abs( (mm_mme.m0 + mm_mme.m1).mass - 91.2) < 10) mm_mmeOffZmask = (mm_mme.m0.charge * mm_mme.m1.charge < 1) & (np.abs( (mm_mme.m0 + mm_mme.m1).mass - 91.2) > 10) mm_mmeZmask = (ak.num(mm_mmeZmask[mm_mmeZmask]) > 0) mm_mmeOffZmask = (ak.num(mm_mmeOffZmask[mm_mmeOffZmask]) > 0) mmpair_mme = (mm_mme.m0 + mm_mme.m1) trilep_mme = mmpair_mme + elec_mme mZ_mme = mmpair_mme.mass mZ_eem = eepair_eem.mass m3l_eem = trilep_eem.mass m3l_mme = trilep_mme.mass # eee and mmm eee = e[(nElec == 3) & (nMuon == 0) & (e.pt > -1)] mmm = mu[(nElec == 0) & (nMuon == 3) & (mu.pt > -1)] eee_leps = ak.combinations(eee, 3, fields=["e0", "e1", "e2"]) mmm_leps = ak.combinations(mmm, 3, fields=["m0", "m1", "m2"]) ee_pairs = ak.combinations(eee, 2, fields=["e0", "e1"]) mm_pairs = ak.combinations(mmm, 2, fields=["m0", "m1"]) ee_pairs_index = ak.argcombinations(eee, 2, fields=["e0", "e1"]) mm_pairs_index = ak.argcombinations(mmm, 2, fields=["m0", "m1"]) mmSFOS_pairs = mm_pairs[ (np.abs(mm_pairs.m0.pdgId) == np.abs(mm_pairs.m1.pdgId)) & (mm_pairs.m0.charge != mm_pairs.m1.charge)] offZmask_mm = ak.all( np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2) > 10., axis=1, keepdims=True) & (ak.num(mmSFOS_pairs) > 0) onZmask_mm = ak.any( np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2) < 10., axis=1, keepdims=True) eeSFOS_pairs = ee_pairs[ (np.abs(ee_pairs.e0.pdgId) == np.abs(ee_pairs.e1.pdgId)) & (ee_pairs.e0.charge != ee_pairs.e1.charge)] offZmask_ee = ak.all( np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2) > 10, axis=1, keepdims=True) & (ak.num(eeSFOS_pairs) > 0) onZmask_ee = ak.any( np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2) < 10, axis=1, keepdims=True) # Create masks **for event selection** eeeOnZmask = (ak.num(onZmask_ee[onZmask_ee]) > 0) eeeOffZmask = (ak.num(offZmask_ee[offZmask_ee]) > 0) mmmOnZmask = (ak.num(onZmask_mm[onZmask_mm]) > 0) mmmOffZmask = (ak.num(offZmask_mm[offZmask_mm]) > 0) # Now we need to create masks for the leptons in order to select leptons from the Z boson candidate (in onZ categories) ZeeMask = ak.argmin(np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2), axis=1, keepdims=True) ZmmMask = ak.argmin(np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2), axis=1, keepdims=True) Zee = eeSFOS_pairs[ZeeMask] Zmm = mmSFOS_pairs[ZmmMask] eZ0 = Zee.e0[ak.num(eeSFOS_pairs) > 0] eZ1 = Zee.e1[ak.num(eeSFOS_pairs) > 0] eZ = eZ0 + eZ1 mZ0 = Zmm.m0[ak.num(mmSFOS_pairs) > 0] mZ1 = Zmm.m1[ak.num(mmSFOS_pairs) > 0] mZ = mZ0 + mZ1 mZ_eee = eZ.mass mZ_mmm = mZ.mass # And for the W boson ZmmIndices = mm_pairs_index[ZmmMask] ZeeIndices = ee_pairs_index[ZeeMask] eW = eee[~ZeeIndices.e0 | ~ZeeIndices.e1] mW = mmm[~ZmmIndices.m0 | ~ZmmIndices.m1] triElec = eee_leps.e0 + eee_leps.e1 + eee_leps.e2 triMuon = mmm_leps.m0 + mmm_leps.m1 + mmm_leps.m2 m3l_eee = triElec.mass m3l_mmm = triMuon.mass # Triggers trig_eeSS = passTrigger(events, 'ee', isData, dataset) trig_mmSS = passTrigger(events, 'mm', isData, dataset) trig_emSS = passTrigger(events, 'em', isData, dataset) trig_eee = passTrigger(events, 'eee', isData, dataset) trig_mmm = passTrigger(events, 'mmm', isData, dataset) trig_eem = passTrigger(events, 'eem', isData, dataset) trig_mme = passTrigger(events, 'mme', isData, dataset) # MET filters # Weights genw = np.ones_like( events['MET_pt']) if isData else events['genWeight'] weights = coffea.analysis_tools.Weights(len(events)) weights.add('norm', genw if isData else (xsec / sow) * genw) eftweights = events['EFTfitCoefficients'] if hasattr( events, "EFTfitCoefficients") else [] # Selections and cuts selections = PackedSelection() channels2LSS = ['eeSSonZ', 'eeSSoffZ', 'mmSSonZ', 'mmSSoffZ', 'emSS'] selections.add('eeSSonZ', (eeonZmask) & (eeSSmask) & (trig_eeSS)) selections.add('eeSSoffZ', (eeoffZmask) & (eeSSmask) & (trig_eeSS)) selections.add('mmSSonZ', (mmonZmask) & (mmSSmask) & (trig_mmSS)) selections.add('mmSSoffZ', (mmoffZmask) & (mmSSmask) & (trig_mmSS)) selections.add('emSS', (emSSmask) & (trig_emSS)) channels3L = ['eemSSonZ', 'eemSSoffZ', 'mmeSSonZ', 'mmeSSoffZ'] selections.add('eemSSonZ', (ee_eemZmask) & (trig_eem)) selections.add('eemSSoffZ', (ee_eemOffZmask) & (trig_eem)) selections.add('mmeSSonZ', (mm_mmeZmask) & (trig_mme)) selections.add('mmeSSoffZ', (mm_mmeOffZmask) & (trig_mme)) channels3L += ['eeeSSonZ', 'eeeSSoffZ', 'mmmSSonZ', 'mmmSSoffZ'] selections.add('eeeSSonZ', (eeeOnZmask) & (trig_eee)) selections.add('eeeSSoffZ', (eeeOffZmask) & (trig_eee)) selections.add('mmmSSonZ', (mmmOnZmask) & (trig_mmm)) selections.add('mmmSSoffZ', (mmmOffZmask) & (trig_mmm)) levels = ['base', '2jets', '4jets', '4j1b', '4j2b'] selections.add('base', (nElec + nMuon >= 2)) selections.add('2jets', (njets >= 2)) selections.add('4jets', (njets >= 4)) selections.add('4j1b', (njets >= 4) & (nbtags >= 1)) selections.add('4j2b', (njets >= 4) & (nbtags >= 2)) # Variables invMass_eeSSonZ = (eeSSonZ.e0 + eeSSonZ.e1).mass invMass_eeSSoffZ = (eeSSoffZ.e0 + eeSSoffZ.e1).mass invMass_mmSSonZ = (mmSSonZ.m0 + mmSSonZ.m1).mass invMass_mmSSoffZ = (mmSSoffZ.m0 + mmSSoffZ.m1).mass invMass_emSS = (emSS.e + emSS.m).mass varnames = {} varnames['met'] = met.pt varnames['ht'] = ht varnames['njets'] = njets varnames['nbtags'] = nbtags varnames['invmass'] = { 'eeSSonZ': invMass_eeSSonZ, 'eeSSoffZ': invMass_eeSSoffZ, 'mmSSonZ': invMass_mmSSonZ, 'mmSSoffZ': invMass_mmSSoffZ, 'emSS': invMass_emSS, 'eemSSonZ': mZ_eem, 'eemSSoffZ': mZ_eem, 'mmeSSonZ': mZ_mme, 'mmeSSoffZ': mZ_mme, 'eeeSSonZ': mZ_eee, 'eeeSSoffZ': mZ_eee, 'mmmSSonZ': mZ_mmm, 'mmmSSoffZ': mZ_mmm, } varnames['m3l'] = { 'eemSSonZ': m3l_eem, 'eemSSoffZ': m3l_eem, 'mmeSSonZ': m3l_mme, 'mmeSSoffZ': m3l_mme, 'eeeSSonZ': m3l_eee, 'eeeSSoffZ': m3l_eee, 'mmmSSonZ': m3l_mmm, 'mmmSSoffZ': m3l_mmm, } varnames['e0pt'] = e0.pt varnames['e0eta'] = e0.eta varnames['m0pt'] = m0.pt varnames['m0eta'] = m0.eta varnames['j0pt'] = j0.pt varnames['j0eta'] = j0.eta varnames['counts'] = np.ones_like(events.MET.pt) # fill Histos hout = self.accumulator.identity() allweights = weights.weight().flatten( ) # Why does it not complain about .flatten() here? hout['SumOfEFTweights'].fill(eftweights, sample=dataset, SumOfEFTweights=varnames['counts'], weight=allweights) for var, v in varnames.items(): for ch in channels2LSS + channels3L: for lev in levels: weight = weights.weight() cuts = [ch] + [lev] cut = selections.all(*cuts) weights_flat = weight[cut].flatten( ) # Why does it not complain about .flatten() here? weights_ones = np.ones_like(weights_flat, dtype=np.int) eftweightsvalues = eftweights[cut] if len( eftweights) > 0 else [] if var == 'invmass': if ch in ['eeeSSoffZ', 'mmmSSoffZ']: continue elif ch in ['eeeSSonZ', 'mmmSSonZ']: continue #values = v[ch] else: values = ak.flatten(v[ch][cut]) hout['invmass'].fill(sample=dataset, channel=ch, cut=lev, invmass=values, weight=weights_flat) elif var == 'm3l': if ch in [ 'eeSSonZ', 'eeSSoffZ', 'mmSSonZ', 'mmSSoffZ', 'emSS', 'eeeSSoffZ', 'mmmSSoffZ', 'eeeSSonZ', 'mmmSSonZ' ]: continue values = ak.flatten(v[ch][cut]) hout['m3l'].fill(eftweightsvalues, sample=dataset, channel=ch, cut=lev, m3l=values, weight=weights_flat) else: values = v[cut] if var == 'ht': hout[var].fill(eftweightsvalues, ht=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) elif var == 'met': hout[var].fill(eftweightsvalues, met=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) elif var == 'njets': hout[var].fill(eftweightsvalues, njets=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) elif var == 'nbtags': hout[var].fill(eftweightsvalues, nbtags=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) elif var == 'counts': hout[var].fill(counts=values, sample=dataset, channel=ch, cut=lev, weight=weights_ones) elif var == 'j0eta': if lev == 'base': continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(eftweightsvalues, j0eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) elif var == 'e0pt': if ch in [ 'mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eftweightsvalues, e0pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat ) # Crashing here, not sure why. Related to values? elif var == 'm0pt': if ch in [ 'eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(eftweightsvalues, m0pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) elif var == 'e0eta': if ch in [ 'mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(eftweightsvalues, e0eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) elif var == 'm0eta': if ch in [ 'eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(eftweightsvalues, m0eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) elif var == 'j0pt': if lev == 'base': continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill(eftweightsvalues, j0pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat) return hout
rp = events["RP"] #Pions pi_cut = abs(rp["mass"] - lp.pi_plus.mass/1000.) < 1e-4 pions = rp[pi_cut] p_cut = pions["p"] > 1. pions = pions[p_cut] #Kaons k_cut = abs(rp["mass"] - lp.K_plus.mass/1000.) < 1e-4 kaons = rp[k_cut] p_cut = kaons["p"] > 1. kaons = kaons[p_cut] Kst = ak.cartesian({"k": kaons, "pi": pions}) Kst_cut = np.sign(Kst["k","charge"]) != np.sign(Kst["pi","charge"]) Kst = Kst[Kst_cut] PDG_K_m = lp.K_plus.mass/1000. PDG_pi_m = lp.pi_plus.mass/1000. Kst["mass"] = kinematics_flat.mass([Kst["k"], Kst["pi"]], [PDG_K_m, PDG_pi_m]) PDG_Kst_m = lp.Kst_892_0.mass/1000. Kst_window = 0.1 Kst_cut = abs(Kst["mass"] - PDG_Kst_m) < Kst_window Kst = Kst[Kst_cut] #Truth-matching Kst_parent_cut_pi = abs(Kst["pi"]["pdg"]) == abs(lp.Kst_892_0.pdgid) Kst_parent_cut_k = abs(Kst["k"]["pdg"]) == abs(lp.Kst_892_0.pdgid)