Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2023-03-17 11:16:38

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