Skip to content

Commit

Permalink
Fix code style issues with Black
Browse files Browse the repository at this point in the history
  • Loading branch information
lint-action committed Oct 24, 2023
1 parent 034e792 commit 147bc94
Show file tree
Hide file tree
Showing 2 changed files with 84 additions and 66 deletions.
141 changes: 80 additions & 61 deletions src/HHbbVV/postprocessing/postprocessing.py
Original file line number Diff line number Diff line change
Expand Up @@ -185,10 +185,10 @@ def get_nonres_vbf_selection_regions(
"VVFatJetParTMD_THWWvsT": [thww_wp, CUT_MAX_VAL],
"vbf_Mass_jj": [500, 10000],
"vbf_dEta_jj": [4, 10000],
"ak8FatJetEta0": [-2.4,2.4],
"ak8FatJetEta1": [-2.4,2.4],
"DijetdEta": [0,2.0],
"DijetdPhi": [2.6,10000],
"ak8FatJetEta0": [-2.4, 2.4],
"ak8FatJetEta1": [-2.4, 2.4],
"DijetdEta": [0, 2.0],
"DijetdPhi": [2.6, 10000],
# add more
},
label="Pass",
Expand All @@ -201,10 +201,10 @@ def get_nonres_vbf_selection_regions(
"VVFatJetParTMD_THWWvsT": [thww_wp, CUT_MAX_VAL],
"vbf_Mass_jj": [500, 10000],
"vbf_dEta_jj": [4, 10000],
"ak8FatJetEta0": [-2.4,2.4],
"ak8FatJetEta1": [-2.4,2.4],
"DijetdEta": [0,2.0],
"DijetdPhi": [2.6,10000],
"ak8FatJetEta0": [-2.4, 2.4],
"ak8FatJetEta1": [-2.4, 2.4],
"DijetdEta": [0, 2.0],
"DijetdPhi": [2.6, 10000],
# add more
},
label="Fail",
Expand Down Expand Up @@ -350,8 +350,8 @@ def get_res_selection_regions(
# TODO: check which of these applies to resonant as well
weight_shifts = {
"pileup": Syst(samples=nonres_sig_keys + res_sig_keys + bg_keys, label="Pileup"),
#"PDFalphaS": Syst(samples=nonres_sig_keys, label="PDF"),
#"QCDscale": Syst(samples=nonres_sig_keys, label="QCDscale"), # commented out
# "PDFalphaS": Syst(samples=nonres_sig_keys, label="PDF"),
# "QCDscale": Syst(samples=nonres_sig_keys, label="QCDscale"), # commented out
"ISRPartonShower": Syst(samples=nonres_sig_keys_ggf + ["V+Jets"], label="ISR Parton Shower"),
"FSRPartonShower": Syst(samples=nonres_sig_keys_ggf + ["V+Jets"], label="FSR Parton Shower"),
"L1EcalPrefiring": Syst(
Expand Down Expand Up @@ -402,34 +402,32 @@ def main(args):
jec_jmsr_shifts=True,
)
print("Loaded BDT preds\n")

# Load VBF Variables (edits events_dict so that all of the events have the appropriate variables and stuff) TODO:
if args.vbf:
if args.vbf:
pt_labels = ["JES", "JER"]
m_labels = ["JMS", "JMR"]
for key,df in events_dict.items():

for key, df in events_dict.items():
# for each JES,JER, and JMS JMR correction
ptlabel = "_JES_up"
mlabel = "_JMS_down"

bb_mask = (
df[("ak8FatJetParticleNetMD_Txbb", 0)]
> df[("ak8FatJetParticleNetMD_Txbb", 1)]
)

_add_vbf_columns(df,bb_mask,ptlabel="",mlabel="")
df[("ak8FatJetParticleNetMD_Txbb", 0)] > df[("ak8FatJetParticleNetMD_Txbb", 1)]
)

_add_vbf_columns(df, bb_mask, ptlabel="", mlabel="")

if key == 'Data':
if key == "Data":
continue

for var in pt_labels:
for direction in ["down", "up"]:
_add_vbf_columns(df,bb_mask,ptlabel=f"_{var}_{direction}",mlabel="")
_add_vbf_columns(df, bb_mask, ptlabel=f"_{var}_{direction}", mlabel="")

for var in m_labels:
for direction in ["down", "up"]:
_add_vbf_columns(df,bb_mask,ptlabel="",mlabel=f"_{var}_{direction}")
_add_vbf_columns(df, bb_mask, ptlabel="", mlabel=f"_{var}_{direction}")

# Control plots
if args.control_plots:
Expand Down Expand Up @@ -561,46 +559,56 @@ def _init(args):

return shape_vars, scan, scan_cuts, scan_wps


# adds all necessary columns to dataframes from events_dict
def _add_vbf_columns(df,bb_mask,ptlabel,mlabel):
def _add_vbf_columns(df, bb_mask, ptlabel, mlabel):
import vector

bbJet = vector.array(
{
"pt": np.where(bb_mask, df[f"ak8FatJetPt{ptlabel}"][0], df[f"ak8FatJetPt{ptlabel}"][1]),
"phi": np.where(bb_mask, df["ak8FatJetPhi"][0], df["ak8FatJetPhi"][1]),
"eta": np.where(bb_mask, df["ak8FatJetEta"][0], df["ak8FatJetEta"][1]),
"M": np.where(bb_mask, df[f"ak8FatJetParticleNetMass{mlabel}"][0], df[f"ak8FatJetParticleNetMass{mlabel}"][1]),
"M": np.where(
bb_mask,
df[f"ak8FatJetParticleNetMass{mlabel}"][0],
df[f"ak8FatJetParticleNetMass{mlabel}"][1],
),
}
)

VVJet = vector.array(
{
"pt": np.where(~bb_mask, df[f"ak8FatJetPt{ptlabel}"][0], df[f"ak8FatJetPt{ptlabel}"][1]),
"pt": np.where(
~bb_mask, df[f"ak8FatJetPt{ptlabel}"][0], df[f"ak8FatJetPt{ptlabel}"][1]
),
"phi": np.where(~bb_mask, df["ak8FatJetPhi"][0], df["ak8FatJetPhi"][1]),
"eta": np.where(~bb_mask, df["ak8FatJetEta"][0], df["ak8FatJetEta"][1]),
"M": np.where(~bb_mask, df[f"ak8FatJetParticleNetMass{mlabel}"][0], df[f"ak8FatJetParticleNetMass{mlabel}"][1]),
"M": np.where(
~bb_mask,
df[f"ak8FatJetParticleNetMass{mlabel}"][0],
df[f"ak8FatJetParticleNetMass{mlabel}"][1],
),
}
)


vbf1 = vector.array(
{
"pt": df[(f'VBFJetPt{ptlabel}', 0)],
"phi": df[('VBFJetPhi', 0)],
"eta": df[('VBFJetEta', 0)],
"M": df[('VBFJetMass', 0)],
}
)
{
"pt": df[(f"VBFJetPt{ptlabel}", 0)],
"phi": df[("VBFJetPhi", 0)],
"eta": df[("VBFJetEta", 0)],
"M": df[("VBFJetMass", 0)],
}
)

