reboost 0.9.1__py3-none-any.whl → 0.10.0a0__py3-none-any.whl

reboost/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.9.1'
-__version_tuple__ = version_tuple = (0, 9, 1)
+__version__ = version = '0.10.0a0'
+__version_tuple__ = version_tuple = (0, 10, 0, 'a0')
 
 __commit_id__ = commit_id = None

reboost/build_hit.py

@@ -323,7 +323,7 @@ def build_hit(
                     if time_dict is not None:
                         start_time = time.time()
 
-                    ak_obj = stps.view_as("ak")
+                    ak_obj = stps.view_as("ak", with_units=True)
 
                     if time_dict is not None:
                         time_dict[proc_name].update_field("conv", start_time)
@@ -349,9 +349,6 @@ def build_hit(
                         if out_detector not in output_tables and files.hit[file_idx] is None:
                             output_tables[out_detector] = None
 
-                        # get the attributes
-                        attrs = utils.copy_units(stps)
-
                         # if we have more than one output detector, make an independent copy.
                         hit_table = (
                             copy.deepcopy(hit_table_layouted)
@@ -376,9 +373,6 @@ def build_hit(
                                 hit_table, outputs=proc_group["outputs"]
                             )
 
-                        # assign units in the output table
-                        hit_table = utils.assign_units(hit_table, attrs)
-
                         # now write
                         if files.hit[file_idx] is not None:
                             # get modes to write with
@@ -443,12 +437,7 @@ def build_hit(
 def _evaluate_operation(
     hit_table, field: str, info: str | dict, local_dict: dict, time_dict: ProfileDict
 ) -> None:
-    if isinstance(info, str):
-        expression = info
-        units = None
-    else:
-        expression = info["expression"]
-        units = info.get("units", None)
+    expression = info if isinstance(info, str) else info["expression"]
 
     # evaluate the expression
     col = core.evaluate_output_column(
@@ -460,7 +449,4 @@ def _evaluate_operation(
         name=field,
     )
 
-    if units is not None:
-        col.attrs["units"] = units
-
     core.add_field_with_nesting(hit_table, field, col)

reboost/core.py

@@ -18,18 +18,9 @@ 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],
-    with_units: bool = False,
+    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.
+    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
     ----------
@@ -77,9 +68,7 @@ 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])
-        units = data_ch.attrs.get("units", None)
-        data_ch = data_ch.view_as("ak")
+        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)]
@@ -96,12 +85,8 @@ 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 data_unflat
+    return ak.unflatten(data_flat, reorder)
 
 
 def evaluate_output_column(

reboost/hpge/psd.py

@@ -18,7 +18,7 @@ from .utils import HPGePulseShapeLibrary, HPGeRZField
 log = logging.getLogger(__name__)
 
 
-def r90(edep: ak.Array, xloc: ak.Array, yloc: ak.Array, zloc: ak.Array) -> Array:
+def r90(edep: ak.Array, xloc: ak.Array, yloc: ak.Array, zloc: ak.Array) -> ak.Array:
     """R90 HPGe pulse shape heuristic.
 
     Parameters
@@ -31,18 +31,22 @@ def r90(edep: ak.Array, xloc: ak.Array, yloc: ak.Array, zloc: ak.Array) -> Array
         array of y coordinate position.
     zloc
         array of z coordinate position.
+
+    Returns
+    -------
+    calculated r90 for each hit
     """
+    pos = [units.units_conv_ak(pos, "mm") for pos in [xloc, yloc, zloc]]
+
+    edep = units.units_conv_ak(edep, "keV")
+
     tot_energy = ak.sum(edep, axis=-1, keepdims=True)
 
     def eweight_mean(field, energy):
         return ak.sum(energy * field, axis=-1, keepdims=True) / tot_energy
 
     # Compute distance of each edep to the weighted mean
-    dist = np.sqrt(
-        (xloc - eweight_mean(edep, xloc)) ** 2
-        + (yloc - eweight_mean(edep, yloc)) ** 2
-        + (zloc - eweight_mean(edep, zloc)) ** 2
-    )
+    dist = np.sqrt(sum((p - eweight_mean(edep, p)) ** 2 for p in pos))
 
     # Sort distances and corresponding edep within each event
     sorted_indices = ak.argsort(dist, axis=-1)
@@ -76,16 +80,16 @@ def r90(edep: ak.Array, xloc: ak.Array, yloc: ak.Array, zloc: ak.Array) -> Array
     r90_indices = ak.argmax(cumsum_edep_corrected >= threshold, axis=-1, keepdims=True)
     r90 = sorted_dist[r90_indices]
 
-    return Array(ak.flatten(r90).to_numpy())
+    return units.attach_units(ak.Array(ak.flatten(r90)), "mm")
 
 
 def drift_time(
-    xloc: ArrayLike,
-    yloc: ArrayLike,
-    zloc: ArrayLike,
+    xloc: ak.Array | VectorOfVectors,
+    yloc: ak.Array | VectorOfVectors,
+    zloc: ak.Array | VectorOfVectors,
     dt_map: HPGeRZField,
     coord_offset: pint.Quantity | pyg4ometry.gdml.Position = (0, 0, 0) * u.m,
-) -> VectorOfVectors:
+) -> ak.Array:
     """Calculates drift times for each step (cluster) in an HPGe detector.
 
     Parameters
@@ -106,7 +110,7 @@ def drift_time(
     # sanitize coord_offset
     coord_offset = units.pg4_to_pint(coord_offset)
 
-    # unit handling (for matching with drift time map units)
+    # unit handling (.r_units) for data in (xloc, yloc)]
     xu, yu = [units.units_convfact(data, dt_map.r_units) for data in (xloc, yloc)]
     zu = units.units_convfact(zloc, dt_map.z_units)
 
@@ -122,7 +126,6 @@ def drift_time(
 
         return None
 
-    # transform coordinates
     xloc = xu * xloc - coord_offset[0].to(dt_map.r_units).m
     yloc = yu * yloc - coord_offset[1].to(dt_map.r_units).m
     zloc = zu * zloc - coord_offset[2].to(dt_map.z_units).m
@@ -133,16 +136,13 @@ def drift_time(
         np.sqrt(xloc**2 + yloc**2),
         zloc,
     )
-    return VectorOfVectors(
-        dt_values,
-        attrs={"units": units.unit_to_lh5_attr(dt_map.φ_units)},
-    )
+    return units.attach_units(ak.Array(dt_values), units.unit_to_lh5_attr(dt_map.φ_units))
 
 
 def drift_time_heuristic(
-    drift_time: ArrayLike,
-    edep: ArrayLike,
-) -> Array:
+    drift_time: ak.Array | VectorOfVectors,
+    edep: ak.Array | VectorOfVectors,
+) -> ak.Array:
     """HPGe drift-time-based pulse-shape heuristic.
 
     See :func:`_drift_time_heuristic_impl` for a description of the algorithm.
