reboost 0.8.5__tar.gz → 0.9.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reboost
-Version: 0.8.5
+Version: 0.9.0
 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.8.5 → reboost-0.9.0}/src/reboost/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.8.5'
-__version_tuple__ = version_tuple = (0, 8, 5)
+__version__ = version = '0.9.0'
+__version_tuple__ = version_tuple = (0, 9, 0)
 
-__commit_id__ = commit_id = 'g5d57bf6e0'
+__commit_id__ = commit_id = 'g0f7fc07b4'

{reboost-0.8.5 → reboost-0.9.0}/src/reboost/hpge/psd.py

@@ -13,7 +13,7 @@ from numpy.typing import ArrayLike, NDArray
 
 from .. import units
 from ..units import ureg as u
-from .utils import HPGeScalarRZField
+from .utils import HPGePulseShapeLibrary, HPGeRZField
 
 log = logging.getLogger(__name__)
 
@@ -83,7 +83,7 @@ def drift_time(
     xloc: ArrayLike,
     yloc: ArrayLike,
     zloc: ArrayLike,
-    dt_map: HPGeScalarRZField,
+    dt_map: HPGeRZField,
     coord_offset: pint.Quantity | pyg4ometry.gdml.Position = (0, 0, 0) * u.m,
 ) -> VectorOfVectors:
     """Calculates drift times for each step (cluster) in an HPGe detector.
@@ -454,33 +454,7 @@ def _get_waveform_value_surface(
     start: float,
     dt: float,
 ) -> tuple[float, float]:
-    """Get the value of the waveform at a certain index.
-
-    Parameters
-    ----------
-    idx
-        the index of the time array to find the waveform at.
-    edep
-        Array of energies for each step
-    drift_time
-        Array of drift times for each step
-    template
-        array of the template for the current waveforms
-    templates_surface
-        The current templates from the surface.
-    activeness_surface
-        The total collected charge for each surface point.
-    dist_step_in_um
-        The binning in distance for the surface pulse library.
-    start
-        first time value of the template
-    dt
-        timestep (in ns) for the template.
-
-    Returns
-    -------
-    Value of the current waveform and the energy.
-    """
+    """Get the value of the waveform at a certain index."""
     n = len(bulk_template)
     out = 0
     etmp = 0
@@ -513,6 +487,7 @@ def _get_waveform_value_surface(
         else:
             out += E * _interpolate_pulse_model(bulk_template, time, start, start + dt * n, dt, mu)
             etmp += E
+
     return out, etmp
 
 
@@ -525,40 +500,78 @@ def _get_waveform_value(
     start: float,
     dt: float,
 ) -> float:
-    """Get the value of the waveform at a certain index.
+    """Get the value of the waveform at a certain index."""
+    out = 0
+    time = start + dt * idx
 
-    Parameters
-    ----------
-    idx
-        the index of the time array to find the waveform at.
-    edep
-        Array of energies for each step
-    drift_time
-        Array of drift times for each step
-    template
-        array of the template for the current waveforms
-    start
-        first time value of the template
-    dt
-        timestep (in ns) for the template.
+    for i in range(len(edep)):
+        n = len(template)
 
