Source code for Corrections.MUO.muCorrections

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/MUO/{}/muon_Z.json.gz"


class muSFRDFProducer():
    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", ["tight_id", "tight_reliso"])

        if self.isMC:
            if not self.isUL:
                raise ValueError("Only implemented for UL datasets")
            if self.year == 2018:
                self.tag = "2018_UL"
                filename = json_path.format(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"'))

            # Declaring two objects, one for tightID and the other for tightRelIso.
            # Of course this should be changed if more WPs are needed, maybe looping over 
            # an input parameter.
            ROOT.gInterpreter.ProcessLine(
                'auto corr_tightid = '
                    'MyCorrections("%s", "NUM_TightID_DEN_TrackerMuons");'
                % filename)
            ROOT.gInterpreter.ProcessLine(
                'auto corr_tightreliso = '
                    'MyCorrections("%s", "NUM_TightRelIso_DEN_TightIDandIPCut");'
                % filename)

            ROOT.gInterpreter.Declare("""
                using Vfloat = const ROOT::RVec<float>&;
                using Vint = const ROOT::RVec<int>&;
                ROOT::RVec<double> get_mu_tight_id_sf(std::string syst, Vfloat eta, Vfloat pt) {
                    ROOT::RVec<double> sf;
                    for (size_t i = 0; i < pt.size(); i++) {
                        if (pt[i] < 15. || fabs(eta[i]) > 2.4) sf.push_back(1.);
                        else sf.push_back(corr_tightid.eval({"%s", fabs(eta[i]), pt[i], syst}));
                    }
                    return sf;
                }
                ROOT::RVec<double> get_mu_tight_reliso_sf(std::string syst, Vfloat eta, Vfloat pt) {
                    ROOT::RVec<double> sf;
                    for (size_t i = 0; i < pt.size(); i++) {
                        if (pt[i] < 15. || fabs(eta[i]) > 2.4) sf.push_back(1.);
                        else sf.push_back(corr_tightreliso.eval({"%s", fabs(eta[i]), pt[i], syst}));
                    }
                    return sf;
                }
            """  % (self.tag, self.tag))

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

        branches = []
        for corr in self.wps:
            for syst_name, syst in [("", "sf"), ("_up", "systup"), ("_down", "systdown")]:
                df = df.Define("musf_%s%s" % (corr, syst_name),
                    'get_mu_%s_sf("%s", Muon_eta, Muon_pt)' % (corr, syst))
                branches.append("musf_%s%s" % (corr, syst_name))

        return df, branches


def muSFRDF(**kwargs):
    """
    Module to obtain muon SFs with their uncertainties.

    :param wps: name of the wps to consider among ``tight_id``, ``tight_reliso``.
        Note: probably more can be considered.
    :type wps: list of str

    YAML sintaxis:

    .. code-block:: yaml

        codename:
            name: muSFRDF
            path: Corrections.MUO.muCorrections
            parameters:
                isMC: self.dataset.process.isMC
                year: self.config.year
                isUL: self.dataset.has_tag('ul')
                wps: [tight_id, tight_reliso]
    """
    return lambda: muSFRDFProducer(**kwargs)