Source code for Corrections.MUO.muonSF_GEMethod

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 MuonSFproducer():
    def __init__(self, *args, **kwargs):
        self.year =   kwargs.pop("year")
        self.isUL =   kwargs.pop("isUL")
        self.isMC =   kwargs.pop("isMC")
        self.preVFP = kwargs.pop("preVFP")

        if self.isMC:
            if not self.isUL:
                raise ValueError("Only implemented for UL datasets")
            if self.year == 2018:
                self.tag = "2018_UL"
            elif self.year == 2017:
                self.tag = "2017_UL"
            elif self.year == 2016:
                if self.preVFP:
                    self.tag = "2016preVFP_UL"
                else:
                    self.tag = "2016postVFP_UL"
            else:
                raise ValueError("Only implemented for RunII (2016, 2017 & 2018)")

            filename = json_path.format(self.tag)

            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 four objects: for HighPtId, for LooseRelTkIso, for the Trigger and for Global/Tracker Reco
            # Of course this should be changed if more WPs are needed, maybe looping over 
            # an input parameter.
            ROOT.gInterpreter.ProcessLine(
                'auto ID_SF = '
                    'MyCorrections("%s", "NUM_HighPtID_DEN_TrackerMuons");'
                % filename)
            ROOT.gInterpreter.ProcessLine(
                'auto iso_SF = '
                    'MyCorrections("%s", "NUM_LooseRelTkIso_DEN_HighPtIDandIPCut");'
                % filename)
            ROOT.gInterpreter.ProcessLine(
                'auto reco_SF = '
                    'MyCorrections("%s", "NUM_GlobalMuons_DEN_genTracks");'
                % filename)
            if self.year == 2016:
                ROOT.gInterpreter.ProcessLine(
                    'auto trig_SF = '
                        'MyCorrections("%s", "NUM_Mu50_or_TkMu50_DEN_CutBasedIdGlobalHighPt_and_TkIsoLoose");'
                    % filename)
            else:
                ROOT.gInterpreter.ProcessLine(
                    'auto trig_SF = '
                        'MyCorrections("%s", "NUM_Mu50_or_OldMu100_or_TkMu100_DEN_CutBasedIdGlobalHighPt_and_TkIsoLoose");'
                    % filename)

            ROOT.gInterpreter.Declare("""
                using Vfloat = const ROOT::RVec<float>&;
                using Vint = const ROOT::RVec<int>&;
                double get_mu_idSF(std::string syst, float eta, float TuneP_relPt, float pt) {
                    double sf;
                    float TuneP_pt = TuneP_relPt*pt;
                    sf = ID_SF.eval({"%s", fabs(eta), TuneP_pt, syst});
                    return sf;
                }
                double get_mu_isoSF(std::string syst, float eta, float TuneP_relPt, float pt) {
                    double sf;
                    float TuneP_pt = TuneP_relPt*pt;
                    sf = iso_SF.eval({"%s", fabs(eta), TuneP_pt, syst});
                    return sf;
                }
                double get_mu_trigSF(std::string syst, float eta, float TuneP_relPt, float pt) {
                    double sf;
                    float TuneP_pt = TuneP_relPt*pt;
                    if(TuneP_pt < 52.0) sf = 1.0;
                    else sf = trig_SF.eval({"%s", fabs(eta), TuneP_pt, syst});
                    return sf;
                }
                double get_mu_recoSF(std::string syst, float eta, float TuneP_relPt, float pt) {
                    double sf;
                    float TuneP_pt = TuneP_relPt*pt;
                    sf = reco_SF.eval({"%s", fabs(eta), TuneP_pt, syst});
                    return sf;
                }

                double get_mu_trigSF_OLDforLegacy2018(float eta, float TuneP_relPt, float pt) {
                    double sf;
                    float TuneP_pt = TuneP_relPt*pt;
                    double SF_trigger[4];
                    if (TuneP_pt < 55){
                       SF_trigger[0] = 0.980; SF_trigger[1] = 0.982; SF_trigger[2] = 0.983; SF_trigger[3] = 0.987;
                    }else if(TuneP_pt < 60){
                       SF_trigger[0] = 0.982; SF_trigger[1] = 0.966; SF_trigger[2] = 1.012; SF_trigger[3] = 0.952;
                    }else if(TuneP_pt < 120){
                       SF_trigger[0] = 0.976; SF_trigger[1] = 0.967; SF_trigger[2] = 1.002; SF_trigger[3] = 0.999;
                    }else if(TuneP_pt < 200){
                       SF_trigger[0] = 0.980; SF_trigger[1] = 0.964; SF_trigger[2] = 1.004; SF_trigger[3] = 1.003;
                    }else if(TuneP_pt < 300){
                       SF_trigger[0] = 0.978; SF_trigger[1] = 0.960; SF_trigger[2] = 1.009; SF_trigger[3] = 1.024;
                    }else if(TuneP_pt < 500){
                       SF_trigger[0] = 0.973; SF_trigger[1] = 0.985; SF_trigger[2] = 0.993; SF_trigger[3] = 0.993;
                    }else {
                       SF_trigger[0] = 0.957; SF_trigger[1] = 0.988; SF_trigger[2] = 1.062; SF_trigger[3] = 1.070;
                    }
                    double etaEdge[4];
                    etaEdge[0] = 0.9; etaEdge[1] = 1.2; etaEdge[2] = 2.1; etaEdge[3] = 2.4;
                    unsigned int ibin;
                    for (ibin = 0; ibin < 4; ibin++) {
                       if (fabs(eta) < etaEdge[ibin]) break;
                    }
                    return sf = SF_trigger[ibin];
                }
                double get_mu_recoSF_OLDforLegacy2018(float eta, float TuneP_relPt, float pt, float phi, float mass) {
                    double sf;
                    float TuneP_pt = TuneP_relPt*pt;
                    float p = TuneP_pt*cosh(eta);
                    double SF_reco[8];
                    if(fabs(eta) < 1.6){
                       SF_reco[0] = 0.9943; SF_reco[1] = 0.9948; SF_reco[2] = 0.9950; SF_reco[3] = 0.994;
                       SF_reco[4] = 0.9914; SF_reco[5] = 0.993;  SF_reco[6] = 0.991;  SF_reco[7] = 1.0;
                    }
                    else {
                       SF_reco[0] = 1.0;   SF_reco[1] = 0.993; SF_reco[2] = 0.990; SF_reco[3] = 0.988;
                       SF_reco[4] = 0.981; SF_reco[5] = 0.983; SF_reco[6] = 0.978; SF_reco[7] = 0.98;
                    }
                    double pEdge[8];
                    pEdge[0] = 100; pEdge[1] = 150; pEdge[2] = 200;  pEdge[3] = 300;
                    pEdge[4] = 400; pEdge[5] = 600; pEdge[6] = 1500; pEdge[7] = 3500;
                    unsigned int ibin;
                    for (ibin = 0; ibin < 8; ibin++) {
                       if (p < pEdge[ibin]) 
                          return sf = SF_reco[ibin];
                    }
                    return sf = SF_reco[7];
                }

                double get_mu_trigSF_OLDforLegacy2017(float eta, float TuneP_relPt, float pt) {
                    double sf;
                    float TuneP_pt = TuneP_relPt*pt;
                    double SF_trigger[4];
                    if (TuneP_pt < 55){
                       SF_trigger[0] = 0.968; SF_trigger[1] = 0.929; SF_trigger[2] = 0.980; SF_trigger[3] = 0.869;
                    }else if(TuneP_pt < 60){
                       SF_trigger[0] = 0.968; SF_trigger[1] = 0.953; SF_trigger[2] = 0.981; SF_trigger[3] = 0.864;
                    }else if(TuneP_pt < 120){
                       SF_trigger[0] = 0.969; SF_trigger[1] = 0.945; SF_trigger[2] = 0.985; SF_trigger[3] = 0.926;
                    }else if(TuneP_pt < 200){
                       SF_trigger[0] = 0.967; SF_trigger[1] = 0.946; SF_trigger[2] = 0.991; SF_trigger[3] = 0.992;
                    }else if(TuneP_pt < 300){
                       SF_trigger[0] = 0.964; SF_trigger[1] = 0.948; SF_trigger[2] = 0.986; SF_trigger[3] = 1.023;
                    }else if(TuneP_pt < 500){
                       SF_trigger[0] = 0.963; SF_trigger[1] = 0.939; SF_trigger[2] = 1.002; SF_trigger[3] = 0.905;
                    }else {
                       SF_trigger[0] = 0.951; SF_trigger[1] = 0.896; SF_trigger[2] = 0.922; SF_trigger[3] = 0.424;
                    }
                    double etaEdge[4];
                    etaEdge[0] = 0.9; etaEdge[1] = 1.2; etaEdge[2] = 2.1; etaEdge[3] = 2.4;
                    unsigned int ibin;
                    for (ibin = 0; ibin < 4; ibin++) {
                       if (fabs(eta) < etaEdge[ibin]) break;
                    }
                    return sf = SF_trigger[ibin];
                }
                double get_mu_recoSF_OLDforLegacy2017(float eta, float TuneP_relPt, float pt, float phi, float mass) {
                    double sf;
                    float TuneP_pt = TuneP_relPt*pt;
                    float p = TuneP_pt*cosh(eta);
                    double SF_reco[8];
                    if(fabs(eta) < 1.6){
                       SF_reco[0] = 0.9938; SF_reco[1] = 0.9950; SF_reco[2] = 0.996; SF_reco[3] = 0.996;
                       SF_reco[4] = 0.994;  SF_reco[5] = 1.003;  SF_reco[6] = 0.987; SF_reco[7] = 0.9;
                    }
                    else {
                       SF_reco[0] = 1.0;   SF_reco[1] = 0.993; SF_reco[2] = 0.989; SF_reco[3] = 0.986;
                       SF_reco[4] = 0.989; SF_reco[5] = 0.983; SF_reco[6] = 0.986; SF_reco[7] = 1.01;
                    }
                    double pEdge[8];
                    pEdge[0] = 100; pEdge[1] = 150; pEdge[2] = 200;  pEdge[3] = 300;
                    pEdge[4] = 400; pEdge[5] = 600; pEdge[6] = 1500; pEdge[7] = 3500;
                    unsigned int ibin;
                    for (ibin = 0; ibin < 8; ibin++) {
                       if (p < pEdge[ibin]) 
                          return sf = SF_reco[ibin];
                    }
                    return sf = SF_reco[7];
                }

                double get_mu_trigSF_OLDforLegacy2016(float eta, float TuneP_relPt, float pt) {
                    double sf;
                    float TuneP_pt = TuneP_relPt*pt;
                    double SF_trigger[4];
                    if (TuneP_pt < 55){
                       SF_trigger[0] = 0.980; SF_trigger[1] = 0.938; SF_trigger[2] = 0.996; SF_trigger[3] = 0.934;
                    }else if(TuneP_pt < 60){
                       SF_trigger[0] = 0.983; SF_trigger[1] = 0.958; SF_trigger[2] = 1.005; SF_trigger[3] = 0.954;
                    }else if(TuneP_pt < 120){
                       SF_trigger[0] = 0.980; SF_trigger[1] = 0.956; SF_trigger[2] = 0.990; SF_trigger[3] = 0.948;
                    }else if(TuneP_pt < 200){
                       SF_trigger[0] = 0.980; SF_trigger[1] = 0.948; SF_trigger[2] = 0.988; SF_trigger[3] = 0.935;
                    }else if(TuneP_pt < 300){
                       SF_trigger[0] = 0.980; SF_trigger[1] = 0.927; SF_trigger[2] = 0.977; SF_trigger[3] = 0.888;
                    }else if(TuneP_pt < 500){
                       SF_trigger[0] = 0.962; SF_trigger[1] = 0.911; SF_trigger[2] = 1.010; SF_trigger[3] = 0.953;
                    }else {
                       SF_trigger[0] = 0.983; SF_trigger[1] = 0.824; SF_trigger[2] = 1.020; SF_trigger[3] = 1.134;
                    }
                    double etaEdge[4];
                    etaEdge[0] = 0.9; etaEdge[1] = 1.2; etaEdge[2] = 2.1; etaEdge[3] = 2.4;
                    unsigned int ibin;
                    for (ibin = 0; ibin < 4; ibin++) {
                       if (fabs(eta) < etaEdge[ibin]) break;
                    }
                    return sf = SF_trigger[ibin];
                }
                double get_mu_recoSF_OLDforLegacy2016(float eta, float TuneP_relPt, float pt, float phi, float mass) {
                    double sf;
                    float TuneP_pt = TuneP_relPt*pt;
                    float p = TuneP_pt*cosh(eta);
                    double SF_reco[8];
                    if(fabs(eta) < 1.6){
                       SF_reco[0] = 0.9914; SF_reco[1] = 0.9936; SF_reco[2] = 0.993; SF_reco[3] = 0.993;
                       SF_reco[4] = 0.990;  SF_reco[5] = 0.990;  SF_reco[6] = 0.989; SF_reco[7] = 0.8;
                    }
                    else {
                       SF_reco[0] = 1.0;   SF_reco[1] = 0.993; SF_reco[2] = 0.991; SF_reco[3] = 0.985;
                       SF_reco[4] = 0.981; SF_reco[5] = 0.979; SF_reco[6] = 0.978; SF_reco[7] = 0.9;
                    }
                    double pEdge[8];
                    pEdge[0] = 100; pEdge[1] = 150; pEdge[2] = 200;  pEdge[3] = 300;
                    pEdge[4] = 400; pEdge[5] = 600; pEdge[6] = 1500; pEdge[7] = 3500;
                    unsigned int ibin;
                    for (ibin = 0; ibin < 8; ibin++) {
                       if (p < pEdge[ibin]) 
                          return sf = SF_reco[ibin];
                    }
                    return sf = SF_reco[7];
                }

                double get_idSF_weight(std::string syst, Vfloat eta, Vfloat TuneP_relPt, Vfloat pt, int mu1_idx, int mu2_idx) {
                    double mu1_sf, mu2_sf, weight;
                    mu1_sf = get_mu_idSF(syst, eta[mu1_idx], TuneP_relPt[mu1_idx], pt[mu1_idx]);
                    mu2_sf = get_mu_idSF(syst, eta[mu2_idx], TuneP_relPt[mu2_idx], pt[mu2_idx]);
                    return weight = mu1_sf*mu2_sf;
                }
                double get_isoSF_weight(std::string syst, Vfloat eta, Vfloat TuneP_relPt, Vfloat pt, int mu1_idx, int mu2_idx) {
                    double mu1_sf, mu2_sf, weight;
                    mu1_sf = get_mu_isoSF(syst, eta[mu1_idx], TuneP_relPt[mu1_idx], pt[mu1_idx]);
                    mu2_sf = get_mu_isoSF(syst, eta[mu2_idx], TuneP_relPt[mu2_idx], pt[mu2_idx]);
                    return weight = mu1_sf*mu2_sf;
                }
            //// For Reco & Trigger SFs from JSON ////
                double get_trigSF_weight(std::string syst, Vfloat eta, Vfloat TuneP_relPt, Vfloat pt, int mu1_idx, int mu2_idx) {
                    double mu1_sf, mu2_sf, weight;
                    mu1_sf = get_mu_trigSF(syst, eta[mu1_idx], TuneP_relPt[mu1_idx], pt[mu1_idx]);
                    mu2_sf = get_mu_trigSF(syst, eta[mu2_idx], TuneP_relPt[mu2_idx], pt[mu2_idx]);
                    return weight = mu1_sf + mu2_sf - mu1_sf*mu2_sf;
                }
                double get_recoSF_weight(std::string syst, Vfloat eta, Vfloat TuneP_relPt, Vfloat pt, int mu1_idx, int mu2_idx) {
                    double mu1_sf, mu2_sf, weight;
                    mu1_sf = get_mu_recoSF(syst, eta[mu1_idx], TuneP_relPt[mu1_idx], pt[mu1_idx]);
                    mu2_sf = get_mu_recoSF(syst, eta[mu2_idx], TuneP_relPt[mu2_idx], pt[mu2_idx]);
                    return weight = mu1_sf*mu2_sf;
                }
            """  % (self.tag, self.tag, self.tag, self.tag))

            ROOT.gInterpreter.Declare("""
            //// For Reco & Trigger SFs in Table format ////
                double get_trigSF_weight_OLDforLegacy(Vfloat eta, Vfloat TuneP_relPt, Vfloat pt, int mu1_idx, int mu2_idx) {
                    double mu1_sf, mu2_sf, weight;
                    mu1_sf = get_mu_trigSF_OLDforLegacy%s(eta[mu1_idx], TuneP_relPt[mu1_idx], pt[mu1_idx]);
                    mu2_sf = get_mu_trigSF_OLDforLegacy%s(eta[mu2_idx], TuneP_relPt[mu2_idx], pt[mu2_idx]);
                    return weight = mu1_sf + mu2_sf - mu1_sf*mu2_sf;
                }
                double get_recoSF_weight_OLDforLegacy(Vfloat eta, Vfloat TuneP_relPt, Vfloat pt, Vfloat phi, Vfloat mass, int mu1_idx, int mu2_idx) {
                    double mu1_sf, mu2_sf, weight;
                    mu1_sf = get_mu_recoSF_OLDforLegacy%s(eta[mu1_idx], TuneP_relPt[mu1_idx], pt[mu1_idx], phi[mu1_idx], mass[mu1_idx]);
                    mu2_sf = get_mu_recoSF_OLDforLegacy%s(eta[mu2_idx], TuneP_relPt[mu2_idx], pt[mu2_idx], phi[mu2_idx], mass[mu2_idx]);
                    return weight = mu1_sf*mu2_sf;
                }
            """  % (self.year, self.year, self.year, self.year))

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

        branches = []
        for corr in ["idSF", "isoSF", "trigSF", "recoSF"]:
            for syst_name, syst in [("", "sf"), ("_up", "systup"), ("_down", "systdown")]:
                df = df.Define("mu_%s_weight%s" % (corr, syst_name),
                               'get_%s_weight("%s", Muon_eta, Muon_tunepRelPt, Muon_pt, mu1_index, mu2_index)' % (corr, syst))
                branches.append("mu_%s_weight%s" % (corr, syst_name))

        ## For Reco & Trigger SFs in Table format ##
        #df = df.Define("mu_trigSF_weight", 'get_trigSF_weight(Muon_eta, Muon_tunepRelPt, Muon_pt, mu1_index, mu2_index)')
        #df = df.Define("mu_recoSF_weight", 'get_recoSF_weight(Muon_eta, Muon_tunepRelPt, Muon_pt, Muon_phi, Muon_mass, mu1_index, mu2_index)')
        #branches.append("mu_trigSF_weight")
        #branches.append("mu_recoSF_weight")

        return df, branches


[docs]def MuonSF(**kwargs): """ Module to obtain high-pt di-muon SFs with their uncertainties for: reco, ID, iso & trigger. Implemented for the full RunII (2016, 2017 & 2018). All SFs are taken from official JSONs. Alternatively, trigger & reco SFs can be read from a table format, taken from the Muon POG twiki. YAML sintaxis: .. code-block:: yaml codename: name: MuonSF path: Corrections.MUO.muonSF_GEMethod parameters: isMC: self.dataset.process.isMC year: self.config.year isUL: self.dataset.has_tag('ul') preVFP: self.dataset.has_tag('preVFP') """ return lambda: MuonSFproducer(**kwargs)