reboost 0.8.3__py3-none-any.whl

reboost/__init__.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+import hdf5plugin
+from lgdo import lh5
+
+from ._version import version as __version__
+from .build_hit import build_hit
+
+__all__ = [
+    "__version__",
+    "build_hit",
+]
+
+lh5.settings.DEFAULT_HDF5_SETTINGS = {"compression": hdf5plugin.Zstd()}

reboost/_version.py

@@ -0,0 +1,34 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '0.8.3'
+__version_tuple__ = version_tuple = (0, 8, 3)
+
+__commit_id__ = commit_id = None

reboost/build_evt.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import logging
+
+import awkward as ak
+import numpy as np
+from dbetto import AttrsDict
+from lgdo import Array, Table, VectorOfVectors, lh5
+
+from . import core, math, shape, utils
+from .shape import group
+
+log = logging.getLogger(__name__)
+
+
+def build_evt(
+    tcm: VectorOfVectors,
+    hitfile: str,
+    outfile: str | None,
+    channel_groups: AttrsDict,
+    pars: AttrsDict,
+    run_part: AttrsDict,
+) -> Table | None:
+    """Build events out of a TCM.
+
+    Parameters
+    ----------
+    tcm
+        the time coincidence map.
+    hitfile
+        file with the hits.
+    outfile
+        the path to the output-file, if `None` with return
+        the events in memory.
+    channel_groups
+        a dictionary of groups of channels. For example:
+
+        .. code-block:: python
+
+            {"det1": "on", "det2": "off", "det3": "ac"}
+
+    pars
+        A dictionary of parameters. The first key should
+        be the run ID, followed by different sets of parameters
+        arranged in groups. Run numbers should be given in the
+        format `"p00-r001"`, etc.
+
+        For example:
+
+        .. code-block:: python
+
+            {"p03-r000": {"reso": {"det1": [1, 2], "det2": [0, 1]}}}
+
+    run_part
+        The run partitioning file giving the number of events
+        for each run. This should be organized as a dictionary
+        with the following format:
+
+        .. code-block:: python
+
+            {"p03-r000": 1000, "p03-r001": 2000}
+
+    Returns
+    -------
+    the event file in memory as a table if no output file is specified.
+    """
+    tcm_tables = utils.get_table_names(tcm)
+    tcm_ak = tcm.view_as("ak")
+
+    # loop over the runs
+    cum_sum = 0
+    tab = None
+
+    for idx, (run_full, n_event) in enumerate(run_part.items()):
+        period, run = run_full.split("-")
+        pars_tmp = pars[run_full]
+
+        # create an output table
+        out_tab = Table(size=n_event)
+
+        tcm_tmp = tcm_ak[cum_sum : cum_sum + n_event]
+
+        # usabilities
+
+        is_off = shape.group.get_isin_group(
+            tcm_tmp.table_key, channel_groups, tcm_tables, group="off"
+        )
+
+        # filter out off channels
+        channels = tcm_tmp.table_key[~is_off]
+        rows = tcm_tmp.row_in_table[~is_off]
+        out_tab.add_field("channel", VectorOfVectors(channels))
+        out_tab.add_field("row_in_table", VectorOfVectors(rows))
+
+        out_tab.add_field("period", Array(np.ones(len(channels)) * int(period[1:])))
+        out_tab.add_field("run", Array(np.ones(len(channels)) * int(run[1:])))
+
+        # now check for channels in ac
+        is_good = group.get_isin_group(channels, channel_groups, tcm_tables, group="on")
+
+        # get energy
+        energy_true = core.read_data_at_channel_as_ak(
+            channels, rows, hitfile, "energy", "hit", tcm_tables
+        )
+
+        energy = math.stats.apply_energy_resolution(
+            energy_true,
+            channels,
+            tcm_tables,
+            pars_tmp.reso,
+            lambda energy, sig0, sig1: np.sqrt(energy * sig1**2 + sig0**2),
+        )
+
+        out_tab.add_field("is_good", VectorOfVectors(is_good[energy > 25]))
+
+        out_tab.add_field("energy", VectorOfVectors(energy[energy > 25]))
+        out_tab.add_field("multiplicity", Array(ak.sum(energy > 25, axis=-1).to_numpy()))
+
+        # write table
+        wo_mode = "of" if idx == 0 else "append"
+
+        # add attrs
+        out_tab.attrs["tables"] = tcm.attrs["tables"]
+
+        if outfile is not None:
+            lh5.write(out_tab, "evt", outfile, wo_mode=wo_mode)
+        else:
+            tab = (
+                ak.concatenate((tab, out_tab.view_as("ak")))
+                if tab is not None
+                else out_tab.view_as("ak")
+            )
+
+    return Table(tab)

reboost/build_glm.py

