Belle II Software light-2406-ragdoll
decayHash.py
1#!/usr/bin/env python3
2
3
10
11import re
12import struct
13import pdg
14import basf2
15import pybasf2
16import uproot
17
18import numpy as np
19
20
21def _bitwiseConversion(value, i='f', o='i'):
22 """
23 Bitwise conversion between to python types
24 This is equivalently to
25 union {
26 Type_i i;
27 Type_o o;
28 } conversion;
29 conversion.i = input;
30 return conversion.o
31 @param i input data type (e.g. f for float)
32 @param o output data type (e.g. i for integer)
33 """
34 s = struct.pack('>' + i, value)
35 return struct.unpack('>' + o, s)[0]
36
37
38def _decayHashFloatToInt(decayHash, decayHashExtended):
39 """
40 Convert decayHash and decayHashExtended 32 bit floats to an 64 bit integer
41 """
42 decayHashInt = _bitwiseConversion(np.float32(decayHash))
43 decayHashExtendedInt = _bitwiseConversion(np.float32(decayHashExtended))
44 decayHashFullInt = decayHashInt << 32
45 decayHashFullInt += decayHashExtendedInt
46 return decayHashFullInt
47
48
50 """
51 DecayHashMap using the C++ implementation of DecayTree and DecayNode
52 """
53
54 def __init__(self, rootfile, removeRadiativeGammaFlag=False):
55 """Constructor"""
56 # Always avoid the top-level 'import ROOT'.
57 import ROOT # noqa
58 with uproot.open(rootfile) as rf:
59 trees = rf.keys()
60 assert len(trees) == 1
61 ntuple = rf[trees[0]].arrays(library='np')
62 # self._removeGammaFlag = removeRadiativeGammaFlag
63
64 self._string = {}
65
66 self._forest = {}
67 for decayHash, decayHashExtended, decayString in zip(
68 ntuple['decayHash'], ntuple['decayHashExtended'], ntuple['decayString']):
69 decayInt = ROOT.Belle2.DecayForest.decayHashFloatToInt(decayHash, decayHashExtended)
70 if decayInt in self._string:
71 continue
72 self._string[decayInt] = decayString
73 self._forest[decayInt] = ROOT.Belle2.DecayForest(decayString, True, removeRadiativeGammaFlag)
74
75 def get_string(self, decayHash, decayHashExtended):
76 """
77 Return DecayString given the decayHash and decayHashExtended
78 @param decayHash output of extraInfo(decayHash)
79 @param decayHashExtended output of extraInfo(decayHashExtended)
80 """
81 # Always avoid the top-level 'import ROOT'.
82 import ROOT # noqa
83 return self._string[ROOT.Belle2.DecayForest.decayHashFloatToInt(decayHash, decayHashExtended)]
84
85 def get_original_decay(self, decayHash, decayHashExtended):
86 """
87 Return original (MC) DecayTree given the decayHash and decayHashExtended
88 @param decayHash output of extraInfo(decayHash)
89 @param decayHashExtended output of extraInfo(decayHashExtended)
90 """
91 # Always avoid the top-level 'import ROOT'.
92 import ROOT # noqa
93 return self._forest[ROOT.Belle2.DecayForest.decayHashFloatToInt(decayHash, decayHashExtended)].getOriginalTree()
94
95 def get_reconstructed_decay(self, decayHash, decayHashExtended):
96 """
97 Return reconstructed DecayTree given the decayHash and decayHashExtended
98 @param decayHash output of extraInfo(decayHash)
99 @param decayHashExtended output of extraInfo(decayHashExtended)
100 """
101 # Always avoid the top-level 'import ROOT'.
102 import ROOT # noqa
103 return self._forest[ROOT.Belle2.DecayForest.decayHashFloatToInt(decayHash, decayHashExtended)].getReconstructedTree()
104
105 def print_hash(self, decayHash, decayHashExtended):
106 """
107 Print the DecayString in a fancy way given the decayHash and decayHashExtended
108 @param decayHash output of extraInfo(decayHash)
109 @param decayHashExtended output of extraInfo(decayHashExtended)
110 """
111 entry = self.get_string(decayHash, decayHashExtended)
112 entries = entry.split('|')
113 all_particles = re.findall(r"(-?[0-9]+)", entries[0])
114
115 if len(all_particles) != len(entries) - 1:
116 print(entry)
117 raise RuntimeError("Bad format of decay string: " +
118 str(len(all_particles)) + " " + str(len(entries)) + " " + str(entries))
119
120 table = []
121 table.append(["Decay ", prettify_pdg_codes(entries[0])])
122 for particle, mc_decay_string in zip(all_particles, entries[1:]):
123 table.append([prettify_pdg_codes(particle), prettify_pdg_codes(mc_decay_string)])
124
125 basf2.pretty_print_table(table, column_widths=[6, '*'])
126
127
128def _pdg_to_name(x):
129 """
130 Convert PDG code to a name
131 @param a pdg code
132 """
133 selected = False
134 if x[0] == '^':
135 selected = True
136 x = x[1:]
137
138 pdg_code = int(x)
139 pdg_string = str(pdg_code)
140 try:
141 pdg_string = pdg.to_name(pdg_code)
142 except BaseException:
143 pass
144
145 if selected:
146 if pybasf2.LogPythonInterface.terminal_supports_colors():
147 return '\x1b[31m' + pdg_string + '\x1b[0m'
148 else:
149 return '^' + pdg_string
150 return pdg_string
151
152
153def prettify_pdg_codes(text):
154 """
155 Prettifiy a string containing PDG codes by replacing PDG codes
156 with their corresponding names.
157 @param text the text
158 """
159 text = re.sub(r"(\^?-?[0-9]+)", lambda x: _pdg_to_name(x.group(0)), text)
160 text = text.replace('gamma', 'g').replace('--> ', '').replace('anti-', 'a-')
161 text = text.replace('Upsilon', 'Y').replace(') ', ')').replace(' (', '(')
162 return text
_forest
Dict Int -> Reconstructed DecayTree.
Definition: decayHash.py:66
_string
Dict Int -> DecayStrings.
Definition: decayHash.py:64
def get_original_decay(self, decayHash, decayHashExtended)
Definition: decayHash.py:85
def get_reconstructed_decay(self, decayHash, decayHashExtended)
Definition: decayHash.py:95
def __init__(self, rootfile, removeRadiativeGammaFlag=False)
Definition: decayHash.py:54
def get_string(self, decayHash, decayHashExtended)
Definition: decayHash.py:75
def print_hash(self, decayHash, decayHashExtended)
Definition: decayHash.py:105
def to_name(pdg)
Definition: pdg.py:87