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
