Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-11-25 02:29:52

0001 # Benedikt Hegner, DESY
0002 # benedikt.hegner@cern.ch
0003 #
0004 # this tool is based on Luca Lista's tree drawer module
0005 
0006 
0007 from builtins import range
0008 class ParticleDecayDrawer(object):
0009     """Draws particle decay tree """
0010     
0011     def __init__(self):
0012         print("Init particleDecayDrawer")
0013         #  booleans: printP4 printPtEtaPhi printVertex   
0014         
0015     def _accept(self, candidate, skipList):
0016         if candidate in skipList: return False;
0017         return self._select(candidate)
0018         
0019     def _select(self, candidate):
0020         return candidate.status() == 3
0021          
0022     def _hasValidDaughters(self, candidate):
0023         nDaughters = candidate.numChildren()
0024         for i in range(nDaughters):
0025             if self._select(candidate.listChildren()[i]): return True
0026         return False        
0027 
0028     def _printP4(self, candidate):
0029         return " "
0030 
0031     def _decay(self, candidate, skipList):
0032 
0033           out = str()
0034           if candidate in skipList:
0035               return ""
0036           skipList.append(candidate)
0037     
0038           id = candidate.pdg_id()
0039           # here the part about the names :-(
0040           out += str(id) + self._printP4(candidate)
0041     
0042           validDau = 0
0043           nOfDaughters = candidate.numChildren()
0044           for i in range(nOfDaughters):
0045               if self._accept(candidate.listChildren()[i], skipList): validDau+=1
0046           if validDau == 0: return out
0047     
0048           out += " ->"
0049     
0050           for i in range(nOfDaughters):
0051               d = candidate.listChildren()[i]
0052               if self._accept(d, skipList):
0053                   decString = self._decay(d, skipList)
0054                   if ("->" in decString):  out += " ( %s ) " %decString
0055                   else:  out += " %s" %decString
0056           return out
0057 
0058     def draw(self, particles): 
0059         """ draw decay tree from list(HepMC.GenParticles)""" 
0060         skipList = []
0061         nodesList = []
0062         momsList = []  
0063         for particle in particles:
0064             if particle.numParents() > 1:
0065                 if self._select(particle):
0066                     skipList.append(particle)
0067                     nodesList.append(particle)
0068                     for j in range(particle.numParents()):
0069                         mom = particle.listParents()[j]
0070                         while (mom.mother()):# != None ):
0071                             mom = mom.mother()
0072                         if self._select(mom):
0073                             momsList.append(mom)
0074 
0075         print("-- decay --")  
0076         if len(momsList) > 0:
0077             if len(momsList) > 1:
0078                 for m in range(len(momsList)):
0079                     decString = self._decay( momsList[m], skipList)
0080                     if len(decString) > 0:
0081                        print("{ %s } " %decString)
0082             else:
0083                 print(self._decay(momsList[0], skipList))   
0084   
0085         if len(nodesList) > 0:
0086             print("-> ")
0087             if len(nodesList) > 1:
0088                 for node in nodesList:
0089                    skipList.remove(node)
0090                    decString = self._decay(node, skipList)
0091                    if len(decString) > 0:
0092                        if "->" in decString:  print(" ( %s ) " %decString)
0093                        else:  print(" " + decString)
0094             else:
0095                 skipList.remove(nodesList[0])
0096                 print(self._decay(nodesList[0], skipList))
0097           
0098         print()
0099