Source code for skim.utils.misc
#!/usr/bin/env python3
##########################################################################
# basf2 (Belle II Analysis Software Framework) #
# Author: The Belle II Collaboration #
# #
# See git log for contributors and copyright holders. #
# This file is licensed under LGPL-3.0, see LICENSE.md. #
##########################################################################
"""
Miscellaneous utility functions for skim experts.
"""
import subprocess
import json
import re
from pathlib import Path
from skim.registry import Registry
[docs]def get_file_metadata(filename):
"""
Retrieve the metadata for a file using ``b2file-metadata-show``.
Parameters:
metadata (str): File to get number of events from.
Returns:
dict: Metadata of file in dict format.
"""
if not Path(filename).exists():
raise FileNotFoundError(f"Could not find file {filename}")
proc = subprocess.run(
["b2file-metadata-show", "--json", str(filename)],
stdout=subprocess.PIPE,
check=True,
)
metadata = json.loads(proc.stdout.decode("utf-8"))
return metadata
[docs]def get_eventN(filename):
"""
Retrieve the number of events in a file using ``b2file-metadata-show``.
Parameters:
filename (str): File to get number of events from.
Returns:
int: Number of events in the file.
"""
return int(get_file_metadata(filename)["nEvents"])
[docs]def resolve_skim_modules(SkimsOrModules, *, LocalModule=None):
"""
Produce an ordered list of skims, by expanding any Python skim module names into a
list of skims in that module. Also produce a dict of skims grouped by Python module.
Raises:
RuntimeError: Raised if a skim is listed twice.
ValueError: Raised if ``LocalModule`` is passed and skims are normally expected
from more than one module.
"""
skims = []
for name in SkimsOrModules:
if name in Registry.names:
skims.append(name)
elif name in Registry.modules:
skims.extend(Registry.get_skims_in_module(name))
duplicates = {skim for skim in skims if skims.count(skim) > 1}
if duplicates:
raise RuntimeError(
f"Skim{'s'*(len(duplicates)>1)} requested more than once: {', '.join(duplicates)}"
)
modules = sorted({Registry.get_skim_module(skim) for skim in skims})
if LocalModule:
if len(modules) > 1:
raise ValueError(
f"Local module {LocalModule} specified, but the combined skim expects "
"skims from more than one module. No steering file written."
)
modules = {LocalModule.rstrip(".py"): sorted(skims)}
else:
modules = {
module: sorted(
[skim for skim in skims if Registry.get_skim_module(skim) == module]
)
for module in modules
}
return skims, modules
class _hashable_list(list):
def __hash__(self):
return hash(tuple(self))
def _sphinxify_decay(decay_string):
"""Format the given decay string by using LaTeX commands instead of plain-text.
Output is formatted for use with Sphinx (ReStructured Text).
This is a utility function for autogenerating skim documentation.
Parameters:
decay_string (str): A decay descriptor.
Returns:
sphinxed_string (str): LaTeX version of the decay descriptor.
"""
decay_string = re.sub("^(B.):generic", "\\1_{\\\\text{had}}", decay_string)
decay_string = decay_string.replace(":generic", "")
decay_string = decay_string.replace(":semileptonic", "_{\\text{SL}}")
decay_string = decay_string.replace(":FSP", "_{FSP}")
decay_string = decay_string.replace(":V0", "_{V0}")
decay_string = re.sub("_[0-9]+", "", decay_string)
# Note: these are applied from top to bottom, so if you have
# both B0 and anti-B0, put anti-B0 first.
substitutes = [
("==>", "\\to"),
("->", "\\to"),
("gamma", "\\gamma"),
("p+", "p"),
("anti-p-", "\\bar{p}"),
("pi+", "\\pi^+"),
("pi-", "\\pi^-"),
("pi0", "\\pi^0"),
("K_S0", "K^0_S"),
("K_L0", "K^0_L"),
("mu+", "\\mu^+"),
("mu-", "\\mu^-"),
("tau+", "\\tau^+"),
("tau-", "\\tau^-"),
("nu", "\\nu"),
("K+", "K^+"),
("K-", "K^-"),
("e+", "e^+"),
("e-", "e^-"),
("J/psi", "J/\\psi"),
("anti-Lambda_c-", "\\Lambda^{-}_{c}"),
("anti-Sigma+", "\\overline{\\Sigma}^{+}"),
("anti-Lambda0", "\\overline{\\Lambda}^{0}"),
("anti-D0*", "\\overline{D}^{0*}"),
("anti-D*0", "\\overline{D}^{0*}"),
("anti-D0", "\\overline{D}^0"),
("anti-B0", "\\overline{B}^0"),
("Sigma+", "\\Sigma^{+}"),
("Lambda_c+", "\\Lambda^{+}_{c}"),
("Lambda0", "\\Lambda^{0}"),
("D+", "D^+"),
("D-", "D^-"),
("D0", "D^0"),
("D*+", "D^{+*}"),
("D*-", "D^{-*}"),
("D*0", "D^{0*}"),
("D_s+", "D^+_s"),
("D_s-", "D^-_s"),
("D_s*+", "D^{+*}_s"),
("D_s*-", "D^{-*}_s"),
("B+", "B^+"),
("B-", "B^-"),
("B0", "B^0"),
("B_s0", "B^0_s"),
("K*0", "K^{0*}"),
]
tex_string = decay_string
for (key, value) in substitutes:
tex_string = tex_string.replace(key, value)
return f":math:`{tex_string}`"
[docs]def fancy_skim_header(SkimClass):
"""Decorator to generate a fancy header to skim documentation and prepend it to the
docstring. Add this just above the definition of a skim.
Also ensures the documentation of the template functions like `BaseSkim.build_lists`
is not repeated in every skim documentation.
.. code-block:: python
@fancy_skim_header
class MySkimName(BaseSkim):
# docstring here describing your skim, and explaining cuts.
"""
SkimName = SkimClass.__name__
SkimCode = Registry.encode_skim_name(SkimName)
authors = SkimClass.__authors__ or ["(no authors listed)"]
description = SkimClass.__description__ or "(no description)"
contact = SkimClass.__contact__ or "(no contact listed)"
category = SkimClass.__category__ or "(no category listed)"
if isinstance(authors, str):
# If we were given a string, split it up at: commas, "and", "&", and newlines
authors = re.split(
r",\s+and\s+|\s+and\s+|,\s+&\s+|\s+&\s+|,\s+|\s*\n\s*", authors
)
# Strip any remaining whitespace either side of an author's name
authors = [re.sub(r"^\s+|\s+$", "", author) for author in authors]
if isinstance(category, list):
category = ", ".join(category)
# If multiple contacts were given, split them up:
contacts = re.split(r",\s+and\s+|\s+and\s+|,\s+&\s+|\s+&\s+|,\s+|\s*\n\s*", contact)
# If the contact is of the form "NAME <EMAIL>" or "NAME (EMAIL)", then make it a link
matches = [re.match("([^<>()`]+) [<(]([^<>()`]+@[^<>()`]+)[>)]", contact) for contact in contacts]
for i, match in enumerate(matches):
if match:
name, email = match[1], match[2]
contacts[i] = f"`{name} <mailto:{email}>`_"
else:
contacts[i] = contacts[i]
header = f"""
Note:
* **Skim description**: {description}
* **Skim name**: {SkimName}
* **Skim LFN code**: {SkimCode}
* **Category**: {category}
* **Author{"s"*(len(authors) > 1)}**: {", ".join(authors)}
* **Contact{"s"*(len(contacts) > 1)}**: {", ".join(contacts)}
"""
if SkimClass.ApplyHLTHadronCut:
HLTLine = "*This skim includes a selection on the HLT flag* ``hlt_hadron``."
header = f"{header.rstrip()}\n\n {HLTLine}\n"
if SkimClass.__doc__:
SkimClass.__doc__ = header + "\n\n" + SkimClass.__doc__.lstrip("\n")
else:
# Handle case where docstring is empty, or was not redefined
SkimClass.__doc__ = header
# If documentation of template functions not redefined, make sure BaseSkim docstring is not repeated
SkimClass.load_standard_lists.__doc__ = SkimClass.load_standard_lists.__doc__ or ""
SkimClass.build_lists.__doc__ = SkimClass.build_lists.__doc__ or ""
SkimClass.validation_histograms.__doc__ = (
SkimClass.validation_histograms.__doc__ or ""
)
SkimClass.additional_setup.__doc__ = SkimClass.additional_setup.__doc__ or ""
return SkimClass
[docs]def dry_run_steering_file(SteeringFile):
"""
Check if the steering file at the given path can be run with the "--dry-run" option.
"""
proc = subprocess.run(
["basf2", "--dry-run", "-i", "i.root", "-o", "o.root", str(SteeringFile)],
stderr=subprocess.PIPE,
stdout=subprocess.PIPE,
)
if proc.returncode != 0:
stdout = proc.stdout.decode("utf-8")
stderr = proc.stderr.decode("utf-8")
raise RuntimeError(
f"An error occured while dry-running steering file {SteeringFile}\n"
f"Script output:\n{stdout}\n{stderr}"
)