reboost 0.5.5__py3-none-any.whl → 0.6.0__py3-none-any.whl

reboost/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import hdf5plugin
 from lgdo import lh5
 
 from ._version import version as __version__
@@ -10,4 +11,4 @@ __all__ = [
     "build_hit",
 ]
 
-lh5.settings.DEFAULT_HDF5_SETTINGS = {"shuffle": True, "compression": "lzf"}
+lh5.settings.DEFAULT_HDF5_SETTINGS = {"compression": hdf5plugin.Zstd()}

reboost/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.5'
-__version_tuple__ = version_tuple = (0, 5, 5)
+__version__ = version = '0.6.0'
+__version_tuple__ = version_tuple = (0, 6, 0)

reboost/build_evt.py

@@ -1,166 +1,134 @@
-"""A program for combining the hits from various detectors, to build events.
+from __future__ import annotations
 
-Is able to parse a config file with the following format config file:
+import logging
 
-.. code-block:: yaml
+import awkward as ak
+import numpy as np
+from dbetto import AttrsDict
+from lgdo import Array, Table, VectorOfVectors, lh5
 
-    channels:
-        geds_on:
-        - det001
-        - det002
-        geds_ac:
-        - det003
+from . import core, math, shape, utils
+from .shape import group
 
-    outputs:
-    - energy
-    - multiplicity
+log = logging.getLogger(__name__)
 
-    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
+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.
 
-     multiplicity:
-        channels: geds_on
-        aggregation_mode: sum
-        expression: "hit.energy > 25"
-        initial: 0
+    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
 
-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.
-"""
+            {"det1": "on", "det2": "off", "det3": "ac"}
 
-from __future__ import annotations
+    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.
 
-import logging
+        For example:
 
-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
+        .. code-block:: python
 
-from . import utils
+            {"p03-r000": {"reso": {"det1": [1, 2], "det2": [0, 1]}}}
 
-log = logging.getLogger(__name__)
+    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
 
-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
+            {"p03-r000": 1000, "p03-r001": 2000}
 
     Returns
     -------
