postprocessing/framework/output.py

0001 from PhysicsTools.NanoAODTools.postprocessing.framework.treeReaderArrayTools import setExtraBranch
0002 from array import array
0003 import ROOT
0004 ROOT.PyConfig.IgnoreCommandLineOptions = True
0005
0006
0007 _rootBranchType2PythonArray = {
0008     'b': 'B',
0009     'B': 'b',
0010     's': 'H',
0011     'S': 'h',
0012     'i': 'I',
0013     'I': 'i',
0014     'F': 'f',
0015     'D': 'd',
0016     'l': 'L',
0017     'L': 'l',
0018     'O': 'B'
0019 }
0020
0021 # Convert data type string returned by TLeaf.GetTypeName() to root branch type
0022 _rootDataType2BranchType ={
0023     'UChar_t'   : 'b',
0024     'Char_t'    : 'B',
0025     'UShort_t'  : 's',
0026     'Short_t'   : 'S',
0027     'UInt_t'    : 'i',
0028     'Int_t'     : 'I',
0029     'Float_t'   : 'F',
0030     'Double_t'  : 'D',
0031     'ULong64_t' : 'l',
0032     'Long64_t'  : 'L',
0033     'Bool_t'    : 'O'
0034 }
0035
0036 class OutputBranch:
0037     def __init__(
0038             self, tree, name, rootBranchType, n=1,
0039             lenVar=None, title=None, limitedPrecision=False
0040     ):
0041         n = int(n)
0042         self.n = n
0043         self.lenVar = lenVar
0044         self.precision = ROOT.MiniFloatConverter.ReduceMantissaToNbitsRounding(
0045             limitedPrecision) if limitedPrecision and rootBranchType == 'F' else lambda x: x
0046         # check if a branch was already there
0047         existingBranch = tree.GetBranch(name)
0048         if (existingBranch):
0049             existingBranchType = _rootDataType2BranchType[existingBranch.FindLeaf(name).GetTypeName()]
0050             if rootBranchType != existingBranchType :
0051                 print(f'Warning: output branch {name} already exists with type {existingBranchType}. Will ignore specified type {rootBranchType}')
0052             self.buff = array(
0053                 _rootBranchType2PythonArray[existingBranchType], n * [0. if existingBranchType in 'FD' else 0])
0054             self.branch = existingBranch
0055             self.branch.SetAddress(self.buff)
0056         else:
0057             self.buff = array(
0058                 _rootBranchType2PythonArray[rootBranchType], n * [0. if rootBranchType in 'FD' else 0])
0059             if lenVar != None:
0060                 self.branch = tree.Branch(
0061                     name, self.buff, "%s[%s]/%s" % (name, lenVar, rootBranchType))
0062             elif n == 1:
0063                 self.branch = tree.Branch(
0064                     name, self.buff, name + "/" + rootBranchType)
0065             else:
0066                 self.branch = tree.Branch(
0067                     name, self.buff, "%s[%d]/%s" % (name, n, rootBranchType))
0068         if title:
0069             self.branch.SetTitle(title)
0070
0071     def fill(self, val):
0072         if self.lenVar:
0073             if len(self.buff) < len(val):  # realloc
0074                 self.buff = array(self.buff.typecode, max(
0075                     len(val), 2 * len(self.buff)) * [0. if self.buff.typecode in 'fd' else 0])
0076                 self.branch.SetAddress(self.buff)
0077             for i, v in enumerate(val):
0078                 self.buff[i] = self.precision(v)
0079         elif self.n == 1:
0080             self.buff[0] = self.precision(val)
0081         else:
0082             if len(val) != self.n:
0083                 raise RuntimeError("Mismatch in filling branch %s of fixed length %d with %d values (%s)" % (
0084                     self.branch.GetName(), self.n, len(val), val))
0085             for i, v in enumerate(val):
0086                 self.buff[i] = v
0087
0088
0089 class OutputTree:
0090     def __init__(self, tfile, ttree, intree):
0091         self._file = tfile
0092         self._tree = ttree
0093         self._intree = intree
0094         self._branches = {}
0095
0096     def branch(
0097             self, name, rootBranchType, n=1, lenVar=None,
0098             title=None, limitedPrecision=False
0099     ):
0100         # and (not self._tree.GetBranch(lenVar)):
0101         if (lenVar != None) and (lenVar not in self._branches):
0102             self._branches[lenVar] = OutputBranch(self._tree, lenVar, "I")
0103         self._branches[name] = OutputBranch(
0104             self._tree, name, rootBranchType, n=n,
0105             lenVar=lenVar, title=title, limitedPrecision=limitedPrecision
0106         )
0107         return self._branches[name]
0108
0109     def fillBranch(self, name, val):
0110         br = self._branches[name]
0111         if br.lenVar and (br.lenVar in self._branches):
0112             self._branches[br.lenVar].buff[0] = len(val)
0113             setExtraBranch(self._intree, br.lenVar, len(val))
0114         br.fill(val)
0115         setExtraBranch(self._intree, name, val)
0116
0117     def tree(self):
0118         return self._tree
0119
0120     def fill(self):
0121         self._tree.Fill()
0122
0123     def write(self):
0124         self._file.cd()
0125         self._tree.Write()
0126
0127
0128 class FullOutput(OutputTree):
0129     def __init__(
0130             self,
0131             inputFile,
0132             inputTree,
0133             outputFile,
0134             branchSelection=None,
0135             outputbranchSelection=None,
0136             fullClone=False,
0137             maxEntries=None,
0138             firstEntry=0,
0139             provenance=False,
0140             jsonFilter=None
0141     ):
0142         outputFile.cd()
0143
0144         self.outputbranchSelection = outputbranchSelection
0145         self.maxEntries = maxEntries
0146         self.firstEntry = firstEntry
0147         if fullClone:
0148             outputTree = inputTree.CopyTree(
0149                 '1', "", maxEntries if maxEntries else ROOT.TVirtualTreePlayer.kMaxEntries, firstEntry)
0150         else:
0151             outputTree = inputTree.CloneTree(0)
0152
0153         # enable back all branches in inputTree, then disable for computation
0154         # the branches as requested in branchSelection
0155         inputTree.SetBranchStatus("*", 1)
0156         if branchSelection:
0157             branchSelection.selectBranches(inputTree)
0158
0159         OutputTree.__init__(self, outputFile, outputTree, inputTree)
0160         self._inputTree = inputTree
0161         self._otherTrees = {}
0162         self._otherObjects = {}
0163         for k in inputFile.GetListOfKeys():
0164             kn = k.GetName()
0165             if kn == "Events":
0166                 continue  # this we are doing
0167             elif kn in ("MetaData", "ParameterSets"):
0168                 if provenance:
0169                     # treat content of other trees as if associated to event 0
0170                     self._otherTrees[kn] = inputFile.Get(
0171                         kn).CopyTree('1' if firstEntry == 0 else '0')
0172             elif kn in ("LuminosityBlocks", "Runs"):
0173                 if not jsonFilter:
0174                     self._otherTrees[kn] = inputFile.Get(
0175                         kn).CopyTree('1' if firstEntry == 0 else '0')
0176                 elif firstEntry == 0:
0177                     _isRun = (kn == "Runs")
0178                     _it = inputFile.Get(kn)
0179                     _ot = _it.CloneTree(0)
0180                     for ev in _it:
0181                         if (jsonFilter.filterRunOnly(ev.run) if _isRun else jsonFilter.filterRunLumi(ev.run, ev.luminosityBlock)):
0182                             _ot.Fill()
0183                     self._otherTrees[kn] = _ot
0184             elif k.GetClassName() == "TTree":
0185                 print("Not copying unknown tree %s" % kn)
0186             else:
0187                 self._otherObjects[kn] = inputFile.Get(kn)
0188
0189     def fill(self):
0190         self._inputTree.readAllBranches()
0191         self._tree.Fill()
0192
0193     def write(self):
0194         if self.outputbranchSelection:
0195             self.outputbranchSelection.selectBranches(self._tree)
0196         self._tree = self.tree().CopyTree('1', "",
0197                                           self.maxEntries if self.maxEntries else ROOT.TVirtualTreePlayer.kMaxEntries, 0)
0198
0199         OutputTree.write(self)
0200         for t in self._otherTrees.values():
0201             t.Write()
0202         for on, ov in self._otherObjects.items():
0203             self._file.WriteTObject(ov, on)
0204
0205
0206 class FriendOutput(OutputTree):
0207     def __init__(self, inputFile, inputTree, outputFile, treeName="Friends"):
0208         outputFile.cd()
0209         outputTree = ROOT.TTree(
0210             treeName, "Friend tree for " + inputTree.GetName())
0211         OutputTree.__init__(self, outputFile, outputTree, inputTree)