from PhysicsTools.NanoAODTools.postprocessing.framework.eventloop import Module
from PhysicsTools.NanoAODTools.postprocessing.framework.datamodel import Collection, Object
from math import sin, cos
from analysis_tools.utils import import_root
ROOT = import_root()
class DummyModule(Module):
def beginFile(self, inputFile, outputFile, inputTree, wrappedOutputTree):
pass
def endFile(self, inputFile, outputFile, inputTree, wrappedOutputTree):
pass
def analyze(self, event):
return True
class Syst(object):
def __init__(self, isMC=True, systs=None, *args, **kwargs):
self.isMC = isMC
self.set_systs(systs)
def set_systs(self, systs):
self.systs = []
for systname, systvalue in systs.items():
setattr(self, systname, ("" if not self.isMC
else ("_%s" % systvalue if systvalue != "" else "")))
if "syst" not in systname:
continue
if systvalue != systs["%s_central" % systname.split("_")[0]]:
self.systs.append(eval("self.%s" % systname))
self.systs = "".join(self.systs)
[docs]class JetLepMetSyst(Syst):
"""
Abstract module that includes systematics and central values for jets, electrons,
muons and taus into other modules.
Note that at most one lepton systematic can be different from its corresponden lepton
central. If this happen, that systematic will be included as ``self.lepton_syst``.
All systematics different from their corresponding central values will be joined
and included into ``self.systs``.
:param jet_central: central appendix for the jets. Default: ``nom`` (smearing).
:type jet_central: str
:param jet_syst: systematic appendix for the jets. Default: ``nom`` (smearing).
:type jet_syst: str
:param muon_central: central appendix for the muons. Default: ``""``.
:type muon_central: str
:param muon_syst: systematic appendix for the muons. Default: ``""``.
:type muon_syst: str
:param electron_central: central appendix for the electrons. Default: ``""``.
:type electron_central: str
:param electron_syst: systematic appendix for the electrons. Default: ``""``.
:type electron_syst: str
:param tau_central: central appendix for the taus. Default: ``corr`` (tes).
:type tau_central: str
:param tau_syst: systematic appendix for the taus. Default: ``corr`` (tes).
:type tau_syst: str
:param met_central: central appendix for the met. Default: ``""``.
:type met_central: str
:param met_syst: systematic appendix for the met. Default: ``""``.
:type met_syst: str
:param met_smear_tag: smearing tag used for the met. Default: ``smeared``
:type met_smear_tag: str
"""
def __init__(self, *args, **kwargs):
kwargs.setdefault("jet_syst", "nom")
kwargs.setdefault("jet_central", "nom")
kwargs.setdefault("muon_syst", "")
kwargs.setdefault("muon_central", "")
kwargs.setdefault("electron_syst", "")
kwargs.setdefault("electron_central", "")
kwargs.setdefault("photon_syst", "")
kwargs.setdefault("photon_central", "")
kwargs.setdefault("tau_syst", "corr")
kwargs.setdefault("tau_central", "corr")
kwargs.setdefault("met_syst", "")
kwargs.setdefault("met_central", "")
kwargs.setdefault("met_smear_tag", "smeared")
super(JetLepMetSyst, self).__init__(kwargs.pop("isMC", True), kwargs)
assert sum([self.muon_syst != self.muon_central,
self.electron_syst != self.electron_central, self.tau_syst != self.tau_central]) <= 1
self.lep_syst = ""
if self.muon_syst != self.muon_central:
self.lep_syst = self.muon_syst
elif self.electron_syst != self.electron_central:
self.lep_syst = self.electron_syst
elif self.tau_syst != self.tau_central:
self.lep_syst = self.tau_syst
class JetLepMetModule(JetLepMetSyst, Module):
def get_daus(self, event, muons, electrons, taus):
if event.pairType == 0:
dau1 = muons[event.dau1_index]
dau2 = taus[event.dau2_index]
dau1_syst = self.muon_syst
elif event.pairType == 1:
dau1 = electrons[event.dau1_index]
dau2 = taus[event.dau2_index]
dau1_syst = self.electron_syst
elif event.pairType == 2:
dau1 = taus[event.dau1_index]
dau2 = taus[event.dau2_index]
dau1_syst = self.tau_syst
else:
raise ValueError("pairType %s is not implemented" % event.pairType)
dau2_syst = self.tau_syst
# build TLorentzVectors
dau1_tlv = ROOT.TLorentzVector()
dau2_tlv = ROOT.TLorentzVector()
dau1_tlv.SetPtEtaPhiM(
eval("dau1.pt%s" % dau1_syst), dau1.eta, dau1.phi, eval("dau1.mass%s" % dau1_syst))
dau2_tlv.SetPtEtaPhiM(
eval("dau2.pt%s" % dau2_syst), dau2.eta, dau2.phi, eval("dau2.mass%s" % dau2_syst))
return dau1, dau2, dau1_tlv, dau2_tlv
def get_bjets(self, event, jets):
bjet1 = jets[event.bjet1_JetIdx]
bjet2 = jets[event.bjet2_JetIdx]
# build TLorentzVectors
bjet1_tlv = ROOT.TLorentzVector()
bjet2_tlv = ROOT.TLorentzVector()
bjet1_tlv.SetPtEtaPhiM(eval("bjet1.pt%s" % self.jet_syst), bjet1.eta,
bjet1.phi, eval("bjet1.mass%s" % self.jet_syst))
bjet2_tlv.SetPtEtaPhiM(eval("bjet2.pt%s" % self.jet_syst), bjet2.eta,
bjet2.phi, eval("bjet2.mass%s" % self.jet_syst))
return bjet1, bjet2, bjet1_tlv, bjet2_tlv
def get_vbfjets(self, event, jets):
if event.VBFjet1_JetIdx != -1 and event.VBFjet2_JetIdx != -1:
vbfjet1 = jets[event.VBFjet1_JetIdx]
vbfjet2 = jets[event.VBFjet2_JetIdx]
# build TLorentzVectors
vbfjet1_tlv = ROOT.TLorentzVector()
vbfjet2_tlv = ROOT.TLorentzVector()
vbfjet1_tlv.SetPtEtaPhiM(eval("vbfjet1.pt%s" % self.jet_syst), vbfjet1.eta,
vbfjet1.phi, eval("vbfjet1.mass%s" % self.jet_syst))
vbfjet2_tlv.SetPtEtaPhiM(eval("vbfjet2.pt%s" % self.jet_syst), vbfjet2.eta,
vbfjet2.phi, eval("vbfjet2.mass%s" % self.jet_syst))
else:
return None, None, None, None
return vbfjet1, vbfjet2, vbfjet1_tlv, vbfjet2_tlv
def get_met(self, event):
met = Object(event, "MET")
met_bis = Object(event, "MET%s" % self.met_smear_tag)
met_tlv = ROOT.TLorentzVector()
met_tlv.SetPxPyPzE(
eval("met_bis.pt%s" % self.met_syst) * cos(eval("met_bis.phi%s" % self.met_syst)),
eval("met_bis.pt%s" % self.met_syst) * sin(eval("met_bis.phi%s" % self.met_syst)),
0,
eval("met_bis.pt%s" % self.met_syst)
)
return met, met_tlv
class VarFromVectorRDFProducer():
def __init__(self, *args, **kwargs):
self.variables = kwargs.pop("variables")
self.index = kwargs.pop("index")
self.input_prefix = kwargs.pop("input_prefix")
self.output_prefix = kwargs.pop("output_prefix")
self.condition = kwargs.pop("condition", True)
def run(self, df):
branches = []
if self.condition:
for var in self.variables:
branches.append(var.replace(self.input_prefix, self.output_prefix))
df = df.Define(branches[-1], "%s.at(%s)" % (var, self.index))
return df, branches
[docs]def VarFromVectorRDF(**kwargs):
"""
Returns the element `index` of whatever input vectors
:param variables: input vectors
:type variables: list of str
:param index: index to get the element from the vectors
:type index: str
:param input_prefix: prefix common to all vectors to be substituted
:type input_prefix: str
:param output_prefix: prefix to use as output for all vectors
:type output_prefix: str
:param condition: (optional) condition used to define or not the new variables
:type condition: str
YAML sintaxis:
.. code-block:: yaml
codename:
name: VarFromVectorRDF
path: Base.Modules.BaseModules
parameters:
variables: [Jet_pt, Jet_eta]
index: bjet_index
input_prefix = Jet
output_prefix = bjet
"""
return lambda: VarFromVectorRDFProducer(**kwargs)