-    Returns
-    -------
-    Value of the current waveform
-    """
-    n = len(template)
+        E = edep[i]
+        mu = drift_time[i]
+
+        out += E * _interpolate_pulse_model(template, time, start, start + dt * n, dt, mu)
+
+    return out
+
+
+@numba.njit(cache=True)
+def _get_waveform_value_pulse_shape_library(
+    idx: int,
+    edep: ak.Array,
+    drift_time: ak.Array,
+    r: ak.Array,
+    z: ak.Array,
+    pulse_shape_library: tuple[np.array, np.array, np.array],
+    start: float,
+    dt: float,
+) -> float:
+    """Get the value of the waveform at a certain index, using the pulse shape library."""
     out = 0
     time = start + dt * idx
 
     for i in range(len(edep)):
+        ri, zi = _get_template_idx(r[i], z[i], pulse_shape_library[1], pulse_shape_library[2])
+
+        n = len(pulse_shape_library[0][ri][zi])
+
         E = edep[i]
         mu = drift_time[i]
 
-        out += E * _interpolate_pulse_model(template, time, start, start + dt * n, dt, mu)
+        out += E * _interpolate_pulse_model(
+            pulse_shape_library[0][ri][zi], time, start, start + dt * n, dt, mu
+        )
 
     return out
 
 
+@numba.njit(cache=True)
+def _get_template_idx(
+    r: float,
+    z: float,
+    r_grid: np.array,
+    z_grid: np.array,
+) -> tuple[int, int]:
+    """Extract the closest template to a given (r,z) point with uniform grid, apart from the first and last point."""
+    if r < r_grid[1]:
+        ri = 0
+    elif r > r_grid[-2]:
+        ri = len(r_grid) - 1
+    else:
+        dr = r_grid[2] - r_grid[1]
+        ri = int((r - r_grid[1]) / dr) + 1
+
+    if z < z_grid[1]:
+        zi = 0
+    elif z > z_grid[-2]:
+        zi = len(z_grid) - 1
+    else:
+        dz = z_grid[2] - z_grid[1]
+        zi = int((z - z_grid[1]) / dz) + 1
+
+    return ri, zi
+
+
 def get_current_template(
     low: float = -1000, high: float = 4000, step: float = 1, mean_aoe: float = 1, **kwargs
 ) -> tuple[NDArray, NDArray]:
@@ -598,7 +611,10 @@ def _get_waveform_maximum_impl(
     t: ArrayLike,
     e: ArrayLike,
     dist: ArrayLike,
+    r: ArrayLike,
+    z: ArrayLike,
     template: ArrayLike,
+    pulse_shape_library: tuple[np.array, np.array, np.array],
     templates_surface: ArrayLike,
     activeness_surface: ArrayLike,
     tmin: float,
@@ -609,18 +625,9 @@ def _get_waveform_maximum_impl(
     time_step: int,
     surface_step_in_um: float,
     include_surface_effects: bool,
+    use_library: bool,
 ):
-    """Basic implementation to get the maximum of the waveform.
-
-    Parameters
-    ----------
-    t
-        drift time for each step.
-    e
-        energy for each step.
-    dist
-        distance to surface for each step.
-    """
+    """Basic implementation to get the maximum of the waveform."""
     max_a = 0
     max_t = 0
     energy = np.sum(e)
@@ -634,8 +641,12 @@ def _get_waveform_maximum_impl(
         if time < tmin or (time > (tmax + time_step)):
             continue
 
-        if not has_surface_hit:
+        if not has_surface_hit and (not use_library):
             val_tmp = _get_waveform_value(j, e, t, template, start=start, dt=1.0)
+        elif use_library:
+            val_tmp = _get_waveform_value_pulse_shape_library(
+                j, e, t, r, z, pulse_shape_library, start=start, dt=1.0
+            )
         else:
             val_tmp, energy = _get_waveform_value_surface(
                 j,
@@ -663,9 +674,13 @@ def _estimate_current_impl(
     edep: ak.Array,
     dt: ak.Array,
     dist_to_nplus: ak.Array,
+    r: ak.Array,
+    z: ak.Array,
     template: np.array,
+    pulse_shape_library: tuple[np.array, np.array, np.array],
     times: np.array,
     include_surface_effects: bool,
+    use_library: bool,
     fccd: float,
     templates_surface: np.array,
     activeness_surface: np.array,
@@ -674,26 +689,13 @@ def _estimate_current_impl(
     """Estimate the maximum current that would be measured in the HPGe detector.
 
     This is based on extracting a waveform with :func:`get_current_waveform` and finding the maxima of it.
