Skip to content

Commit

Permalink
addressing pylint errors
Browse files Browse the repository at this point in the history
  • Loading branch information
anpicci committed Dec 19, 2024
1 parent 71d3485 commit 6f70518
Show file tree
Hide file tree
Showing 3 changed files with 136 additions and 124 deletions.
8 changes: 4 additions & 4 deletions analysis/topeft_run2/analysis_processor.py
Original file line number Diff line number Diff line change
Expand Up @@ -363,7 +363,7 @@ def process(self, events):
obj_correction_syst_lst.append("TESDown")
obj_correction_syst_lst.append("FESUp")
obj_correction_syst_lst.append("FESDown")

wgt_correction_syst_lst = [
"lepSF_muonUp","lepSF_muonDown","lepSF_elecUp","lepSF_elecDown",f"btagSFbc_{year}Up",f"btagSFbc_{year}Down","btagSFbc_corrUp","btagSFbc_corrDown",f"btagSFlight_{year}Up",f"btagSFlight_{year}Down","btagSFlight_corrUp","btagSFlight_corrDown","PUUp","PUDown","PreFiringUp","PreFiringDown",f"triggerSF_{year}Up",f"triggerSF_{year}Down", # Exp systs
"FSRUp","FSRDown","ISRUp","ISRDown","renormUp","renormDown","factUp","factDown", # Theory systs
Expand All @@ -373,7 +373,7 @@ def process(self, events):
wgt_correction_syst_lst.append("lepSF_taus_realDown")
wgt_correction_syst_lst.append("lepSF_taus_fakeUp")
wgt_correction_syst_lst.append("lepSF_taus_fakeDown")

data_syst_lst = [
"FFUp","FFDown","FFptUp","FFptDown","FFetaUp","FFetaDown",f"FFcloseEl_{year}Up",f"FFcloseEl_{year}Down",f"FFcloseMu_{year}Up",f"FFcloseMu_{year}Down"
]
Expand Down Expand Up @@ -453,12 +453,12 @@ def process(self, events):
if self.tau_h_analysis:
tau["pt"], tau["mass"] = ApplyTESSystematic(year, tau, isData, syst_var)
tau["pt"], tau["mass"] = ApplyFESSystematic(year, tau, isData, syst_var)

events_cache = events.caches[0]
cleanedJets = ApplyJetCorrections(year, corr_type='jets', isData=isData, era=run_era).build(cleanedJets, lazy_cache=events_cache) #Run3 ready
cleanedJets = ApplyJetSystematics(year,cleanedJets,syst_var)
met = ApplyJetCorrections(year, corr_type='met', isData=isData, era=run_era).build(met_raw, cleanedJets, lazy_cache=events_cache)

cleanedJets["isGood"] = tc_os.is_tight_jet(getattr(cleanedJets, jetptname), cleanedJets.eta, cleanedJets.jetId, pt_cut=30., eta_cut=get_te_param("eta_j_cut"), id_cut=get_te_param("jet_id_cut"))
cleanedJets["isFwd"] = te_os.isFwdJet(getattr(cleanedJets, jetptname), cleanedJets.eta, cleanedJets.jetId, jetPtCut=40.)
goodJets = cleanedJets[cleanedJets.isGood]
Expand Down
2 changes: 1 addition & 1 deletion analysis/topeft_run2/datacards_post_processing.py
Original file line number Diff line number Diff line change
Expand Up @@ -144,7 +144,7 @@ def main():
os.mkdir(ptzlj0pt_path)
if args.set_up_top22006:
print(f"\nCopying TOP-22-006 relevant files to {ptzlj0pt_path}...")

if args.set_up_offZdivision:
print(f"\nCopying 3l-offZ-division relevant files to {ptzlj0pt_path}...")
for fname in datacard_files:
Expand Down
250 changes: 131 additions & 119 deletions topeft/modules/corrections.py
Original file line number Diff line number Diff line change
Expand Up @@ -568,91 +568,95 @@ def ApplyTES(year, taus, isData, tagger, syst_name, vsJetWP):
fes = np.where(whereFlag, SFevaluator['TauFES_{year}'.format(year=year)](eta,dm), 1)

if False:
## Correction-lib implementation - MUST BE TESTED WHEN TAU IN THE MASTER BRANCH PROCESSOR
arg_tau = ["pt", "eta", "decayMode", "genPartFlav"]
pt_mask_flat = ak.flatten((pt>0) & (pt<1000))

deep_tau_cuts = [
("DeepTau2017v2p1VSjet", ak.flatten(padded_taus[f"is{vsJetWP}"]>0), ("pt", "decayMode", "genPartFlav", vsJetWP)),
("DeepTau2017v2p1VSe", ak.flatten(padded_taus["iseTight"]>0), ("eta", "genPartFlav", "VVLoose")),
("DeepTau2017v2p1VSmu", ak.flatten(padded_taus["ismTight"]>0), ("eta", "genPartFlav", "Loose")),
]

DT_sf_list = []
DT_up_list = []
DT_do_list = []

for deep_tau_cut in deep_tau_cuts:
discr = deep_tau_cut[0]
id_mask = deep_tau_cut[1]
arg_list = deep_tau_cut[2]
args = []
for ttag in arg_list:
args.append(
ak.flatten(padded_taus[ttag])
)
tau_mask = id_mask & pt_mask_flat

args_sf = args + ["sf"]
DT_sf_list.append(
ak.where(
~tau_mask,
1,
ceval[discr].evaluate(*args_sf)
)
)
args_up = args + ["up"]
DT_up_list.append(
ak.where(
~tau_mask,
1,
ceval[discr].evaluate(*args_up)
)
)
args_down = args + ["down"]
DT_do_list.append(
ak.where(
~tau_mask,
1,
ceval[discr].evaluate(*args_down)
)
)

DT_sf_flat = None
DT_up_flat = None
DT_do_flat = None
for idr, DT_sf_discr in enumerate(DT_sf_list):
DT_sf_discr = ak.to_numpy(DT_sf_discr)
DT_up_discr = ak.to_numpy(DT_up_list[idr])
DT_do_discr = ak.to_numpy(DT_do_list[idr])
if idr == 0:
DT_sf_flat = DT_sf_discr
DT_up_flat = DT_up_discr
DT_do_flat = DT_do_discr
else:
DT_sf_flat *= DT_sf_discr
DT_up_flat *= DT_up_discr
DT_do_flat *= DT_do_discr

DT_sf = ak.unflatten(DT_sf_flat, ak.num(pt))
DT_up = ak.unflatten(DT_up_flat, ak.num(pt))
DT_do = ak.unflatten(DT_do_flat, ak.num(pt))
## end of correction-lib implementation
## Correction-lib implementation - MUST BE TESTED WHEN TAU IN THE MASTER BRANCH PROCESSOR
clib_year = clib_year_map[year]
json_path = topcoffea_path(f"data/POG/TAU/{clib_year}/tau.json.gz")
ceval = correctionlib.CorrectionSet.from_file(json_path)

arg_tau = ["pt", "eta", "decayMode", "genPartFlav"]
pt_mask_flat = ak.flatten((pt>0) & (pt<1000))

deep_tau_cuts = [
("DeepTau2017v2p1VSjet", ak.flatten(padded_taus[f"is{vsJetWP}"]>0), ("pt", "decayMode", "genPartFlav", vsJetWP)),
("DeepTau2017v2p1VSe", ak.flatten(padded_taus["iseTight"]>0), ("eta", "genPartFlav", "VVLoose")),
("DeepTau2017v2p1VSmu", ak.flatten(padded_taus["ismTight"]>0), ("eta", "genPartFlav", "Loose")),
]

DT_sf_list = []
DT_up_list = []
DT_do_list = []

for deep_tau_cut in deep_tau_cuts:
discr = deep_tau_cut[0]
id_mask = deep_tau_cut[1]
arg_list = deep_tau_cut[2]
args = []
for ttag in arg_list:
args.append(
ak.flatten(padded_taus[ttag])
)
tau_mask = id_mask & pt_mask_flat

args_sf = args + ["sf"]
DT_sf_list.append(
ak.where(
~tau_mask,
1,
ceval[discr].evaluate(*args_sf)
)
)
args_up = args + ["up"]
DT_up_list.append(
ak.where(
~tau_mask,
1,
ceval[discr].evaluate(*args_up)
)
)
args_down = args + ["down"]
DT_do_list.append(
ak.where(
~tau_mask,
1,
ceval[discr].evaluate(*args_down)
)
)

DT_sf_flat = None
DT_up_flat = None
DT_do_flat = None
for idr, DT_sf_discr in enumerate(DT_sf_list):
DT_sf_discr = ak.to_numpy(DT_sf_discr)
DT_up_discr = ak.to_numpy(DT_up_list[idr])
DT_do_discr = ak.to_numpy(DT_do_list[idr])
if idr == 0:
DT_sf_flat = DT_sf_discr
DT_up_flat = DT_up_discr
DT_do_flat = DT_do_discr
else:
DT_sf_flat *= DT_sf_discr
DT_up_flat *= DT_up_discr
DT_do_flat *= DT_do_discr

DT_sf = ak.unflatten(DT_sf_flat, ak.num(pt))
DT_up = ak.unflatten(DT_up_flat, ak.num(pt))
DT_do = ak.unflatten(DT_do_flat, ak.num(pt))
## end of correction-lib implementation

return (taus.pt*tes*fes, taus.mass*tes*fes)

def ApplyTESSystematic(year, Taus, isData, syst_name):
def ApplyTESSystematic(year, taus, isData, syst_name):
if not syst_name.startswith('TES'):
return (Taus.pt)
return (taus.pt)
if isData:
return (Taus.pt)
return (taus.pt)

pt = Taus.pt
dm = Taus.decayMode
gen = Taus.genPartFlav
pt = taus.pt
dm = taus.decayMode
gen = taus.genPartFlav

kinFlag = (pt>20) & (pt<205) & (gen==5)
dmFlag = ((Taus.decayMode==0) | (Taus.decayMode==1) | (Taus.decayMode==10) | (Taus.decayMode==11))
dmFlag = ((taus.decayMode==0) | (taus.decayMode==1) | (taus.decayMode==10) | (taus.decayMode==11))
whereFlag = kinFlag & dmFlag
syst_lab = f'TauTES_{year}'

Expand All @@ -662,21 +666,21 @@ def ApplyTESSystematic(year, Taus, isData, syst_name):
syst_lab += '_down'

tes_syst = np.where(whereFlag, SFevaluator['TauTES_{year}'.format(year=year)](dm,pt), 1)
return (Taus.pt*tes_syst, Taus.mass*tes_syst)
def ApplyFESSystematic(year, Taus, isData, syst_name):
return (taus.pt*tes_syst, taus.mass*tes_syst)

def ApplyFESSystematic(year, taus, isData, syst_name):
if not syst_name.startswith('FES'):
return (Taus.pt)
return (taus.pt)
if isData:
return (Taus.pt)
return (taus.pt)

pt = Taus.pt
eta = Taus.eta
dm = Taus.decayMode
gen = Taus.genPartFlav
pt = taus.pt
eta = taus.eta
dm = taus.decayMode
gen = taus.genPartFlav

kinFlag = (pt>20) & (pt<205) & (gen==5)
dmFlag = ((Taus.decayMode==0) | (Taus.decayMode==1))
dmFlag = ((taus.decayMode==0) | (taus.decayMode==1))
whereFlag = kinFlag & dmFlag

syst_lab = f'TauFES_{year}'
Expand All @@ -687,65 +691,73 @@ def ApplyFESSystematic(year, Taus, isData, syst_name):
syst_lab += '_down'

fes_syst = np.where(whereFlag, SFevaluator['TauFES_{year}'.format(year=year)](eta,dm), 1)
return (Taus.pt*fes_syst, Taus.mass*fes_syst)
def AttachTauSF(events, Taus, year, vsJetWP="Loose"):
padded_Taus = ak.pad_none(Taus,1)
padded_Taus = ak.with_name(padded_Taus, "TauCandidate")
padded_Taus["sf_tau"] = 1.0
padded_Taus["sf_tau_up"] = 1.0
padded_Taus["sf_tau_down"] = 1.0
return (taus.pt*fes_syst, taus.mass*fes_syst)

def AttachTauSF(events, taus, year, vsJetWP="Loose"):
padded_taus = ak.pad_none(taus,1)
padded_taus = ak.with_name(padded_taus, "TauCandidate")
padded_taus["sf_tau"] = 1.0
padded_taus["sf_tau_up"] = 1.0
padded_taus["sf_tau_down"] = 1.0

clib_year = clib_year_map[year]
json_path = topcoffea_path(f"data/POG/TAU/{clib_year}/tau.json.gz")
ceval = correctionlib.CorrectionSet.from_file(json_path)

pt = padded_Taus.pt
dm = padded_Taus.decayMode
wp = padded_Taus.idDeepTau2017v2p1VSjet
eta = padded_Taus.eta
gen = padded_Taus.genPartFlav
mass= padded_Taus.mass
pt = padded_taus.pt
dm = padded_taus.decayMode
wp = padded_taus.idDeepTau2017v2p1VSjet
eta = padded_taus.eta
gen = padded_taus.genPartFlav
mass= padded_taus.mass

## legacy
whereFlag = ((pt>20) & (pt<205) & (gen==5) & (padded_Taus[f"is{vsJetWP}"]>0))
whereFlag = ((pt>20) & (pt<205) & (gen==5) & (padded_taus[f"is{vsJetWP}"]>0))
real_sf_loose = np.where(whereFlag, SFevaluator[f'TauSF_{year}_{vsJetWP}'](dm,pt), 1)
real_sf_loose_up = np.where(whereFlag, SFevaluator[f'TauSF_{year}_{vsJetWP}_up'](dm,pt), 1)
real_sf_loose_down = np.where(whereFlag, SFevaluator[f'TauSF_{year}_{vsJetWP}_down'](dm,pt), 1)

whereFlag = ((pt>20) & (pt<205) & ((gen==1)|(gen==3)) & (padded_Taus["iseTight"]>0))
whereFlag = ((pt>20) & (pt<205) & ((gen==1)|(gen==3)) & (padded_taus["iseTight"]>0))
fake_elec_sf = np.where(whereFlag, SFevaluator[f'Tau_elecFakeSF_{year}'](np.abs(eta)), 1)
fake_elec_sf_up = np.where(whereFlag, SFevaluator[f'Tau_elecFakeSF_{year}_up'](np.abs(eta)), 1)
fake_elec_sf_down = np.where(whereFlag, SFevaluator[f'Tau_elecFakeSF_{year}_down'](np.abs(eta)), 1)
whereFlag = ((pt>20) & (pt<205) & ((gen==2)|(gen==4)) & (padded_Taus["ismTight"]>0))
whereFlag = ((pt>20) & (pt<205) & ((gen==2)|(gen==4)) & (padded_taus["ismTight"]>0))
fake_muon_sf = np.where(whereFlag, SFevaluator[f'Tau_muonFakeSF_{year}'](np.abs(eta)), 1)
fake_muon_sf_up = np.where(whereFlag, SFevaluator[f'Tau_muonFakeSF_{year}_up'](np.abs(eta)), 1)
fake_muon_sf_down = np.where(whereFlag, SFevaluator[f'Tau_muonFakeSF_{year}_down'](np.abs(eta)), 1)

whereFlag = ((pt>20) & (pt<205) & (gen!=5) & (gen!=4) & (gen!=3) & (gen!=2) & (gen!=1) & (padded_Taus["isLoose"]>0))
whereFlag = ((pt>20) & (pt<205) & (gen!=5) & (gen!=4) & (gen!=3) & (gen!=2) & (gen!=1) & (padded_taus["isLoose"]>0))
new_fake_sf = np.where(whereFlag, SFevaluator['TauFakeSF'](pt), 1)
new_fake_sf_up = np.where(whereFlag, SFevaluator['TauFakeSF_up'](pt), 1)
new_fake_sf_down = np.where(whereFlag, SFevaluator['TauFakeSF_down'](pt), 1)

real_sf = real_sf_loose
real_sf_up = real_sf_loose_up
real_sf_down = real_sf_loose_down
padded_Taus["sf_tau_real"] = real_sf
padded_Taus["sf_tau_real_up"] = real_sf_up
padded_Taus["sf_tau_real_down"] = real_sf_down
padded_Taus["sf_tau_fake"] = fake_elec_sf*fake_muon_sf*new_fake_sf
padded_Taus["sf_tau_fake_up"] = fake_elec_sf_up*fake_muon_sf_up*new_fake_sf_up
padded_Taus["sf_tau_fake_down"] = fake_elec_sf_down*fake_muon_sf_down*new_fake_sf_down

events["sf_2l_taus_real"] = padded_Taus.sf_tau_real[:,0]
events["sf_2l_taus_real_hi"] = padded_Taus.sf_tau_real_up[:,0]
events["sf_2l_taus_real_lo"] = padded_Taus.sf_tau_real_down[:,0]
events["sf_2l_taus_fake"] = padded_Taus.sf_tau_fake[:,0]
events["sf_2l_taus_fake_hi"] = padded_Taus.sf_tau_fake_up[:,0]
events["sf_2l_taus_fake_lo"] = padded_Taus.sf_tau_fake_down[:,0]
padded_taus["sf_tau_real"] = real_sf
padded_taus["sf_tau_real_up"] = real_sf_up
padded_taus["sf_tau_real_down"] = real_sf_down
padded_taus["sf_tau_fake"] = fake_elec_sf*fake_muon_sf*new_fake_sf
padded_taus["sf_tau_fake_up"] = fake_elec_sf_up*fake_muon_sf_up*new_fake_sf_up
padded_taus["sf_tau_fake_down"] = fake_elec_sf_down*fake_muon_sf_down*new_fake_sf_down

events["sf_2l_taus_real"] = padded_taus.sf_tau_real[:,0]
events["sf_2l_taus_real_hi"] = padded_taus.sf_tau_real_up[:,0]
events["sf_2l_taus_real_lo"] = padded_taus.sf_tau_real_down[:,0]
events["sf_2l_taus_fake"] = padded_taus.sf_tau_fake[:,0]
events["sf_2l_taus_fake_hi"] = padded_taus.sf_tau_fake_up[:,0]
events["sf_2l_taus_fake_lo"] = padded_taus.sf_tau_fake_down[:,0]

if False:
## Correction-lib implementation - MUST BE TESTED WHEN TAU IN THE MASTER BRANCH PROCESSOR
padded_taus["sf_tau"] = 1.0
padded_taus["sf_tau_up"] = 1.0
padded_taus["sf_tau_down"] = 1.0

clib_year = clib_year_map[year]
json_path = topcoffea_path(f"data/POG/TAU/{clib_year}/tau.json.gz")
ceval = correctionlib.CorrectionSet.from_file(json_path)

DT_sf_list = []
DT_up_list = []
DT_do_list = []
Expand Down Expand Up @@ -822,7 +834,7 @@ def AttachTauSF(events, Taus, year, vsJetWP="Loose"):
events["sf_2l_taus_hi"] = padded_taus.sf_tau_up[:,0]
events["sf_2l_taus_lo"] = padded_taus.sf_tau_down[:,0]
## end of correction-lib implementation

def AttachPerLeptonFR(leps, flavor, year):
# Get the flip rates lookup object
if year not in clib_year_map.keys():
Expand Down

0 comments on commit 6f70518

Please sign in to comment.