def get_selected_electrons(electrons, pt_cut_leading, pt_cut_subleading, aeta_cut, iso_cut): this_worker = get_worker_wrapper() NUMPY_LIB = this_worker.NUMPY_LIB backend = this_worker.backend passes_iso = electrons.pfRelIso03_all > iso_cut passes_id = electrons.pfId == True passes_subleading_pt = electrons.pt > pt_cut_subleading passes_leading_pt = electrons.pt > pt_cut_leading passes_aeta = NUMPY_LIB.abs(electrons.eta) < aeta_cut evs_all = NUMPY_LIB.ones(electrons.numevents(), dtype=bool) els_all = NUMPY_LIB.ones(electrons.numobjects(), dtype=bool) selected_electrons = passes_iso & passes_id & passes_subleading_pt & passes_aeta selected_electrons_leading = selected_electrons & passes_leading_pt ev_1el = (ha_kernels.sum_in_offsets( this_worker.backend, electrons.offsets, selected_electrons_leading, evs_all, electrons.masks["all"], ) >= 1) return selected_electrons, ev_1el
def get_selected_muons(muons, pt_cut_leading, pt_cut_subleading, aeta_cut, iso_cut): this_worker = get_worker_wrapper() NUMPY_LIB = this_worker.NUMPY_LIB backend = this_worker.backend passes_iso = muons.pfRelIso03_all > iso_cut passes_id = muons.tightId == True passes_subleading_pt = muons.pt > pt_cut_subleading passes_leading_pt = muons.pt > pt_cut_leading passes_aeta = NUMPY_LIB.abs(muons.eta) < aeta_cut selected_muons = passes_iso & passes_id & passes_subleading_pt & passes_aeta selected_muons_leading = selected_muons & passes_leading_pt evs_all = NUMPY_LIB.ones(muons.numevents(), dtype=bool) # select events with at least 1 good muon evs_1mu = (ha_kernels.sum_in_offsets( this_worker.backend, muons.offsets, selected_muons_leading, evs_all, muons.masks["all"], ) >= 1) return selected_muons, evs_1mu
def analyze_data_function(data, parameters): ret = Results() num_events = data["num_events"] muons = data["Muon"] mu_pt = nplib.sqrt(muons.Px**2 + muons.Py**2) muons.attrs_data["pt"] = mu_pt mask_events = nplib.ones(muons.numevents(), dtype=nplib.bool) mask_muons_passing_pt = muons.pt > parameters["muons_ptcut"] num_muons_event = kernels.sum_in_offsets(backend, muons.offsets, mask_muons_passing_pt, mask_events, muons.masks["all"], nplib.int8) mask_events_dimuon = num_muons_event == 2 #get the leading muon pt in events that have exactly two muons inds = nplib.zeros(num_events, dtype=nplib.int32) leading_muon_pt = kernels.get_in_offsets(backend, muons.offsets, muons.pt, inds, mask_events_dimuon, mask_muons_passing_pt) #compute a weighted histogram weights = nplib.ones(num_events, dtype=nplib.float32) bins = nplib.linspace(0, 300, 101, dtype=nplib.float32) hist_muons_pt = Histogram(*kernels.histogram_from_vector( backend, leading_muon_pt[mask_events_dimuon], weights[mask_events_dimuon], bins)) #save it to the output ret["hist_leading_muon_pt"] = hist_muons_pt return ret
def test_kernel_sum_in_offsets(dataset): muons = dataset.structs["Muon"][0] sel_ev = nplib.ones(muons.numevents(), dtype=nplib.bool) sel_mu = nplib.ones(muons.numobjects(), dtype=nplib.bool) z = kernels.sum_in_offsets(backend, muons.offsets, muons.pt, sel_ev, sel_mu, dtype=nplib.float32)
def test_kernel_sum_in_offsets(self): dataset = self.dataset muons = dataset.structs["Muon"][0] sel_ev = self.NUMPY_LIB.ones(muons.numevents(), dtype=self.NUMPY_LIB.bool) sel_mu = self.NUMPY_LIB.ones(muons.numobjects(), dtype=self.NUMPY_LIB.bool) z = kernels.sum_in_offsets(self.ha, muons.offsets, muons.pt, sel_ev, sel_mu, dtype=self.NUMPY_LIB.float32) return muons.numevents()
def run_analysis(dataset, out, dnnmodel, use_cuda, ismc): from keras.backend.tensorflow_backend import set_session this_worker = get_worker_wrapper() NUMPY_LIB = this_worker.NUMPY_LIB backend = this_worker.backend hists = {} histo_bins = { "nmu": np.array([0, 1, 2, 3], dtype=np.float32), "njet": np.array([0, 1, 2, 3, 4, 5, 6, 7], dtype=np.float32), "mu_pt": np.linspace(0, 300, 20), "mu_eta": np.linspace(-5, 5, 20), "mu_phi": np.linspace(-5, 5, 20), "mu_iso": np.linspace(0, 1, 20), "mu_charge": np.array([-1, 0, 1], dtype=np.float32), "met_pt": np.linspace(0, 200, 20), "jet_pt": np.linspace(0, 400, 20), "jet_eta": np.linspace(-5, 5, 20), "jet_phi": np.linspace(-5, 5, 20), "jet_btag": np.linspace(0, 1, 20), "dnnpred_m": np.linspace(0, 1, 20), "dnnpred_s": np.linspace(0, 0.2, 20), "inv_mass": np.linspace(150, 200, 20), "sumpt": np.linspace(0, 1000, 20), } t0 = time.time() i = 0 mu = dataset.structs["Muon"][i] el = dataset.structs["Electron"][i] jets = dataset.structs["Jet"][i] evvars = dataset.eventvars[i] mu.hepaccelerate_backend = backend el.hepaccelerate_backend = backend jets.hepaccelerate_backend = backend evs_all = NUMPY_LIB.ones(dataset.numevents(), dtype=NUMPY_LIB.bool) print("Lepton selection") sel_mu, sel_ev_mu = get_selected_muons(mu, 40, 20, 2.4, 0.1) sel_ev_mu = sel_ev_mu & (evvars["HLT_IsoMu24"] == True) mu.masks["selected"] = sel_mu sel_el, sel_ev_el = get_selected_electrons(el, 40, 20, 2.4, 0.1) el.masks["selected"] = sel_el nmu = ha_kernels.sum_in_offsets( backend, mu.offsets, mu.masks["selected"], evs_all, mu.masks["all"], dtype=NUMPY_LIB.int32, ) nel = ha_kernels.sum_in_offsets( backend, el.offsets, el.masks["selected"], evs_all, el.masks["all"], dtype=NUMPY_LIB.int32, ) # get contiguous arrays of the first two muons for all events mu1 = mu.select_nth(0, object_mask=sel_mu) mu2 = mu.select_nth(1, object_mask=sel_mu) el1 = el.select_nth(0, object_mask=sel_el) el2 = el.select_nth(1, object_mask=sel_el) weight_ev_mu = apply_lepton_corrections(mu, sel_mu, this_worker.electron_weights) weight_ev_el = apply_lepton_corrections(el, sel_el, this_worker.electron_weights) weights = {"nominal": weight_ev_mu * weight_ev_el} weights_jet = {} for k in weights.keys(): weights_jet[k] = NUMPY_LIB.zeros_like(jets.pt) ha_kernels.broadcast(backend, jets.offsets, weights["nominal"], weights_jet[k]) all_jecs = [("nominal", "", None)] if ismc: for i in range(this_worker.jecs_up.shape[1]): all_jecs += [(i, "up", this_worker.jecs_up[:, i])] all_jecs += [(i, "down", this_worker.jecs_down[:, i])] jets_pt_orig = NUMPY_LIB.copy(jets.pt) # per-event histograms fill_histograms_several( hists, "nominal", "hist__all__", [ (evvars["MET_pt"], "met_pt", histo_bins["met_pt"]), ], evs_all, weights, use_cuda, ) fill_histograms_several( hists, "nominal", "hist__all__", [ (jets.pt, "jets_pt", histo_bins["jet_pt"]), ], jets.masks["all"], weights_jet, use_cuda, ) print("Jet selection") # loop over the jet corrections for ijec, sdir, jec in all_jecs: systname = "nominal" if ijec != "nominal": systname = ("jec{0}".format(ijec), sdir) if not jec is None: jet_pt_corr = apply_jec(jets_pt_orig, this_worker.jecs_bins, jec) # compute the corrected jet pt jets.pt = jets_pt_orig * NUMPY_LIB.abs(jet_pt_corr) print("jec", ijec, sdir, jets.pt.mean()) # get selected jets sel_jet, sel_bjet = select_jets(jets, mu, el, sel_mu, sel_el, 40, 2.0, 0.3, 0.4) # compute the number of jets per event njet = ha_kernels.sum_in_offsets( backend, jets.offsets, sel_jet, evs_all, jets.masks["all"], dtype=NUMPY_LIB.int32, ) nbjet = ha_kernels.sum_in_offsets( backend, jets.offsets, sel_bjet, evs_all, jets.masks["all"], dtype=NUMPY_LIB.int32, ) inv_mass_3j = NUMPY_LIB.zeros(jets.numevents(), dtype=NUMPY_LIB.float32) best_comb_3j = NUMPY_LIB.zeros((jets.numevents(), 3), dtype=NUMPY_LIB.int32) if use_cuda: this_worker.kernels.comb_3_invmass_closest[32, 256]( jets.pt, jets.eta, jets.phi, jets.mass, jets.offsets, 172.0, inv_mass_3j, best_comb_3j, ) cuda.synchronize() else: this_worker.kernels.comb_3_invmass_closest( jets.pt, jets.eta, jets.phi, jets.mass, jets.offsets, 172.0, inv_mass_3j, best_comb_3j, ) best_btag = NUMPY_LIB.zeros(jets.numevents(), dtype=NUMPY_LIB.float32) if use_cuda: this_worker.kernels.max_val_comb[32, 1024](jets.btag, jets.offsets, best_comb_3j, best_btag) cuda.synchronize() else: this_worker.kernels.max_val_comb(jets.btag, jets.offsets, best_comb_3j, best_btag) # get the events with at least three jets sel_ev_jet = njet >= 3 sel_ev_bjet = nbjet >= 1 selected_events = (sel_ev_mu | sel_ev_el) & sel_ev_jet & sel_ev_bjet print("Selected {0} events".format(selected_events.sum())) # get contiguous vectors of the first two jet data jet1 = jets.select_nth(0, object_mask=sel_jet) jet2 = jets.select_nth(1, object_mask=sel_jet) jet3 = jets.select_nth(2, object_mask=sel_jet) # create a mask vector for the first two jets first_two_jets = NUMPY_LIB.zeros_like(sel_jet) inds = NUMPY_LIB.zeros_like(evs_all, dtype=NUMPY_LIB.int32) targets = NUMPY_LIB.ones_like(evs_all, dtype=NUMPY_LIB.int32) inds[:] = 0 ha_kernels.set_in_offsets( backend, jets.offsets, first_two_jets, inds, targets, selected_events, sel_jet, ) inds[:] = 1 ha_kernels.set_in_offsets( backend, jets.offsets, first_two_jets, inds, targets, selected_events, sel_jet, ) # compute the invariant mass of the first two jets dijet_inv_mass, dijet_pt = compute_inv_mass(jets, selected_events, sel_jet & first_two_jets, use_cuda) sumpt_jets = ha_kernels.sum_in_offsets(backend, jets.offsets, jets.pt, selected_events, sel_jet) # create a keras-like array arr = NUMPY_LIB.vstack([ nmu, nel, njet, dijet_inv_mass, dijet_pt, mu1["pt"], mu1["eta"], mu1["phi"], mu1["charge"], mu1["pfRelIso03_all"], mu2["pt"], mu2["eta"], mu2["phi"], mu2["charge"], mu2["pfRelIso03_all"], el1["pt"], el1["eta"], el1["phi"], el1["charge"], el1["pfRelIso03_all"], el2["pt"], el2["eta"], el2["phi"], el2["charge"], el2["pfRelIso03_all"], jet1["pt"], jet1["eta"], jet1["phi"], jet1["btag"], jet2["pt"], jet2["eta"], jet2["phi"], jet2["btag"], inv_mass_3j, best_btag, sumpt_jets, ]).T # print("evaluating DNN model") with this_worker.graph.as_default(): set_session(this_worker.session) pred = dnnmodel.eval(arr, use_cuda) pred = NUMPY_LIB.vstack(pred).T pred_m = NUMPY_LIB.mean(pred, axis=1) pred_s = NUMPY_LIB.std(pred, axis=1) fill_histograms_several( hists, systname, "hist__nmu1_njetge3_nbjetge1__", [ (pred_m, "pred_m", histo_bins["dnnpred_m"]), (pred_s, "pred_s", histo_bins["dnnpred_s"]), (nmu, "nmu", histo_bins["nmu"]), (nel, "nel", histo_bins["nmu"]), (njet, "njet", histo_bins["njet"]), (mu1["pt"], "mu1_pt", histo_bins["mu_pt"]), (mu1["eta"], "mu1_eta", histo_bins["mu_eta"]), (mu1["phi"], "mu1_phi", histo_bins["mu_phi"]), (mu1["charge"], "mu1_charge", histo_bins["mu_charge"]), (mu1["pfRelIso03_all"], "mu1_iso", histo_bins["mu_iso"]), (mu2["pt"], "mu2_pt", histo_bins["mu_pt"]), (mu2["eta"], "mu2_eta", histo_bins["mu_eta"]), (mu2["phi"], "mu2_phi", histo_bins["mu_phi"]), (mu2["charge"], "mu2_charge", histo_bins["mu_charge"]), (mu2["pfRelIso03_all"], "mu2_iso", histo_bins["mu_iso"]), (el1["pt"], "el1_pt", histo_bins["mu_pt"]), (el1["eta"], "el1_eta", histo_bins["mu_eta"]), (el1["phi"], "el1_phi", histo_bins["mu_phi"]), (el1["charge"], "el1_charge", histo_bins["mu_charge"]), (el1["pfRelIso03_all"], "el1_iso", histo_bins["mu_iso"]), (el2["pt"], "el2_pt", histo_bins["mu_pt"]), (el2["eta"], "el2_eta", histo_bins["mu_eta"]), (el2["phi"], "el2_phi", histo_bins["mu_phi"]), (el2["charge"], "el2_charge", histo_bins["mu_charge"]), (el2["pfRelIso03_all"], "el2_iso", histo_bins["mu_iso"]), (jet1["pt"], "j1_pt", histo_bins["jet_pt"]), (jet1["eta"], "j1_eta", histo_bins["jet_eta"]), (jet1["phi"], "j1_phi", histo_bins["jet_phi"]), (jet1["btag"], "j1_btag", histo_bins["jet_btag"]), (jet2["pt"], "j2_pt", histo_bins["jet_pt"]), (jet2["eta"], "j2_eta", histo_bins["jet_eta"]), (jet2["phi"], "j2_phi", histo_bins["jet_phi"]), (jet2["btag"], "j2_btag", histo_bins["jet_btag"]), (inv_mass_3j, "inv_mass_3j", histo_bins["inv_mass"]), (best_btag, "best_btag", histo_bins["jet_btag"]), (sumpt_jets, "sumpt", histo_bins["sumpt"]), ], selected_events, weights, use_cuda, ) # save the array for the first jet correction scenario only if save_arrays and ijec == 0: outfile_arr = "{0}_arrs.npy".format(out) print("Saving array with shape {0} to {1}".format( arr.shape, outfile_arr)) with open(outfile_arr, "wb") as fi: np.save(fi, NUMPY_LIB.asnumpy(arr)) t1 = time.time() res = Results({}) for hn in hists.keys(): hists[hn] = Results(hists[hn]) res["hists"] = Results(hists) res["numevents"] = dataset.numevents() speed = dataset.numevents() / (t1 - t0) print("run_analysis: {0:.2E} events in {1:.2f} seconds, speed {2:.2E} Hz". format(dataset.numevents(), t1 - t0, speed)) return res
def analyze_data_function(data, parameters): ret = Results() ha = parameters["ha"] num_events = data["num_events"] lep = data["Lep"] lep.hepaccelerate_backend = ha lep.attrs_data["pt"] = lep.lep_pt lep.attrs_data["eta"] = lep.lep_eta lep.attrs_data["phi"] = lep.lep_phi lep.attrs_data["charge"] = lep.lep_charge lep.attrs_data["type"] = lep.lep_type lep_mass = np.zeros_like(lep["pt"], dtype=nplib.float32) lep_mass = np.where(lep["type"] == 11, 0.511, lep_mass) lep_mass = np.where(lep["type"] == 13, 105.65837, lep_mass) lep.attrs_data["mass"] = lep_mass mask_events = nplib.ones(lep.numevents(), dtype=nplib.bool) lep_ele = lep["type"] == 11 lep_muon = lep["type"] == 13 ele_Iso = np.logical_and( lep_ele, np.logical_and(lep.lep_ptcone30 / lep.pt < 0.15, lep.lep_etcone20 / lep.pt < 0.20)) muon_Iso = np.logical_and( lep_muon, np.logical_and(lep.lep_ptcone30 / lep.pt < 0.15, lep.lep_etcone20 / lep.pt < 0.30)) pass_iso = np.logical_or(ele_Iso, muon_Iso) lep.attrs_data["pass_iso"] = pass_iso num_lep_event = kernels.sum_in_offsets( backend, lep.offsets, lep.masks["all"], mask_events, lep.masks["all"], nplib.int8, ) mask_events_4lep = num_lep_event == 4 lep_attrs = ["pt", "eta", "phi", "charge", "type", "mass", "pass_iso"] #, "ptcone30", "etcone20"] lep0 = lep.select_nth(0, mask_events_4lep, lep.masks["all"], attributes=lep_attrs) lep1 = lep.select_nth(1, mask_events_4lep, lep.masks["all"], attributes=lep_attrs) lep2 = lep.select_nth(2, mask_events_4lep, lep.masks["all"], attributes=lep_attrs) lep3 = lep.select_nth(3, mask_events_4lep, lep.masks["all"], attributes=lep_attrs) mask_event_sumchg_zero = (lep0["charge"] + lep1["charge"] + lep2["charge"] + lep3["charge"] == 0) sum_lep_type = lep0["type"] + lep1["type"] + lep2["type"] + lep3["type"] all_pass_iso = (lep0["pass_iso"] & lep1["pass_iso"] & lep2["pass_iso"] & lep3["pass_iso"]) mask_event_sum_lep_type = np.logical_or( (sum_lep_type == 44), np.logical_or((sum_lep_type == 48), (sum_lep_type == 52))) mask_events = mask_events & mask_event_sumchg_zero & mask_events_4lep & mask_event_sum_lep_type & all_pass_iso mask_lep1_passing_pt = lep1["pt"] > parameters["leading_lep_ptcut"] mask_lep2_passing_pt = lep2["pt"] > parameters["lep_ptcut"] mask_events = mask_events & mask_lep1_passing_pt & mask_lep2_passing_pt l0 = to_cartesian(lep0) l1 = to_cartesian(lep1) l2 = to_cartesian(lep2) l3 = to_cartesian(lep3) llll = {k: l0[k] + l1[k] + l2[k] + l3[k] for k in ["px", "py", "pz", "e"]} llll_sph = to_spherical(llll) llll_sph["mass"] = llll_sph["mass"] / 1000. # Convert to GeV #import pdb;pdb.set_trace(); # compute a weighted histogram weights = nplib.ones(num_events, dtype=nplib.float32) ## Add xsec weights based on sample name if parameters["is_mc"]: weights = data['eventvars']['mcWeight'] * data['eventvars'][ 'scaleFactor_PILEUP'] * data['eventvars']['scaleFactor_ELE'] * data[ 'eventvars']['scaleFactor_MUON'] * data['eventvars'][ 'scaleFactor_LepTRIGGER'] info = infofile.infos[parameters["sample"]] weights *= (lumi * 1000 * info["xsec"]) / (info["sumw"] * info["red_eff"]) bins = nplib.linspace(110, 150, 11, dtype=nplib.float32) hist_m4lep = Histogram(*kernels.histogram_from_vector( backend, llll_sph["mass"][mask_events], weights[mask_events], bins, )) # save it to the output ret["hist_m4lep"] = hist_m4lep return ret