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