File indexing completed on 2023-03-17 11:15:45
0001
0002
0003
0004
0005
0006
0007
0008
0009 from builtins import range
0010 class NTupleVariable:
0011 """Branch containing an individual variable (either of the event or of an object), created with a name and a function to compute it
0012 - name, type, help, default: obvious
0013 - function: a function that taken an object computes the value to fill (e.g. lambda event : len(event.goodVertices))
0014 """
0015 def __init__(self, name, function, type=float, help="", default=-99, mcOnly=False, filler=None):
0016 self.name = name
0017 self.function = function
0018 self.type = type
0019 self.help = help
0020 self.default = default
0021 self.mcOnly = mcOnly
0022 self.filler = filler
0023 def __call__(self,object):
0024 ret = self.function(object)
0025 return ret
0026 def makeBranch(self,treeNumpy,isMC):
0027 if self.mcOnly and not isMC: return
0028 treeNumpy.var(self.name, type=self.type, default=self.default, title=self.help, filler=self.filler)
0029 def fillBranch(self,treeNumpy,object,isMC):
0030 if self.mcOnly and not isMC: return
0031 treeNumpy.fill(self.name, self(object))
0032 def __repr__(self):
0033 return "<NTupleVariable[%s]>" % self.name
0034
0035
0036 class NTupleObjectType:
0037 """Type defining a collection of variables associated to a single object. Contans NTupleVariable and NTupleSubObject"""
0038 def __init__(self,name,baseObjectTypes=[],mcOnly=[],variables=[]):
0039 self.name = name
0040 self.baseObjectTypes = baseObjectTypes
0041 self.mcOnly = mcOnly
0042 self.variables = []
0043 self.subObjects = []
0044 for v in variables:
0045 if issubclass(v.__class__,NTupleSubObject):
0046 self.subObjects.append(v)
0047 else:
0048 self.variables.append(v)
0049 self._subObjectVars = {}
0050 def ownVars(self,isMC):
0051 """Return only my vars, not including the ones from the bases"""
0052 vars = [ v for v in self.variables if (isMC or not v.mcOnly) ]
0053 if self.subObjects:
0054 if isMC not in self._subObjectVars:
0055 subvars = []
0056 for so in self.subObjects:
0057 if so.mcOnly and not isMC: continue
0058 for subvar in so.objectType.allVars(isMC):
0059 subvars.append(NTupleVariable(so.name+"_"+subvar.name,
0060
0061 lambda object, subvar=subvar, so=so : subvar(so(object)),
0062
0063
0064 type = subvar.type, help = subvar.help, default = subvar.default, mcOnly = subvar.mcOnly,
0065 filler = subvar.filler))
0066 self._subObjectVars[isMC] = subvars
0067 vars += self._subObjectVars[isMC]
0068 return vars
0069 def allVars(self,isMC):
0070 """Return all vars, including the base ones. Duplicate bases are not added twice"""
0071 ret = []; names = {}
0072 if not isMC and self.mcOnly: return []
0073 for base in self.allBases():
0074 if not isMC and base.mcOnly: continue
0075 for var in base.ownVars(isMC):
0076 if var.name in names: raise RuntimeError("Duplicate definition of variable %s from %s and %s" % (var.name, base.name, names[var.name]))
0077 names[var.name] = base.name
0078 ret.append(var)
0079 for var in self.ownVars(isMC):
0080 if var.name in names: raise RuntimeError("Duplicate definition of variable %s from %s and %s" % (var.name, self.name, names[var.name]))
0081 names[var.name] = self.name
0082 ret.append(var)
0083 return ret
0084 def allBases(self):
0085 ret = []
0086 for b in self.baseObjectTypes:
0087 if b not in ret:
0088 ret.append(b)
0089 for b2 in b.allBases():
0090 if b2 not in ret:
0091 ret.append(b2)
0092 return ret
0093 def addVariables(self,newvars):
0094 currentnames = [v.name for v in self.allVars(True)]
0095 uniquenewvars = []
0096 for var in newvars:
0097 if var.name in uniquenewvars: raise RuntimeError("Duplicate definition of variable %s while adding variables to object type %s" % (var.name,self.name))
0098 uniquenewvars.append(var.name)
0099 if var.name not in currentnames:
0100 self.variables.append(var)
0101 else:
0102 raise RuntimeError("Variable %s is already present in object type %s" % (var.name,self.name))
0103 def addSubObjects(self,sos):
0104 currentnames = [v.name for v in self.subObjects]
0105 uniquenewobjs = []
0106 for ob in sos:
0107 if ob.name in uniquenewobjs: raise RuntimeError("Duplicate definition of sub-object %s while adding it to object type %s" % (ob.name,self.name))
0108 uniquenewobjs.append(ob.name)
0109 if ob.name not in currentnames:
0110 self.subObjects.append(ob)
0111 else:
0112 raise RuntimeError("Sub-object %s is already present in object type %s" % (ob.name,self.name))
0113 self._subObjectVars.clear()
0114 mynewvars = self.allVars(True)
0115 mynewvars = self.allVars(False)
0116 def removeVariable(self,name):
0117 self.variables = [ v for v in self.variables if v.name != name]
0118 def __repr__(self):
0119 return "<NTupleObjectType[%s]>" % self.name
0120
0121
0122
0123
0124 class NTupleSubObject:
0125 """Type to add a sub-object within an NTupleObjectType, given a name (used as prefix), a function to extract the sub-object and NTupleObjectType to define tye type"""
0126 def __init__(self,name,function,objectType,mcOnly=False):
0127 self.name = name
0128 self.function = function
0129 self.objectType = objectType
0130 self.mcOnly = mcOnly
0131 def __call__(self,object):
0132 return self.function(object)
0133
0134 class NTupleObject:
0135 """Type defining a set of branches associated to a single object (i.e. an instance of NTupleObjectType)"""
0136 def __init__(self, name, objectType, help="", mcOnly=False):
0137 self.name = name
0138 self.objectType = objectType
0139 self.mcOnly = mcOnly
0140 self.help = ""
0141 def makeBranches(self,treeNumpy,isMC):
0142 if not isMC and self.mcOnly: return
0143 allvars = self.objectType.allVars(isMC)
0144 for v in allvars:
0145 h = v.help
0146 if self.help: h = "%s for %s" % ( h if h else v.name, self.help )
0147 treeNumpy.var("%s_%s" % (self.name, v.name), type=v.type, default=v.default, title=h, filler=v.filler)
0148 def fillBranches(self,treeNumpy,object,isMC):
0149 if self.mcOnly and not isMC: return
0150 allvars = self.objectType.allVars(isMC)
0151 for v in allvars:
0152 treeNumpy.fill("%s_%s" % (self.name, v.name), v(object))
0153 def __repr__(self):
0154 return "<NTupleObject[%s]>" % self.name
0155
0156
0157 class NTupleCollection:
0158 """Type defining a set of branches associated to a list of objects (i.e. an instance of NTupleObjectType)"""
0159 def __init__(self, name, objectType, maxlen, help="", mcOnly=False, sortAscendingBy=None, sortDescendingBy=None, filter=None):
0160 self.name = name
0161 self.objectType = objectType
0162 self.maxlen = maxlen
0163 self.help = help
0164 if objectType.mcOnly and mcOnly == False:
0165
0166 mcOnly = True
0167 self.mcOnly = mcOnly
0168 if sortAscendingBy != None and sortDescendingBy != None:
0169 raise RuntimeError("Cannot specify two sort conditions")
0170 self.filter = filter
0171 self.sortAscendingBy = sortAscendingBy
0172 self.sortDescendingBy = sortDescendingBy
0173 def makeBranchesScalar(self,treeNumpy,isMC):
0174 if not isMC and self.objectType.mcOnly: return
0175 treeNumpy.var("n"+self.name, int)
0176 allvars = self.objectType.allVars(isMC)
0177 for v in allvars:
0178 for i in range(1,self.maxlen+1):
0179 h = v.help
0180 if self.help: h = "%s for %s [%d]" % ( h if h else v.name, self.help, i-1 )
0181 treeNumpy.var("%s%d_%s" % (self.name, i, v.name), type=v.type, default=v.default, title=h, filler=v.filler)
0182 def makeBranchesVector(self,treeNumpy,isMC):
0183 if not isMC and self.objectType.mcOnly: return
0184 treeNumpy.var("n"+self.name, int)
0185 allvars = self.objectType.allVars(isMC)
0186 for v in allvars:
0187 h = v.help
0188 if self.help: h = "%s for %s" % ( h if h else v.name, self.help )
0189 name="%s_%s" % (self.name, v.name) if v.name != "" else self.name
0190 treeNumpy.vector(name, "n"+self.name, self.maxlen, type=v.type, default=v.default, title=h, filler=v.filler)
0191 def fillBranchesScalar(self,treeNumpy,collection,isMC):
0192 if not isMC and self.objectType.mcOnly: return
0193 if self.filter != None: collection = [ o for o in collection if self.filter(o) ]
0194 if self.sortAscendingBy != None: collection = sorted(collection, key=self.sortAscendingBy)
0195 if self.sortDescendingBy != None: collection = sorted(collection, key=self.sortDescendingBy, reverse=True)
0196 num = min(self.maxlen,len(collection))
0197 treeNumpy.fill("n"+self.name, num)
0198 allvars = self.objectType.allVars(isMC)
0199 for i in range(num):
0200 o = collection[i]
0201 for v in allvars:
0202 treeNumpy.fill("%s%d_%s" % (self.name, i+1, v.name), v(o))
0203 def fillBranchesVector(self,treeNumpy,collection,isMC):
0204 if not isMC and self.objectType.mcOnly: return
0205 if self.filter != None: collection = [ o for o in collection if self.filter(o) ]
0206 if self.sortAscendingBy != None: collection = sorted(collection, key=self.sortAscendingBy)
0207 if self.sortDescendingBy != None: collection = sorted(collection, key=self.sortDescendingBy, reverse=True)
0208 num = min(self.maxlen,len(collection))
0209 treeNumpy.fill("n"+self.name, num)
0210 allvars = self.objectType.allVars(isMC)
0211 for v in allvars:
0212 name="%s_%s" % (self.name, v.name) if v.name != "" else self.name
0213 treeNumpy.vfill(name, [ v(collection[i]) for i in range(num) ])
0214 def __repr__(self):
0215 return "<NTupleCollection[%s]>" % self.name
0216
0217 def get_cpp_declaration(self, isMC):
0218 s = []
0219 for v in self.objectType.allVars(isMC):
0220 s += ["{0} {1}__{2}[{3}];".format(v.type.__name__, self.name, v.name, self.maxlen)]
0221 return "\n".join(s)
0222
0223 def get_cpp_wrapper_class(self, isMC):
0224 s = "class %s {\n" % self.name
0225 s += "public:\n"
0226 for v in self.objectType.allVars(isMC):
0227 s += " {0} {1};\n".format(v.type.__name__, v.name)
0228 s += "};\n"
0229 return s
0230
0231 def get_py_wrapper_class(self, isMC):
0232 s = "class %s:\n" % self.name
0233 s += " def __init__(self, tree, n):\n"
0234 for v in self.objectType.allVars(isMC):
0235 if len(v.name)>0:
0236 s += " self.{0} = tree.{1}_{2}[n];\n".format(v.name, self.name, v.name)
0237 else:
0238 s += " self.{0} = tree.{0}[n];\n".format(self.name)
0239
0240 s += " @staticmethod\n"
0241 s += " def make_array(event):\n"
0242 s += " return [{0}(event.input, i) for i in range(event.input.n{0})]\n".format(self.name)
0243 return s
0244
0245