reboost 0.8.3__py3-none-any.whl

reboost/units.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import awkward as ak
+import pint
+import pyg4ometry as pg4
+from lgdo import LGDO
+
+log = logging.getLogger(__name__)
+
+ureg = pint.get_application_registry()
+"""The physical units registry."""
+
+# default pretty printing of physical units
+ureg.formatter.default_format = "~P"
+
+
+def pg4_to_pint(obj: pint.Quantity | pg4.gdml.Defines.VectorBase) -> pint.Quantity:
+    """Convert pyg4ometry object to pint Quantity."""
+    if isinstance(obj, pint.Quantity):
+        return obj
+    if isinstance(obj, pg4.gdml.Defines.VectorBase):
+        return [getattr(obj, field).eval() for field in ("x", "y", "z")] * ureg(obj.unit)
+    msg = f"I don't know how to convert object of type {type(obj)} to pint object"
+    raise ValueError(msg)
+
+
+def units_convfact(data: Any | LGDO | ak.Array, target_units: pint.Unit | str) -> float:
+    """Calculate numeric conversion factor to reach `target_units`.
+
+    Parameters
+    ----------
+    data
+        starting data structure. If an :class:`LGDO` or :class:`ak.Array`, try to
+        determine units by peeking into its attributes. Otherwise, just return 1.
+    target_units
+        units you wish to convert data to.
+    """
+    if isinstance(data, LGDO) and "units" in data.attrs:
+        return ureg(data.attrs["units"]).to(target_units).magnitude
+    if isinstance(data, ak.Array) and "units" in ak.parameters(data):
+        return ureg(ak.parameters(data)["units"]).to(target_units).magnitude
+    return 1
+
+
+def units_conv_ak(data: Any | LGDO | ak.Array, target_units: pint.Unit | str) -> ak.Array:
+    """Calculate numeric conversion factor to reach `target_units`, and apply to data converted to ak.
+
+    Parameters
+    ----------
+    data
+        starting data structure. If an :class:`LGDO` or :class:`ak.Array`, try to
+        determine units by peeking into its attributes. Otherwise, return the data
+        unchanged.
+    target_units
+        units you wish to convert data to.
+    """
+    fact = units_convfact(data, target_units)
+    if isinstance(data, LGDO) and fact != 1:
+        return ak.without_parameters(data.view_as("ak") * fact)
+    if isinstance(data, ak.Array) and fact != 1:
+        return ak.without_parameters(data * fact)
+    return data.view_as("ak") if isinstance(data, LGDO) else data
+
+
+def unwrap_lgdo(data: Any | LGDO | ak.Array, library: str = "ak") -> tuple[Any, pint.Unit | None]:
+    """Return a view of the data held by the LGDO and its physical units.
+
+    Parameters
+    ----------
+    data
+        the data container. If not an :class:`LGDO` or :class:`ak.Array`, it will be
+        returned as is with ``None`` units.
+    library
+        forwarded to :meth:`LGDO.view_as`.
+
+    Returns
+    -------
+    A tuple of the un-lgdo'd data and the data units.
+    """
+    ret_data = data
+    ret_units = None
+    if isinstance(data, LGDO):
+        ret_data = data.view_as(library)
+        if "units" in data.attrs:
+            ret_units = ureg(data.attrs["units"]).u
+
+    if isinstance(data, ak.Array):
+        if library != "ak":
+            msg = "cannot unwrap an awkward array as a non-awkward type"
+            raise ValueError(msg)
+
+        if "units" in ak.parameters(data):
+            ret_units = ureg(ak.parameters(data)["units"]).u
+            ret_data = ak.without_parameters(data)
+
+    return ret_data, ret_units
+
+
+def unit_to_lh5_attr(unit: pint.Unit) -> str:
+    """Convert Pint unit to a string that can be used as attrs["units"] in an LGDO."""
+    # TODO: we should check if this can be always parsed by Unitful.jl
+    if isinstance(unit, pint.Unit):
+        return f"{unit:~C}"
+    return unit

reboost/utils.py

