Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2021-07-23 02:25:52

0001 import numpy
0002 from ROOT import TTree
0003 import ROOT
0004 
0005 class Tree(object):
0006     
0007     def __init__(self, name, title, defaultFloatType="D", defaultIntType="I"):
0008         self.vars = {}
0009         self.vecvars = {}
0010         self.tree = TTree(name, title)
0011         self.defaults = {}
0012         self.vecdefaults = {}
0013         self.defaultFloatType = defaultFloatType
0014         self.defaultIntType = defaultIntType
0015         self.fillers = {}
0016 
0017     def setDefaultFloatType(self, defaultFloatType):
0018         self.defaultFloatType = defaultFloatType
0019 
0020     def setDefaultIntType(self, defaultIntType):
0021         self.defaultIntType = defaultIntType
0022         
0023     def copyStructure(self, tree):
0024         for branch in tree.GetListOfBranches():
0025             name = branch.GetName() 
0026             typeName = branch.GetListOfLeaves()[0].GetTypeName()
0027             type = float
0028             if typeName == 'Int_t':
0029                 type = int
0030             self.var(name, type)            
0031     
0032     def branch_(self, selfmap, varName, type, len, postfix="", storageType="default", title=None):
0033         """Backend function used to create scalar and vector branches. 
0034            Users should call "var" and "vector", not this function directly."""
0035         if storageType == "default": 
0036             storageType = self.defaultIntType if type is int else self.defaultFloatType
0037         if type is float  :
0038             if storageType == "F": 
0039                 selfmap[varName]=numpy.zeros(len,numpy.float32)
0040                 self.tree.Branch(varName,selfmap[varName],varName+postfix+'/F')
0041             elif storageType == "D":
0042                 selfmap[varName]=numpy.zeros(len,numpy.float64)
0043                 self.tree.Branch(varName,selfmap[varName],varName+postfix+'/D')
0044             else:
0045                 raise RuntimeError('Unknown storage type %s for branch %s' % (storageType, varName))
0046         elif type is int: 
0047             dtypes = {
0048                 "i" : numpy.uint32,
0049                 "s" : numpy.uint16,
0050                 "b" : numpy.uint8,
0051                 "l" : numpy.uint64,
0052                 "I" : numpy.int32,
0053                 "S" : numpy.int16,
0054                 "B" : numpy.int8,
0055                 "L" : numpy.int64,
0056             }
0057             if storageType not in dtypes: 
0058                 raise RuntimeError('Unknown storage type %s for branch %s' % (storageType, varName))
0059             selfmap[varName]=numpy.zeros(len,dtypes[storageType])
0060             self.tree.Branch(varName,selfmap[varName],varName+postfix+'/'+storageType)
0061         else:
0062             raise RuntimeError('Unknown type %s for branch %s' % (type, varName))
0063         if title:
0064             self.tree.GetBranch(varName).SetTitle(title)
0065 
0066     def var(self, varName,type=float, default=-99, title=None, storageType="default", filler=None ):
0067         if type in [int, float]:
0068             self.branch_(self.vars, varName, type, 1, title=title, storageType=storageType)
0069             self.defaults[varName] = default
0070         elif __builtins__['type'](type) == str:
0071             # create a value, looking up the type from ROOT and calling the default constructor
0072             self.vars[varName] = getattr(ROOT,type)()
0073             if type in [ "TLorentzVector" ]: # custom streamer classes
0074                 self.tree.Branch(varName+".", type, self.vars[varName], 8000,-1)
0075             else:
0076                 self.tree.Branch(varName+".", type, self.vars[varName])
0077             if filler is None:
0078                 raise RuntimeError("Error: when brancing with an object, filler should be set to a function that takes as argument an object instance and a value, and set the instance to the value (as otherwise python assignment of objects changes the address as well)")
0079             self.fillers[varName] = filler
0080         else:
0081             raise RuntimeError('Unknown type %s for branch %s: it is not int, float or a string' % (type, varName))
0082         self.defaults[varName] = default
0083 
0084     def vector(self, varName, lenvar, maxlen=None, type=float, default=-99, title=None, storageType="default", filler=None ):
0085         """either lenvar is a string, and maxlen an int (variable size array), or lenvar is an int and maxlen is not specified (fixed array)"""
0086         if type in [int, float]:
0087             if __builtins__['type'](lenvar) == int:  # need the __builtins__ since 'type' is a variable here :-/
0088                 self.branch_(self.vecvars, varName, type, lenvar, postfix="[%d]" % lenvar, title=title, storageType=storageType)
0089             else:
0090                 if maxlen == None: RuntimeError, 'You must specify a maxlen if making a dynamic array';
0091                 self.branch_(self.vecvars, varName, type, maxlen, postfix="[%s]" % lenvar, title=title, storageType=storageType)
0092         elif __builtins__['type'](type) == str:
0093             self.vecvars[varName] = ROOT.TClonesArray(type,(lenvar if __builtins__['type'](lenvar) == int else maxlen))
0094             if type in [ "TLorentzVector" ]: # custom streamer classes
0095                 self.tree.Branch(varName+".", self.vecvars[varName], 32000, -1)
0096             else:
0097                 self.tree.Branch(varName+".", self.vecvars[varName])
0098             if filler is None:
0099                 raise RuntimeError("Error: when brancing with an object, filler should be set to a function that takes as argument an object instance and a value, and set the instance to the value (as otherwise python assignment of objects changes the address as well)")
0100             self.fillers[varName] = filler
0101         self.vecdefaults[varName] = default
0102 
0103     def reset(self):
0104         for name,value in self.vars.items():
0105             if name in self.fillers:
0106                 self.fillers[name](value, self.defaults[name])
0107             else:
0108                 value[0]=self.defaults[name]
0109         for name,value in self.vecvars.items():
0110             if isinstance(value, numpy.ndarray):
0111                 value.fill(self.vecdefaults[name])
0112             else:
0113                 if isinstance(value, ROOT.TObject) and value.ClassName() == "TClonesArray":
0114                     value.ExpandCreateFast(0)
0115             
0116     def fill(self, varName, value ):
0117         if isinstance(self.vars[varName], numpy.ndarray):
0118             self.vars[varName][0]=value
0119         else:
0120             self.fillers[varName](self.vars[varName],value)
0121 
0122     def vfill(self, varName, values ):
0123         a = self.vecvars[varName]
0124         if isinstance(a, numpy.ndarray):
0125             for (i,v) in enumerate(values):
0126                 a[i]=v
0127         else:
0128             if isinstance(a, ROOT.TObject) and a.ClassName() == "TClonesArray":
0129                 a.ExpandCreateFast(len(values))
0130             fillit = self.fillers[varName]
0131             for (i,v) in enumerate(values):
0132                 fillit(a[i],v)