def do_jet_pt_resolution(self, event):
indices = get_bin_indices(
[
event.Jet_eta.content,
event_to_object_var(event.fixedGridRhoFastjetAll,
event.Jet_pt.starts, event.Jet_pt.stops)
],
[self.jers["eta_low"].values, self.jers["rho_low"].values],
[self.jers["eta_high"].values, self.jers["rho_high"].values],
1,
)[:, 0]
df = self.jers.iloc[indices]
params = df[["param0", "param1", "param2", "param3"]].values
ptbounds = df[["pt_low", "pt_high"]].values
event.Jet_ptResolution = awk.JaggedArray(
event.Jet_pt.starts,
event.Jet_pt.stops,
jer_formula(
np.minimum(np.maximum(event.Jet_pt.content, ptbounds[:, 0]),
ptbounds[:, 1]),
params[:, 0],
params[:, 1],
params[:, 2],
params[:, 3],
),
)
def do_jes_correction(self, event, source):
df = self.jesuncs[self.jesuncs["source"] == source]
indices = get_bin_indices(
[event.Jet_eta.content],
[df["eta_low"].values],
[df["eta_high"].values],
1,
)[:, 0]
pt = np.array(list(df.iloc[indices]["pt"].values))
corr_up = np.array(list(df.iloc[indices]["corr_up"].values))
corr_down = np.array(list(df.iloc[indices]["corr_down"].values))
corr_up = interpolate(event.Jet_ptJESOnly.content, pt, corr_up)
corr_down = interpolate(event.Jet_ptJESOnly.content, pt, corr_down)
starts = event.Jet_eta.starts
stops = event.Jet_eta.stops
setattr(event, "Jet_JECjes{}Up".format(source),
awk.JaggedArray(
starts,
stops,
corr_up,
))
setattr(event, "Jet_JECjes{}Down".format(source),
awk.JaggedArray(
starts,
stops,
-1. * corr_down,
))
def fevaluate_pu(ev, evidx, nsig, source, var_):
mins = corrs["nTrueInt"].values.astype(np.float32)
maxs = mins[:] + 1
maxs[-1] = np.inf
indices = get_bin_indices([getattr(ev, var_)], [mins], [maxs], 1)[:, 0]
ev_corrs = corrs.iloc[indices]
nominal = ev_corrs["corr"].values
up = (ev_corrs["corr_up"].values / nominal - 1.) * (source == "pileup")
down = (ev_corrs["corr_down"].values / nominal - 1.) * (source
== "pileup")
return weight_numba(nominal, nsig, up, down)
def get_efficiencies(event, obj, selection, effmap):
indices = get_bin_indices(
[
getattr(event, "{}_eta".format(obj))[selection].content,
getattr(event,
"{}_ptShift".format(obj))(event)[selection].content.astype(
np.float32)
],
[effmap["xlow"].values, effmap["ylow"].values],
[effmap["xupp"].values, effmap["yupp"].values],
1,
)[:, 0]
df = effmap.iloc[indices]
return df["content"].values, df["error"].values
def event(self, event):
if self.parent not in self.input_paths:
weights = np.ones(event.size)
event.WeightQcdEwkNominal = weights
for variation in self.variations[1:]:
setattr(
event,
"WeightQcdEwk_{}".format(variation),
np.zeros(event.size),
)
else:
indices = get_bin_indices(
[event.GenPartBoson_pt],
[self.input_df.index.get_level_values("bin_min").values],
[self.input_df.index.get_level_values("bin_max").values],
1,
)[:, 0]
corrections = self.input_df.iloc[indices]
event.WeightQcdEwkNominal = corrections[""].values
for variation in self.variations[1:]:
setattr(event, "WeightQcdEwk_{}".format(variation),
(corrections[variation] / corrections[""]).values - 1.)
def fevaluate_object_weights(ev, evidx, nsig, source, name_):
event_vars = [v(ev) for v in bins_vars]
for v in event_vars:
v.content[np.isnan(v.content)] = 0.
indices = get_bin_indices(
[event_vars[idx].content.astype(np.float32) for idx in range(len(event_vars))],
[df["bin{}_low".format(idx)].values.astype(np.float32) for idx in range(len(event_vars))],
[df["bin{}_upp".format(idx)].values.astype(np.float32) for idx in range(len(event_vars))],
df["weight"].unique().shape[0],
).ravel()
dfw = df.iloc[indices]
sf, sfup, sfdown = weighted_mean_numba(
event_vars[0].content, dfw["weight"].values, dfw["corr"].values,
dfw["unc_up"].values, dfw["unc_down"].values,
add_syst(ev).content, df["weight"].unique().shape[0],
)
return awk.JaggedArray(
event_vars[0].starts, event_vars[0].stops,
weight_numba(sf, nsig, sfup, sfdown),
)
def fevaluate_btagsf(ev, evidx, nsig, source, attrs_):
jet_flavour = dict_apply(h2f, ev.Jet.hadronFlavour.content)
# Create mask
mask = np.ones((jet_flavour.shape[0], df.shape[0]), dtype=np.bool8)
# Flavour mask
event_attrs = [jet_flavour.astype(np.float32)]
mins = [df["jetFlavor"].values.astype(np.float32)]
maxs = [(df["jetFlavor"].values + 1).astype(np.float32)]
for jet_attr, df_attr in attrs_:
obj_attr = getattr(ev.Jet, jet_attr)
if callable(obj_attr):
obj_attr = obj_attr(ev)
event_attrs.append(obj_attr.content.astype(np.float32))
mins.append(df[df_attr + "Min"].values.astype(np.float32))
maxs.append(df[df_attr + "Max"].values.astype(np.float32))
# Create indices from mask
indices = get_bin_indices(event_attrs, mins, maxs, 3)
idx_central = indices[:, 0]
idx_down = indices[:, 1]
idx_up = indices[:, 2]
jpt = ev.Jet.ptShift(ev)
sf = btag_formula(jpt.content, df.iloc[idx_central])
sf_up = btag_formula(jpt.content, df.iloc[idx_up])
sf_down = btag_formula(jpt.content, df.iloc[idx_down])
sf_up = (source == "btagSF") * (sf_up / sf - 1.)
sf_down = (source == "btagSF") * (sf_down / sf - 1.)
return awk.JaggedArray(
jpt.starts,
jpt.stops,
weight_numba(sf, nsig, sf_up, sf_down),
)
def do_jer_correction(self, event):
indices = get_bin_indices(
[event.Jet_eta.content],
[self.jersfs["eta_low"].values],
[self.jersfs["eta_high"].values],
1,
)[:, 0]
ressfs = self.jersfs.iloc[indices][["corr", "corr_up",
"corr_down"]].values
jersfs = np.ones_like(event.Jet_pt.content, dtype=np.float32)
jersfs_up = np.ones_like(event.Jet_pt.content, dtype=np.float32)
jersfs_down = np.ones_like(event.Jet_pt.content, dtype=np.float32)
# match gen jets
gidx = event.Jet_genJetIdx
gsize = event.GenJet_pt.counts
mask = (gidx >= 0) & (gidx < gsize)
indices = (event.GenJet_pt.starts + gidx[mask]).content
gpt_matched = event.GenJet_pt.content[indices]
gen_var = 1. - gpt_matched / event.Jet_pt[mask].content
jersfs[mask.content] = 1. + (ressfs[mask.content, 0] - 1.) * gen_var
jersfs_up[mask.content] = 1. + (ressfs[mask.content, 1] - 1.) * gen_var
jersfs_down[
mask.content] = 1. + (ressfs[mask.content, 2] - 1.) * gen_var
# unmatched gen jets
gaus_var = np.random.normal(0.,
event.Jet_ptResolution[gidx < 0].content)
ressfs_mod = ressfs[(gidx < 0).content]**2 - 1.
ressfs_mod[ressfs_mod < 0.] = 0.
jersfs[(gidx < 0).content] = 1. + gaus_var * np.sqrt(ressfs_mod[:, 0])
jersfs_up[(gidx <
0).content] = 1. + gaus_var * np.sqrt(ressfs_mod[:, 1])
jersfs_down[(gidx <
0).content] = 1. + gaus_var * np.sqrt(ressfs_mod[:, 2])
# negative checks
jersfs[jersfs < 0.] = 0.
jersfs_up[jersfs_up < 0.] = 0.
jersfs_down[jersfs_down < 0.] = 0.
# write to event
starts, stops = event.Jet_pt.starts, event.Jet_pt.stops
event.Jet_JECjerSF = awk.JaggedArray(starts, stops, jersfs)
event.Jet_JECjerSFUp = awk.JaggedArray(
starts,
stops,
np.divide(jersfs_up,
jersfs,
out=np.ones_like(jersfs),
where=(jersfs != 0.)) - 1.,
)
event.Jet_JECjerSFDown = awk.JaggedArray(
starts,
stops,
np.divide(jersfs_down,
jersfs,
out=np.ones_like(jersfs),
where=(jersfs != 0.)) - 1.,
)
if self.apply_jer_corrections:
event.Jet_ptJESOnly = event.Jet_pt[:, :]
event.MET_ptJESOnly = event.MET_pt[:]
event.MET_phiJESOnly = event.MET_phi[:]
event.Jet_pt = (event.Jet_pt * event.Jet_JECjerSF)[:, :]
met, mephi = met_shift(event, self.unclust_threshold)
event.MET_pt = met[:]
event.MET_phi = mephi[:]