@@ -0,0 +1,503 @@
+from __future__ import annotations
+
+import importlib
+import itertools
+import logging
+import re
+import time
+from collections.abc import Iterable, Mapping
+from contextlib import contextmanager
+from pathlib import Path
+
+import h5py
+from dbetto import AttrsDict
+from lgdo import lh5
+from lgdo.types import Struct, Table, VectorOfVectors
+
+from .profile import ProfileDict
+
+log = logging.getLogger(__name__)
+
+
+def get_table_names(tcm: VectorOfVectors) -> dict:
+    """Extract table names from tcm.attrs['tables'] and return them as a dictionary."""
+    raw = tcm.attrs["tables"]
+    cleaned = raw.strip("[]").replace(" ", "").replace("'", "")
+    tables = cleaned.split(",")
+    tables = [tab.split("/")[-1] for tab in tables]
+
+    return {name: idx for idx, name in enumerate(tables)}
+
+
+def get_wo_mode(
+    group: int, out_det: int, in_det: int, chunk: int, new_hit_file: bool, overwrite: bool = False
+) -> str:
+    """Get the mode for lh5 file writing.
+
+    If all indices are 0 and we are writing a new output file
+    then the mode "overwrite_file" is used (if the overwrite) flag
+    is set, otherwise the mode "write_safe" is used.
+
+    Otherwise the code choses between "append_column" if this is the
+    first time a group is being written to the file, or "append"
+
+    Parameters
+    ----------
+    group
+        the index of the processing group
+    out_det
+        the index of the output detector
+    in_det
+        the index of the input detector
+    chunk
+        the chunk index
+    new_hit_file
+        a flag of whether we are writing a new hit file. This does not indicate whether
+        the file already exists on disk, but whether the file name is different from the
+        last written chunk for this detector.
+    overwrite
+        a flag of whether to overwrite the old file.
+
+    Returns
+    -------
+    the mode for IO
+    """
+    indices = [group, out_det, in_det, chunk]
+
+    good_idx = all(i == 0 for i in indices)
+
+    if good_idx and new_hit_file:
+        return "overwrite_file" if overwrite else "write_safe"
+
+    # if we have a detector not the first and chunk 0 append column
+    is_ac = ((in_det > 0) or (out_det > 0)) & (chunk == 0)
+    is_ac = is_ac or (in_det == 0 and out_det == 0 and chunk == 0 and (group > 0))
+
+    if is_ac and new_hit_file:
+        return "append_column"
+    return "append"
+
+
+def get_wo_mode_forwarded(
+    written_tables: set[str], new_hit_file: bool, overwrite: bool = False
+) -> str:
+    """Get the mode for lh5 file writing for forwarded tables tahat will be copied without chunking.
+
+    If we are writing a new output file and no other tables had been written yet, then
+    the mode "overwrite_file" is used if the overwrite flag is set, otherwise the mode
+    "write_safe" is used.
+
+    Otherwise "append" is used.
+
+    Parameters
+    ----------
+    written_tables
+        a set of already written table names, also including other table names of
+        non-forwarded (i.e. processed) tables.
+    new_hit_file
+        a flag of whether we are writing a new hit file. This does not indicate whether
+        the file already exists on disk, but whether the file name is different from the
+        last written chunk for this forwarded table.
+    overwrite
+        a flag of whether to overwrite the old file.
+
+    Returns
+    -------
+    the mode for IO
+    """
+    if not new_hit_file:
+        return "append"
+    if overwrite and len(written_tables) == 0:
+        return "overwrite_file"
+    return "write_safe"
+
+
+def is_new_hit_file(files: AttrsDict, file_idx: int) -> bool:
+    """Return whether the hit file with the given index is a "new" hit file.
+
+    A new file is either the first file written, or when the previous file index has a
+    different file name.
+    """
+    return (file_idx == 0) or (files.hit[file_idx] != files.hit[file_idx - 1])
+
+
+def get_file_dict(
+    stp_files: list[str] | str,
+    glm_files: list[str] | str | None,
+    hit_files: list[str] | str | None = None,
+) -> AttrsDict:
+    """Get the file info as a AttrsDict.
+
+    Creates an :class:`dbetto.AttrsDict` with keys `stp_files`,
+    `glm_files` and `hit_files`. Each key contains a list of
+    file-paths (or `None`).
+
+    Parameters
+    ----------
+    stp_files
+        string or list of strings of the stp files.
+    glm_files
+        string or list of strings of the glm files, or None in which
+        case the glm will be created in memory.
+    hit_files
+        string or list of strings of the hit files, if None the output
+        files will be created in memory.
+    """
+    # make a list of the right length
+    if isinstance(stp_files, str):
+        stp_files = [stp_files]
+
+    glm_files_list = [None] * len(stp_files) if glm_files is None else glm_files
+
+    # make a list of files in case
+    # 1) hit_files is a str and stp_files is a list
+    # 2) hit_files and stp_files are both lists of different length
+
+    hit_is_list = isinstance(hit_files, list)
+
+    if not hit_is_list:
+        hit_files_list = [hit_files] * len(stp_files)
+    elif hit_is_list and len(hit_files) == 1 and len(stp_files) > 1:
+        hit_files_list = [hit_files[0]] * len(stp_files)
+    else:
+        hit_files_list = hit_files
+
+    files = {}
+
+    for file_type, file_list in zip(
+        ["stp", "glm", "hit"], [stp_files, glm_files_list, hit_files_list], strict=True
+    ):
+        if isinstance(file_list, str):
+            files[file_type] = [file_list]
+        else:
+            files[file_type] = file_list
+
+    return AttrsDict(files)
+
+
+def get_file_list(path: str | None, threads: int | None = None) -> list[str]:
+    """Get a list of files accounting for the multithread index."""
+    if threads is None or path is None:
+        return path
+    return [f"{(Path(path).with_suffix(''))}_t{idx}.lh5" for idx in range(threads)]
+
+
+def copy_units(tab: Table) -> dict:
+    """Extract a dictionary of attributes (i.e. units).
+
+    Parameters
+    ----------
+    tab
+        Table to get the units from.
+
+    Returns
+    -------
+    a dictionary with the units for each field
+    in the table.
+    """
+    units = {}
+
+    for field in tab:
+        if "units" in tab[field].attrs:
+            units[field] = tab[field].attrs["units"]
+
+    return units
+
+
+def assign_units(tab: Table, units: Mapping) -> Table:
+    """Copy the attributes from the map of attributes to the table.
+
+    Parameters
+    ----------
+    tab
+        Table to add attributes to.
+    units
+        mapping (dictionary like) of units of each field
+
+    Returns
+    -------
+    an updated table with LGDO attributes.
+    """
+    for field in tab:
+        if field in units:
+            if not isinstance(tab[field], VectorOfVectors):
+                tab[field].attrs["units"] = units[field]
+            else:
+                tab[field].flattened_data.attrs["units"] = units[field]
+
+    return tab
+
+
+def _search_string(string: str):
+    """Capture the characters matching the pattern for a function call."""
+    pattern = r"\b([a-zA-Z_][a-zA-Z0-9_\.]*)\s*\("
+    return re.findall(pattern, string)
+
+
+def get_function_string(expr: str, aliases: dict | None = None) -> tuple[str, dict]:
+    """Get a function call to evaluate.
+
+    Search for any patterns matching the pattern for a function call.
+    We also detect any cases of aliases being used, by default
+    just for `numpy` as `np` and `awkward` as `ak`. In this
+    case, the full name is replaces with the alias in the expression
+    and also in the output globals dictionary.
+
+    It is possible to chain together functions eg:
+
+    .. code-block:: python
+
+        ak.num(np.array([1, 2]))
+
+    and all packages will be imported.
+
+    Parameters
+    ----------
+    expr
+        expression to evaluate.
+    aliases
+        dictionary of package aliases for names used in dictionary. These allow to
+        give shorter names to packages. This is combined with two defaults `ak` for
+        `awkward` and `np` for `numpy`. If `None` is supplied only these are used.
+
+    Returns
+    -------
+    a tuple of call string and dictionary of the imported global packages.
+    """
+    # aliases for easier lookup
+    aliases = (
+        {"numpy": "np", "awkward": "ak"}
+        if aliases is None
+        else aliases | {"numpy": "np", "awkward": "ak"}
+    )
+
+    # move to only alias names
+    for name, short_name in aliases.items():
+        expr = expr.replace(name, short_name)
+
+    globs = {}
+    # search on the whole expression
+
+    funcs = _search_string(expr.strip())
+    for func_call in funcs:
+        # no "." then can't be a module
+        if "." not in func_call:
+            continue
+
+        subpackage, _func = func_call.rsplit(".", 1)
+        package = subpackage.split(".")[0]
+
+        # import the subpackage
+        for name, short_name in aliases.items():
+            subpackage = subpackage.replace(short_name, name)
+
+        # handle the aliases
+        package_import = package
+        for name, short_name in aliases.items():
+            if package == short_name:
+                package_import = name
+
+        # build globals
+        try:
+            importlib.import_module(subpackage, package=__package__)
+
+            globs = globs | {
+                package: importlib.import_module(package_import),
+            }
+        except Exception as e:
+            # for imports of our own package, raise the error back to the user
+            if package_import == "reboost":
+                raise e
+            msg = f"Function {package_import} cannot be imported"
+            log.debug(msg)
+            continue
+
+    return expr, globs
+
+
+def get_channels_from_groups(names: list | str | None, groupings: dict | None = None) -> list:
+    """Get a list of channels from a list of groups.
+
+    Parameters
+    ----------
+    names
+        list of channel names
+    groupings
+        dictionary of the groupings of channels
+
+    Returns
+    -------
+    list of channels
+    """
+    if names is None:
+        channels_e = []
+    elif isinstance(names, str):
+        channels_e = groupings[names]
+    elif isinstance(names, list):
+        channels_e = list(itertools.chain.from_iterable([groupings[e] for e in names]))
+    else:
+        msg = f"names {names} must be list or str or `None`"
+        raise ValueError(msg)
+
+    return channels_e
+
+
+def merge_dicts(dict_list: list) -> dict:
+    """Merge a list of dictionaries, concatenating the items where they exist.
+
+    Parameters
+    ----------
+    dict_list
+        list of dictionaries to merge
+
+    Returns
+    -------
+    a new dictionary after merging.
+
+    Examples
+    --------
+    >>> merge_dicts([{"a":[1,2,3],"b":[2]},{"a":[4,5,6],"c":[2]}])
+    {"a":[1,2,3,4,5,6],"b":[2],"c":[2]}
+    """
+    merged = {}
+
+    for tmp_dict in dict_list:
+        for key, item in tmp_dict.items():
+            if key in merged:
+                merged[key].extend(item)
+            else:
+                merged[key] = item
+
+    return merged
+
+
+@contextmanager
+def filter_logging(level):
+    logger = logging.getLogger("root")
+    old_level = logger.getEffectiveLevel()
+    logger.setLevel(level)
+    try:
+        yield
+    finally:
+        logger.setLevel(old_level)
+
+
+def _check_input_file(parser, file: str | Iterable[str], descr: str = "input") -> None:
+    file = (file,) if isinstance(file, str) else file
+    not_existing = [f for f in file if not Path(f).exists()]
+    if not_existing != []:
+        parser.error(f"{descr} file(s) {''.join(not_existing)} missing")
+
+
+def _check_output_file(parser, file: str | Iterable[str] | None, optional: bool = False) -> None:
+    if file is None and optional:
+        return
+
+    file = (file,) if isinstance(file, str) else file
+    for f in file:
+        if Path(f).exists():
+            parser.error(f"output file {f} already exists")
+
+
+def write_lh5(
+    hit_table: Table,
+    file: str,
+    time_dict: ProfileDict,
+    out_field: str,
+    out_detector: str,
+    wo_mode: str,
+):
+    """Write the lh5 file. This function handles writing first the data as a struct and then appending to this.
+
+    Parameters
+    ----------
+    hit_table
+        the table to write
+    file
+        the file to write to
+    time_dict
+        the dictionary of timing information to update.
+    out_field
+        output field
+    out_detector
+        output detector name
+    wo_mode
+        the mode to pass to `lh5.write`
+    """
+    if time_dict is not None:
+        start_time = time.time()
+
+    if wo_mode not in ("a", "append"):
+        lh5.write(
+            Struct({out_detector: hit_table}),
+            out_field,
+            file,
+            wo_mode=wo_mode,
+        )
+    else:
+        lh5.write(
+            hit_table,
+            f"{out_field}/{out_detector}",
+            file,
+            wo_mode=wo_mode,
+        )
+    if time_dict is not None:
+        time_dict.update_field("write", start_time)
+
+
+def get_remage_detector_uids(h5file: str | Path) -> dict:
+    """Get mapping of detector names to UIDs from a remage output file.
+
+    The remage LH5 output files contain a link structure that lets the user
+    access detector tables by UID. For example:
+
+    .. code-block:: text
+
+        ├── stp · struct{det1,det2,optdet1,optdet2,scint1,scint2}
+        └── __by_uid__ · struct{det001,det002,det011,det012,det101,det102}
+            ├── det001 -> /stp/scint1
+            ├── det002 -> /stp/scint2
+            ├── det011 -> /stp/det1
+            ├── det012 -> /stp/det2
+            ├── det101 -> /stp/optdet1
+            └── det102 -> /stp/optdet2
+
+    This function analyzes this structure and returns:
+
+    .. code-block:: text
+
+        {1: 'scint1',
+         2: 'scint2',
+         11: 'det1',
+         12: 'det2',
+         101: 'optdet1',
+         102: 'optdet2'g
+
+    Parameters
+    ----------
+    h5file
+        path to remage output file.
+    """
+    if isinstance(h5file, Path):
+        h5file = h5file.as_posix()
+
+    out = {}
+    with h5py.File(h5file, "r") as f:
+        g = f["/stp/__by_uid__"]
+        # loop over links
+        for key in g:
+            # is this a link?
+            link = g.get(key, getlink=True)
+            if isinstance(link, h5py.SoftLink):
+                m = re.fullmatch(r"det(\d+)", key)
+                if m is None:
+                    msg = rf"'{key}' is not formatted as expected, i.e. 'det(\d+)', skipping"
+                    log.warning(msg)
+                    continue
+
+                # get the name of the link target without trailing groups (to
+                # i.e. remove /stp)
+                name = link.path.split("/")[-1]
+
+                out[int(m.group(1))] = name
+    return out

