reboost 0.8.3__py3-none-any.whl

reboost/math/functions.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+import logging
+
+import awkward as ak
+import numpy as np
+from lgdo import Array, VectorOfVectors
+from lgdo.types import LGDO
+
+log = logging.getLogger(__name__)
+
+
+def piecewise_linear_activeness(
+    distances: VectorOfVectors | ak.Array, fccd: float, dlf: float
+) -> VectorOfVectors | Array:
+    r"""Piecewise linear HPGe activeness model.
+
+    Based on:
+
+    .. math::
+
+        f(d) =
+        \begin{cases}
+        0 & \text{if } d < f*l, \\
+        \frac{x-f*l}{f - f*l} & \text{if } t \leq d < f, \\
+        1 & \text{otherwise.}
+        \end{cases}
+
+    Where:
+
+    - `d`: Distance to surface,
+    - `l`: Dead layer fraction, the fraction of the FCCD which is fully inactive
+    - `f`: Full charge collection depth (FCCD).
+
+    In addition, any distance of `np.nan` (for example if the calculation
+    was not performed for some steps) is assigned an activeness of one.
+
+    Parameters
+    ----------
+    distances
+        the distance from each step to the detector surface. Can be either a
+        `numpy` or `awkward` array, or a LGDO `VectorOfVectors` or `Array`. The computation
+        is performed for each element and the shape preserved in the output.
+
+    fccd
+        the value of the FCCD
+    dlf
+        the fraction of the FCCD which is fully inactive.
+
+    Returns
+    -------
+    a :class:`VectorOfVectors` or :class:`Array` of the activeness
+    """
+    # convert to ak
+    if isinstance(distances, LGDO):
+        distances_ak = distances.view_as("ak")
+    elif not isinstance(distances, ak.Array):
+        distances_ak = ak.Array(distances)
+    else:
+        distances_ak = distances
+
+    dl = fccd * dlf
+    distances_flat = (
+        ak.flatten(distances_ak).to_numpy() if distances_ak.ndim > 1 else distances_ak.to_numpy()
+    )
+
+    # compute the linear piecewise
+    results = np.full_like(distances_flat, np.nan, dtype=np.float64)
+    lengths = ak.num(distances_ak) if distances_ak.ndim > 1 else len(distances_ak)
+
+    mask1 = (distances_flat > fccd) | np.isnan(distances_flat)
+    mask2 = (distances_flat <= dl) & (~mask1)
+    mask3 = ~(mask1 | mask2)
+
+    # assign the values
+    results[mask1] = 1
+    results[mask2] = 0
+    results[mask3] = (distances_flat[mask3] - dl) / (fccd - dl)
+
+    # reshape
+    results = ak.unflatten(ak.Array(results), lengths) if distances_ak.ndim > 1 else results
+
+    return VectorOfVectors(results) if results.ndim > 1 else Array(results)
+
+
+def vectorised_active_energy(
+    distances: VectorOfVectors | ak.Array,
+    edep: VectorOfVectors | ak.Array,
+    fccd: float | list,
+    dlf: float | list,
+) -> VectorOfVectors | Array:
+    r"""Energy after piecewise linear HPGe activeness model vectorised over FCCD or dead layer fraction.
+
+    Based on the same linear activeness function as :func:`piecewise_linear_activeness`. However,
+    this function vectorises the calculation to provide a range of output energies varying the fccd or
+    dead layer fraction. Either fccd or dlf can be a list. This adds an extra dimension to the
+    output, with the same length as the input fccd or dlf list.
+
+    .. warning:
+        It is not currently implemented to vary both dlf and fccd.
+
+    Parameters
+    ----------
+    distances
+        the distance from each step to the detector surface. Can be either a
+        `awkward` array, or a LGDO `VectorOfVectors` . The computation
+        is performed for each element and the first dimension is preserved, a
+        new dimension is added vectorising over the FCCD or DLF.
+    edep
+        the energy for each step.
+    fccd
+        the value of the FCCD, can be a list.
+    dlf
+        the fraction of the FCCD which is fully inactive, can be a list.
+
+    Returns
+    -------
+    a :class:`VectorOfVectors` or :class:`Array` of the activeness
+    """
+    # add checks on fccd, dlf
+    fccd = np.array(fccd)
+    dlf = np.array(dlf)
+
+    if (fccd.ndim + dlf.ndim) > 1:
+        msg = "Currently only one of FCCD and dlf can be varied"
+        raise NotImplementedError(msg)
+
+    # convert fccd and or dlf to the right shape
+    if fccd.ndim == 0:
+        if dlf.ndim == 0:
+            dlf = dlf[np.newaxis]
+        fccd = np.full_like(dlf, fccd)
+
+    dl = fccd * dlf
+
+    def _convert(field):
+        # convert to ak
+        if isinstance(field, VectorOfVectors):
+            field_ak = field.view_as("ak")
+        elif not isinstance(field, ak.Array):
+            field_ak = ak.Array(field)
+        else:
+            msg = f"{field} must be an awkward array or VectorOfVectors"
+            raise TypeError(msg)
+
+        return field_ak, ak.flatten(field_ak).to_numpy()[:, np.newaxis]
+
+    distances_ak, distances_flat = _convert(distances)
+    _, edep_flat = _convert(edep)
+    runs = ak.num(distances_ak, axis=-1)
+
+    # vectorise fccd or tl
+
+    fccd_list = np.tile(fccd, (len(distances_flat), 1))
+    dl_list = np.tile(dl, (len(distances_flat), 1))
+    distances_shaped = np.tile(distances_flat, (1, len(dl)))
+
+    # compute the linear piecewise
+    results = np.full_like(fccd_list, np.nan, dtype=np.float64)
+
+    # Masks
+    mask1 = (distances_shaped > fccd_list) | np.isnan(distances_shaped)
+    mask2 = ((distances_shaped <= dl_list) | (fccd_list == dl_list)) & ~mask1
+    mask3 = ~(mask1 | mask2)  # Safe, avoids recomputing anything expensive
+
+    # Assign values
+    results[mask1] = 1.0
+    results[mask2] = 0.0
+    results[mask3] = (distances_shaped[mask3] - dl_list[mask3]) / (
+        fccd_list[mask3] - dl_list[mask3]
+    )
+
+    energy = ak.sum(ak.unflatten(results * edep_flat, runs), axis=-2)
+
+    return VectorOfVectors(energy) if energy.ndim > 1 else Array(energy.to_numpy())

