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/TAU/{}/tau.json.gz"
class TauSFRDFProducer():
def __init__(self, *args, **kwargs):
self.year = kwargs.pop("year")
self.isUL = kwargs.pop("isUL")
self.isMC = kwargs.pop("isMC")
self.vsjet_wps = kwargs.pop("vsjet_wps", ["Medium", "Tight"])
self.vse_wps = kwargs.pop("vse_wps", ["VVLoose", "Tight"])
self.vsmu_wps = kwargs.pop("vsmu_wps", ["VLoose", "Tight"])
self.vsjet_vse_wp = kwargs.pop("vsjet_vse_wp", "VVLoose")
if self.isMC:
if self.year == 2018:
if self.isUL:
filename = json_path.format("2018_UL")
else:
filename = json_path.format("2018_ReReco")
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_vsjet = MyCorrections("%s", "DeepTau2017v2p1VSjet");' % filename)
ROOT.gInterpreter.ProcessLine(
'auto corr_vse = MyCorrections("%s", "DeepTau2017v2p1VSe");' % filename)
ROOT.gInterpreter.ProcessLine(
'auto corr_vsmu = MyCorrections("%s", "DeepTau2017v2p1VSmu");' % filename)
ROOT.gInterpreter.ProcessLine(
'auto corr_tes = MyCorrections("%s", "tau_energy_scale");' % filename)
if self.isUL:
ROOT.gInterpreter.Declare("""
using Vfloat = const ROOT::RVec<float>&;
using Vint = const ROOT::RVec<int>&;
ROOT::RVec<double> get_deeptau_vsjet_sf(
Vfloat pt, Vint dm, Vint genmatch,
std::string wp, std::string wp_vse, std::string syst, std::string flag) {
ROOT::RVec<double> sf;
for (size_t i = 0; i < pt.size(); i++) {
if (flag == "dm" && (pt[i] <= 40 || (dm[i] > 2 && dm[i] < 10))) {
sf.push_back(1.);
}
else
sf.push_back(corr_vsjet.eval({pt[i], dm[i], genmatch[i], wp, wp_vse, syst, flag}));
}
return sf;
}
""")
else:
ROOT.gInterpreter.Declare("""
using Vfloat = const ROOT::RVec<float>&;
using Vint = const ROOT::RVec<int>&;
ROOT::RVec<double> get_deeptau_vsjet_sf(
Vfloat pt, Vint dm, Vint genmatch,
std::string wp, std::string syst, std::string flag) {
ROOT::RVec<double> sf;
for (size_t i = 0; i < pt.size(); i++) {
if (flag == "dm" && (pt[i] <= 40 || (dm[i] > 2 && dm[i] < 10))) {
sf.push_back(1.);
}
else
sf.push_back(corr_vsjet.eval({pt[i], dm[i], genmatch[i], wp, syst, flag}));
}
return sf;
}
""")
ROOT.gInterpreter.Declare("""
ROOT::RVec<double> get_deeptau_vse_sf(
Vfloat eta, Vint genmatch, std::string wp, std::string syst) {
ROOT::RVec<double> sf;
for (size_t i = 0; i < eta.size(); i++) {
sf.push_back(corr_vse.eval({eta[i], genmatch[i], wp, syst}));
}
return sf;
}
ROOT::RVec<double> get_deeptau_vsmu_sf(
Vfloat eta, Vint genmatch, std::string wp, std::string syst) {
ROOT::RVec<double> sf;
for (size_t i = 0; i < eta.size(); i++) {
sf.push_back(corr_vsmu.eval({eta[i], genmatch[i], wp, syst}));
}
return sf;
}
ROOT::RVec<double> get_tes(
Vfloat pt, Vfloat eta, Vint dm, Vint genmatch, std::string id, std::string syst) {
ROOT::RVec<double> tes_factor;
for (size_t i = 0; i < eta.size(); i++) {
if (dm[i] == 5 || dm[i] == 6) tes_factor.push_back(1.);
else tes_factor.push_back(corr_tes.eval({pt[i], eta[i], dm[i], genmatch[i], id, syst}));
}
return tes_factor;
}
""")
def run(self, df):
if not self.isMC:
return df, []
branches = []
for syst_name, syst in [("", "nom"), ("_up", "up"), ("_down", "down")]:
for wp in self.vsjet_wps:
if self.isUL:
df = df.Define("Tau_sfDeepTau2017v2p1VSjet_pt_binned_%s%s" % (wp, syst_name),
'get_deeptau_vsjet_sf('
'Tau_pt, Tau_decayMode, Tau_genPartFlav, "%s", "%s", "%s", "pt")' % (
wp, self.vsjet_vse_wp, syst))
df = df.Define(
"Tau_sfDeepTau2017v2p1VSjet_dm_binned_%s%s" % (wp, syst_name),
'get_deeptau_vsjet_sf('
'Tau_pt, Tau_decayMode, Tau_genPartFlav, "%s", "%s", "%s", "dm")' % (
wp, self.vsjet_vse_wp, syst))
else:
df = df.Define("Tau_sfDeepTau2017v2p1VSjet_pt_binned_%s%s" % (wp, syst_name),
'get_deeptau_vsjet_sf('
'Tau_pt, Tau_decayMode, Tau_genPartFlav, "%s", "%s", "pt")' % (
wp, syst))
df = df.Define(
"Tau_sfDeepTau2017v2p1VSjet_dm_binned_%s%s" % (wp, syst_name),
'get_deeptau_vsjet_sf('
'Tau_pt, Tau_decayMode, Tau_genPartFlav, "%s", "%s", "dm")' % (
wp, syst))
branches.append("Tau_sfDeepTau2017v2p1VSjet_pt_binned_%s%s" % (wp, syst_name))
branches.append("Tau_sfDeepTau2017v2p1VSjet_dm_binned_%s%s" % (wp, syst_name))
for wp in self.vse_wps:
df = df.Define("Tau_sfDeepTau2017v2p1VSe_%s%s" % (wp, syst_name),
'get_deeptau_vse_sf(Tau_eta, Tau_genPartFlav, "%s", "%s")' % (wp, syst))
branches.append("Tau_sfDeepTau2017v2p1VSe_%s%s" % (wp, syst_name))
for wp in self.vsmu_wps:
df = df.Define("Tau_sfDeepTau2017v2p1VSmu_%s%s" % (wp, syst_name),
'get_deeptau_vsmu_sf(Tau_eta, Tau_genPartFlav, "%s", "%s")' % (wp, syst))
branches.append("Tau_sfDeepTau2017v2p1VSmu_%s%s" % (wp, syst_name))
df = df.Define("tes_factor%s" % syst_name,
'get_tes(Tau_pt, Tau_eta, Tau_decayMode, Tau_genPartFlav, '
'"DeepTau2017v2p1", "%s")' % syst)
df = df.Define("Tau_pt_corr%s" % syst_name, "Tau_pt * tes_factor%s" % syst_name)
df = df.Define("Tau_mass_corr%s" % syst_name, "Tau_mass * tes_factor%s" % syst_name)
branches.append("tes_factor%s" % syst_name)
branches.append("Tau_pt_corr%s" % syst_name)
branches.append("Tau_mass_corr%s" % syst_name)
return df, branches
[docs]def tauSFRDF(**kwargs):
"""
Module to obtain DeepTau2017v2p1 SFs and TES with their up/down uncertainties.
:param vsjet_wps: DeepTau2017v2p1VSjet WPs to consider. Default: `[Medium, Tight]`
:type vsjet_wps: list of str
:param vse_wps: DeepTau2017v2p1VSe WPs to consider. Default: `[VVLoose, Tight]`
:type vse_wps: list of str
:param vsmu_wps: DeepTau2017v2p1VSmu WPs to consider. Default: `[VLoose, Tight]`
:type vsmu_wps: list of str
:param vsjet_vse_wp: DeepTau2017v2p1VSeWP used to compute the DeepTau2017v2p1VSjet scale factor.
Default: `VVLoose`
:type vsjet_vse_wp: str
YAML sintaxis:
.. code-block:: yaml
codename:
name: tauSFRDF
path: Corrections.TAU.tauCorrections
parameters:
isMC: self.dataset.process.isMC
year: self.config.year
isUL: self.dataset.has_tag('ul')
vsjet_wps: [Medium, Tight]
vse_wps: [VVLoose, Tight]
vsmu_wps: [VLoose, Tight]
"""
return lambda: TauSFRDFProducer(**kwargs)