Belle II Software  light-2205-abys
decayHash.py
1 #!/usr/bin/env python3
2 
3 
10 
11 import re
12 import struct
13 import pdg
14 import basf2
15 import pybasf2
16 
17 # inspect is also used by LogPythonInterface. Do not remove
18 import numpy as np
19 
20 
21 def _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 
38 def _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  import root_numpy
57  import ROOT
58  ntuple = root_numpy.root2array(rootfile)
59  # self._removeGammaFlag = removeRadiativeGammaFlag
60 
61  self._string_string = {}
62 
63  self._forest_forest = {}
64  for decayHash, decayHashExtended, decayString in ntuple:
65  decayInt = ROOT.Belle2.DecayForest.decayHashFloatToInt(decayHash, decayHashExtended)
66  if decayInt in self._string_string:
67  continue
68  self._string_string[decayInt] = decayString
69  self._forest_forest[decayInt] = ROOT.Belle2.DecayForest(decayString, True, removeRadiativeGammaFlag)
70 
71  def get_string(self, decayHash, decayHashExtended):
72  """
73  Return DecayString given the decayHash and decayHashExtended
74  @param decayHash output of extraInfo(decayHash)
75  @param decayHashExtended output of extraInfo(decayHashExtended)
76  """
77  import ROOT
78  return self._string_string[ROOT.Belle2.DecayForest.decayHashFloatToInt(decayHash, decayHashExtended)]
79 
80  def get_original_decay(self, decayHash, decayHashExtended):
81  """
82  Return original (MC) DecayTree given the decayHash and decayHashExtended
83  @param decayHash output of extraInfo(decayHash)
84  @param decayHashExtended output of extraInfo(decayHashExtended)
85  """
86  import ROOT
87  return self._forest_forest[ROOT.Belle2.DecayForest.decayHashFloatToInt(decayHash, decayHashExtended)].getOriginalTree()
88 
89  def get_reconstructed_decay(self, decayHash, decayHashExtended):
90  """
91  Return reconstructed DecayTree given the decayHash and decayHashExtended
92  @param decayHash output of extraInfo(decayHash)
93  @param decayHashExtended output of extraInfo(decayHashExtended)
94  """
95  import ROOT
96  return self._forest_forest[ROOT.Belle2.DecayForest.decayHashFloatToInt(decayHash, decayHashExtended)].getReconstructedTree()
97 
98  def print_hash(self, decayHash, decayHashExtended):
99  """
100  Print the DecayString in a fancy way given the decayHash and decayHashExtended
101  @param decayHash output of extraInfo(decayHash)
102  @param decayHashExtended output of extraInfo(decayHashExtended)
103  """
104  entry = self.get_stringget_string(decayHash, decayHashExtended)
105  entries = entry.split('|')
106  all_particles = re.findall(r"(-?[0-9]+)", entries[0])
107 
108  if len(all_particles) != len(entries) - 1:
109  print(entry)
110  raise RuntimeError("Bad format of decay string: " +
111  str(len(all_particles)) + " " + str(len(entries)) + " " + str(entries))
112 
113  table = []
114  table.append(["Decay ", prettify_pdg_codes(entries[0])])
115  for particle, mc_decay_string in zip(all_particles, entries[1:]):
116  table.append([prettify_pdg_codes(particle), prettify_pdg_codes(mc_decay_string)])
117 
118  basf2.pretty_print_table(table, column_widths=[6, '*'])
119 
120 
121 def _pdg_to_name(x):
122  """
123  Convert PDG code to a name
124  @param a pdg code
125  """
126  selected = False
127  if x[0] == '^':
128  selected = True
129  x = x[1:]
130 
131  pdg_code = int(x)
132  pdg_string = str(pdg_code)
133  try:
134  pdg_string = pdg.to_name(pdg_code)
135  except BaseException:
136  pass
137 
138  if selected:
139  if pybasf2.LogPythonInterface.terminal_supports_colors():
140  return '\x1b[31m' + pdg_string + '\x1b[0m'
141  else:
142  return '^' + pdg_string
143  return pdg_string
144 
145 
146 def prettify_pdg_codes(text):
147  """
148  Prettifiy a string containing PDG codes by replacing PDG codes
149  with their corresponding names.
150  @param text the text
151  """
152  text = re.sub(r"(\^?-?[0-9]+)", lambda x: _pdg_to_name(x.group(0)), text)
153  text = text.replace('gamma', 'g').replace('--> ', '').replace('anti-', 'a-')
154  text = text.replace('Upsilon', 'Y').replace(') ', ')').replace(' (', '(')
155  return text
_forest
Dict Int -> Reconstructed DecayTree.
Definition: decayHash.py:63
_string
Dict Int -> DecayStrings.
Definition: decayHash.py:61
def get_original_decay(self, decayHash, decayHashExtended)
Definition: decayHash.py:80
def get_reconstructed_decay(self, decayHash, decayHashExtended)
Definition: decayHash.py:89
def __init__(self, rootfile, removeRadiativeGammaFlag=False)
Definition: decayHash.py:54
def get_string(self, decayHash, decayHashExtended)
Definition: decayHash.py:71
def print_hash(self, decayHash, decayHashExtended)
Definition: decayHash.py:98
def to_name(pdg)
Definition: pdg.py:86