File indexing completed on 2024-04-06 12:23:30
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
0072 self.vars[varName] = getattr(ROOT,type)()
0073 if type in [ "TLorentzVector" ]:
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:
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" ]:
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)
