Merge pull request uniovi-hepex#15 from peruzzim/nested_modules

Expose values calculated by previous eventLoop modules

Author: arizzi

python/postprocessing/framework/eventloop.py

@@ -1,4 +1,5 @@
 from PhysicsTools.NanoAODTools.postprocessing.framework.datamodel import Event
+from PhysicsTools.NanoAODTools.postprocessing.framework.treeReaderArrayTools import clearExtraBranches
 import sys, time
 
 class Module:
@@ -26,6 +27,7 @@ def eventLoop(modules, inputFile, outputFile, inputTree, wrappedOutputTree, maxE
     for i in xrange(entries) if eventRange == None else eventRange:
         if maxEvents > 0 and i >= maxEvents-1: break
         e = Event(inputTree,i)
+        clearExtraBranches(inputTree)
         doneEvents += 1
         ret = True
         for m in modules: 

python/postprocessing/framework/output.py

@@ -2,6 +2,8 @@
 import ROOT
 ROOT.PyConfig.IgnoreCommandLineOptions = True
 
+from PhysicsTools.NanoAODTools.postprocessing.framework.treeReaderArrayTools import setExtraBranch
+
 _rootBranchType2PythonArray = { 'b':'B', 'B':'b', 'i':'I', 'I':'i', 'F':'f', 'D':'d', 'l':'L', 'L':'l', 'O':'B' }
 
 class OutputBranch:
@@ -30,9 +32,10 @@ def fill(self, val):
             for i,v in enumerate(val): self.buff[i] = v
 
 class OutputTree:
-    def __init__(self, tfile, ttree):
+    def __init__(self, tfile, ttree, intree):
         self._file = tfile
         self._tree = ttree
+        self._intree = intree
         self._branches = {} 
     def branch(self, name, rootBranchType, n=1, lenVar=None, title=None):
         if (lenVar != None) and (lenVar not in self._branches) and (not self._tree.GetBranch(lenVar)):
@@ -43,7 +46,9 @@ def fillBranch(self, name, val):
         br = self._branches[name]
         if br.lenVar and (br.lenVar in self._branches):
             self._branches[br.lenVar].buff[0] = len(val)
+            setExtraBranch(self._intree,br.lenVar,len(val))
         br.fill(val)
+        setExtraBranch(self._intree,name,val)
     def tree(self):
         return self._tree
     def fill(self):
@@ -58,7 +63,7 @@ def __init__(self, inputFile, inputTree, outputFile, branchSelection = None, ful
         if branchSelection: 
             branchSelection.selectBranches(inputTree)
         outputTree = inputTree.CopyTree('1') if fullClone else inputTree.CloneTree(0)
-        OutputTree.__init__(self, outputFile, outputTree)
+        OutputTree.__init__(self, outputFile, outputTree, inputTree)
         self._inputTree = inputTree
         self._otherTrees = {}
         self._otherObjects = {}
@@ -95,5 +100,5 @@ class FriendOutput(OutputTree):
     def __init__(self, inputFile, inputTree, outputFile, treeName="Friends"):
         outputFile.cd()
         outputTree = ROOT.TTree(treeName,"Friend tree for "+inputTree.GetName())
-        OutputTree.__init__(self, outputFile, outputTree)
+        OutputTree.__init__(self, outputFile, outputTree, inputTree)
 

python/postprocessing/framework/treeReaderArrayTools.py

@@ -19,6 +19,7 @@ def InputTree(tree,entrylist=None):
     tree.gotoEntry = types.MethodType(_gotoEntry, tree)
     tree.readAllBranches = types.MethodType(_readAllBranches, tree)
     tree.entries = tree._ttreereader.GetEntries(False)
+    tree._extrabranches={}
     return tree
 
 def getArrayReader(tree, branchName):
@@ -41,11 +42,18 @@ def getValueReader(tree, branchName):
        tree._ttrvs[branchName] = _makeValueReader(tree, typ, branchName)
     return tree._ttrvs[branchName]
 
+def clearExtraBranches(tree):
+    tree._extrabranches = {}
+
+def setExtraBranch(tree,name,val):
+    tree._extrabranches[name] = val
 
 def readBranch(tree, branchName):
     """Return the branch value if the branch is a value, and a TreeReaderArray if the branch is an array"""
     if tree._ttreereader._isClean: raise RuntimeError, "readBranch must not be called before calling gotoEntry"
-    if branchName in tree._ttras: 
+    if branchName in tree._extrabranches:
+        return tree._extrabranches[branchName]
+    elif branchName in tree._ttras:
         return tree._ttras[branchName]
     elif branchName in tree._ttrvs: 
         ret = tree._ttrvs[branchName].Get()[0]