Adding command line rootslimtree for removing branches from a tree.

main/python/cmdLineUtils.py

@@ -747,7 +747,15 @@ def rootCp(sourceList, destFileName, destPathSplit, \
 ##########
 # ROOTEVENTSELECTOR
 
-def _copyTreeSubset(sourceFile,sourcePathSplit,destFile,destPathSplit,firstEvent,lastEvent,selectionString):
+def _setBranchStatus(tree,branchSelectionString,status=0):
+    """This is used by _copyTreeSubset() to turn on/off branches"""
+    for branchToModify in branchSelectionString.split(","):
+        logging.info("Setting branch status to %d for %s"%(status,branchToModify)  )
+        tree.SetBranchStatus(branchToModify,status)
+    return tree
+
+def _copyTreeSubset(sourceFile,sourcePathSplit,destFile,destPathSplit,firstEvent,lastEvent,selectionString,
+    branchinclude, branchexclude):
     """Copy a subset of the tree from (sourceFile,sourcePathSplit)
     to (destFile,destPathSplit) according to options in optDict"""
     retcode = changeDirectory(sourceFile,sourcePathSplit[:-1])
@@ -757,27 +765,31 @@ def _copyTreeSubset(sourceFile,sourcePathSplit,destFile,destPathSplit,firstEvent
     # changeDirectory for the small tree not to be memory-resident
     retcode = changeDirectory(destFile,destPathSplit)
     if retcode != 0: return retcode
-    smallTree = bigTree.CloneTree(0)
+
     if lastEvent == -1:
         lastEvent = nbrEntries-1
-    isNtuple = bigTree.InheritsFrom(ROOT.TNtuple.Class())
-    for i in range(firstEvent, lastEvent+1):
-        bigTree.GetEntry(i)
-        if isNtuple:
-            super(ROOT.TNtuple,smallTree).Fill()
-        else:
-            smallTree.Fill()
-    if selectionString:
-        if isNtuple:
-            smallSkimmedTree = super(ROOT.TNtuple,smallTree).CopyTree(selectionString)
-        else:
-            smallSkimmedTree = smallTree.CopyTree(selectionString)
-        smallSkimmedTree.Write()
-    else:
-        smallTree.Write()
+    numberOfEntries = (lastEvent-firstEvent)+1
+
+    # "Skim" events based on branch values using selectionString
+    # as well as selecting a range of events by index
+    outputTree = bigTree.CopyTree(selectionString,"",numberOfEntries,firstEvent)
+
+    # "Slim" tree by removing branches - 
+    # This is done after the skimming to allow for the user to skim on a 
+    # branch they no longer need to keep
+    if branchexclude:
+        _setBranchStatus(outputTree,branchexclude,0)
+    if branchinclude:
+        _setBranchStatus(outputTree,branchinclude,1)
+    if branchexclude or branchinclude:
+        outputTree = outputTree.CloneTree()
+
+    outputTree.Write()
     return retcode
 
-def _copyTreeSubsets(fileName, pathSplitList, destFile, destPathSplit, first, last, selectionString):
+
+def _copyTreeSubsets(fileName, pathSplitList, destFile, destPathSplit, first, last, selectionString,
+    branchinclude, branchexclude):
     retcode = 0
     destFileName = destFile.GetName()
     rootFile = openROOTFile(fileName) \
@@ -787,12 +799,13 @@ def _copyTreeSubsets(fileName, pathSplitList, destFile, destPathSplit, first, la
     for pathSplit in pathSplitList:
         if isTree(rootFile,pathSplit):
             retcode += _copyTreeSubset(rootFile,pathSplit, \
-            destFile,destPathSplit,first,last,selectionString)
+            destFile,destPathSplit,first,last,selectionString,branchinclude, branchexclude)
     if fileName != destFileName: rootFile.Close()
     return retcode
 
 def rootEventselector(sourceList, destFileName, destPathSplit, \
-                      compress=None, recreate=False, first=0, last=-1, selectionString=""):
+                      compress=None, recreate=False, first=0, last=-1, selectionString="",
+                      branchinclude="", branchexclude=""):
     # Check arguments
     if sourceList == [] or destFileName == "": return 1
     if recreate and destFileName in sourceList:
@@ -807,7 +820,7 @@ def rootEventselector(sourceList, destFileName, destPathSplit, \
     retcode = 0
     for fileName, pathSplitList in sourceList:
         retcode += _copyTreeSubsets(fileName, pathSplitList, destFile, destPathSplit, \
-                                    first, last, selectionString)
+                                    first, last, selectionString, branchinclude, branchexclude)
     destFile.Close()
     return retcode
 

main/python/rooteventselector.py

@@ -5,6 +5,11 @@
 # Mail: [email protected]
 # Date: 20/08/15
 
+# Additions
+# Author: Lawrence Lee
+# Mail: [email protected]
+# Date: 1/4/16
+
 """Command line to copy subsets of trees from source ROOT files to new trees on a destination ROOT file"""
 
 import cmdLineUtils
@@ -34,6 +39,12 @@
 
 - rooteventselector -s "(branch1Value > 100)&&( branch2Value )" source.root:tree dest.root
   Copy the tree 'tree' from 'source.root' to 'dest.root' and apply a selection to the output tree.
+
+- rooteventselector -e "muon_*" source.root:tree dest.root
+  Copy the tree 'tree' from 'source.root' to 'dest.root' and remove branches matching "muon_*"
+
+- rooteventselector -e "*" -i "muon_*" source.root:tree dest.root
+  Copy the tree 'tree' from 'source.root' to 'dest.root' and only write branches matching "muon_*"
 """
 
 def execute():
@@ -44,13 +55,18 @@ def execute():
     parser.add_argument("-f","--first", type=int, default=0, help=FIRST_EVENT_HELP)
     parser.add_argument("-l","--last", type=int, default=-1, help=LAST_EVENT_HELP)
     parser.add_argument("-s","--selection", default="")
+    parser.add_argument("-i","--branchinclude", default="")
+    parser.add_argument("-e","--branchexclude", default="")
 
     # Put arguments in shape
     sourceList, destFileName, destPathSplit, optDict = cmdLineUtils.getSourceDestListOptDict(parser)
 
     # Process rootEventselector
     return cmdLineUtils.rootEventselector(sourceList, destFileName, destPathSplit, \
                                           compress=optDict["compress"], recreate=optDict["recreate"], \
-                                          first=optDict["first"], last=optDict["last"], selectionString=optDict["selection"])
+                                          first=optDict["first"], last=optDict["last"], \
+                                          selectionString=optDict["selection"], \
+                                          branchinclude=optDict["branchinclude"],\
+                                          branchexclude=optDict["branchexclude"])
 
 sys.exit(execute())

main/python/rootslimtree.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env @python@
+
+# ROOT command line tools: rootslimtree
+# Author: Lawrence Lee via Julien Ripoche's rooteventselector
+# Mail: [email protected]
+# Date: 4/4/16
+
+"""Command line to copy trees with subset of branches from source ROOT files to new trees on a destination ROOT file"""
+
+import cmdLineUtils
+import sys
+
+# Help strings
+COMMAND_HELP = "Copy trees with a subset of branches from source ROOT files"
+
+EPILOG="""Examples:
+- rootslimtree source.root:tree dest.root
+  Copy the tree 'tree' from 'source.root' to 'dest.root'.
+
+- rootslimtree --recreate source.root:tree dest.root
+  Recreate the destination file 'dest.root' and copy the tree 'tree' from 'source.root' to 'dest.root'.
+
+- rootslimtree -c 1 source.root:tree dest.root
+  Change the compression factor of the destination file 'dest.root' and  copy the tree 'tree' from 'source.root' to 'dest.root'. For more information about compression settings of ROOT file, please look at the reference guide available on the ROOT site.
+
+- rootslimtree -e "muon_*" source.root:tree dest.root
+  Copy the tree 'tree' from 'source.root' to 'dest.root' and remove branches matching "muon_*"
+
+- rootslimtree -e "*" -i "muon_*" source.root:tree dest.root
+  Copy the tree 'tree' from 'source.root' to 'dest.root' and only write branches matching "muon_*"
+"""
+
+def execute():
+    # Collect arguments with the module argparse
+    parser = cmdLineUtils.getParserSourceDest(COMMAND_HELP, EPILOG)
+    parser.add_argument("-c","--compress", type=int, help=cmdLineUtils.COMPRESS_HELP)
+    parser.add_argument("--recreate", help=cmdLineUtils.RECREATE_HELP, action="store_true")
+    parser.add_argument("-i","--branchinclude", default="")
+    parser.add_argument("-e","--branchexclude", default="")
+
+    # Put arguments in shape
+    sourceList, destFileName, destPathSplit, optDict = cmdLineUtils.getSourceDestListOptDict(parser)
+
+    # Process rootEventselector in simplified slimtree mode
+    return cmdLineUtils.rootEventselector(sourceList, destFileName, destPathSplit, \
+                                          compress=optDict["compress"], recreate=optDict["recreate"], \
+                                          first=0, last=-1, \
+                                          selectionString="", \
+                                          branchinclude=optDict["branchinclude"],\
+                                          branchexclude=optDict["branchexclude"])
+
+sys.exit(execute())