PythonAnalysis/python/ParticleDecayDrawer.py

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