postprocessing/framework/datamodel.py

0001 from PhysicsTools.NanoAODTools.postprocessing.framework.treeReaderArrayTools import InputTree
0002 import ROOT
0003 import math
0004 ROOT.PyConfig.IgnoreCommandLineOptions = True
0005
0006 statusflags = { # GenPart_statusFlags, stored bitwise (powers of 2):
0007   # https://github.com/cms-sw/cmssw/edit/master/PhysicsTools/NanoAOD/python/genparticles_cff.py
0008   # https://cms-nanoaod-integration.web.cern.ch/integration/master-106X/mc106Xul18_doc.html#GenPart
0009   'isPrompt':                      (1 << 0),   'fromHardProcess':                    (1 <<  8),
0010   'isDecayedLeptonHadron':         (1 << 1),   'isHardProcessTauDecayProduct':       (1 <<  9),
0011   'isTauDecayProduct':             (1 << 2),   'isDirectHardProcessTauDecayProduct': (1 << 10),
0012   'isPromptTauDecayProduct':       (1 << 3),   'fromHardProcessBeforeFSR':           (1 << 11),
0013   'isDirectTauDecayProduct':       (1 << 4),   'isFirstCopy':                        (1 << 12),
0014   'isDirectPromptTauDecayProduct': (1 << 5),   'isLastCopy':                         (1 << 13),
0015   'isDirectHadronDecayProduct':    (1 << 6),   'isLastCopyBeforeFSR':                (1 << 14),
0016   'isHardProcess':                 (1 << 7),
0017 }
0018
0019
0020 class Event:
0021     """Class that allows seeing an entry of a PyROOT TTree as an Event"""
0022
0023     def __init__(self, tree, entry):
0024         self._tree = tree
0025         self._entry = entry
0026         self._tree.gotoEntry(entry)
0027
0028     def __getattr__(self, name):
0029         if name in self.__dict__:
0030             return self.__dict__[name]
0031         return self._tree.readBranch(name)
0032
0033     def __getitem__(self, attr):
0034         return self.__getattr__(attr)
0035
0036     def eval(self, expr):
0037         """Evaluate an expression, as TTree::Draw would do.
0038
0039            This is added for convenience, but it may perform poorly and the implementation is not bulletproof,
0040            so it's better to rely on reading values, collections or objects directly
0041         """
0042         if not hasattr(self._tree, '_exprs'):
0043             self._tree._exprs = {}
0044             # remove useless warning about EvalInstance()
0045             import warnings
0046             warnings.filterwarnings(action='ignore', category=RuntimeWarning,
0047                                     message='creating converter for unknown type "const char\*\*"$')
0048             warnings.filterwarnings(action='ignore', category=RuntimeWarning,
0049                                     message='creating converter for unknown type "const char\*\[\]"$')
0050         if expr not in self._tree._exprs:
0051             formula = ROOT.TTreeFormula(expr, expr, self._tree)
0052             if formula.IsInteger():
0053                 formula.go = formula.EvalInstance64
0054             else:
0055                 formula.go = formula.EvalInstance
0056             self._tree._exprs[expr] = formula
0057             # force sync, to be safe
0058             self._tree.GetEntry(self._entry)
0059             self._tree.entry = self._entry
0060             # self._tree._exprs[expr].SetQuickLoad(False)
0061         else:
0062             self._tree.gotoEntry(self._entry)
0063             formula = self._tree._exprs[expr]
0064         if "[" in expr:  # unclear why this is needed, but otherwise for some arrays x[i] == 0 for all i > 0
0065             formula.GetNdata()
0066         return formula.go()
0067
0068
0069 class Object:
0070     """Class that allows seeing a set branches plus possibly an index as an Object"""
0071
0072     def __init__(self, event, prefix, index=None):
0073         self._event = event
0074         self._prefix = prefix + "_"
0075         self._index = index
0076
0077     def __getattr__(self, name):
0078         if name in self.__dict__:
0079             return self.__dict__[name]
0080         if name[:2] == "__" and name[-2:] == "__":
0081             raise AttributeError
0082         val = getattr(self._event, self._prefix + name)
0083         if self._index != None:
0084             val = val[self._index]
0085         # convert char to integer number
0086         val = ord(val) if type(val) == str else val
0087         self.__dict__[name] = val  # cache
0088         return val
0089
0090     def __getitem__(self, attr):
0091         return self.__getattr__(attr)
0092
0093     def p4(self, corr_pt=None):
0094         """Create TLorentzVector for this particle."""
0095         ret = ROOT.TLorentzVector()
0096         if corr_pt == None:
0097             ret.SetPtEtaPhiM(self.pt, self.eta, self.phi, self.mass)
0098         else:
0099             ret.SetPtEtaPhiM(corr_pt, self.eta, self.phi, self.mass)
0100         return ret
0101
0102     def DeltaR(self, other):
0103         if isinstance(other, ROOT.TLorentzVector):
0104             deta = abs(other.Eta() - self.eta)
0105             dphi = abs(other.Phi() - self.phi)
0106         else:
0107             deta = abs(other.eta - self.eta)
0108             dphi = abs(other.phi - self.phi)
0109         while dphi > math.pi:
0110             dphi = abs(dphi - 2 * math.pi)
0111         return math.sqrt(dphi**2 + deta**2)
0112
0113     def statusflag(self, flag):
0114         """Find if bit for statusflag is set (for GenPart only)."""
0115         return (self.statusFlags & statusflags[flag])==statusflags[flag]
0116
0117     def subObj(self, prefix):
0118         return Object(self._event, self._prefix + prefix)
0119
0120     def __repr__(self):
0121         return ("<%s[%s]>" % (self._prefix[:-1], self._index)) if self._index != None else ("<%s>" % self._prefix[:-1])
0122
0123     def __str__(self):
0124         return self.__repr__()
0125
0126
0127 class Collection:
0128     def __init__(self, event, prefix, lenVar=None):
0129         self._event = event
0130         self._prefix = prefix
0131         if lenVar != None:
0132             self._len = getattr(event, lenVar)
0133         else:
0134             self._len = getattr(event, "n" + prefix)
0135         self._cache = {}
0136
0137     def __getitem__(self, index):
0138         if type(index) == int and index in self._cache:
0139             return self._cache[index]
0140         if index >= self._len:
0141             raise IndexError("Invalid index %r (len is %r) at %s" % (index, self._len, self._prefix))
0142         elif index < 0:
0143             raise IndexError("Invalid index %r (negative) at %s" % (index, self._prefix))
0144         ret = Object(self._event, self._prefix, index=index)
0145         if type(index) == int:
0146             self._cache[index] = ret
0147         return ret
0148
0149     def __len__(self):
0150         return self._len