Belle II Software  release-05-01-25
registry.py
1 # !/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 """"""
5 
6 from importlib import import_module
7 import pandas as pd
8 
9 from basf2 import B2ERROR
10 
11 from tabulate import tabulate
12 
13 
14 _RegisteredSkims = [
15  # --- WG0: Systematics ---
16  ("10000000", "systematics", "Random"),
17  # ("10600100", "systematics", "Systematics"), renamed to SystematicsDstar.
18  ("10600300", "systematics", "SystematicsTracking"),
19  ("10600400", "systematics", "Resonance"),
20  ("10600500", "systematics", "SystematicsRadMuMu"),
21  ("10600600", "systematics", "SystematicsEELL"),
22  ("10600700", "systematics", "SystematicsRadEE"),
23  ("10620200", "systematics", "SystematicsLambda"),
24  ("11640100", "systematics", "SystematicsPhiGamma"),
25  ("10600800", "systematics", "SystematicsFourLeptonFromHLTFlag"),
26  ("10600900", "systematics", "SystematicsRadMuMuFromHLTFlag"),
27  ("10611000", "systematics", "SystematicsJpsi"),
28  ("10611100", "systematics", "SystematicsKshort"),
29  ("10601200", "systematics", "SystematicsBhabha"),
30  ("10601300", "systematics", "SystematicsCombinedHadronic"),
31  ("10601400", "systematics", "SystematicsCombinedLowMulti"),
32  ("10601500", "systematics", "SystematicsDstar"),
33 
34 
35  # --- WG1: SL + missing energy ---
36  ("11110100", "semileptonic", "PRsemileptonicUntagged"),
37  ("11130300", "leptonic", "LeptonicUntagged"),
38  ("11130301", "leptonic", "dilepton"),
39  ("11160200", "semileptonic", "SLUntagged"),
40  ("11160201", "semileptonic", "B0toDstarl_Kpi_Kpipi0_Kpipipi"),
41  ("11180100", "fei", "feiHadronicB0"),
42  ("11180200", "fei", "feiHadronicBplus"),
43  ("11180300", "fei", "feiSLB0"),
44  ("11180400", "fei", "feiSLBplus"),
45  ("11180500", "fei", "feiHadronic"),
46  ("11180600", "fei", "feiSL"),
47 
48  # --- WG2: Electroweak penguins ---
49  ("12160100", "ewp", "BtoXgamma"),
50  ("12160200", "ewp", "BtoXll"),
51  ("12160300", "ewp", "BtoXll_LFV"),
52  ("12160400", "ewp", "inclusiveBplusToKplusNuNu"),
53 
54  # --- WG3: Time-dependent CP violation ---
55  ("13160200", "tdcpv", "TDCPV_ccs"),
56  ("13160300", "tdcpv", "TDCPV_qqs"),
57 
58  # --- WG4: Charmed B decays ---
59  ("14120300", "btocharm", "BtoD0h_Kspi0"),
60  ("14120400", "btocharm", "BtoD0h_Kspipipi0"),
61  # B0 -> D-(k+ ""- pi-)pi+ # ("14140500", "", "BtoD0h_Kspi0pi0"),
62  # Add when skim script is ready
63  ("14120600", "btocharm", "B0toDpi_Kpipi"),
64  ("14120601", "btocharm", "B0toDpi_Kspi"), # B0 -> D-(Ks pi-)pi+
65  # B0 -> D*-(anti-D0 pi-)pi+ With anti-D0 -> k+ pi-
66  ("14120700", "btocharm", "B0toDstarPi_D0pi_Kpi"),
67  # merge B0 -> D*-(anti-D0 pi-)pi+ with anti-D0 -> k- pi+ pi+ pi-
68  # and anti-D0 -> K- pi+ pi0
69  ("14120800", "btocharm", "B0toDstarPi_D0pi_Kpipipi_Kpipi0"),
70  ("14121100", "btocharm", "B0toDrho_Kpipi"),
71  ("14121101", "btocharm", "B0toDrho_Kspi"),
72  ("14121200", "btocharm", "B0toDstarRho_D0pi_Kpi"),
73  ("14121201", "btocharm", "B0toDstarRho_D0pi_Kpipipi_Kpipi0"),
74  ("14140100", "btocharm", "BtoD0h_hh"),
75  ("14140101", "btocharm", "BtoD0h_Kpi"),
76  # B+ -> anti-D0/anti-D0* (K- pi+ pi+ pi-, K- ""+ pi0) h+
77  ("14140102", "btocharm", "BtoD0h_Kpipipi_Kpipi0"),
78  ("14140200", "btocharm", "BtoD0h_Kshh"),
79  ("14141000", "btocharm", "BtoD0rho_Kpi"),
80  ("14141001", "btocharm", "BtoD0rho_Kpipipi_Kpipi0"),
81  ("14141002", "btocharm", "B0toDD_Kpipi_Kspi"),
82  ("14141003", "btocharm", "B0toDstarD"),
83  ("14121300", "btocharm", "B0toD0Kpipi0_pi0"),
84 
85 
86  # --- WG5: Quarkonium ---
87  ("15410300", "quarkonium", "InclusiveLambda"),
88  ("15420100", "quarkonium", "BottomoniumEtabExclusive"),
89  ("15440100", "quarkonium", "BottomoniumUpsilon"),
90  # ("16460100", "quarkonium", "ISRpipicc"), Subset of 16460200, deleted.
91  ("16460200", "quarkonium", "CharmoniumPsi"),
92 
93  # --- WG7: Charm physics ---
94  ("17230100", "charm", "XToD0_D0ToHpJm"), # D0 -> K pi/pi pi/K K
95  # D0 -> pi0 pi0/Ks pi0/Ks Ks # ("17230100", "", "D0ToHpJm"),
96  # D0 -> K pi/pi pi/K K
97  ("17230200", "charm", "XToD0_D0ToNeutrals"),
98  ("17230300", "charm", "DstToD0Pi_D0ToRare"), # D0 -> g g/e e/mu mu
99  ("17230400", "charm", "XToDp_DpToKsHp"), # D+ -> Ks h+
100  ("17230500", "charm", "XToDp_DpToHpHmJp"), # D+ -> h+ h- j+
101  ("17230600", "charm", "LambdacTopHpJm"), # Lambda_c+ -> proton h- j+
102  ("17240100", "charm", "DstToD0Pi_D0ToHpJm"), # D* -> D0 -> K pi/pi pi/K K
103  # D* -> D0 -> K- pi+ pi0 (""+WS)
104  ("17240200", "charm", "DstToD0Pi_D0ToHpJmPi0"),
105  ("17240300", "charm", "DstToD0Pi_D0ToHpHmPi0"), # D* -> D0 -> h h pi0
106  # D* -> D0 -> Ks omega / Ks eta -> Ks pi+ pi- pi0
107  ("17240400", "charm", "DstToD0Pi_D0ToKsOmega"),
108  # D* -> D0 -> K- pi+ eta (""+WS)
109  ("17240500", "charm", "DstToD0Pi_D0ToHpJmEta"),
110  # D* -> D0 -> pi0 pi0/Ks pi0/Ks Ks
111  ("17240600", "charm", "DstToD0Pi_D0ToNeutrals"),
112  ("17240700", "charm", "DstToD0Pi_D0ToHpJmKs"), # D* -> D0 -> h h Ks
113  # D* -> D0 -> K- pi+ pi0 (""+WS)
114  ("17240800", "charm", "EarlyData_DstToD0Pi_D0ToHpJmPi0"),
115  ("17240900", "charm", "EarlyData_DstToD0Pi_D0ToHpHmPi0"), # D* -> D0 -> h h pi0
116  ("17241000", "charm", "DstToDpPi0_DpToHpPi0"), # D*+ -> D+ pi0, D+ -> h+ pi0
117  ("17241100", "charm", "DstToD0Pi_D0ToHpHmHpJm"), # D* -> D0 -> h h h j
118 
119  # --- WG8: Dark matter searches and tau physics ---
120  ("18020100", "dark", "SinglePhotonDark"),
121  ("18020200", "dark", "GammaGammaControlKLMDark"),
122  ("18020300", "dark", "ALP3Gamma"),
123  ("18020400", "dark", "EGammaControlDark"),
124  ("18000000", "dark", "InelasticDarkMatter"),
125  ("18000001", "dark", "RadBhabhaV0Control"),
126  ("18360100", "taupair", "TauLFV"),
127  ("18520100", "dark", "DimuonPlusMissingEnergy"),
128  ("18520200", "dark", "ElectronMuonPlusMissingEnergy"),
129  ("18520300", "dark", "DielectronPlusMissingEnergy"),
130  ("18520400", "dark", "LFVZpVisible"),
131  ("18570600", "taupair", "TauGeneric"),
132  ("18570700", "taupair", "TauThrust"),
133  ("18530100", "lowMulti", "TwoTrackLeptonsForLuminosity"),
134  ("18520500", "lowMulti", "LowMassTwoTrack"),
135  ("18530200", "lowMulti", "SingleTagPseudoScalar"),
136 
137  # --- WG9: Charmless B decays ---
138  ("19120100", "btocharmless", "BtoPi0Pi0"),
139  ("19130201", "btocharmless", "BtoHadTracks"),
140  ("19130300", "btocharmless", "BtoHad1Pi0"),
141  ("19130310", "btocharmless", "BtoHad3Tracks1Pi0"),
142  ("19120400", "btocharmless", "BtoRhopRhom"),
143 ]
144 """
145 A list of all official registered skims and their skim code and parent module. Entries
146 must be of the form ``(code, module, name)``.
147 """
148 
149 
150 def _add_skim_registry_table(SkimRegistry):
151  """
152  Decorator to add a Sphinx table to the docstring of the skim registry.
153 
154  Inserts table wherever '<TABLE>' is in the docstring.
155  """
156 
157  df = pd.DataFrame(_RegisteredSkims, columns=["Skim code", "Module", "Skim name"])
158  df = df[["Module", "Skim name", "Skim code"]].sort_values(by=["Module", "Skim code"])
159  table = tabulate(df, showindex="never", tablefmt="grid", headers=df.columns)
160 
161  # Manual text manipulation (read: filthy hack) to make the table hierarchical
162  OriginalLines = table.split("\n")
163  header, OriginalLines, footer = OriginalLines[:2], OriginalLines[2:-1], OriginalLines[-1]
164  CurrentModule = ""
165  lines = []
166  lines.append("\n ".join(header))
167  for BorderLine, TextLine in zip(OriginalLines[::2], OriginalLines[1::2]):
168  segments = TextLine.split("|")
169  module = segments[1].lstrip().rstrip()
170  if CurrentModule == module:
171  segments[1] = " " * len(segments[1])
172  BorderLine = "|" + " " * len(segments[1]) + BorderLine.lstrip("+").lstrip("-")
173  else:
174  CurrentModule = module
175  lines.append(BorderLine)
176  lines.append("|".join(segments))
177  lines.append(footer)
178 
179  SkimRegistry.__doc__ = SkimRegistry.__doc__.replace("<TABLE>", "\n ".join(lines))
180 
181  return SkimRegistry
182 
183 
184 @_add_skim_registry_table
186  """
187  Class containing information on all official registered skims. This class also
188  contains helper functions for getting information from the registry. For
189  convenience, an instance of this class is provided: `skim.registry.Registry`.
190 
191  The table below lists all registered skims and their skim codes:
192 
193  <TABLE>
194  """
195  _registry = _RegisteredSkims
196 
197  def __init__(self):
198  self._codes = [code for code, _, _ in self._registry]
199  self._modules = list({module for _, module, _ in self._registry})
200  self._names = [names for _, _, names in self._registry]
201 
202  @property
203  def names(self):
204  """A list of all registered skim names."""
205  return self._names
206 
207  @property
208  def codes(self):
209  """A list of all registered skim codes."""
210  return self._codes
211 
212  @property
213  def modules(self):
214  """A list of all registered skim modules."""
215  return self._modules
216 
217  def get_skim_module(self, SkimName):
218  """Retrieve the skim module name from the registry which contains the given
219  skim.
220 
221  Parameters:
222  SkimName (str): Name of the skim as it appears in the skim registry.
223 
224  Returns:
225  The name of the skim module which contains the skim.
226  """
227  lookup = {name: module for _, module, name in self._registry}
228  try:
229  return lookup[SkimName]
230  except KeyError:
231  B2ERROR(
232  f"Unrecognised skim name {SkimName}. "
233  "Please add your skim to the list in `skim/scripts/skim/registry.py`."
234  )
235  raise LookupError(SkimName)
236 
237  def get_skims_in_module(self, SkimModule):
238  """Retrieve a list of the skims listed in the registry as existing in
239  the given skim module.
240 
241  Parameters:
242  SkimModule (str): The name of the module, *e.g.* ``btocharmless`` (not
243  ``skim.btocharmless`` or ``btocharmless.py``).
244 
245  Returns:
246  The skims listed in the registry as belonging to ``SkimModule``.
247  """
248  if SkimModule not in self.modules:
249  B2ERROR(f"Unrecognised skim module {SkimModule}.")
250  raise LookupError(SkimModule)
251 
252  ModuleLookup = {name: module for _, module, name in self._registry}
253  NameLookup = {
254  module: [name for name in self.names if ModuleLookup[name] == module]
255  for module in self.modules
256  }
257  return NameLookup[SkimModule]
258 
259  def get_skim_function(self, SkimName):
260  """Get the skim class constructor for the given skim.
261 
262  This is achieved by importing the module listed alongside the skim name in the
263  skim registry.
264 
265  Parameters:
266  SkimName (str): Name of the skim to be found.
267 
268  Returns:
269  The class constructor for the given skim.
270  """
271  ModuleName = self.get_skim_module(SkimName)
272  SkimModule = import_module(f"skim.{ModuleName}")
273  return getattr(SkimModule, SkimName)
274 
275  def encode_skim_name(self, SkimName):
276  """Find the 8 digit skim code assigned to the skim with the provided name.
277 
278  Parameters:
279  SkimName (str): Name of the corresponding skim as it appears in the skim registry.
280 
281  Returns:
282  8 digit skim code assigned to the given skim.
283  """
284  lookup = {name: code for code, _, name in self._registry}
285  try:
286  return lookup[SkimName]
287  except KeyError:
288  B2ERROR(
289  f"Unrecognised skim name {SkimName}. "
290  "Please add your skim to the list in `skim/scripts/skim/registry.py`."
291  )
292  raise LookupError(SkimName)
293 
294  def decode_skim_code(self, SkimCode):
295  """Find the name of the skim which corresponds to the provided skim code.
296 
297  This is useful to determine the skim script used to produce a specific uDST
298  file, given the 8-digit code name of the file itself.
299 
300  Parameters:
301  SkimCode (str): 8 digit skim code assigned to some skim.
302 
303  Returns:
304  Name of the corresponding skim as it appears in the skim registry.
305  """
306  lookup = {code: name for code, _, name in self._registry}
307  try:
308  return lookup[SkimCode]
309  except KeyError:
310  B2ERROR(
311  f"Unrecognised skim code {SkimCode}. "
312  "Please add your skim to the list in `skim/scripts/skim/registry.py`."
313  )
314  raise LookupError(SkimCode)
315 
316 
317 Registry = SkimRegistryClass()
318 """
319 An instance of `skim.registry.SkimRegistryClass`. Use this in your script to get
320 information from the registry.
321 
322  >>> from skim.registry import Registry
323  >>> Registry.encode_skim_name("SinglePhotonDark")
324  18020100
325 """
skim.registry.SkimRegistryClass._names
_names
Definition: registry.py:200
skim.registry.SkimRegistryClass.get_skim_module
def get_skim_module(self, SkimName)
Definition: registry.py:217
skim.registry.SkimRegistryClass.names
def names(self)
Definition: registry.py:203
skim.registry.SkimRegistryClass.decode_skim_code
def decode_skim_code(self, SkimCode)
Definition: registry.py:294
skim.registry.SkimRegistryClass._modules
_modules
Definition: registry.py:199
skim.registry.SkimRegistryClass
Definition: registry.py:185
skim.registry.SkimRegistryClass.modules
def modules(self)
Definition: registry.py:213
skim.registry.SkimRegistryClass.encode_skim_name
def encode_skim_name(self, SkimName)
Definition: registry.py:275
skim.registry.SkimRegistryClass._registry
_registry
Definition: registry.py:195
skim.registry.SkimRegistryClass.get_skims_in_module
def get_skims_in_module(self, SkimModule)
Definition: registry.py:237
skim.registry.SkimRegistryClass.codes
def codes(self)
Definition: registry.py:208
skim.registry.SkimRegistryClass._codes
_codes
Definition: registry.py:198
skim.registry.SkimRegistryClass.get_skim_function
def get_skim_function(self, SkimName)
Definition: registry.py:259