Getting up-to-date

core/python/LooperBase.py

@@ -34,7 +34,7 @@ def __init__(self, variables):
 
         # Internal state for running
         self.position = -1
-        self.eList = None
+        self._eList = None
 
         self.classUUIDs = []
 

core/python/Sample.py

@@ -27,7 +27,7 @@ def new_name():
 
 def check_equal_(vals):
     if not len(set(vals)) == 1:
-        raise ValueError( "These values should be identical but are not: %r"%vals )
+        raise ValueError( "Sample combine check failed on: %r"%vals )
     else:
         return vals[0]
 
@@ -404,6 +404,28 @@ def fromCMGCrabDirectory(cls, name, baseDirectory, treeFilename = 'tree.root', t
                 selectionString = selectionString, weightString = weightString, 
                 isData = isData, color = color, texName = texName )
 
+    def split( self, n, clear = True):
+        ''' Split sample into n sub-samples
+        '''
+        if not n>=1:
+            raise ValueError( "Can not split into: '%r'" % n )
+
+        chunks = helpers.partition( self.files, min(n , len(self.files) ) ) 
+
+        if clear: self.clear() # Kill yourself.
+
+        return [ Sample( 
+                name            = self.name+"_%i" % n_sample, 
+                treeName        = self.treeName, 
+                files           = chunks[n_sample], 
+                normalization   = self.normalization, 
+                selectionString = self.selectionString, 
+                weightString    = self.weightString, 
+                isData          = self.isData, 
+                color           = self.color, 
+                texName         = self.texName ) for n_sample in xrange(len(chunks)) ]
+
+
     # Handle loading of chain -> load it when first used 
     @property
     def chain(self):
@@ -430,7 +452,7 @@ def __loadChain(self):
                     else:
                         logger.error( "Check of root file failed. Skipping. File: %s", f )
                 except IOError as e:
-                    logger.warning( "Could not load file %s", f )
+                    logger.error( "Could not load file %s", f )
                     raise e
 
             logger.debug( "Loaded %i files for sample '%s'.", counter, self.name )
@@ -525,23 +547,18 @@ def combineWithSampleWeight(self, weightString):
                 self.name, weightString )
             return weightString
 
-    def getEventList(self, selectionString=None, name=None):
+    def getEventList(self, selectionString=None):
         ''' Get a TEventList from a selectionString (combined with self.selectionString, if exists).
         '''
 
         selectionString_ = self.combineWithSampleSelection( selectionString )
 
         tmp=str(uuid.uuid4())
-        logger.debug( "Making eList for sample %s and selectionString %s", self.name, selectionString_ )
+        logger.debug( "Making event list for sample %s and selectionString %s", self.name, selectionString_ )
         self.chain.Draw('>>'+tmp, selectionString_ if selectionString else "(1)")
         elistTMP_t = ROOT.gDirectory.Get(tmp)
 
-        if not name:
-            return elistTMP_t
-        else:
-            elistTMP = elistTMP_t.Clone(name)
-            del elistTMP_t
-            return elistTMP
+        return elistTMP_t
 
     def getYieldFromDraw(self, selectionString = None, weightString = None):
         ''' Get yield from self.chain according to a selectionString and a weightString

core/python/TreeMaker.py

@@ -33,6 +33,12 @@ def __init__(self, variables, sequence = [], treeName = "Events"):
         # function to fill the event 
         self.sequence = sequence
 
+    def debugBranchAddresses( self, prefix = ""):
+        ''' If strings are empty, there is an issue with memory. Only for debugging purposes.
+        '''
+        for b in self.branches:
+            print prefix, b.GetName(), repr(b.GetAddress())
+
     def cloneWithoutCompile(self, externalTree = None):
         ''' make a deep copy of self to e.g. avoid re-compilation of class in a loop. 
             Reset TTree as to not create a memory leak.
@@ -54,6 +60,7 @@ def cloneWithoutCompile(self, externalTree = None):
             res.tree = ROOT.TTree( treeName, treeName )
 
         res.makeBranches()
+        #self.debugBranchAddresses(prefix = "cloneWithoutCompile")
 
         return res
 
@@ -81,14 +88,15 @@ def makeBranches(self):
                 vectorCount+=1
             else:
                 raise ValueError( "Don't know what variable %r is." % s )
