reboost 0.10.0a0__tar.gz → 0.10.0a2__tar.gz

{reboost-0.10.0a0/src/reboost.egg-info → reboost-0.10.0a2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reboost
-Version: 0.10.0a0
+Version: 0.10.0a2
 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.10.0a0 → reboost-0.10.0a2}/src/reboost/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.10.0a0'
-__version_tuple__ = version_tuple = (0, 10, 0, 'a0')
+__version__ = version = '0.10.0a2'
+__version_tuple__ = version_tuple = (0, 10, 0, 'a2')
 
-__commit_id__ = commit_id = 'gd31e91903'
+__commit_id__ = commit_id = 'g8f2cabc6a'

{reboost-0.10.0a0 → reboost-0.10.0a2}/src/reboost/build_hit.py

@@ -51,14 +51,15 @@ A :func:`build_hit` to parse the following configuration file:
 
             t0: ak.fill_none(ak.firsts(HITS.time, axis=-1), np.nan)
 
-            evtid: ak.fill_none(ak.firsts(HITS.__evtid, axis=-1), np.nan)
+            evtid: ak.fill_none(ak.firsts(HITS.evtid, axis=-1), np.nan)
 
             # distance to the nplus surface in mm
             distance_to_nplus_surface_mm: reboost.hpge.distance_to_surface(
-                HITS.__xloc, HITS.__yloc, HITS.__zloc,
+                HITS.xloc, HITS.yloc, HITS.zloc,
                 DETECTOR_OBJECTS.pyobj,
                 DETECTOR_OBJECTS.phyvol.position.eval(),
-                surface_type='nplus')
+                surface_type='nplus',
+                unit='m')
 
             # activness based on FCCD (no TL)
             activeness: ak.where(
@@ -75,7 +76,7 @@ A :func:`build_hit` to parse the following configuration file:
                 )
 
             # summed energy of the hit accounting for activeness
-            energy_raw: ak.sum(HITS.__edep * HITS.activeness, axis=-1)
+            energy_raw: ak.sum(HITS.edep * HITS.activeness, axis=-1)
 
             # energy with smearing
             energy: reboost.math.sample_convolve(
@@ -92,7 +93,7 @@ A :func:`build_hit` to parse the following configuration file:
                 )
 
             # example of low level reduction on clusters
-            energy_clustered: ak.sum(ak.unflatten(HITS.__edep, HITS.clusters_lengths), axis=-1)
+            energy_clustered: ak.sum(ak.unflatten(HITS.edep, HITS.clusters_lengths), axis=-1)
 
             # example of using a reboost helper
             steps_clustered: reboost.shape.reduction.energy_weighted_average(HITS, HITS.clusters_lengths)
@@ -115,7 +116,7 @@ A :func:`build_hit` to parse the following configuration file:
             - num_scint_ph_lar
 
           operations:
-            tot_edep_wlsr: ak.sum(HITS.edep[np.abs(HITS.zloc) < 3000], axis=-1)
+            tot_edep_wlsr: ak.sum(HITS.edep[np.abs(HITS.zloc) < 3], axis=-1)
 
         - name: spms
 
@@ -180,6 +181,8 @@ from dbetto import AttrsDict
 from lgdo import lh5
 from lgdo.lh5.exceptions import LH5EncodeError
 
+from reboost import units
+
 from . import core, utils
 from .iterator import GLMIterator
 from .profile import ProfileDict
@@ -334,6 +337,7 @@ def build_hit(
                             expression=proc_group["hit_table_layout"],
                             time_dict=time_dict[proc_name],
                         )
+
                     else:
                         hit_table_layouted = copy.deepcopy(stps)
 
@@ -448,5 +452,8 @@ def _evaluate_operation(
         time_dict=time_dict,
         name=field,
     )
+    if not isinstance(info, str) and "units" in info:
+        col = units.attach_units(col, info["units"])
+    units.move_units_to_flattened_data(col)
 
     core.add_field_with_nesting(hit_table, field, col)