reboost/math/stats.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+
+import awkward as ak
+import numpy as np
+from lgdo import Array
+from numpy.typing import ArrayLike
+
+log = logging.getLogger(__name__)
+
+
+def get_resolution(
+    energies: ak.Array, channels: ak.Array, tcm_tables: dict, reso_pars: dict, reso_func: Callable
+) -> ak.Array:
+    """Get the resolution for each energy.
+
+    Parameters
+    ----------
+    energies
+        the energies to smear
+    channels
+        the channel index for each energy
+    tcm_tables
+        the mapping from indices to channel names.
+    reso_pars
+        the pars for each channel.
+    reso_func
+        the function to compute the resolution.
+    """
+    n_pars = len(reso_pars[next(iter(reso_pars))])
+
+    pars_shaped = []
+
+    for _ in range(n_pars):
+        pars_shaped.append(np.zeros(len(ak.flatten(channels))))
+
+    num = ak.num(channels, axis=-1)
+
+    for key, value in tcm_tables.items():
+        for i in range(n_pars):
+            pars_shaped[i][ak.flatten(channels) == value] = reso_pars[key][i]
+
+    ch_reso = reso_func(ak.flatten(energies), *pars_shaped)
+    return ak.unflatten(ch_reso, num)
+
+
+def apply_energy_resolution(
+    energies: ak.Array, channels: ak.Array, tcm_tables: dict, reso_pars: dict, reso_func: Callable
+):
+    """Apply the energy resolution sampling to an array with many channels.
+
+    Parameters
+    ----------
+    energies
+        the energies to smear
+    channels
+        the channel index for each energy
+    tcm_tables
+        the mapping from indices to channel names.
+    reso_pars
+        the pars for each channel.
+    reso_func
+        the function to compute the resolution.
+    """
+    num = ak.num(channels, axis=-1)
+
+    ch_reso = get_resolution(energies, channels, tcm_tables, reso_pars, reso_func)
+    energies_flat_smear = gaussian_sample(ak.flatten(energies), ak.flatten(ch_reso))
+
+    return ak.unflatten(energies_flat_smear, num)
+
+
+def gaussian_sample(mu: ArrayLike, sigma: ArrayLike | float, *, seed: int | None = None) -> Array:
+    r"""Generate samples from a gaussian.
+
+    Based on:
+
+    .. math::
+
+        y_i \sim \mathcal{N}(\mu_i,\sigma_i)
+
+    where $y_i$ is the output, $x_i$ the input (mu) and $\sigma$ is the standard
+    deviation for each point.
+
+    Parameters
+    ----------
+    mu
+        the mean positions to sample from, should be a flat (ArrayLike) object.
+    sigma
+        the standard deviation for each input value, can also be a single float.
+    seed
+        the random seed.
+
+    Returns
+    -------
+    sampled values.
+    """
+    # convert inputs
+
+    if isinstance(mu, Array):
+        mu = mu.view_as("np")
+    elif isinstance(mu, ak.Array):
+        mu = mu.to_numpy()
+    elif not isinstance(mu, np.ndarray):
+        mu = np.array(mu)
+
+    # similar for sigma
+    if isinstance(sigma, Array):
+        sigma = sigma.view_as("np")
+    elif isinstance(sigma, ak.Array):
+        sigma = sigma.to_numpy()
+    elif not isinstance(sigma, float | int | np.ndarray):
+        sigma = np.array(sigma)
+
+    rng = np.random.default_rng(seed=seed)  # Create a random number generator
+
+    return Array(rng.normal(loc=mu, scale=sigma))

