Source code for Corrections.EGM.eleCorrections

import os

from analysis_tools.utils import import_root
ROOT = import_root()

import correctionlib
correctionlib.register_pyroot_binding()

json_path = "/cvmfs/cms.cern.ch/rsync/cms-nanoAOD/jsonpog-integration/POG/EGM/{}/electron.json.gz"


class eleSFRDFProducer():
    def __init__(self, *args, **kwargs):
        self.year = kwargs.pop("year")
        self.isUL = kwargs.pop("isUL")
        self.isMC = kwargs.pop("isMC")
        self.wps = kwargs.pop("wps", ["wp80iso"])
        
        if self.isMC:
            if not self.isUL:
                raise ValueError("Only implemented for UL datasets")
            if self.year == 2018:
                self.tag = "2018"
                filename = json_path.format("%s_UL" % self.tag)
            else:
                raise ValueError("2016 and 2017 not yet implemented")

            if "/libCorrectionsWrapper.so" not in ROOT.gSystem.GetLibraries():
                ROOT.gInterpreter.Load("libCorrectionsWrapper.so")
            ROOT.gInterpreter.Declare(os.path.expandvars(
                '#include "$CMSSW_BASE/src/Corrections/Wrapper/interface/custom_sf.h"'))
            ROOT.gInterpreter.ProcessLine(
                'auto corr = MyCorrections("%s", "UL-Electron-ID-SF");' % filename)
                
            ROOT.gInterpreter.Declare("""
                using Vfloat = const ROOT::RVec<float>&;
                using Vint = const ROOT::RVec<int>&;
                ROOT::RVec<double> get_ele_sf(
                        std::string syst, std::string wp, Vfloat eta, Vfloat pt) {
                    ROOT::RVec<double> sf;
                    for (size_t i = 0; i < pt.size(); i++) {
                        if (pt[i] < 10.) sf.push_back(1.);
                        else sf.push_back(corr.eval({"%s", syst, wp, eta[i], pt[i]}));
                    }
                    return sf;
                }
            """ % self.tag)

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

        branches = []
        for syst_name, syst in [("", "sf"), ("_up", "sfup"), ("_down", "sfdown")]:
            for wp in self.wps:
                df = df.Define("elesf_%s%s" % (wp, syst_name),
                    'get_ele_sf("%s", "%s", Electron_eta, Electron_pt)' % (syst, wp))
                branches.append("elesf_%s%s" % (wp, syst_name))

        return df, branches


def eleSFRDF(**kwargs):
    """
    Module to obtain electron SFs with their uncertainties.

    :param wps: name of the wps to consider among ``Loose``, ``Medium``,
        ``RecoAbove20``, ``RecoBelow20``, ``Tight``, ``Veto``, ``wp80iso`` (default),
        ``wp80noiso``, ``wp90iso,`` ``wp90noiso``.
    :type wps: list of str

    YAML sintaxis:

    .. code-block:: yaml

        codename:
            name: eleSFRDF
            path: Corrections.EGM.eleCorrections
            parameters:
                isMC: self.dataset.process.isMC
                year: self.config.year
                isUL: self.dataset.has_tag('ul')
                wps: [wp80iso, ...]
    """

    return lambda: eleSFRDFProducer(**kwargs)