diff --git a/python/postprocessing/modules/kit/jecHelper.py b/python/postprocessing/modules/kit/jecHelper.py new file mode 100644 index 00000000000..fdc0dc1b96c --- /dev/null +++ b/python/postprocessing/modules/kit/jecHelper.py @@ -0,0 +1,237 @@ +from PhysicsTools.NanoAODTools.postprocessing.framework.datamodel import Collection + +import correctionlib +import numpy as np +import sys +import os + +class jecHelper: + def __init__(self, dataEra, runType, isData, jecNames, splitJER, addHEMIssue): + + # get tags from inbuilt function + self.jesTag, self.jerTag = self.getJECTags(dataEra, runType, isData) + + # read the correct json file + self.json = correctionlib.CorrectionSet.from_file( + self.getInputPath(dataEra, runType)) + + # get nominal compound correction + self.nomJES = self.json.compound[f"{self.jesTag}_L1L2L3Res_AK4PFchs"] + # get nominal jer + self.nomJER = self.json[f"{self.jerTag}_ScaleFactor_AK4PFchs"] + # get pt resolution + self.ptRes = self.json[f"{self.jerTag}_PtResolution_AK4PFchs"] + + # get jes sources + self.jesSource = {} + for jes in jecNames: + self.jesSource[jes] = self.json[f"{self.jesTag}_{jes}_AK4PFchs"] + # save some settings + self.jecNames = jecNames + self.splitJER = splitJER + self.addHEMIssue = addHEMIssue + + def getInputPath(self, era, runType): + #basePath = "/cvmfs/cms.cern.ch/rsync/cms-nanoAOD/jsonpog-integration/POG/JME" + basePath = "/nfs/dust/cms/user/vdlinden/nanoDev/JME" + if runType == "ul": + path = os.path.join(basePath, era+"_UL", "jet_jerc.json.gz") + else: + sys.exit(f"invalid runtype for JER tag: {runType}") + return path + + + def getJECTags(self, era, runType, isData): + if runType == "ul": + if era == "2018": + jertag="Summer19UL18_JRV2" + jestag="Summer19UL18_V5" + elif era == "2017": + jertag="Summer19UL17_JRV2" + jestag="Summer19UL17_V5" + else: + sys.exit(f"invalid era for ul JER tag: {era}") + else: + sys.exit(f"invalid runtype for JER tag: {runType}") + + if isData: + return jestag+"_DATA", jertag+"_DATA" + else: + return jestag+"_MC", jertag+"_MC" + + def getJERsplitID(self, pt, eta): + if not self.splitJER: + return "" + if abs(eta) < 1.93: + return 0 + elif abs(eta) < 2.5: + return 1 + + def set(self, event, branch, value): + event.__setattr__(branch, value) + + def applyJEC(self, event): + #print("applying JECs to current event") + + # loading jet collections + jets = Collection(event, "Jet") + genjets = Collection(event, "GenJet") + + # loading jet features + jet_pt = np.array([jet.pt for jet in jets]) + jet_mass = np.array([jet.mass for jet in jets]) + jet_eta = np.array([jet.eta for jet in jets]) + jet_area = np.array([jet.area for jet in jets]) + jet_rawFactor = np.array([jet.rawFactor for jet in jets]) + jet_genJetIdx = np.array([jet.genJetIdx for jet in jets]) + + # loading genJet pt + genJet_pt_orig = np.array([jet.pt for jet in genjets]) + # get gen jet pts of jets + genJet_pt = np.array([genJet_pt_orig[i] if 0<=i0.) + genJetCorrUp = (upJERsf-1.)*dPtRel*(genJet_pt>0.) + genJetCorrDn = (dnJERsf-1.)*dPtRel*(genJet_pt>0.) + # if no genjet is available: + # randomized smearing based on jetpt resolution + rnd = np.random.normal(0., res) + # only apply smear factor when SF is larger than 1 and no genJet + smearCorr = (rnd*np.sqrt(np.fmax(nomJERsf,1.)**2-1.))*(genJet_pt==0) + smearCorrUp = (rnd*np.sqrt(np.fmax(upJERsf,1.)**2-1.))*(genJet_pt==0) + smearCorrDn = (rnd*np.sqrt(np.fmax(dnJERsf,1.)**2-1.))*(genJet_pt==0) + + #print("genJetCorr", genJetCorr) + #print("smearCorr", smearCorr) + # final smear factor: + smearFactor = 1. + genJetCorr + smearCorr + smearFactorUp = 1. + genJetCorrUp + smearCorrUp + smearFactorDn = 1. + genJetCorrDn + smearCorrDn + #print("final smear factor", smearFactor) + + # check the updated jet vectors in case the energy gets negative + # TODO dont know how this could happen but is also done in nanoAOD-tools + for i, jet in enumerate(jets): + jet.pt = jet_pt_corr[i] + jet.mass = jet_mass_corr[i] + energy = jet.p4().E() + if energy*smearFactor[i] < 1e-2: + #print(f"energy too low {energy}") + smearFactor[i] = 1e-2/energy + if energy*smearFactorUp[i] < 1e-2: + #print(f"energy too low {energy}") + smearFactorUp[i] = 1e-2/energy + if energy*smearFactorDn[i] < 1e-2: + #print(f"energy too low {energy}") + smearFactorDn[i] = 1e-2/energy + + + + + #print("original jetpt", jet_pt) + #print("jes corr jetpt", jet_pt_corr) + jet_pt_nom = jet_pt_corr*smearFactor + setattr(event, "Jet_pt_nom", jet_pt_nom) + jet_mass_nom = jet_mass_corr*smearFactor + setattr(event, "Jet_mass_nom", jet_mass_nom) + + # jer variations + jet_pt_jerUp = jet_pt_corr*smearFactorUp + setattr(event, "Jet_pt_jerUp", jet_pt_jerUp) + jet_mass_jerUp = jet_mass_corr*smearFactorUp + setattr(event, "Jet_mass_jerUp", jet_mass_jerUp) + + jet_pt_jerDown = jet_pt_corr*smearFactorDn + setattr(event, "Jet_pt_jerDown", jet_pt_jerDown) + jet_mass_jerDown = jet_mass_corr*smearFactorDn + setattr(event, "Jet_mass_jerDown", jet_mass_jerDown) + if self.splitJER: + # split JER per ID + # jer0: eta < 1.93 + # jer1: eta < 2.5 + # jer2-5: eta > 2.5 TODO implement also this split, not needed for now + # jer2: eta < 3 and pt < 50 + # jer3: eta < 3 and pt >= 50 + # jer4: eta >= 3 and pt < 50 + # jer5: eta >= 3 and pt >= 50 + jerID = 0.+(np.abs(jet_eta)>=1.93)*(1.+(np.abs(jet_eta)>=2.5)) + + jet_pt_jer0Up = jet_pt_corr*(1.+(smearFactorUp-1.)*(jerID==0)) + setattr(event, "Jet_pt_jer0Up", jet_pt_jer0Up) + jet_mass_jer0Up = jet_mass_corr*(1.+(smearFactorUp-1.)*(jerID==0)) + setattr(event, "Jet_mass_jer0Up", jet_mass_jer0Up) + + jet_pt_jer1Up = jet_pt_corr*(1.+(smearFactorUp-1.)*(jerID==1)) + setattr(event, "Jet_pt_jer1Up", jet_pt_jer1Up) + jet_mass_jer1Up = jet_mass_corr*(1.+(smearFactorUp-1.)*(jerID==1)) + setattr(event, "Jet_mass_jer1Up", jet_mass_jer1Up) + + jet_pt_jer0Down = jet_pt_corr*(1.+(smearFactorDn-1.)*(jerID==0)) + setattr(event, "Jet_pt_jer0Down", jet_pt_jer0Down) + jet_mass_jer0Down = jet_mass_corr*(1.+(smearFactorDn-1.)*(jerID==0)) + setattr(event, "Jet_mass_jer0Down", jet_mass_jer0Down) + + jet_pt_jer1Down = jet_pt_corr*(1.+(smearFactorDn-1.)*(jerID==1)) + setattr(event, "Jet_pt_jer1Down", jet_pt_jer1Down) + jet_mass_jer1Down = jet_mass_corr*(1.+(smearFactorDn-1.)*(jerID==1)) + setattr(event, "Jet_mass_jer1Down", jet_mass_jer1Down) + + + + #print("jes+jer corr jetpt", jet_pt_nom) + #print("jet eta", jet_eta) + #print("jer id", jerID) + #print("jerUp", jet_pt_jerUp) + #print("jerDown", jet_pt_jerDown) + #print("jer0Up", jet_pt_jer0Up) + #print("jer0Down", jet_pt_jer0Down) + #print("jer1Up", jet_pt_jer1Up) + #print("jer1Down", jet_pt_jer1Down) + + + # jes uncertainties + for jes in self.jecNames: + #print(jes) + # read the uncertainties with the jer+jes corrected pt + delta = self.jesSource[jes].evaluate(jet_pt_nom, jet_eta) + jet_pt_jesUp = jet_pt_nom*(1.+delta) + setattr(event, f"Jet_pt_jes{jes}Up", jet_pt_jesUp) + jet_pt_jesDn = jet_pt_nom*(1.-delta) + setattr(event, f"Jet_pt_jes{jes}Down", jet_pt_jesDn) + jet_mass_jesUp = jet_mass_nom*(1.+delta) + setattr(event, f"Jet_mass_jes{jes}Up", jet_mass_jesUp) + jet_mass_jesDn = jet_mass_nom*(1.-delta) + setattr(event, f"Jet_mass_jes{jes}Down", jet_mass_jesDn) + #print("ptup", jet_pt_jesUp) + #print("ptdn", jet_pt_jesDn) + diff --git a/python/postprocessing/modules/kit/snape.py b/python/postprocessing/modules/kit/snape.py index e47975d08f9..bfa707ee341 100644 --- a/python/postprocessing/modules/kit/snape.py +++ b/python/postprocessing/modules/kit/snape.py @@ -11,6 +11,7 @@ from PhysicsTools.snape.container.CutFlow import CutFlow from PhysicsTools.snape.modules.BaseCalculator import BaseCalculator as base +from .jecHelper import jecHelper import os import importlib @@ -25,7 +26,8 @@ def __init__(self): def setup(self, configName = "config_testAnalysis", dataEra = "2017", runType = "legacy", isData = False, isTopNanoAOD = False, - sampleName = None, jecs = [""]): + jecs = ["nom"], splitJER = False, addHEMIssue = False, + sampleName = None): # load config module configImport = "PhysicsTools.snape.configs.{}".format(configName) @@ -37,7 +39,9 @@ def setup(self, configName = "config_testAnalysis", dataEra = "2017", isData = isData, isTopNanoAOD = isTopNanoAOD, sampleName = sampleName, - jecs = jecs) + jecs = jecs, + splitJER = splitJER, + addHEMIssue = addHEMIssue) # init snape module self.snape = VariableCalculator() @@ -46,6 +50,15 @@ def setup(self, configName = "config_testAnalysis", dataEra = "2017", # init branches in variable calculator self.snape.init_values() + # save jec settings + self.jecNames = jecs + self.jecs = self.snape.config.jecs + self.splitJER = splitJER + self.addHEMIssue = addHEMIssue + + # load jec stuff + self.jecHelper = jecHelper(dataEra, runType, isData, jecs, splitJER, addHEMIssue) + def beginJob(self, cutflowpath): self.cutflowpath = os.path.join(cutflowpath, "cutflow.txt") self.cutflow.output_file = self.cutflowpath @@ -126,6 +139,10 @@ def analyze(self, event): anyJECSelected = False # reset object container base.oc.reset() + + # add loading of jec sources here for now, change structure later + self.jecHelper.applyJEC(event) + # loop over jec sources for sys in self.snape.config.jecs: # set members of snape