@@ -160,11 +160,11 @@ def drift_time_heuristic(
     edep, e_units = units.unwrap_lgdo(edep)
 
     # we want to attach the right units to the dt heuristic, if possible
-    attrs = {}
     if t_units is not None and e_units is not None:
-        attrs["units"] = units.unit_to_lh5_attr(t_units / e_units)
+        unit = units.unit_to_lh5_attr(t_units / e_units)
+        return units.attach_units(ak.Array(_drift_time_heuristic_impl(drift_time, edep)), unit)
 
-    return Array(_drift_time_heuristic_impl(drift_time, edep), attrs=attrs)
+    return ak.Array(_drift_time_heuristic_impl(drift_time, edep))
 
 
 @numba.njit(cache=True)
@@ -828,15 +828,15 @@ def maximum_current(
     activeness_surface: ArrayLike | None = None,
     surface_step_in_um: float = 10,
     return_mode: str = "current",
-) -> Array:
+) -> ak.Array:
     """Estimate the maximum current in the HPGe detector based on :func:`_estimate_current_impl`.
 
     Parameters
     ----------
     edep
-        Array of energies for each step.
+       Energies for each step.
     drift_time
-        Array of drift times for each step.
+        Drift times for each step.
     dist_to_nplus
         Distance to n-plus electrode, only needed if surface heuristics are enabled.
     r
@@ -844,7 +844,7 @@ def maximum_current(
     z
         z coordinate (only needed if a full PSS library is used).
     template
-        array of the bulk pulse template, can also be a :class:`HPGePulseShapeLibrary`.
+        Array of the bulk pulse template
     times
         time-stamps for the bulk pulse template
     fccd_in_um
@@ -864,12 +864,12 @@ def maximum_current(
     -------
     An Array of the maximum current/ time / energy for each hit.
     """
