switch all hpge processors to use ak.Array and unit handelling

src/reboost/hpge/psd.py

@@ -141,8 +141,8 @@ def _ak_dt_map(layouts, **_kwargs):
 
 
 def drift_time_heuristic(
-    drift_time: ArrayLike,
-    edep: ArrayLike,
+    drift_time: ak.Array,
+    edep: ak.Array,
 ) -> ak.Array:
     """HPGe drift-time-based pulse-shape heuristic.
 
@@ -165,7 +165,9 @@ def drift_time_heuristic(
     if t_units is not None and e_units is not None:
         attrs["units"] = units.unit_to_lh5_attr(t_units / e_units)
 
-    return ak.Array(_drift_time_heuristic_impl(drift_time, edep), attrs=attrs)
+    return units.units_conv_ak(
+        ak.Array(_drift_time_heuristic_impl(drift_time, edep)), attrs["units"]
+    )
 
 
 @numba.njit(cache=True)

src/reboost/hpge/surface.py

@@ -104,7 +104,7 @@ def distance_to_surface(
             local_positions[indices], surface_indices=surface_indices
         )
 
-    return ak.Array(ak.unflatten(distances, size), dtype=np.float32, attrs={"units": "mm"})
+    return ak.Array(ak.unflatten(distances, size), attrs={"units": "mm"})
 
 
 @numba.njit(cache=True)

src/reboost/math/functions.py

@@ -50,6 +50,8 @@ def piecewise_linear_activeness(distances: ak.Array, fccd_in_mm: float, dlf: flo
     a :class:`VectorOfVectors` or :class:`Array` of the activeness
     """
     # convert to ak
+    distances = ak.Array(distances)
+
     distances_ak = units.units_conv_ak(distances, "mm")
 
     dl = fccd_in_mm * dlf

src/reboost/units.py

@@ -45,6 +45,19 @@ def units_convfact(data: Any | LGDO | ak.Array, target_units: pint.Unit | str) -
     return 1
 
 
