21def _bitwiseConversion(value, i='f', o='i'):
23 Bitwise conversion between to python types
24 This is equivalently to
31 @param i input data type (e.g. f for float)
32 @param o output data type (e.g. i for integer)
34 s = struct.pack(
'>' + i, value)
35 return struct.unpack(
'>' + o, s)[0]
38def _decayHashFloatToInt(decayHash, decayHashExtended):
40 Convert decayHash and decayHashExtended 32 bit floats to an 64 bit integer
42 decayHashInt = _bitwiseConversion(np.float32(decayHash))
43 decayHashExtendedInt = _bitwiseConversion(np.float32(decayHashExtended))
44 decayHashFullInt = decayHashInt << 32
45 decayHashFullInt += decayHashExtendedInt
46 return decayHashFullInt
51 DecayHashMap using the C++ implementation of DecayTree and DecayNode
54 def __init__(self, rootfile, removeRadiativeGammaFlag=False):
56 # Always avoid the top-level 'import ROOT'.
58 with uproot.open(rootfile) as rf:
60 assert len(trees) == 1
61 ntuple = rf[trees[0]].arrays(library='np')
62 # self._removeGammaFlag = removeRadiativeGammaFlag
63 ## Dict Int -> DecayStrings
65 ## Dict Int -> Reconstructed DecayTree
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:
72 self._string[decayInt] = decayString
73 self._forest[decayInt] = ROOT.Belle2.DecayForest(decayString, True, removeRadiativeGammaFlag)
75 def get_string(self, decayHash, decayHashExtended):
77 Return DecayString given the decayHash and decayHashExtended
78 @param decayHash output of extraInfo(decayHash)
79 @param decayHashExtended output of extraInfo(decayHashExtended)
81 # Always avoid the top-level 'import ROOT'.
83 return self._string[ROOT.Belle2.DecayForest.decayHashFloatToInt(decayHash, decayHashExtended)]
85 def get_original_decay(self, decayHash, decayHashExtended):
87 Return original (MC) DecayTree given the decayHash and decayHashExtended
88 @param decayHash output of extraInfo(decayHash)
89 @param decayHashExtended output of extraInfo(decayHashExtended)
91 # Always avoid the top-level 'import ROOT'.
93 return self._forest[ROOT.Belle2.DecayForest.decayHashFloatToInt(decayHash, decayHashExtended)].getOriginalTree()
95 def get_reconstructed_decay(self, decayHash, decayHashExtended):
97 Return reconstructed DecayTree given the decayHash and decayHashExtended
98 @param decayHash output of extraInfo(decayHash)
99 @param decayHashExtended output of extraInfo(decayHashExtended)
101 # Always avoid the top-level 'import ROOT'.
103 return self._forest[ROOT.Belle2.DecayForest.decayHashFloatToInt(decayHash, decayHashExtended)].getReconstructedTree()
105 def print_hash(self, decayHash, decayHashExtended):
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)
111 entry = self.get_string(decayHash, decayHashExtended)
112 entries = entry.split('|')
113 all_particles = re.findall(r"(-?[0-9]+)", entries[0])
115 if len(all_particles) != len(entries) - 1:
117 raise RuntimeError("Bad format of decay string: " +
118 str(len(all_particles)) + " " + str(len(entries)) + " " + str(entries))
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)])
125 basf2.pretty_print_table(table, column_widths=[6, '*'])
130 Convert PDG code to a name
139 pdg_string = str(pdg_code)
141 pdg_string = pdg.to_name(pdg_code)
142 except BaseException:
146 if pybasf2.LogPythonInterface.terminal_supports_colors():
147 return '\x1b[31m' + pdg_string + '\x1b[0m'
149 return '^' + pdg_string
153def prettify_pdg_codes(text):
155 Prettifiy a string containing PDG codes by replacing PDG codes
156 with their corresponding names.
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(' (', '(')
__init__(self, rootfile, removeRadiativeGammaFlag=False)