reboost 0.2.2__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.2'
-__version_tuple__ = version_tuple = (0, 2, 2)
+__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
@@ -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.
@@ -294,19 +294,25 @@ def build_hit(
                     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],
@@ -314,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,
@@ -327,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 (
@@ -338,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/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/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.2.dist-info → reboost-0.2.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reboost
-Version: 0.2.2
+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.2.dist-info → reboost-0.2.3.dist-info}/RECORD

@@ -1,15 +1,15 @@
 reboost/__init__.py,sha256=RVNl3Qgx_hTUeBGXaWYmiTcmXUDhTfvlAGGC8bo_jP8,316
-reboost/_version.py,sha256=OjGGK5TcHVG44Y62aAqeJH4CskkZoY9ydbHOtCDew50,511
-reboost/build_evt.py,sha256=5Q3T0LCl8xMtyRRhcs6layC1xh4vp2f26PgB1yab2zs,4798
-reboost/build_glm.py,sha256=5cuHnMogRKGaVCi4eqTrgDroLFy4s8mGp2nKcKfKQGI,9243
-reboost/build_hit.py,sha256=yAjYmBJUZvnyTuU_kpuJle_TKeSU027OIJr7No9z2Ms,13476
-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=7Nclc6RUCOSJ1CWVAX0rFNJGM1LEgqvc4tD04CxEAtg,10766
-reboost/iterator.py,sha256=cqfh3c0uLP67S0YGaw05-McZQzdMb8BISULIm3PEbKA,3990
+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
@@ -26,11 +26,11 @@ reboost/optmap/numba_pdg.py,sha256=y8cXR5PWE2Liprp4ou7vl9do76dl84vXU52ZJD9_I7A,7
 reboost/optmap/optmap.py,sha256=j4rfbQ84PYSpE-BvP4Rdt96ZjPdwy8P4e4eZz1mATys,12817
 reboost/shape/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 reboost/shape/cluster.py,sha256=RIvBlhHzp88aaUZGofp5SD9bimnoiqIOddhQ84jiwoM,8135
-reboost/shape/group.py,sha256=bSmFCl_yi1hGaKudjiicDEJsiBNyAHiKYdr8ZuH4pSM,4406
+reboost/shape/group.py,sha256=Q3DhEPxbhw3p4bwvpswSd0A-p224l5vRZnfQIEkOVJE,4475
 reboost/shape/reduction.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-reboost-0.2.2.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-reboost-0.2.2.dist-info/METADATA,sha256=F4NsVNDavleMAL-xEYSpR5ElKHtjQXLevD_-XAaT3y4,44219
-reboost-0.2.2.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-reboost-0.2.2.dist-info/entry_points.txt,sha256=DxhD6BidSWNot9BrejHJjQ7RRLmrMaBIl52T75oWTwM,93
-reboost-0.2.2.dist-info/top_level.txt,sha256=q-IBsDepaY_AbzbRmQoW8EZrITXRVawVnNrB-_zyXZs,8
-reboost-0.2.2.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.2.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