{reboost-0.10.0a0 → reboost-0.10.0a2}/src/reboost/core.py

@@ -11,16 +11,25 @@ from dbetto import AttrsDict
 from lgdo import lh5
 from lgdo.types import LGDO, Table
 
-from . import utils
+from . import units, utils
 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]
+    channels: ak.Array,
+    rows: ak.Array,
+    file: str,
+    field: str,
+    group: str,
+    tab_map: dict[int, str],
+    with_units: bool = False,
 ) -> 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.
+    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
     ----------
@@ -68,7 +77,9 @@ def read_data_at_channel_as_ak(
         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")
+        data_ch = lh5.read(f"{group}/{tab_name}/{field}", file, idx=idx[arg_idx])
+        units = data_ch.attrs.get("units", None)
+        data_ch = data_ch.view_as("ak")
 
         # sort back to order for tcm
         data_ch = data_ch[np.argsort(arg_idx)]
@@ -85,8 +96,12 @@ def read_data_at_channel_as_ak(
 
     # sort the final data
     data_flat = data_flat[np.argsort(tcm_rows_full)]
+    data_unflat = ak.unflatten(data_flat, reorder)
+
+    if with_units and units is not None:
+        return ak.with_parameter(data_unflat, "units", units)
 
-    return ak.unflatten(data_flat, reorder)
+    return data_unflat
 
 
 def evaluate_output_column(
@@ -441,6 +456,10 @@ def evaluate_hit_table_layout(
 
     res = eval(group_func, globs, locs)
 
+    if isinstance(res, Table):
+        for data in res.values():
+            units.move_units_to_flattened_data(data)
+
     if time_dict is not None:
         time_dict.update_field(name="hit_layout", time_start=time_start)
 
@@ -519,8 +538,5 @@ def remove_columns(tab: Table, outputs: list) -> Table:
 
 def merge(hit_table: Table, output_table: ak.Array | None):
     """Merge the table with the array."""
-    return (
-        hit_table.view_as("ak")
-        if output_table is None
-        else ak.concatenate((output_table, hit_table.view_as("ak")))
-    )
+    hit_table = hit_table.view_as("ak", with_units=True)
+    return hit_table if output_table is None else ak.concatenate((output_table, hit_table))

{reboost-0.10.0a0 → reboost-0.10.0a2}/src/reboost/optmap/evt.py

@@ -3,10 +3,8 @@ from __future__ import annotations
 import logging
 from collections import OrderedDict
 from collections.abc import Generator, Iterable
-from pathlib import Path
 
 import numpy as np
-from lgdo import lh5
 from lgdo.lh5 import LH5Iterator
 from lgdo.types import Table
 
@@ -105,27 +103,6 @@ def generate_optmap_evt(
     assert had_last_chunk, "did not reach last chunk in optmap-evt building"
 
 
-def build_optmap_evt(
-    lh5_in_file: str, lh5_out_file: str, detectors: Iterable[str | int], buffer_len: int = int(5e6)
-) -> None:
-    """Create a faster map for lookup of the hits in each detector, for each primary event."""
-    lh5_out_file = Path(lh5_out_file)
-    lh5_out_file_tmp = lh5_out_file.with_stem(".evt-tmp." + lh5_out_file.stem)
-    if lh5_out_file_tmp.exists():
-        msg = f"temporary output file {lh5_out_file_tmp} already exists"
-        raise RuntimeError(msg)
-
-    for vert_it_count, chunk in enumerate(generate_optmap_evt(lh5_in_file, detectors, buffer_len)):
-        log.info("store evt file %s (%d)", lh5_out_file_tmp, vert_it_count - 1)
-        lh5.write(Table(chunk), name=EVT_TABLE_NAME, lh5_file=lh5_out_file_tmp, wo_mode="append")
-
-    # after finishing the output file, rename to the actual output file name.
-    if lh5_out_file.exists():
-        msg = f"output file {lh5_out_file_tmp} already exists after writing tmp output file"
-        raise RuntimeError(msg)
-    lh5_out_file_tmp.rename(lh5_out_file)
-
-
 def get_optical_detectors_from_geom(geom_fn) -> dict[int, str]:
     import pyg4ometry
     import pygeomtools

{reboost-0.10.0a0 → reboost-0.10.0a2}/src/reboost/units.py

@@ -7,7 +7,7 @@ import awkward as ak
 import numpy as np
 import pint
 import pyg4ometry as pg4
-from lgdo import LGDO
+from lgdo import LGDO, VectorOfVectors
 
 log = logging.getLogger(__name__)
 
@@ -69,6 +69,23 @@ def attach_units(data: ak.Array | LGDO, unit: str | None) -> ak.Array | LGDO:
     return data
 
 
+def move_units_to_flattened_data(data: LGDO) -> None:
+    """If `data` is a VectorOfVectors move units from attrs to flattened data attrs.
+
+    Parameters
+    ----------
+    data
+        the nested data structure
+    """
+    if isinstance(data, VectorOfVectors) and ("units" in data.attrs):
+        unit = data.attrs.pop("units")
+        if isinstance(data.flattened_data, VectorOfVectors):
+            data.flattened_data.attrs |= {"units": unit}
+            move_units_to_flattened_data(data.flattened_data)
+        else:
+            data.flattened_data.attrs |= {"units": unit}
+
+
 def units_conv_ak(data: Any | LGDO | ak.Array, target_units: pint.Unit | str) -> Any | ak.Array:
     """Calculate numeric conversion factor to reach `target_units`, and apply to data converted to ak.
 

{reboost-0.10.0a0 → reboost-0.10.0a2}/src/reboost/utils.py

@@ -5,7 +5,7 @@ import itertools
 import logging
 import re
 import time
-from collections.abc import Iterable, Mapping
+from collections.abc import Iterable
 from contextlib import contextmanager
 from pathlib import Path
 
@@ -182,52 +182,6 @@ 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:
-            if not isinstance(tab[field], VectorOfVectors):
-                tab[field].attrs["units"] = units[field]
-            else:
-                tab[field].flattened_data.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*\("

{reboost-0.10.0a0 → reboost-0.10.0a2/src/reboost.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reboost
-Version: 0.10.0a0
+Version: 0.10.0a2
 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.10.0a0 → reboost-0.10.0a2}/tests/hit/configs/basic.yaml

@@ -6,11 +6,16 @@ processing_groups:
 
     outputs:
       - t0
+      - t0_u
       - evtid
       - energy
+      - xloc
 
     operations:
       t0: ak.fill_none(ak.firsts(HITS.time, axis=-1), np.nan)
+      t0_u:
+        expression: ak.fill_none(ak.firsts(HITS.time, axis=-1), np.nan)
+        units: ns
       energy: ak.sum(HITS.edep, axis=-1)
 
 forward:

{reboost-0.10.0a0 → reboost-0.10.0a2}/tests/hit/test_build_hit.py

@@ -11,6 +11,7 @@ import pytest
 from lgdo import Array, Struct, Table, VectorOfVectors, lh5
 
 import reboost
+from reboost import units
 
 
 @pytest.fixture(scope="module")
@@ -136,12 +137,24 @@ def test_reshape(test_gen_lh5_flat, tmptestdir):
     )  # m
     tab = Table(data)
 
+    for field in data.values():
+        units.move_units_to_flattened_data(field)
+
     output = lh5.read("stp/det1", outfile)
 
     # check the outputs
     assert output == tab
     assert lh5.read("vtx", outfile) == Table({"evtid": Array([0, 1])})
 
+    # test units
+    for field, unit in zip(
+        ["edep", "time", "xloc", "yloc", "zloc", "dist_to_surf"],
+        ["keV", "ns", "m", "m", "m", "m"],
+        strict=True,
+    ):
+        assert "units" not in output[field].attrs
+        assert output[field].flattened_data.attrs["units"] == unit
+
 
 def test_only_forward(test_gen_lh5_flat, tmptestdir):
     outfile = f"{tmptestdir}/basic_hit_foward_only.lh5"
@@ -181,15 +194,20 @@ def test_basic(test_gen_lh5, tmptestdir):
             == b"Zstandard compression: http://www.zstd.net"
         )
 
-    hits = lh5.read("hit/det1", outfile).view_as("ak")
+    hits = lh5.read("hit/det1", outfile).view_as("ak", with_units=True)
 
     assert ak.all(hits.energy == [300, 330])
     assert ak.all(hits.t0 == [0, 0.1])
+
     assert hits.evtid[0] == 0
     assert hits.evtid[1] == 1
 
     assert len(hits) == 2
 
+    assert ak.parameters(hits.t0) == {}
+    assert ak.parameters(hits.t0_u)["units"] == "ns"
+    assert ak.parameters(hits.xloc)["units"] == "m"
+
     # test in memory
 
     hits, time_dict = reboost.build_hit(
@@ -205,6 +223,10 @@ def test_basic(test_gen_lh5, tmptestdir):
     assert ak.all(hits["det1"].evtid[0] == [0, 0])
     assert ak.all(hits["det1"].evtid[1] == [1, 1, 1])
 
+    assert ak.parameters(hits["det1"].t0) == {}
+    assert ak.parameters(hits["det1"].t0_u)["units"] == "ns"
+    assert ak.parameters(hits["det1"].xloc)["units"] == "m"
+
     assert set(time_dict.keys()) == {"global_objects", "geds"}
     assert set(time_dict["geds"].keys()) == {
         "detector_objects",
@@ -213,7 +235,7 @@ def test_basic(test_gen_lh5, tmptestdir):
         "expressions",
     }
     assert set(time_dict["geds"]["read"].keys()) == {"stp"}
-    assert set(time_dict["geds"]["expressions"].keys()) == {"t0", "energy"}
+    assert set(time_dict["geds"]["expressions"].keys()) == {"t0", "t0_u", "energy"}
 
 
 def test_file_merging(test_gen_lh5, tmptestdir):

{reboost-0.10.0a0 → reboost-0.10.0a2}/tests/test_core.py

@@ -59,13 +59,13 @@ def hitfiles(tmptestdir):
     # make some hit tier files
     channel1 = Table(
         {
-            "energy": Array([100, 200, 400, 300]),
+            "energy": Array([100, 200, 400, 300], attrs={"units": "keV"}),
             "times": VectorOfVectors([[0.1], [0.2, 0.3], [0.4, 98], [2]]),
         }
     )
     channel2 = Table(
         {
-            "energy": Array([10, 70, 0, 56, 400, 400]),
+            "energy": Array([10, 70, 0, 56, 400, 400], attrs={"units": "keV"}),
             "times": VectorOfVectors([[12], [], [-0.4, 0.4], [89], [1], [2]]),
         }
     )
@@ -83,12 +83,13 @@ def hitfiles(tmptestdir):
 
 def test_read_data_at_channel(hitfiles):
     # make a TCM
-    tcm_channels = ak.Array([[0], [0], [0, 1], [1], [1], [0, 1], [1], [1]])
-    tcm_rows = ak.Array([[0], [1], [2, 0], [1], [2], [3, 3], [4], [5]])
+    tcm_channels = ak.Array([[0], [0], [0, 1], [1], [1], [0, 1], [1], [], [1]])
+    tcm_rows = ak.Array([[0], [1], [2, 0], [1], [2], [3, 3], [4], [], [5]])
 
     energy = reboost.core.read_data_at_channel_as_ak(
         tcm_channels, tcm_rows, hitfiles[2], "energy", "hit", {"det001": 0, "det002": 1}
     )
+    assert "units" not in ak.parameters(energy)
 
     # check the same
     assert len(energy) == len(tcm_channels)
@@ -105,6 +106,18 @@ def test_read_data_at_channel(hitfiles):
     )
     assert len(times) == len(tcm_channels)
 
+    energy = reboost.core.read_data_at_channel_as_ak(
+        tcm_channels,
+        tcm_rows,
+        hitfiles[2],
+        "energy",
+        "hit",
+        {"det001": 0, "det002": 1},
+        with_units=True,
+    )
+    assert "units" in ak.parameters(energy)
+    assert ak.parameters(energy)["units"] == "keV"
+
 
 def test_get_objects(test_data_configs, make_gdml):
     # check basic eval

{reboost-0.10.0a0 → reboost-0.10.0a2}/tests/test_optmap.py

@@ -13,7 +13,6 @@ from reboost.optmap.create import (
     merge_optical_maps,
     rebin_optical_maps,
 )
-from reboost.optmap.evt import build_optmap_evt
 from reboost.optmap.optmap import OpticalMap
 
 
@@ -53,16 +52,6 @@ def tbl_hits(tmptestdir):
     return (str(hit_file),)
 
 
-def test_optmap_evt(tbl_hits, tmptestdir):
-    evt_out_file = tmptestdir / "evt-out.lh5"
-    build_optmap_evt(
-        tbl_hits[0],
-        str(evt_out_file),
-        detectors=("1", "002", "003"),
-        buffer_len=20,  # note: shorter window sizes (e.g. 10) do not work.
-    )
-
-
 @pytest.fixture
 def tbl_evt_fns(tmptestdir) -> tuple[str]:
     evt_count = 100

{reboost-0.10.0a0 → reboost-0.10.0a2}/tests/test_units.py

@@ -4,7 +4,7 @@ import awkward as ak
 import pint
 import pyg4ometry as pg4
 import pytest
-from lgdo.types import Array
+from lgdo.types import Array, VectorOfVectors
 
 from reboost import units
 from reboost.units import ureg as u
@@ -92,3 +92,11 @@ def test_attach_units():
     data_units = units.attach_units(Array([1, 2, 3]), "mm")
 
     assert units.get_unit_str(data_units) == "mm"
+
+
+def test_move_units_to_flattened_data():
+    data = VectorOfVectors([[1, 2, 3]], attrs={"units": "mm"})
+    units.move_units_to_flattened_data(data)
+
+    assert data.attrs.get("units") is None
+    assert data.flattened_data.attrs.get("units") == "mm"

{reboost-0.10.0a0 → reboost-0.10.0a2}/tests/test_utils.py

@@ -6,14 +6,11 @@ from pathlib import Path
 
 import pytest
 import yaml
-from lgdo.types import Array, Table, VectorOfVectors
+from lgdo.types import VectorOfVectors
 
 import reboost
 from reboost import utils
-from reboost.shape import group
 from reboost.utils import (
-    assign_units,
-    copy_units,
     get_file_dict,
     get_function_string,
     get_table_names,
@@ -204,28 +201,6 @@ def test_get_files_dict():
     assert files.glm == ["glm1.lh5", "glm2.lh5"]
 
 
-def test_units():
-    table = Table({"a": Array([1, 2, 3]), "b": Array([4, 5, 6]), "evtid": Array([0, 0, 1])})
-
-    table.a.attrs = {"datatype": "array<1>{real}", "units": "ns"}
-    table.b.attrs = {"datatype": "array<1>{real}", "units": "keV"}
-
-    units = copy_units(table)
-
-    assert units["a"] == "ns"
-    assert units["b"] == "keV"
-    reshaped = group.group_by_evtid(table.view_as("ak"))
-
-    # also add an array field
-    units["c"] = "keV"
-    reshaped["c"] = Array([1, 2])
-    reshaped = assign_units(reshaped, units)
-
-    assert reshaped.a.flattened_data.attrs["units"] == "ns"
-    assert reshaped.b.flattened_data.attrs["units"] == "keV"
-    assert reshaped.c.attrs["units"] == "keV"
-
-
 def test_table_names():
     names = "['hit/det001','hit/det002']"