vbf2 = vector.array(
{
"pt": df[(f'VBFJetPt{ptlabel}', 1)],
"phi": df[('VBFJetPhi', 1)],
"eta": df[('VBFJetEta', 1)],
"M": df[('VBFJetMass', 1)],
}
)
{
"pt": df[(f"VBFJetPt{ptlabel}", 1)],
"phi": df[("VBFJetPhi", 1)],
"eta": df[("VBFJetEta", 1)],
"M": df[("VBFJetMass", 1)],
}
)

jj = vbf1 + vbf2

Expand All @@ -613,17 +621,27 @@ def _add_vbf_columns(df,bb_mask,ptlabel,mlabel):
# Adding variables defined in HIG-20-005 that show strong differentiation for VBF signal events and background

# seperation between both ak8 higgs jets
if "vbf_dR_HH" not in df.columns: df[f"vbf_dR_HH"] = VVJet.deltaR(bbJet)
if "vbf_dR_j0_HVV" not in df.columns: df[f"vbf_dR_j0_HVV"] = vbf1.deltaR(VVJet)
if "vbf_dR_j1_HVV" not in df.columns: df[f"vbf_dR_j1_HVV"] = vbf2.deltaR(VVJet)
if "vbf_dR_j0_Hbb" not in df.columns: df[f"vbf_dR_j0_Hbb"] = vbf1.deltaR(bbJet)
if "vbf_dR_j1_Hbb" not in df.columns: df[f"vbf_dR_j1_Hbb"] = vbf2.deltaR(bbJet)
if "vbf_dR_jj" not in df.columns: df[f"vbf_dR_jj"] = vbf1.deltaR(vbf2)
if "vbf_Mass_jj{ptlabel}" not in df.columns: df[f"vbf_Mass_jj{ptlabel}"] = jj.M
if "vbf_dEta_jj" not in df.columns: df[f"vbf_dEta_jj"] = np.abs(vbf1.eta - vbf2.eta)

if "DijetdEta" not in df.columns: df[f"DijetdEta"] = np.abs(bbJet.eta - VVJet.eta)
if "DijetdPhi" not in df.columns: df[f"DijetdPhi"] = np.abs(bbJet.phi - VVJet.phi)
if "vbf_dR_HH" not in df.columns:
df[f"vbf_dR_HH"] = VVJet.deltaR(bbJet)
if "vbf_dR_j0_HVV" not in df.columns:
df[f"vbf_dR_j0_HVV"] = vbf1.deltaR(VVJet)
if "vbf_dR_j1_HVV" not in df.columns:
df[f"vbf_dR_j1_HVV"] = vbf2.deltaR(VVJet)
if "vbf_dR_j0_Hbb" not in df.columns:
df[f"vbf_dR_j0_Hbb"] = vbf1.deltaR(bbJet)
if "vbf_dR_j1_Hbb" not in df.columns:
df[f"vbf_dR_j1_Hbb"] = vbf2.deltaR(bbJet)
if "vbf_dR_jj" not in df.columns:
df[f"vbf_dR_jj"] = vbf1.deltaR(vbf2)
if "vbf_Mass_jj{ptlabel}" not in df.columns:
df[f"vbf_Mass_jj{ptlabel}"] = jj.M
if "vbf_dEta_jj" not in df.columns:
df[f"vbf_dEta_jj"] = np.abs(vbf1.eta - vbf2.eta)