-    ak.Array of the evt tier data (if the data is not saved to disk)
+    the event file in memory as a table if no output file is specified.
     """
-    # create the objects needed for evaluate expression
+    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]
 
-    file_info = {
-        "hit": (hit_file, "hit", "det{:03}"),
-        "evt": (evt_file, "evt"),
-    }
+        # create an output table
+        out_tab = Table(size=n_event)
 
-    # iterate through the TCM
+        tcm_tmp = tcm_ak[cum_sum : cum_sum + n_event]
 
-    out_ak = ak.Array([])
-    mode = "overwrite_file"
+        # usabilities
 
-    # get channel groupings
-    channels = {}
-    for group, info in config["channels"].items():
-        if isinstance(info, str):
-            channels[group] = [info]
+        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:])))
 
-        elif isinstance(info, list):
-            channels[group] = info
+        # now check for channels in ac
+        is_good = group.get_isin_group(channels, channel_groups, tcm_tables, group="on")
 
-    for tcm_lh5 in LH5Iterator(tcm_file, "tcm", buffer_len=buffer):
-        tcm_lh5_sel = tcm_lh5
-        tcm_ak = tcm_lh5_sel.view_as("ak")
+        # get energy
+        energy_true = core.read_data_at_channel_as_ak(
+            channels, rows, hitfile, "energy", "hit", tcm_tables
+        )
 
-        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))),
+        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),
         )
 
-        n_rows = len(tcm.cumulative_length)
-        out_tab = Table(size=n_rows)
+        out_tab.add_field("is_good", VectorOfVectors(is_good[energy > 25]))
 
-        for name, info in config["operations"].items():
-            msg = f"computing field {name}"
-            log.debug(msg)
+        out_tab.add_field("energy", VectorOfVectors(energy[energy > 25]))
+        out_tab.add_field("multiplicity", Array(ak.sum(energy > 25, axis=-1).to_numpy()))
 
-            defaultv = info.get("initial", np.nan)
-            if isinstance(defaultv, str) and (defaultv in ["np.nan", "np.inf", "-np.inf"]):
-                defaultv = eval(defaultv)
+        # write table
+        wo_mode = "of" if idx == 0 else "append"
 
-            channels_use = utils.get_channels_from_groups(info.get("channels", []), channels)
-            channels_exclude = utils.get_channels_from_groups(
-                info.get("exclude_channels", []), channels
-            )
+        # add attrs
+        out_tab.attrs["tables"] = tcm.attrs["tables"]
 
-            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"
+        if outfile is not None:
+            lh5.write(out_tab, "evt", outfile, wo_mode=wo_mode)
         else:
-            out_ak = ak.concatenate((out_ak, out_tab.view_as("ak")))
+            tab = (
+                ak.concatenate((tab, out_tab.view_as("ak")))
+                if tab is not None
+                else out_tab.view_as("ak")
+            )
 
-    if evt_file is None:
-        return out_ak
-    return None
+    return Table(tab)

reboost/core.py

@@ -5,7 +5,9 @@ import time
 from typing import Any
 
 import awkward as ak
+import numpy as np
 from dbetto import AttrsDict
+from lgdo import lh5
 from lgdo.types import LGDO, Table
 
 from . import utils
@@ -14,6 +16,78 @@ from .profile import ProfileDict
 log = logging.getLogger(__name__)
 
 
+def read_data_at_channel_as_ak(
+    channels: ak.Array, rows: ak.Array, file: str, field: str, group: str, tab_map: dict[int, str]
+) -> ak.Array:
+    r"""Read the data from a particular field to an awkward array. This replaces the TCM like object defined by the channels and rows with the corresponding data field.
+
+    Parameters
+    ----------
+    channels
+        Array of the channel indices (uids).
+    rows
+        Array of the rows in the files to gather data from.
+    file
+        File to read the data from.
+    field
+        the field to read.
+    group
+        the group to read data from (eg. `hit` or `stp`.)
+    tab_map
+        mapping between indices and table names. Of the form:
+
+        .. code:: python
+
+            {NAME: UID}
+
+        For example:
+
+        .. code:: python
+
+            {"det001": 1, "det002": 2}
+
+    Returns
+    -------
+    an array with the data, of the same same as the channels and rows.
+    """
+    # initialise the output
+    data_flat = None
+    tcm_rows_full = None
+
+    # save the unflattening
+    reorder = ak.num(rows)
+
+    for tab_name, key in tab_map.items():
+        # get the rows to read
+
+        idx = ak.flatten(rows[channels == key]).to_numpy()
+        arg_idx = np.argsort(idx)
+
+        # get the rows in the flattened data we want to append to
+        tcm_rows = np.where(ak.flatten(channels == key))[0]
+
+        # read the data with sorted idx
+        data_ch = lh5.read(f"{group}/{tab_name}/{field}", file, idx=idx[arg_idx]).view_as("ak")
+
+        # sort back to order for tcm
+        data_ch = data_ch[np.argsort(arg_idx)]
+
+        # append to output
+        data_flat = ak.concatenate((data_flat, data_ch)) if data_flat is not None else data_ch
+        tcm_rows_full = (
+            np.concatenate((tcm_rows_full, tcm_rows)) if tcm_rows_full is not None else tcm_rows
+        )
+
+    if len(data_flat) != len(tcm_rows_full):
+        msg = "every index in the tcm should have been read"
+        raise ValueError(msg)
+
+    # sort the final data
+    data_flat = data_flat[np.argsort(tcm_rows_full)]
+
+    return ak.unflatten(data_flat, reorder)
+
+
 def evaluate_output_column(
     hit_table: Table,
     expression: str,

reboost/math/stats.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+from typing import Callable
 
 import awkward as ak
 import numpy as np
@@ -10,7 +11,68 @@ from numpy.typing import ArrayLike
 log = logging.getLogger(__name__)
 
 
-def gaussian_sample(mu: ArrayLike, sigma: ArrayLike | float, *, seed: int = 999) -> Array:
+def get_resolution(
+    energies: ak.Array, channels: ak.Array, tcm_tables: dict, reso_pars: dict, reso_func: Callable
+) -> ak.Array:
+    """Get the resolution for each energy.
+
+    Parameters
+    ----------
+    energies
+        the energies to smear
+    channels
+        the channel index for each energy
+    tcm_tables
+        the mapping from indices to channel names.
+    reso_pars
+        the pars for each channel.
+    reso_func
+        the function to compute the resolution.
+    """
+    n_pars = len(reso_pars[next(iter(reso_pars))])
+
+    pars_shaped = []
+
+    for _ in range(n_pars):
+        pars_shaped.append(np.zeros(len(ak.flatten(channels))))
+
+    num = ak.num(channels, axis=-1)
+
+    for key, value in tcm_tables.items():
+        for i in range(n_pars):
+            pars_shaped[i][ak.flatten(channels) == value] = reso_pars[key][i]
+
+    ch_reso = reso_func(ak.flatten(energies), *pars_shaped)
+    return ak.unflatten(ch_reso, num)
+
+
+def apply_energy_resolution(
+    energies: ak.Array, channels: ak.Array, tcm_tables: dict, reso_pars: dict, reso_func: Callable
+):
+    """Apply the energy resolution sampling to an array with many channels.
+
+    Parameters
+    ----------
+    energies
+        the energies to smear
+    channels
+        the channel index for each energy
+    tcm_tables
+        the mapping from indices to channel names.
+    reso_pars
+        the pars for each channel.
+    reso_func
+        the function to compute the resolution.
+    """
+    num = ak.num(channels, axis=-1)
+
+    ch_reso = get_resolution(energies, channels, tcm_tables, reso_pars, reso_func)
+    energies_flat_smear = gaussian_sample(ak.flatten(energies), ak.flatten(ch_reso))
+
+    return ak.unflatten(energies_flat_smear, num)
+
+
+def gaussian_sample(mu: ArrayLike, sigma: ArrayLike | float, *, seed: int | None = None) -> Array:
     r"""Generate samples from a gaussian.
 
     Based on:

