reboost 0.9.0__tar.gz → 0.10.0a0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reboost
-Version: 0.9.0
+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.0 → reboost-0.10.0a0}/pyproject.toml

@@ -37,7 +37,7 @@ dependencies = [
     "numpy",
     "scipy",
     "numba >=0.60",
-    "legend-pydataobj >=1.15.1",
+    "legend-pydataobj >=1.17.2",
     "legend-pygeom-optics >=0.15.0",
     "legend-pygeom-tools >=0.0.26",
     "legend-pygeom-hpges",

{reboost-0.9.0 → reboost-0.10.0a0}/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.10.0a0'
+__version_tuple__ = version_tuple = (0, 10, 0, 'a0')
 
-__commit_id__ = commit_id = 'g0f7fc07b4'
+__commit_id__ = commit_id = 'gd31e91903'

{reboost-0.9.0 → reboost-0.10.0a0}/src/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-0.9.0 → reboost-0.10.0a0}/src/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-0.9.0 → reboost-0.10.0a0}/src/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-0.9.0 → reboost-0.10.0a0}/src/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-0.9.0 → reboost-0.10.0a0}/src/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))