Source code for Corrections.JME.jec

import os
import tarfile
import tempfile
import itertools

from Base.Modules.baseModules import JetLepMetSyst

from analysis_tools.utils import import_root, randomize

ROOT = import_root()

class jecProviderRDFProducer(JetLepMetSyst):
    def __init__(self, isMC,
            jerInputFileName="Spring16_25nsV10_MC_PtResolution_AK4PFchs.txt",
            jec_sources=[], *args, **kwargs):

        self.isMC = isMC
        self.pt = kwargs.pop("pt", "")
        self.eta = kwargs.pop("eta", "")
        self.mass = kwargs.pop("mass", "")

        if self.isMC:
            if not os.getenv("JEC_PATH"):
                self.jerInputArchivePath = os.environ['CMSSW_BASE'] + \
                    "/src/Corrections/JME/data/"
                self.jerTag = jerInputFileName[:jerInputFileName.find('_MC_') + len('_MC')]
                self.jerArchive = tarfile.open(
                    self.jerInputArchivePath + self.jerTag + ".tar.gz", "r:gz")
                self.jerInputFilePath = tempfile.mkdtemp()
                self.jerArchive.extractall(self.jerInputFilePath)
                self.jerInputFileName = "RegroupedV2_%s_UncertaintySources_AK4PFchs.txt"\
                    % self.jerTag
                assert(os.path.isfile(os.path.join(self.jerInputFilePath, self.jerInputFileName)))
                os.environ["JEC_PATH"] = os.path.join(self.jerInputFilePath, self.jerInputFileName)

            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"))

            try:
                assert ROOT.jecProvider
            except AttributeError:
                ROOT.gROOT.ProcessLine(".L {}/interface/jecProvider.h".format(base))

            self.jec_sources = jec_sources
            jec_sources_str = ", ".join(['"%s"' % jec_source for jec_source in self.jec_sources])
            self.jec_provider = randomize("jec_provider")

            ROOT.gInterpreter.Declare(
                'auto %s = jecProvider("%s", {%s});' % (
                    self.jec_provider, os.getenv("JEC_PATH"), jec_sources_str
                )
            )
        super(jecProviderRDFProducer, self).__init__(self, isMC=isMC, *args, **kwargs)

    def run(self, df):
        if not self.isMC:
            return df, []

        jecs = randomize("jec_uncertainties")

        if not self.pt:
            branches = ["%s_%s" % (name, d) for (d, name) in itertools.product(
                ["up", "down"], ["jec_factor_%s" % jec_source for jec_source in self.jec_sources])]

            df = df.Define(jecs, "%s.get_jec_uncertainties("
                "nJet, Jet_pt%s, Jet_eta)" % (self.jec_provider, self.jet_syst)).Define(
                "%s_up" % jecs, "%s[0]" % jecs).Define("%s_down" % jecs, "%s[1]" % jecs)
            for ib, branch in enumerate(branches):
                if ib < len(self.jec_sources):
                    df = df.Define(branch, "%s_up[%s]" % (jecs, ib))
                else:
                    df = df.Define(branch, "%s_down[%s]" % (jecs, ib - len(self.jec_sources)))
        else:
            assert self.pt and self.eta
            prefix = self.pt[:self.pt.find("pt")]
            df = df.Define(jecs, "%s.get_single_jec_uncertainties("
                "%s, %s)" % (self.jec_provider, self.pt, self.eta)).Define(
                "%s_up" % jecs, "%s[0]" % jecs).Define(
                "%s_down" % jecs, "%s[1]" % jecs)

            branches = []
            for d in ["up", "down"]:
                for isource, jec_source in enumerate(self.jec_sources):
                    df = df.Define("%spt_%s_%s" % (prefix, jec_source, d),
                        "%s * %s_%s[%s]" % (self.pt, jecs, d, isource))
                    branches.append("%spt_%s_%s" % (prefix, jec_source, d))
                    if self.mass:
                        df = df.Define("%smass_%s_%s" % (prefix, jec_source, d),
                            "%s * %s_%s[%s]" % (self.mass, jecs, d, isource))
                        branches.append("%smass_%s_%s" % (prefix, jec_source, d))            

        return df, branches