reboost/shape/group.py

@@ -4,11 +4,49 @@ import logging
 
 import awkward as ak
 import numpy as np
+from dbetto import AttrsDict
 from lgdo import Table, VectorOfVectors
+from numpy.typing import ArrayLike
 
 log = logging.getLogger(__name__)
 
 
+def isin(channels: ak.Array, chan_list: list):
+    """Check if each element of the awkward array channels is in the channel list."""
+    num_channels = ak.num(channels, axis=-1)
+    channels_flat = ak.flatten(channels)
+    isin = np.isin(channels_flat, chan_list)
+
+    # unflatten
+    return ak.unflatten(isin, num_channels)
+
+
+def get_isin_group(
+    channels: ArrayLike, groups: AttrsDict, tcm_tables: dict, group: str = "off"
+) -> ak.Array:
+    """For each channel check if it is in the group.
+
+    Parameters
+    ----------
+    channels
+        Array of the channel indices.
+    groups
+        A mapping of the group for every channel name.
+    tcm_tables
+        the mapping of indices to table names
+    group
+        the group to select.
+
+    Returns
+    -------
+    an awkward array of the same shape of channels of booleans.
+    """
+    usability = {uid: groups[name] for name, uid in tcm_tables.items()}
+    group_idx = [key for key, item in usability.items() if item == group]
+
+    return isin(channels, group_idx)
+
+
 def _sort_data(obj: ak.Array, *, time_name: str = "time", evtid_name: str = "evtid") -> ak.Array:
     """Sort the data by evtid then time.
 

reboost/utils.py

@@ -18,6 +18,16 @@ from .profile import ProfileDict
 log = logging.getLogger(__name__)
 
 
+def get_table_names(tcm: VectorOfVectors) -> dict:
+    """Extract table names from tcm.attrs['tables'] and return them as a dictionary."""
+    raw = tcm.attrs["tables"]
+    cleaned = raw.strip("[]").replace(" ", "").replace("'", "")
+    tables = cleaned.split(",")
+    tables = [tab.split("/")[-1] for tab in tables]
+
+    return {name: idx for idx, name in enumerate(tables)}
+
+
 def get_wo_mode(
     group: int, out_det: int, in_det: int, chunk: int, new_hit_file: bool, overwrite: bool = False
 ) -> str:

{reboost-0.5.5.dist-info → reboost-0.6.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reboost
-Version: 0.5.5
+Version: 0.6.0
 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
@@ -696,6 +696,7 @@ Classifier: Topic :: Scientific/Engineering
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: hdf5plugin
 Requires-Dist: colorlog
 Requires-Dist: numpy
 Requires-Dist: scipy

{reboost-0.5.5.dist-info → reboost-0.6.0.dist-info}/RECORD

@@ -1,22 +1,22 @@
-reboost/__init__.py,sha256=3cYLf7XEyFRX5GK8f50gY4ecGR5O5HORITpDthOFpOg,265
-reboost/_version.py,sha256=tKngLgUb_iFMdwvy40ZSx3eTf2ZsHvVnKan3fsexL-g,511
-reboost/build_evt.py,sha256=yH0bf4bwbp4feWV3JgvSAD5RcvhOX6c9PhH8FAe3Xv4,4710
+reboost/__init__.py,sha256=VZz9uo7i2jgAx8Zi15SptLZnE_qcnGuNWwqkD3rYHFA,278
+reboost/_version.py,sha256=jF9TuoEIJRaca3ScKo6qaz6PzaMlu7jjuSQIrJ3nX4U,511
+reboost/build_evt.py,sha256=VXIfK_pfe_Cgym6gI8dESwONZi-v_4fll0Pn09vePQY,3767
 reboost/build_glm.py,sha256=IerSLQfe51ZO7CQP2kmfPnOIVaDtcfw3byOM02Vaz6o,9472
 reboost/build_hit.py,sha256=23JL5B7qThdHZqAK_HWoytqcEOWDhGsk4n5UMtojJ1c,15513
 reboost/cli.py,sha256=HZgqUZK0tSmnlGqoXjrbmLitW_i001TzibxvDrRxLLg,6324
-reboost/core.py,sha256=pUco_IaTKf50PTVrtyFwoYveJVS58mqs9P3TUrtEyjs,12827
+reboost/core.py,sha256=WGbWe2rcfMDEaehVyw7peqAHoTFWoCu5J6CdWHC5aWA,14974
 reboost/iterator.py,sha256=fATFDxu2PUc0e48OdJJujZo2kwykfRLH1oBtcB-s5pM,6905
 reboost/log_utils.py,sha256=VqS_9OC5NeNU3jcowVOBB0NJ6ssYvNWnirEY-JVduEA,766
 reboost/profile.py,sha256=EOTmjmS8Rm_nYgBWNh6Rntl2XDsxdyed7yEdWtsZEeg,2598
 reboost/units.py,sha256=3EH8XlpbsObdu5vLgxhm1600L6UNYD5jng4SjJT_1QE,2202
-reboost/utils.py,sha256=FXHX7n4YnaYK-MSUDfe9ssEewHjcajkjDb9zeqCzopU,10560
+reboost/utils.py,sha256=WT2jlyRT3LEMBEBRJtYIat3KIIZKocv4rHlTViBbhYM,10938
 reboost/hpge/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 reboost/hpge/psd.py,sha256=jAUAoQ_PMz76wyA1NXYHNKtOwoCnRT3My8_LCFrKi-U,13860
 reboost/hpge/surface.py,sha256=lbWcFnFFWKxtFKs755GyM9US_IfyxaoM6MpOIZgIMM0,7478
 reboost/hpge/utils.py,sha256=0Rx4HubCOm8JMECjWcAJXfAch9OkSlRpUkdsSlzwZ2E,2830
 reboost/math/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 reboost/math/functions.py,sha256=OymiYTcA0NXxxm-MBDw5kqyNwHoLCmuv4J48AwnSrbU,5633
-reboost/math/stats.py,sha256=iiOEi87x93kqPWeSmlRiA5Oe-R8XR-plm6Z532PhC9M,1401
+reboost/math/stats.py,sha256=cG-6mQx33Dzpv3ABkHLEwC104WJ_PMgbWtmjg37SBj4,3164
 reboost/optmap/__init__.py,sha256=imvuyld-GLw8qdwqW-lXCg2feptcTyQo3wIzPvDHwmY,93
 reboost/optmap/cli.py,sha256=SzbPRgsbR5Llm3aSJubH02Ym8FQyTH7kvuLjK7faLiY,9572
 reboost/optmap/convolve.py,sha256=_volpLmhW5mOPA0KkzXRyHyqkj4_zDSnvfHv1Dtuxm8,14390
@@ -27,11 +27,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=_z2qCOret3E-kj-nrp1-J5j2lEwQpgfYdQp2pgpDHR8,4449
+reboost/shape/group.py,sha256=gOCYgir2gZqmW1JXtbNRPlQqP0gmUcbe7RVb9CbY1pU,5540
 reboost/shape/reduction.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-reboost-0.5.5.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-reboost-0.5.5.dist-info/METADATA,sha256=BgZAUrcCJTX0ZHssvfS3w-uFe7Zg9Gt5vJUdMNOp8pE,44248
-reboost-0.5.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-reboost-0.5.5.dist-info/entry_points.txt,sha256=DxhD6BidSWNot9BrejHJjQ7RRLmrMaBIl52T75oWTwM,93
-reboost-0.5.5.dist-info/top_level.txt,sha256=q-IBsDepaY_AbzbRmQoW8EZrITXRVawVnNrB-_zyXZs,8
-reboost-0.5.5.dist-info/RECORD,,
+reboost-0.6.0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+reboost-0.6.0.dist-info/METADATA,sha256=QGF9b0vthakVY3Wl9OUQnPIyAv3IYpui4WK9Zkp7Nwo,44274
+reboost-0.6.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+reboost-0.6.0.dist-info/entry_points.txt,sha256=DxhD6BidSWNot9BrejHJjQ7RRLmrMaBIl52T75oWTwM,93
+reboost-0.6.0.dist-info/top_level.txt,sha256=q-IBsDepaY_AbzbRmQoW8EZrITXRVawVnNrB-_zyXZs,8
+reboost-0.6.0.dist-info/RECORD,,