-
-    Parameters
-    ----------
-    edep
-        Array of energies for each step.
-    dt
-        Array of drift times for each step.
-    dist_to_nplus
-        Array of distance to nplus contact for each step (can be `None`, in which case no surface effects are included.)
-    template
-        array of the bulk pulse template
-    times
-        time-stamps for the bulk pulse template
     """
     A = np.zeros(len(dt))
     maximum_t = np.zeros(len(dt))
     energy = np.zeros(len(dt))
 
     time_step = 1
-    n = len(template)
+    n = len(times)
     start = times[0]
 
     if include_surface_effects:
@@ -707,6 +709,9 @@ def _estimate_current_impl(
     for i in range(len(dt)):
         t = np.asarray(dt[i])
         e = np.asarray(edep[i])
+        r_tmp = np.asarray(r[i])
+        z_tmp = np.asarray(z[i])
+
         dist = np.asarray(dist_to_nplus[i])
 
         # get the expected maximum
@@ -720,7 +725,6 @@ def _estimate_current_impl(
             for j, d in enumerate(dist):
                 dtmp = int(d / surface_step_in_um)
 
-                # Use branchless selection
                 use_offset = dtmp <= ncols
                 offset_val = offsets[dtmp] if use_offset else 0.0
                 time_tmp = t[j] + offset_val * use_offset
@@ -737,9 +741,12 @@ def _estimate_current_impl(
                 t,
                 e,
                 dist,
-                template,
-                templates_surface,
-                activeness_surface,
+                r=r_tmp,
+                z=z_tmp,
+                template=template,
+                pulse_shape_library=pulse_shape_library,
+                templates_surface=templates_surface,
+                activeness_surface=activeness_surface,
                 tmin=tmin,
                 tmax=tmax,
                 start=start,
@@ -748,17 +755,73 @@ def _estimate_current_impl(
                 time_step=time_step,
                 surface_step_in_um=surface_step_in_um,
                 include_surface_effects=include_surface_effects,
+                use_library=use_library,
             )
 
     return A, maximum_t, energy
 
 
+def prepare_surface_inputs(
+    dist_to_nplus: ak.Array,
+    edep: ak.Array,
+    templates_surface: ArrayLike,
+    activeness_surface,
+    template: ArrayLike,
+) -> tuple:
+    """Prepare the inputs needed for surface sims."""
+    include_surface_effects = False
+
+    # prepare surface templates
+    if templates_surface is not None:
+        if dist_to_nplus is None:
+            msg = "Surface effects requested but distance not provided"
+            raise ValueError(msg)
+
+        include_surface_effects = True
+    else:
+        # convert types to keep numba happy
+        templates_surface = np.zeros((1, len(template)))
+        dist_to_nplus = ak.full_like(edep, np.nan)
+
+    # convert types for numba
+    if activeness_surface is None:
+        activeness_surface = np.zeros(len(template))
+
+    return include_surface_effects, dist_to_nplus, templates_surface, activeness_surface
+
+
+def prepare_pulse_shape_library(
+    template: ArrayLike | HPGePulseShapeLibrary,
+    times: ArrayLike,
+    edep: ak.Array,
+    r: ak.Array,
+    z: ak.Array,
+):
+    """Prepare the inputs for the full pulse shape library."""
+    use_library = False
+    if isinstance(template, HPGePulseShapeLibrary):
+        # convert to a form we can use
+        times = template.t
+        pulse_shape_library = (template.waveforms, template.r, template.z)
+        template = np.zeros_like(template.waveforms[0][0])
+        use_library = True
+
+    else:
+        pulse_shape_library = (np.zeros((1, 1, len(template))), np.zeros(1), np.zeros(1))
+        r = ak.full_like(edep, np.nan)
+        z = ak.full_like(edep, np.nan)
+
+    return use_library, pulse_shape_library, template, times, r, z
+
+
 def maximum_current(
     edep: ArrayLike,
     drift_time: ArrayLike,
     dist_to_nplus: ArrayLike | None = None,
+    r: ArrayLike | None = None,
+    z: ArrayLike | None = None,
     *,
-    template: np.array,
+    template: np.array | HPGePulseShapeLibrary,
     times: np.array,
     fccd_in_um: float = 0,
     templates_surface: ArrayLike | None = None,
@@ -776,16 +839,22 @@ def maximum_current(
         Array of drift times for each step.
     dist_to_nplus
         Distance to n-plus electrode, only needed if surface heuristics are enabled.
+    r
+        Radial coordinate (only needed if a full PSS library is used)
+    z
+        z coordinate (only needed if a full PSS library is used).
     template
-        array of the bulk pulse template
+        array of the bulk pulse template, can also be a :class:`HPGePulseShapeLibrary`.
     times
         time-stamps for the bulk pulse template
-    fccd
+    fccd_in_um
         Value of the full-charge-collection depth, if `None` no surface corrections are performed.
-    surface_library
+    templates_surface
         2D array (distance, time) of the rate of charge arriving at the p-n junction. Each row
         should be an array of length 10000 giving the charge arriving at the p-n junction for each timestep
         (in ns). This is produced by :func:`.hpge.surface.get_surface_response` or other libraries.
+    activeness_surface
+        An array of the activeness at each surface point.
     surface_step_in_um
         Distance step for the surface library.
     return_mode
@@ -796,40 +865,35 @@ 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)
 
-    include_surface_effects = False
-
-    if templates_surface is not None:
-        if dist_to_nplus is None:
-            msg = "Surface effects requested but distance not provided"
-            raise ValueError(msg)
-
-        include_surface_effects = True
-    else:
-        # convert types to keep numba happy
-        templates_surface = np.zeros((1, len(template)))
-        dist_to_nplus = ak.full_like(edep, np.nan)
-
-    # convert types for numba
-    if activeness_surface is None:
-        activeness_surface = np.zeros(len(template))
+    # prepare inputs for surface sims
+    include_surface_effects, dist_to_nplus, templates_surface, activeness_surface = (
+        prepare_surface_inputs(dist_to_nplus, edep, templates_surface, activeness_surface, template)
+    )
 
-    if not ak.all(ak.num(edep, axis=-1) == ak.num(drift_time, axis=-1)):
-        msg = "edep and drift time must have the same shape"
-        raise ValueError(msg)
+    # and for the full PS library
+    use_library, pulse_shape_library, template, times, r, z = prepare_pulse_shape_library(
+        template, times, edep, r, z
+    )
 
+    # and now compute the current
     curr, time, energy = _estimate_current_impl(
         ak.values_astype(ak.Array(edep), np.float64),
         ak.values_astype(ak.Array(drift_time), np.float64),
         ak.values_astype(ak.Array(dist_to_nplus), np.float64),
+        r=ak.values_astype(ak.Array(r), np.float64),
+        z=ak.values_astype(ak.Array(z), np.float64),
         template=template,
+        pulse_shape_library=pulse_shape_library,
         times=times,
         fccd=fccd_in_um,
         include_surface_effects=include_surface_effects,
+        use_library=use_library,
         templates_surface=templates_surface,
         activeness_surface=activeness_surface,
         surface_step_in_um=surface_step_in_um,

reboost-0.9.0/src/reboost/hpge/utils.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import NamedTuple
+
+import lgdo
+import numpy as np
+import pint
+from dbetto import AttrsDict
+from lgdo import lh5
+from scipy.interpolate import RegularGridInterpolator
+
+
+class HPGePulseShapeLibrary(NamedTuple):
+    """A set of templates defined in the cylindrical-like (r, z) HPGe plane."""
+
+    waveforms: np.array
+    "Field, function of the coordinates (r, z)."
+    r_units: pint.Unit
+    "Physical units of the coordinate `r`."
+    z_units: pint.Unit
+    "Physical units of the coordinate `z`."
+    t_units: pint.Unit
+    "Physical units of the times."
+    r: np.array
+    "One dimensional arrays specifying the radial coordinates"
+    z: np.array
+    "One dimensional arrays specifying the z coordinates"
+    t: np.array
+    "Times used to define the waveforms"
+
+
+def get_hpge_pulse_shape_library(
+    filename: str, obj: str, field: str, out_of_bounds_val: int | float = np.nan
+) -> HPGePulseShapeLibrary:
+    """Create the pulse shape library, holding simulated waveforms.
+
+    Reads from disk the following data structure: ::
+
+        FILENAME/
+        └── OBJ · struct{r,z,dt,t0,FIELD}
+            ├── r · array<1>{real} ── {'units': 'UNITS'}
+            ├── z · array<1>{real} ── {'units': 'UNITS'}
+            ├── dt · real ── {'units': 'UNITS'}
+            ├── t0 · real ── {'units': 'UNITS'}
+            └── FIELD · array<3>{real} ── {'units': 'UNITS'}
+
+    The conventions follow those used for :func:`get_hpge_rz_field`.
+    For the FIELD the first and second dimensions are `r` and `z`, respectively, with the last
+    dimension representing the waveform. dt and t0 define the timestamps for the waveforms.
+
+
+    Parameters
+    ----------
+    filename
+        name of the LH5 file containing the gridded scalar field.
+    obj
+        name of the HDF5 dataset where the data is saved.
+    field
+        name of the HDF5 dataset holding the waveforms.
+    out_of_bounds_val
+        value to use to replace NaNs in the field values.
+    """
+    data = lh5.read(obj, filename)
+
+    if not isinstance(data, lgdo.Struct):
+        msg = f"{obj} in {filename} is not an LGDO Struct"
+        raise ValueError(msg)
+
+    t0 = data["t0"].value
+    dt = data["dt"].value
+
+    t0_u = data["t0"].attrs["units"]
+    dt_u = data["dt"].attrs["units"]
+
+    if t0_u != dt_u:
+        msg = "t0 and dt must have the same units"
+        raise ValueError(msg)
+
+    tu = t0_u
+
+    data = AttrsDict(
+        {
+            k: np.nan_to_num(data[k].view_as("np", with_units=True), nan=out_of_bounds_val)
+            for k in ("r", "z", field)
+        }
+    )
+
+    times = t0 + np.arange(np.shape(data[field].m)[2]) * dt
+
+    return HPGePulseShapeLibrary(data[field].m, data.r.u, data.z.u, tu, data.r.m, data.z.m, times)
+
+
+class HPGeRZField(NamedTuple):
+    """A field defined in the cylindrical-like (r, z) HPGe plane."""
+
+    φ: Callable
+    "Field, function of the coordinates (r, z)."
+    r_units: pint.Unit
+    "Physical units of the coordinate `r`."
+    z_units: pint.Unit
+    "Physical units of the coordinate `z`."
+    φ_units: pint.Unit
+    "Physical units of the field."
+    ndim: int
+    "Number of dimensions for the field"
+
+
+def get_hpge_rz_field(
+    filename: str, obj: str, field: str, out_of_bounds_val: int | float = np.nan, **kwargs
+) -> HPGeRZField:
+    """Create an interpolator for a gridded HPGe field defined on `(r, z)`.
+
+    Reads from disk the following data structure: ::
+
+        FILENAME/
+        └── OBJ · struct{r,z,FIELD}
+            ├── r · array<1>{real} ── {'units': 'UNITS'}
+            ├── z · array<1>{real} ── {'units': 'UNITS'}
+            └── FIELD · array<N+2>{real} ── {'units': 'UNITS'}
+
+    where ``FILENAME``, ``OBJ`` and ``FIELD`` are provided as
+    arguments to this function. `obj` is a :class:`~lgdo.types.struct.Struct`,
+    `r` and `z` are one dimensional arrays specifying the radial and z
+    coordinates of the rectangular grid — not the coordinates of each single
+    grid point. In this coordinate system, the center of the p+ contact surface
+    is at `(0, 0)`, with the p+ contact facing downwards. `field` is instead a
+    ndim plus two-dimensional array specifying the field value at each grid point. The
+    first and second dimensions are `r` and `z`, respectively, with the latter dimensions
+    representing the dimensions of the output field.
+
+    NaN values are interpreted as points outside the detector profile in the `(r, z)` plane.
+
+    Before returning a :class:`HPGeScalarRZField`, the gridded field is fed to
+    :class:`scipy.interpolate.RegularGridInterpolator`.
+
+    Parameters
+    ----------
+    filename
+        name of the LH5 file containing the gridded scalar field.
+    obj
+        name of the HDF5 dataset where the data is saved.
+    field
+        name of the HDF5 dataset holding the field values.
+    out_of_bounds_val
+        value to use to replace NaNs in the field values.
+    """
+    data = lh5.read(obj, filename)
+
+    if not isinstance(data, lgdo.Struct):
+        msg = f"{obj} in {filename} is not an LGDO Struct"
+        raise ValueError(msg)
+
+    data = AttrsDict(
+        {
+            k: np.nan_to_num(data[k].view_as("np", with_units=True), nan=out_of_bounds_val)
+            for k in ("r", "z", field)
+        }
+    )
+    ndim = data[field].m.ndim - 2
+    interpolator = RegularGridInterpolator(
+        (data.r.m, data.z.m), data[field].m, **(kwargs | {"fill_value": out_of_bounds_val})
+    )
+
+    return HPGeRZField(interpolator, data.r.u, data.z.u, data[field].u, ndim)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reboost
-Version: 0.8.5
+Version: 0.9.0
 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.8.5 → reboost-0.9.0}/tests/conftest.py

@@ -7,8 +7,12 @@ from pathlib import Path
 from tempfile import gettempdir
 
 import numba
+import numpy as np
 import pytest
 from legendtestdata import LegendTestData
+from lgdo import Array, Scalar, Struct, lh5
+
+from reboost.hpge import psd
 
 _tmptestdir = Path(gettempdir()) / f"reboost-tests-{getuser()}-{uuid.uuid4()!s}"
 
@@ -49,4 +53,44 @@ def patch_numba_for_tests():
     numba.njit = njit_patched
 
 
+@pytest.fixture(scope="module")
+def test_pulse_shape_library(tmptestdir):
+    model, _ = psd.get_current_template(
+        -1000,
+        3000,
+        1.0,
+        amax=1,
+        mean_aoe=1,
+        mu=0,
+        sigma=100,
+        tau=100,
+        tail_fraction=0.65,
+        high_tail_fraction=0.1,
+        high_tau=10,
+    )
+
+    # loop
+    r = z = np.linspace(0, 100, 200)
+    waveforms = np.zeros((200, 200, 4001))
+    for i in range(200):
+        for j in range(200):
+            waveforms[i, j] = model
+
+    t0 = -1000
+    dt = 1
+
+    res = Struct(
+        {
+            "r": Array(r, attrs={"units": "mm"}),
+            "z": Array(z, attrs={"units": "mm"}),
+            "waveforms": Array(waveforms, attrs={"units": ""}),
+            "dt": Scalar(dt, attrs={"units": "ns"}),
+            "t0": Scalar(t0, attrs={"units": "ns"}),
+        }
+    )
+    lh5.write(res, "V01", f"{tmptestdir}/pulse_shape_lib.lh5")
+
+    return f"{tmptestdir}/pulse_shape_lib.lh5"
+
+
 patch_numba_for_tests()

{reboost-0.8.5 → reboost-0.9.0}/tests/hpge/test_current.py

@@ -6,6 +6,8 @@ import pytest
 from lgdo import Array, VectorOfVectors
 
 from reboost.hpge import psd, surface
+from reboost.hpge.psd import _get_template_idx
+from reboost.hpge.utils import get_hpge_pulse_shape_library
 from reboost.shape import cluster
 
 
@@ -46,6 +48,14 @@ def test_model():
     return model, x
 
 
+def test_get_template(test_pulse_shape_library):
+    lib = get_hpge_pulse_shape_library(test_pulse_shape_library, "V01", "waveforms")
+
+    ri, zi = _get_template_idx(10, 10, lib.r, lib.z)
+
+    assert len(lib.waveforms[ri][zi]) == 4001
+
+
 def test_maximum_current(test_model):
     model, x = test_model
 
@@ -183,3 +193,28 @@ def test_maximum_current_surface(test_model):
 
         # surface effects reduce the current
         assert np.all(curr_surf < curr_bulk)
+
+
+def test_maximum_current_library(test_pulse_shape_library):
+    lib = get_hpge_pulse_shape_library(test_pulse_shape_library, "V01", "waveforms")
+
+    model = lib.waveforms[0][0]
+    x = lib.t
+
+    edep = VectorOfVectors(
+        ak.Array([[100.0, 300.0, 50.0], [10.0, 0.0, 100.0], [500.0]]), attrs={"unit": "keV"}
+    )
+    times = VectorOfVectors(
+        ak.Array([[400, 500, 700], [800, 0, 1500], [700]], attrs={"unit": "ns"})
+    )
+    r = VectorOfVectors(
+        ak.Array([[20.0, 10.0, 5.0], [10.0, 1.0, 0.0], [70.0]]), attrs={"unit": "mm"}
+    )
+    z = VectorOfVectors(
+        ak.Array([[40.0, 2.0, 25.0], [22.0, 4.0, 1.2], [20.0]]), attrs={"unit": "mm"}
+    )
+
+    curr = psd.maximum_current(edep, times, template=model, times=x).view_as("ak")
+    curr2 = psd.maximum_current(edep, times, r=r, z=z, template=lib, times=x).view_as("ak")
+
+    assert ak.all(curr == curr2)

{reboost-0.8.5 → reboost-0.9.0}/tests/hpge/test_dt_heuristic.py

@@ -80,7 +80,7 @@ gamma_stp = Table(
 
 @pytest.fixture(scope="module")
 def dt_map(legendtestdata):
-    return utils.get_hpge_scalar_rz_field(
+    return utils.get_hpge_rz_field(
         legendtestdata["lh5/hpge-drift-time-maps.lh5"],
         "V99000A",
         "drift_time",
@@ -114,12 +114,7 @@ def dt_map_dummy(legendtestdata):
         drift_time,
     )
 
-    return utils.HPGeScalarRZField(
-        interpolator,
-        data.r.u,
-        data.z.u,
-        u.us,
-    )
+    return utils.HPGeRZField(interpolator, data.r.u, data.z.u, u.us, 1)
 
 
 def test_drift_time_dummy(dt_map_dummy):

{reboost-0.8.5 → reboost-0.9.0}/tests/hpge/test_hpge_map.py

@@ -1,21 +1,27 @@
 from __future__ import annotations
 
+import numpy as np
 import pytest
 from scipy.interpolate import RegularGridInterpolator
 
-from reboost.hpge.utils import HPGeScalarRZField, get_hpge_scalar_rz_field
+from reboost.hpge.utils import (
+    HPGePulseShapeLibrary,
+    HPGeRZField,
+    get_hpge_pulse_shape_library,
+    get_hpge_rz_field,
+)
 from reboost.units import ureg as u
 
 
 def test_read_hpge_map(legendtestdata):
-    dt_map = get_hpge_scalar_rz_field(
+    dt_map = get_hpge_rz_field(
         legendtestdata["lh5/hpge-drift-time-maps.lh5"],
         "V99000A",
         "drift_time",
         out_of_bounds_val=0,
     )
 
-    assert isinstance(dt_map, HPGeScalarRZField)
+    assert isinstance(dt_map, HPGeRZField)
 
     assert dt_map.r_units == u.m
     assert dt_map.z_units == u.m
@@ -28,3 +34,11 @@ def test_read_hpge_map(legendtestdata):
 
     assert dt_map.φ((0, 0)) == 0
     assert dt_map.φ([(0, 0.01), (0.03, 0.03)]) == pytest.approx([135, 695])
+
+
+def test_read_pulse_shape_library(test_pulse_shape_library):
+    # check th reading works
+    lib = get_hpge_pulse_shape_library(test_pulse_shape_library, "V01", "waveforms")
+    assert isinstance(lib, HPGePulseShapeLibrary)
+
+    assert np.shape(lib.waveforms) == (200, 200, 4001)

reboost-0.8.5/src/reboost/hpge/utils.py

@@ -1,79 +0,0 @@
-from __future__ import annotations
-
-from collections.abc import Callable
-from typing import NamedTuple
-
-import lgdo
-import numpy as np
-import pint
-from dbetto import AttrsDict
-from lgdo import lh5
-from scipy.interpolate import RegularGridInterpolator
-
-
-class HPGeScalarRZField(NamedTuple):
-    """A scalar field defined in the cylindrical-like (r, z) HPGe plane."""
-
-    φ: Callable
-    "Scalar field, function of the coordinates (r, z)."
-    r_units: pint.Unit
-    "Physical units of the coordinate `r`."
-    z_units: pint.Unit
-    "Physical units of the coordinate `z`."
-    φ_units: pint.Unit
-    "Physical units of the field."
-
-
-def get_hpge_scalar_rz_field(
-    filename: str, obj: str, field: str, out_of_bounds_val: int | float = np.nan, **kwargs
-) -> HPGeScalarRZField:
-    """Create an interpolator for a gridded scalar HPGe field defined on `(r, z)`.
-
-    Reads from disk the following data structure: ::
-
-        FILENAME/
-        └── OBJ · struct{r,z,FIELD}
-            ├── r · array<1>{real} ── {'units': 'UNITS'}
-            ├── z · array<1>{real} ── {'units': 'UNITS'}
-            └── FIELD · array<2>{real} ── {'units': 'UNITS'}
-
-    where ``FILENAME``, ``OBJ`` and ``FIELD`` are provided as
-    arguments to this function. `obj` is a :class:`~lgdo.types.struct.Struct`,
-    `r` and `z` are one dimensional arrays specifying the radial and z
-    coordinates of the rectangular grid — not the coordinates of each single
-    grid point. In this coordinate system, the center of the p+ contact surface
-    is at `(0, 0)`, with the p+ contact facing downwards. `field` is instead a
-    two-dimensional array specifying the field value at each grid point. The
-    first and second dimensions are `r` and `z`, respectively. NaN values are
-    interpreted as points outside the detector profile in the `(r, z)` plane.
-
-    Before returning a :class:`HPGeScalarRZField`, the gridded field is fed to
-    :class:`scipy.interpolate.RegularGridInterpolator`.
-
-    Parameters
-    ----------
-    filename
-        name of the LH5 file containing the gridded scalar field.
-    obj
-        name of the HDF5 dataset where the data is saved.
-    field
-        name of the HDF5 dataset holding the field values.
-    out_of_bounds_val
-        value to use to replace NaNs in the field values.
-    """
-    data = lh5.read(obj, filename)
-
-    if not isinstance(data, lgdo.Struct):
-        msg = f"{obj} in {filename} is not an LGDO Struct"
-        raise ValueError(msg)
-
-    data = AttrsDict(
-        {
-            k: np.nan_to_num(data[k].view_as("np", with_units=True), nan=out_of_bounds_val)
-            for k in ("r", "z", field)
-        }
-    )
-
-    interpolator = RegularGridInterpolator((data.r.m, data.z.m), data[field].m, **kwargs)
-
-    return HPGeScalarRZField(interpolator, data.r.u, data.z.u, data[field].u)