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

reboost/__init__.py

@@ -1,6 +1,17 @@
 from __future__ import annotations
 
-from reboost import build_hit, core, iterator, math, optmap, shape
+from reboost import build_hit, core, iterator, math, shape
 from reboost._version import version as __version__
 
-__all__ = ["__version__", "build_hit", "core", "iterator", "math", "optmap", "shape"]
+__all__ = [
+    "__version__",
+    "build_glm",
+    "build_hit",
+    "build_hit",
+    "build_tcm",
+    "core",
+    "iterator",
+    "math",
+    "optmap",
+    "shape",
+]

reboost/_version.py

@@ -1,8 +1,13 @@
-# file generated by setuptools_scm
+# file generated by setuptools-scm
 # don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
 TYPE_CHECKING = False
 if TYPE_CHECKING:
-    from typing import Tuple, Union
+    from typing import Tuple
+    from typing import Union
+
     VERSION_TUPLE = Tuple[Union[int, str], ...]
 else:
     VERSION_TUPLE = object
@@ -12,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.1'
-__version_tuple__ = version_tuple = (0, 1, 1)
+__version__ = version = '0.2.1'
+__version_tuple__ = version_tuple = (0, 2, 1)

reboost/build_evt.py

@@ -1,9 +1,170 @@
+"""A program for combining the hits from various detectors, to build events.
+
+Is able to parse a config file with the following format config file:
+
+.. code-block:: yaml
+
+    channels:
+        geds_on:
+        - det001
+        - det002
+        geds_ac:
+        - det003
+
+    outputs:
+    - energy
+    - multiplicity
+
+    operations:
+     energy_id:
+        channels: geds_on
+        aggregation_mode: gather
+        query: "hit.energy > 25"
+        expression: tcm.channel_id
+
+     energy:
+        aggregation_mode: keep_at_ch:evt.energy_id
+        expression: "hit.energy > 25"
+        channels: geds_on
+
+     multiplicity:
+        channels: geds_on
+        aggregation_mode: sum
+        expression: "hit.energy > 25"
+        initial: 0
+
+
+Must contain:
+- "channels": dictionary of channel groupings
+- "outputs": fields for the output file
+- "operations": operations to perform see :func:`pygama.evt.build_evt.evaluate_expression` for more details.
+"""
+
 from __future__ import annotations
 
 import logging
 
+import awkward as ak
+import numpy as np
+from lgdo import Table
+from lgdo.lh5 import LH5Iterator, write
+from pygama.evt.build_evt import evaluate_expression
+from pygama.evt.utils import TCMData
+
+from . import utils
+
 log = logging.getLogger(__name__)
 
 
-def build_evt() -> None:
-    pass
+def build_evt(
+    hit_file: str, tcm_file: str, evt_file: str | None, config: dict, buffer: int = int(5e6)
+) -> ak.Array | None:
+    """Generates the event tier from the hit and tcm.
+
+    Parameters
+    ----------
+    hit_file
+        path to the hit tier file
+    tcm_file
+        path to the tcm tier file
+    evt_file
+        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
+
+    Returns
+    -------
+    ak.Array of the evt tier data (if the data is not saved to disk)
+    """
+    msg = "... beginning the evt tier processing"
+    log.info(msg)
+
+    # create the objects needed for evaluate expression
+
+    file_info = {
+        "hit": (hit_file, "hit", "det{:03}"),
+        "evt": (evt_file, "evt"),
+    }
+
+    # iterate through the TCM
+
+    out_ak = ak.Array([])
+    mode = "of"
+
+    # get channel groupings
+    channels = {}
+    for group, info in config["channels"].items():
+        if isinstance(info, str):
+            channels[group] = [info]
+
+        elif isinstance(info, list):
+            channels[group] = info
+
+    for tcm_lh5, _, n_rows_read in LH5Iterator(tcm_file, "tcm", buffer_len=buffer):
+        tcm_lh5_sel = tcm_lh5
+        tcm_ak = tcm_lh5_sel.view_as("ak")[:n_rows_read]
+
+        tcm = TCMData(
+            id=np.array(ak.flatten(tcm_ak.array_id)),
+            idx=np.array(ak.flatten(tcm_ak.array_idx)),
+            cumulative_length=np.array(np.cumsum(ak.num(tcm_ak.array_id, axis=-1))),
+        )
+
+        n_rows = len(tcm.cumulative_length)
+        out_tab = Table(size=n_rows)
+
+        for name, info in config["operations"].items():
+            msg = f"computing field {name}"
+            log.debug(msg)
+
+            defaultv = info.get("initial", np.nan)
+            if isinstance(defaultv, str) and (defaultv in ["np.nan", "np.inf", "-np.inf"]):
+                defaultv = eval(defaultv)
+
+            channels_use = utils.get_channels_from_groups(info.get("channels", []), channels)
+            channels_exclude = utils.get_channels_from_groups(
+                info.get("exclude_channels", []), channels
+            )
+
+            if "aggregation_mode" not in info:
+                field = out_tab.eval(
+                    info["expression"].replace("evt.", ""), info.get("parameters", {})
+                )
+            else:
+                field = evaluate_expression(
+                    file_info,
+                    tcm,
+                    channels_use,
+                    table=out_tab,
+                    mode=info["aggregation_mode"],
+                    expr=info["expression"],
+                    query=info.get("query", None),
+                    sorter=info.get("sort", None),
+                    channels_skip=channels_exclude,
+                    default_value=defaultv,
+                    n_rows=n_rows,
+                )
+
+            msg = f"field {field}"
+            log.debug(msg)
+            out_tab.add_field(name, field)
+
+        # remove fields if necessary
+        existing_cols = list(out_tab.keys())
+        for col in existing_cols:
+            if col not in config["outputs"]:
+                out_tab.remove_column(col, delete=True)
+
+        # write
+        if evt_file is not None:
+            write(out_tab, "evt", evt_file, wo_mode=mode)
+            mode = "append"
+        else:
+            out_ak = ak.concatenate((out_ak, out_tab.view_as("ak")))
+
+    if evt_file is None:
+        return out_ak
+    return None

