def process(self, events): output = self.accumulator.identity() dataset = events.metadata['dataset'] integratedLuminosity = 137.19*1000 # fb^{-1} to pb^{-1} ht = events.HT weights = integratedLuminosity*events.CrossSection[ht > 1200]/len(events) GenParticles = events.GenParticles finalParticles = (GenParticles.Status == 1) & (GenParticles.pt > 1) & (abs(GenParticles.eta) < 2.5) & (GenParticles.Charge != 0) nTracksGEN = ak.sum(finalParticles[ht > 1200], axis=1) tracks = events.Tracks tracks_pt = np.sqrt(tracks.x**2 + tracks.y**2) tracks_eta = np.arcsinh(tracks.z / tracks_pt) track_cut = (tracks_pt > 1.) & (abs(tracks_eta) < 2.5) & (tracks.fromPV0 >= 2) & tracks.matchedToPFCandidate nTracksRECO = ak.to_numpy(ak.sum(track_cut[ht > 1200], axis=1)) output["sumw"][dataset] += len(events) output["nTracksHist"].fill( dataset="CMSSW GEN", nTracks=nTracksGEN, weight=weights ) output["nTracksHist"].fill( dataset="CMSSW RECO", nTracks=nTracksRECO, weight=weights ) return output
class CutStats(): def __init__(self, **kwargs): self.debug = kwargs.get("debug") self.event_stats = {} self.object_stats = {} self.objects = {} def add_initial_events(self, sample, n_events): if sample not in self.event_stats.keys(): self.event_stats[sample] = { "n_events_initial" : n_events } else: self.event_stats[sample]["n_events_initial"] += n_events def add_initial_objects(self, sample, object, n_objects): if object not in self.object_stats.keys(): self.object_stats[object] = {} if sample not in self.object_stats[object].keys(): self.object_stats[object][sample] = { "n_objects_initial" : n_objects_initial } else: self.object_stats[object][sample]["n_objects_initial"] += n_objects_initial def add_event_cuts(self, events cuts, names, sample): if sample not in self.event_stats.keys(): self.event_stats[sample] = { "n_events_initial" : n_events } else: self.event_stats[sample]["n_events_initial"] += n_events for cut, name in zip(cuts, names): if name not in self.event_stats[sample].keys(): self.event_stats[sample][name] = awkward.sum(events[cut]) else: self.event_stats[sample][name] += awkward.sum(events[cut])
def rescale(accumulator, xsecs=xsecs, lumi=lumi, data="BTagMu"): #def rescale(accumulator, xsecs=xsecs, data="BTagMu"): """Scale by lumi""" #lumi = 1000*lumi # Convert lumi to pb^-1 from coffea import hist scale = {} sumxsecs = ak.sum(xsecs.values()) #N_data = accumulator['nbtagmu_event_level'][data] #print("N_data =", N_data) for dataset, N_mc in collections.OrderedDict(sorted(accumulator['sumw'].items())).items(): if dataset in xsecs: print(" ", dataset, "\t", N_mc, "events\t", xsecs[dataset], "pb") #scale[dataset] = (xsecs[dataset]/sumxsecs)*(N_data/N_mc) scale[dataset] = (xsecs[dataset]*lumi)/N_mc else: print(" ", "X ", dataset) scale[dataset] = 0#lumi / N_mc print(scale) datasets_mc = [item for item in list(xsecs.keys()) if not 'GluGlu' in item] for h in accumulator.values(): if isinstance(h, hist.Hist): h.scale(scale, axis="dataset") N_data = ak.sum(h[data].values().values()) N_mc = ak.sum(h[datasets_mc].sum('dataset', 'flavor').values().values()) #scaletodata = dict(zip(scale.keys(), len(scale)*[1./N_data])) scaletodata = dict(zip(scale.keys(), len(scale)*[N_data/N_mc])) h.scale(scaletodata, axis="dataset") return accumulator
def ggTauTau_inclusive_preselection(events, photons, electrons, muons, taus, options, debug): cut_diagnostics = utils.CutDiagnostics(events = events, debug = debug, cut_set = "[analysis_selections.py : ggTauTau_inclusive_preselection]") # Get number of electrons, muons, taus selected_electrons = electrons[lepton_selections.select_electrons(events, photons, electrons, options, debug)] selected_muons = muons[lepton_selections.select_muons(events, photons, muons, options, debug)] selected_taus = taus[tau_selections.select_taus(events, photons, selected_muons, selected_electrons, taus, options, debug)] n_electrons = awkward.num(selected_electrons) n_muons = awkward.num(selected_muons) n_taus = awkward.num(selected_taus) # Require >= 1 lep/tau n_leptons_and_taus = n_electrons + n_muons + n_taus lep_tau_cut = n_leptons_and_taus >= options["n_leptons_and_taus"] # Require OS leptons/taus for events with 2 leptons/taus sum_charge = awkward.sum(selected_electrons.charge, axis=1) + awkward.sum(selected_muons.charge, axis=1) + awkward.sum(selected_taus.charge, axis=1) charge_cut = sum_charge == 0 two_leptons = n_leptons_and_taus == 2 not_two_leptons = n_leptons_and_taus != 2 os_cut = (two_leptons & charge_cut) | not_two_leptons # only require 2 OS leptons if there are ==2 leptons in the event all_cuts = lep_tau_cut & os_cut cut_diagnostics.add_cuts([lep_tau_cut, os_cut, all_cuts], ["N_leptons + N_taus >= 1", "OS dileptons", "all"]) return events[all_cuts], photons[all_cuts], selected_electrons[all_cuts], selected_muons[all_cuts], selected_taus[all_cuts]
def compute_eff_algo(tau_1, tau_2, a, Pt_thr): eff_presel = ak.sum(ak.num(tau_1, axis=-1)) deepTau_thr_1 = deep_thr(tau_1, a, Pt_thr) deepTau_thr_2 = deep_thr(tau_2, a, Pt_thr) deepTau_mask_1 = tau_1.deepTau_VSjet > deepTau_thr_1 deepTau_mask_2 = tau_2.deepTau_VSjet > deepTau_thr_2 eff = ak.sum(ditau_selection(deepTau_mask_1, deepTau_mask_2)) return compute_eff_witherr(eff, eff_presel)
def test_more(): nparray = np.array( [np.datetime64("2021-06-03T10:00"), np.datetime64("2021-06-03T11:00")] ) akarray = ak.Array(nparray) assert (akarray[1:] - akarray[:-1]).tolist() == [np.timedelta64(60, "m")] assert ak.sum(akarray[1:] - akarray[:-1]) == np.timedelta64(60, "m") assert ak.sum(akarray[1:] - akarray[:-1], axis=0) == [np.timedelta64(60, "m")]
def test(): array = ak.Array([ [0, 1, None], [None, 3], [], None, [4, 5, None, 6], ]) assert ak.sum(array, axis=-1).tolist() == [ 1, # 0 + 1 3, # 3 0, # list is empty None, # list is missing 15, # 4 + 5 + 6 ] assert ak.sum(array, axis=-2).tolist() == [ 4, 9, 0, 6, # 0+4 # 1+3+5 # no data # 6 ] assert ak.min(array, axis=-1).tolist() == [ 0, # min([0, 1]) 3, # min([3]) None, # list is empty None, # list is missing 4, # min([4, 5, 6]) ] assert ak.min(array, axis=-2).tolist() == [ 0, 1, None, 6, # 0,4 # 1,3,5 # no data # 6 ] # first bug-fix: single '?' assert str(ak.min(array, axis=-1).type) == "5 * ?int64" # second bug-fix: correct mask_identity=False behavior assert ak.ptp(array, axis=-1).tolist() == [1, 0, None, None, 2] assert ak.ptp(array, axis=-1, mask_identity=False).tolist() == [1, 0, 0, None, 2] assert ak.ptp(array, axis=-2).tolist() == [4, 4, None, 0] assert ak.ptp(array, axis=-2, mask_identity=False).tolist() == [4, 4, 0, 0]
def sum(self, axis=-1): """Sum an array of vectors elementwise using `x` and `y` components""" return awkward.zip( { "x": awkward.sum(self.x, axis=axis), "y": awkward.sum(self.y, axis=axis), }, with_name="TwoVector", behavior=self.behavior, )
def sum(self, axis=-1): """Sum an array of vectors elementwise using `x` and `y` components""" out = awkward.zip( { "x": awkward.sum(self.x, axis=axis), "y": awkward.sum(self.y, axis=axis), }, with_name="TwoVector", highlevel=False, ) return awkward.Array(out, behavior=self.behavior)
def sum(self, axis=-1): """Sum an array of vectors elementwise using `x`, `y`, `z`, `t`, and `charge` components""" return awkward.zip( { "x": awkward.sum(self.x, axis=axis), "y": awkward.sum(self.y, axis=axis), "z": awkward.sum(self.z, axis=axis), "t": awkward.sum(self.t, axis=axis), "charge": awkward.sum(self.charge, axis=axis), }, with_name="Candidate", )
def sum(self, axis=-1): """Sum an array of vectors elementwise using `x`, `y`, `z`, and `t` components""" return awkward.zip( { "x": awkward.sum(self.x, axis=axis), "y": awkward.sum(self.y, axis=axis), "z": awkward.sum(self.z, axis=axis), "t": awkward.sum(self.t, axis=axis), }, with_name="LorentzVector", behavior=self.behavior, )
def corrected_polar_met(met_pt, met_phi, jet_pt, jet_phi, jet_pt_orig, deltas=None): sj, cj = numpy.sin(jet_phi), numpy.cos(jet_phi) x = met_pt * numpy.cos(met_phi) + awkward.sum( jet_pt * cj - jet_pt_orig * cj, axis=1 ) y = met_pt * numpy.sin(met_phi) + awkward.sum( jet_pt * sj - jet_pt_orig * sj, axis=1 ) if deltas: positive, dx, dy = deltas x = x + dx if positive else x - dx y = y + dy if positive else y - dy return awkward.zip({"pt": numpy.hypot(x, y), "phi": numpy.arctan2(y, x)})
def add_object_cuts(self, objects, object_name, cuts, names, sample): if object_name not in self.object_stats.keys(): self.object_stats[object_name] = {} if sample not in self.object_stats[object_name].keys(): self.object_stats[object_name][sample] = { "n_objects_initial" : awkward.sum(awkward.num(objects)) } else: self.object_stats[object_name][sample]["n_objects_initial"] += awkward.sum(awkward.num(objects)) for cut, name in zip(cuts, names): if name not in self.object_stats[object_name][sample].keys(): self.object_stats[object_name][sample][name] = awkward.sum(awkward.num(objects[cut])) else: self.object_stats[object_name][sample][name] += awkward.sum(awkward.num(objects[cut]))
def Phi_mpi_pi(array_phi): """ Periodically (T=2*pi, shift step done only once) bring values of the given array into [-pi, pi] range. Arguments: array_phi: awkward array, values assumed to be radian measure of phi angle Returns: Awkward array, values of input array brought to [-pi, pi] range """ array_phi = ak.where(array_phi <= np.pi, array_phi, array_phi - 2*np.pi) array_phi = ak.where(array_phi >= -np.pi, array_phi, array_phi + 2*np.pi) assert ak.sum(array_phi > np.pi) + ak.sum(array_phi < -np.pi) == 0 return array_phi
def get_norm(file, weight): with uproot.open(file_map[file]) as f: tree = f["Events"] events = tree.arrays(["genWeight", "LHEScaleWeight"], library="ak", how="zip") nominal_weights = events.jagged0.LHEScaleWeight[:, 4] var_weights = events.jagged0.LHEScaleWeight[:, weight_map[weight]] sum_nominal = awkward.sum(nominal_weights) sum_weight = awkward.sum(var_weights) print(sum_nominal, sum_weight) return sum_nominal / sum_weight
def add_pdf_weight(weights, pdf_weights): nom = np.ones(len(weights.weight())) up = np.ones(len(weights.weight())) down = np.ones(len(weights.weight())) # NNPDF31_nnlo_hessian_pdfas # https://lhapdfsets.web.cern.ch/current/NNPDF31_nnlo_hessian_pdfas/NNPDF31_nnlo_hessian_pdfas.info if pdf_weights is not None and "306000 - 306102" in pdf_weights.__doc__: # Hessian PDF weights # Eq. 21 of https://arxiv.org/pdf/1510.03865v1.pdf arg = pdf_weights[:, 1:-2] - np.ones((len(weights.weight()), 100)) summed = ak.sum(np.square(arg), axis=1) pdf_unc = np.sqrt((1. / 99.) * summed) weights.add('PDF_weight', nom, pdf_unc + nom) # alpha_S weights # Eq. 27 of same ref as_unc = 0.5 * (pdf_weights[:, 102] - pdf_weights[:, 101]) weights.add('aS_weight', nom, as_unc + nom) # PDF + alpha_S weights # Eq. 28 of same ref pdfas_unc = np.sqrt(np.square(pdf_unc) + np.square(as_unc)) weights.add('PDFaS_weight', nom, pdfas_unc + nom) else: weights.add('aS_weight', nom, up, down) weights.add('PDF_weight', nom, up, down) weights.add('PDFaS_weight', nom, up, down)
def process(self, events): output = self.accumulator.identity() dataset = events.metadata['dataset'] integratedLuminosity = 137.19 * 1000 # fb^{-1} to pb^{-1} ht = events.HT weights = integratedLuminosity * events.CrossSection[ht > 1200] / len( events) GenParticles_pt = events.GenParticles.pt GenParticles_eta = events.GenParticles.eta GenParticles_Status = events.GenParticles.Status GenParticles_PdgId = events.GenParticles.PdgId GenParticles_Charge = events.GenParticles.Charge finalParticles = (GenParticles_Status == 1) & (GenParticles_pt > 1) & ( abs(GenParticles_eta) < 2.5) & (GenParticles_Charge != 0) nTracks = ak.sum(finalParticles[ht > 1200], axis=1) output["sumw"][dataset] += len(events) output["nTracks"].fill(dataset=dataset, nTracks=nTracks, weight=weights) return output
def compute_rate(tau_1, tau_2, Nev_den, Pt_thr, a, L1rate): mask_1 = (tau_1.deepTau_VSjet > deep_thr( tau_1, a, Pt_thr)) & reco_tau_selection(tau_1, minPt=Pt_thr) mask_2 = (tau_2.deepTau_VSjet > deep_thr( tau_2, a, Pt_thr)) & reco_tau_selection(tau_2, minPt=Pt_thr) Nev_num = ak.sum(ditau_selection(mask_1, mask_2)) return Nev_num / Nev_den * L1rate
def select_deltaR(events, obj1, obj2, threshold, debug): if debug > 1: start = time.time() nEvents = len(events) n_obj1 = numpy.array(awkward.num(obj1), numpy.int64) n_obj2 = numpy.array(awkward.num(obj2), numpy.int64) mask_offsets = numpy.zeros(nEvents + 1, numpy.int64) mask_contents = numpy.empty(awkward.sum(n_obj1), numpy.bool) mask_offsets, mask_contents = select_deltaR_raw(events, nEvents, obj1, n_obj1, obj2, n_obj2, threshold, mask_offsets, mask_contents, debug) if debug > 1: elapsed_time = time.time() - start print( "[object_selections.py] PERFORMANCE: time to execute select_deltaR: %.6f" % (elapsed_time)) mask = awkward.Array( awkward.layout.ListOffsetArray64( awkward.layout.Index64(mask_offsets), awkward.layout.NumpyArray(mask_contents))) return mask
def generate_epd_hit_matrix(self): ring_sum = np.zeros((32, len(self.nMip))) print("Filling array of dimension", ring_sum.shape) for i in range(32): ring_i = ak.sum(ak.where(self.row == i+1, self.nMip, 0), axis=-1) ring_sum[i] = ring_i return ring_sum
def process(self, events): # Note: This is required to ensure that behaviors are registered # when running this code in a remote task. ak.behavior.update(candidate.behavior) output = self.accumulator.identity() dataset = events.metadata['dataset'] muons = ak.zip( { "pt": events.Muon_pt, "eta": events.Muon_eta, "phi": events.Muon_phi, "mass": events.Muon_mass, "charge": events.Muon_charge, }, with_name="PtEtaPhiMCandidate") cut = (ak.num(muons) == 2) & (ak.sum(muons.charge) == 0) # add first and second muon in every event together dimuon = muons[cut][:, 0] + muons[cut][:, 1] output["sumw"][dataset] += len(events) output["mass"].fill( dataset=dataset, mass=dimuon.mass, ) return output
def load_file(self): """ Function which reads a root or csv file and removes interactions and events without energy depositions. Returns: interactions: awkward array n_simulated_events: Total number of simulated events in the opened file. (This includes events removed by cuts) """ if self.file.endswith(".root"): interactions, n_simulated_events, start, stop = self._load_root_file( ) elif self.file.endswith(".csv"): interactions, n_simulated_events, start, stop = self._load_csv_file( ) else: raise ValueError( f'Cannot load events from file "{self.file}": .root or .cvs file needed.' ) if np.any(interactions['ed'] < 0): warnings.warn('At least one of the energy deposits is negative!') # Removing all events with zero energy deposit m = interactions['ed'] > 0 if self.cut_by_eventid: # ufunc does not work here... m2 = (interactions['evtid'] >= start) & (interactions['evtid'] < stop) m = m & m2 interactions = interactions[m] if self.cut_nr_only: m = ((interactions['type'] == "neutron") & (interactions['edproc'] == "hadElastic")) | ( interactions['edproc'] == "ionIoni") e_dep_er = ak.sum(interactions[~m]['ed'], axis=1) e_dep_nr = ak.sum(interactions[m]['ed'], axis=1) interactions = interactions[(e_dep_er < 10) & (e_dep_nr > 0)] # Removing all events with no interactions: m = ak.num(interactions['ed']) > 0 interactions = interactions[m] return interactions, n_simulated_events
def process(self, events): has2jets = ak.sum(events.Jet.pt > 40, axis=1) >= 2 return (hist.Hist.new.Reg(100, 0, 200, name="met", label="$E_{T}^{miss}$ [GeV]").Double().fill( events[has2jets].MET.pt))
def flip_weight(self, electron): #f_1 = self.evaluator['el'](electron.pt[:,0:1], abs(electron.eta[:,0:1])) #f_2 = self.evaluator['el'](electron.pt[:,1:2], abs(electron.eta[:,1:2])) # For custom measurements f_1 = yahist_2D_lookup(self.ratio, electron.pt[:, 0:1], abs(electron.eta[:, 0:1])) f_2 = yahist_2D_lookup(self.ratio, electron.pt[:, 1:2], abs(electron.eta[:, 1:2])) # I'm using ak.prod and ak.sum to replace empty arrays by 1 and 0, respectively weight = ak.sum(f_1 / (1 - f_1), axis=1) * ak.prod( 1 - f_2 / (1 - f_2), axis=1) + ak.sum( f_2 / (1 - f_2), axis=1) * ak.prod(1 - f_1 / (1 - f_1), axis=1) return weight
def compute_rate_base(tau_1, tau_2, Nev_den, Pt_thr, L1rate): mask_1 = iso_tau_selection(tau_1, "mediumIsoAbs", "mediumIsoRel") & reco_tau_selection( tau_1, minPt=Pt_thr) mask_2 = iso_tau_selection(tau_2, "mediumIsoAbs", "mediumIsoRel") & reco_tau_selection( tau_2, minPt=Pt_thr) Nev_num = ak.sum(ditau_selection(mask_1, mask_2)) return Nev_num / Nev_den * L1rate
def process_event(weights, evt_vars): #, bjets): out_hists = {} out_hists["MET"] = Hist1D(ak.to_numpy(evt_vars.MET_pt), bins=bin_met, label="MET", weights=ak.to_numpy(genWeight)) #out_hists["njets"] = Hist1D(ak.to_numpy(evt_vars.nJet), bins = bin_njet, label="njet", weights = ak.to_numpy(genWeight)) #out_hists["nbjets"] = Hist1D(ak.num(bjets), bins = bin_bjet) out_hists["weight"] = ak.sum(weights) return out_hists
def add_cuts(self, cuts, names): for cut, name in zip(cuts, names): if self.debug > 0: n_objects_cut = awkward.sum(awkward.num(self.objects[cut])) if self.n_objects_initial == 0: return print( "%s ObjectCutDiagnostics: Applying cut %s would have an eff of %.4f" % (self.cut_set, name, float(n_objects_cut) / float(self.n_objects_initial)))
def FR_weight(self, lepton): f_1 = yahist_2D_lookup(self.ratio, lepton.conePt[:, 0:1], abs(lepton.eta[:, 0:1])) #breakpoint() # I'm using ak.prod and ak.sum to replace empty arrays by 1 and 0, respectively weight = ak.sum( f_1 / (1 - f_1), axis=1 ) #*ak.prod(1-f_2/(1-f_2), axis=1) + ak.sum(f_2/(1-f_2), axis=1)*ak.prod(1-f_1/(1-f_1), axis=1) return weight
def set_genZ(events, selection_options, debug): electron_idxs = abs(events.GenPart_pdgId) == 11 muon_idxs = abs(events.GenPart_pdgId) == 13 tau_idxs = abs(events.GenPart_pdgId) == 15 motherOfElectrons = events.GenPart_genPartIdxMother[electron_idxs] motherOfMuons = events.GenPart_genPartIdxMother[muon_idxs] motherOfTaus = events.GenPart_genPartIdxMother[tau_idxs] ZToEleEvents = ak.sum( (events.GenPart_pdgId[motherOfElectrons] == 23), axis=1) >= 2 ZToMuEvents = ak.sum( (events.GenPart_pdgId[motherOfMuons] == 23), axis=1) >= 2 ZToTauEvents = ak.sum( (events.GenPart_pdgId[motherOfTaus] == 23), axis=1) >= 2 events[ "genZ_decayMode"] = 1 * ZToEleEvents + 2 * ZToMuEvents + 3 * ZToTauEvents return events
def test_reducers(): # axis=None reducers are implemented in NumPy. assert ak.sum(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]])) == 6 + 6j assert ak.prod(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]])) == -12 + 12j # axis != None reducers are implemented in libawkward; this should be ReducerSum. assert ak.sum(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]), axis=1).tolist() == [ 3 + 3j, 0 + 0j, 3 + 3j, ] # And this is in ReducerProd. assert ak.prod(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]), axis=1).tolist() == [ 0 + 4j, 1 + 0j, 3 + 3j, ] # ReducerCount, ReducerCountNonzero, ReducerAny, and ReducerAll work. assert ak.count(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]), axis=1).tolist() == [2, 0, 1] assert ak.count_nonzero(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]), axis=1).tolist() == [2, 0, 1] assert ak.any(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]), axis=1).tolist() == [ True, False, True, ] assert ak.all(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]), axis=1).tolist() == [ True, True, True, ] assert ak.any(ak.from_iter([[1 + 1j, 2 + 2j, 0 + 0j], [], [3 + 3j]]), axis=1).tolist() == [True, False, True] assert ak.all(ak.from_iter([[1 + 1j, 2 + 2j, 0 + 0j], [], [3 + 3j]]), axis=1).tolist() == [False, True, True]