@@ -0,0 +1,305 @@
+from __future__ import annotations
+
+import copy
+import logging
+
+import awkward as ak
+import numpy as np
+from lgdo import Array, Table, lh5
+from lgdo.lh5 import LH5Iterator, LH5Store
+from numpy.typing import ArrayLike
+
+from . import utils
+
+log = logging.getLogger(__name__)
+
+
+def get_glm_rows(stp_evtids: ArrayLike, vert: ArrayLike, *, start_row: int = 0) -> ak.Array:
+    """Get the rows of the Geant4 event lookup map (glm).
+
+    Parameters
+    ----------
+    stp_evtids
+        Array of evtids for the steps
+    vert
+        Array of simulated evtid for the vertices.
+    start_row
+        The index of the first element of stp_evtids.
+
+    Returns
+    -------
+    an awkward array of the `glm`.
+    """
+    # convert inputs
+    if stp_evtids is None:
+        output = ak.Array({"evtid": vert})
+        output["n_rows"] = np.array([0] * len(vert), dtype=float)
+        output["start_row"] = np.array([np.nan] * len(vert), dtype=float)
+
+        return output
+
+    if not isinstance(stp_evtids, np.ndarray):
+        stp_evtids = (
+            stp_evtids.to_numpy() if isinstance(stp_evtids, ak.Array) else np.array(stp_evtids)
+        )
+
+    if not isinstance(vert, np.ndarray):
+        vert = vert.to_numpy() if isinstance(vert, ak.Array) else np.array(vert)
+
+    # check that the steps and vertices are sorted or the algorithm will fail
+
+    if not np.all(vert[:-1] <= vert[1:]):
+        msg = "The vertices must be sorted"
+        raise ValueError(msg)
+
+    if len(stp_evtids) > 0 and not np.all(stp_evtids[:-1] <= stp_evtids[1:]):
+        msg = "The steps must be sorted"
+        raise ValueError(msg)
+
+    # convert to numpy
+    stp_indices = np.arange(len(stp_evtids)) + start_row
+
+    # cut the arrays
+    sel = (stp_evtids >= vert[0]) & (stp_evtids <= vert[-1])
+    stp_evtids = stp_evtids[sel]
+    stp_indices = stp_indices[sel]
+
+    # build output
+    output = ak.Array({"evtid": vert})
+
+    # restructure to jagged array
+    counts = ak.run_lengths(stp_evtids)
+    steps = ak.unflatten(stp_evtids, counts)
+    indices = ak.unflatten(stp_indices, counts)
+
+    ak_tmp = ak.Array({"evtid": ak.fill_none(ak.firsts(steps), np.nan), "indices": indices})
+
+    # find indices to insert the new entries
+    positions = np.searchsorted(output.evtid, ak_tmp.evtid)
+
+    # add a check that every ak_tmp evtid is in output.evtid ?
+    if not np.all(np.isin(ak_tmp.evtid, output.evtid)):
+        bad_evtid = ak_tmp.evtid[~np.isin(ak_tmp.evtid, output.evtid)]
+        msg = f"Error not every evtid in the stp table is present in the vertex table {bad_evtid} are not"
+        raise ValueError(msg)
+
+    # get the start row
+    start_row = np.array([np.nan] * len(vert), dtype=float)
+    start_row[positions] = ak.fill_none(ak.firsts(ak_tmp.indices), np.nan)
+
+    n_row = np.array([0] * len(vert), dtype=float)
+    n_row[positions] = ak.num(ak_tmp.indices)
+
+    # add to the  output
+    output["n_rows"] = n_row
+    output["start_row"] = start_row
+
+    return output
+
+
+def get_stp_evtids(
+    lh5_table: str,
+    stp_file: str,
+    id_name: str,
+    start_row: int,
+    last_vertex_evtid: int,
+    stp_buffer: int,
+) -> tuple[int, int, ak.Array]:
+    """Extracts the rows of a stp file corresponding to a particular range of `evtid`.
+
+    The reading starts at `start_row` to allow for iterating through
+    the file. The iteration stops when the `evtid` being read is larger than
+    `last_vertex_evtid`.
+
+    Parameters
+    ----------
+    lh5_table
+        the table name to read.
+    stp_file
+        the file name path.
+    id_name
+        the name of the `evtid` field.
+    start_row
+        the row to begin reading.
+    last_vertex_evtid
+        the last evtid to read up to.
+    stp_buffer
+        the number of rows to read at once.
+
+    Returns
+    -------
+         a tuple of the updated `start_row`, the first row for the chunk and an awkward Array of the steps.
+    """
+    # make a LH5Store
+    store = LH5Store()
+
+    # some outputs
+    evtids_proc = None
+    last_evtid = 0
+    chunk_start = 0
+
+    # get the total number of rows
+    n_rows_tot = store.read_n_rows(f"{lh5_table}/{id_name}", stp_file)
+
+    # iterate over the file
+    # stop when the entire file is read
+
+    while start_row < n_rows_tot:
+        # read the file
+        lh5_obj = store.read(
+            f"{lh5_table}/{id_name}",
+            stp_file,
+            start_row=start_row,
+            n_rows=stp_buffer,
+        )
+        evtids = lh5_obj.view_as("ak")
+        n_read = len(evtids)
+
+        # pick the first evtid
+        if evtids.ndim > 1:
+            evtids = ak.fill_none(ak.firsts(evtids, axis=-1), -1)
+
+        # if the evtids_proc is not set then this is the first valid chunk
+        if evtids_proc is None:
+            evtids_proc = evtids
+            chunk_start = start_row
+        elif evtids[0] <= last_vertex_evtid:
+            evtids_proc = ak.concatenate((evtids_proc, evtids))
+
+        # get the last evtid
+        last_evtid = evtids[-1]
+
+        # if the last evtid is greater than the last vertex we can stop reading
+        if last_evtid > last_vertex_evtid or (start_row + n_read >= n_rows_tot):
+            break
+
+        # increase rhe start row for the next read
+
+        if start_row + n_read <= n_rows_tot:
+            start_row += n_read
+
+    return start_row, chunk_start, evtids_proc
+
+
+def build_glm(
+    stp_files: str | list[str],
+    glm_files: str | list[str] | None,
+    lh5_groups: list | None = None,
+    *,
+    out_table_name: str = "glm",
+    id_name: str = "g4_evtid",
+    evtid_buffer: int = int(1e7),
+    stp_buffer: int = int(1e7),
+) -> ak.Array | None:
+    """Builds a g4_evtid look up (glm) from the stp data.
+
+    This object is used by `reboost` to efficiency iterate through the data.
+    It consists of a :class:`lgdo.VectorOfVectors` for each lh5_table in the input files.
+    The rows of this :class:`lgdo.VectorOfVectors` correspond to the `id_name` while the data
+    are the `stp` indices for this event.
+
+    Parameters
+    ----------
+    stp_files
+        path to the stp (input) file.
+    glm_files
+        path to the glm data, can also be `None` in which case an `ak.Array` is returned in memory.
+    out_table_name
+        name for the output table.
+    id_name
+        name of the evtid file, default `g4_evtid`.
+    stp_buffer
+        the number of rows of the step file to read at a time
+    evtid_buffer
+        the number of evtids to read at a time
+
+    Returns
+    -------
+    either `None` or an `ak.Array`
+    """
+    store = LH5Store()
+    files = utils.get_file_dict(stp_files=stp_files, glm_files=glm_files)
+
+    # loop over files
+    glm_sum = {}
+
+    for file_idx, stp_file in enumerate(files.stp):
+        msg = f"start generating glm for {stp_file} "
+        log.debug(msg)
+
+        # loop over the lh5_tables
+        lh5_table_list = [
+            det
+            for det in lh5.ls(stp_file, "stp/")
+            if lh5_groups is None or det.split("/")[1] in lh5_groups
+        ]
+
+        # get rows in the table
+        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
+        else:
+            glm_sum = None
+
+        # start row for each table
+        start_row = dict.fromkeys(lh5_table_list, 0)
+
+        vfield = f"vtx/{id_name}"
+
+        # iterate over the vertex table
+        for vidx, vert_obj in enumerate(LH5Iterator(stp_file, vfield, buffer_len=evtid_buffer)):
+            # range of vertices
+            vert_ak = vert_obj.view_as("ak")
+
+            msg = f"... read chunk {vidx}"
+            log.debug(msg)
+
+            for idx, lh5_table in enumerate(lh5_table_list):
+                # create the output table
+                out_tab = Table(size=len(vert_ak))
+
+                # read the stp rows starting from `start_row` until the
+                # evtid is larger than that in the vertices
+
+                start_row_tmp, chunk_row, evtids = get_stp_evtids(
+                    lh5_table,
+                    stp_file,
+                    id_name,
+                    start_row[lh5_table],
+                    last_vertex_evtid=vert_ak[-1],
+                    stp_buffer=stp_buffer,
+                )
+
+                # set the start row for the next chunk
+                start_row[lh5_table] = start_row_tmp
+
+                # now get the glm rows
+                glm = get_glm_rows(evtids, vert_ak, start_row=chunk_row)
+
+                for field in ["evtid", "n_rows", "start_row"]:
+                    out_tab.add_field(field, Array(glm[field].to_numpy()))
+
+                # write the output file
+                mode = "overwrite_file" if (vidx == 0 and idx == 0) else "append"
+
+                lh5_subgroup = lh5_table.replace("stp/", "")
+
+                if files.glm[file_idx] is not None:
+                    store.write(
+                        out_tab,
+                        f"{out_table_name}/{lh5_subgroup}",
+                        files.glm[file_idx],
+                        wo_mode=mode,
+                    )
+                else:
+                    glm_sum[lh5_subgroup] = (
+                        copy.deepcopy(glm)
+                        if glm_sum[lh5_subgroup] is None
+                        else ak.concatenate((glm_sum[lh5_subgroup], glm))
+                    )
+
+    # return if it was requested to keep glm in memory
+    if glm_sum is not None:
+        return ak.Array(glm_sum)
+    return None