-    # extract LGDO data and units
-    drift_time, _ = units.unwrap_lgdo(drift_time)
-    edep, _ = units.unwrap_lgdo(edep)
-    dist_to_nplus, _ = units.unwrap_lgdo(dist_to_nplus)
-    r, _ = units.unwrap_lgdo(r)
-    z, _ = units.unwrap_lgdo(z)
+    # convert to target units
+    drift_time = units.units_conv_ak(drift_time, "ns")
+    edep = units.units_conv_ak(edep, "keV")
+    dist_to_nplus = units.units_conv_ak(dist_to_nplus, "mm")
+    r = units.units_conv_ak(r, "mm")
+    z = units.units_conv_ak(z, "mm")
 
     # prepare inputs for surface sims
     include_surface_effects, dist_to_nplus, templates_surface, activeness_surface = (
@@ -901,11 +901,12 @@ def maximum_current(
 
     # return
     if return_mode == "max_time":
-        return Array(time)
+        return units.attach_units(ak.Array(time), "ns")
     if return_mode == "current":
-        return Array(curr)
+        # current has no unit (depends on the template)
+        return ak.Array(curr)
     if return_mode == "energy":
-        return Array(energy)
+        return units.attach_units(ak.Array(energy), "keV")
 
     msg = f"Return mode {return_mode} is not implemented."
     raise ValueError(msg)

reboost/hpge/surface.py

@@ -5,27 +5,32 @@ import logging
 import awkward as ak
 import numba
 import numpy as np
+import pint
+import pyg4ometry
 import pygeomhpges
 from lgdo import VectorOfVectors
 from lgdo.types import LGDO
-from numpy.typing import ArrayLike
+from numpy.typing import NDArray
 from scipy import stats
 
+from reboost.units import ureg as u
+
+from .. import units
+
 log = logging.getLogger(__name__)
 
 
 def distance_to_surface(
-    positions_x: VectorOfVectors,
-    positions_y: VectorOfVectors,
-    positions_z: VectorOfVectors,
+    positions_x: ak.Array | VectorOfVectors,
+    positions_y: ak.Array | VectorOfVectors,
+    positions_z: ak.Array | VectorOfVectors,
     hpge: pygeomhpges.base.HPGe,
-    det_pos: ArrayLike,
+    det_pos: pint.Quantity | pyg4ometry.gdml.Position | tuple = (0, 0, 0) * u.m,
     *,
     surface_type: str | None = None,
-    unit: str = "mm",
-    distances_precompute: VectorOfVectors | None = None,
+    distances_precompute: ak.Array | None = None,
     precompute_cutoff: float | None = None,
-) -> VectorOfVectors:
+) -> ak.Array:
     """Computes the distance from each step to the detector surface.
 
     The calculation can be performed for any surface type `nplus`, `pplus`,
@@ -44,32 +49,36 @@ def distance_to_surface(
     hpge
         HPGe object.
     det_pos
-        position of the detector origin, must be a 3 component array corresponding to `(x,y,z)`.
+        position of the detector origin, must be a 3 component array corresponding to `(x,y,z)`. If no units
+        are specified mm is assumed.
     surface_type
         string of which surface to use, can be `nplus`, `pplus` `passive` or None (in which case the distance to any surface is calculated).
     unit
         unit for the hit tier positions table.
     distances_precompute
-        VectorOfVectors of distance to any surface computed by remage.
+        Distance to any surface computed by remage.
     precompute_cutoff
         cutoff on distances_precompute to not compute the distance for (in mm)
 
     Returns
     -------
-    VectorOfVectors with the same shape as `positions_x/y/z` of the distance to the surface.
+    Distance to the surface for each hit with the same shape as `positions_x/y/z`.
 
     Note
     ----
     `positions_x/positions_y/positions_z` must all have the same shape.
     """
-    factor = np.array([1, 100, 1000])[unit == np.array(["mm", "cm", "m"])][0]
-
     # compute local positions
     pos = []
     sizes = []
 
+    if not isinstance(det_pos, pint.Quantity | pyg4ometry.gdml.Position):
+        det_pos = det_pos * u.mm
+
+    det_pos = units.pg4_to_pint(det_pos)
+
     for idx, pos_tmp in enumerate([positions_x, positions_y, positions_z]):
-        local_pos_tmp = ak.Array(pos_tmp) * factor - det_pos[idx]
+        local_pos_tmp = units.units_conv_ak(pos_tmp, "mm") - det_pos[idx].to("mm").m
         local_pos_flat_tmp = ak.flatten(local_pos_tmp).to_numpy()
         pos.append(local_pos_flat_tmp)
         sizes.append(ak.num(local_pos_tmp, axis=1))
@@ -106,7 +115,7 @@ def distance_to_surface(
             local_positions[indices], surface_indices=surface_indices
         )
 
-    return VectorOfVectors(ak.unflatten(distances, size), dtype=np.float32)
+    return units.attach_units(ak.Array(ak.unflatten(distances, size)), "mm")
 
 
 @numba.njit(cache=True)
@@ -222,7 +231,7 @@ def get_surface_response(
     factor: float = 0.29,
     nsteps: int = 10000,
     delta_x: float = 10,
-):
+) -> NDArray:
     """Extract the surface response current pulse based on diffusion.
 
     This extracts the amount of charge arrived (cumulative) at the p-n
@@ -248,6 +257,10 @@ def get_surface_response(
         the number of time steps.
     delta_x
         the width of each position bin.
+
+    Returns
+    -------
+    2D array of the amount of charge arriving at the p-n junction as a function of time for each depth.
     """
     # number of position steps
     nx = int(fccd / delta_x)

reboost/math/functions.py

@@ -5,14 +5,13 @@ import logging
 import awkward as ak
 import numpy as np
 from lgdo import Array, VectorOfVectors
-from lgdo.types import LGDO
+
+from .. import units
 
 log = logging.getLogger(__name__)
 
 
-def piecewise_linear_activeness(
-    distances: VectorOfVectors | ak.Array, fccd: float, dlf: float
-) -> VectorOfVectors | Array:
+def piecewise_linear_activeness(distances: ak.Array, fccd_in_mm: float, dlf: float) -> ak.Array:
     r"""Piecewise linear HPGe activeness model.
 
     Based on:
@@ -38,11 +37,10 @@ def piecewise_linear_activeness(
     Parameters
     ----------
     distances
-        the distance from each step to the detector surface. Can be either a
-        `numpy` or `awkward` array, or a LGDO `VectorOfVectors` or `Array`. The computation
+        the distance from each step to the detector surface. The computation
         is performed for each element and the shape preserved in the output.
 
-    fccd
+    fccd_in_mm
         the value of the FCCD
     dlf
         the fraction of the FCCD which is fully inactive.
@@ -52,14 +50,9 @@ def piecewise_linear_activeness(
     a :class:`VectorOfVectors` or :class:`Array` of the activeness
     """
     # convert to ak
-    if isinstance(distances, LGDO):
-        distances_ak = distances.view_as("ak")
-    elif not isinstance(distances, ak.Array):
-        distances_ak = ak.Array(distances)
-    else:
-        distances_ak = distances
+    distances_ak = units.units_conv_ak(distances, "mm")
 
-    dl = fccd * dlf
+    dl = fccd_in_mm * dlf
     distances_flat = (
         ak.flatten(distances_ak).to_numpy() if distances_ak.ndim > 1 else distances_ak.to_numpy()
     )