reboost/optmap/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from .optmap import OpticalMap
+
+__all__ = ["OpticalMap"]

reboost/optmap/cli.py

@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+import argparse
+import logging
+
+import dbetto
+
+from ..log_utils import setup_log
+from ..utils import _check_input_file, _check_output_file
+
+log = logging.getLogger(__name__)
+
+
+def optical_cli() -> None:
+    parser = argparse.ArgumentParser(
+        prog="reboost-optical",
+        description="%(prog)s command line interface",
+    )
+
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="count",
+        default=0,
+        help="""Increase the program verbosity""",
+    )
+
+    parser.add_argument(
+        "--bufsize",
+        action="store",
+        type=int,
+        default=int(5e6),
+        help="""Row count for input table buffering (only used if applicable). default: %(default)e""",
+    )
+
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    # STEP 1a: build map file from evt tier
+    map_parser = subparsers.add_parser("createmap", help="build optical map from evt file(s)")
+    map_parser.add_argument(
+        "--settings",
+        action="store",
+        help="""Select a config file for binning.""",
+        required=True,
+    )
+    map_parser.add_argument(
+        "--detectors",
+        help=(
+            "file that contains a list of detector ids that will be produced as additional output maps."
+            + "By default, all channels will be included."
+        ),
+    )
+    map_parser_det_group = map_parser.add_mutually_exclusive_group(required=True)
+    map_parser_det_group.add_argument(
+        "--geom",
+        help="GDML geometry file",
+    )
+    map_parser.add_argument(
+        "--n-procs",
+        "-N",
+        type=int,
+        default=1,
+        help="number of worker processes to use. default: %(default)e",
+    )
+    map_parser.add_argument(
+        "--check",
+        action="store_true",
+        help="""Check map statistics after creation. default: %(default)s""",
+    )
+    map_parser.add_argument(
+        "input", help="input stp or optmap-evt LH5 file", metavar="INPUT_EVT", nargs="+"
+    )
+    map_parser.add_argument("output", help="output map LH5 file", metavar="OUTPUT_MAP")
+
+    # STEP 1b: view maps
+    mapview_parser = subparsers.add_parser(
+        "viewmap",
+        help="view optical map (arrows: navigate slices/axes, 'c': channel selector)",
+        formatter_class=argparse.RawTextHelpFormatter,
+        description=(
+            "Interactively view optical maps stored in LH5 files.\n\n"
+            "Keyboard controls:\n"
+            "  left/right  - previous/next slice along the current axis\n"
+            "  up/down     - switch slicing axis (x, y, z)\n"
+            "  c           - open channel selector overlay to switch detector map\n\n"
+            "Display notes:\n"
+            "  - Cells where no primary photons were simulated are shown in white.\n"
+            "  - Cells where no photons were detected are shown in grey.\n"
+            "  - Cells with values above the colormap maximum are shown in red.\n"
+            "  - Use --hist to choose which histogram to display. 'prob_unc_rel' shows the\n"
+            "    relative uncertainty prob_unc / prob where defined.\n"
+            "  - Use --divide to show the ratio of two map files (this/other)."
+        ),
+        epilog=(
+            "Examples:\n"
+            "  reboost-optical viewmap mymap.lh5\n"
+            "  reboost-optical viewmap mymap.lh5 --channel _1067205\n"
+            "  reboost-optical viewmap mymap.lh5 --hist prob_unc_rel --min 0 --max 1\n"
+            "  reboost-optical viewmap mymap.lh5 --divide other.lh5 --title 'Comparison'"
+        ),
+    )
+    mapview_parser.add_argument("input", help="input map LH5 file", metavar="INPUT_MAP")
+    mapview_parser.add_argument(
+        "--channel",
+        action="store",
+        default="all",
+        help="channel to display ('all' or '_<detid>'). Press 'c' in the viewer to switch. default: %(default)s",
+    )
+    mapview_parser.add_argument(
+        "--hist",
+        choices=("_nr_gen", "_nr_det", "prob", "prob_unc", "prob_unc_rel"),
+        action="store",
+        default="prob",
+        help="select optical map histogram to show. default: %(default)s",
+    )
+    mapview_parser.add_argument(
+        "--divide",
+        action="store",
+        help="divide by another map file before display (ratio). default: none",
+    )
+    mapview_parser.add_argument(
+        "--min",
+        default=1e-4,
+        type=(lambda s: s if s == "auto" else float(s)),
+        help="colormap min value; use 'auto' for automatic scaling. default: %(default)e",
+    )
+    mapview_parser.add_argument(
+        "--max",
+        default=1e-2,
+        type=(lambda s: s if s == "auto" else float(s)),
+        help="colormap max value; use 'auto' for automatic scaling. default: %(default)e",
+    )
+    mapview_parser.add_argument("--title", help="title of figure. default: stem of filename")
+
+    # STEP 1c: merge maps
+    mapmerge_parser = subparsers.add_parser("mergemap", help="merge optical maps")
+    mapmerge_parser.add_argument(
+        "input", help="input map LH5 files", metavar="INPUT_MAP", nargs="+"
+    )
+    mapmerge_parser.add_argument("output", help="output map LH5 file", metavar="OUTPUT_MAP")
+    mapmerge_parser.add_argument(
+        "--settings",
+        action="store",
+        help="""Select a config file for binning.""",
+        required=True,
+    )
+    mapmerge_parser.add_argument(
+        "--n-procs",
+        "-N",
+        type=int,
+        default=1,
+        help="number of worker processes to use. default: %(default)e",
+    )
+    mapmerge_parser.add_argument(
+        "--check",
+        action="store_true",
+        help="""Check map statistics after creation. default: %(default)s""",
+    )
+
+    # STEP 1d: check map
+    checkmap_parser = subparsers.add_parser("checkmap", help="check optical maps")
+    checkmap_parser.add_argument("input", help="input map LH5 file", metavar="INPUT_MAP")
+
+    # STEP X: rebin maps
+    rebin_parser = subparsers.add_parser("rebin", help="rebin optical maps")
+    rebin_parser.add_argument("input", help="input map LH5 files", metavar="INPUT_MAP")
+    rebin_parser.add_argument("output", help="output map LH5 file", metavar="OUTPUT_MAP")
+    rebin_parser.add_argument("--factor", type=int, help="integer scale-down factor")
+
+    args = parser.parse_args()
+
+    log_level = (None, logging.INFO, logging.DEBUG)[min(args.verbose, 2)]
+    setup_log(log_level)
+
+    # STEP 1a: build map file from evt tier
+    if args.command == "createmap":
+        from .create import create_optical_maps
+
+        _check_input_file(parser, args.input)
+        _check_output_file(parser, args.output)
+
+        # load settings for binning from config file.
+        _check_input_file(parser, args.input, "settings")
+        settings = dbetto.utils.load_dict(args.settings)
+
+        chfilter = "*"
+        if args.detectors is not None:
+            # load detector ids from a JSON/YAML array
+            chfilter = dbetto.utils.load_dict(args.detectors)
+
+        create_optical_maps(
+            args.input,
+            settings,
+            args.bufsize,
+            chfilter=chfilter,
+            output_lh5_fn=args.output,
+            check_after_create=args.check,
+            n_procs=args.n_procs,
+            geom_fn=args.geom,
+        )
+
+    # STEP 1b: view maps
+    if args.command == "viewmap":
+        from .mapview import view_optmap
+
+        _check_input_file(parser, args.input)
+        if args.divide is not None:
+            _check_input_file(parser, args.divide)
+        view_optmap(
+            args.input,
+            args.channel,
+            args.divide,
+            cmap_min=args.min,
+            cmap_max=args.max,
+            title=args.title,
+            histogram_choice=args.hist,
+        )
+
+    # STEP 1c: merge maps
+    if args.command == "mergemap":
+        from .create import merge_optical_maps
+
+        # load settings for binning from config file.
+        _check_input_file(parser, args.input, "settings")
+        settings = dbetto.utils.load_dict(args.settings)
+
+        _check_input_file(parser, args.input)
+        _check_output_file(parser, args.output)
+        merge_optical_maps(
+            args.input, args.output, settings, check_after_create=args.check, n_procs=args.n_procs
+        )
+
+    # STEP 1d: check maps
+    if args.command == "checkmap":
+        from .create import check_optical_map
+
+        _check_input_file(parser, args.input)
+        check_optical_map(args.input)
+
+    # STEP X: rebin maps
+    if args.command == "rebin":
+        from .create import rebin_optical_maps
+
+        _check_input_file(parser, args.input)
+        _check_output_file(parser, args.output)
+        rebin_optical_maps(args.input, args.output, args.factor)