00001 import logging
00002 import os.path
00003
00004 from PyQt4.QtGui import QColor
00005
00006 from Vispa.Share.BasicDataAccessor import BasicDataAccessor
00007 from Vispa.Share.RelativeDataAccessor import RelativeDataAccessor
00008 from Vispa.Share.ParticleDataAccessor import ParticleDataAccessor
00009 from Vispa.Plugins.EventBrowser.EventFileAccessor import EventFileAccessor
00010 from Vispa.Main.Exceptions import exception_traceback
00011 from Vispa.Plugins.EdmBrowser.ParticleDataList import defaultParticleDataList
00012
00013 def eq(self,other):
00014 return id(self)==id(other)
00015 def ne(self,other):
00016 return id(self)!=id(other)
00017
00018 def all(container):
00019
00020 if hasattr(container,'GetEntries'):
00021 try:
00022 entries = container.GetEntries()
00023 for entry in xrange(entries):
00024 yield entry
00025 except:
00026 raise cmserror("Looping of %s failed" %container)
00027
00028 elif hasattr(container, 'size'):
00029
00030 if hasattr(container, 'ids'):
00031 container = container.ids()
00032 try:
00033 entries = container.size()
00034 for entry in xrange(entries):
00035 yield container[entry]
00036 except:
00037 pass
00038
00039
00040
00041
00042
00043
00044
00045
00046 class BranchDummy(object):
00047 def __init__(self,branchtuple):
00048 self.branchtuple=branchtuple
00049
00050 class EdmDataAccessor(BasicDataAccessor, RelativeDataAccessor, ParticleDataAccessor, EventFileAccessor):
00051
00052 def __init__(self):
00053 logging.debug(__name__ + ": __init__")
00054
00055 self._dataObjects = []
00056 self._edmLabel={}
00057 self._edmParent={}
00058 self._edmChildren={}
00059 self._edmMotherRelations={}
00060 self._edmDaughterRelations={}
00061 self._edmChildrenObjects={}
00062
00063 self._eventIndex = 0
00064 self._numEvents = 0
00065
00066 self._filename=""
00067 self._branches=[]
00068 self._filteredBranches=[]
00069 self._events=None
00070 self._readOnDemand=True
00071 self._underscore=False
00072 self._filterBranches=True
00073 self.maxLevels=2
00074 self.maxDaughters=1000
00075
00076 def isRead(self,object,levels=1):
00077 if not id(object) in self._edmChildrenObjects.keys():
00078 return False
00079 if levels>1 and id(object) in self._edmChildren.keys():
00080 for child in self._edmChildren[id(object)]:
00081 if not self.isRead(child, levels-1):
00082 return False
00083 return True
00084
00085 def children(self,object):
00086 """ Get children of an object """
00087 if id(object) in self._edmChildren.keys() and self.isRead(object):
00088 return self._edmChildren[id(object)]
00089 else:
00090 return ()
00091
00092 def isContainer(self,object):
00093 """ Get children of an object """
00094 if id(object) in self._edmChildren.keys() and self.isRead(object):
00095 return len(self._edmChildren[id(object)])>0
00096 else:
00097 return True
00098
00099 def motherRelations(self,object):
00100 """ Get motherRelations of an object """
00101 if id(object) in self._edmMotherRelations.keys():
00102 return self._edmMotherRelations[id(object)]
00103 else:
00104 return ()
00105
00106 def daughterRelations(self,object):
00107 """ Get daughterRelations of an object """
00108 if id(object) in self._edmDaughterRelations.keys():
00109 return self._edmDaughterRelations[id(object)]
00110 else:
00111 return ()
00112
00113 def label(self,object):
00114 return self.getShortLabel(object)
00115
00116 def getShortLabel(self,object):
00117 if id(object) in self._edmLabel.keys():
00118 splitlabel=self._edmLabel[id(object)].strip(".").split(".")
00119 return splitlabel[len(splitlabel)-1]
00120 else:
00121 return ""
00122
00123 def getShortLabelWithType(self,object):
00124 return self.getShortLabel(object)+" <"+self.getShortType(object)+">"
00125
00126 def getObjectLabel(self,object):
00127 splitlabel=self._edmLabel[id(object)].strip(".").split(".")
00128 return ".".join(splitlabel[1:-1])
00129
00130 def getType(self,object):
00131 typ=str(object.__class__)
00132 if "\'" in typ:
00133 typ=typ.split("\'")[1]
00134 if "." in typ:
00135 typ=typ.split(".")[len(typ.split("."))-1]
00136 return typ.strip(" ")
00137
00138 def getShortType(self,object):
00139 typ=self.getType(object).split("<")[0].strip(" ")
00140 return typ
00141
00142 def getBranch(self,object):
00143 entry=object
00144 while id(entry) in self._edmParent.keys() and self._edmParent[id(entry)]!=None:
00145 entry=self._edmParent[id(entry)]
00146 return entry
00147
00148 def getDepth(self,object):
00149 entry=object
00150 i=0
00151 while id(entry) in self._edmParent.keys() and self._edmParent[id(entry)]!=None:
00152 entry=self._edmParent[id(entry)]
00153 i+=1
00154 return i
00155
00156 def getObjectProperties(self,object):
00157 """ get all method properties of an object """
00158 objects=[]
00159 for attr in dir(object):
00160 prop=getattr(object,attr)
00161 if not attr.startswith("__") and (self._underscore or attr.strip("_")==attr):
00162 objects+=[(attr,prop)]
00163 return objects
00164
00165 def getObjectRef(self,object):
00166 """ get object and resolve references """
00167 typshort=self.getShortType(object)
00168 ref_types=["edm::Ref","edm::RefProd","edm::RefToBase","edm::RefToBaseProd","edm::Ptr"]
00169 value=object
00170 ref=False
00171 if typshort in ref_types:
00172 try:
00173 if hasattr(object, "isNull") and object.isNull():
00174 value="ERROR: "+self.getType(object)+" object is null"
00175 elif hasattr(object, "isAvailable") and not object.isAvailable():
00176 value="ERROR: "+self.getType(object)+" object is not available"
00177 else:
00178 value=object.get()
00179 if type(value)==type(None):
00180 value="ERROR: Could not get "+self.getType(object)
00181 else:
00182 ref=True
00183 except Exception, message:
00184 value="ERROR: "+str(message)
00185 return value,ref
00186
00187 def getObjectContent(self,object):
00188 """ get string value of a method """
00189 if not callable(object):
00190 return object
00191 else:
00192 typ=""
00193 if not object.__doc__ or str(object.__doc__)=="":
00194 return "ERROR: Empty __doc__ string"
00195 docs=str(object.__doc__).split("\n")
00196 for doc in docs:
00197 parameters=[]
00198 for p in doc[doc.find("(")+1:doc.find(")")].split(","):
00199 if p!="" and not "=" in p:
00200 parameters+=[p]
00201 if len(parameters)!=0:
00202 continue
00203 typestring=doc[:doc.find("(")]
00204 split_typestring=typestring.split(" ")
00205 templates=0
00206 end_typestring=0
00207 for i in reversed(range(len(split_typestring))):
00208 templates+=split_typestring[i].count("<")
00209 templates-=split_typestring[i].count(">")
00210 if templates==0:
00211 end_typestring=i
00212 break
00213 typ=" ".join(split_typestring[:end_typestring])
00214 hidden_types=["iterator","Iterator"]
00215 root_types=["ROOT::"]
00216 if typ=="" or "void" in typ or True in [t in typ for t in hidden_types]:
00217 return None
00218 from ROOT import TClass
00219 if True in [t in typ for t in root_types] and TClass.GetClass(typ)==None:
00220 return "ERROR: Cannot display object of type "+typ
00221 try:
00222 object=object()
00223 value=object
00224 except Exception, message:
00225 value="ERROR: "+str(message)
00226 if "Buffer" in str(type(value)):
00227 return "ERROR: Cannot display object of type "+typ
00228 else:
00229 return value
00230
00231 def isVectorObject(self,object):
00232 typ=self.getShortType(object)
00233 return typ=="list" or typ[-6:].lower()=="vector" or typ[-3:].lower()=="map" or typ[-10:].lower()=="collection" or hasattr(object,"size")
00234
00235 def compareObjects(self,a,b):
00236 same=False
00237 if hasattr(a,"px") and hasattr(a,"py") and hasattr(a,"pz") and hasattr(a,"energy") and \
00238 hasattr(b,"px") and hasattr(b,"py") and hasattr(b,"pz") and hasattr(b,"energy"):
00239 same=a.px()==b.px() and a.py()==b.py() and a.pz()==b.pz() and a.energy()==b.energy()
00240 return same
00241
00242 def getDaughterObjects(self,object):
00243 """ get list of daughter objects from properties """
00244 objects=[]
00245
00246 objectdict={}
00247 hidden_attr=["front","back","IsA","clone","masterClone","masterClonePtr","mother","motherRef","motherPtr","daughter","daughterRef","daughterPtr","is_back_safe"]
00248 broken_attr=[]
00249 for attr1,property1 in self.getObjectProperties(object):
00250 if attr1 in hidden_attr:
00251 pass
00252 elif attr1 in broken_attr:
00253 objectdict[attr1]=("ERROR: Cannot read property",False)
00254 else:
00255 (value,ref)=self.getObjectRef(self.getObjectContent(property1))
00256 if not isinstance(value,type(None)) and (not self.isVectorObject(object) or self._propertyType(value)!=None):
00257 objectdict[attr1]=(value,ref)
00258 for name in sorted(objectdict.keys()):
00259 objects+=[(name,objectdict[name][0],objectdict[name][1],self._propertyType(objectdict[name][0]))]
00260
00261 if self.isVectorObject(object):
00262 n=0
00263 for o in all(object):
00264 (value,ref)=self.getObjectRef(o)
00265 typ=self._propertyType(value)
00266 if typ!=None:
00267 name="["+str(n)+"]"
00268 elif "GenParticle" in str(value):
00269 name=defaultParticleDataList.getNameFromId(value.pdgId())
00270 else:
00271 name=self.getType(value)+" ["+str(n)+"]"
00272 objects+=[(name,value,ref,typ)]
00273 n+=1
00274
00275 for name,mother,ref,propertyType in objects:
00276 if hasattr(mother,"numberOfDaughters") and hasattr(mother,"daughter"):
00277 try:
00278 for n in range(mother.numberOfDaughters()):
00279 daughter=mother.daughter(n)
00280 found=False
00281 for na,da,re,st in objects:
00282 if self.compareObjects(daughter,da):
00283 daughter=da
00284 found=True
00285 if not id(mother) in self._edmDaughterRelations.keys():
00286 self._edmDaughterRelations[id(mother)]=[]
00287 self._edmDaughterRelations[id(mother)]+=[daughter]
00288 if not id(daughter) in self._edmMotherRelations.keys():
00289 self._edmMotherRelations[id(daughter)]=[]
00290 self._edmMotherRelations[id(daughter)]+=[mother]
00291 except Exception, message:
00292 logging.error("Cannot read candidate relations: "+str(message))
00293 return objects
00294
00295 def _propertyType(self,value):
00296 if type(value) in (bool,):
00297 return "Boolean"
00298 elif type(value) in (int, long):
00299 return "Integer"
00300 elif type(value) in (float,):
00301 return "Double"
00302 elif type(value) in (complex,str,unicode):
00303 return "String"
00304 else:
00305 return None
00306
00307 def properties(self,object):
00308 """ Make list of all properties """
00309 logging.debug(__name__ + ": properties: "+self.label(object))
00310 properties=[]
00311
00312 objectproperties={}
00313 objectproperties_sorted=[]
00314 if id(object) in self._edmChildrenObjects.keys():
00315 for name,value,ref,propertyType in self._edmChildrenObjects[id(object)]:
00316 if propertyType!=None:
00317 objectproperties[name]=(value,propertyType)
00318 objectproperties_sorted+=[name]
00319
00320 properties+=[("Category","Object info","")]
00321 shortlabel=self.getShortLabel(object)
00322 properties+=[("String","label",shortlabel)]
00323 properties+=[("String","type",self.getType(object))]
00324 objectlabel=self.getObjectLabel(object)
00325 if objectlabel!="":
00326 properties+=[("String","object",objectlabel)]
00327 branchlabel=self.label(self.getBranch(object))
00328 if shortlabel.strip(".")!=branchlabel.strip("."):
00329 properties+=[("String","branch",branchlabel)]
00330 else:
00331 properties+=[("Category","Branch info","")]
00332 properties+=[("String","Type",branchlabel.split("_")[0])]
00333 properties+=[("String","Label",branchlabel.split("_")[1])]
00334 properties+=[("String","Product",branchlabel.split("_")[2])]
00335 properties+=[("String","Process",branchlabel.split("_")[3])]
00336
00337 for property in ["pdgId","charge","status"]:
00338 if property in objectproperties.keys():
00339 properties+=[(objectproperties[property][1],property,objectproperties[property][0])]
00340 del objectproperties[property]
00341
00342 if "px" in objectproperties.keys():
00343 properties+=[("Category","Vector","")]
00344 for property in ["energy","px","py","pz","mass","pt","eta","phi","p","theta","y","rapidity","et","mt","mtSqr","massSqr"]:
00345 if property in objectproperties.keys():
00346 properties+=[(objectproperties[property][1],property,objectproperties[property][0])]
00347 del objectproperties[property]
00348
00349 if "x" in objectproperties.keys():
00350 properties+=[("Category","Vector","")]
00351 for property in ["x","y","z"]:
00352 if property in objectproperties.keys():
00353 properties+=[(objectproperties[property][1],property,objectproperties[property][0])]
00354 del objectproperties[property]
00355
00356 if False in [str(value[0]).startswith("ERROR") for value in objectproperties.values()]:
00357 properties+=[("Category","Values","")]
00358 for property in objectproperties_sorted:
00359 if property in objectproperties.keys():
00360 if not str(objectproperties[property][0]).startswith("ERROR"):
00361 properties+=[(objectproperties[property][1],property,objectproperties[property][0])]
00362 del objectproperties[property]
00363
00364 if len(objectproperties.keys())>0:
00365 properties+=[("Category","Errors","")]
00366 for property in objectproperties_sorted:
00367 if property in objectproperties.keys():
00368 properties+=[(objectproperties[property][1],property,objectproperties[property][0])]
00369
00370 return tuple(properties)
00371
00372 def readObjectsRecursive(self,mother,label,edmobject,levels=1):
00373 """ read edm objects recursive """
00374 logging.debug(__name__ + ": readObjectsRecursive (levels="+str(levels)+"): "+label)
00375
00376 if not id(edmobject) in self._edmLabel.keys():
00377 if not isinstance(edmobject,(int,float,long,complex,str,unicode,bool)):
00378
00379 try:
00380 type(edmobject).__eq__=eq
00381 type(edmobject).__ne__=ne
00382 except:
00383 pass
00384 self._edmLabel[id(edmobject)]=label
00385 self._edmParent[id(edmobject)]=mother
00386 self._edmChildren[id(edmobject)]=[]
00387 if not id(mother) in self._edmChildren.keys():
00388 self._edmChildren[id(mother)]=[]
00389 self._edmChildren[id(mother)]+=[edmobject]
00390 if levels==0:
00391
00392 return [edmobject],True
00393 else:
00394
00395 return self.readDaughtersRecursive(edmobject,[edmobject],levels)
00396
00397 def readDaughtersRecursive(self,edmobject,objects,levels=1):
00398 """ read daughter objects of an edmobject """
00399 logging.debug(__name__ + ": readDaughtersRecursive (levels="+str(levels)+"): "+str(edmobject))
00400
00401 if not id(edmobject) in self._edmChildrenObjects.keys():
00402 self._edmChildrenObjects[id(edmobject)]=self.getDaughterObjects(edmobject)
00403
00404 ok=True
00405 daughters=self._edmChildrenObjects[id(edmobject)]
00406 i=0
00407 for name,daughter,ref,propertyType in daughters:
00408
00409 if propertyType==None:
00410 if ref:
00411 label="* "+name
00412 else:
00413 label=name
00414 if id(edmobject) in self._edmLabel.keys() and self._edmLabel[id(edmobject)]!="":
00415 label=self._edmLabel[id(edmobject)]+"."+label
00416 (res,ok)=self.readObjectsRecursive(edmobject,label,daughter,levels-1)
00417 objects+=res
00418 i+=1
00419 if i>self.maxDaughters:
00420 logging.warning("Did not read all daughter objects. Maximum is set to "+str(self.maxDaughters)+".")
00421 return objects,False
00422 return objects,ok
00423
00424 def read(self,object,levels=1):
00425 """ reads contents of a branch """
00426 logging.debug(__name__ + ": read")
00427 if isinstance(object,BranchDummy):
00428 if hasattr(object,"product"):
00429 return object.product
00430 if not self._events:
00431 return object
00432 try:
00433 self._events.getByLabel(object.branchtuple[2],object.branchtuple[3],object.branchtuple[4],object.branchtuple[1])
00434 if object.branchtuple[1].isValid():
00435 product=object.branchtuple[1].product()
00436 if not isinstance(product,(int,float,long,complex,str,unicode,bool)):
00437
00438 try:
00439 type(product).__eq__=eq
00440 type(product).__ne__=ne
00441 except:
00442 pass
00443 self._dataObjects.insert(self._dataObjects.index(object),product)
00444 self._dataObjects.remove(object)
00445 self._edmLabel[id(product)]=object.branchtuple[0]
00446 object.product=product
00447 object=product
00448 else:
00449 self._edmChildrenObjects[id(object)]=[("ERROR","ERROR: Branch is not valid.",False,True)]
00450 logging.info("Branch is not valid: "+object.branchtuple[0]+".")
00451 object.invalid=True
00452 return object
00453 except Exception, e:
00454 self._edmChildrenObjects[id(object)]=[("ERROR","ERROR: Unable to read branch : "+str(e),False,True)]
00455 object.unreadable=True
00456 logging.warning("Unable to read branch "+object.branchtuple[0]+" : "+exception_traceback())
00457 return object
00458 if self.isRead(object,levels):
00459 return object
00460 if levels>0:
00461 self.readDaughtersRecursive(object,[],levels)
00462 return object
00463
00464 def goto(self, index):
00465 """ Goto event number index in file.
00466 """
00467 self._eventIndex=index-1
00468 self._edmLabel={}
00469 self._edmChildren={}
00470 self._edmMotherRelations={}
00471 self._edmDaughterRelations={}
00472 self._edmChildrenObjects={}
00473 if self._events:
00474 self._events.to(self._eventIndex)
00475 self._dataObjects=[]
00476 i=0
00477 for branchtuple in self._filteredBranches:
00478 branch=BranchDummy(branchtuple)
00479 self._dataObjects+=[branch]
00480 self._edmLabel[id(branch)]=branchtuple[0]
00481 if not self._readOnDemand:
00482 self.read(branch,self.maxLevels)
00483 i+=1
00484 if self._filterBranches and self._events:
00485 self.setFilterBranches(True)
00486 return True
00487
00488 def eventNumber(self):
00489 return self._eventIndex+1
00490
00491 def numberOfEvents(self):
00492 return self._numEvents
00493
00494 def topLevelObjects(self):
00495 return self._dataObjects
00496
00497 def open(self, filename=None):
00498 """ Open edm file and show first event """
00499 self._filename=filename
00500 self._branches=[]
00501 if os.path.splitext(filename)[1].lower()==".txt":
00502 file = open(filename)
00503 for line in file.readlines():
00504 if "\"" in line:
00505 linecontent=[l.strip(" \n").rstrip(".") for l in line.split("\"")]
00506 self._branches+=[(linecontent[0]+"_"+linecontent[1]+"_"+linecontent[3]+"_"+linecontent[5],None,linecontent[1],linecontent[3],linecontent[5])]
00507 else:
00508 linecontent=line.strip("\n").split(" ")[0].split("_")
00509 if len(linecontent)>3:
00510 self._branches+=[(linecontent[0]+"_"+linecontent[1]+"_"+linecontent[2]+"_"+linecontent[3],None,linecontent[1],linecontent[2],linecontent[3])]
00511 elif os.path.splitext(filename)[1].lower()==".root":
00512 from DataFormats.FWLite import Events, Handle
00513 self._events = Events(self._filename)
00514 self._numEvents=self._events.size()
00515 branches=self._events.object().getBranchDescriptions()
00516 for branch in branches:
00517 try:
00518 branchname=branch.friendlyClassName()+"_"+branch.moduleLabel()+"_"+branch.productInstanceName()+"_"+branch.processName()
00519 handle=Handle(branch.fullClassName())
00520 self._branches+=[(branchname,handle,branch.moduleLabel(),branch.productInstanceName(),branch.processName())]
00521 except Exception, e:
00522 logging.warning("Cannot read branch "+branchname+":"+str(e))
00523 self._branches.sort(lambda x, y: cmp(x[0], y[0]))
00524 self._filteredBranches=self._branches[:]
00525 return self.goto(1)
00526
00527 def particleId(self, object):
00528 charge=self.property(object,"pdgId")
00529 if charge==None:
00530 charge=0
00531 return charge
00532
00533 def isQuark(self, object):
00534 particleId = self.particleId(object)
00535 if not particleId:
00536 return False
00537 return defaultParticleDataList.isQuarkId(particleId)
00538
00539 def isLepton(self, object):
00540 particleId = self.particleId(object)
00541 if not particleId:
00542 return False
00543 return defaultParticleDataList.isLeptonId(particleId)
00544
00545 def isGluon(self, object):
00546 particleId = self.particleId(object)
00547 if not particleId:
00548 return False
00549 return defaultParticleDataList.isGluonId(particleId)
00550
00551 def isBoson(self, object):
00552 particleId = self.particleId(object)
00553 if not particleId:
00554 return False
00555 return defaultParticleDataList.isBosonId(particleId)
00556
00557 def isPhoton(self, object):
00558 particleId = self.particleId(object)
00559 if not particleId:
00560 return False
00561 if not hasattr(defaultParticleDataList,"isPhotonId"):
00562 return False
00563 return defaultParticleDataList.isPhotonId(particleId)
00564
00565 def isHiggs(self, object):
00566 particleId = self.particleId(object)
00567 if not particleId:
00568 return False
00569 if not hasattr(defaultParticleDataList,"isHiggsId"):
00570 return False
00571 return defaultParticleDataList.isHiggsId(particleId)
00572
00573 def lineStyle(self, object):
00574 particleId = self.particleId(object)
00575 if hasattr(defaultParticleDataList,"isPhotonId") and defaultParticleDataList.isPhotonId(particleId):
00576 return self.LINE_STYLE_WAVE
00577 elif defaultParticleDataList.isGluonId(particleId):
00578 return self.LINE_STYLE_SPIRAL
00579 elif defaultParticleDataList.isBosonId(particleId):
00580 return self.LINE_STYLE_DASH
00581 return self.LINE_STYLE_SOLID
00582
00583 def color(self, object):
00584 particleId = self.particleId(object)
00585 if defaultParticleDataList.isLeptonId(particleId):
00586 return QColor(244, 164, 96)
00587 elif defaultParticleDataList.isQuarkId(particleId):
00588 return QColor(0, 100, 0)
00589 elif hasattr(defaultParticleDataList,"isHiggsId") and defaultParticleDataList.isHiggsId(particleId):
00590 return QColor(247, 77, 251)
00591 elif defaultParticleDataList.isBosonId(particleId):
00592 return QColor(253, 74, 74)
00593 return QColor(176, 179, 177)
00594
00595 def charge(self, object):
00596 charge=self.property(object,"charge")
00597 if charge==None:
00598 charge=0
00599 return charge
00600
00601 def linkMother(self, object, mother):
00602 pass
00603
00604 def linkDaughter(self, object, daughter):
00605 pass
00606
00607 def underscoreProperties(self):
00608 return self._underscore
00609
00610 def setUnderscoreProperties(self,check):
00611 self._underscore=check
00612
00613 def filterBranches(self):
00614 return self._filterBranches
00615
00616 def setFilterBranches(self,check):
00617 if not self._events:
00618 return True
00619 self._filterBranches=check
00620 if check:
00621 for branch in self._dataObjects[:]:
00622 result=self.read(branch,0)
00623 if isinstance(result,BranchDummy):
00624 self._dataObjects.remove(result)
00625 if hasattr(result,"invalid"):
00626 self._filteredBranches.remove(result.branchtuple)
00627 return True
00628 else:
00629 self._filteredBranches=self._branches[:]
00630 self.goto(self.eventNumber())
00631 return False
00632
00633 def filteredBranches(self):
00634 return self._filteredBranches
00635