[docs]def jecProviderRDF(**kwargs): """ Module to compute jec uncertainty factors. If parameters `pt` and `eta` are included, the module will output the pt (and mass if it's included as a parameter) of the object after applying the uncertainties. Note: UL2016_preVFP is not implemented :param jec_sources: names of the systematic sources to consider. They depend on the year: - 2016: ``FlavorQCD``, ``RelativeBal``, ``HF``, ``BBEC1``, ``EC2``, ``Absolute``,\ ``BBEC1_2016``, ``EC2_2016``, ``Absolute_2016``, ``HF_2016``, ``RelativeSample_2016``,\ ``Total`` - 2017: ``FlavorQCD``, ``RelativeBal``, ``HF``, ``BBEC1``, ``EC2``, ``Absolute``,\ ``BBEC1_2017``, ``EC2_2017``, ``Absolute_2017``, ``HF_2017``, ``RelativeSample_2017``,\ ``Total`` - 2018: ``FlavorQCD``, ``RelativeBal``, ``HF``, ``BBEC1``, ``EC2``, ``Absolute``,\ ``BBEC1_2018``, ``EC2_2018``, ``Absolute_2018``, ``HF_2018``, ``RelativeSample_2018``,\ ``Total`` Note: probably more are available, to be checked. :type jec_sources: list of str :param pt: pt of the object to shift :type pt: str :param eta: pt of the object to shift :type eta: str :param mass: mass of the object to shift :type mass: str YAML sintaxis: .. code-block:: yaml codename: name: jecProviderRDF path: Base.Modules.jec 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) jec_sources: [] (pt: bjet1_pt_nom) (eta: bjet1_eta) (mass: bjet1_mass_nom) """ isMC = kwargs.pop("isMC") isUL = kwargs.pop("isUL") year = kwargs.pop("year") isPreVFP = kwargs.pop("ispreVFP", False) jetType = kwargs.pop("jetType", "AK4PFchs") # jerTag = kwargs.pop("jerTag", "") jec_sources = kwargs.pop("jec_sources", []) jerTagsMC = { '2016': 'Summer16_07Aug2017_V11_MC', '2017': 'Fall17_17Nov2017_V32_MC', '2018': 'Autumn18_V19_MC', 'UL2016_preVFP': 'Summer19UL16APV_V7_MC', # julia 'UL2016': 'Summer19UL16_V7_MC', # julia 'UL2017': 'Summer19UL17_V5_MC', 'UL2018': 'Summer19UL18_V5_MC', } campaign_tag = "%s%s" % (("UL" if isUL else ""), year) if isPreVFP: campaign_tag += "_preVFP" jerTag = jerTagsMC[campaign_tag] jerInputFileName = jerTag + "_PtResolution_" + jetType + ".txt" jec_sources_per_year = { 2016: ["FlavorQCD", "RelativeBal", "HF", "BBEC1", "EC2", "Absolute", "BBEC1_2016", "EC2_2016", "Absolute_2016", "HF_2016", "RelativeSample_2016", "Total"], 2017: ["FlavorQCD", "RelativeBal", "HF", "BBEC1", "EC2", "Absolute", "BBEC1_2017", "EC2_2017", "Absolute_2017", "HF_2017", "RelativeSample_2017", "Total"], 2018: ["FlavorQCD", "RelativeBal", "HF", "BBEC1", "EC2", "Absolute", "BBEC1_2018", "EC2_2018", "Absolute_2018", "HF_2018", "RelativeSample_2018", "Total"] } if len(jec_sources) == 0: jec_sources = jec_sources_per_year[int(year)] else: for jec_source in jec_sources: assert jec_source in jec_sources_per_year[int(year)] return lambda: jecProviderRDFProducer(isMC, jerInputFileName, jec_sources, **kwargs)
class jecVarRDFProducer(): def __init__(self, isMC, jec_sources=[]): self.isMC = isMC self.jec_sources = jec_sources def run(self, df): if not self.isMC: return df, [] branches = [] for d in ["up", "down"]: for isource, jec_source in enumerate(self.jec_sources): branches.append("Jet_pt_%s_%s" % (jec_source, d)) branches.append("Jet_mass_%s_%s" % (jec_source, d)) df = df.Define("Jet_pt_%s_%s" % (jec_source, d), "Jet_pt_nom * jec_factor_%s_%s" %(jec_source, d)) df = df.Define("Jet_mass_%s_%s" % (jec_source, d), "Jet_mass_nom * jec_factor_%s_%s" %(jec_source, d)) return df, branches
[docs]def jecVarRDF(**kwargs): """ Module to compute jet pt and mass after applying jec uncertainty factors Required RDFModules: :ref:`JME_smearing_jetSmearerRDF`. Note: UL2016_preVFP is not implemented :param jec_sources: names of the systematic sources to consider. They depend on the year: - 2016: ``FlavorQCD``, ``RelativeBal``, ``HF``, ``BBEC1``, ``EC2``, ``Absolute``,\ ``BBEC1_2016``, ``EC2_2016``, ``Absolute_2016``, ``HF_2016``, ``RelativeSample_2016``,\ ``Total`` - 2017: ``FlavorQCD``, ``RelativeBal``, ``HF``, ``BBEC1``, ``EC2``, ``Absolute``,\ ``BBEC1_2017``, ``EC2_2017``, ``Absolute_2017``, ``HF_2017``, ``RelativeSample_2017``,\ ``Total`` - 2018: ``FlavorQCD``, ``RelativeBal``, ``HF``, ``BBEC1``, ``EC2``, ``Absolute``,\ ``BBEC1_2018``, ``EC2_2018``, ``Absolute_2018``, ``HF_2018``, ``RelativeSample_2018``,\ ``Total`` Note: probably more are available, to be checked. :type jec_sources: list of str YAML sintaxis: .. code-block:: yaml codename: name: jecVarRDF path: Base.Modules.jec parameters: year: self.config.year isMC: self.dataset.process.isMC jec_sources: [] """ isMC = kwargs.pop("isMC") year = kwargs.pop("year") jec_sources = kwargs.pop("jec_sources", []) jec_sources_per_year = { 2016: ["FlavorQCD", "RelativeBal", "HF", "BBEC1", "EC2", "Absolute", "BBEC1_2016", "EC2_2016", "Absolute_2016", "HF_2016", "RelativeSample_2016", "Total"], 2017: ["FlavorQCD", "RelativeBal", "HF", "BBEC1", "EC2", "Absolute", "BBEC1_2017", "EC2_2017", "Absolute_2017", "HF_2017", "RelativeSample_2017", "Total"], 2018: ["FlavorQCD", "RelativeBal", "HF", "BBEC1", "EC2", "Absolute", "BBEC1_2018", "EC2_2018", "Absolute_2018", "HF_2018", "RelativeSample_2018", "Total"] } if len(jec_sources) == 0: jec_sources = jec_sources_per_year[int(year)] else: for jec_source in jec_sources: assert jec_source in jec_sources_per_year[int(year)] return lambda: jecVarRDFProducer(isMC, jec_sources, **kwargs)
class jecMETRDFProducer(): def __init__(self, isMC, jec_sources=[]): self.isMC = isMC self.jec_sources = jec_sources 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_jec = metShift();') def run(self, df): if not self.isMC: return df, [] branches = [] for d in ["up", "down"]: for isource, jec_source in enumerate(self.jec_sources): branches.append("MET_smeared_pt_%s_%s" % (jec_source, d)) branches.append("MET_smeared_phi_%s_%s" % (jec_source, d)) df = df.Define("smeared_met_%s_%s" % (jec_source, d), "met_shifter_jec.get_shifted_met(" "Jet_eta, Jet_phi, Jet_pt_nom, Jet_mass_nom, Jet_pt_{0}_{1}, Jet_mass_{0}_{1}, MET_smeared_pt, MET_smeared_phi)" .format(jec_source, d)) df = df.Define("MET_smeared_pt_%s_%s" % (jec_source, d), "smeared_met_%s_%s[0]" % (jec_source, d)) df = df.Define("MET_smeared_phi_%s_%s" % (jec_source, d), "smeared_met_%s_%s[1]" % (jec_source, d)) return df, branches
[docs]def jecMETRDF(**kwargs): """ Module to compute MET pt and phi after applying jec uncertainty factors Required RDFModules: :ref:`JME_smearing_jetSmearerRDF`. Note: UL2016_preVFP is not implemented :param jec_sources: names of the systematic sources to consider. They depend on the year: - 2016: ``FlavorQCD``, ``RelativeBal``, ``HF``, ``BBEC1``, ``EC2``, ``Absolute``,\ ``BBEC1_2016``, ``EC2_2016``, ``Absolute_2016``, ``HF_2016``, ``RelativeSample_2016``,\ ``Total`` - 2017: ``FlavorQCD``, ``RelativeBal``, ``HF``, ``BBEC1``, ``EC2``, ``Absolute``,\ ``BBEC1_2017``, ``EC2_2017``, ``Absolute_2017``, ``HF_2017``, ``RelativeSample_2017``,\ ``Total`` - 2018: ``FlavorQCD``, ``RelativeBal``, ``HF``, ``BBEC1``, ``EC2``, ``Absolute``,\ ``BBEC1_2018``, ``EC2_2018``, ``Absolute_2018``, ``HF_2018``, ``RelativeSample_2018``,\ ``Total`` Note: probably more are available, to be checked. :type jec_sources: list of str YAML sintaxis: .. code-block:: yaml codename: name: jecMETRDF path: Base.Modules.jec parameters: year: self.config.year isMC: self.dataset.process.isMC jec_sources: [] """ isMC = kwargs.pop("isMC") year = kwargs.pop("year") jec_sources = kwargs.pop("jec_sources", []) jec_sources_per_year = { 2016: ["FlavorQCD", "RelativeBal", "HF", "BBEC1", "EC2", "Absolute", "BBEC1_2016", "EC2_2016", "Absolute_2016", "HF_2016", "RelativeSample_2016", "Total"], 2017: ["FlavorQCD", "RelativeBal", "HF", "BBEC1", "EC2", "Absolute", "BBEC1_2017", "EC2_2017", "Absolute_2017", "HF_2017", "RelativeSample_2017", "Total"], 2018: ["FlavorQCD", "RelativeBal", "HF", "BBEC1", "EC2", "Absolute", "BBEC1_2018", "EC2_2018", "Absolute_2018", "HF_2018", "RelativeSample_2018", "Total"] } if len(jec_sources) == 0: jec_sources = jec_sources_per_year[int(year)] else: for jec_source in jec_sources: assert jec_source in jec_sources_per_year[int(year)] return lambda: jecMETRDFProducer(isMC, jec_sources, **kwargs)