-
+        #self.debugBranchAddresses(prefix = "makeBranches")
         logger.debug( "TreeMaker created %i new scalars and %i new vectors.", scalerCount, vectorCount )
 
     def clear(self):
         if self.tree: self.tree.IsA().Destructor( self.tree )
 
     def fill(self):
         # Write to TTree
+        #self.debugBranchAddresses( prefix = "Filling")
         if self.treeIsExternal:
             for b in self.branches:
                 b.Fill()

core/python/TreeReader.py

@@ -45,8 +45,8 @@ def __init__(self, sample, variables=[], sequence = [], selectionString = None,
         # Whether all branches are to be read or whether that information should come from the variables
         self.allBranchesActive = allBranchesActive
 
-        # Read branch information from the chain
-        self.readLeafInfo()
+        ## Read branch information from the chain -> Will be useful when using auto-type
+        #self.readLeafInfo()
 
         # 'variables' are read from the chain
         super(TreeReader, self).__init__( variables = list(set(variables)) ) 
@@ -59,18 +59,19 @@ def __init__(self, sample, variables=[], sequence = [], selectionString = None,
 
         # set the addresses of the branches corresponding to 'variables'
         self.setAddresses()
-        
+
         # make eList from cutString
         # Turn on everything for flexibility with the selectionString
         logger.debug("Initializing TreeReader for sample %s", self.sample.name)
         self.activateAllBranches()
-        self.eList = self.sample.getEventList(selectionString = self.selectionString) if self.selectionString is not None else None
+        self._eList = self.sample.getEventList(selectionString = self.selectionString) if self.selectionString is not None else None
         self.activateBranches()
-        self.nEvents = self.eList.GetN() if  self.eList else self.sample.chain.GetEntries()
+        self.nEvents = self._eList.GetN() if  self._eList else self.sample.chain.GetEntries()
         logger.debug("Found %i events in  %s", self.nEvents, self.sample.name)
 
         #  default event range of the reader
         self.eventRange = (0, self.nEvents)
+
 
     def setAddresses(self):
         ''' Set all the branch addresses to the members in the class instance
@@ -84,26 +85,25 @@ def setAddresses(self):
                     self.sample.chain.SetBranchAddress(comp.name, ROOT.AddressOf(self.event, comp.name ))
             else:
                 raise ValueError( "Don't know what variable %r is." % s )
-
 
     def cloneTree(self, branchList = [], newTreename = None, rootfile = None):
         '''Clone tree after preselection and event range
         '''
         selectionString = self.selectionString if self.selectionString is not None else "1"
-        if self.eList:
+        if self._eList:
             # If there is an eList, first restrict it to the event range, then clone
+            import uuid
+            name = str(uuid.uuid4())
             list_to_copy = ROOT.TEventList("tmp","tmp")
             for i_ev in xrange(*self.eventRange):
-                list_to_copy.Enter(self.eList.GetEntry(i_ev))
+                list_to_copy.Enter(self._eList.GetEntry(i_ev))
+
+            self.sample.chain.GetEntry(list_to_copy.GetEntry(0)) #This is needed to keep branch addresses when running over a few events and >=1 file
 
             # activate branches that we want to copy, disable the ones we only need for reading
             self.activateBranches( turnOnReadBranches = False, branchList = branchList )
-            # preserving current event list
-            tmpEventList = self.sample.chain.GetEventList()
-            tmpEventList = 0 if not self.sample.chain.GetEventList() else tmpEventList
 
             # Copy only the selected events
-            self.sample.chain.SetEventList( list_to_copy ) 
 
             # Create the new tree in a file (if there is one)
             tmp_directory = ROOT.gDirectory
@@ -113,37 +113,19 @@ def cloneTree(self, branchList = [], newTreename = None, rootfile = None):
 
             # Copying tree
 
-            # 1. causes trouble with few events and few files. 
-            #res =  self.sample.chain.CopyTree( "(1)", "" )
-
-            # 2. Doesn't take into account event list
-            #res =  self.sample.chain.CloneTree( 0 )
-            #res.CopyEntries( self.sample.chain )
-
-            # 3. Loop (no problems observed)
             logger.debug("Copying %i events in a loop.", list_to_copy.GetN())
-            tree = self.sample.chain.GetTree()
-            res =  self.sample.chain.GetTree().CloneTree( 0 )
+            res = self.sample.chain.GetTree().CloneTree( 0 )
             for i_event in xrange(list_to_copy.GetN()):
-                tree.GetEntry( list_to_copy.GetEntry(i_event) )
+                self.sample.chain.GetEntry( list_to_copy.GetEntry(i_event) )
                 res.Fill()
-            # Needed?
-            res.Write()
 
-            ## 4. Doesn't take into account event list 
-            #tree =  self.sample.chain.GetTree()
-            #tree.SetEventList( tmpEventList )
-            #res = tree.CloneTree( 0 )
-            #res.CopyEntries( tree )
+            res.Write()
 
             logger.debug("Number of events: list_to_copy %i res.GetEntries() %i", list_to_copy.GetN(), res.GetEntries())
 
             # Change back to previous gDirectory
             tmp_directory.cd()
 
-            # restoring event list
-            self.sample.chain.SetEventList( tmpEventList ) 
-
             # activate what we read, don't activate the ones we just copied
             self.activateBranches( turnOnReadBranches = True, branchList = [] )
 
@@ -192,37 +174,32 @@ def activateAllBranches(self):
         '''
         self.sample.chain.SetBranchStatus("*", 1)
 
-    def readLeafInfo(self):
-        ''' Read information on the leaves from the chain and store in dict
-            FIXME: Pretty sure this doesn't yet work for fixed sized vectors
-        '''
-        leafInfo = []
-        for s in self.sample.chain.GetListOfLeaves():
-            leaf = {'type':shortTypeDict[s.GetTypeName()], 'name':s.GetName()}
-            countval = array.array('i',[-9999])
-            pointer = s.GetLeafCounter(countval)
-            leaf['dim'] = 'scalar'
-            if pointer:
-                # Vector
-                # GetLeafCounter returns countval==1 if a counter leaf was found and pointer points to that leaf
-                # https://root.cern.ch/doc/master/classTLeaf.html#a062bd89a11fd1f922c096e66d2601ab6
-                if countval[0]==1:
-                    leaf['counterInt'] = pointer.GetName()
-                leaf['dim'] = 'vector'
-            else:
-                # For fixed size arrays, the pointer is zero and the countval is the size of the array
-                leaf['nMax'] = countval[0]
-            leafInfo.append(leaf)
-        return leafInfo
+#    def readLeafInfo(self):
+#        ''' Read information on the leaves from the chain and store in dict
+#            FIXME: Pretty sure this doesn't yet work for fixed sized vectors
+#        '''
+#        leafInfo = []
+#        for s in self.sample.chain.GetListOfLeaves():
+#            leaf = {'type':shortTypeDict[s.GetTypeName()], 'name':s.GetName()}
+#            countval = array.array('i',[-9999])
+#            pointer = s.GetLeafCounter(countval)
+#            leaf['dim'] = 'scalar'
+#            if pointer:
+#                # Vector
+#                # GetLeafCounter returns countval==1 if a counter leaf was found and pointer points to that leaf
+#                # https://root.cern.ch/doc/master/classTLeaf.html#a062bd89a11fd1f922c096e66d2601ab6
+#                if countval[0]==1:
+#                    leaf['counterInt'] = pointer.GetName()
+#                leaf['dim'] = 'vector'
+#            else:
+#                # For fixed size arrays, the pointer is zero and the countval is the size of the array
+#                leaf['nMax'] = countval[0]
+#            leafInfo.append(leaf)
+#        return leafInfo
 
     def getEventRanges(self, maxFileSizeMB = None, maxNEvents = None, nJobs = None, minJobs = None):
         '''For convinience: Define splitting of sample according to various criteria
         '''
-        def chunks(l, n):
-            """Yield successive n-sized chunks from l."""
-            for i in xrange(0, len(l), n):
-                yield (i, i+n)
-
         if maxFileSizeMB is not None:
             nSplit = sum( os.path.getsize(f) for f in self.sample.files ) / ( 1024**2*maxFileSizeMB )
         elif maxNEvents is not None:
@@ -247,14 +224,14 @@ def setEventRange( self, evtRange ):
         self.eventRange = ( max(0, evtRange[0]), min( self.nEvents, evtRange[1]) ) 
         logger.debug( "[setEventRange] Set eventRange %r (was: %r) for reader of sample %s", self.eventRange, old_eventRange, self.sample.name )
 
-    def setEventList( self, evtList ):
-        ''' Specify an event list that the reader will run over. 
-        '''
-        self.sample.chain.SetEventList( evtList ) 
-        self.eList = evtList 
-        self.nEvents = self.eList.GetN()
-        self.eventRange = (0, self.nEvents)
-        logger.debug( "[setEventList] Set eventRange %r for reader of sample %s", self.eventRange, self.sample.name )
+    #def setEventList( self, evtList ):
+    #    ''' Specify an event list that the reader will run over. 
+    #    '''
+    #    self.sample.chain.SetEventList( evtList ) 
+    #    self.eList = evtList 
+    #    self.nEvents = self.eList.GetN()
+    #    self.eventRange = (0, self.nEvents)
+    #    logger.debug( "[setEventList] Set eventRange %r for reader of sample %s", self.eventRange, self.sample.name )
 
     def reduceEventRange( self, reduction_factor ):
         ''' Reduce event range by a given factor. 
@@ -294,7 +271,7 @@ def _execute(self):
         # get entry 
         errorLevel = ROOT.gErrorIgnoreLevel
         ROOT.gErrorIgnoreLevel = 3000
-        self.sample.chain.GetEntry ( self.eList.GetEntry( self.position ) ) if self.eList else self.sample.chain.GetEntry( self.position )
+        self.sample.chain.GetEntry ( self._eList.GetEntry( self.position ) ) if self._eList else self.sample.chain.GetEntry( self.position )
         ROOT.gErrorIgnoreLevel = errorLevel
 
         # sequence

core/python/helpers.py

@@ -5,11 +5,19 @@
 import logging
 logger = logging.getLogger(__name__)
 
+# Error handling
 class EmptySampleError(Exception):
     '''Accessing a sample without ROOT files.
     '''
     pass
 
+# List helper
+def partition(lst, n):
+    ''' Partition list into chunks of approximately equal size'''
+    # http://stackoverflow.com/questions/2659900/python-slicing-a-list-into-n-nearly-equal-length-partitions
+    n_division = len(lst) / float(n)
+    return [ lst[int(round(n_division * i)): int(round(n_division * (i + 1)))] for i in xrange(n) ]
+
 # Translation of short types to ROOT C types
 cStringTypeDict = {
     'b': 'UChar_t',
@@ -49,7 +57,6 @@ def decorate(func):
         return func
     return decorate
 
-
 def checkRootFile(f, checkForObjects=[] ):
     ''' Checks whether a root file exists, was not recoverd or otherwise broken and
     contains the objects in 'checkForObjects'

examples/example_plot.py

@@ -51,7 +51,7 @@
 stack = Stack( [ s0, s1], [ s2 ] )
 
 # Let's use a trivial weight. All functions will 
-plot_weight   = lambda data:1
+plot_weight   = lambda event, sample : 1
 
 # Two selection strings
 selectionString = "nJet>0"

fwlite/python/FWLiteSample.py

@@ -85,11 +85,10 @@ def fromDirectory(cls, name, directory, color = 0, texName = None, maxN = None):
         maxN = maxN if maxN is not None and maxN>0 else None
         files = files[:maxN]
 
-        sample =  cls(name = name, files = files, color=color, texName = texName)
-        return sample
+        return cls(name = name, files = files, color=color, texName = texName)
 
     @classmethod
-    def fromDAS(cls, name, dataset, instance = 'global', prefix='root://cms-xrd-global.cern.ch/', maxN = None):
+    def fromDAS(cls, name, dataset, instance = 'global', prefix='root://cms-xrd-global.cern.ch/', texName = None, maxN = None):
         ''' Make sample from DAS. 
         '''
         # https://github.com/CERN-PH-CMG/cmg-cmssw/blob/0f1d3bf62e7ec91c2e249af1555644b7f414ab50/CMGTools/Production/python/dataset.py#L437
@@ -119,8 +118,24 @@ def _dasPopen(dbs):
             line = line.rstrip()
             files.append(prefix+line)
 
-        return cls(name, files=files)
+        return cls(name, files=files, texName = texName)
+
+    @classmethod
+    def combine(cls, name, samples, texName = None, maxN = None, color = 0):
+        '''Make new sample from a list of samples.
+        '''
+        if not (type(samples) in [type([]), type(())]) or len(samples)<1:
+            raise ValueError( "Need non-empty list of samples. Got %r"% samples)
+
+        files = sum([s.files for s in samples], [])
+        maxN = maxN if maxN is not None and maxN>0 else None
+        files = files[:maxN]
 
+        return cls(name = name, \
+                   files = files,
+                   color = color, 
+                   texName = texName
+            )
 
     def fwliteReader(self, **kwargs):
         ''' Return a FWLiteReader class for the sample