reboost-0.8.3.dist-info/METADATA

@@ -0,0 +1,82 @@
+Metadata-Version: 2.4
+Name: reboost
+Version: 0.8.3
+Summary: New LEGEND Monte-Carlo simulation post-processing
+Author-email: Manuel Huber <info@manuelhu.de>, Toby Dixon <toby.dixon.23@ucl.ac.uk>, Luigi Pertoldi <gipert@pm.me>
+Maintainer: The LEGEND Collaboration
+License-Expression: GPL-3.0
+Project-URL: Homepage, https://github.com/legend-exp/reboost
+Project-URL: Bug Tracker, https://github.com/legend-exp/reboost/issues
+Project-URL: Discussions, https://github.com/legend-exp/reboost/discussions
+Project-URL: Changelog, https://github.com/legend-exp/reboost/releases
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: MacOS
+Classifier: Operating System :: POSIX
+Classifier: Operating System :: Unix
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Topic :: Scientific/Engineering
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: hdf5plugin
+Requires-Dist: colorlog
+Requires-Dist: numpy
+Requires-Dist: scipy
+Requires-Dist: numba>=0.60
+Requires-Dist: legend-pydataobj>=1.15.1
+Requires-Dist: legend-pygeom-optics>=0.15.0
+Requires-Dist: legend-pygeom-tools>=0.0.25
+Requires-Dist: legend-pygeom-hpges
+Requires-Dist: hist
+Requires-Dist: dbetto
+Requires-Dist: particle
+Requires-Dist: pandas
+Requires-Dist: matplotlib
+Requires-Dist: pygama
+Requires-Dist: pyg4ometry
+Provides-Extra: all
+Requires-Dist: reboost[docs,test]; extra == "all"
+Provides-Extra: docs
+Requires-Dist: furo; extra == "docs"
+Requires-Dist: myst-parser; extra == "docs"
+Requires-Dist: sphinx; extra == "docs"
+Requires-Dist: sphinx-copybutton; extra == "docs"
+Provides-Extra: test
+Requires-Dist: pre-commit; extra == "test"
+Requires-Dist: pytest>=6.0; extra == "test"
+Requires-Dist: pytest-cov; extra == "test"
+Requires-Dist: legend-pygeom-hpges; extra == "test"
+Requires-Dist: pylegendtestdata>=0.6; extra == "test"
+Dynamic: license-file
+
+# reboost
+
+[![PyPI](https://img.shields.io/pypi/v/reboost?logo=pypi)](https://pypi.org/project/reboost/)
+[![conda-forge](https://img.shields.io/conda/vn/conda-forge/reboost.svg)](https://anaconda.org/conda-forge/reboost)
+![GitHub tag (latest by date)](https://img.shields.io/github/v/tag/legend-exp/reboost?logo=git)
+[![GitHub Workflow Status](https://img.shields.io/github/checks-status/legend-exp/reboost/main?label=main%20branch&logo=github)](https://github.com/legend-exp/reboost/actions)
+[![pre-commit](https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit&logoColor=white)](https://github.com/pre-commit/pre-commit)
+[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
+[![Codecov](https://img.shields.io/codecov/c/github/legend-exp/reboost?logo=codecov)](https://app.codecov.io/gh/legend-exp/reboost)
+![GitHub issues](https://img.shields.io/github/issues/legend-exp/reboost?logo=github)
+![GitHub pull requests](https://img.shields.io/github/issues-pr/legend-exp/reboost?logo=github)
+![License](https://img.shields.io/github/license/legend-exp/reboost)
+[![Read the Docs](https://img.shields.io/readthedocs/reboost?logo=readthedocs)](https://reboost.readthedocs.io)
+
+_reboost_ is a package to post-process
+[remage](https://remage.readthedocs.io/en/stable/) simulations. Post processing
+is the step of applying a detector response model to the (idealised) _remage_ /
+_Geant4_ simulations to ''boost" them allowing comparison to data.
+
+_reboost_ provides tools to:
+
+- apply a HPGe detector response model to the simulations,
+- dedicated tools to generate optical maps,
+- functionality to control the full post-processing chain with configuration
+  files.
+
+For more information see our dedicated
+[documentation](https://reboost.readthedocs.io/en/stable/)!