if "DijetdEta" not in df.columns:
df[f"DijetdEta"] = np.abs(bbJet.eta - VVJet.eta)
if "DijetdPhi" not in df.columns:
df[f"DijetdPhi"] = np.abs(bbJet.phi - VVJet.phi)

# Subleading VBF-jet cos(θ) in the HH+2j center of mass frame:
# https://github.com/scikit-hep/vector/blob/main/src/vector/_methods.py#L916
Expand All @@ -633,17 +651,18 @@ def _add_vbf_columns(df,bb_mask,ptlabel,mlabel):
# Leading VBF-jet cos(θ) in the HH+2j center of mass frame:
thetab1 = 2 * np.arctan(np.exp(-j1_CMF.eta))
thetab1 = np.cos(thetab1) # 12

if f"vbf_cos_j1{ptlabel}{mlabel}" not in df.columns: df[f"vbf_cos_j1{ptlabel}{mlabel}"] = np.abs(thetab1)

if f"vbf_cos_j1{ptlabel}{mlabel}" not in df.columns:
df[f"vbf_cos_j1{ptlabel}{mlabel}"] = np.abs(thetab1)

# Subleading VBF-jet cos(θ) in the HH+2j center of mass frame:
j2_CMF = vbf2.boostCM_of_p4(system_4vec)
thetab2 = 2 * np.arctan(np.exp(-j2_CMF.eta))
thetab2 = np.cos(thetab2)
if f"vbf_cos_j2{ptlabel}{mlabel}" not in df.columns: df[f"vbf_cos_j2{ptlabel}{mlabel}"] = np.abs(thetab2)
if f"vbf_cos_j2{ptlabel}{mlabel}" not in df.columns:
df[f"vbf_cos_j2{ptlabel}{mlabel}"] = np.abs(thetab2)


if "vbf_prod_centrality" not in df.columns:
if "vbf_prod_centrality" not in df.columns:
# H1-centrality * H2-centrality:
delta_eta = vbf1.eta - vbf2.eta
avg_eta = (vbf1.eta + vbf2.eta) / 2
Expand Down Expand Up @@ -1454,7 +1473,7 @@ def get_templates(
if sample in wsyst.samples and year in wsyst.years:
# print(wshift)
for skey, shift in [("Down", "down"), ("Up", "up")]:
print(sample,shift, "debugging on line 1319")
print(sample, shift, "debugging on line 1319")
if "QCDscale" in wshift:
# QCDscale7pt/QCDscale4
# https://github.com/LPC-HH/HHLooper/blob/master/python/prepare_card_SR_final.py#L263-L288
Expand Down Expand Up @@ -1516,7 +1535,7 @@ def get_templates(

if sig_splits is None:
sig_splits = [plot_sig_keys]
print(sig_splits,plot_sig_keys)
print(sig_splits, plot_sig_keys)
for i, shape_var in enumerate(shape_vars):
for j, p_sig_keys in enumerate(sig_splits):
split_str = "" if len(sig_splits) == 1 else f"sigs{j}_"
Expand Down
9 changes: 4 additions & 5 deletions src/HHbbVV/processors/bbVVSkimmer.py
Original file line number Diff line number Diff line change
Expand Up @@ -389,7 +389,6 @@ def process(self, events: ak.Array):
skimmed_events["nGoodVBFJets"] = np.array(ak.sum(vbf_jet_mask, axis=1))
# Temp
skimmed_events["gen_weights"] = gen_weights


# VBF ak4 Jet vars (pt, eta, phi, M, nGoodJets)
# if self._vbf_search:
Expand Down Expand Up @@ -633,15 +632,15 @@ def process(self, events: ak.Array):
systematics += list(weights.variations)

# TODO: need to be careful about the sum of gen weights used for the LHE/QCDScale uncertainties
logger.debug("weights ", weights._weights.keys())
logger.debug("weights ", weights._weights.keys())

# TEMP: save each individual weight
skimmed_events["L1EcalPrefiring"] = ak.to_numpy(events.L1PreFiringWeight.Nom)

for key in weights._weights.keys():
print(f"single_weight_{key}")
skimmed_events[f"single_weight_{key}"] = weights.partial_weight([key])

for systematic in systematics:
if systematic in weights.variations:
weight = weights.weight(modifier=systematic)
Expand Down

0 comments on commit 147bc94

Please sign in to comment.