reboost 0.9.0__tar.gz → 0.9.1__tar.gz

{reboost-0.9.0/src/reboost.egg-info → reboost-0.9.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reboost
-Version: 0.9.0
+Version: 0.9.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

{reboost-0.9.0 → reboost-0.9.1}/src/reboost/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.9.0'
-__version_tuple__ = version_tuple = (0, 9, 0)
+__version__ = version = '0.9.1'
+__version_tuple__ = version_tuple = (0, 9, 1)
 
-__commit_id__ = commit_id = 'g0f7fc07b4'
+__commit_id__ = commit_id = 'g366fa2abd'

{reboost-0.9.0 → reboost-0.9.1}/src/reboost/core.py

@@ -18,9 +18,18 @@ 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(

{reboost-0.9.0 → reboost-0.9.1/src/reboost.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reboost
-Version: 0.9.0
+Version: 0.9.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

{reboost-0.9.0 → reboost-0.9.1}/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