Source code for Corrections.JME.smearing

import os
import tarfile
import tempfile

from analysis_tools.utils import import_root
from Base.Modules.baseModules import JetLepMetSyst

ROOT = import_root()


class jetSmearerRDFProducer():
    def __init__(self, isMC,ispreVFP,
            jerInputFileName="Spring16_25nsV10_MC_PtResolution_AK4PFchs.txt",
            jerUncertaintyInputFileName="Spring16_25nsV10_MC_SF_AK4PFchs.txt",
            *args, **kwargs):

        self.isMC = isMC
        self.ispreVFP = ispreVFP
        if self.isMC:
            self.jerInputArchivePath = os.environ['CMSSW_BASE'] + \
                "/src/PhysicsTools/NanoAODTools/data/jme/"
            self.jerTag = jerInputFileName[:jerInputFileName.find('_MC_') + len('_MC')]
            self.jerArchive = tarfile.open(
                self.jerInputArchivePath + self.jerTag + ".tgz", "r:gz")
            self.jerInputFilePath = tempfile.mkdtemp()
            self.jerArchive.extractall(self.jerInputFilePath)
            self.jerInputFileName = jerInputFileName
            self.jerUncertaintyInputFileName = jerUncertaintyInputFileName

            if not os.getenv("_jetSmearer"):
                os.environ["_jetSmearer"] = "jetSmearer"
                if "/libCorrectionsJME.so" not in ROOT.gSystem.GetLibraries():
                    ROOT.gSystem.Load("libCorrectionsJME.so")
                base = "{}/{}/src/Corrections/JME".format(
                    os.getenv("CMT_CMSSW_BASE"), os.getenv("CMT_CMSSW_VERSION"))
                ROOT.gROOT.ProcessLine(".L {}/interface/jetSmearer.h".format(base))
                ROOT.gInterpreter.Declare(
                    'auto jet_smearer = jetSmearer("%s", "%s", "%s");' % (
                        self.jerInputFilePath, self.jerInputFileName, self.jerUncertaintyInputFileName)
                )

    def run(self, df):
        if not self.isMC:
            return df, []
        all_branches = df.GetColumnNames()
        fixedGridRhoFastjetAll = "fixedGridRhoFastjetAll"
        if fixedGridRhoFastjetAll not in all_branches:
            fixedGridRhoFastjetAll = "Rho_fixedGridRhoFastjetAll"

        branches = ["jet_smear_factor", "jet_smear_factor_down", "jet_smear_factor_up",
            "jet_pt_resolution"]
        df = df.Define("__jer_results", "jet_smearer.get_smear_vals("
            "run, luminosityBlock, event, Jet_pt, Jet_eta, Jet_phi, Jet_mass, "
            "GenJet_pt, GenJet_eta, GenJet_phi, GenJet_mass, %s)" % fixedGridRhoFastjetAll)
        for ib, branch in enumerate(branches):
            df = df.Define(branch, "__jer_results[%s]" % ib)
        return df, branches