@@ -68,24 +61,26 @@ def piecewise_linear_activeness(
     results = np.full_like(distances_flat, np.nan, dtype=np.float64)
     lengths = ak.num(distances_ak) if distances_ak.ndim > 1 else len(distances_ak)
 
-    mask1 = (distances_flat > fccd) | np.isnan(distances_flat)
+    mask1 = (distances_flat > fccd_in_mm) | np.isnan(distances_flat)
     mask2 = (distances_flat <= dl) & (~mask1)
     mask3 = ~(mask1 | mask2)
 
     # assign the values
     results[mask1] = 1
     results[mask2] = 0
-    results[mask3] = (distances_flat[mask3] - dl) / (fccd - dl)
+    results[mask3] = (distances_flat[mask3] - dl) / (fccd_in_mm - dl)
 
     # reshape
     results = ak.unflatten(ak.Array(results), lengths) if distances_ak.ndim > 1 else results
 
-    return VectorOfVectors(results) if results.ndim > 1 else Array(results)
+    results = ak.Array(results)
+
+    return units.attach_units(results, "mm")
 
 
 def vectorised_active_energy(
-    distances: VectorOfVectors | ak.Array,
-    edep: VectorOfVectors | ak.Array,
+    distances: ak.Array,
+    edep: ak.Array,
     fccd: float | list,
     dlf: float | list,
 ) -> VectorOfVectors | Array:
@@ -115,7 +110,7 @@ def vectorised_active_energy(
 
     Returns
     -------
-    a :class:`VectorOfVectors` or :class:`Array` of the activeness
+    Activeness for each set of parameters
     """
     # add checks on fccd, dlf
     fccd = np.array(fccd)
@@ -133,20 +128,14 @@ def vectorised_active_energy(
 
     dl = fccd * dlf
 
-    def _convert(field):
+    def _convert(field, unit):
         # convert to ak
-        if isinstance(field, VectorOfVectors):
-            field_ak = field.view_as("ak")
-        elif not isinstance(field, ak.Array):
-            field_ak = ak.Array(field)
-        else:
-            msg = f"{field} must be an awkward array or VectorOfVectors"
-            raise TypeError(msg)
+        field_ak = units.units_conv_ak(field, unit)
 
         return field_ak, ak.flatten(field_ak).to_numpy()[:, np.newaxis]
 
-    distances_ak, distances_flat = _convert(distances)
-    _, edep_flat = _convert(edep)
+    distances_ak, distances_flat = _convert(distances, "mm")
+    _, edep_flat = _convert(edep, "keV")
     runs = ak.num(distances_ak, axis=-1)
 
     # vectorise fccd or tl
@@ -172,4 +161,4 @@ def vectorised_active_energy(
 
     energy = ak.sum(ak.unflatten(results * edep_flat, runs), axis=-2)
 
-    return VectorOfVectors(energy) if energy.ndim > 1 else Array(energy.to_numpy())
+    return units.attach_units(energy, "keV")

reboost/math/stats.py

@@ -6,7 +6,6 @@ from collections.abc import Callable
 import awkward as ak
 import numpy as np
 from lgdo import Array
-from numpy.typing import ArrayLike
 
 log = logging.getLogger(__name__)
 
@@ -72,7 +71,7 @@ def apply_energy_resolution(
     return ak.unflatten(energies_flat_smear, num)
 
 
-def gaussian_sample(mu: ArrayLike, sigma: ArrayLike | float, *, seed: int | None = None) -> Array:
+def gaussian_sample(mu: ak.Array, sigma: ak.Array | float, *, seed: int | None = None) -> ak.Array:
     r"""Generate samples from a gaussian.
 
     Based on:
@@ -99,9 +98,7 @@ def gaussian_sample(mu: ArrayLike, sigma: ArrayLike | float, *, seed: int | None
     """
     # convert inputs
 
-    if isinstance(mu, Array):
-        mu = mu.view_as("np")
-    elif isinstance(mu, ak.Array):
+    if isinstance(mu, ak.Array):
         mu = mu.to_numpy()
     elif not isinstance(mu, np.ndarray):
         mu = np.array(mu)
@@ -116,4 +113,4 @@ def gaussian_sample(mu: ArrayLike, sigma: ArrayLike | float, *, seed: int | None
 
     rng = np.random.default_rng(seed=seed)  # Create a random number generator
 
-    return Array(rng.normal(loc=mu, scale=sigma))
+    return ak.Array(rng.normal(loc=mu, scale=sigma))

reboost/shape/cluster.py

@@ -7,12 +7,12 @@ import numba
 import numpy as np
 from lgdo import VectorOfVectors
 
+from .. import units
+
 log = logging.getLogger(__name__)
 
 
-def apply_cluster(
-    cluster_run_lengths: VectorOfVectors | ak.Array, field: ak.Array | VectorOfVectors
-) -> VectorOfVectors:
+def apply_cluster(cluster_run_lengths: ak.Array, field: ak.Array) -> ak.Array:
     """Apply clustering to a field.
 
     Parameters
@@ -26,24 +26,26 @@ def apply_cluster(
         cluster_run_lengths = cluster_run_lengths.view_as("ak")
 
     if isinstance(field, VectorOfVectors):
-        field = field.view_as("ak")
+        field = field.view_as("ak", with_units=True)
 
     n_cluster = ak.num(cluster_run_lengths, axis=-1)
     clusters = ak.unflatten(ak.flatten(field), ak.flatten(cluster_run_lengths))
 
     # reshape into cluster oriented
-    return VectorOfVectors(ak.unflatten(clusters, n_cluster))
+    return units.attach_units(
+        ak.Array(ak.unflatten(clusters, n_cluster)), units.get_unit_str(field)
+    )
 
 
 def cluster_by_step_length(
-    trackid: ak.Array | VectorOfVectors,
-    pos_x: ak.Array | VectorOfVectors,
-    pos_y: ak.Array | VectorOfVectors,
-    pos_z: ak.Array | VectorOfVectors,
-    dist: ak.Array | VectorOfVectors | None = None,
+    trackid: ak.Array,
+    pos_x: ak.Array,
+    pos_y: ak.Array,
+    pos_z: ak.Array,
+    dist: ak.Array | None = None,
     surf_cut: float | None = None,
-    threshold: float = 0.1,
-    threshold_surf: float | None = None,
+    threshold_in_mm: float = 0.1,
+    threshold_surf_in_mm: float | None = None,
 ) -> VectorOfVectors:
     """Perform clustering based on the step length.
 
@@ -59,20 +61,20 @@ def cluster_by_step_length(
     Parameters
     ----------
     trackid
-        index of the track.
+        index of the tracks.
     pos_x
-        x position of the step.
+        x position of the steps.
     pos_y
-        y position of the step.
+        y position of the steps.
     pos_z
-        z position of the step.
+        z position of the steps.
     dist
         distance to the detector surface. Can be `None` in which case all steps are treated as being in the "bulk".
     surf_cut
         Size of the surface region (in mm), if `None` no selection is applied (default).
-    threshold
+    threshold_in_mm
         Distance threshold in mm to combine steps in the bulk.
-    threshold_surf
+    threshold_surf_in_mm
         Distance threshold in mm to combine steps in the surface.
 
     Returns
@@ -80,37 +82,22 @@ def cluster_by_step_length(
     Array of the run lengths of each cluster within a hit.
     """
     # type conversions
-    if isinstance(pos_x, VectorOfVectors):
-        pos_x = pos_x.view_as("ak")
-
-    if isinstance(pos_y, VectorOfVectors):
-        pos_y = pos_y.view_as("ak")
-
-    if isinstance(pos_z, VectorOfVectors):
-        pos_z = pos_z.view_as("ak")
-
-    if isinstance(trackid, VectorOfVectors):
-        trackid = trackid.view_as("ak")
-
-    if isinstance(dist, VectorOfVectors):
-        dist = dist.view_as("ak")
 
     pos = np.vstack(
         [
-            ak.flatten(pos_x).to_numpy().astype(np.float64),
-            ak.flatten(pos_y).to_numpy().astype(np.float64),
-            ak.flatten(pos_z).to_numpy().astype(np.float64),
+            ak.flatten(units.units_conv_ak(p, "mm")).to_numpy().astype(np.float64)
+            for p in [pos_x, pos_y, pos_z]
         ]
     ).T
 
-    indices_flat = cluster_by_distance_numba(
+    indices_flat = _cluster_by_distance_numba(
         ak.flatten(ak.local_index(trackid)).to_numpy(),
         ak.flatten(trackid).to_numpy(),
         pos,
         dist_to_surf=ak.flatten(dist).to_numpy() if dist is not None else dist,
         surf_cut=surf_cut,
-        threshold=threshold,
-        threshold_surf=threshold_surf,
+        threshold=threshold_in_mm,
+        threshold_surf=threshold_surf_in_mm,
     )
 
     # reshape into being event oriented
@@ -119,11 +106,11 @@ def cluster_by_step_length(
     # number of steps per cluster
     counts = ak.run_lengths(indices)
 
-    return VectorOfVectors(counts)
+    return ak.Array(counts)
 
 
 @numba.njit
-def cluster_by_distance_numba(
+def _cluster_by_distance_numba(
     local_index: np.ndarray,
     trackid: np.ndarray,
     pos: np.ndarray,
@@ -210,10 +197,10 @@ def cluster_by_distance_numba(
 
 
 def step_lengths(
-    x_cluster: ak.Array | VectorOfVectors,
-    y_cluster: ak.Array | VectorOfVectors,
-    z_cluster: ak.Array | VectorOfVectors,
-) -> VectorOfVectors:
+    x_cluster: ak.Array,
+    y_cluster: ak.Array,
+    z_cluster: ak.Array,
+) -> ak.Array:
     """Compute the distance between consecutive steps.
 
     This is based on calculating the distance between consecutive steps in the same track,
@@ -242,13 +229,14 @@ def step_lengths(
     data = [x_cluster, y_cluster, z_cluster]
 
     for idx, var in enumerate(data):
-        if isinstance(var, VectorOfVectors):
-            data[idx] = var.view_as("ak")
         # check shape
-        if data[idx].ndim != 3:
+        if var.ndim != 3:
             msg = f"The input array for step lengths must be 3 dimensional not {data[idx.dim]}"
             raise ValueError(msg)
 
+        # type convert
+        data[idx] = units.units_conv_ak(data[idx], "mm")
+
     counts = ak.num(data[0], axis=-1)
     data = np.vstack([ak.flatten(ak.flatten(var)).to_numpy() for var in data])
     dist = np.append(np.sqrt(np.sum(np.diff(data, axis=1) ** 2, axis=0)), 0)
@@ -257,4 +245,4 @@ def step_lengths(
     clusters = ak.unflatten(ak.Array(dist), ak.flatten(counts))
 
     out = ak.unflatten(clusters, n_cluster)
-    return VectorOfVectors(out[:, :, :-1])
+    return units.attach_units(ak.Array(out[:, :, :-1]), "mm")

reboost/shape/group.py

@@ -8,6 +8,8 @@ from dbetto import AttrsDict
 from lgdo import Table, VectorOfVectors
 from numpy.typing import ArrayLike
 
+from reboost import units
+
 log = logging.getLogger(__name__)
 
 
@@ -63,13 +65,19 @@ def _sort_data(obj: ak.Array, *, time_name: str = "time", evtid_name: str = "evt
     -------
     sorted awkward array
     """
+    units_dict = {field: units.get_unit_str(obj[field]) for field in obj.fields}
+
     obj = obj[ak.argsort(obj[evtid_name])]
     obj_unflat = ak.unflatten(obj, ak.run_lengths(obj[evtid_name]))
 
     indices = ak.argsort(obj_unflat[time_name], axis=-1)
     sorted_obj = obj_unflat[indices]
 
-    return ak.flatten(sorted_obj)
+    out = {}
+    for field in sorted_obj.fields:
+        out[field] = units.attach_units(ak.flatten(sorted_obj[field]), units_dict[field])
+
+    return ak.Array(out)
 
 
 def group_by_evtid(data: Table | ak.Array, *, evtid_name: str = "evtid") -> Table:
@@ -95,7 +103,9 @@ def group_by_evtid(data: Table | ak.Array, *, evtid_name: str = "evtid") -> Tabl
     The input table must be sorted (by `evtid`).
     """
     # convert to awkward
-    obj_ak = data.view_as("ak") if isinstance(data, Table) else data
+    obj_ak = data.view_as("ak", with_units=True) if isinstance(data, Table) else data
+
+    units_dict = {field: units.get_unit_str(obj_ak[field]) for field in obj_ak.fields}
 
     # extract cumulative lengths
     counts = ak.run_lengths(obj_ak[evtid_name])
@@ -111,10 +121,14 @@ def group_by_evtid(data: Table | ak.Array, *, evtid_name: str = "evtid") -> Tabl
     for f in obj_ak.fields:
         out_tbl.add_field(
             f,
-            VectorOfVectors(
-                cumulative_length=cumulative_length, flattened_data=obj_ak[f].to_numpy()
+            units.attach_units(
+                VectorOfVectors(
+                    cumulative_length=cumulative_length, flattened_data=obj_ak[f].to_numpy()
+                ),
+                units_dict[f],
             ),
         )
+
     return out_tbl
 
 
@@ -139,7 +153,7 @@ def group_by_time(
     data
         :class:`lgdo.Table` or `ak.Array` which must contain the time_name and evtid_name fields
     window
-        time window in us used to search for coincident hits
+        time window in us used to search for coincident hits.
     time_name
         name of the timing field
     evtid_name
@@ -155,15 +169,22 @@ def group_by_time(
     ----
     The input table must be sorted (first by `evtid` then `time`).
     """
-    obj = data.view_as("ak") if isinstance(data, Table) else data
+    obj = data.view_as("ak", with_units=True) if isinstance(data, Table) else data
+    units_dict = {field: units.get_unit_str(obj[field]) for field in obj.fields}
+
+    window = window * 1000.0  # convert to ns
     obj = _sort_data(obj, time_name=time_name, evtid_name=evtid_name)
 
+    # convert to ns
+    obj[time_name] = units.units_conv_ak(obj[time_name], "ns")
+    units_dict[time_name] = "ns"
+
     # get difference
     time_diffs = np.diff(obj[time_name])
     index_diffs = np.diff(obj[evtid_name])
 
     # index of the last element in each run
-    time_change = (time_diffs > window * 1000) & (index_diffs == 0)
+    time_change = (time_diffs > window) & (index_diffs == 0)
     index_change = index_diffs > 0
 
     # cumulative length is just the index of changes plus 1
@@ -183,7 +204,12 @@ def group_by_time(
     for f in fields:
         out_tbl.add_field(
             f,
-            VectorOfVectors(cumulative_length=cumulative_length, flattened_data=obj[f].to_numpy()),
+            units.attach_units(
+                VectorOfVectors(
+                    cumulative_length=cumulative_length, flattened_data=obj[f].to_numpy()
+                ),
+                units_dict[f],
+            ),
         )
 
     return out_tbl

reboost/units.py

@@ -4,6 +4,7 @@ import logging
 from typing import Any
 
 import awkward as ak
+import numpy as np
 import pint
 import pyg4ometry as pg4
 from lgdo import LGDO
@@ -45,7 +46,30 @@ def units_convfact(data: Any | LGDO | ak.Array, target_units: pint.Unit | str) -
     return 1
 
 
-def units_conv_ak(data: Any | LGDO | ak.Array, target_units: pint.Unit | str) -> ak.Array:
+def attach_units(data: ak.Array | LGDO, unit: str | None) -> ak.Array | LGDO:
+    """Convenience function to attach units to `ak.Array` or LGDO.
+
+    Parameters
+    ----------
+    data
+        the array to add units to
+    unit
+        the unit
+    """
+    if isinstance(data, ak.Array):
+        if unit is not None:
+            return ak.with_parameter(data, parameter="units", value=unit)
+    elif isinstance(data, LGDO):
+        if unit is not None:
+            data.attrs["units"] = unit
+    else:
+        msg = f"to attach units data must be an ak.Array or LGDO not {type(data)}"
+        raise TypeError(msg)
+
+    return data
+
+
+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.
 
     Parameters
@@ -62,7 +86,13 @@ def units_conv_ak(data: Any | LGDO | ak.Array, target_units: pint.Unit | str) ->
         return ak.without_parameters(data.view_as("ak") * fact)
     if isinstance(data, ak.Array) and fact != 1:
         return ak.without_parameters(data * fact)
-    return data.view_as("ak") if isinstance(data, LGDO) else data
+
+    # try to return ak.Array if possible
+    if isinstance(data, LGDO):
+        return data.view_as("ak")
+    if isinstance(data, np.ndarray):
+        return ak.Array(data)
+    return data
 
 
 def unwrap_lgdo(data: Any | LGDO | ak.Array, library: str = "ak") -> tuple[Any, pint.Unit | None]:
@@ -99,6 +129,31 @@ def unwrap_lgdo(data: Any | LGDO | ak.Array, library: str = "ak") -> tuple[Any,
     return ret_data, ret_units
 
 
+def get_unit_str(data: ak.Array | LGDO) -> str | None:
+    """Get the units as a string for an awkward array with attached units.
+
+    Parameters
+    ----------
+    data
+        the array with attached units
+
+    Returns
+    -------
+    a string of the units.
+    """
+    if isinstance(data, ak.Array):
+        attrs = ak.parameters(data)
+    elif isinstance(data, LGDO):
+        attrs = data.attrs
+    else:
+        msg = f"Cannot extract units: {data} is not an LGDO or ak.Array"
+        raise ValueError(msg)
+
+    if "units" in attrs:
+        return attrs["units"]
+    return None
+
+
 def unit_to_lh5_attr(unit: pint.Unit) -> str:
     """Convert Pint unit to a string that can be used as attrs["units"] in an LGDO."""
     # TODO: we should check if this can be always parsed by Unitful.jl

{reboost-0.9.1.dist-info → reboost-0.10.0a0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reboost
-Version: 0.9.1
+Version: 0.10.0a0
 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
@@ -26,7 +26,7 @@ Requires-Dist: colorlog
 Requires-Dist: numpy
 Requires-Dist: scipy
 Requires-Dist: numba>=0.60
-Requires-Dist: legend-pydataobj>=1.15.1
+Requires-Dist: legend-pydataobj>=1.17.2
 Requires-Dist: legend-pygeom-optics>=0.15.0
 Requires-Dist: legend-pygeom-tools>=0.0.26
 Requires-Dist: legend-pygeom-hpges

{reboost-0.9.1.dist-info → reboost-0.10.0a0.dist-info}/RECORD

@@ -1,26 +1,26 @@
 reboost/__init__.py,sha256=VZz9uo7i2jgAx8Zi15SptLZnE_qcnGuNWwqkD3rYHFA,278
 reboost/__main__.py,sha256=42koSxY2st4mMIRSAnKz06nP5HppMPxBVFf2jaHljGs,95
-reboost/_version.py,sha256=LwGndsRSpclYq-j3wgRr2nzOXwUYj0Jtg7Kof7R0BEw,704
+reboost/_version.py,sha256=uvugpfFX1pEScMzJ2JeDcYh7lOMZajzrzSxwt8iaOLU,714
 reboost/build_evt.py,sha256=VXIfK_pfe_Cgym6gI8dESwONZi-v_4fll0Pn09vePQY,3767
 reboost/build_glm.py,sha256=IerSLQfe51ZO7CQP2kmfPnOIVaDtcfw3byOM02Vaz6o,9472
-reboost/build_hit.py,sha256=N_nxvH69SvILVNmyvVfhQwQdD_PDW8tlsqj2ciO5nKE,17409
+reboost/build_hit.py,sha256=fg0vDPnWCQN7Px0KKLuQvwdsA4DBT2z6IoAb_KjG9CY,17034
 reboost/cli.py,sha256=68EzKiWTHJ2u1RILUv7IX9HaVq6nTTM80_W_MUnWRe4,6382
-reboost/core.py,sha256=NtCDTZ6QQyFhj0BiGuIXBu6WBSrrY7C4GTfaw1u5K-w,15501
+reboost/core.py,sha256=TPxvZgUaHZdxfQSDdX2zIerQXt3Gq-zQaA6AeXZKNvA,15232
 reboost/iterator.py,sha256=qlEqRv5qOh8eIs-dyVOLYTvH-ZpQDx9fLckpcAdtWjs,6975
 reboost/log_utils.py,sha256=VqS_9OC5NeNU3jcowVOBB0NJ6ssYvNWnirEY-JVduEA,766
 reboost/profile.py,sha256=EOTmjmS8Rm_nYgBWNh6Rntl2XDsxdyed7yEdWtsZEeg,2598
-reboost/units.py,sha256=LUwl6swLQoG09Rt9wcDdu6DTrwDsy-C751BNGzX4sz8,3651
+reboost/units.py,sha256=Jj1OjYoY6xoWTVGId_w8Cn5ugN0QLGluVhWvLeqeyII,5019
 reboost/utils.py,sha256=vl-_BUOeXcazNs4zN-9k-OVEptdf3FtCeej2QZhClKc,14599
 reboost/daq/__init__.py,sha256=rNPhxx1Yawt3tENYhmOYSum9_TdV57ZU5kjxlWFAGuo,107
 reboost/daq/core.py,sha256=Rs6Q-17fzEod2iX_2WqEmnqKnNRFoWTYURl3wYhFihU,9915
 reboost/daq/utils.py,sha256=KcH6zvlInmD2YiF6V--DSYBTYudJw3G-hp2JGOcES2o,1042
 reboost/hpge/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-reboost/hpge/psd.py,sha256=P7dUJQPvxW6vndJ79r0j7ANZvSuV_IuauERhvWD74j0,26989
-reboost/hpge/surface.py,sha256=feH-kxRRp3HkikRRJ-LCu6zvVONEOYVd3THx12emGTM,8494
+reboost/hpge/psd.py,sha256=ORGZepsevHq2m5QVPuPLJ-AKLChLpCZZHkQpLR5TTcM,27246
+reboost/hpge/surface.py,sha256=9c0Ks1xJidqtnwMl3tI0sufSZTekGfSAhuk_b57daWU,8886
 reboost/hpge/utils.py,sha256=0hHu5S1lDOiMMVPgfHY03R5ggyeWX2OwONHVZeFmcpU,5652
 reboost/math/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-reboost/math/functions.py,sha256=OymiYTcA0NXxxm-MBDw5kqyNwHoLCmuv4J48AwnSrbU,5633
-reboost/math/stats.py,sha256=Rq4Wdzv-3aoSK7EsPZCuOEHfnOz3w0moIzCEHbC07xw,3173
+reboost/math/functions.py,sha256=Nw6p4fxF68TmpHYcpDriisXvudOB8Ep26TlckVFGrME,5045
+reboost/math/stats.py,sha256=FOC9Qyixt8L6ccqGoavADU3LqWkbJkYfURiR-XIwSpE,3080
 reboost/optmap/__init__.py,sha256=imvuyld-GLw8qdwqW-lXCg2feptcTyQo3wIzPvDHwmY,93
 reboost/optmap/__main__.py,sha256=DfzkXQ7labOe53hd7jH5pAbTW491jjQYSMLyl72L4Rk,111
 reboost/optmap/cli.py,sha256=_7WBlx55eRyW_wWB-ELbFaWXin2d3xsh6Q5bFoNJaHE,8694
@@ -31,14 +31,14 @@ reboost/optmap/mapview.py,sha256=fswwXolA6au8u8gljBKy8PSXC2W7Cy_GwOV86-duYG8,688
 reboost/optmap/numba_pdg.py,sha256=y8cXR5PWE2Liprp4ou7vl9do76dl84vXU52ZJD9_I7A,731
 reboost/optmap/optmap.py,sha256=3clc1RA8jA4YJte83w085MY8zLpG-G7DBkpZ2UeKPpM,12825
 reboost/shape/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-reboost/shape/cluster.py,sha256=nwR1Dnf00SDICGPqpXeM1Q7_DwTtO9uP3wmuML45c3g,8195
-reboost/shape/group.py,sha256=gOCYgir2gZqmW1JXtbNRPlQqP0gmUcbe7RVb9CbY1pU,5540
+reboost/shape/cluster.py,sha256=EHPgjUrbF-uMXQ3D7VPKkGKgLnea6ZBR-R_1DSuGmLU,7671
+reboost/shape/group.py,sha256=PY7FGUnralj1BkRl4YtguCnzvOr-ymMd4ZFVmzPCwmM,6352
 reboost/shape/reduction.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 reboost/spms/__init__.py,sha256=8I6WT8i_kUPqEDnSD0aCf6A26cjKjQQZSNrvwZ3o-Ac,415
 reboost/spms/pe.py,sha256=LwqrK1HOZWzGcNZnntaqI6r4rnDww4KW9Mao4xLFbDE,8226
-reboost-0.9.1.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-reboost-0.9.1.dist-info/METADATA,sha256=nW1jNoPEgO6PPTKtW0U_eDQ8ua35E3MVgumUrkYb1Y4,3877
-reboost-0.9.1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-reboost-0.9.1.dist-info/entry_points.txt,sha256=DxhD6BidSWNot9BrejHJjQ7RRLmrMaBIl52T75oWTwM,93
-reboost-0.9.1.dist-info/top_level.txt,sha256=q-IBsDepaY_AbzbRmQoW8EZrITXRVawVnNrB-_zyXZs,8
-reboost-0.9.1.dist-info/RECORD,,
+reboost-0.10.0a0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+reboost-0.10.0a0.dist-info/METADATA,sha256=o7lnoQcj8MkJ-VW7NQKwE75vlHx4qJAV4ycWgb8Yvuc,3880
+reboost-0.10.0a0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+reboost-0.10.0a0.dist-info/entry_points.txt,sha256=DxhD6BidSWNot9BrejHJjQ7RRLmrMaBIl52T75oWTwM,93
+reboost-0.10.0a0.dist-info/top_level.txt,sha256=q-IBsDepaY_AbzbRmQoW8EZrITXRVawVnNrB-_zyXZs,8
+reboost-0.10.0a0.dist-info/RECORD,,

{reboost-0.9.1.dist-info → reboost-0.10.0a0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.10.2)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any