WIP: Update backend

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<=i<len(genJet_pt_orig) else -1 for i in jet_genJetIdx])
+
+        # remove already applied corrections
+        jet_pt_raw   = jet_pt   * (1.-jet_rawFactor)
+        jet_mass_raw = jet_mass * (1.-jet_rawFactor)
+
+        # get rho value
+        # TODO is this the correct one?
+        rho = getattr(event, "fixedGridRhoFastjetAll")
+
+        # get nominal JEC correction
+        nomJESsf = self.nomJES.evaluate(jet_area, jet_eta, jet_pt_raw, rho)
+
+        # apply the nominal jec SFs
+        jet_pt_corr = jet_pt_raw*nomJESsf
+        jet_mass_corr = jet_mass_raw*nomJESsf
+
+        # get nominal JER correction
+        nomJERsf = self.nomJER.evaluate(jet_eta, "nom")
+        upJERsf  = self.nomJER.evaluate(jet_eta, "up")
+        dnJERsf  = self.nomJER.evaluate(jet_eta, "down")
+
+        # get jet pt resolution
+        # TODO which pt to use here? raw or jes corrected?
+        res = self.ptRes.evaluate(jet_eta, jet_pt_corr, rho)
+
+        # if jet has a genJet:
+        # get dPt = jet_pt - genJet_pt
+        # then apply smear factor as
+        # sf = 1 + (nomJERsf-1.)*dPt/jet_pt
+        # i.e. the larger the jet/genjet difference is the more weight is put on the smearing factor
+        dPtRel = (jet_pt_corr-genJet_pt)/jet_pt_corr
+        genJetCorr   = (nomJERsf-1.)*dPtRel*(genJet_pt>0.)
+        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)
+

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