Back to home page

Project CMSSW displayed by LXR
 
 

    

File indexing completed on 2023-03-17 11:15:45 

0001 #!/bin/env python
0002 #
0003 # Objects to be used with AutoFillTreeProducer
0004 # 
0005 # the variable declaration contains both the booking information and a function to fill the variable
0006 #
0007 # TODO: more documentation needed here!
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                                   #DebugComposer(so,subvar),#lambda object : subvar(so(object)),
0061                                   lambda object, subvar=subvar, so=so : subvar(so(object)), 
0062                                   # ^-- lambda object : subvar(so(object)) doesn't work due to scoping, see
0063                                   #     http://stackoverflow.com/questions/2295290/what-do-lambda-function-closures-capture-in-python/2295372#2295372
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)] # require no conflict with all variables, including mcOnly ones
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() # clear and update cache of subobj variables
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             #print "collection %s is set to mcOnly since the type %s is mcOnly" % (name, objectType.name)
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
[ Source navigation ] [ Diff markup ] [ Identifier search ] [ general search ]

This page was automatically generated by the 2.2.1 LXR engine.

The LXR team
Valid CSS 2.1! Valid HTML 4.01!