+def attach_units(data: ak.Array, unit: str) -> ak.Array:
+    """Convenience function to attach units to `ak.Array`.
+
+    Parameters
+    ----------
+    data
+        the array to add units to
+    unit
+        the unit
+    """
+    return ak.with_parameter(data, parameter="units", value=unit)
+
+
 def units_conv_ak(data: Any | LGDO | ak.Array, target_units: pint.Unit | str) -> ak.Array:
     """Calculate numeric conversion factor to reach `target_units`, and apply to data converted to ak.
 

tests/hit/configs/hit_config.yaml

@@ -36,7 +36,7 @@ processing_groups:
 
       activeness:
         reboost.math.functions.piecewise_linear_activeness(HITS.distance_to_nplus,
-        fccd=DETECTOR_OBJECTS.det_pars.fccd_in_mm,
+        fccd_in_mm=DETECTOR_OBJECTS.det_pars.fccd_in_mm,
         dlf=DETECTOR_OBJECTS.det_pars.dlf)
 
       active_energy: ak.sum(HITS.edep*HITS.activeness, axis=-1)

tests/hpge/test_current.py

@@ -4,6 +4,7 @@
 import numpy as np
 import pytest
 
+from reboost import units
 from reboost.hpge import psd, surface
 from reboost.shape import cluster
 
@@ -48,8 +49,8 @@ def test_model():
 def test_maximum_current(test_model):
     model, x = test_model
 
-    edep = ak.Array([[100.0, 300.0, 50.0], [10.0, 0.0, 100.0], [500.0]], attrs={"unit": "keV"})
-    times = ak.Array([[400, 500, 700], [800, 0, 1500], [700]], attrs={"unit": "ns"})
+    edep = units.attach_units(ak.Array([[100.0, 300.0, 50.0], [10.0, 0.0, 100.0], [500.0]]), "keV")
+    times = units.attach_units(ak.Array([[400, 500, 700], [800, 0, 1500], [700]]), "ns")
 
     curr = psd.maximum_current(edep, times, template=model, times=x)
     assert isinstance(curr, ak.Array)
@@ -92,17 +93,17 @@ def test_maximum_current(test_model):
 def test_with_cluster(test_model):
     model, x = test_model
 
-    edep = ak.Array([[100.0, 300.0, 50.0], [10.0, 1.0, 100.0], [500.0]], attrs={"unit": "keV"})
-    times = ak.Array([[400, 410, 420], [800, 0, 1500], [700]], attrs={"unit": "ns"})
-    xloc = ak.Array([[1, 1.1, 1.2], [0, 50, 80], [100]], attrs={"unit": "mm"})
-    dist = ak.Array([[50, 40, 0.2], [300, 0.4, 0.2], [0.8]], attrs={"unit": "ns"})
+    edep = units.attach_units(ak.Array([[100.0, 300.0, 50.0], [10.0, 1.0, 100.0], [500.0]]), "keV")
+    times = units.attach_units(ak.Array([[400, 410, 420], [800, 0, 1500], [700]]), "ns")
+    xloc = units.attach_units(ak.Array([[1, 1.1, 1.2], [0, 50, 80], [100]]), "mm")
+    dist = units.attach_units(ak.Array([[50, 40, 0.2], [300, 0.4, 0.2], [0.8]]), "mm")
 
     yloc = ak.full_like(xloc, 0.0)
     zloc = ak.full_like(xloc, 0.0)
     trackid = ak.full_like(xloc, 0)
 
     clusters = cluster.cluster_by_step_length(
-        trackid, xloc, yloc, zloc, dist, threshold=1, threshold_surf=1, surf_cut=0
+        trackid, xloc, yloc, zloc, dist, threshold_in_mm=1, threshold_surf_in_mm=1, surf_cut=0
     )
     cluster_edep = cluster.apply_cluster(clusters, edep)
     cluster_times = cluster.apply_cluster(clusters, times)
@@ -123,21 +124,23 @@ def test_maximum_current_surface(test_model):
 
     # test for both input types
     for dtype in [np.float64, np.float32]:
-        edep = ak.Array(
-            ak.values_astype(ak.Array([[100.0, 300.0, 50.0], [10.0, 0.0, 100.0], [500.0]]), dtype),
-            attrs={"unit": "keV"},
+        edep = units.attach_units(
+            ak.Array(
+                ak.values_astype(
+                    ak.Array([[100.0, 300.0, 50.0], [10.0, 0.0, 100.0], [500.0]]), dtype
+                ),
+            ),
+            "keV",
         )
 
-        times = ak.Array(
-            ak.values_astype(
-                ak.Array([[400, 500, 700], [800, 0, 1500], [700]], attrs={"unit": "ns"}), dtype
-            )
+        times = units.attach_units(
+            ak.Array(ak.values_astype(ak.Array([[400, 500, 700], [800, 0, 1500], [700]]), dtype)),
+            "ns",
         )
 
-        dist = ak.Array(
-            ak.values_astype(
-                ak.Array([[50, 40, 0.2], [300, 0.4, 0.2], [0.8]], attrs={"unit": "ns"}), dtype
-            )
+        dist = units.attach_units(
+            ak.Array(ak.values_astype(ak.Array([[50, 40, 0.2], [300, 0.4, 0.2], [0.8]]), dtype)),
+            "mm",
         )
 
         surface_models = surface.get_surface_library(1002, 10)

tests/hpge/test_dt_heuristic.py

@@ -4,16 +4,17 @@
 import numpy as np
 import pytest
 from dbetto import AttrsDict
-from lgdo import Table, VectorOfVectors, lh5
+from lgdo import lh5
 from scipy.interpolate import RegularGridInterpolator
 
+from reboost import units
 from reboost.hpge import psd, utils
 from reboost.units import ureg as u
 
 # data from remage sim in ./simulation
 # fmt: off
 gamma_stp = ak.Array({
-        "edep": ak.Array(
+        "edep": units.attach_units(ak.Array(
             [
                 [763, 20.1, 30.2, 40.5, 36.5, 110],
                 [0.0158, 526, 111],
@@ -26,9 +27,8 @@
                 [728, 16.1, 256],
                 [613, 0.179, 9.03, 101, 26.9, 22.9, 26.1, 31.5, 11.1, 5.56, 31.7, 27.2, 23, 33.7, 37.6],
             ],
-            attrs={"units": "keV"},
-        ),
-        "xloc": ak.Array(
+        ), "keV"),
+        "xloc": units.attach_units(ak.Array(
             [
                 [-0.0022, 0.0115, 0.0141, 0.0169, 0.0165, 0.0165],
                 [0.0106, 0.0104, 0.0103],
@@ -41,9 +41,9 @@
                 [-0.00369, -0.00343, 0.00682],
                 [-0.0162, -0.0224, -0.0224, -0.0224, -0.0224, -0.0224, -0.0224, -0.0224, -0.022, -0.022, -0.022, -0.022, -0.022, -0.022, -0.022],
             ],
-            attrs={"units": "m"},
-        ),
-        "yloc": ak.Array(
+
+        ), "mm"),
+        "yloc": units.attach_units(ak.Array(
             [
                 [-0.00951, -0.0201, -0.0199, -0.014, -0.00924, -0.0092],
                 [-0.0264, -0.0264, -0.0264],
@@ -56,9 +56,8 @@
                 [-0.000536, -0.000655, -0.00369],
                 [-0.00355, -0.0134, -0.0134, -0.0134, -0.0134, -0.0134, -0.0134, -0.0134, -0.0133, -0.0133, -0.0133, -0.0133, -0.0133, -0.0133, -0.0133],
             ],
-            attrs={"units": "m"},
-        ),
-        "zloc": ak.Array(
+        ),"mm"),
+        "zloc": units.attach_units(ak.Array(
             [
                 [0.0281, 0.0124, 0.0112, 0.00313, 0.0052, 0.00495],
                 [0.029, 0.0292, 0.0293],
@@ -71,8 +70,7 @@
                 [0.0197, 0.0195, 0.0117],
                 [0.0029, 0.00108, 0.00108, 0.00107, 0.00106, 0.00106, 0.00106, 0.00106, 0.00128, 0.00128, 0.00129, 0.0013, 0.0013, 0.0013, 0.0013],
             ],
-            attrs={"units": "m"},
-        ),
+        ), "mm"),
     }
 )
 # fmt: on
@@ -123,12 +121,12 @@ def dt_map_dummy(legendtestdata):
 
 
 def test_drift_time_dummy(dt_map_dummy):
-    gamma_stp_shift = Table(
+    gamma_stp_shift = ak.Array(
         {
             "edep": gamma_stp.edep,
-            "xloc": VectorOfVectors(gamma_stp.xloc.view_as("ak") + 10, attrs={"units": "m"}),
-            "yloc": VectorOfVectors(gamma_stp.yloc.view_as("ak") + 10, attrs={"units": "m"}),
-            "zloc": VectorOfVectors(gamma_stp.zloc.view_as("ak") + 10, attrs={"units": "m"}),
+            "xloc": ak.Array(gamma_stp.xloc + 10, attrs={"units": "m"}),
+            "yloc": ak.Array(gamma_stp.yloc + 10, attrs={"units": "m"}),
+            "zloc": ak.Array(gamma_stp.zloc + 10, attrs={"units": "m"}),
         }
     )
 
@@ -146,11 +144,11 @@ def test_drift_time_dummy(dt_map_dummy):
         * u.m,
     )
     # turn into an Awkward array so we can index
-    dt_arr = dt_values.view_as("ak")
+    dt_arr = dt_values
 
     # helper to pull out expected from the RegularGridInterpolator
     def expected_dt(event, step):
-        stp = gamma_stp.view_as("ak")
+        stp = gamma_stp
         x = stp.xloc[event][step]
         y = stp.yloc[event][step]
         z = stp.zloc[event][step]
@@ -176,9 +174,12 @@ def test_drift_time(dt_map):
         gamma_stp.zloc,
         dt_map,
     )
+
+    dt_values, unit = units.unwrap_lgdo(dt_values)
     assert isinstance(dt_values, ak.Array)
     assert dt_values.ndim == 2
-    assert dt_values.attrs["units"] == "ns"
+
+    assert unit == "ns"
 
 
 def test_drift_time_heuristics(dt_map):
@@ -193,6 +194,7 @@ def test_drift_time_heuristics(dt_map):
         dt_values,
         gamma_stp.edep,
     )
-
+    dt_heu, unit = units.unwrap_lgdo(dt_heu)
     assert isinstance(dt_heu, ak.Array)
-    assert dt_heu.attrs["units"] == "ns/keV"
+
+    assert unit == "ns/keV"

tests/hpge/test_r90.py

@@ -116,8 +116,8 @@ def test_r90():
     elect_r90 = 0.0001816
     gamma_r90 = 0.0157370
 
-    assert round(r90_output.nda[0], 7) == elect_r90
-    assert round(r90_output.nda[1], 7) == gamma_r90
+    assert round(r90_output[0], 7) == elect_r90
+    assert round(r90_output[1], 7) == gamma_r90
 
     edep = [gamma_edep]
     xloc = [gamma_xloc]

tests/hpge/test_surface.py

@@ -6,7 +6,6 @@
 import pytest
 from legendhpges import make_hpge
 from legendtestdata import LegendTestData
-from lgdo import types
 
 from reboost.hpge.surface import distance_to_surface, get_surface_response
 
@@ -46,7 +45,7 @@ def test_distance_to_surface(test_data_configs):
     dist_full = distance_to_surface(
         pos.xloc, pos.yloc, pos.zloc, gedet, det_pos=dist, surface_type=None
     )
-    assert isinstance(dist_full, types.LGDO)
+    assert isinstance(dist_full, ak.Array)
 
     # check skipping the calculation for points > 5 mm
     dist = distance_to_surface(

tests/test_math.py

@@ -12,9 +12,7 @@ def test_hpge_activeness():
     # test with VectorOfVectors and ak.Array input
     distances = ak.Array([[0.2], [0.6], [2]])
 
-    activeness = functions.piecewise_linear_activeness(
-        VectorOfVectors(distances), fccd_in_mm=1, dlf=0.4
-    )
+    activeness = functions.piecewise_linear_activeness(ak.Array(distances), fccd_in_mm=1, dlf=0.4)
 
     # first point should be 0
     assert activeness[0][0] == 0
@@ -24,7 +22,7 @@ def test_hpge_activeness():
 
     # test with ak.Array input
     distances = ak.Array([[0.2], [0.6], [2]])
-    activeness = functions.piecewise_linear_activeness(distances, fccd=1, dlf=0.4)
+    activeness = functions.piecewise_linear_activeness(distances, fccd_in_mm=1, dlf=0.4)
 
     # first point should be 0
     assert activeness[0][0] == 0
@@ -33,19 +31,9 @@ def test_hpge_activeness():
     assert activeness[2][0] == 1
 
     # test with Array
-    activeness = functions.piecewise_linear_activeness([[0.2, 0.6, 2]], fccd=1, dlf=0.4)
-
-    assert np.allclose(activeness.view_as("np"), [0, 1 / 3.0, 1])
+    activeness = functions.piecewise_linear_activeness([[0.2, 0.6, 2]], fccd_in_mm=1, dlf=0.4)
 
-    # test with array
-    distances = [0.2, 0.6, 2]
-    activeness = functions.piecewise_linear_activeness(Array(distances), fccd=1, dlf=0.4)
-
-    # first point should be 0
-    assert activeness[0] == 0
-    # second should be 0.1/0.5 = 0.2
-    assert activeness[1] == pytest.approx(1 / 3.0)
-    assert activeness[2] == 1
+    assert np.allclose(activeness, [0, 1 / 3.0, 1])
 
 
 def test_vectorised_activeness():
@@ -54,18 +42,14 @@ def test_vectorised_activeness():
     edep = ak.Array([[100, 200], [100], [100], [100]])
 
     # simple case
-    energy = functions.vectorised_active_energy(
-        VectorOfVectors(distances), VectorOfVectors(edep), fccd=1, dlf=0.4
-    )
+    energy = functions.vectorised_active_energy(distances, edep, fccd=1, dlf=0.4)
 
     assert energy[0] == 200
     assert energy[1] == pytest.approx(100 / 3.0)
     assert energy[2] == 100
 
     # now vectorised over fccd
-    energy_fccd = functions.vectorised_active_energy(
-        VectorOfVectors(distances), VectorOfVectors(edep), fccd=[0, 0.5, 1, 2, 3], dlf=1
-    )
+    energy_fccd = functions.vectorised_active_energy(distances, edep, fccd=[0, 0.5, 1, 2, 3], dlf=1)
 
     # with distance of 0.2 only fccd of 0 gives energy
     assert np.allclose(energy_fccd[0], [300, 200, 200, 200, 200])

tests/test_shape.py

@@ -135,11 +135,11 @@ def test_cluster_by_step_length():
         trackid, x, y, z, dist, threshold_in_mm=0.2, threshold_surf_in_mm=0.02, surf_cut=0.15
     )
 
-    assert ak.all(clusters.view_as("ak") == ak.Array([[1, 1, 1, 1, 3, 1], [1, 1, 1, 2, 1], [1]]))
+    assert ak.all(clusters == ak.Array([[1, 1, 1, 1, 3, 1], [1, 1, 1, 2, 1], [1]]))
 
     # test without surface region
 
-    clusters = cluster.cluster_by_step_length(trackid, x, y, z, threshold=0.2)
+    clusters = cluster.cluster_by_step_length(trackid, x, y, z, threshold_in_mm=0.2)
     assert ak.all(clusters == ak.Array([[2, 1, 1, 3, 1], [1, 1, 1, 2, 1], [1]]))
 
 

tests/test_utils.py

@@ -78,11 +78,9 @@ def test_get_function_string():
     distances = awkward.Array([[0.2], [2], [0.6]])
 
     # just check it runs
-    assert eval(func_string, {"distances": distances}, globals_dict).view_as("ak")[0][0] == 0
-    assert eval(func_string, {"distances": distances}, globals_dict).view_as("ak")[1][0] == 1
-    assert eval(func_string, {"distances": distances}, globals_dict).view_as("ak")[2][
-        0
-    ] == pytest.approx(0.2)
+    assert eval(func_string, {"distances": distances}, globals_dict)[0][0] == 0
+    assert eval(func_string, {"distances": distances}, globals_dict)[1][0] == 1
+    assert eval(func_string, {"distances": distances}, globals_dict)[2][0] == pytest.approx(0.2)
 
     # try a more compliated expression
     expression = (