[docs]def jetSmearerRDF(**kwargs): """ Module to compute jet smearing factors (nom, up, down) and jet_pt_resolution Note: UL2016_preVFP is not implemented YAML sintaxis: .. code-block:: yaml codename: name: jetSmearerRDF path: Corrections.JME.smearing parameters: year: self.config.year isMC: self.dataset.process.isMC jerTag: self.config.year isUL: self.dataset.has_tag('ul') ispreVFP: self.config.get_aux("isPreVFP", False) """ isMC = kwargs.pop("isMC") isUL = kwargs.pop("isUL") ispreVFP = kwargs.pop("ispreVFP", False) year = kwargs.pop("year") jetType = kwargs.pop("jetType", "AK4PFchs") jerTag = kwargs.pop("jerTag", "") jerTagsMC = { '2016': 'Summer16_25nsV1_MC', '2017': 'Fall17_V3_MC', '2018': 'Autumn18_V7b_MC', 'UL2016_preVFP': 'Summer20UL16APV_JRV3_MC', 'UL2016': 'Summer20UL16_JRV3_MC', 'UL2017': 'Summer19UL17_JRV2_MC', 'UL2018': 'Summer19UL18_JRV2_MC', } jerInputFileName = "" jerUncertaintyInputFileName = "" if jerTag != "": jerTag = jerTagsMC["%s%s" % (("UL" if isUL else ""), year)] jerInputFileName = jerTag + "_PtResolution_" + jetType + ".txt" jerUncertaintyInputFileName = jerTag + "_SF_" + jetType + ".txt" if ispreVFP: jerTag = jerTagsMC["UL2016_preVFP"] else: raise ValueError("Empty jerTag, deprecated") if year == 2016: jerInputFileName = "Summer16_25nsV1_MC_PtResolution_" + jetType + ".txt" jerUncertaintyInputFileName = "Summer16_25nsV1_MC_SF_" + jetType + ".txt" elif year == 2017 or year == 2018: # use 2017 JER for 2018 for the time being jerInputFileName = "Fall17_V3_MC_PtResolution_" + jetType + ".txt" jerUncertaintyInputFileName = "Fall17_V3_MC_SF_" + jetType + ".txt" elif year == 2018 and False: # jetSmearer not working with 2018 JERs yet jerInputFileName = "Autumn18_V7_MC_PtResolution_" + jetType + ".txt" jerUncertaintyInputFileName = "Autumn18_V7_MC_SF_" + jetType + ".txt" return lambda: jetSmearerRDFProducer( isMC=isMC,ispreVFP=ispreVFP, jerInputFileName=jerInputFileName, jerUncertaintyInputFileName=jerUncertaintyInputFileName, **kwargs)
class jetVarRDFProducer(): def __init__(self, *args, **kwargs): self.isMC = kwargs.pop("isMC") def run(self, df): if not self.isMC: return df, [] df = df.Define("Jet_pt_nom", "Jet_pt * jet_smear_factor") df = df.Define("Jet_mass_nom", "Jet_mass * jet_smear_factor") df = df.Define("Jet_pt_smeared_up", "Jet_pt * jet_smear_factor_up") df = df.Define("Jet_mass_smeared_up", "Jet_mass * jet_smear_factor_up") df = df.Define("Jet_pt_smeared_down", "Jet_pt * jet_smear_factor_down") df = df.Define("Jet_mass_smeared_down", "Jet_mass * jet_smear_factor_down") return df, ["Jet_pt_nom", "Jet_mass_nom", "Jet_pt_smeared_up", "Jet_mass_smeared_up", "Jet_pt_smeared_down", "Jet_mass_smeared_down"]
[docs]def jetVarRDF(**kwargs): """ Module to compute jet pt and mass after applying smearing factors YAML sintaxis: .. code-block:: yaml codename: name: jetVarRDF path: Base.Modules.smearing parameters: isMC: self.dataset.process.isMC """ return lambda: jetVarRDFProducer(**kwargs)
class metSmearerRDFProducer(JetLepMetSyst): def __init__(self, isMC): self.isMC = isMC if self.isMC: if "/libCorrectionsJME.so" not in ROOT.gSystem.GetLibraries(): ROOT.gSystem.Load("libCorrectionsJME.so") base = "{}/{}/src/Corrections/JME".format( os.getenv("CMT_CMSSW_BASE"), os.getenv("CMT_CMSSW_VERSION")) if not os.getenv("_METShift"): os.environ["_METShift"] = "METShift" ROOT.gROOT.ProcessLine(".L {}/interface/metShift.h".format(base)) ROOT.gInterpreter.Declare('auto met_shifter = metShift();') def run(self, df): if not self.isMC: return df, [] branches = ["MET_smeared_pt", "MET_smeared_phi"] branches_up = ["MET_smeared_pt_up", "MET_smeared_phi_up"] branches_down = ["MET_smeared_pt_down", "MET_smeared_phi_down"] df = df.Define("smeared_met", "met_shifter.get_shifted_met(" "Jet_eta, Jet_phi, Jet_pt, Jet_mass, Jet_pt_nom, Jet_mass_nom, MET_pt, MET_phi)") df = df.Define("smeared_met_up", "met_shifter.get_shifted_met(" "Jet_eta, Jet_phi, Jet_pt, Jet_mass, Jet_pt_smeared_up, Jet_mass_smeared_up, MET_pt, MET_phi)") df = df.Define("smeared_met_down", "met_shifter.get_shifted_met(" "Jet_eta, Jet_phi, Jet_pt, Jet_mass, Jet_pt_smeared_down, Jet_mass_smeared_down, MET_pt, MET_phi)") for ib, branch in enumerate(branches): df = df.Define(branch, "smeared_met[%s]" % ib) for ib, branch in enumerate(branches_up): df = df.Define(branch, "smeared_met_up[%s]" % ib) for ib, branch in enumerate(branches_down): df = df.Define(branch, "smeared_met_down[%s]" % ib) return df, branches + branches_up + branches_down
[docs]def metSmearerRDF(**kwargs): """ Module to compute MET pt and phi after applying smearing to all jets and up and down variations YAML sintaxis: .. code-block:: yaml codename: name: metSmearerRDF path: Base.Modules.smearing parameters: isMC: self.dataset.process.isMC """ isMC = kwargs.pop("isMC") return lambda: metSmearerRDFProducer(isMC, **kwargs)
class metJetTauSmearerRDFProducer(metSmearerRDFProducer): def run(self, df): if not self.isMC: return df, [] branches = ["MET_smeared_pt", "MET_smeared_phi"] branches_up = ["MET_smeared_pt_up", "MET_smeared_phi_up"] branches_down = ["MET_smeared_pt_down", "MET_smeared_phi_down"] branches_tes_up = ["MET_smeared_pt_corr_up", "MET_smeared_phi_corr_up"] branches_tes_down = ["MET_smeared_pt_corr_down", "MET_smeared_phi_corr_down"] df = df.Define("smeared_met", "met_shifter.get_shifted_met(" "Jet_eta, Jet_phi, Jet_pt, Jet_mass, Jet_pt_nom, Jet_mass_nom, MET_pt, MET_phi)") df = df.Define("tes_smeared_met", "met_shifter.get_shifted_met(" "Tau_eta, Tau_phi, Tau_pt, Tau_mass, Tau_pt_corr, Tau_mass_corr, smeared_met[0], smeared_met[1])") df = df.Define("MET_smeared_pt", "tes_smeared_met[0]") df = df.Define("MET_smeared_phi", "tes_smeared_met[1]") df = df.Define("smeared_met_up", "met_shifter.get_shifted_met(" "Jet_eta, Jet_phi, Jet_pt_nom, Jet_mass_nom, Jet_pt_smeared_up, Jet_mass_smeared_up, MET_smeared_pt, MET_smeared_phi)") df = df.Define("smeared_met_down", "met_shifter.get_shifted_met(" "Jet_eta, Jet_phi, Jet_pt_nom, Jet_mass_nom, Jet_pt_smeared_down, Jet_mass_smeared_down, MET_smeared_pt, MET_smeared_phi)") df = df.Define("smeared_met_corr_up", "met_shifter.get_shifted_met(" "Tau_eta, Tau_phi, Tau_pt_corr, Tau_mass_corr, Tau_pt_corr_up, Tau_mass_corr_up, MET_smeared_pt, MET_smeared_phi)") df = df.Define("smeared_met_corr_down", "met_shifter.get_shifted_met(" "Tau_eta, Tau_phi, Tau_pt_corr, Tau_mass_corr, Tau_pt_corr_down, Tau_mass_corr_down, MET_smeared_pt, MET_smeared_phi)") df = df.Define(branches_up[0], "smeared_met_up[0]") df = df.Define(branches_up[1], "smeared_met_up[1]") df = df.Define(branches_down[0], "smeared_met_down[0]") df = df.Define(branches_down[1], "smeared_met_down[1]") df = df.Define(branches_tes_up[0], "smeared_met_corr_up[0]") df = df.Define(branches_tes_up[1], "smeared_met_corr_up[1]") df = df.Define(branches_tes_down[0], "smeared_met_corr_down[0]") df = df.Define(branches_tes_down[1], "smeared_met_corr_down[1]") return df, branches + branches_up + branches_down + branches_tes_up + branches_tes_down def metJetTauSmearerRDF(**kwargs): """ Module to compute MET pt and phi after applying smearing to all jets and tau energy scale and up and down variations YAML sintaxis: .. code-block:: yaml codename: name: metJetTauSmearerRDF path: Base.Modules.smearing parameters: isMC: self.dataset.process.isMC """ isMC = kwargs.pop("isMC") return lambda: metJetTauSmearerRDFProducer(isMC, **kwargs)