reboost/build_glm.py

@@ -221,7 +221,7 @@ def build_glm(
         log.info(msg)
 
         # loop over the lh5_tables
-        lh5_table_list = [table for table in lh5.ls(stp_file, "stp/") if table != "stp/vertices"]
+        lh5_table_list = list(lh5.ls(stp_file, "stp/"))
 
         # get rows in the table
         if files.glm is None:

reboost/build_hit.py

@@ -195,7 +195,7 @@ def build_hit(
     config
         dictionary or path to YAML file containing the processing chain.
     args
-        dictionary or :class:`legendmeta.AttrsDict` of the global arguments.
+        dictionary or :class:`dbetto.AttrsDict` of the global arguments.
     stp_files
         list of strings or string of the stp file path.
     glm_files
@@ -246,7 +246,9 @@ def build_hit(
         # loop over processing groups
         for group_idx, proc_group in enumerate(config["processing_groups"]):
             proc_name = proc_group.get("name", "default")
-            time_dict[proc_name] = ProfileDict()
+
+            if proc_name not in time_dict:
+                time_dict[proc_name] = ProfileDict()
 
             # extract the output detectors and the mapping to input detectors
             detectors_mapping = utils.merge_dicts(

reboost/build_tcm.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+import logging
+import re
+
+import awkward as ak
+from lgdo import Table, lh5
+
+from reboost import shape
+
+log = logging.getLogger(__name__)
+
+
+def build_tcm(
+    hit_file: str,
+    out_file: str,
+    channels: list[str],
+    time_name: str = "t0",
+    idx_name: str = "global_evtid",
+    time_window_in_us: float = 10,
+) -> None:
+    """Build the (Time Coincidence Map) TCM from the hit tier.
+
+    Parameters
+    ----------
+    hit_file
+        path to hit tier file.
+    out_file
+        output path for tcm.
+    channels
+        list of channel names to include.
+    time_name
+        name of the hit tier field used for time grouping.
+    idx_name
+        name of the hit tier field used for index grouping.
+    time_window_in_us
+        time window used to define the grouping.
+    """
+    hash_func = r"\d+"
+
+    msg = "start building time-coincidence map"
+    log.info(msg)
+
+    chan_ids = [re.search(hash_func, channel).group() for channel in channels]
+
+    hit_data = []
+    for channel in channels:
+        hit_data.append(
+            lh5.read(f"{channel}/hit", hit_file, field_mask=[idx_name, time_name]).view_as("ak")
+        )
+    tcm = get_tcm_from_ak(
+        hit_data, chan_ids, window=time_window_in_us, time_name=time_name, idx_name=idx_name
+    )
+
+    if tcm is not None:
+        lh5.write(tcm, "tcm", out_file, wo_mode="of")
+
+
+def get_tcm_from_ak(
+    hit_data: list[ak.Array],
+    channels: list[int],
+    *,
+    window: float = 10,
+    time_name: str = "t0",
+    idx_name: str = "global_evtid",
+) -> Table:
+    """Builds a time-coincidence map from a hit of hit data Tables.
+
+    - build an ak.Array of the data merging channels with fields base on "time_name", and "idx_name" and adding a field `rawid` from the channel idx, also add the row (`hit_idx`)
+    - sorts this array by "idx_name" then "time_name" fields
+    - group by "idx_name" and "time_name" based on the window parameter
+
+    Parameters
+    ----------
+    hit_data
+        list of hit tier data for each channel
+    channels
+        list of channel indices
+    window
+        time window for selecting coincidences (in us)
+    time_name
+        name of the field for time information
+    idx_name
+        name of the decay index field
+
+    Returns
+    -------
+    an LGDO.VectorOfVectors containing the time-coincidence map
+    """
+    # build ak_obj for sorting
+    sort_objs = []
+
+    for ch_idx, data_tmp in zip(channels, hit_data):
+        obj_tmp = ak.copy(data_tmp)
+        obj_tmp = obj_tmp[[time_name, idx_name]]
+        hit_idx = ak.local_index(obj_tmp)
+
+        obj_tmp = ak.with_field(obj_tmp, hit_idx, "array_idx")
+
+        obj_tmp["array_id"] = int(ch_idx)
+        sort_objs.append(obj_tmp)
+
+    obj_tot = ak.concatenate(sort_objs)
+
+    return shape.group.group_by_time(
+        obj_tot,
+        time_name=time_name,
+        evtid_name=idx_name,
+        window=window,
+        fields=["array_id", "array_idx"],
+    )

reboost/core.py

@@ -104,6 +104,9 @@ def evaluate_object(
     -------
     the evaluated object.
     """
+    msg = f"Evaluating object with expression {expression} and {local_dict}"
+    log.debug(msg)
+
     func_call, globals_dict = utils.get_function_string(expression)
 
     if "pyg4ometry" in globals_dict:

reboost/iterator.py

@@ -131,5 +131,5 @@ class GLMIterator:
                 vert_rows = None
             # vertex table should have same structure as glm
 
-            return (stp_rows, vert_rows, self.current_i_entry, n_steps)
+            return (stp_rows, vert_rows, self.current_i_entry - 1, n_steps)
         return (None, None, self.current_i_entry, 0)

reboost/optmap/__init__.py

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

reboost/optmap/cli.py

@@ -82,9 +82,11 @@ def optical_cli() -> None:
         help="default: %(default)s",
     )
     mapview_parser.add_argument(
-        "--display-error",
-        action="store_true",
-        help="display error instead of magnitude. default: %(default)s",
+        "--hist",
+        choices=("nr_gen", "nr_det", "p_det", "p_det_err", "p_det_err_rel"),
+        action="store",
+        default="p_det",
+        help="select optical map histogram to show. default: %(default)s",
     )
     mapview_parser.add_argument(
         "--divide",
@@ -168,6 +170,12 @@ def optical_cli() -> None:
     )
     convolve_parser.add_argument("--output", help="output hit LH5 file", metavar="OUTPUT_HIT")
 
+    # 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)]
@@ -229,7 +237,7 @@ def optical_cli() -> None:
             cmap_min=args.min,
             cmap_max=args.max,
             title=args.title,
-            show_error=args.display_error,
+            histogram_choice=args.hist,
         )
 
     # STEP 2c: merge maps
@@ -261,3 +269,11 @@ def optical_cli() -> None:
         _check_input_file(parser, [args.map, args.edep])
         _check_output_file(parser, args.output)
         convolve(args.map, args.edep, args.edep_lgdo, args.material, args.output, args.bufsize)
+
+    # STEP X: rebin maps
+    if args.command == "rebin":
+        from reboost.optmap.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)

reboost/optmap/create.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 import gc
 import logging
 import multiprocessing as mp
@@ -430,7 +431,7 @@ def check_optical_map(map_l5_file: str):
     for submap in list_optical_maps(map_l5_file):
         try:
             om = OpticalMap.load_from_file(map_l5_file, submap)
-        except BaseException:
+        except Exception:
             log.exception("error while loading optical map %s", submap)
             continue
         om.check_histograms(include_prefix=True)
@@ -439,3 +440,42 @@ def check_optical_map(map_l5_file: str):
             log.error("edges of optical map %s differ", submap)
         else:
             all_binning = om.binning
+
+
+def rebin_optical_maps(map_l5_file: str, output_lh5_file: str, factor: int):
+    """Rebin the optical map by an integral factor.
+
+    .. note ::
+
+        the factor has to divide the bincounts on all axes.
+    """
+    if not isinstance(factor, int) or factor <= 1:
+        msg = f"invalid rebin factor {factor}"
+        raise ValueError(msg)
+
+    def _rebin_map(large: NDArray, factor: int) -> NDArray:
+        factor = np.full(3, factor, dtype=int)
+        sh = np.column_stack([np.array(large.shape) // factor, factor]).ravel()
+        return large.reshape(sh).sum(axis=(1, 3, 5))
+
+    for submap in list_optical_maps(map_l5_file):
+        log.info("rebinning optical map group: %s", submap)
+
+        om = OpticalMap.load_from_file(map_l5_file, submap)
+
+        settings = om.get_settings()
+        if not all(b % factor == 0 for b in settings["bins"]):
+            msg = f"invalid factor {factor}, not a divisor"
+            raise ValueError(msg)
+        settings = copy.copy(settings)
+        settings["bins"] = [b // factor for b in settings["bins"]]
+
+        om_new = OpticalMap.create_empty(om.name, settings)
+        om_new.h_vertex = _rebin_map(om.h_vertex, factor)
+        om_new.h_hits = _rebin_map(om.h_hits, factor)
+        om_new.create_probability()
+        om_new.write_lh5(lh5_file=output_lh5_file, group=submap, wo_mode="write_safe")
+
+    # just copy hitcounts exponent.
+    for dset in ("_hitcounts_exp", "_hitcounts"):
+        lh5.write(lh5.read(dset, lh5_file=map_l5_file), dset, lh5_file=output_lh5_file)

reboost/optmap/evt.py

@@ -60,38 +60,40 @@ def build_optmap_evt(
     # This function follows the assumption, that the output event ids are at least "somewhat"
     # monotonic, i.e. later chunks do not contain lower evtids than the previous chunk(s).
     # Going back is not implemented.
-    def _ensure_vert_df(vert_it: LH5Iterator, evtid: int):
+    def _ensure_vert_df(vert_it: LH5Iterator, evtid: int) -> None:
         nonlocal vert_df, vert_df_bounds, vert_it_count, hits_expected
 
-        if vert_df_bounds is not None and vert_df is not None:
-            if evtid < vert_df_bounds[0]:
-                msg = "non-monotonic evtid encountered, but cannot go back"
-                raise KeyError(msg)
-            if evtid >= vert_df_bounds[0] and evtid <= vert_df_bounds[1]:
-                return  # vert_df already contains the given evtid.
-
-        # here, evtid > vert_df_bounds[1] (or vert_df_bounds is still None). We need to fetch
-        # the next event table chunk.
-
-        vert_it_count += 1
-        # we might have filled a dataframe, save it to disk.
-        _store_vert_df()
-
-        # read the next vertex chunk into memory.
-        (vert_lgdo, vert_entry, vert_n_rows) = next(vert_it)
-        vert_df = vert_lgdo.view_as("pd").iloc[0:vert_n_rows]
-
-        # prepare vertex coordinates.
-        vert_df = vert_df.set_index("evtid", drop=True).drop(["n_part", "time"], axis=1)
-        vert_df_bounds = [vert_df.index.min(), vert_df.index.max()]
-        hits_expected = 0
-        # add columns for all detectors.
-        for d in detectors:
-            vert_df[d] = hit_count_type(0)
+        # skipping multiple chunks is possible in sparsely populated simulations.
+        while vert_df_bounds is None or evtid > vert_df_bounds[1] or evtid < vert_df_bounds[0]:
+            if vert_df_bounds is not None and vert_df is not None:
+                if evtid < vert_df_bounds[0]:
+                    msg = "non-monotonic evtid encountered, but cannot go back"
+                    raise KeyError(msg)
+                if evtid >= vert_df_bounds[0] and evtid <= vert_df_bounds[1]:
+                    return  # vert_df already contains the given evtid.
+
+            # here, evtid > vert_df_bounds[1] (or vert_df_bounds is still None). We need to fetch
+            # the next event table chunk.
+
+            vert_it_count += 1
+            # we might have filled a dataframe, save it to disk.
+            _store_vert_df()
+
+            # read the next vertex chunk into memory.
+            (vert_lgdo, vert_entry, vert_n_rows) = next(vert_it)
+            vert_df = vert_lgdo.view_as("pd").iloc[0:vert_n_rows]
+
+            # prepare vertex coordinates.
+            vert_df = vert_df.set_index("evtid", drop=True).drop(["n_part", "time"], axis=1)
+            vert_df_bounds = [vert_df.index.min(), vert_df.index.max()]
+            hits_expected = 0
+            # add columns for all detectors.
+            for d in detectors:
+                vert_df[d] = hit_count_type(0)
 
     log.info("prepare evt table")
-    # use smaller integer type uint16 to spare RAM when storing types.
-    hit_count_type = np.uint16
+    # use smaller integer type uint8 to spare RAM when storing types.
+    hit_count_type = np.uint8
     for opti_it_count, (opti_lgdo, opti_entry, opti_n_rows) in enumerate(opti_it):
         assert (opti_it_count == 0) == (opti_entry == 0)
         opti_df = opti_lgdo.view_as("pd").iloc[0:opti_n_rows]

reboost/optmap/mapview.py

@@ -92,12 +92,13 @@ def _channel_selector(fig) -> None:
 def _read_data(
     optmap_fn: str,
     detid: str = "all",
-    show_error: bool = False,
+    histogram_choice: str = "p_det",
 ) -> tuple[tuple[NDArray], NDArray]:
-    optmap_all = lh5.read(f"/{detid}/p_det", optmap_fn)
+    histogram = histogram_choice if histogram_choice != "p_det_err_rel" else "p_det"
+    optmap_all = lh5.read(f"/{detid}/{histogram}", optmap_fn)
     optmap_edges = tuple([b.edges for b in optmap_all.binning])
     optmap_weights = optmap_all.weights.nda.copy()
-    if show_error:
+    if histogram_choice == "p_det_err_rel":
         optmap_err = lh5.read(f"/{detid}/p_det_err", optmap_fn)
         divmask = optmap_weights > 0
         optmap_weights[divmask] = optmap_err.weights.nda[divmask] / optmap_weights[divmask]
@@ -111,13 +112,13 @@ def _prepare_data(
     divide_fn: str | None = None,
     cmap_min: float | Literal["auto"] = 1e-4,
     cmap_max: float | Literal["auto"] = 1e-2,
-    show_error: bool = False,
+    histogram_choice: str = "p_det",
     detid: str = "all",
 ) -> tuple[tuple[NDArray], NDArray]:
-    optmap_edges, optmap_weights = _read_data(optmap_fn, detid, show_error)
+    optmap_edges, optmap_weights = _read_data(optmap_fn, detid, histogram_choice)
 
     if divide_fn is not None:
-        divide_edges, divide_map = _read_data(divide_fn, detid, show_error)
+        divide_edges, divide_map = _read_data(divide_fn, detid, histogram_choice)
         divmask = divide_map > 0
         optmap_weights[divmask] = optmap_weights[divmask] / divide_map[divmask]
         optmap_weights[~divmask] = -1
@@ -157,12 +158,12 @@ def view_optmap(
     start_axis: int = 2,
     cmap_min: float | Literal["auto"] = 1e-4,
     cmap_max: float | Literal["auto"] = 1e-2,
-    show_error: bool = False,
+    histogram_choice: str = "p_det",
     title: str | None = None,
 ) -> None:
     available_dets = list_optical_maps(optmap_fn)
 
-    prepare_args = (optmap_fn, divide_fn, cmap_min, cmap_max, show_error)
+    prepare_args = (optmap_fn, divide_fn, cmap_min, cmap_max, histogram_choice)
     edges, weights, cmap_min, cmap_max = _prepare_data(*prepare_args, detid)
 
     fig = plt.figure(figsize=(10, 10))

reboost/optmap/optmap.py

@@ -48,6 +48,7 @@ class OpticalMap:
                 for i in range(3)
             ]
 
+    @staticmethod
     def create_empty(name: str, settings: Mapping[str, str]) -> OpticalMap:
         om = OpticalMap(name, settings)
         om.h_vertex = om._prepare_hist()
@@ -115,7 +116,7 @@ class OpticalMap:
             assert ax.is_range
             assert ax.closedleft
             oor_mask &= (ax.first <= col) & (col < ax.last)
-            idx_s = np.floor((col - ax.first).astype(np.float64) / ax.step).astype(np.int64)
+            idx_s = np.floor((col.astype(np.float64) - ax.first) / ax.step).astype(np.int64)
             assert np.all(idx_s[oor_mask] < self._single_shape[dim])
             idx += s * idx_s
 
@@ -215,9 +216,11 @@ class OpticalMap:
         return ratio_0, ratio_err_0
 
     def create_probability(self) -> None:
+        """Compute probability map (and map uncertainty) from vertex and hit map."""
         self.h_prob, self.h_prob_uncert = self._divide_hist(self.h_hits, self.h_vertex)
 
     def write_lh5(self, lh5_file: str, group: str = "all", wo_mode: str = "write_safe") -> None:
+        """Write this map to a LH5 file."""
         if wo_mode not in ("write_safe", "overwrite_file"):
             msg = f"invalid wo_mode {wo_mode} for optical map"
             raise ValueError(msg)
@@ -237,6 +240,25 @@ class OpticalMap:
         write_hist(self.h_prob, "p_det", lh5_file, group, "write_safe")
         write_hist(self.h_prob_uncert, "p_det_err", lh5_file, group, "write_safe")
 
+    def get_settings(self) -> dict:
+        """Get the binning settings that were used to create this optical map instance."""
+        if self.settings is not None:
+            return self.settings
+
+        range_in_m = []
+        bins = []
+        for b in self.binning:
+            if not b.is_range:
+                msg = "cannot get binning settings for variable binning map"
+                raise RuntimeError(msg)
+            if b.get_binedgeattrs().get("units") != "m":
+                msg = "invalid units. can only work with optical maps in meter"
+                raise RuntimeError(msg)
+            range_in_m.append([b.first, b.last])
+            bins.append(b.nbins)
+
+        return {"range_in_m": np.array(range_in_m), "bins": np.array(bins)}
+
     def check_histograms(self, include_prefix: bool = False) -> None:
         log_prefix = "" if not include_prefix else self.name + " - "
 

reboost/shape/cluster.py

@@ -3,47 +3,203 @@ from __future__ import annotations
 import logging
 
 import awkward as ak
+import numba
 import numpy as np
 from lgdo import VectorOfVectors
 
 log = logging.getLogger(__name__)
 
 
-def cluster_by_sorted_field(
-    cluster_variable: ak.Array | VectorOfVectors, field: ak.Array | VectorOfVectors
+def apply_cluster(
+    cluster_run_lengths: VectorOfVectors | ak.Array, field: ak.Array | VectorOfVectors
 ) -> VectorOfVectors:
-    """Extract the cumulative lengths of the clusters per event.
+    """Apply clustering to a field.
 
-    Example processing block:
+    Parameters
+    ----------
+    cluster_ids
+        run lengths of each cluster
+    field
+        the field to cluster
+    """
+    if isinstance(cluster_run_lengths, VectorOfVectors):
+        cluster_run_lengths = cluster_run_lengths.view_as("ak")
 
-    .. code-block:: yaml
+    if isinstance(field, VectorOfVectors):
+        field = field.view_as("ak")
+
+    n_cluster = ak.num(cluster_run_lengths, axis=-1)
+    clusters = ak.unflatten(ak.flatten(field), ak.flatten(cluster_run_lengths))
+
+    # reshape into cluster oriented
+    return VectorOfVectors(ak.unflatten(clusters, n_cluster))
+
+
+def cluster_by_step_length(
+    trackid: ak.Array | VectorOfVectors,
+    pos_x: ak.Array | VectorOfVectors,
+    pos_y: ak.Array | VectorOfVectors,
+    pos_z: ak.Array | VectorOfVectors,
+    dist: ak.Array | VectorOfVectors,
+    surf_cut: float = 2,
+    threshold: float = 0.1,
+    threshold_surf: float = 0.0,
+) -> VectorOfVectors:
+    """Perform clustering based on the step length.
 
-        trackid_cluster_edep: reboost.shape.cluster.cluster_by_sorted_field(HITS.trackid,HITS.edep)
+    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.
+
+    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
+    parameters to be different for bulk and surface.
 
     Parameters
     ----------
     trackid
-        the id of the tracks.
-    field
-        another field to cluster.
+        index of the track.
+    pos_x
+        x position of the step.
+    pos_y
+        y position of the step.
+    pos_z
+        z position of the step.
+    dist
+        distance to the detector surface.
+    surf_cut
+        Size of the surface region (in mm)
+    threshold
+        Distance threshold in mm to combine steps in the bulk.
+    threshold_surf
+        Distance threshold in mm to combine steps in the surface.
 
     Returns
     -------
-    A VectorOfVectors with the clustered field, an additional axis is present due to the clustering
+    Array of the run lengths of each cluster within a hit.
     """
-    if isinstance(cluster_variable, VectorOfVectors):
-        cluster_variable = cluster_variable.view_as("ak")
+    # type conversions
+    if isinstance(pos_x, VectorOfVectors):
+        pos_x = pos_x.view_as("ak")
+
+    if isinstance(pos_y, VectorOfVectors):
+        pos_y = pos_y.view_as("ak")
+
+    if isinstance(pos_z, VectorOfVectors):
+        pos_z = pos_z.view_as("ak")
+
+    if isinstance(trackid, VectorOfVectors):
+        trackid = trackid.view_as("ak")
+
+    if isinstance(dist, VectorOfVectors):
+        dist = dist.view_as("ak")
+
+    pos = np.vstack(
+        [
+            ak.flatten(pos_x).to_numpy(),
+            ak.flatten(pos_y).to_numpy(),
+            ak.flatten(pos_z).to_numpy(),
+        ]
+    ).T
+
+    indices_flat = cluster_by_distance_numba(
+        ak.flatten(ak.local_index(trackid)).to_numpy(),
+        ak.flatten(trackid).to_numpy(),
+        pos,
+        ak.flatten(dist).to_numpy(),
+        surf_cut=surf_cut,
+        threshold=threshold,
+        threshold_surf=threshold_surf,
+    )
+
+    # reshape into being event oriented
+    indices = ak.unflatten(indices_flat, ak.num(ak.local_index(trackid)))
+
+    # number of steps per cluster
+    counts = ak.run_lengths(indices)
+
+    return VectorOfVectors(counts)
+
+
+@numba.njit
+def cluster_by_distance_numba(
+    local_index: np.ndarray,
+    trackid: np.ndarray,
+    pos: np.ndarray,
+    dist_to_surf: np.ndarray,
+    surf_cut: float = 2,
+    threshold: float = 0.1,
+    threshold_surf: float = 0.0,
+) -> np.ndarray:
+    """Cluster steps by the distance between points in the same track.
+
+    This function gives the basic numerical calculations for
+    :func:`cluster_by_step_length`.
 
-    if isinstance(field, VectorOfVectors):
-        field = cluster_variable.view_as("ak")
+    Parameters
+    ----------
+    local_index
+        1D array of the local index within each hit (step group)
+    trackid
+        1D array of index of the track
+    pos
+        `(n,3)` size array of the positions
+    dist_to_surf
+        1D array of the distance to the detector surface.
+    surf_cut
+        Size of the surface region (in mm)
+    threshold
+        Distance threshold in mm to combine steps in the bulk.
+    threshold_surf
+        Distance threshold in mm to combine steps in the surface.
+
+    Returns
+    -------
+    np.ndarray
+        1D array of cluster indices
+    """
 
-    # run length of each cluster
-    counts = ak.run_lengths(cluster_variable)
+    def _dist(a, b):
+        return np.sqrt(np.sum((a - b) ** 2))
 
-    n_cluster = ak.num(counts, axis=-1)
-    clusters = ak.unflatten(ak.flatten(field), ak.flatten(counts))
+    n = len(local_index)
+    out = np.zeros(n, dtype=numba.int32)
 
-    return VectorOfVectors(ak.unflatten(clusters, n_cluster))
+    trackid_prev = -1
+    pos_prev = np.zeros(3, dtype=numba.float64)
+    cluster_idx = -1
+    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)
+        )
+
+        # New hit, reset cluster index
+        if idx == 0 or local_index[idx] == 0:
+            cluster_idx = 0
+            pos_prev = pos[idx]
+
+        # either new track, moving from surface to bulk,
+        # moving from bulk to surface, or stepping more than
+        # the threshold. Start a new cluster.
+        elif new_cluster:
+            cluster_idx += 1
+            pos_prev = pos[idx, :]
+
+        out[idx] = cluster_idx
+
+        # Update previous values
+        trackid_prev = trackid[idx]
+        is_surf_prev = dist_to_surf[idx] < surf_cut
+
+    return out
 
 
 def step_lengths(

reboost/shape/group.py

@@ -25,8 +25,11 @@ def _sort_data(obj: ak.Array, *, time_name: str = "time", evtid_name: str = "evt
     -------
     sorted awkward array
     """
-    indices = np.lexsort((obj[time_name], obj[evtid_name]))
-    return obj[indices]
+    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]
+
+    return ak.flatten(sorted_obj)
 
 
 def group_by_evtid(data: Table | ak.Array, *, evtid_name: str = "evtid") -> Table:

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: reboost
-Version: 0.1.1
+Version: 0.2.1
 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
@@ -707,6 +707,7 @@ Requires-Dist: dbetto
 Requires-Dist: particle
 Requires-Dist: pandas
 Requires-Dist: matplotlib
+Requires-Dist: pygama
 Provides-Extra: all
 Requires-Dist: reboost[docs,test]; extra == "all"
 Provides-Extra: docs
@@ -722,6 +723,7 @@ Requires-Dist: legend-pygeom-hpges; extra == "test"
 Requires-Dist: legend-pygeom-tools; extra == "test"
 Requires-Dist: pyg4ometry; extra == "test"
 Requires-Dist: pylegendtestdata; extra == "test"
+Dynamic: license-file
 
 # reboost
 
@@ -735,4 +737,17 @@ Requires-Dist: pylegendtestdata; extra == "test"
 ![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)
 
-New LEGEND Monte-Carlo simulation post-processing [WIP]
+_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/)!

reboost-0.2.1.dist-info/RECORD

@@ -0,0 +1,36 @@
+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/cli.py,sha256=HTZ05DRnDodcf_D6BJCCavx5HqhKDadJCgf-oh8HTJk,6365
+reboost/core.py,sha256=AamREubQsAqJ-y10NRx18r-PuqlQj3iTl2PzOTWXGQI,10540
+reboost/iterator.py,sha256=cqfh3c0uLP67S0YGaw05-McZQzdMb8BISULIm3PEbKA,3990
+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/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/stats.py,sha256=iiOEi87x93kqPWeSmlRiA5Oe-R8XR-plm6Z532PhC9M,1401
+reboost/optmap/__init__.py,sha256=imvuyld-GLw8qdwqW-lXCg2feptcTyQo3wIzPvDHwmY,93
+reboost/optmap/cli.py,sha256=wBexh-zrr5ABherEyk9xigxdArvOAKiiRQwAYon9Sro,9408
+reboost/optmap/convolve.py,sha256=5FksUrVIG8ysn42QbWBrAx8M1HfAVJtaJJyE8oJ1NGM,12043
+reboost/optmap/create.py,sha256=Nm5-xEe8M9q2GFQnUv8oN8qpAz9nZArIrQcPboqRmCQ,17153
+reboost/optmap/evt.py,sha256=m3NWuLEk4zDQJO5vXq_XLLnqmkQwmtdKI3fqmZQBBvc,4707
+reboost/optmap/mapview.py,sha256=73kpe0_SKDj9bIhEx1ybX1sBP8TyvufiLfps84A_ijA,6798
+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/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.1.1.dist-info → reboost-0.2.1.dist-info}/WHEEL

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

reboost-0.1.1.dist-info/RECORD

@@ -1,35 +0,0 @@
-reboost/__init__.py,sha256=yMYemxriXanK9QewWEIKVh5JbOzmnQhS_B4EjW494R0,242
-reboost/_version.py,sha256=PKIMyjdUACH4-ONvtunQCnYE2UhlMfp9su83e3HXl5E,411
-reboost/build_evt.py,sha256=jCmczvzk1aJ8-EWTz4WCcx0FyDqY6qtCwoLiEX-Yc0I,122
-reboost/build_glm.py,sha256=NQzdV6EFQkWjVcXbdMoOJQt6KmkJ6XWGW133xtUMtYA,9295
-reboost/build_hit.py,sha256=o7iR5SLA9fM3zrJ1M_GWSx6ONeM_nVXDYb5jK3HtxiU,13249
-reboost/cli.py,sha256=HTZ05DRnDodcf_D6BJCCavx5HqhKDadJCgf-oh8HTJk,6365
-reboost/core.py,sha256=6dxron2cUVqNr-UUh2MHQ4HVGWFVmmR_vAPhv-7XKHw,10443
-reboost/iterator.py,sha256=E0nn2Sg_zXe05affnoKr0-z3Y-eQl70tgT687UGgmy0,3986
-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/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/stats.py,sha256=iiOEi87x93kqPWeSmlRiA5Oe-R8XR-plm6Z532PhC9M,1401
-reboost/optmap/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-reboost/optmap/cli.py,sha256=GQTRrLO1fBxzpUPl560DWbxsUIYmuuStzNN2dTor97Y,8675
-reboost/optmap/convolve.py,sha256=5FksUrVIG8ysn42QbWBrAx8M1HfAVJtaJJyE8oJ1NGM,12043
-reboost/optmap/create.py,sha256=BJuAJ_U9iZ0V5-CCxmrR30JovqUwu_GAWkBSWL-N5bM,15616
-reboost/optmap/evt.py,sha256=RmEOap9gnclgT5O0ls097TO5HDkWXcS9JzY7MqseWbE,4439
-reboost/optmap/mapview.py,sha256=lSB8brbL3rQPik7gCCwkjXXv-qmEOTsVoCPaoWwWwCM,6643
-reboost/optmap/numba_pdg.py,sha256=y8cXR5PWE2Liprp4ou7vl9do76dl84vXU52ZJD9_I7A,731
-reboost/optmap/optmap.py,sha256=NvxIdu28qOSHemfg0lFYF7DcTZAse9mjZQyMq3a3feY,11913
-reboost/shape/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-reboost/shape/cluster.py,sha256=T9YoP7Tx0uS1_nhE3PZ80LcutXw-GuWSFMCP_pXYvCA,3029
-reboost/shape/group.py,sha256=CSA2OG4rg5Qwea2qfpMRy5UtDDUFvNknXwfXNJeFn7c,4293
-reboost/shape/reduction.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-reboost-0.1.1.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-reboost-0.1.1.dist-info/METADATA,sha256=SrHP-kabqtTe1kk_9-rz8mIU1oQmzMUqHRqSnFu4Olg,43645
-reboost-0.1.1.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-reboost-0.1.1.dist-info/entry_points.txt,sha256=DxhD6BidSWNot9BrejHJjQ7RRLmrMaBIl52T75oWTwM,93
-reboost-0.1.1.dist-info/top_level.txt,sha256=q-IBsDepaY_AbzbRmQoW8EZrITXRVawVnNrB-_zyXZs,8
-reboost-0.1.1.dist-info/RECORD,,