def process(self, df): ## objects used for cuts inpObj_noCut = ob.inpObj(df, self.scaleFactor) # Our preselection cuts = bl.cutList(df, inpObj_noCut) if self.setupNPArr is None: self.setupNPArray(cuts, variables) output = self.accumulator.identity() # run cut loop for name, cut in cuts.items(): # defining objects inpObj = {} for key, item in inpObj_noCut.items(): inpObj[key] = item[cut] if len(inpObj['evtw']) > 0: if name == "_npz": varValDict = utl.varGetter(inpObj) for varName, varDetail in varValDict.items(): if variables[varName][4] == 1: output['{}'.format(varName)] += col_accumulator( varDetail[0]) elif variables[varName][4] == 2: output['{}'.format(varName)] += col_accumulator( np.repeat( ak.to_numpy(varDetail[0]), ak.to_numpy(varValDict["njetsAK8"][0]))) return output
def test_numpyarray(): for dtype1 in ("i1", "i2", "i4", "i8", "u1", "u2", "u4", "u8", "f4", "f8", "?"): for dtype2 in ("i1", "i2", "i4", "i8", "u1", "u2", "u4", "u8", "f4", "f8", "?"): for dtype3 in ("i1", "i2", "i4", "i8", "u1", "u2", "u4", "u8", "f4", "f8", "?"): for dtype4 in ("i1", "i2", "i4", "i8", "u1", "u2", "u4", "u8", "f4", "f8", "?"): one = numpy.array([0, 1, 2], dtype=dtype1) two = numpy.array([3, 0], dtype=dtype2) three = numpy.array([], dtype=dtype3) four = numpy.array([4, 5, 0, 6, 7], dtype=dtype4) combined = numpy.concatenate([one, two, three, four]) ak_combined = awkward1.layout.NumpyArray(one).mergemany([ awkward1.layout.NumpyArray(two), awkward1.layout.NumpyArray(three), awkward1.layout.NumpyArray(four), ]) assert awkward1.to_list(ak_combined) == combined.tolist() assert awkward1.to_numpy( ak_combined).dtype == combined.dtype ak_combined = awkward1.layout.NumpyArray(one).mergemany([ awkward1.layout.NumpyArray(two), awkward1.layout.EmptyArray(), awkward1.layout.NumpyArray(four), ]) assert awkward1.to_list(ak_combined) == combined.tolist() assert awkward1.to_numpy( ak_combined).dtype == numpy.concatenate( [one, two, four]).dtype
def _transpose_column(data, fields): first_field = awkward1.to_numpy(data[fields[0]]) event_T = np.expand_dims(first_field, axis=1) for f in fields[1:]: column = _check_lorentz_vector(data[f]) column = awkward1.to_numpy(column) event_T = _append_column(event_T, column) return event_T
def yahist_2D_lookup(h, ar1, ar2): ''' takes a yahist 2D histogram (which has a lookup function) and an awkward array. ''' return ak.unflatten( h.lookup( ak.to_numpy(ak.flatten(ar1)), ak.to_numpy(ak.flatten(ar2)), ), ak.num(ar1) )
def test(): empty1 = awkward1.Array(awkward1.layout.EmptyArray(), check_valid=True) empty2 = awkward1.Array(awkward1.layout.ListOffsetArray64( awkward1.layout.Index64(numpy.array([0, 0, 0, 0], dtype=numpy.int64)), awkward1.layout.EmptyArray()), check_valid=True) array = awkward1.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]], check_valid=True) awkward1.to_numpy(empty1).dtype.type is numpy.float64 awkward1.to_list(array[empty1]) == [] awkward1.to_list(array[empty1, ]) == [] awkward1.to_list(array[empty2]) == [[], [], []] awkward1.to_list(array[empty2, ]) == [[], [], []]
def _ensure_flat(array, allow_missing=False): """Normalize an array to a flat numpy array or raise ValueError""" if isinstance(array, awkward.AwkwardArray): array = awkward1.from_awkward0(array) elif not isinstance(array, (awkward1.Array, numpy.ndarray)): raise ValueError("Expected a numpy or awkward array, received: %r" % array) aktype = awkward1.type(array) if not isinstance(aktype, awkward1.types.ArrayType): raise ValueError("Expected an array type, received: %r" % aktype) isprimitive = isinstance(aktype.type, awkward1.types.PrimitiveType) isoptionprimitive = isinstance( aktype.type, awkward1.types.OptionType) and isinstance( aktype.type.type, awkward1.types.PrimitiveType) if allow_missing and not (isprimitive or isoptionprimitive): raise ValueError( "Expected an array of type N * primitive or N * ?primitive, received: %r" % aktype) if not (allow_missing or isprimitive): raise ValueError( "Expected an array of type N * primitive, received: %r" % aktype) if isinstance(array, awkward1.Array): array = awkward1.to_numpy(array, allow_missing=allow_missing) return array
def zip_rle(output, dataset): return ak.to_numpy( ak.zip([ output['%s_run' % dataset].value.astype(int), output['%s_lumi' % dataset].value.astype(int), output['%s_event' % dataset].value.astype(int), ]))
def process(self, events): # Initialize accumulator out = self.accumulator.identity() # Event selection: opposite charged same flavor Electron = events.Electron Electron_mask = (Electron.pt > 20) & (np.abs(Electron.eta) < 2.5) & (Electron.cutBased > 1) Ele_channel_mask = ak.num(Electron[Electron_mask]) > 1 Ele_channel_events = events[Ele_channel_mask] Ele = Ele_channel_events.Electron # All possible pairs of Electron in each event ele_pairs = ak.combinations(Ele, 2, axis=1) # TLorentz vector sum of ele_pairs ele_left, ele_right = ak.unzip(ele_pairs) diele = ele_left + ele_right diffsign_diele = diele[diele.charge == 0] leading_diffsign_diele = diffsign_diele[ak.argmax(diffsign_diele.pt, axis=1, keepdims=True)] #Mee = ak.flatten(leading_diffsign_diele.mass) # This makes type error ( primitive expected but ?float given ) Mee = ak.to_numpy(leading_diffsign_diele.mass) Mee = Mee.flatten() out.fill(dataset=events.metadata["dataset"], mass=Mee) return out
def doAwkwardLookup(h, ar): ''' takes a ya_hist histogram (which has a lookup function) and an awkward array. ''' return ak.unflatten( h.lookup( ak.to_numpy( ak.flatten(ar) ) ), ak.num(ar) )
def test_numpy_array(): assert numpy.array_equal(numpy.asarray(awkward1.Array([1.1, 2.2, 3.3, 4.4, 5.5], check_valid=True)), numpy.array([1.1, 2.2, 3.3, 4.4, 5.5])) assert numpy.array_equal(numpy.asarray(awkward1.Array(numpy.array([1.1, 2.2, 3.3, 4.4, 5.5]), check_valid=True)), numpy.array([1.1, 2.2, 3.3, 4.4, 5.5])) assert numpy.array_equal(numpy.asarray(awkward1.Array([[1.1, 2.2], [3.3, 4.4], [5.5, 6.6]], check_valid=True)), numpy.array([[1.1, 2.2], [3.3, 4.4], [5.5, 6.6]])) assert numpy.array_equal(numpy.asarray(awkward1.Array(numpy.array([[1.1, 2.2], [3.3, 4.4], [5.5, 6.6]]), check_valid=True)), numpy.array([[1.1, 2.2], [3.3, 4.4], [5.5, 6.6]])) assert numpy.array_equal(numpy.asarray(awkward1.Array(["one", "two", "three"], check_valid=True)), numpy.array(["one", "two", "three"])) assert numpy.array_equal(numpy.asarray(awkward1.Array([b"one", b"two", b"three"], check_valid=True)), numpy.array([b"one", b"two", b"three"])) assert numpy.array_equal(numpy.asarray(awkward1.Array([], check_valid=True)), numpy.array([])) content0 = awkward1.layout.NumpyArray(numpy.array([1.1, 2.2, 3.3, 4.4, 5.5], dtype=numpy.float64)) content1 = awkward1.layout.NumpyArray(numpy.array([1, 2, 3], dtype=numpy.int64)) tags = awkward1.layout.Index8(numpy.array([0, 1, 1, 0, 0, 0, 1, 0], dtype=numpy.int8)) index = awkward1.layout.Index64(numpy.array([0, 0, 1, 1, 2, 3, 2, 4], dtype=numpy.int64)) array = awkward1.Array(awkward1.layout.UnionArray8_64(tags, index, [content0, content1]), check_valid=True) assert numpy.array_equal(numpy.asarray(array), numpy.array([1.1, 1, 2, 2.2, 3.3, 4.4, 3, 5.5])) assert awkward1.to_numpy(awkward1.Array([1.1, 2.2, None, None, 3.3], check_valid=True)).tolist() == [1.1, 2.2, None, None, 3.3] assert awkward1.to_numpy(awkward1.Array([[1.1, 2.2], [None, None], [3.3, 4.4]], check_valid=True)).tolist() == [[1.1, 2.2], [None, None], [3.3, 4.4]] assert awkward1.to_numpy(awkward1.Array([[1.1, 2.2], None, [3.3, 4.4]], check_valid=True)).tolist() == [[1.1, 2.2], [None, None], [3.3, 4.4]] assert numpy.array_equal(numpy.asarray(awkward1.Array([{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}, {"x": 3, "y": 3.3}], check_valid=True)), numpy.array([(1, 1.1), (2, 2.2), (3, 3.3)], dtype=[("x", numpy.int64), ("y", numpy.float64)])) assert numpy.array_equal(numpy.asarray(awkward1.Array([(1, 1.1), (2, 2.2), (3, 3.3)], check_valid=True)), numpy.array([(1, 1.1), (2, 2.2), (3, 3.3)], dtype=[("0", numpy.int64), ("1", numpy.float64)]))
def expr_to_vals(expr): vals = eval(expr, dict(), loc) # if varexp is a simple constant, broadcast it to an array if _array_ndim(vals) == 0: vals = vals * np.ones(len(df)) if sel: if _array_ndim(vals) < _array_ndim(globalmask): vals, _ = awkward1.broadcast_arrays(vals, globalmask) vals = vals[globalmask] if _array_ndim(vals) > 1: vals = awkward1.flatten(vals) vals = awkward1.to_numpy(vals) return vals
def process(self, events): output = self.accumulator.identity() # use a very loose preselection to filter the events presel = ak.num(events.Jet) > 2 ev = events[presel] dataset = ev.metadata['dataset'] # load the config - probably not needed anymore cfg = loadConfig() output['totalEvents']['all'] += len(events) output['skimmedEvents']['all'] += len(ev) ## Muons muon = Collections(ev, "Muon", "tightSSTTH").get() vetomuon = Collections(ev, "Muon", "vetoTTH").get() dimuon = choose(muon, 2) SSmuon = ak.any((dimuon['0'].charge * dimuon['1'].charge) > 0, axis=1) OSmuon = ak.any((dimuon['0'].charge * dimuon['1'].charge) < 0, axis=1) leading_muon_idx = ak.singletons(ak.argmax(muon.pt, axis=1)) leading_muon = muon[leading_muon_idx] ## Electrons electron = Collections(ev, "Electron", "tightSSTTH").get() vetoelectron = Collections(ev, "Electron", "vetoTTH").get() dielectron = choose(electron, 2) SSelectron = ak.any( (dielectron['0'].charge * dielectron['1'].charge) > 0, axis=1) OSelectron = ak.any( (dielectron['0'].charge * dielectron['1'].charge) < 0, axis=1) leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1)) leading_electron = electron[leading_electron_idx] ## Merge electrons and muons - this should work better now in ak1 lepton = ak.concatenate([muon, electron], axis=1) dilepton = cross(muon, electron) SSlepton = ak.any((dilepton['0'].charge * dilepton['1'].charge) > 0, axis=1) OSlepton = ak.any((dilepton['0'].charge * dilepton['1'].charge) < 0, axis=1) leading_lepton_idx = ak.singletons(ak.argmax(lepton.pt, axis=1)) leading_lepton = lepton[leading_lepton_idx] trailing_lepton_idx = ak.singletons(ak.argmin(lepton.pt, axis=1)) trailing_lepton = lepton[trailing_lepton_idx] second_lepton = lepton[~(trailing_lepton_idx & leading_lepton_idx)] ## Jets jet = getJets(ev, minPt=25, maxEta=4.7, pt_var='pt_nom') jet = jet[ak.argsort( jet.pt_nom, ascending=False )] # need to sort wrt smeared and recorrected jet pt jet = jet[~match(jet, muon, deltaRCut=0.4)] # remove jets that overlap with muons jet = jet[~match( jet, electron, deltaRCut=0.4)] # remove jets that overlap with electrons central = jet[(abs(jet.eta) < 2.4)] btag = getBTagsDeepFlavB( jet, year=self.year) # should study working point for DeepJet light = getBTagsDeepFlavB(jet, year=self.year, invert=True) fwd = getFwdJet(light) fwd_noPU = getFwdJet(light, puId=False) ## forward jets high_p_fwd = fwd[ak.singletons(ak.argmax( fwd.p, axis=1))] # highest momentum spectator high_pt_fwd = fwd[ak.singletons(ak.argmax( fwd.pt_nom, axis=1))] # highest transverse momentum spectator high_eta_fwd = fwd[ak.singletons(ak.argmax(abs( fwd.eta), axis=1))] # most forward spectator ## Get the two leading b-jets in terms of btag score high_score_btag = central[ak.argsort(central.btagDeepFlavB)][:, :2] jf = cross(high_p_fwd, jet) mjf = (jf['0'] + jf['1']).mass deltaEta = abs(high_p_fwd.eta - jf[ak.singletons(ak.argmax(mjf, axis=1))]['1'].eta) deltaEtaMax = ak.max(deltaEta, axis=1) mjf_max = ak.max(mjf, axis=1) jj = choose(jet, 2) mjj_max = ak.max((jj['0'] + jj['1']).mass, axis=1) ## MET -> can switch to puppi MET met_pt = ev.MET.pt met_phi = ev.MET.phi ## other variables ht = ak.sum(jet.pt, axis=1) st = met_pt + ht + ak.sum(muon.pt, axis=1) + ak.sum(electron.pt, axis=1) lt = met_pt + ak.sum(muon.pt, axis=1) + ak.sum(electron.pt, axis=1) ht_central = ak.sum(central.pt, axis=1) # define the weight weight = Weights(len(ev)) if not re.search(re.compile('MuonEG|DoubleMuon|DoubleEG|EGamma'), dataset): # lumi weight weight.add("weight", ev.weight * cfg['lumi'][self.year]) # PU weight - not in the babies... weight.add("PU", ev.puWeight, weightUp=ev.puWeightUp, weightDown=ev.puWeightDown, shift=False) # b-tag SFs weight.add( "btag", self.btagSF.Method1a(btag, light, b_direction='central', c_direction='central')) # lepton SFs weight.add("lepton", self.leptonSF.get(electron, muon)) sel = Selection( dataset=dataset, events=ev, year=self.year, ele=electron, ele_veto=vetoelectron, mu=muon, mu_veto=vetomuon, jet_all=jet, jet_central=central, jet_btag=btag, jet_fwd=fwd, met=ev.MET, ) BL = sel.dilep_baseline(SS=False) BL_minusNb = sel.dilep_baseline(SS=False, omit=['N_btag>0']) output['N_b'].fill(dataset=dataset, multiplicity=ak.num(btag)[BL_minusNb], weight=weight.weight()[BL_minusNb]) if re.search(re.compile('MuonEG|DoubleMuon|DoubleEG|EGamma'), dataset): #rle = ak.to_numpy(ak.zip([ev.run, ev.luminosityBlock, ev.event])) run_ = ak.to_numpy(ev.run) lumi_ = ak.to_numpy(ev.luminosityBlock) event_ = ak.to_numpy(ev.event) output['%s_run' % dataset] += processor.column_accumulator( run_[BL]) output['%s_lumi' % dataset] += processor.column_accumulator( lumi_[BL]) output['%s_event' % dataset] += processor.column_accumulator( event_[BL]) # Now, take care of systematic unceratinties if not re.search(re.compile('MuonEG|DoubleMuon|DoubleEG|EGamma'), dataset): alljets = getJets(ev, minPt=0, maxEta=4.7) alljets = alljets[(alljets.jetId > 1)] for var in self.variations: # get the collections that change with the variations btag = getBTagsDeepFlavB( jet, year=self.year) # should study working point for DeepJet weight = Weights(len(ev)) weight.add("weight", ev.weight * cfg['lumi'][self.year]) weight.add("PU", ev.puWeight, weightUp=ev.puWeightUp, weightDown=ev.puWeightDown, shift=False) if var == 'centralUp': weight.add( "btag", self.btagSF.Method1a(btag, light, b_direction='central', c_direction='up')) elif var == 'centralDown': weight.add( "btag", self.btagSF.Method1a(btag, light, b_direction='central', c_direction='down')) elif var == 'upCentral': weight.add( "btag", self.btagSF.Method1a(btag, light, b_direction='up', c_direction='central')) elif var == 'downCentral': weight.add( "btag", self.btagSF.Method1a(btag, light, b_direction='down', c_direction='central')) weight.add("lepton", self.leptonSF.get(electron, muon)) met = ev.MET sel = Selection( dataset=dataset, events=ev, year=self.year, ele=electron, ele_veto=vetoelectron, mu=muon, mu_veto=vetomuon, jet_all=jet, jet_central=central, jet_btag=btag, jet_fwd=fwd, met=met, ) BL = sel.dilep_baseline(SS=False) BL_minusNb = sel.dilep_baseline(SS=False, omit=['N_btag>0']) output['N_b_' + var].fill( dataset=dataset, multiplicity=ak.num(btag)[BL_minusNb], weight=weight.weight()[BL_minusNb]) return output
def get_simplified_spec( spec: Dict[str, Any], ylds: yields.Yields, allowed_modifiers: List[str], prune_channels: List[str], include_signal: bool = False, ) -> pyhf.workspace: newspec = { 'channels': [{ 'name': channel['name'], 'samples': [{ 'name': 'Bkg', 'data': ylds.yields[channel['name']].sum(axis=0).flatten().tolist(), "modifiers": [{ "data": { "hi_data": (ylds.yields[channel['name']].sum(axis=0) + ak.to_numpy(ylds.uncertainties[channel['name']]) ).flatten().tolist(), "lo_data": (ylds.yields[channel['name']].sum(axis=0) - ak.to_numpy(ylds.uncertainties[channel['name']]) ).flatten().tolist(), }, "name": "totalError", "type": "histosys", }], }], } for channel in spec['channels'] if channel['name'] not in prune_channels], 'measurements': [ { 'name': measurement['name'], 'config': { 'parameters': [{ "auxdata": [1.0], "bounds": [[0.915, 1.085]], "fixed": True, # this is the important part "inits": [1.0], "name": "lumi", "sigmas": [0.017], }] + [ dict( parameter, name=parameter['name'], ) for parameter in measurement['config']['parameters'] if parameter['name'] in allowed_modifiers ], 'poi': 'mu_Sig', }, } for measurement in spec['measurements'] ], 'observations': [ dict( copy.deepcopy(observation), name=observation['name'], ) for observation in spec['observations'] ], 'version': spec['version'], } if include_signal: channels_with_signal = [{ 'name': c['name'], 'samples': c['samples'] + [{ "name": "Signal", "data": [0] * len(ylds.yields[c['name']].sum(axis=0).flatten().tolist()), "modifiers": [{ "data": None, "name": "mu_Sig", "type": "normfactor" }], }], } for c in newspec['channels']] newspec['channels'] = channels_with_signal return newspec
def process(self, events): output = self.accumulator.identity() # use a very loose preselection to filter the events presel = ak.num(events.Jet)>2 ev = events[presel] dataset = ev.metadata['dataset'] # load the config - probably not needed anymore cfg = loadConfig() output['totalEvents']['all'] += len(events) output['skimmedEvents']['all'] += len(ev) ## Muons muon = Collections(ev, "Muon", "tightSSTTH").get() vetomuon = Collections(ev, "Muon", "vetoTTH").get() dimuon = choose(muon, 2) SSmuon = ak.any((dimuon['0'].charge * dimuon['1'].charge)>0, axis=1) OSmuon = ak.any((dimuon['0'].charge * dimuon['1'].charge)<0, axis=1) leading_muon_idx = ak.singletons(ak.argmax(muon.pt, axis=1)) leading_muon = muon[leading_muon_idx] ## Electrons electron = Collections(ev, "Electron", "tightSSTTH").get() vetoelectron = Collections(ev, "Electron", "vetoTTH").get() dielectron = choose(electron, 2) SSelectron = ak.any((dielectron['0'].charge * dielectron['1'].charge)>0, axis=1) OSelectron = ak.any((dielectron['0'].charge * dielectron['1'].charge)<0, axis=1) leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1)) leading_electron = electron[leading_electron_idx] ## Merge electrons and muons - this should work better now in ak1 lepton = ak.concatenate([muon, electron], axis=1) dilepton = cross(muon, electron) SSlepton = ak.any((dilepton['0'].charge * dilepton['1'].charge)>0, axis=1) OSlepton = ak.any((dilepton['0'].charge * dilepton['1'].charge)<0, axis=1) leading_lepton_idx = ak.singletons(ak.argmax(lepton.pt, axis=1)) leading_lepton = lepton[leading_lepton_idx] trailing_lepton_idx = ak.singletons(ak.argmin(lepton.pt, axis=1)) trailing_lepton = lepton[trailing_lepton_idx] ## Jets jet = getJets(ev, minPt=25, maxEta=4.7, pt_var='pt_nom') jet = jet[ak.argsort(jet.pt_nom, ascending=False)] # need to sort wrt smeared and recorrected jet pt jet = jet[~match(jet, muon, deltaRCut=0.4)] # remove jets that overlap with muons jet = jet[~match(jet, electron, deltaRCut=0.4)] # remove jets that overlap with electrons central = jet[(abs(jet.eta)<2.4)] btag = getBTagsDeepFlavB(jet, year=self.year) # should study working point for DeepJet light = getBTagsDeepFlavB(jet, year=self.year, invert=True) fwd = getFwdJet(light) fwd_noPU = getFwdJet(light, puId=False) ## forward jets high_p_fwd = fwd[ak.singletons(ak.argmax(fwd.p, axis=1))] # highest momentum spectator high_pt_fwd = fwd[ak.singletons(ak.argmax(fwd.pt_nom, axis=1))] # highest transverse momentum spectator high_eta_fwd = fwd[ak.singletons(ak.argmax(abs(fwd.eta), axis=1))] # most forward spectator ## Get the two leading b-jets in terms of btag score high_score_btag = central[ak.argsort(central.btagDeepFlavB)][:,:2] jf = cross(high_p_fwd, jet) mjf = (jf['0']+jf['1']).mass deltaEta = abs(high_p_fwd.eta - jf[ak.singletons(ak.argmax(mjf, axis=1))]['1'].eta) deltaEtaMax = ak.max(deltaEta, axis=1) mjf_max = ak.max(mjf, axis=1) jj = choose(jet, 2) mjj_max = ak.max((jj['0']+jj['1']).mass, axis=1) ## MET -> can switch to puppi MET met_pt = ev.MET.pt met_phi = ev.MET.phi ## other variables ht = ak.sum(jet.pt, axis=1) st = met_pt + ht + ak.sum(muon.pt, axis=1) + ak.sum(electron.pt, axis=1) ht_central = ak.sum(central.pt, axis=1) # define the weight weight = Weights( len(ev) ) if not re.search(re.compile('MuonEG|DoubleMuon|DoubleEG|EGamma'), dataset): # lumi weight weight.add("weight", ev.weight*cfg['lumi'][self.year]) # PU weight - not in the babies... weight.add("PU", ev.puWeight, weightUp=ev.puWeightUp, weightDown=ev.puWeightDown, shift=False) # b-tag SFs weight.add("btag", self.btagSF.Method1a(btag, light)) # lepton SFs weight.add("lepton", self.leptonSF.get(electron, muon)) cutflow = Cutflow(output, ev, weight=weight) sel = Selection( dataset = dataset, events = ev, year = self.year, ele = electron, ele_veto = vetoelectron, mu = muon, mu_veto = vetomuon, jet_all = jet, jet_central = central, jet_btag = btag, jet_fwd = fwd, met = ev.MET, ) BL = sel.dilep_baseline(cutflow=cutflow, SS=False) # first, make a few super inclusive plots output['PV_npvs'].fill(dataset=dataset, multiplicity=ev.PV[BL].npvs, weight=weight.weight()[BL]) output['PV_npvsGood'].fill(dataset=dataset, multiplicity=ev.PV[BL].npvsGood, weight=weight.weight()[BL]) output['N_jet'].fill(dataset=dataset, multiplicity=ak.num(jet)[BL], weight=weight.weight()[BL]) BL_minusNb = sel.dilep_baseline(SS=False, omit=['N_btag>0']) output['N_b'].fill(dataset=dataset, multiplicity=ak.num(btag)[BL_minusNb], weight=weight.weight()[BL_minusNb]) output['N_central'].fill(dataset=dataset, multiplicity=ak.num(central)[BL], weight=weight.weight()[BL]) output['N_ele'].fill(dataset=dataset, multiplicity=ak.num(electron)[BL], weight=weight.weight()[BL]) output['N_mu'].fill(dataset=dataset, multiplicity=ak.num(electron)[BL], weight=weight.weight()[BL]) BL_minusFwd = sel.dilep_baseline(SS=False, omit=['N_fwd>0']) output['N_fwd'].fill(dataset=dataset, multiplicity=ak.num(fwd)[BL_minusFwd], weight=weight.weight()[BL_minusFwd]) BL_minusMET = sel.dilep_baseline(SS=False, omit=['MET>50']) output['MET'].fill( dataset = dataset, pt = ev.MET[BL_minusMET].pt, phi = ev.MET[BL_minusMET].phi, weight = weight.weight()[BL_minusMET] ) #output['electron'].fill( # dataset = dataset, # pt = ak.to_numpy(ak.flatten(electron[BL].pt)), # eta = ak.to_numpy(ak.flatten(electron[BL].eta)), # phi = ak.to_numpy(ak.flatten(electron[BL].phi)), # weight = weight.weight()[BL] #) # #output['muon'].fill( # dataset = dataset, # pt = ak.to_numpy(ak.flatten(muon[BL].pt)), # eta = ak.to_numpy(ak.flatten(muon[BL].eta)), # phi = ak.to_numpy(ak.flatten(muon[BL].phi)), # weight = weight.weight()[BL] #) output['lead_lep'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(leading_lepton[BL].pt)), eta = ak.to_numpy(ak.flatten(leading_lepton[BL].eta)), phi = ak.to_numpy(ak.flatten(leading_lepton[BL].phi)), weight = weight.weight()[BL] ) output['trail_lep'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(trailing_lepton[BL].pt)), eta = ak.to_numpy(ak.flatten(trailing_lepton[BL].eta)), phi = ak.to_numpy(ak.flatten(trailing_lepton[BL].phi)), weight = weight.weight()[BL] ) output['fwd_jet'].fill( dataset = dataset, pt = ak.flatten(high_p_fwd[BL].pt_nom), eta = ak.flatten(high_p_fwd[BL].eta), phi = ak.flatten(high_p_fwd[BL].phi), weight = weight.weight()[BL] ) output['b1'].fill( dataset = dataset, pt = ak.flatten(high_score_btag[:, 0:1][BL].pt_nom), eta = ak.flatten(high_score_btag[:, 0:1][BL].eta), phi = ak.flatten(high_score_btag[:, 0:1][BL].phi), weight = weight.weight()[BL] ) output['b2'].fill( dataset = dataset, pt = ak.flatten(high_score_btag[:, 1:2][BL].pt_nom), eta = ak.flatten(high_score_btag[:, 1:2][BL].eta), phi = ak.flatten(high_score_btag[:, 1:2][BL].phi), weight = weight.weight()[BL] ) output['j1'].fill( dataset = dataset, pt = ak.flatten(jet.pt_nom[:, 0:1][BL]), eta = ak.flatten(jet.eta[:, 0:1][BL]), phi = ak.flatten(jet.phi[:, 0:1][BL]), weight = weight.weight()[BL] ) output['j2'].fill( dataset = dataset, pt = ak.flatten(jet[:, 1:2][BL].pt_nom), eta = ak.flatten(jet[:, 1:2][BL].eta), phi = ak.flatten(jet[:, 1:2][BL].phi), weight = weight.weight()[BL] ) output['j3'].fill( dataset = dataset, pt = ak.flatten(jet[:, 2:3][BL].pt_nom), eta = ak.flatten(jet[:, 2:3][BL].eta), phi = ak.flatten(jet[:, 2:3][BL].phi), weight = weight.weight()[BL] ) if re.search(re.compile('MuonEG|DoubleMuon|DoubleEG|EGamma'), dataset): #rle = ak.to_numpy(ak.zip([ev.run, ev.luminosityBlock, ev.event])) run_ = ak.to_numpy(ev.run) lumi_ = ak.to_numpy(ev.luminosityBlock) event_ = ak.to_numpy(ev.event) output['%s_run'%dataset] += processor.column_accumulator(run_[BL]) output['%s_lumi'%dataset] += processor.column_accumulator(lumi_[BL]) output['%s_event'%dataset] += processor.column_accumulator(event_[BL]) # Now, take care of systematic unceratinties if not re.search(re.compile('MuonEG|DoubleMuon|DoubleEG|EGamma'), dataset): alljets = getJets(ev, minPt=0, maxEta=4.7) alljets = alljets[(alljets.jetId>1)] for var in self.variations: # get the collections that change with the variations jet = getPtEtaPhi(alljets, pt_var=var) jet = jet[(jet.pt>25)] jet = jet[~match(jet, muon, deltaRCut=0.4)] # remove jets that overlap with muons jet = jet[~match(jet, electron, deltaRCut=0.4)] # remove jets that overlap with electrons central = jet[(abs(jet.eta)<2.4)] btag = getBTagsDeepFlavB(jet, year=self.year) # should study working point for DeepJet light = getBTagsDeepFlavB(jet, year=self.year, invert=True) fwd = getFwdJet(light) fwd_noPU = getFwdJet(light, puId=False) ## forward jets high_p_fwd = fwd[ak.singletons(ak.argmax(fwd.p, axis=1))] # highest momentum spectator high_pt_fwd = fwd[ak.singletons(ak.argmax(fwd.pt, axis=1))] # highest transverse momentum spectator high_eta_fwd = fwd[ak.singletons(ak.argmax(abs(fwd.eta), axis=1))] # most forward spectator ## Get the two leading b-jets in terms of btag score high_score_btag = central[ak.argsort(central.btagDeepFlavB)][:,:2] met = ev.MET met['pt'] = getattr(met, var) sel = Selection( dataset = dataset, events = ev, year = self.year, ele = electron, ele_veto = vetoelectron, mu = muon, mu_veto = vetomuon, jet_all = jet, jet_central = central, jet_btag = btag, jet_fwd = fwd, met = met, ) BL = sel.dilep_baseline(SS=False) # get the modified selection -> more difficult #selection.add('N_jet>2_'+var, (ak.num(jet.pt)>=3)) # stupid bug here... #selection.add('N_btag=2_'+var, (ak.num(btag)==2) ) #selection.add('N_central>1_'+var, (ak.num(central)>=2) ) #selection.add('N_fwd>0_'+var, (ak.num(fwd)>=1) ) #selection.add('MET>30_'+var, (getattr(ev.MET, var)>30) ) ### Don't change the selection for now... #bl_reqs = os_reqs + ['N_jet>2_'+var, 'MET>30_'+var, 'N_btag=2_'+var, 'N_central>1_'+var, 'N_fwd>0_'+var] #bl_reqs_d = { sel: True for sel in bl_reqs } #BL = selection.require(**bl_reqs_d) # the OS selection remains unchanged output['N_jet_'+var].fill(dataset=dataset, multiplicity=ak.num(jet)[BL], weight=weight.weight()[BL]) BL_minusFwd = sel.dilep_baseline(SS=False, omit=['N_fwd>0']) output['N_fwd_'+var].fill(dataset=dataset, multiplicity=ak.num(fwd)[BL_minusFwd], weight=weight.weight()[BL_minusFwd]) BL_minusNb = sel.dilep_baseline(SS=False, omit=['N_btag>0']) output['N_b_'+var].fill(dataset=dataset, multiplicity=ak.num(btag)[BL_minusNb], weight=weight.weight()[BL_minusNb]) output['N_central_'+var].fill(dataset=dataset, multiplicity=ak.num(central)[BL], weight=weight.weight()[BL]) # We don't need to redo all plots with variations. E.g., just add uncertainties to the jet plots. output['j1_'+var].fill( dataset = dataset, pt = ak.flatten(jet.pt[:, 0:1][BL]), eta = ak.flatten(jet.eta[:, 0:1][BL]), phi = ak.flatten(jet.phi[:, 0:1][BL]), weight = weight.weight()[BL] ) output['b1_'+var].fill( dataset = dataset, pt = ak.flatten(high_score_btag[:, 0:1].pt[:, 0:1][BL]), eta = ak.flatten(high_score_btag[:, 0:1].eta[:, 0:1][BL]), phi = ak.flatten(high_score_btag[:, 0:1].phi[:, 0:1][BL]), weight = weight.weight()[BL] ) output['fwd_jet_'+var].fill( dataset = dataset, pt = ak.flatten(high_p_fwd[BL].pt), #p = ak.flatten(high_p_fwd[BL].p), eta = ak.flatten(high_p_fwd[BL].eta), phi = ak.flatten(high_p_fwd[BL].phi), weight = weight.weight()[BL] ) BL_minusMET = sel.dilep_baseline(SS=False, omit=['MET>50']) output['MET_'+var].fill( dataset = dataset, pt = getattr(ev.MET, var)[BL_minusMET], phi = ev.MET[BL_minusMET].phi, weight = weight.weight()[BL_minusMET] ) return output
def test_allow_missing(): array = awkward1.Array([1.1, 2.2, None, 3.3, None, None, 4.4, 5.5]) awkward1.to_numpy(array) with pytest.raises(ValueError): awkward1.to_numpy(array, allow_missing=False)
# acceptance calculation full_hh4b_samples = uproot4.concatenate("data/hh4b/nano_*.root:Events") genpart_pt = full_hh4b_samples['GenPart_pt'] genpart_status = full_hh4b_samples['GenPart_status'] full_hh4b_samples['GenPart_status'] full_hh4b_samples['GenPart_statusFlags'] higgses = (full_hh4b_samples['GenPart_pdgId'] == 25) full_hh4b_samples['GenPart_status'][higgses] np.unique( ak.to_numpy( ak.pad_none(full_hh4b_samples['GenPart_status'][higgses], 40, axis=1))) ak.pad_none(hh4b_samples['GenPart_status'][higgses], 40, axis=1) higgs_pt = genpart_pt[(full_hh4b_samples['GenPart_pdgId'] == 25)] gt_300 = ak.sort(higgs_pt, axis=1)[:, -1] > 300 num_gt_300 = ak.sum(ak.sort(higgs_pt, axis=1)[:, -1] > 300) num_gt_300 / len(higgs_pt) fhiggs_pt = genpart_pt[(full_hh4b_samples['GenPart_pdgId'] == 25 )][full_hh4b_samples['GenPart_status'][higgses] == 22] jet_pt = ak.pad_none(full_hh4b_samples['FatJet_pt'], 2, axis=1)[:, :2][gt_300] jet_msd = ak.pad_none(full_hh4b_samples['FatJet_msoftdrop'], 2, axis=1)[:, :2][gt_300]
def process(self, events): output = self.accumulator.identity() # we can use a very loose preselection to filter the events. nothing is done with this presel, though presel = ak.num(events.Jet)>=2 ev = events[presel] dataset = ev.metadata['dataset'] # load the config - probably not needed anymore cfg = loadConfig() output['totalEvents']['all'] += len(events) output['skimmedEvents']['all'] += len(ev) ## Electrons electron = Collections(ev, "Electron", "tightFCNC", 0, self.year).get() electron = electron[(electron.pt > 15) & (np.abs(electron.eta) < 2.4)] electron = electron[(electron.genPartIdx >= 0)] electron = electron[(np.abs(electron.matched_gen.pdgId)==11)] #from here on all leptons are gen-matched electron = electron[( (electron.genPartFlav==1) | (electron.genPartFlav==15) )] #and now they are all prompt leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1)) leading_electron = electron[leading_electron_idx] trailing_electron_idx = ak.singletons(ak.argmin(electron.pt, axis=1)) trailing_electron = electron[trailing_electron_idx] leading_parent = find_first_parent(leading_electron.matched_gen) trailing_parent = find_first_parent(trailing_electron.matched_gen) is_flipped = ( ( (electron.matched_gen.pdgId*(-1) == electron.pdgId) | (find_first_parent(electron.matched_gen)*(-1) == electron.pdgId) ) & (np.abs(electron.pdgId) == 11) ) flipped_electron = electron[is_flipped] flipped_electron = flipped_electron[(ak.fill_none(flipped_electron.pt, 0)>0)] flipped_electron = flipped_electron[~(ak.is_none(flipped_electron))] n_flips = ak.num(flipped_electron) ##Muons muon = Collections(ev, "Muon", "tightFCNC").get() muon = muon[(muon.pt > 15) & (np.abs(muon.eta) < 2.4)] muon = muon[(muon.genPartIdx >= 0)] muon = muon[(np.abs(muon.matched_gen.pdgId)==13)] #from here, all muons are gen-matched muon = muon[( (muon.genPartFlav==1) | (muon.genPartFlav==15) )] #and now they are all prompt ##Leptons lepton = ak.concatenate([muon, electron], axis=1) SSlepton = (ak.sum(lepton.charge, axis=1) != 0) & (ak.num(lepton)==2) OSlepton = (ak.sum(lepton.charge, axis=1) == 0) & (ak.num(lepton)==2) emulepton = (ak.num(electron) == 1) & (ak.num(muon) == 1) no_mumu = (ak.num(muon) <= 1) leading_lepton_idx = ak.singletons(ak.argmax(lepton.pt, axis=1)) leading_lepton = lepton[leading_lepton_idx] trailing_lepton_idx = ak.singletons(ak.argmin(lepton.pt, axis=1)) trailing_lepton = lepton[trailing_lepton_idx] #jets jet = getJets(ev, minPt=40, maxEta=2.4, pt_var='pt') jet = jet[ak.argsort(jet.pt, ascending=False)] # need to sort wrt smeared and recorrected jet pt jet = jet[~match(jet, muon, deltaRCut=0.4)] # remove jets that overlap with muons jet = jet[~match(jet, electron, deltaRCut=0.4)] ## MET -> can switch to puppi MET met_pt = ev.MET.pt met_phi = ev.MET.phi # setting up the various weights weight = Weights( len(ev) ) weight2 = Weights( len(ev)) if not dataset=='MuonEG': # generator weight weight.add("weight", ev.genWeight) weight2.add("weight", ev.genWeight) weight2.add("charge flip", self.charge_flip_ratio.flip_weight(electron)) #selections filters = getFilters(ev, year=self.year, dataset=dataset) ss = (SSlepton) os = (OSlepton) jet_all = (ak.num(jet) >= 2) diele = (ak.num(electron) == 2) emu = (emulepton) flips = (n_flips == 1) no_flips = (n_flips == 0) nmm = no_mumu selection = PackedSelection() selection.add('filter', (filters) ) selection.add('ss', ss ) selection.add('os', os ) selection.add('jet', jet_all ) selection.add('ee', diele) selection.add('emu', emu) selection.add('flip', flips) selection.add('nflip', no_flips) selection.add('no_mumu', nmm) bl_reqs = ['filter'] + ['jet'] bl_reqs_d = { sel: True for sel in bl_reqs } baseline = selection.require(**bl_reqs_d) f_reqs = bl_reqs + ['flip'] + ['ss'] + ['ee'] f_reqs_d = {sel: True for sel in f_reqs} flip_sel = selection.require(**f_reqs_d) f2_reqs = bl_reqs + ['flip'] + ['ss'] + ['emu'] f2_reqs_d = {sel: True for sel in f2_reqs} flip_sel2 = selection.require(**f2_reqs_d) f3_reqs = bl_reqs + ['flip'] + ['ss'] + ['no_mumu'] f3_reqs_d = {sel: True for sel in f3_reqs} flip_sel3 = selection.require(**f3_reqs_d) nf_reqs = bl_reqs + ['nflip'] + ['os'] + ['ee'] nf_reqs_d = {sel: True for sel in nf_reqs} n_flip_sel = selection.require(**nf_reqs_d) nf2_reqs = bl_reqs + ['nflip'] + ['os'] + ['emu'] nf2_reqs_d = {sel: True for sel in nf2_reqs} n_flip_sel2 = selection.require(**nf2_reqs_d) nf3_reqs = bl_reqs + ['nflip'] + ['os'] + ['no_mumu'] nf3_reqs_d = {sel: True for sel in nf3_reqs} n_flip_sel3 = selection.require(**nf3_reqs_d) s_reqs = bl_reqs + ['ss'] + ['no_mumu'] s_reqs_d = { sel: True for sel in s_reqs } ss_sel = selection.require(**s_reqs_d) o_reqs = bl_reqs + ['os'] + ['no_mumu'] o_reqs_d = {sel: True for sel in o_reqs } os_sel = selection.require(**o_reqs_d) ees_reqs = bl_reqs + ['ss'] + ['ee'] ees_reqs_d = { sel: True for sel in ees_reqs } eess_sel = selection.require(**ees_reqs_d) eeo_reqs = bl_reqs + ['os'] + ['ee'] eeo_reqs_d = {sel: True for sel in eeo_reqs } eeos_sel = selection.require(**eeo_reqs_d) ems_reqs = bl_reqs + ['ss'] + ['emu'] ems_reqs_d = { sel: True for sel in ems_reqs } emss_sel = selection.require(**ems_reqs_d) emo_reqs = bl_reqs + ['os'] + ['emu'] emo_reqs_d = {sel: True for sel in emo_reqs } emos_sel = selection.require(**emo_reqs_d) #outputs output['N_jet'].fill(dataset=dataset, multiplicity=ak.num(jet)[baseline], weight=weight.weight()[baseline]) output['N_ele'].fill(dataset=dataset, multiplicity=ak.num(lepton)[ss_sel], weight=weight.weight()[ss_sel]) output['N_ele2'].fill(dataset=dataset, multiplicity=ak.num(lepton)[os_sel], weight=weight2.weight()[os_sel]) output['electron_flips'].fill(dataset=dataset, multiplicity = n_flips[flip_sel], weight=weight.weight()[flip_sel]) output['electron_flips2'].fill(dataset=dataset, multiplicity = n_flips[n_flip_sel], weight=weight2.weight()[n_flip_sel]) output['electron_flips3'].fill(dataset=dataset, multiplicity = n_flips[flip_sel2], weight=weight.weight()[flip_sel2]) output['electron_flips4'].fill(dataset=dataset, multiplicity = n_flips[n_flip_sel2], weight=weight2.weight()[n_flip_sel2]) output["electron"].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(leading_electron[flip_sel3].pt)), eta = np.abs(ak.to_numpy(ak.flatten(leading_electron[flip_sel3].eta))), weight = weight.weight()[flip_sel3] ) output["electron2"].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(leading_electron[n_flip_sel3].pt)), eta = np.abs(ak.to_numpy(ak.flatten(leading_electron[n_flip_sel3].eta))), weight = weight2.weight()[n_flip_sel3] ) output["flipped_electron"].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(leading_electron[flip_sel].pt)), eta = np.abs(ak.to_numpy(ak.flatten(leading_electron[flip_sel].eta))), weight = weight.weight()[flip_sel] ) output["flipped_electron2"].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(leading_electron[n_flip_sel].pt)), eta = np.abs(ak.to_numpy(ak.flatten(leading_electron[n_flip_sel].eta))), weight = weight2.weight()[n_flip_sel] ) output["flipped_electron3"].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(leading_electron[flip_sel2].pt)), eta = np.abs(ak.to_numpy(ak.flatten(leading_electron[flip_sel2].eta))), weight = weight.weight()[flip_sel2] ) output["flipped_electron4"].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(leading_electron[n_flip_sel2].pt)), eta = np.abs(ak.to_numpy(ak.flatten(leading_electron[n_flip_sel2].eta))), weight = weight2.weight()[n_flip_sel2] ) #output["lepton_parent"].fill( # dataset = dataset, # pdgID = np.abs(ak.to_numpy(ak.flatten(leading_parent[ss_sel]))), # weight = weight.weight()[ss_sel] #) # #output["lepton_parent2"].fill( # dataset = dataset, # pdgID = np.abs(ak.to_numpy(ak.flatten(trailing_parent[ss_sel]))), # weight = weight.weight()[ss_sel] #) return output
def process(self, events): output = self.accumulator.identity() # use a very loose preselection to filter the events presel = ak.num(events.Jet) > 2 ev = events[presel] dataset = ev.metadata['dataset'] # load the config - probably not needed anymore cfg = loadConfig() output['totalEvents']['all'] += len(events) output['skimmedEvents']['all'] += len(ev) ## Muons muon = Collections(ev, "Muon", "tightTTH").get() vetomuon = Collections(ev, "Muon", "vetoTTH").get() dimuon = choose(muon, 2) SSmuon = ak.any((dimuon['0'].charge * dimuon['1'].charge) > 0, axis=1) OSmuon = ak.any((dimuon['0'].charge * dimuon['1'].charge) < 0, axis=1) leading_muon_idx = ak.singletons(ak.argmax(muon.pt, axis=1)) leading_muon = muon[leading_muon_idx] ## Electrons electron = Collections(ev, "Electron", "tightTTH").get() vetoelectron = Collections(ev, "Electron", "vetoTTH").get() dielectron = choose(electron, 2) SSelectron = ak.any( (dielectron['0'].charge * dielectron['1'].charge) > 0, axis=1) OSelectron = ak.any( (dielectron['0'].charge * dielectron['1'].charge) < 0, axis=1) leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1)) leading_electron = electron[leading_electron_idx] ## Merge electrons and muons - this should work better now in ak1 lepton = ak.concatenate([muon, electron], axis=1) dilepton = cross(muon, electron) SSlepton = ak.any((dilepton['0'].charge * dilepton['1'].charge) > 0, axis=1) OSlepton = ak.any((dilepton['0'].charge * dilepton['1'].charge) < 0, axis=1) leading_lepton_idx = ak.singletons(ak.argmax(lepton.pt, axis=1)) leading_lepton = lepton[leading_lepton_idx] trailing_lepton_idx = ak.singletons(ak.argmin(lepton.pt, axis=1)) trailing_lepton = lepton[trailing_lepton_idx] ## Jets jet = getJets(ev, minPt=25, maxEta=4.7, pt_var='pt_nom') jet = jet[ak.argsort( jet.pt_nom, ascending=False )] # need to sort wrt smeared and recorrected jet pt jet = jet[~match(jet, muon, deltaRCut=0.4)] # remove jets that overlap with muons jet = jet[~match( jet, electron, deltaRCut=0.4)] # remove jets that overlap with electrons central = jet[(abs(jet.eta) < 2.4)] btag = getBTagsDeepFlavB( jet, year=self.year) # should study working point for DeepJet light = getBTagsDeepFlavB(jet, year=self.year, invert=True) fwd = getFwdJet(light) fwd_noPU = getFwdJet(light, puId=False) ## forward jets high_p_fwd = fwd[ak.singletons(ak.argmax( fwd.p, axis=1))] # highest momentum spectator high_pt_fwd = fwd[ak.singletons(ak.argmax( fwd.pt_nom, axis=1))] # highest transverse momentum spectator high_eta_fwd = fwd[ak.singletons(ak.argmax(abs( fwd.eta), axis=1))] # most forward spectator ## Get the two leading b-jets in terms of btag score high_score_btag = central[ak.argsort(central.btagDeepFlavB)][:, :2] jf = cross(high_p_fwd, jet) mjf = (jf['0'] + jf['1']).mass deltaEta = abs(high_p_fwd.eta - jf[ak.singletons(ak.argmax(mjf, axis=1))]['1'].eta) deltaEtaMax = ak.max(deltaEta, axis=1) mjf_max = ak.max(mjf, axis=1) jj = choose(jet, 2) mjj_max = ak.max((jj['0'] + jj['1']).mass, axis=1) ## MET -> can switch to puppi MET met_pt = ev.MET.pt met_phi = ev.MET.phi ## other variables ht = ak.sum(jet.pt, axis=1) st = met_pt + ht + ak.sum(muon.pt, axis=1) + ak.sum(electron.pt, axis=1) ht_central = ak.sum(central.pt, axis=1) ## event selectors filters = getFilters(ev, year=self.year, dataset=dataset) triggers = getTriggers(ev, year=self.year, dataset=dataset) dilep = ((ak.num(electron) == 1) & (ak.num(muon) == 1)) lep0pt = ((ak.num(electron[(electron.pt > 25)]) + ak.num(muon[(muon.pt > 25)])) > 0) lep1pt = ((ak.num(electron[(electron.pt > 20)]) + ak.num(muon[(muon.pt > 20)])) > 1) lepveto = ((ak.num(vetoelectron) + ak.num(vetomuon)) == 2) # define the weight weight = Weights(len(ev)) if not dataset == 'MuonEG': # lumi weight weight.add("weight", ev.weight * cfg['lumi'][self.year]) # PU weight - not in the babies... weight.add("PU", ev.puWeight, weightUp=ev.puWeightUp, weightDown=ev.puWeightDown, shift=False) # b-tag SFs weight.add("btag", self.btagSF.Method1a(btag, light)) # lepton SFs weight.add("lepton", self.leptonSF.get(electron, muon)) selection = PackedSelection() selection.add('lepveto', lepveto) selection.add('dilep', dilep) selection.add('trigger', (triggers)) selection.add('filter', (filters)) selection.add('p_T(lep0)>25', lep0pt) selection.add('p_T(lep1)>20', lep1pt) selection.add('OS', OSlepton) selection.add('N_btag=2', (ak.num(btag) == 2)) selection.add('N_jet>2', (ak.num(jet) >= 3)) selection.add('N_central>1', (ak.num(central) >= 2)) selection.add('N_fwd>0', (ak.num(fwd) >= 1)) selection.add('MET>30', (ev.MET.pt > 30)) os_reqs = [ 'lepveto', 'dilep', 'trigger', 'filter', 'p_T(lep0)>25', 'p_T(lep1)>20', 'OS' ] bl_reqs = os_reqs + [ 'N_btag=2', 'N_jet>2', 'N_central>1', 'N_fwd>0', 'MET>30' ] os_reqs_d = {sel: True for sel in os_reqs} os_selection = selection.require(**os_reqs_d) bl_reqs_d = {sel: True for sel in bl_reqs} BL = selection.require(**bl_reqs_d) cutflow = Cutflow(output, ev, weight=weight) cutflow_reqs_d = {} for req in bl_reqs: cutflow_reqs_d.update({req: True}) cutflow.addRow(req, selection.require(**cutflow_reqs_d)) # first, make a few super inclusive plots output['PV_npvs'].fill(dataset=dataset, multiplicity=ev.PV[os_selection].npvs, weight=weight.weight()[os_selection]) output['PV_npvsGood'].fill(dataset=dataset, multiplicity=ev.PV[os_selection].npvsGood, weight=weight.weight()[os_selection]) output['N_jet'].fill(dataset=dataset, multiplicity=ak.num(jet)[os_selection], weight=weight.weight()[os_selection]) output['N_b'].fill(dataset=dataset, multiplicity=ak.num(btag)[os_selection], weight=weight.weight()[os_selection]) output['N_central'].fill(dataset=dataset, multiplicity=ak.num(central)[os_selection], weight=weight.weight()[os_selection]) output['N_ele'].fill(dataset=dataset, multiplicity=ak.num(electron)[os_selection], weight=weight.weight()[os_selection]) output['N_mu'].fill(dataset=dataset, multiplicity=ak.num(electron)[os_selection], weight=weight.weight()[os_selection]) output['N_fwd'].fill(dataset=dataset, multiplicity=ak.num(fwd)[os_selection], weight=weight.weight()[os_selection]) output['MET'].fill(dataset=dataset, pt=ev.MET[os_selection].pt, phi=ev.MET[os_selection].phi, weight=weight.weight()[os_selection]) output['electron'].fill(dataset=dataset, pt=ak.to_numpy(ak.flatten(electron[BL].pt)), eta=ak.to_numpy(ak.flatten(electron[BL].eta)), phi=ak.to_numpy(ak.flatten(electron[BL].phi)), weight=weight.weight()[BL]) output['muon'].fill(dataset=dataset, pt=ak.to_numpy(ak.flatten(muon[BL].pt)), eta=ak.to_numpy(ak.flatten(muon[BL].eta)), phi=ak.to_numpy(ak.flatten(muon[BL].phi)), weight=weight.weight()[BL]) output['lead_lep'].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(leading_lepton[BL].pt)), eta=ak.to_numpy(ak.flatten(leading_lepton[BL].eta)), phi=ak.to_numpy(ak.flatten(leading_lepton[BL].phi)), weight=weight.weight()[BL]) output['trail_lep'].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(trailing_lepton[BL].pt)), eta=ak.to_numpy(ak.flatten(trailing_lepton[BL].eta)), phi=ak.to_numpy(ak.flatten(trailing_lepton[BL].phi)), weight=weight.weight()[BL]) output['fwd_jet'].fill(dataset=dataset, pt=ak.flatten(high_p_fwd[BL].pt_nom), eta=ak.flatten(high_p_fwd[BL].eta), phi=ak.flatten(high_p_fwd[BL].phi), weight=weight.weight()[BL]) output['b1'].fill(dataset=dataset, pt=ak.flatten(high_score_btag[:, 0:1][BL].pt_nom), eta=ak.flatten(high_score_btag[:, 0:1][BL].eta), phi=ak.flatten(high_score_btag[:, 0:1][BL].phi), weight=weight.weight()[BL]) output['b2'].fill(dataset=dataset, pt=ak.flatten(high_score_btag[:, 1:2][BL].pt_nom), eta=ak.flatten(high_score_btag[:, 1:2][BL].eta), phi=ak.flatten(high_score_btag[:, 1:2][BL].phi), weight=weight.weight()[BL]) output['j1'].fill(dataset=dataset, pt=ak.flatten(jet.pt_nom[:, 0:1][BL]), eta=ak.flatten(jet.eta[:, 0:1][BL]), phi=ak.flatten(jet.phi[:, 0:1][BL]), weight=weight.weight()[BL]) output['j2'].fill(dataset=dataset, pt=ak.flatten(jet[:, 1:2][BL].pt_nom), eta=ak.flatten(jet[:, 1:2][BL].eta), phi=ak.flatten(jet[:, 1:2][BL].phi), weight=weight.weight()[BL]) output['j3'].fill(dataset=dataset, pt=ak.flatten(jet[:, 2:3][BL].pt_nom), eta=ak.flatten(jet[:, 2:3][BL].eta), phi=ak.flatten(jet[:, 2:3][BL].phi), weight=weight.weight()[BL]) # Now, take care of systematic unceratinties if not dataset == 'MuonEG': alljets = getJets(ev, minPt=0, maxEta=4.7) alljets = alljets[(alljets.jetId > 1)] for var in self.variations: # get the collections that change with the variations jet = getPtEtaPhi(alljets, pt_var=var) jet = jet[(jet.pt > 25)] jet = jet[~match( jet, muon, deltaRCut=0.4)] # remove jets that overlap with muons jet = jet[~match( jet, electron, deltaRCut=0.4)] # remove jets that overlap with electrons central = jet[(abs(jet.eta) < 2.4)] btag = getBTagsDeepFlavB( jet, year=self.year) # should study working point for DeepJet light = getBTagsDeepFlavB(jet, year=self.year, invert=True) fwd = getFwdJet(light) fwd_noPU = getFwdJet(light, puId=False) ## forward jets high_p_fwd = fwd[ak.singletons(ak.argmax( fwd.p, axis=1))] # highest momentum spectator high_pt_fwd = fwd[ak.singletons(ak.argmax( fwd.pt, axis=1))] # highest transverse momentum spectator high_eta_fwd = fwd[ak.singletons( ak.argmax(abs(fwd.eta), axis=1))] # most forward spectator ## Get the two leading b-jets in terms of btag score high_score_btag = central[ak.argsort( central.btagDeepFlavB)][:, :2] # get the modified selection -> more difficult selection.add('N_jet>2_' + var, (ak.num(jet.pt) >= 3)) # stupid bug here... selection.add('N_btag=2_' + var, (ak.num(btag) == 2)) selection.add('N_central>1_' + var, (ak.num(central) >= 2)) selection.add('N_fwd>0_' + var, (ak.num(fwd) >= 1)) selection.add('MET>30_' + var, (getattr(ev.MET, var) > 30)) ## Don't change the selection for now... bl_reqs = os_reqs + [ 'N_jet>2_' + var, 'MET>30_' + var, 'N_btag=2_' + var, 'N_central>1_' + var, 'N_fwd>0_' + var ] bl_reqs_d = {sel: True for sel in bl_reqs} BL = selection.require(**bl_reqs_d) # the OS selection remains unchanged output['N_jet_' + var].fill( dataset=dataset, multiplicity=ak.num(jet)[os_selection], weight=weight.weight()[os_selection]) output['N_fwd_' + var].fill( dataset=dataset, multiplicity=ak.num(fwd)[os_selection], weight=weight.weight()[os_selection]) output['N_b_' + var].fill( dataset=dataset, multiplicity=ak.num(btag)[os_selection], weight=weight.weight()[os_selection]) output['N_central_' + var].fill( dataset=dataset, multiplicity=ak.num(central)[os_selection], weight=weight.weight()[os_selection]) # We don't need to redo all plots with variations. E.g., just add uncertainties to the jet plots. output['j1_' + var].fill(dataset=dataset, pt=ak.flatten(jet.pt[:, 0:1][BL]), eta=ak.flatten(jet.eta[:, 0:1][BL]), phi=ak.flatten(jet.phi[:, 0:1][BL]), weight=weight.weight()[BL]) output['b1_' + var].fill( dataset=dataset, pt=ak.flatten(high_score_btag[:, 0:1].pt[:, 0:1][BL]), eta=ak.flatten(high_score_btag[:, 0:1].eta[:, 0:1][BL]), phi=ak.flatten(high_score_btag[:, 0:1].phi[:, 0:1][BL]), weight=weight.weight()[BL]) output['fwd_jet_' + var].fill( dataset=dataset, pt=ak.flatten(high_p_fwd[BL].pt), eta=ak.flatten(high_p_fwd[BL].eta), phi=ak.flatten(high_p_fwd[BL].phi), weight=weight.weight()[BL]) output['MET_' + var].fill(dataset=dataset, pt=getattr(ev.MET, var)[os_selection], phi=ev.MET[os_selection].phi, weight=weight.weight()[os_selection]) return output
tracks_bst_leadPt_ak4_isr) sphTensor_noLowMult = eventShapesUtilities.sphericityTensor( tracks_bst_noLowMult) sph_allTracks[ievt] = eventShapesUtilities.sphericity(sphTensor_allTracks) sph_highMult[ievt] = eventShapesUtilities.sphericity(sphTensor_highMult) sph_leadPt[ievt] = eventShapesUtilities.sphericity(sphTensor_leadPt) sph_leadPt_ak4_suep[ievt] = eventShapesUtilities.sphericity( sphTensor_leadPt_ak4_suep) sph_leadPt_ak4_isr[ievt] = eventShapesUtilities.sphericity( sphTensor_leadPt_ak4_isr) sph_noLowMult[ievt] = eventShapesUtilities.sphericity(sphTensor_noLowMult) trackMultiplicity[ievt] = tracks.size # Getting cross section and HT and keeping only processed events CrossSection = ak.to_numpy(events['CrossSection']) HT = ak.to_numpy(events['HT']) CrossSection = CrossSection[:N_events] HT = HT[:N_events] # Filter HT < 1200 out CrossSection = CrossSection[HT >= 1200] sph_allTracks = sph_allTracks[HT >= 1200] sph_dPhi = sph_dPhi[HT >= 1200] sph_relE = sph_relE[HT >= 1200] sph_highMult = sph_highMult[HT >= 1200] sph_leadPt = sph_leadPt[HT >= 1200] sph_leadPt_ak4_suep = sph_leadPt_ak4_suep[HT >= 1200] sph_leadPt_ak4_isr = sph_leadPt_ak4_isr[HT >= 1200] sph_noLowMult = sph_noLowMult[HT >= 1200] beta_v = beta_v[HT >= 1200]
def process(self, events): output = self.accumulator.identity() # we can use a very loose preselection to filter the events. nothing is done with this presel, though presel = ak.num(events.Jet) > 0 ev = events[presel] dataset = ev.metadata['dataset'] # load the config - probably not needed anymore cfg = loadConfig() output['totalEvents']['all'] += len(events) output['skimmedEvents']['all'] += len(ev) ## Electrons electron = Collections(ev, "Electron", "tight").get() electron = electron[(electron.miniPFRelIso_all < 0.12) & (electron.pt > 20) & (abs(electron.eta) < 2.4)] gen_matched_electron = electron[((electron.genPartIdx >= 0) & (abs( electron.matched_gen.pdgId) == 11))] n_gen = ak.num(gen_matched_electron) is_flipped = ((gen_matched_electron.matched_gen.pdgId * (-1) == gen_matched_electron.pdgId) & (abs(gen_matched_electron.pdgId) == 11)) #is_flipped = (abs(ev.GenPart[gen_matched_electron.genPartIdx].pdgId) == abs(gen_matched_electron.pdgId))&(ev.GenPart[gen_matched_electron.genPartIdx].pdgId/abs(ev.GenPart[gen_matched_electron.genPartIdx].pdgId) != gen_matched_electron.pdgId/abs(gen_matched_electron.pdgId)) flipped_electron = gen_matched_electron[is_flipped] n_flips = ak.num(flipped_electron) sielectron = choose(electron, 1) dielectron = choose(electron, 2) SSelectron = ak.any( (dielectron['0'].charge * dielectron['1'].charge) > 0, axis=1) leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1)) leading_electron = electron[leading_electron_idx] ## MET -> can switch to puppi MET met_pt = ev.MET.pt met_phi = ev.MET.phi # setting up the various weights weight = Weights(len(ev)) weight2 = Weights(len(ev)) if not dataset == 'MuonEG': # generator weight weight.add("weight", ev.genWeight) weight2.add("weight", ev.genWeight) weight2.add("charge flip", self.charge_flip_ratio.flip_ratio(sielectron['0'])) #selections filters = getFilters(ev, year=self.year, dataset=dataset) electr = ((ak.num(electron) == 2)) ss = (SSelectron) gen = (n_gen >= 1) flip = (n_flips >= 1) selection = PackedSelection() selection.add('filter', (filters)) selection.add('electr', electr) selection.add('ss', ss) selection.add('flip', flip) selection.add('gen', gen) bl_reqs = ['filter', 'electr'] bl_reqs_d = {sel: True for sel in bl_reqs} baseline = selection.require(**bl_reqs_d) s_reqs = bl_reqs + ['ss'] s_reqs_d = {sel: True for sel in s_reqs} ss_sel = selection.require(**s_reqs_d) f_reqs = bl_reqs + ['gen', 'flip'] f_reqs_d = {sel: True for sel in f_reqs} flip_sel = selection.require(**f_reqs_d) #outputs output['N_ele'].fill(dataset=dataset, multiplicity=ak.num(electron)[flip_sel], weight=weight.weight()[flip_sel]) output['electron_flips'].fill(dataset=dataset, multiplicity=n_flips[flip_sel], weight=weight.weight()[flip_sel]) output['N_ele2'].fill(dataset=dataset, multiplicity=ak.num(electron)[baseline], weight=weight2.weight()[baseline]) output['electron_flips2'].fill(dataset=dataset, multiplicity=n_flips[baseline], weight=weight2.weight()[baseline]) output["electron"].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(flipped_electron[flip_sel].pt)), eta=ak.to_numpy(ak.flatten(abs(flipped_electron[flip_sel].eta))), #phi = ak.to_numpy(ak.flatten(leading_electron[baseline].phi)), weight=weight.weight()[flip_sel]) output["electron2"].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(leading_electron[baseline].pt)), eta=ak.to_numpy(ak.flatten(abs(leading_electron[baseline].eta))), #phi = ak.to_numpy(ak.flatten(leading_electron[baseline].phi)), weight=weight2.weight()[baseline]) return output
def process(self, events): # get meta infos dataset = events.metadata["dataset"] isRealData = not hasattr(events, "genWeight") n_events = len(events) selection = processor.PackedSelection() weights = processor.Weights(n_events) output = self.accumulator.identity() # weights if not isRealData: output['sumw'][dataset] += awkward1.sum(events.genWeight) # trigger triggers = {} for channel in ["e","mu"]: trigger = np.zeros(len(events), dtype='bool') for t in self._trigger[channel]: try: trigger = trigger | events.HLT[t] except: warnings.warn("Missing trigger %s" % t, RuntimeWarning) triggers[channel] = trigger # met filter met_filters = ["goodVertices", "globalSuperTightHalo2016Filter", "HBHENoiseFilter", "HBHENoiseIsoFilter", "EcalDeadCellTriggerPrimitiveFilter", "BadPFMuonFilter", ] met_filters_mask = np.ones(len(events), dtype='bool') for t in met_filters: met_filters_mask = met_filters_mask & events.Flag[t] selection.add("met_filter", awkward1.to_numpy(met_filters_mask)) # load objects muons = events.Muon electrons = events.Electron jets = events.Jet fatjets = events.FatJet subjets = events.SubJet fatjetsLS = events.FatJetLS met = events.MET # muons goodmuon = ( (muons.mediumId) & (muons.miniPFRelIso_all <= 0.2) & (muons.pt >= 27) & (abs(muons.eta) <= 2.4) & (abs(muons.dz) < 0.1) & (abs(muons.dxy) < 0.05) & (muons.sip3d < 4) ) good_muons = muons[goodmuon] ngood_muons = awkward1.sum(goodmuon, axis=1) # electrons goodelectron = ( (electrons.mvaFall17V2noIso_WP90) & (electrons.pt >= 30) & (abs(electrons.eta) <= 1.479) & (abs(electrons.dz) < 0.1) & (abs(electrons.dxy) < 0.05) & (electrons.sip3d < 4) ) good_electrons = electrons[goodelectron] ngood_electrons = awkward1.sum(goodelectron, axis=1) # good leptons good_leptons = awkward1.concatenate([good_muons, good_electrons], axis=1) good_leptons = good_leptons[awkward1.argsort(good_leptons.pt)] # lepton candidate candidatelep = awkward1.firsts(good_leptons) # lepton channel selection selection.add("ch_e", awkward1.to_numpy((triggers["e"]) & (ngood_electrons==1) & (ngood_muons==0))) # not sure if need to require 0 muons or 0 electrons in the next line selection.add("ch_mu", awkward1.to_numpy((triggers["mu"]) & (ngood_electrons==0) & (ngood_muons==1))) # jets ht = awkward1.sum(jets[jets.pt > 30].pt,axis=1) selection.add("ht_400", awkward1.to_numpy(ht>=400)) goodjet = ( (jets.isTight) & (jets.pt > 30) & (abs(jets.eta) <= 2.5) ) good_jets = jets[goodjet] # fat jets jID = "isTight" # TODO: add mass correction # a way to get the first two subjets # cart = awkward1.cartesian([fatjets, subjets], nested=True) # idxes = awkward1.pad_none(awkward1.argsort(cart['0'].delta_r(cart['1'])), 2, axis=2) # sj1 = subjets[idxes[:,:,0]] # sj2 = subjets[idxes[:,:,1]] good_fatjet = ( (getattr(fatjets, jID)) & (abs(fatjets.eta) <= 2.4) & (fatjets.pt > 50) & (fatjets.msoftdrop > 30) & (fatjets.msoftdrop < 210) #& (fatjets.pt.copy(content=fatjets.subjets.content.counts) == 2) # TODO: require 2 subjets? # this can probably be done w FatJet_subJetIdx1 or FatJet_subJetIdx2 & (awkward1.all(fatjets.subjets.pt >= 20)) & (awkward1.all(abs(fatjets.subjets.eta) <= 2.4)) ) good_fatjets = fatjets[good_fatjet] # hbb candidate mask_hbb = ( (good_fatjets.pt > 200) & (good_fatjets.delta_r(candidatelep) > 2.0) ) candidateHbb = awkward1.firsts(good_fatjets[mask_hbb]) # b-tag #& (good_fatjets.particleNetMD_Xbb > 0.9) selection.add('hbb_btag',awkward1.to_numpy(candidateHbb.deepTagMD_ZHbbvsQCD >= 0.8)) # score would be larger for tight category (0.97) # No AK4 b-tagged jets away from bb jet jets_HbbV = jets[good_jets.delta_r(candidateHbb) >= 1.2] selection.add('hbb_vetobtagaway', awkward1.to_numpy(awkward1.max(jets_HbbV.btagDeepB, axis=1, mask_identity=False) > BTagEfficiency.btagWPs[self._year]['medium'])) # fat jets Lepton Subtracted # wjj candidate mask_wjj = ( (fatjetsLS.pt > 50) & (fatjetsLS.delta_r(candidatelep) > 1.2) # need to add 2 subjets w pt > 20 & eta<2.4 # need to add ID? ) candidateWjj = awkward1.firsts(fatjetsLS[mask_wjj][awkward1.argmin(fatjetsLS[mask_wjj].delta_r(candidatelep),axis=1,keepdims=True)]) # add t2/t1 <= 0.75 (0.45 HP) selection.add('hww_mass', awkward1.to_numpy(candidateWjj.mass >= 10)) print('met ',met) # wjjlnu info #HSolverLiInfo hwwInfoLi; # qqSDmass = candidateWjj.msoftdrop # hwwLi = hSolverLi->minimize(candidatelep.p4(), met.p4(), wjjcand.p4(), qqSDmass, hwwInfoLi) #neutrino = hwwInfoLi.neutrino; #wlnu = hwwInfoLi.wlnu; #wqq = hwwInfoLi.wqqjet; #hWW = hwwInfoLi.hWW; #wwDM = PhysicsUtilities::deltaR( wlnu,wqq) * hWW.pt()/2.0; # add dlvqq <= 11 (2.5 HP) # in the meantime let's add the mass ''' mm = (candidatejet - candidatelep).mass2 jmass = (mm>0)*np.sqrt(np.maximum(0, mm)) + (mm<0)*candidatejet.mass joffshell = jmass < 62.5 massassumption = 80.*joffshell + (125 - 80.)*~joffshell x = massassumption**2/(2*candidatelep.pt*met.pt) + np.cos(candidatelep.phi - met.phi) met_eta = ( (x < 1)*np.arcsinh(x*np.sinh(candidatelep.eta)) + (x > 1)*( candidatelep.eta - np.sign(candidatelep.eta)*np.arccosh(candidatelep.eta) ) ) met_p4 = TLorentzVectorArray.from_ptetaphim(np.array([0.]),np.array([0.]),np.array([0.]),np.array([0.])) if met.size > 0: met_p4 = TLorentzVectorArray.from_ptetaphim(met.pt, met_eta.fillna(0.), met.phi, np.zeros(met.size)) # hh system candidateHH = candidateWjj + met_p4 + candidateHbb selection.add('hh_mass', candidateHH.mass >= 700) selection.add('hh_centrality', candidateHH.pt/candidateHH.mass >= 0.3) ''' channels = {"e": ["met_filter","ch_e","ht_400","hbb_btag","hbb_vetobtagaway","hww_mass"], #,"hh_mass","hh_centrality"], "mu": ["met_filter","ch_mu","ht_400","hbb_btag","hbb_vetobtagaway","hww_mass"] #,"hh_mass","hh_centrality"], } # need to add gen info if not isRealData: weights.add('genweight', events.genWeight) add_pileup_weight(weights, events.Pileup.nPU, self._year, dataset) for channel, cuts in channels.items(): allcuts = set() output['cutflow'].fill(dataset=dataset, channel=channel, cut=0, weight=weights.weight()) for i, cut in enumerate(cuts): allcuts.add(cut) cut = selection.all(*allcuts) output['cutflow'].fill(dataset=dataset, channel=channel, cut=i + 1, weight=weights.weight()[cut]) return output
def process(self, events): output = self.accumulator.identity() # use a very loose preselection to filter the events presel = ak.num(events.Jet)>2 ev = events[presel] dataset = ev.metadata['dataset'] # load the config - probably not needed anymore cfg = loadConfig() output['totalEvents']['all'] += len(events) output['skimmedEvents']['all'] += len(ev) ## Generated leptons gen_lep = ev.GenL leading_gen_lep = gen_lep[ak.singletons(ak.argmax(gen_lep.pt, axis=1))] trailing_gen_lep = gen_lep[ak.singletons(ak.argmin(gen_lep.pt, axis=1))] ## Muons muon = Collections(ev, "Muon", "tightTTH").get() vetomuon = Collections(ev, "Muon", "vetoTTH").get() leading_muon_idx = ak.singletons(ak.argmax(muon.pt, axis=1)) leading_muon = muon[leading_muon_idx] ## Electrons electron = Collections(ev, "Electron", "tightTTH").get() vetoelectron = Collections(ev, "Electron", "vetoTTH").get() leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1)) leading_electron = electron[leading_electron_idx] ## Merge electrons and muons - this should work better now in ak1 dilepton = cross(muon, electron) dimuon = choose(muon,2) OS_dimuon = dimuon[(dimuon['0'].charge*dimuon['1'].charge < 0)] dielectron = choose(electron) OS_dielectron = dielectron[(dielectron['0'].charge*dielectron['1'].charge < 0)] OS_dimuon_bestZmumu = OS_dimuon[ak.singletons(ak.argmin(abs(OS_dimuon.mass-91.2), axis=1))] OS_dielectron_bestZee = OS_dielectron[ak.singletons(ak.argmin(abs(OS_dielectron.mass-91.2), axis=1))] OS_dilepton_mass = ak.fill_none(ak.pad_none(ak.concatenate([OS_dimuon_bestZmumu.mass, OS_dielectron_bestZee.mass], axis=1), 1, clip=True), -1) lepton = ak.concatenate([muon, electron], axis=1) leading_lepton_idx = ak.singletons(ak.argmax(lepton.pt, axis=1)) leading_lepton = lepton[leading_lepton_idx] trailing_lepton_idx = ak.singletons(ak.argmin(lepton.pt, axis=1)) trailing_lepton = lepton[trailing_lepton_idx] ## Jets jet = getJets(ev, minPt=25, maxEta=4.7, pt_var='pt_nom') jet = jet[ak.argsort(jet.pt_nom, ascending=False)] # need to sort wrt smeared and recorrected jet pt jet = jet[~match(jet, muon, deltaRCut=0.4)] # remove jets that overlap with muons jet = jet[~match(jet, electron, deltaRCut=0.4)] # remove jets that overlap with electrons central = jet[(abs(jet.eta)<2.4)] btag = getBTagsDeepFlavB(jet, year=self.year) # should study working point for DeepJet light = getBTagsDeepFlavB(jet, year=self.year, invert=True) fwd = getFwdJet(light) fwd_noPU = getFwdJet(light, puId=False) ## forward jets j_fwd = fwd[ak.singletons(ak.argmax(fwd.p, axis=1))] # highest momentum spectator jf = cross(j_fwd, jet) mjf = (jf['0']+jf['1']).mass j_fwd2 = jf[ak.singletons(ak.argmax(mjf, axis=1))]['1'] # this is the jet that forms the largest invariant mass with j_fwd delta_eta = abs(j_fwd2.eta - j_fwd.eta) ## MET -> can switch to puppi MET met_pt = ev.MET.pt met_phi = ev.MET.phi ## other variables ht = ak.sum(jet.pt, axis=1) st = met_pt + ht + ak.sum(muon.pt, axis=1) + ak.sum(electron.pt, axis=1) # define the weight weight = Weights( len(ev) ) if not dataset=='MuonEG': # lumi weight weight.add("weight", ev.weight*cfg['lumi'][self.year]) #weight.add("weight", ev.genWeight*cfg['lumi'][self.year]*mult) # PU weight - not in the babies... weight.add("PU", ev.puWeight, weightUp=ev.puWeightUp, weightDown=ev.puWeightDown, shift=False) # b-tag SFs weight.add("btag", self.btagSF.Method1a(btag, light)) ## lepton SFs #weight.add("lepton", self.leptonSF.get(electron, muon)) cutflow = Cutflow(output, ev, weight=weight) sel = Selection( dataset = dataset, events = ev, year = self.year, ele = electron, ele_veto = vetoelectron, mu = muon, mu_veto = vetomuon, jet_all = jet, jet_central = central, jet_btag = btag, jet_fwd = fwd, met = ev.MET, ) BL = sel.trilep_baseline(cutflow=cutflow) # first, make a few super inclusive plots output['PV_npvs'].fill(dataset=dataset, multiplicity=ev.PV[BL].npvs, weight=weight.weight()[BL]) output['PV_npvsGood'].fill(dataset=dataset, multiplicity=ev.PV[BL].npvsGood, weight=weight.weight()[BL]) output['N_jet'].fill(dataset=dataset, multiplicity=ak.num(jet)[BL], weight=weight.weight()[BL]) output['N_b'].fill(dataset=dataset, multiplicity=ak.num(btag)[BL], weight=weight.weight()[BL]) output['N_central'].fill(dataset=dataset, multiplicity=ak.num(central)[BL], weight=weight.weight()[BL]) output['N_ele'].fill(dataset=dataset, multiplicity=ak.num(electron)[BL], weight=weight.weight()[BL]) output['N_mu'].fill(dataset=dataset, multiplicity=ak.num(electron)[BL], weight=weight.weight()[BL]) output['N_fwd'].fill(dataset=dataset, multiplicity=ak.num(fwd)[BL], weight=weight.weight()[BL]) output['nLepFromTop'].fill(dataset=dataset, multiplicity=ev[BL].nLepFromTop, weight=weight.weight()[BL]) output['nLepFromTau'].fill(dataset=dataset, multiplicity=ev.nLepFromTau[BL], weight=weight.weight()[BL]) output['nLepFromZ'].fill(dataset=dataset, multiplicity=ev.nLepFromZ[BL], weight=weight.weight()[BL]) output['nLepFromW'].fill(dataset=dataset, multiplicity=ev.nLepFromW[BL], weight=weight.weight()[BL]) output['nGenTau'].fill(dataset=dataset, multiplicity=ev.nGenTau[BL], weight=weight.weight()[BL]) output['nGenL'].fill(dataset=dataset, multiplicity=ak.num(ev.GenL[BL], axis=1), weight=weight.weight()[BL]) # make a plot of the dilepton mass, but without applying the cut on the dilepton mass itself (N-1 plot) output['dilep_mass'].fill(dataset=dataset, mass=ak.flatten(OS_dilepton_mass[sel.trilep_baseline(omit=['offZ'])]), weight=weight.weight()[sel.trilep_baseline(omit=['offZ'])]) output['MET'].fill( dataset = dataset, pt = ev.MET[BL].pt, phi = ev.MET[BL].phi, weight = weight.weight()[BL] ) output['lead_gen_lep'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(leading_gen_lep[BL].pt)), eta = ak.to_numpy(ak.flatten(leading_gen_lep[BL].eta)), phi = ak.to_numpy(ak.flatten(leading_gen_lep[BL].phi)), weight = weight.weight()[BL] ) output['trail_gen_lep'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(trailing_gen_lep[BL].pt)), eta = ak.to_numpy(ak.flatten(trailing_gen_lep[BL].eta)), phi = ak.to_numpy(ak.flatten(trailing_gen_lep[BL].phi)), weight = weight.weight()[BL] ) output['lead_lep'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(leading_lepton[BL].pt)), eta = ak.to_numpy(ak.flatten(leading_lepton[BL].eta)), phi = ak.to_numpy(ak.flatten(leading_lepton[BL].phi)), weight = weight.weight()[BL] ) output['trail_lep'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(trailing_lepton[BL].pt)), eta = ak.to_numpy(ak.flatten(trailing_lepton[BL].eta)), phi = ak.to_numpy(ak.flatten(trailing_lepton[BL].phi)), weight = weight.weight()[BL] ) output['j1'].fill( dataset = dataset, pt = ak.flatten(jet.pt_nom[:, 0:1][BL]), eta = ak.flatten(jet.eta[:, 0:1][BL]), phi = ak.flatten(jet.phi[:, 0:1][BL]), weight = weight.weight()[BL] ) output['j2'].fill( dataset = dataset, pt = ak.flatten(jet[:, 1:2][BL].pt_nom), eta = ak.flatten(jet[:, 1:2][BL].eta), phi = ak.flatten(jet[:, 1:2][BL].phi), weight = weight.weight()[BL] ) #output['j3'].fill( # dataset = dataset, # pt = ak.flatten(jet[:, 2:3][BL].pt_nom), # eta = ak.flatten(jet[:, 2:3][BL].eta), # phi = ak.flatten(jet[:, 2:3][BL].phi), # weight = weight.weight()[BL] #) output['fwd_jet'].fill( dataset = dataset, pt = ak.flatten(j_fwd[BL].pt), eta = ak.flatten(j_fwd[BL].eta), phi = ak.flatten(j_fwd[BL].phi), weight = weight.weight()[BL] ) output['high_p_fwd_p'].fill(dataset=dataset, p = ak.flatten(j_fwd[BL].p), weight = weight.weight()[BL]) return output
def process(self, events): output = self.accumulator.identity() # use a very loose preselection to filter the events presel = ak.num(events.Jet)>2 ev = events[presel] dataset = ev.metadata['dataset'] # load the config - probably not needed anymore cfg = loadConfig() output['totalEvents']['all'] += len(events) output['skimmedEvents']['all'] += len(ev) if not re.search(re.compile('MuonEG|DoubleMuon|DoubleEG|EGamma'), dataset): ## Generated leptons gen_lep = ev.GenL leading_gen_lep = gen_lep[ak.singletons(ak.argmax(gen_lep.pt, axis=1))] trailing_gen_lep = gen_lep[ak.singletons(ak.argmin(gen_lep.pt, axis=1))] ## Muons muon = Collections(ev, "Muon", "tightSSTTH").get() vetomuon = Collections(ev, "Muon", "vetoTTH").get() dimuon = choose(muon, 2) SSmuon = ak.any((dimuon['0'].charge * dimuon['1'].charge)>0, axis=1) leading_muon_idx = ak.singletons(ak.argmax(muon.pt, axis=1)) leading_muon = muon[leading_muon_idx] ## Electrons electron = Collections(ev, "Electron", "tightSSTTH").get() vetoelectron = Collections(ev, "Electron", "vetoTTH").get() dielectron = choose(electron, 2) SSelectron = ak.any((dielectron['0'].charge * dielectron['1'].charge)>0, axis=1) leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1)) leading_electron = electron[leading_electron_idx] ## Merge electrons and muons - this should work better now in ak1 dilepton = cross(muon, electron) SSlepton = ak.any((dilepton['0'].charge * dilepton['1'].charge)>0, axis=1) lepton = ak.concatenate([muon, electron], axis=1) leading_lepton_idx = ak.singletons(ak.argmax(lepton.pt, axis=1)) leading_lepton = lepton[leading_lepton_idx] trailing_lepton_idx = ak.singletons(ak.argmin(lepton.pt, axis=1)) trailing_lepton = lepton[trailing_lepton_idx] dilepton_mass = (leading_lepton+trailing_lepton).mass dilepton_pt = (leading_lepton+trailing_lepton).pt dilepton_dR = delta_r(leading_lepton, trailing_lepton) lepton_pdgId_pt_ordered = ak.fill_none(ak.pad_none(lepton[ak.argsort(lepton.pt, ascending=False)].pdgId, 2, clip=True), 0) if not re.search(re.compile('MuonEG|DoubleMuon|DoubleEG|EGamma'), dataset): n_nonprompt = getNonPromptFromFlavour(electron) + getNonPromptFromFlavour(muon) n_chargeflip = getChargeFlips(electron, ev.GenPart) + getChargeFlips(muon, ev.GenPart) mt_lep_met = mt(lepton.pt, lepton.phi, ev.MET.pt, ev.MET.phi) min_mt_lep_met = ak.min(mt_lep_met, axis=1) ## Tau and other stuff tau = getTaus(ev) track = getIsoTracks(ev) ## Jets jet = getJets(ev, minPt=25, maxEta=4.7, pt_var='pt_nom') jet = jet[ak.argsort(jet.pt_nom, ascending=False)] # need to sort wrt smeared and recorrected jet pt jet = jet[~match(jet, muon, deltaRCut=0.4)] # remove jets that overlap with muons jet = jet[~match(jet, electron, deltaRCut=0.4)] # remove jets that overlap with electrons central = jet[(abs(jet.eta)<2.4)] btag = getBTagsDeepFlavB(jet, year=self.year) # should study working point for DeepJet light = getBTagsDeepFlavB(jet, year=self.year, invert=True) fwd = getFwdJet(light) fwd_noPU = getFwdJet(light, puId=False) high_score_btag = central[ak.argsort(central.btagDeepFlavB)][:,:2] bl = cross(lepton, high_score_btag) bl_dR = delta_r(bl['0'], bl['1']) min_bl_dR = ak.min(bl_dR, axis=1) ## forward jets j_fwd = fwd[ak.singletons(ak.argmax(fwd.p, axis=1))] # highest momentum spectator jf = cross(j_fwd, jet) mjf = (jf['0']+jf['1']).mass j_fwd2 = jf[ak.singletons(ak.argmax(mjf, axis=1))]['1'] # this is the jet that forms the largest invariant mass with j_fwd delta_eta = abs(j_fwd2.eta - j_fwd.eta) ## MET -> can switch to puppi MET met_pt = ev.MET.pt met_phi = ev.MET.phi ## other variables ht = ak.sum(jet.pt, axis=1) st = met_pt + ht + ak.sum(muon.pt, axis=1) + ak.sum(electron.pt, axis=1) # define the weight weight = Weights( len(ev) ) if not re.search(re.compile('MuonEG|DoubleMuon|DoubleEG|EGamma'), dataset): # lumi weight weight.add("weight", ev.weight*cfg['lumi'][self.year]) #weight.add("weight", ev.genWeight*cfg['lumi'][self.year]*mult) # PU weight - not in the babies... weight.add("PU", ev.puWeight, weightUp=ev.puWeightUp, weightDown=ev.puWeightDown, shift=False) # b-tag SFs weight.add("btag", self.btagSF.Method1a(btag, light)) # lepton SFs weight.add("lepton", self.leptonSF.get(electron, muon)) cutflow = Cutflow(output, ev, weight=weight) sel = Selection( dataset = dataset, events = ev, year = self.year, ele = electron, ele_veto = vetoelectron, mu = muon, mu_veto = vetomuon, jet_all = jet, jet_central = central, jet_btag = btag, jet_fwd = fwd, met = ev.MET, ) BL = sel.dilep_baseline(cutflow=cutflow, SS=True) weight_BL = weight.weight()[BL] if True: # define the inputs to the NN # this is super stupid. there must be a better way. NN_inputs = np.stack([ ak.to_numpy(ak.num(jet[BL])), ak.to_numpy(ak.num(tau[BL])), ak.to_numpy(ak.num(track[BL])), ak.to_numpy(st[BL]), ak.to_numpy(ev.MET[BL].pt), ak.to_numpy(ak.max(mjf[BL], axis=1)), ak.to_numpy(pad_and_flatten(delta_eta[BL])), ak.to_numpy(pad_and_flatten(leading_lepton[BL].pt)), ak.to_numpy(pad_and_flatten(leading_lepton[BL].eta)), ak.to_numpy(pad_and_flatten(trailing_lepton[BL].pt)), ak.to_numpy(pad_and_flatten(trailing_lepton[BL].eta)), ak.to_numpy(pad_and_flatten(dilepton_mass[BL])), ak.to_numpy(pad_and_flatten(dilepton_pt[BL])), ak.to_numpy(pad_and_flatten(j_fwd[BL].pt)), ak.to_numpy(pad_and_flatten(j_fwd[BL].p)), ak.to_numpy(pad_and_flatten(j_fwd[BL].eta)), ak.to_numpy(pad_and_flatten(jet[:, 0:1][BL].pt)), ak.to_numpy(pad_and_flatten(jet[:, 1:2][BL].pt)), ak.to_numpy(pad_and_flatten(jet[:, 0:1][BL].eta)), ak.to_numpy(pad_and_flatten(jet[:, 1:2][BL].eta)), ak.to_numpy(pad_and_flatten(high_score_btag[:, 0:1][BL].pt)), ak.to_numpy(pad_and_flatten(high_score_btag[:, 1:2][BL].pt)), ak.to_numpy(pad_and_flatten(high_score_btag[:, 0:1][BL].eta)), ak.to_numpy(pad_and_flatten(high_score_btag[:, 1:2][BL].eta)), ak.to_numpy(min_bl_dR[BL]), ak.to_numpy(min_mt_lep_met[BL]), ]) NN_inputs = np.moveaxis(NN_inputs, 0, 1) model, scaler = load_onnx_model('v8') try: NN_inputs_scaled = scaler.transform(NN_inputs) NN_pred = predict_onnx(model, NN_inputs_scaled) best_score = np.argmax(NN_pred, axis=1) except ValueError: #print ("Empty NN_inputs") NN_pred = np.array([]) best_score = np.array([]) NN_inputs_scaled = NN_inputs #k.clear_session() output['node'].fill(dataset=dataset, multiplicity=best_score, weight=weight_BL) output['node0_score_incl'].fill(dataset=dataset, score=NN_pred[:,0] if np.shape(NN_pred)[0]>0 else np.array([]), weight=weight_BL) output['node0_score'].fill(dataset=dataset, score=NN_pred[best_score==0][:,0] if np.shape(NN_pred)[0]>0 else np.array([]), weight=weight_BL[best_score==0]) output['node1_score'].fill(dataset=dataset, score=NN_pred[best_score==1][:,1] if np.shape(NN_pred)[0]>0 else np.array([]), weight=weight_BL[best_score==1]) output['node2_score'].fill(dataset=dataset, score=NN_pred[best_score==2][:,2] if np.shape(NN_pred)[0]>0 else np.array([]), weight=weight_BL[best_score==2]) output['node3_score'].fill(dataset=dataset, score=NN_pred[best_score==3][:,3] if np.shape(NN_pred)[0]>0 else np.array([]), weight=weight_BL[best_score==3]) output['node4_score'].fill(dataset=dataset, score=NN_pred[best_score==4][:,4] if np.shape(NN_pred)[0]>0 else np.array([]), weight=weight_BL[best_score==4]) SR_sel_pp = ((best_score==0) & ak.flatten((leading_lepton[BL].pdgId<0))) SR_sel_mm = ((best_score==0) & ak.flatten((leading_lepton[BL].pdgId>0))) leading_lepton_BL = leading_lepton[BL] output['lead_lep_SR_pp'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(leading_lepton_BL[SR_sel_pp].pt)), weight = weight_BL[SR_sel_pp] ) output['lead_lep_SR_mm'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(leading_lepton_BL[SR_sel_mm].pt)), weight = weight_BL[SR_sel_mm] ) del model del scaler del NN_inputs, NN_inputs_scaled, NN_pred # first, make a few super inclusive plots output['PV_npvs'].fill(dataset=dataset, multiplicity=ev.PV[BL].npvs, weight=weight_BL) output['PV_npvsGood'].fill(dataset=dataset, multiplicity=ev.PV[BL].npvsGood, weight=weight_BL) output['N_jet'].fill(dataset=dataset, multiplicity=ak.num(jet)[BL], weight=weight_BL) output['N_b'].fill(dataset=dataset, multiplicity=ak.num(btag)[BL], weight=weight_BL) output['N_central'].fill(dataset=dataset, multiplicity=ak.num(central)[BL], weight=weight_BL) output['N_ele'].fill(dataset=dataset, multiplicity=ak.num(electron)[BL], weight=weight_BL) output['N_mu'].fill(dataset=dataset, multiplicity=ak.num(electron)[BL], weight=weight_BL) output['N_fwd'].fill(dataset=dataset, multiplicity=ak.num(fwd)[BL], weight=weight_BL) output['ST'].fill(dataset=dataset, pt=st[BL], weight=weight_BL) output['HT'].fill(dataset=dataset, pt=ht[BL], weight=weight_BL) if not re.search(re.compile('MuonEG|DoubleMuon|DoubleEG|EGamma'), dataset): output['nLepFromTop'].fill(dataset=dataset, multiplicity=ev[BL].nLepFromTop, weight=weight_BL) output['nLepFromTau'].fill(dataset=dataset, multiplicity=ev.nLepFromTau[BL], weight=weight_BL) output['nLepFromZ'].fill(dataset=dataset, multiplicity=ev.nLepFromZ[BL], weight=weight_BL) output['nLepFromW'].fill(dataset=dataset, multiplicity=ev.nLepFromW[BL], weight=weight_BL) output['nGenTau'].fill(dataset=dataset, multiplicity=ev.nGenTau[BL], weight=weight_BL) output['nGenL'].fill(dataset=dataset, multiplicity=ak.num(ev.GenL[BL], axis=1), weight=weight_BL) output['chargeFlip_vs_nonprompt'].fill(dataset=dataset, n1=n_chargeflip[BL], n2=n_nonprompt[BL], n_ele=ak.num(electron)[BL], weight=weight_BL) output['MET'].fill( dataset = dataset, pt = ev.MET[BL].pt, phi = ev.MET[BL].phi, weight = weight_BL ) if not re.search(re.compile('MuonEG|DoubleMuon|DoubleEG|EGamma'), dataset): output['lead_gen_lep'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(leading_gen_lep[BL].pt)), eta = ak.to_numpy(ak.flatten(leading_gen_lep[BL].eta)), phi = ak.to_numpy(ak.flatten(leading_gen_lep[BL].phi)), weight = weight_BL ) output['trail_gen_lep'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(trailing_gen_lep[BL].pt)), eta = ak.to_numpy(ak.flatten(trailing_gen_lep[BL].eta)), phi = ak.to_numpy(ak.flatten(trailing_gen_lep[BL].phi)), weight = weight_BL ) output['lead_lep'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(leading_lepton[BL].pt)), eta = ak.to_numpy(ak.flatten(leading_lepton[BL].eta)), phi = ak.to_numpy(ak.flatten(leading_lepton[BL].phi)), weight = weight_BL ) output['trail_lep'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(trailing_lepton[BL].pt)), eta = ak.to_numpy(ak.flatten(trailing_lepton[BL].eta)), phi = ak.to_numpy(ak.flatten(trailing_lepton[BL].phi)), weight = weight_BL ) output['j1'].fill( dataset = dataset, pt = ak.flatten(jet.pt_nom[:, 0:1][BL]), eta = ak.flatten(jet.eta[:, 0:1][BL]), phi = ak.flatten(jet.phi[:, 0:1][BL]), weight = weight_BL ) output['j2'].fill( dataset = dataset, pt = ak.flatten(jet[:, 1:2][BL].pt_nom), eta = ak.flatten(jet[:, 1:2][BL].eta), phi = ak.flatten(jet[:, 1:2][BL].phi), weight = weight_BL ) output['j3'].fill( dataset = dataset, pt = ak.flatten(jet[:, 2:3][BL].pt_nom), eta = ak.flatten(jet[:, 2:3][BL].eta), phi = ak.flatten(jet[:, 2:3][BL].phi), weight = weight_BL ) output['fwd_jet'].fill( dataset = dataset, pt = ak.flatten(j_fwd[BL].pt), eta = ak.flatten(j_fwd[BL].eta), phi = ak.flatten(j_fwd[BL].phi), weight = weight_BL ) output['high_p_fwd_p'].fill(dataset=dataset, p = ak.flatten(j_fwd[BL].p), weight = weight_BL) return output
if not (tracks_boosted_minus3.size == 0): s3 = eventShapesUtilities.sphericityTensor(tracks_boosted_minus3) if np.isfinite(s3).all(): evtShape3[ievt] = eventShapesUtilities.sphericity(s3) if not (tracks_boosted_minus4.size == 0): s4 = eventShapesUtilities.sphericityTensor(tracks_boosted_minus4) if np.isfinite(s4).all(): evtShape4[ievt] = eventShapesUtilities.sphericity(s4) if not (tracks_boosted_minus5.size == 0): s5 = eventShapesUtilities.sphericityTensor(tracks_boosted_minus5) if np.isfinite(s5).all(): evtShape5[ievt] = eventShapesUtilities.sphericity(s5) CrossSection = ak.to_numpy(events['CrossSection'][events['HT'] > 1200]) evtShape = evtShape[events['HT'] > 1200] evtShape1 = evtShape1[events['HT'] > 1200] evtShape2 = evtShape2[events['HT'] > 1200] evtShape3 = evtShape3[events['HT'] > 1200] evtShape4 = evtShape4[events['HT'] > 1200] evtShape5 = evtShape5[events['HT'] > 1200] all_shapes = ((evtShape > -1) & (evtShape1 > -1) & (evtShape2 > -1) & (evtShape3 > -1) & (evtShape4 > -1) & (evtShape5 > -1)) CrossSection = CrossSection[all_shapes] evtShape = evtShape[all_shapes] evtShape1 = evtShape1[all_shapes] evtShape2 = evtShape2[all_shapes] evtShape3 = evtShape3[all_shapes]
def process(self, events): output = self.accumulator.identity() # we can use a very loose preselection to filter the events. nothing is done with this presel, though presel = ak.num(events.Jet) >= 2 ev = events[presel] dataset = ev.metadata['dataset'] # load the config - probably not needed anymore cfg = loadConfig() output['totalEvents']['all'] += len(events) output['skimmedEvents']['all'] += len(ev) ## Electrons electron = Collections(ev, "Electron", "tight").get() electron = electron[(electron.pt > 20) & (abs(electron.eta) < 2.4)] electron = electron[((electron.genPartIdx >= 0) & (np.abs(electron.matched_gen.pdgId) == 11) )] #from here on all leptons are gen-matched ##Muons muon = Collections(ev, "Muon", "tight").get() muon = muon[(muon.pt > 20) & (abs(muon.eta) < 2.4)] muon = muon[((muon.genPartIdx >= 0) & (np.abs(muon.matched_gen.pdgId) == 13))] ##Leptons lepton = ak.concatenate([muon, electron], axis=1) SSlepton = (ak.sum(lepton.charge, axis=1) != 0) & (ak.num(lepton) == 2) OSlepton = (ak.sum(lepton.charge, axis=1) == 0) & (ak.num(lepton) == 2) leading_lepton_idx = ak.singletons(ak.argmax(lepton.pt, axis=1)) leading_lepton = lepton[leading_lepton_idx] #jets jet = getJets(ev, minPt=40, maxEta=2.4, pt_var='pt') jet = jet[ak.argsort( jet.pt, ascending=False )] # need to sort wrt smeared and recorrected jet pt jet = jet[~match(jet, muon, deltaRCut=0.4)] # remove jets that overlap with muons jet = jet[~match(jet, electron, deltaRCut=0.4)] ## MET -> can switch to puppi MET met_pt = ev.MET.pt met_phi = ev.MET.phi # setting up the various weights weight = Weights(len(ev)) weight2 = Weights(len(ev)) if not dataset == 'MuonEG': # generator weight weight.add("weight", ev.genWeight) weight2.add("weight", ev.genWeight) weight2.add("charge flip", self.charge_flip_ratio.flip_weight(electron)) #selections filters = getFilters(ev, year=self.year, dataset=dataset) ss = (SSlepton) os = (OSlepton) jet_all = (ak.num(jet) >= 2) selection = PackedSelection() selection.add('filter', (filters)) selection.add('ss', ss) selection.add('os', os) selection.add('jet', jet_all) bl_reqs = ['filter', 'jet'] bl_reqs_d = {sel: True for sel in bl_reqs} baseline = selection.require(**bl_reqs_d) s_reqs = bl_reqs + ['ss'] s_reqs_d = {sel: True for sel in s_reqs} ss_sel = selection.require(**s_reqs_d) o_reqs = bl_reqs + ['os'] o_reqs_d = {sel: True for sel in o_reqs} os_sel = selection.require(**o_reqs_d) #outputs output['N_jet'].fill(dataset=dataset, multiplicity=ak.num(jet)[baseline], weight=weight.weight()[baseline]) output['N_ele'].fill(dataset=dataset, multiplicity=ak.num(lepton)[ss_sel], weight=weight.weight()[ss_sel]) output['N_ele2'].fill(dataset=dataset, multiplicity=ak.num(lepton)[os_sel], weight=weight2.weight()[os_sel]) output["electron"].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(leading_lepton[ss_sel].pt)), eta=abs(ak.to_numpy(ak.flatten(leading_lepton[ss_sel].eta))), phi=ak.to_numpy(ak.flatten(leading_lepton[ss_sel].phi)), weight=weight.weight()[ss_sel]) output["electron2"].fill( dataset=dataset, pt=ak.to_numpy(ak.flatten(leading_lepton[os_sel].pt)), eta=abs(ak.to_numpy(ak.flatten(leading_lepton[os_sel].eta))), phi=ak.to_numpy(ak.flatten(leading_lepton[os_sel].phi)), weight=weight2.weight()[os_sel]) return output
def run(self, decay, ncpu=5): nevents_real = 0 import time start_time = time.time() processEvents = {} for pr in self.processes: fin = self.baseDir + pr + '.root' #input file if not os.path.isfile(fin): print('file ', fin, ' does not exist. exit') exit(3) tfin = ROOT.TFile.Open(fin) tfin.cd() found = False for key in tfin.GetListOfKeys(): if 'eventsProcessed' == key.GetName(): events = tfin.eventsProcessed.GetVal() processEvents[pr] = events found = True if not found: processEvents[pr] = 1 tfin.Close() for pr in self.processes: print(' running over process : ', pr) fin = self.baseDir + pr + '.root' #input file fout = self.baseDir + pr + '.root' #output file for tree fhisto = self.baseDir + pr + '_histo.root' #output file for histograms file = uproot.open(fin) tree = file['events'] #Number of events to keep and analyse n_events = 1000 #Container for the reco particles p_c = 'RP' events = tree.arrays(library="ak", how="zip", filter_name=f"{p_c}*")[:n_events] print('events loaded ', events) p = events[p_c] print('fin ', fin) p_c = 'RP' #Number of events to keep and analyse n_events = 10000 file = uproot.open(fin) tree = file[self.treename] print('here-1.2') events = tree.arrays(library="ak", how="zip", filter_name=f"{p_c}*")[:n_events] print('here0') #Container for the reco particles p = events[p_c][:n_events] print('here1') p["p"] = kinematics_flat.calc_p(p) p_cut = p["p"] > 1. p = p[p_cut] print('here2') pi_cut = abs(p["mass"] - lp.pi_plus.mass / 1000.) < 1e-4 pi = p[pi_cut] print('here3') k_cut = abs(p["mass"] - lp.K_plus.mass / 1000.) < 1e-4 k = p[k_cut] print('here4') D = ak.cartesian({"k": k, "pi": pi}) D_cut = np.sign(D["k", "charge"]) != np.sign(D["pi", "charge"]) D = D[D_cut] print('here5') PDG_K_m = lp.K_plus.mass / 1000. PDG_pi_m = lp.pi_plus.mass / 1000. D["mass"] = kinematics_flat.mass([D["k"], D["pi"]], [PDG_K_m, PDG_pi_m]) print('here6') PDG_D_m = lp.D_0.mass / 1000. D_window = 0.05 D_cut = abs(D["mass"] - PDG_D_m) < D_window D = D[D_cut] print('here7') B = ak.cartesian({"D_k": D["k"], "D_pi": D["pi"], "pi": pi}) B_cut = np.sign(B["D_k", "charge"]) == np.sign(B["pi", "charge"]) B = B[B_cut] B["mass"] = kinematics_flat.mass([B["D_k"], B["D_pi"], B["pi"]], [PDG_K_m, PDG_pi_m, PDG_pi_m]) print('here8') data_np = ak.to_numpy(ak.flatten(B["mass"])) print(data_np) validfile = self.testfile(fout) if not validfile: continue nevents_real += df_cut.Count().GetValue() tf = ROOT.TFile.Open(fhisto, 'RECREATE') for v in self.variables: model = ROOT.RDF.TH1DModel( v, ";{};".format(self.variables[v]["title"]), self.variables[v]["bin"], self.variables[v]["xmin"], self.variables[v]["xmax"]) h = snapshot_tdf.Histo1D(model, self.variables[v]["name"]) try: h.Scale(1. * self.procDict[pr]["crossSection"] * self.procDict[pr]["kfactor"] * self.procDict[pr]["matchingEfficiency"] / processEvents[pr]) except KeyError: h.Scale(1. / h.Integral(0, -1)) h.Write() tf.Close() elapsed_time = time.time() - start_time print( '==============================SUMMARY==============================' ) print('Elapsed time (H:M:S) : ', time.strftime("%H:%M:%S", time.gmtime(elapsed_time))) print('Events Processed/Second : ', int(nevents_real / elapsed_time)) print('Total Events Processed : ', nevents_real) print( '===================================================================' )
def fill_numpy(record): """this function reads the awkward array and edits for our needs: Projecting into 32x30 grid""" #defined binning binX = np.arange(-81, 82, 5.088333) #binZ = np.arange(-77, 78, 5.088333) #new z-shifted bin in Z #binZ = np.arange(-70, 83, 5.088333) #new z- asymmetric bin #binZ = np.arange(-90, 65, 5.088333) #binZ = np.arange(-80, 75, 5.088333) #binZ = np.arange(-65, 90, 5.088333) #binZ = np.arange(-45, 110, 5.088333) #binZ = np.arange(-35, 121, 5.088333) #binZ with 40 #binZ = np.arange(-100, 105, 5.088333) #binZ = np.arange(-104, 106, 5.088333) #binZ = np.arange(-97, 111, 5.088333) #binZ = np.arange(-104.3, 101.2, 5.088333) binZ = np.arange(-115, 100, 5.088333) ## Unable to escape using python list here. But we can live with that. l = [] E = [] theta = [] for i in range(0, nevents): #Get hits and convert them into numpy array z = ak.to_numpy(record[i].z) x = ak.to_numpy(record[i].x) y = ak.to_numpy(record[i].y) e = ak.to_numpy(record[i].e) #Get indicent photon energies incE = ak.to_numpy(record[i].E) E.append(incE.compressed()[0]) #Get polar angle theta of the gun inTh = ak.to_numpy(record[i].theta) theta.append(inTh.compressed()[0]) layers = [] #loop over layers and project them into 2d grid. for j in range(0, 30): idx = np.where((y <= (hmap[j] + 0.9999)) & (y > (hmap[j] + 0.0001))) xlayer = x.take(idx)[0] zlayer = z.take(idx)[0] elayer = e.take(idx)[0] H, xedges, yedges = np.histogram2d(xlayer, zlayer, bins=(binX, binZ), weights=elayer) layers.append(H) l.append(layers) ## convert them into numpy array shower = np.asarray(l) e0 = np.reshape(np.asarray(E), (-1, 1)) t0 = np.reshape(np.asarray(theta), (-1, 1)) return shower, e0, t0
fig.savefig(f'valplots/{hist_spec[0]}.png') print("Checking Coffea/NanoEvents compatibility:") import coffea import uproot import awkward1 as ak import numpy as np from coffea.nanoevents import NanoEventsFactory for fn in check_files: print(f" {fn}") def nano_evts(fname): factory = NanoEventsFactory.from_file( fname, entry_start=0, entry_stop=10000, metadata={"dataset": ""}, ) return factory.events() evts = nano_evts(fn) print(" Subjets pt ordered (should be false):", ak.all(evts.SubJet.pt[:, :-1] - evts.SubJet.pt[:, 1:] >= 0)) print(" SDmass from subjets hist:") print( np.histogram(ak.to_numpy( ak.flatten(evts.FatJet.subjets.sum().mass, None)), bins=np.linspace(40, 200, 20))[0])
def from_uproot( ntuple_paths: List[pathlib.Path], pos_in_file: str, variable: str, bins: np.ndarray, weight: Optional[str] = None, selection_filter: Optional[str] = None, ) -> Tuple[np.ndarray, np.ndarray]: """Reads an ntuple with uproot, and fills a histogram with the observable. The paths may contain wildcards. Args: ntuple_paths (List[pathlib.Path]): list of paths to ntuples pos_in_file (str): name of tree within ntuple variable (str): variable to bin histogram in bins (numpy.ndarray): bin edges for histogram weight (Optional[str], optional): event weight to extract, defaults to None (no weights applied) selection_filter (Optional[str], optional): filter to be applied on events, defaults to None (no filter) Returns: Tuple[np.ndarray, np.ndarray]: - yield per bin - stat. uncertainty per bin """ # concatenate the path to file and location within file with ":" paths_with_trees = [str(path) + ":" + pos_in_file for path in ntuple_paths] # determine whether the weight is a float or an expression # (for which a branch needs to be read) if weight is not None: try: float(weight) weight_is_expression = False except ValueError: # weight is not a float, need to evaluate the expression weight_is_expression = True else: # no weight specified, all weights are 1.0 weight_is_expression = False weight = "1.0" if weight_is_expression: # need to read observables and weights array_generator = uproot.iterate( paths_with_trees, expressions=[variable, weight], cut=selection_filter, ) obs_list = [] weight_list = [] for arr in array_generator: obs_list.append(ak.to_numpy(arr[variable])) weight_list.append(ak.to_numpy(arr[weight])) observables = np.concatenate(obs_list) weights = np.concatenate(weight_list) else: # only need to read the observables array_generator = uproot.iterate( paths_with_trees, expressions=[variable], cut=selection_filter, ) obs_list = [] for arr in array_generator: obs_list.append(ak.to_numpy(arr[variable])) observables = np.concatenate(obs_list) weights = np.ones_like(observables) * float(weight) yields, stdev = _bin_data(observables, weights, bins) return yields, stdev
def process(self, events): output = self.accumulator.identity() # use a very loose preselection to filter the events presel = ak.num(events.Jet)>2 ev = events[presel] dataset = ev.metadata['dataset'] # load the config - probably not needed anymore cfg = loadConfig() output['totalEvents']['all'] += len(events) output['skimmedEvents']['all'] += len(ev) ## Generated leptons gen_lep = ev.GenL leading_gen_lep = gen_lep[ak.singletons(ak.argmax(gen_lep.pt, axis=1))] trailing_gen_lep = gen_lep[ak.singletons(ak.argmin(gen_lep.pt, axis=1))] ## Muons muon = Collections(ev, "Muon", "tightSSTTH").get() vetomuon = Collections(ev, "Muon", "vetoTTH").get() dimuon = choose(muon, 2) SSmuon = ak.any((dimuon['0'].charge * dimuon['1'].charge)>0, axis=1) leading_muon_idx = ak.singletons(ak.argmax(muon.pt, axis=1)) leading_muon = muon[leading_muon_idx] ## Electrons electron = Collections(ev, "Electron", "tightSSTTH").get() vetoelectron = Collections(ev, "Electron", "vetoTTH").get() dielectron = choose(electron, 2) SSelectron = ak.any((dielectron['0'].charge * dielectron['1'].charge)>0, axis=1) leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1)) leading_electron = electron[leading_electron_idx] ## Merge electrons and muons - this should work better now in ak1 dilepton = cross(muon, electron) SSlepton = ak.any((dilepton['0'].charge * dilepton['1'].charge)>0, axis=1) lepton = ak.concatenate([muon, electron], axis=1) leading_lepton_idx = ak.singletons(ak.argmax(lepton.pt, axis=1)) leading_lepton = lepton[leading_lepton_idx] trailing_lepton_idx = ak.singletons(ak.argmin(lepton.pt, axis=1)) trailing_lepton = lepton[trailing_lepton_idx] n_nonprompt = getNonPromptFromFlavour(electron) + getNonPromptFromFlavour(muon) n_chargeflip = getChargeFlips(electron, ev.GenPart) + getChargeFlips(muon, ev.GenPart) ## Jets jet = getJets(ev, minPt=25, maxEta=4.7, pt_var='pt_nom') jet = jet[ak.argsort(jet.pt_nom, ascending=False)] # need to sort wrt smeared and recorrected jet pt jet = jet[~match(jet, muon, deltaRCut=0.4)] # remove jets that overlap with muons jet = jet[~match(jet, electron, deltaRCut=0.4)] # remove jets that overlap with electrons central = jet[(abs(jet.eta)<2.4)] btag = getBTagsDeepFlavB(jet, year=self.year) # should study working point for DeepJet light = getBTagsDeepFlavB(jet, year=self.year, invert=True) fwd = getFwdJet(light) fwd_noPU = getFwdJet(light, puId=False) ## forward jets j_fwd = fwd[ak.singletons(ak.argmax(fwd.p, axis=1))] # highest momentum spectator jf = cross(j_fwd, jet) mjf = (jf['0']+jf['1']).mass j_fwd2 = jf[ak.singletons(ak.argmax(mjf, axis=1))]['1'] # this is the jet that forms the largest invariant mass with j_fwd delta_eta = abs(j_fwd2.eta - j_fwd.eta) ## MET -> can switch to puppi MET met_pt = ev.MET.pt met_phi = ev.MET.phi ## other variables ht = ak.sum(jet.pt, axis=1) st = met_pt + ht + ak.sum(muon.pt, axis=1) + ak.sum(electron.pt, axis=1) ## event selectors filters = getFilters(ev, year=self.year, dataset=dataset) dilep = ((ak.num(electron) + ak.num(muon))==2) pos_charge = ((ak.sum(electron.pdgId, axis=1) + ak.sum(muon.pdgId, axis=1))<0) neg_charge = ((ak.sum(electron.pdgId, axis=1) + ak.sum(muon.pdgId, axis=1))>0) lep0pt = ((ak.num(electron[(electron.pt>25)]) + ak.num(muon[(muon.pt>25)]))>0) lep0pt_40 = ((ak.num(electron[(electron.pt>40)]) + ak.num(muon[(muon.pt>40)]))>0) lep0pt_100 = ((ak.num(electron[(electron.pt>100)]) + ak.num(muon[(muon.pt>100)]))>0) lep1pt = ((ak.num(electron[(electron.pt>20)]) + ak.num(muon[(muon.pt>20)]))>1) lep1pt_30 = ((ak.num(electron[(electron.pt>30)]) + ak.num(muon[(muon.pt>30)]))>1) lepveto = ((ak.num(vetoelectron) + ak.num(vetomuon))==2) # define the weight weight = Weights( len(ev) ) #mult = 1 #if dataset=='inclusive': mult = 0.0478/47.448 #if dataset=='plus': mult = 0.0036/7.205 if not dataset=='MuonEG': # lumi weight weight.add("weight", ev.weight*cfg['lumi'][self.year]) #weight.add("weight", ev.genWeight*cfg['lumi'][self.year]*mult) # PU weight - not in the babies... weight.add("PU", ev.puWeight, weightUp=ev.puWeightUp, weightDown=ev.puWeightDown, shift=False) # b-tag SFs weight.add("btag", self.btagSF.Method1a(btag, light)) # lepton SFs weight.add("lepton", self.leptonSF.get(electron, muon)) selection = PackedSelection() selection.add('lepveto', lepveto) selection.add('dilep', dilep ) selection.add('filter', (filters) ) selection.add('p_T(lep0)>25', lep0pt ) selection.add('p_T(lep0)>40', lep0pt_40 ) selection.add('p_T(lep1)>20', lep1pt ) selection.add('p_T(lep1)>30', lep1pt_30 ) selection.add('SS', ( SSlepton | SSelectron | SSmuon) ) selection.add('pos', ( pos_charge ) ) selection.add('neg', ( neg_charge ) ) selection.add('N_jet>3', (ak.num(jet)>=4) ) selection.add('N_jet>4', (ak.num(jet)>=5) ) selection.add('N_central>2', (ak.num(central)>=3) ) selection.add('N_central>3', (ak.num(central)>=4) ) selection.add('N_btag>0', (ak.num(btag)>=1) ) selection.add('MET>50', (ev.MET.pt>50) ) selection.add('ST', (st>600) ) selection.add('N_fwd>0', (ak.num(fwd)>=1 )) selection.add('delta_eta', (ak.any(delta_eta>2, axis=1) ) ) selection.add('fwd_p>500', (ak.any(j_fwd.p>500, axis=1) ) ) ss_reqs = ['lepveto', 'dilep', 'SS', 'filter', 'p_T(lep0)>25', 'p_T(lep1)>20', 'N_jet>3', 'N_central>2', 'N_btag>0'] bl_reqs = ss_reqs + ['N_fwd>0', 'N_jet>4', 'N_central>3', 'ST', 'MET>50', 'delta_eta'] sr_reqs = bl_reqs + ['fwd_p>500', 'p_T(lep0)>40', 'p_T(lep1)>30'] ss_reqs_d = { sel: True for sel in ss_reqs } ss_selection = selection.require(**ss_reqs_d) bl_reqs_d = { sel: True for sel in bl_reqs } BL = selection.require(**bl_reqs_d) sr_reqs_d = { sel: True for sel in sr_reqs } SR = selection.require(**sr_reqs_d) cutflow = Cutflow(output, ev, weight=weight) cutflow_reqs_d = {} for req in sr_reqs: cutflow_reqs_d.update({req: True}) cutflow.addRow( req, selection.require(**cutflow_reqs_d) ) # first, make a few super inclusive plots output['PV_npvs'].fill(dataset=dataset, multiplicity=ev.PV[ss_selection].npvs, weight=weight.weight()[ss_selection]) output['PV_npvsGood'].fill(dataset=dataset, multiplicity=ev.PV[ss_selection].npvsGood, weight=weight.weight()[ss_selection]) output['N_jet'].fill(dataset=dataset, multiplicity=ak.num(jet)[ss_selection], weight=weight.weight()[ss_selection]) output['N_b'].fill(dataset=dataset, multiplicity=ak.num(btag)[ss_selection], weight=weight.weight()[ss_selection]) output['N_central'].fill(dataset=dataset, multiplicity=ak.num(central)[ss_selection], weight=weight.weight()[ss_selection]) output['N_ele'].fill(dataset=dataset, multiplicity=ak.num(electron)[ss_selection], weight=weight.weight()[ss_selection]) output['N_mu'].fill(dataset=dataset, multiplicity=ak.num(electron)[ss_selection], weight=weight.weight()[ss_selection]) output['N_fwd'].fill(dataset=dataset, multiplicity=ak.num(fwd)[ss_selection], weight=weight.weight()[ss_selection]) output['nLepFromTop'].fill(dataset=dataset, multiplicity=ev[BL].nLepFromTop, weight=weight.weight()[BL]) output['nLepFromTau'].fill(dataset=dataset, multiplicity=ev.nLepFromTau[BL], weight=weight.weight()[BL]) output['nLepFromZ'].fill(dataset=dataset, multiplicity=ev.nLepFromZ[BL], weight=weight.weight()[BL]) output['nLepFromW'].fill(dataset=dataset, multiplicity=ev.nLepFromW[BL], weight=weight.weight()[BL]) output['nGenTau'].fill(dataset=dataset, multiplicity=ev.nGenTau[BL], weight=weight.weight()[BL]) output['nGenL'].fill(dataset=dataset, multiplicity=ak.num(ev.GenL[BL], axis=1), weight=weight.weight()[BL]) output['chargeFlip_vs_nonprompt'].fill(dataset=dataset, n1=n_chargeflip[ss_selection], n2=n_nonprompt[ss_selection], n_ele=ak.num(electron)[ss_selection], weight=weight.weight()[ss_selection]) output['MET'].fill( dataset = dataset, pt = ev.MET[ss_selection].pt, phi = ev.MET[ss_selection].phi, weight = weight.weight()[ss_selection] ) output['lead_gen_lep'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(leading_gen_lep[BL].pt)), eta = ak.to_numpy(ak.flatten(leading_gen_lep[BL].eta)), phi = ak.to_numpy(ak.flatten(leading_gen_lep[BL].phi)), weight = weight.weight()[BL] ) output['trail_gen_lep'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(trailing_gen_lep[BL].pt)), eta = ak.to_numpy(ak.flatten(trailing_gen_lep[BL].eta)), phi = ak.to_numpy(ak.flatten(trailing_gen_lep[BL].phi)), weight = weight.weight()[BL] ) output['lead_lep'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(leading_lepton[BL].pt)), eta = ak.to_numpy(ak.flatten(leading_lepton[BL].eta)), phi = ak.to_numpy(ak.flatten(leading_lepton[BL].phi)), weight = weight.weight()[BL] ) output['trail_lep'].fill( dataset = dataset, pt = ak.to_numpy(ak.flatten(trailing_lepton[BL].pt)), eta = ak.to_numpy(ak.flatten(trailing_lepton[BL].eta)), phi = ak.to_numpy(ak.flatten(trailing_lepton[BL].phi)), weight = weight.weight()[BL] ) output['j1'].fill( dataset = dataset, pt = ak.flatten(jet.pt_nom[:, 0:1][BL]), eta = ak.flatten(jet.eta[:, 0:1][BL]), phi = ak.flatten(jet.phi[:, 0:1][BL]), weight = weight.weight()[BL] ) output['j2'].fill( dataset = dataset, pt = ak.flatten(jet[:, 1:2][BL].pt_nom), eta = ak.flatten(jet[:, 1:2][BL].eta), phi = ak.flatten(jet[:, 1:2][BL].phi), weight = weight.weight()[BL] ) output['j3'].fill( dataset = dataset, pt = ak.flatten(jet[:, 2:3][BL].pt_nom), eta = ak.flatten(jet[:, 2:3][BL].eta), phi = ak.flatten(jet[:, 2:3][BL].phi), weight = weight.weight()[BL] ) return output