reboost 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

reboost/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.1'
-__version_tuple__ = version_tuple = (0, 2, 1)
+__version__ = version = '0.2.3'
+__version_tuple__ = version_tuple = (0, 2, 3)

reboost/build_evt.py

@@ -71,7 +71,6 @@ def build_evt(
         path to the evt tier (output) file, if `None` the :class:`Table` is returned in memory
     config
         dictionary of the configuration.
-
     buffer
         number of events to process simultaneously
 

reboost/build_glm.py

@@ -224,7 +224,7 @@ def build_glm(
         lh5_table_list = list(lh5.ls(stp_file, "stp/"))
 
         # get rows in the table
-        if files.glm is None:
+        if files.glm[file_idx] is None:
             for lh5_table in lh5_table_list:
                 if lh5_table.replace("stp/", "") not in glm_sum:
                     glm_sum[lh5_table.replace("stp/", "")] = None
@@ -232,7 +232,7 @@ def build_glm(
             glm_sum = None
 
         # start row for each table
-        start_row = {lh5_tab: 0 for lh5_tab in lh5_table_list}
+        start_row = dict.fromkeys(lh5_table_list, 0)
 
         vfield = f"stp/vertices/{id_name}"
 
@@ -274,7 +274,7 @@ def build_glm(
 
                 lh5_subgroup = lh5_table.replace("stp/", "")
 
-                if files.glm is not None:
+                if files.glm[file_idx] is not None:
                     store.write(
                         out_tab,
                         f"{out_table_name}/{lh5_subgroup}",

reboost/build_hit.py

@@ -179,7 +179,7 @@ def build_hit(
     config: Mapping | str,
     args: Mapping | AttrsDict,
     stp_files: str | list[str],
-    glm_files: str | list[str],
+    glm_files: str | list[str] | None,
     hit_files: str | list[str] | None,
     *,
     start_evtid: int = 0,
@@ -199,7 +199,7 @@ def build_hit(
     stp_files
         list of strings or string of the stp file path.
     glm_files
-        list of strings or string of the glm file path.
+        list of strings or string of the glm file path, if `None` will be build in memory.
     hit_files
         list of strings or string of the hit file path. The `hit` file can also be `None` in which
         case the hits are returned as an `ak.Array` in memory.
@@ -246,6 +246,8 @@ def build_hit(
         # loop over processing groups
         for group_idx, proc_group in enumerate(config["processing_groups"]):
             proc_name = proc_group.get("name", "default")
+            msg = f"... starting group {proc_name}"
+            log.info(msg)
 
             if proc_name not in time_dict:
                 time_dict[proc_name] = ProfileDict()
@@ -261,9 +263,11 @@ def build_hit(
                     for mapping in proc_group.get("detector_mapping")
                 ]
             )
-
             # loop over detectors
             for in_det_idx, (in_detector, out_detectors) in enumerate(detectors_mapping.items()):
+                msg = f"... processing {in_detector} (to {out_detectors})"
+                log.info(msg)
+
                 # get detector objects
                 det_objects = core.get_detector_objects(
                     output_detectors=out_detectors,
@@ -286,23 +290,29 @@ def build_hit(
                     time_dict=time_dict[proc_name],
                 )
                 for stps, _, chunk_idx, _ in glm_it:
-                    # converting to awwkard
+                    # converting to awkward
                     if stps is None:
                         continue
 
-                    # produce the hit table
                     ak_obj = stps.view_as("ak")
 
+                    # produce the hit table
                     for out_det_idx, out_detector in enumerate(out_detectors):
                         # loop over the rows
                         if out_detector not in output_tables and files.hit is None:
                             output_tables[out_detector] = None
 
-                        hit_table = core.evaluate_hit_table_layout(
-                            copy.deepcopy(ak_obj),
-                            expression=proc_group["hit_table_layout"],
-                            time_dict=time_dict[proc_name],
-                        )
+                        # get the attributes
+                        attrs = utils.copy_units(stps)
+
+                        if "hit_table_layout" in proc_group:
+                            hit_table = core.evaluate_hit_table_layout(
+                                copy.deepcopy(ak_obj),
+                                expression=proc_group["hit_table_layout"],
+                                time_dict=time_dict[proc_name],
+                            )
+                        else:
+                            hit_table = copy.deepcopy(stps)
 
                         local_dict = {
                             "DETECTOR_OBJECTS": det_objects[out_detector],
@@ -310,7 +320,7 @@ def build_hit(
                             "DETECTOR": out_detector,
                         }
                         # add fields
-                        for field, expression in proc_group["operations"].items():
+                        for field, expression in proc_group.get("operations", {}).items():
                             # evaluate the expression
                             col = core.evaluate_output_column(
                                 hit_table,
@@ -323,10 +333,20 @@ def build_hit(
                             hit_table.add_field(field, col)
 
                         # remove unwanted fields
-                        hit_table = core.remove_columns(hit_table, outputs=proc_group["outputs"])
+                        if "outputs" in proc_group:
+                            hit_table = core.remove_columns(
+                                hit_table, outputs=proc_group["outputs"]
+                            )
+
+                        # assign units in the output table
+                        hit_table = utils.assign_units(hit_table, attrs)
 
                         # get the IO mode
 
+                        new_hit_file = (file_idx == 0) or (
+                            files.hit[file_idx] != files.hit[file_idx - 1]
+                        )
+
                         wo_mode = (
                             "of"
                             if (
@@ -334,6 +354,7 @@ def build_hit(
                                 and out_det_idx == 0
                                 and in_det_idx == 0
                                 and chunk_idx == 0
+                                and new_hit_file
                             )
                             else "append"
                         )

reboost/build_tcm.py

@@ -6,7 +6,7 @@ import re
 import awkward as ak
 from lgdo import Table, lh5
 
-from reboost import shape
+from reboost.shape import group
 
 log = logging.getLogger(__name__)
 
@@ -102,7 +102,7 @@ def get_tcm_from_ak(
 
     obj_tot = ak.concatenate(sort_objs)
 
-    return shape.group.group_by_time(
+    return group.group_by_time(
         obj_tot,
         time_name=time_name,
         evtid_name=idx_name,

reboost/core.py

@@ -59,6 +59,7 @@ def evaluate_output_column(
     expr = expression.replace(f"{table_name}.", "")
 
     # get func call and modules to import
+
     func_call, globals_dict = utils.get_function_string(expr)
 
     msg = f"evaluating table with command {expr} and local_dict {local_dict.keys()}"
@@ -153,7 +154,7 @@ def get_global_objects(
 
 
 def get_detectors_mapping(
-    output_detector_expression: str,
+    output_detector_expression: str | list,
     objects: AttrsDict | None = None,
     input_detector_name: str | None = None,
 ) -> dict:
@@ -210,15 +211,21 @@ def get_detectors_mapping(
                                 input_detector_name = "dets",objects=objs)
     {'dets': ['ch0', 'ch1', 'ch2']}
     """
-    func, globs = utils.get_function_string(output_detector_expression)
     out_names = []
+    if isinstance(output_detector_expression, str):
+        out_list = [output_detector_expression]
+    else:
+        out_list = list(output_detector_expression)
+
+    for expression_tmp in out_list:
+        func, globs = utils.get_function_string(expression_tmp)
 
-    # if no package was imported its just a name
-    try:
-        objs = evaluate_object(output_detector_expression, local_dict={"OBJECTS": objects})
-        out_names.extend(objs)
-    except Exception:
-        out_names.append(output_detector_expression)
+        # if no package was imported its just a name
+        try:
+            objs = evaluate_object(expression_tmp, local_dict={"OBJECTS": objects})
+            out_names.extend(objs)
+        except Exception:
+            out_names.append(expression_tmp)
 
     # simple one to one mapping
     if input_detector_name is None:
@@ -273,19 +280,19 @@ def get_detector_objects(
 
     det_objects_dict = {}
     for output_detector in output_detectors:
-        det_objects_dict[output_detector] = AttrsDict(
-            {
-                obj_name: evaluate_object(
-                    obj_expression,
-                    local_dict={
-                        "ARGS": args,
-                        "DETECTOR": output_detector,
-                        "OBJECTS": global_objects,
-                    },
-                )
-                for obj_name, obj_expression in expressions.items()
-            }
-        )
+        obj_dict = {}
+        for obj_name, obj_expression in expressions.items():
+            obj_dict[obj_name] = evaluate_object(
+                obj_expression,
+                local_dict={
+                    "ARGS": args,
+                    "DETECTOR": output_detector,
+                    "OBJECTS": global_objects,
+                    "DETECTOR_OBJECTS": AttrsDict(obj_dict),
+                },
+            )
+
+        det_objects_dict[output_detector] = AttrsDict(obj_dict)
     res = AttrsDict(det_objects_dict)
 
     if time_dict is not None:

reboost/iterator.py

@@ -5,7 +5,9 @@ import time
 import typing
 
 from lgdo.lh5 import LH5Store
-from lgdo.types import LGDO
+from lgdo.types import LGDO, Table
+
+from reboost import build_glm
 
 log = logging.getLogger(__name__)
 
@@ -15,7 +17,7 @@ class GLMIterator:
 
     def __init__(
         self,
-        glm_file: str,
+        glm_file: str | None,
         stp_file: str,
         lh5_group: str,
         start_row: int,
@@ -31,7 +33,8 @@ class GLMIterator:
         Parameters
         ----------
         glm_file
-            the file containing the event lookup map.
+            the file containing the event lookup map, if `None` the glm will
+            be created in memory.
         stp_file
             the file containing the steps to read.
         lh5_group
@@ -65,6 +68,11 @@ class GLMIterator:
         self.sto = LH5Store()
         self.n_rows_read = 0
         self.time_dict = time_dict
+        self.glm = None
+
+        # build the glm in memory
+        if self.glm_file is None:
+            self.glm = build_glm.build_glm(stp_file, None, out_table_name="glm", id_name="evtid")
 
     def __iter__(self) -> typing.Iterator:
         self.current_i_entry = 0
@@ -83,10 +91,21 @@ class GLMIterator:
         if self.time_dict is not None:
             time_start = time.time()
 
-        # read the glm rows
-        glm_rows, n_rows_read = self.sto.read(
-            f"glm/{self.lh5_group}", self.glm_file, start_row=self.start_row_tmp, n_rows=n_rows
-        )
+        # read the glm rows]
+        if self.glm_file is not None:
+            glm_rows, n_rows_read = self.sto.read(
+                f"glm/{self.lh5_group}", self.glm_file, start_row=self.start_row_tmp, n_rows=n_rows
+            )
+        else:
+            # get the maximum row to read
+            max_row = self.start_row_tmp + n_rows
+            max_row = min(len(self.glm[self.lh5_group]), max_row)
+
+            if max_row != self.start_row_tmp:
+                glm_rows = Table(self.glm[self.lh5_group][self.start_row_tmp : max_row])
+
+            n_rows_read = max_row - self.start_row_tmp
+
         if self.time_dict is not None:
             self.time_dict.update_field("read/glm", time_start)
 
@@ -106,7 +125,6 @@ class GLMIterator:
             # extract range of stp rows to read
             start = glm_ak.start_row[0]
             n = sum(glm_ak.n_rows)
-
             if self.time_dict is not None:
                 time_start = time.time()
 

reboost/math/functions.py

@@ -11,7 +11,7 @@ log = logging.getLogger(__name__)
 
 
 def piecewise_linear_activeness(
-    distances: VectorOfVectors | ak.Array, fccd: float, tl: float
+    distances: VectorOfVectors | ak.Array, fccd: float, dlf: float
 ) -> VectorOfVectors | Array:
     r"""Piecewise linear HPGe activeness model.
 
@@ -21,14 +21,15 @@ def piecewise_linear_activeness(
 
         f(d) =
         \begin{cases}
-        0 & \text{if } d < t, \\
-        \frac{x-l}{f - l} & \text{if } t \leq d < f, \\
+        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`: Depth of transition layer start
+    - `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
@@ -43,8 +44,8 @@ def piecewise_linear_activeness(
 
     fccd
         the value of the FCCD
-    tl
-        the start of the transition layer.
+    dlf
+        the fraction of the FCCD which is fully inactive.
 
     Returns
     -------
@@ -58,10 +59,117 @@ def piecewise_linear_activeness(
     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 = ak.where(
-        (distances_ak > fccd) | np.isnan(distances_ak),
-        1,
-        ak.where(distances_ak <= tl, 0, (distances_ak - tl) / (fccd - tl)),
+    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]
     )
-    return VectorOfVectors(results) if results.ndim > 1 else Array(results.to_numpy())
+
+    energy = ak.sum(ak.unflatten(results * edep_flat, runs), axis=-2)
+
+    return VectorOfVectors(energy) if energy.ndim > 1 else Array(energy.to_numpy())

reboost/shape/cluster.py

@@ -40,17 +40,17 @@ def cluster_by_step_length(
     pos_x: ak.Array | VectorOfVectors,
     pos_y: ak.Array | VectorOfVectors,
     pos_z: ak.Array | VectorOfVectors,
-    dist: ak.Array | VectorOfVectors,
-    surf_cut: float = 2,
+    dist: ak.Array | VectorOfVectors | None = None,
+    surf_cut: float | None = None,
     threshold: float = 0.1,
-    threshold_surf: float = 0.0,
+    threshold_surf: float | None = None,
 ) -> VectorOfVectors:
     """Perform clustering based on the step length.
 
     Steps are clustered based on distance, if either:
      - a step is in a new track,
      - a step moves from surface to bulk region (or visa versa),
-     - the distance between the first step and the cluster and the current is above a threshold.
+     - the distance between the current step and the first step of the current cluster is above a threshold.
 
     Then a new cluster is started. The surface region is defined as the volume
     less than surf_cut distance to the surface. This allows for a fine tuning of the
@@ -67,9 +67,9 @@ def cluster_by_step_length(
     pos_z
         z position of the step.
     dist
-        distance to the detector surface.
+        distance to the detector surface. Can be `None` in which case all steps are treated as being in the "bulk".
     surf_cut
-        Size of the surface region (in mm)
+        Size of the surface region (in mm), if `None` no selection is applied (default).
     threshold
         Distance threshold in mm to combine steps in the bulk.
     threshold_surf
@@ -107,7 +107,7 @@ def cluster_by_step_length(
         ak.flatten(ak.local_index(trackid)).to_numpy(),
         ak.flatten(trackid).to_numpy(),
         pos,
-        ak.flatten(dist).to_numpy(),
+        dist_to_surf=ak.flatten(dist).to_numpy() if dist is not None else dist,
         surf_cut=surf_cut,
         threshold=threshold,
         threshold_surf=threshold_surf,
@@ -127,10 +127,10 @@ def cluster_by_distance_numba(
     local_index: np.ndarray,
     trackid: np.ndarray,
     pos: np.ndarray,
-    dist_to_surf: np.ndarray,
-    surf_cut: float = 2,
+    dist_to_surf: np.ndarray | None,
+    surf_cut: float | None = None,
     threshold: float = 0.1,
-    threshold_surf: float = 0.0,
+    threshold_surf: float | None = None,
 ) -> np.ndarray:
     """Cluster steps by the distance between points in the same track.
 
@@ -146,9 +146,9 @@ def cluster_by_distance_numba(
     pos
         `(n,3)` size array of the positions
     dist_to_surf
-        1D array of the distance to the detector surface.
+        1D array of the distance to the detector surface. Can be `None` in which case all steps are treated as being in the bulk.
     surf_cut
-        Size of the surface region (in mm)
+        Size of the surface region (in mm), if `None` no selection is applied.
     threshold
         Distance threshold in mm to combine steps in the bulk.
     threshold_surf
@@ -172,14 +172,20 @@ def cluster_by_distance_numba(
     is_surf_prev = False
 
     for idx in range(n):
-        thr = threshold if dist_to_surf[idx] > surf_cut else threshold_surf
-
-        new_cluster = (
-            (trackid[idx] != trackid_prev)
-            or (is_surf_prev and (dist_to_surf[idx] > surf_cut))
-            or ((not is_surf_prev) and (dist_to_surf[idx] < surf_cut))
-            or (_dist(pos[idx, :], pos_prev) > thr)
-        )
+        # consider a surface and a bulk region
+        if dist_to_surf is not None:
+            thr = threshold if dist_to_surf[idx] > surf_cut else threshold_surf
+
+            new_cluster = (
+                (trackid[idx] != trackid_prev)
+                or (is_surf_prev and (dist_to_surf[idx] > surf_cut))
+                or ((not is_surf_prev) and (dist_to_surf[idx] < surf_cut))
+                or (_dist(pos[idx, :], pos_prev) > thr)
+            )
+        # basic clustering without split into surface / bulk
+        else:
+            thr = threshold
+            new_cluster = (trackid[idx] != trackid_prev) or (_dist(pos[idx, :], pos_prev) > thr)
 
         # New hit, reset cluster index
         if idx == 0 or local_index[idx] == 0:
@@ -197,7 +203,8 @@ def cluster_by_distance_numba(
 
         # Update previous values
         trackid_prev = trackid[idx]
-        is_surf_prev = dist_to_surf[idx] < surf_cut
+        if dist_to_surf is not None:
+            is_surf_prev = dist_to_surf[idx] < surf_cut
 
     return out
 

reboost/shape/group.py

@@ -25,7 +25,9 @@ def _sort_data(obj: ak.Array, *, time_name: str = "time", evtid_name: str = "evt
     -------
     sorted awkward array
     """
+    obj = obj[ak.argsort(obj[evtid_name])]
     obj_unflat = ak.unflatten(obj, ak.run_lengths(obj[evtid_name]))
+
     indices = ak.argsort(obj_unflat[time_name], axis=-1)
     sorted_obj = obj_unflat[indices]
 
@@ -120,9 +122,9 @@ def group_by_time(
 
     # get difference
     time_diffs = np.diff(obj[time_name])
-    index_diffs = np.diff(obj[evtid_name])
+    index_diffs = np.array(np.diff(obj[evtid_name]), dtype=np.int32)
 
-    # index of thhe last element in each run
+    # index of the last element in each run
     time_change = (time_diffs > window * 1000) & (index_diffs == 0)
     index_change = index_diffs > 0
 

reboost/utils.py

@@ -1,25 +1,62 @@
 from __future__ import annotations
 
 import importlib
+import itertools
 import logging
 import re
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
 from contextlib import contextmanager
 from pathlib import Path
 
 from dbetto import AttrsDict
+from lgdo.types import Table
 
 log = logging.getLogger(__name__)
 
 
 def get_file_dict(
     stp_files: list[str] | str,
-    glm_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."""
+    """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
+    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)
+    stp_is_list = isinstance(stp_files, list)
+
+    make_files_list = (not hit_is_list and stp_is_list) or (
+        hit_is_list and stp_is_list and len(hit_files) == 1 and len(stp_files) > 1
+    )
+
+    hit_files_list = [hit_files] * len(stp_files) if (make_files_list) else hit_files
+
     files = {}
-    for file_type, file_list in zip(["stp", "glm", "hit"], [stp_files, glm_files, hit_files]):
+
+    for file_type, file_list in zip(
+        ["stp", "glm", "hit"], [stp_files, glm_files_list, hit_files_list]
+    ):
         if isinstance(file_list, str):
             files[file_type] = [file_list]
         else:
@@ -35,6 +72,48 @@ def get_file_list(path: str | None, threads: int | None = None) -> list[str]:
     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:
+            tab[field].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*\("
@@ -119,6 +198,33 @@ def get_function_string(expr: str, aliases: dict | None = None) -> tuple[str, di
     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.
 

{reboost-0.2.1.dist-info → reboost-0.2.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reboost
-Version: 0.2.1
+Version: 0.2.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

{reboost-0.2.1.dist-info → reboost-0.2.3.dist-info}/RECORD

@@ -1,20 +1,20 @@
 reboost/__init__.py,sha256=RVNl3Qgx_hTUeBGXaWYmiTcmXUDhTfvlAGGC8bo_jP8,316
-reboost/_version.py,sha256=UoNvMtd4wCG76RwoSpNCUtaFyTwakGcZolfjXzNVSMY,511
-reboost/build_evt.py,sha256=5Q3T0LCl8xMtyRRhcs6layC1xh4vp2f26PgB1yab2zs,4798
-reboost/build_glm.py,sha256=kxQN6MYl-vfKnOHufPqf5ifEoaToqaR30iMXISxIhYQ,9253
-reboost/build_hit.py,sha256=sJR2qXup-qP1IoDLRxquuTyeI7DP_1S80QZ-w94qaZY,13293
-reboost/build_tcm.py,sha256=N1rZwht88ZaKWmURch1VrVUbQROXfP56D0aj_JLsRhU,2951
+reboost/_version.py,sha256=wD8hnA5gV5UmPkQnpT3xR6V2csgj9K5NEADogbLK79M,511
+reboost/build_evt.py,sha256=zj3wG_kaV3EoRMQ33AkCNa_2Fv8cLtRuhyRyRmSrOYQ,4797
+reboost/build_glm.py,sha256=LQkM6x6mMOE92-c78uoclOvP9zp3vdMuLQCSP2f2Zk4,9263
+reboost/build_hit.py,sha256=KKfTJgoR5JnAMQVru58B76zPcqZxiexIUlWPUhb1zmU,14260
+reboost/build_tcm.py,sha256=-PawBHoHj0zsm4XsZu5bco9d9a09STicZchduefSNfI,2951
 reboost/cli.py,sha256=HTZ05DRnDodcf_D6BJCCavx5HqhKDadJCgf-oh8HTJk,6365
-reboost/core.py,sha256=AamREubQsAqJ-y10NRx18r-PuqlQj3iTl2PzOTWXGQI,10540
-reboost/iterator.py,sha256=cqfh3c0uLP67S0YGaw05-McZQzdMb8BISULIm3PEbKA,3990
+reboost/core.py,sha256=7Nclc6RUCOSJ1CWVAX0rFNJGM1LEgqvc4tD04CxEAtg,10766
+reboost/iterator.py,sha256=72AyoRTgMpWghZt2UOqRj0RGiNzaiBAwgNIUZdduK2s,4698
 reboost/log_utils.py,sha256=VqS_9OC5NeNU3jcowVOBB0NJ6ssYvNWnirEY-JVduEA,766
 reboost/profile.py,sha256=EOTmjmS8Rm_nYgBWNh6Rntl2XDsxdyed7yEdWtsZEeg,2598
-reboost/utils.py,sha256=PMnHvSD5MpIzJyA3IQD_fLAK-O1RMY68DPGbQJp7Yww,4967
+reboost/utils.py,sha256=T9GIknSKWsKAalbQT9Ny3u9UTYEvy8gghFhmoCs41Io,7751
 reboost/hpge/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 reboost/hpge/psd.py,sha256=vFs8Y5XVW261pB6aOvWmIDzqOaBg-gEOLhL9PbjlEKI,2113
 reboost/hpge/surface.py,sha256=SZyTmOCTipf27jYaJhtdInzGF1RZ2wKpbtf6HlOQYwM,3662
 reboost/math/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-reboost/math/functions.py,sha256=ZgQpm87pGE0wH4Ekjm-8SbEmzfZ5MlAxS-fTw0RsNMc,1875
+reboost/math/functions.py,sha256=OymiYTcA0NXxxm-MBDw5kqyNwHoLCmuv4J48AwnSrbU,5633
 reboost/math/stats.py,sha256=iiOEi87x93kqPWeSmlRiA5Oe-R8XR-plm6Z532PhC9M,1401
 reboost/optmap/__init__.py,sha256=imvuyld-GLw8qdwqW-lXCg2feptcTyQo3wIzPvDHwmY,93
 reboost/optmap/cli.py,sha256=wBexh-zrr5ABherEyk9xigxdArvOAKiiRQwAYon9Sro,9408
@@ -25,12 +25,12 @@ reboost/optmap/mapview.py,sha256=73kpe0_SKDj9bIhEx1ybX1sBP8TyvufiLfps84A_ijA,679
 reboost/optmap/numba_pdg.py,sha256=y8cXR5PWE2Liprp4ou7vl9do76dl84vXU52ZJD9_I7A,731
 reboost/optmap/optmap.py,sha256=j4rfbQ84PYSpE-BvP4Rdt96ZjPdwy8P4e4eZz1mATys,12817
 reboost/shape/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-reboost/shape/cluster.py,sha256=Cj4V1maPR-q_w6rKwF_hLW3Zmsv6zHva_I5oA2mm3PY,7442
-reboost/shape/group.py,sha256=bSmFCl_yi1hGaKudjiicDEJsiBNyAHiKYdr8ZuH4pSM,4406
+reboost/shape/cluster.py,sha256=RIvBlhHzp88aaUZGofp5SD9bimnoiqIOddhQ84jiwoM,8135
+reboost/shape/group.py,sha256=Q3DhEPxbhw3p4bwvpswSd0A-p224l5vRZnfQIEkOVJE,4475
 reboost/shape/reduction.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-reboost-0.2.1.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-reboost-0.2.1.dist-info/METADATA,sha256=zSFjGcEzPVcoBtwwbVyRdbI95BooP7swkYvojFfqNjU,44219
-reboost-0.2.1.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-reboost-0.2.1.dist-info/entry_points.txt,sha256=DxhD6BidSWNot9BrejHJjQ7RRLmrMaBIl52T75oWTwM,93
-reboost-0.2.1.dist-info/top_level.txt,sha256=q-IBsDepaY_AbzbRmQoW8EZrITXRVawVnNrB-_zyXZs,8
-reboost-0.2.1.dist-info/RECORD,,
+reboost-0.2.3.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+reboost-0.2.3.dist-info/METADATA,sha256=9cPQ0Bz2cGao4exSC1XwIqp61rfSHNIDLrIwD8SdbB8,44219
+reboost-0.2.3.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
+reboost-0.2.3.dist-info/entry_points.txt,sha256=DxhD6BidSWNot9BrejHJjQ7RRLmrMaBIl52T75oWTwM,93
+reboost-0.2.3.dist-info/top_level.txt,sha256=q-IBsDepaY_AbzbRmQoW8EZrITXRVawVnNrB-_zyXZs,8
+reboost-0.2.3.dist-info/RECORD,,

{reboost-0.2.1.dist-info → reboost-0.2.3.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (80.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any