reboost 0.8.4__py3-none-any.whl → 0.9.0__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.8.4'
-__version_tuple__ = version_tuple = (0, 8, 4)
+__version__ = version = '0.9.0'
+__version_tuple__ = version_tuple = (0, 9, 0)
 
 __commit_id__ = commit_id = None

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/hpge/utils.py

@@ -11,23 +11,105 @@ from lgdo import lh5
 from scipy.interpolate import RegularGridInterpolator
 
 
-class HPGeScalarRZField(NamedTuple):
-    """A scalar field defined in the cylindrical-like (r, z) HPGe plane."""
+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
-    "Scalar field, function of the coordinates (r, z)."
+    "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_scalar_rz_field(
+def get_hpge_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)`.
+) -> HPGeRZField:
+    """Create an interpolator for a gridded HPGe field defined on `(r, z)`.
 
     Reads from disk the following data structure: ::
 
@@ -35,7 +117,7 @@ def get_hpge_scalar_rz_field(
         └── OBJ · struct{r,z,FIELD}
             ├── r · array<1>{real} ── {'units': 'UNITS'}
             ├── z · array<1>{real} ── {'units': 'UNITS'}
-            └── FIELD · array<2>{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`,
@@ -43,9 +125,11 @@ def get_hpge_scalar_rz_field(
     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.
+    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`.
@@ -73,7 +157,9 @@ def get_hpge_scalar_rz_field(
             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})
+    )
 
-    interpolator = RegularGridInterpolator((data.r.m, data.z.m), data[field].m, **kwargs)
-
-    return HPGeScalarRZField(interpolator, data.r.u, data.z.u, data[field].u)
+    return HPGeRZField(interpolator, data.r.u, data.z.u, data[field].u, ndim)

reboost/optmap/convolve.py

@@ -156,15 +156,70 @@ def iterate_stepwise_depositions_scintillate(
         raise ValueError(msg)
 
     rng = np.random.default_rng() if rng is None else rng
-    output_list = _iterate_stepwise_depositions_scintillate(edep_hits, rng, scint_mat_params, mode)
+    counts = ak.num(edep_hits.edep)
+    output_array = _iterate_stepwise_depositions_scintillate(
+        edep_hits, rng, scint_mat_params, mode, ak.sum(counts)
+    )
 
-    # convert the numba result back into an awkward array.
-    builder = ak.ArrayBuilder()
-    for r in output_list:
-        with builder.list():
-            builder.extend(r)
+    return ak.unflatten(output_array, counts)
 
-    return builder.snapshot()
+
+def iterate_stepwise_depositions_numdet(
+    edep_hits: ak.Array,
+    optmap: OptmapForConvolve,
+    det: str,
+    map_scaling: float = 1,
+    map_scaling_sigma: float = 0,
+    rng: np.random.Generator | None = None,
+):
+    if edep_hits.xloc.ndim == 1:
+        msg = "the pe processors only support already reshaped output"
+        raise ValueError(msg)
+
+    rng = np.random.default_rng() if rng is None else rng
+    counts = ak.num(edep_hits.num_scint_ph)
+    output_array, res = _iterate_stepwise_depositions_numdet(
+        edep_hits,
+        rng,
+        np.where(optmap.dets == det)[0][0],
+        map_scaling,
+        map_scaling_sigma,
+        optmap.edges,
+        optmap.weights,
+        ak.sum(counts),
+    )
+
+    if res["det_no_stats"] > 0:
+        log.warning(
+            "had edep out in voxels without stats: %d",
+            res["det_no_stats"],
+        )
+    if res["oob"] > 0:
+        log.warning(
+            "had edep out of map bounds: %d (%.2f%%)",
+            res["oob"],
+            (res["oob"] / (res["ib"] + res["oob"])) * 100,
+        )
+
+    return ak.unflatten(output_array, counts)
+
+
+def iterate_stepwise_depositions_times(
+    edep_hits: ak.Array,
+    scint_mat_params: sc.ComputedScintParams,
+    rng: np.random.Generator | None = None,
+):
+    if edep_hits.particle.ndim == 1:
+        msg = "the pe processors only support already reshaped output"
+        raise ValueError(msg)
+
+    rng = np.random.default_rng() if rng is None else rng
+    counts = ak.sum(edep_hits.num_det_ph, axis=1)
+    output_array = _iterate_stepwise_depositions_times(
+        edep_hits, rng, scint_mat_params, ak.sum(counts)
+    )
+
+    return ak.unflatten(output_array, counts)
 
 
 _pdg_func = numba_pdgid_funcs()
@@ -270,17 +325,15 @@ def _iterate_stepwise_depositions_pois(
 # - cache=True does not work with outer prange, i.e. loading the cached file fails (numba bug?)
 @njit(parallel=False, nogil=True, cache=True)
 def _iterate_stepwise_depositions_scintillate(
-    edep_hits, rng, scint_mat_params: sc.ComputedScintParams, mode: str
+    edep_hits, rng, scint_mat_params: sc.ComputedScintParams, mode: str, output_length: int
 ):
     pdgid_map = {}
-    output_list = []
+    output = np.empty(shape=output_length, dtype=np.int64)
 
+    output_index = 0
     for rowid in range(len(edep_hits)):  # iterate hits
         hit = edep_hits[rowid]
-        hit_output = []
-
-        # iterate steps inside the hit
-        for si in range(len(hit.particle)):
+        for si in range(len(hit.particle)):  # iterate steps inside the hit
             # get the particle information.
             particle = hit.particle[si]
             if particle not in pdgid_map:
@@ -295,12 +348,106 @@ def _iterate_stepwise_depositions_scintillate(
                 rng,
                 emission_term_model=("poisson" if mode == "no-fano" else "normal_fano"),
             )
-            hit_output.append(num_phot)
+            output[output_index] = num_phot
+            output_index += 1
 
-        assert len(hit_output) == len(hit.particle)
-        output_list.append(hit_output)
+    assert output_index == output_length
+    return output
+
+
+# - run with NUMBA_FULL_TRACEBACKS=1 NUMBA_BOUNDSCHECK=1 for testing/checking
+# - cache=True does not work with outer prange, i.e. loading the cached file fails (numba bug?)
+@njit(parallel=False, nogil=True, cache=True)
+def _iterate_stepwise_depositions_numdet(
+    edep_hits,
+    rng,
+    detidx: int,
+    map_scaling: float,
+    map_scaling_sigma: float,
+    optmap_edges,
+    optmap_weights,
+    output_length: int,
+):
+    oob = ib = det_no_stats = 0
+    output = np.empty(shape=output_length, dtype=np.int64)
+
+    output_index = 0
+    for rowid in range(len(edep_hits)):  # iterate hits
+        hit = edep_hits[rowid]
+
+        map_scaling_evt = map_scaling
+        if map_scaling_sigma > 0:
+            map_scaling_evt = rng.normal(loc=map_scaling, scale=map_scaling_sigma)
+
+        # iterate steps inside the hit
+        for si in range(len(hit.xloc)):
+            loc = np.array([hit.xloc[si], hit.yloc[si], hit.zloc[si]], dtype=np.float64)
+            # coordinates -> bins of the optical map.
+            bins = np.empty(3, dtype=np.int64)
+            for j in range(3):
+                edges = optmap_edges[j].astype(np.float64)
+                start = edges[0]
+                width = edges[1] - edges[0]
+                nbins = edges.shape[0] - 1
+                bins[j] = int((loc[j] - start) / width)
+
+                if bins[j] < 0 or bins[j] >= nbins:
+                    bins[j] = -1  # normalize all out-of-bounds bins just to one end.
+
+            if bins[0] == -1 or bins[1] == -1 or bins[2] == -1:
+                detp = 0.0  # out-of-bounds of optmap
+                oob += 1
+            else:
+                # get probabilities from map.
+                detp = optmap_weights[detidx, bins[0], bins[1], bins[2]] * map_scaling_evt
+                if detp < 0:
+                    det_no_stats += 1
+                ib += 1
+
+            pois_cnt = 0 if detp <= 0.0 else rng.poisson(lam=hit.num_scint_ph[si] * detp)
+            output[output_index] = pois_cnt
+            output_index += 1
+
+    assert output_index == output_length
+    return output, {"oob": oob, "ib": ib, "det_no_stats": det_no_stats}
+
+
+# - run with NUMBA_FULL_TRACEBACKS=1 NUMBA_BOUNDSCHECK=1 for testing/checking
+# - cache=True does not work with outer prange, i.e. loading the cached file fails (numba bug?)
+# - the output dictionary is not threadsafe, so parallel=True is not working with it.
+@njit(parallel=False, nogil=True, cache=True)
+def _iterate_stepwise_depositions_times(
+    edep_hits, rng, scint_mat_params: sc.ComputedScintParams, output_length: int
+):
+    pdgid_map = {}
+    output = np.empty(shape=output_length, dtype=np.float64)
+
+    output_index = 0
+    for rowid in range(len(edep_hits)):  # iterate hits
+        hit = edep_hits[rowid]
+
+        assert len(hit.particle) == len(hit.num_det_ph)
+        # iterate steps inside the hit
+        for si in range(len(hit.particle)):
+            pois_cnt = hit.num_det_ph[si]
+            if pois_cnt <= 0:
+                continue
+
+            # get the particle information.
+            particle = hit.particle[si]
+            if particle not in pdgid_map:
+                pdgid_map[particle] = (_pdgid_to_particle(particle), _pdg_func.charge(particle))
+            part, _charge = pdgid_map[particle]
+
+            # get time spectrum.
+            # note: we assume "immediate" propagation after scintillation.
+            scint_times = sc.scintillate_times(scint_mat_params, part, pois_cnt, rng) + hit.time[si]
+            assert len(scint_times) == pois_cnt
+            output[output_index : output_index + len(scint_times)] = scint_times
+            output_index += len(scint_times)
 
-    return output_list
+    assert output_index == output_length
+    return output
 
 
 def _get_scint_params(material: str):

reboost/spms/__init__.py

@@ -1,10 +1,19 @@
 from __future__ import annotations
 
-from .pe import detected_photoelectrons, emitted_scintillation_photons, load_optmap, load_optmap_all
+from .pe import (
+    detected_photoelectrons,
+    emitted_scintillation_photons,
+    load_optmap,
+    load_optmap_all,
+    number_of_detected_photoelectrons,
+    photoelectron_times,
+)
 
 __all__ = [
     "detected_photoelectrons",
     "emitted_scintillation_photons",
     "load_optmap",
     "load_optmap_all",
+    "number_of_detected_photoelectrons",
+    "photoelectron_times",
 ]

reboost/spms/pe.py

@@ -109,7 +109,12 @@ def detected_photoelectrons(
     map_scaling: float = 1,
     map_scaling_sigma: float = 0,
 ) -> VectorOfVectors:
-    """Derive the number of detected photoelectrons (p.e.) from scintillator hits using an optical map.
+    """Derive the number and arrival times of detected photoelectrons (p.e.) from scintillator hits using an optical map.
+
+    .. deprecated :: 0.8.5
+        Use the other, more fine-granular and split processors
+        :func:`number_of_detected_photoelectrons` and :func:`photoelectron_times` to
+        replace this legacy processor.
 
     Parameters
     ----------
@@ -181,3 +186,77 @@ def emitted_scintillation_photons(
     scint_mat_params = convolve._get_scint_params(material)
     ph = convolve.iterate_stepwise_depositions_scintillate(hits, scint_mat_params)
     return VectorOfVectors(ph)
+
+
+def number_of_detected_photoelectrons(
+    xloc: ak.Array,
+    yloc: ak.Array,
+    zloc: ak.Array,
+    num_scint_ph: ak.Array,
+    optmap: convolve.OptmapForConvolve,
+    spm_detector: str,
+    map_scaling: float = 1,
+    map_scaling_sigma: float = 0,
+) -> VectorOfVectors:
+    """Derive the number of detected photoelectrons (p.e.) from scintillator hits using an optical map.
+
+    Parameters
+    ----------
+    xloc
+        array of x coordinate position of scintillation events.
+    yloc
+        array of y coordinate position of scintillation events.
+    zloc
+        array of z coordinate position of scintillation events.
+    num_scint_ph
+        array of emitted scintillation photons, as generated by
+        :func:`emitted_scintillation_photons`.
+    detp
+    """
+    hits = ak.Array(
+        {
+            "xloc": units_conv_ak(xloc, "m"),
+            "yloc": units_conv_ak(yloc, "m"),
+            "zloc": units_conv_ak(zloc, "m"),
+            "num_scint_ph": units_conv_ak(num_scint_ph, "dimensionless"),
+        }
+    )
+
+    ph = convolve.iterate_stepwise_depositions_numdet(
+        hits, optmap, spm_detector, map_scaling, map_scaling_sigma
+    )
+    return VectorOfVectors(ph)
+
+
+def photoelectron_times(
+    num_det_ph: ak.Array,
+    particle: ak.Array,
+    time: ak.Array,
+    material: str,
+) -> VectorOfVectors:
+    """Derive the arrival times of scintillation photons.
+
+    Parameters
+    ----------
+    num_det_ph
+        array of detected scintillation photons, as generated by
+        :func:`emitted_scintillation_photons`.
+    particle
+        array of particle PDG IDs of scintillation events.
+    time
+        array of timestamps of scintillation events.
+    material
+        scintillating material name.
+    """
+    hits = ak.Array(
+        {
+            "num_det_ph": units_conv_ak(num_det_ph, "dimensionless"),
+            "particle": units_conv_ak(particle, "dimensionless"),
+            "time": units_conv_ak(time, "ns"),
+        }
+    )
+
+    scint_mat_params = convolve._get_scint_params(material)
+    pe = convolve.iterate_stepwise_depositions_times(hits, scint_mat_params)
+
+    return VectorOfVectors(pe, attrs={"units": "ns"})

reboost/utils.py

@@ -445,7 +445,7 @@ def write_lh5(
         time_dict.update_field("write", start_time)
 
 
-def get_remage_detector_uids(h5file: str | Path) -> dict:
+def get_remage_detector_uids(h5file: str | Path, *, lh5_table: str = "stp") -> dict:
     """Get mapping of detector names to UIDs from a remage output file.
 
     The remage LH5 output files contain a link structure that lets the user
@@ -483,7 +483,7 @@ def get_remage_detector_uids(h5file: str | Path) -> dict:
 
     out = {}
     with h5py.File(h5file, "r") as f:
-        g = f["/stp/__by_uid__"]
+        g = f[f"/{lh5_table}/__by_uid__"]
         # loop over links
         for key in g:
             # is this a link?

{reboost-0.8.4.dist-info → reboost-0.9.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reboost
-Version: 0.8.4
+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.4.dist-info → reboost-0.9.0.dist-info}/RECORD

@@ -1,6 +1,6 @@
 reboost/__init__.py,sha256=VZz9uo7i2jgAx8Zi15SptLZnE_qcnGuNWwqkD3rYHFA,278
 reboost/__main__.py,sha256=42koSxY2st4mMIRSAnKz06nP5HppMPxBVFf2jaHljGs,95
-reboost/_version.py,sha256=6jY56wQpNbuVX1aDDCpFgsBq74u7lloIJcrJSefeVec,704
+reboost/_version.py,sha256=TvxBYkx8Rz_Q1S3JFp831BRT8Wo0Yxt6TJMtgZKenTo,704
 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
@@ -10,21 +10,21 @@ 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/utils.py,sha256=-4315U6m1M-rwoaI_inI1wIo_l20kvoGmC84D_QOhkE,14563
+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=qgX1KWVApd3GOIkRzUwJ5MNMIlzVavv35MGZSWTUZHo,24565
+reboost/hpge/psd.py,sha256=P7dUJQPvxW6vndJ79r0j7ANZvSuV_IuauERhvWD74j0,26989
 reboost/hpge/surface.py,sha256=feH-kxRRp3HkikRRJ-LCu6zvVONEOYVd3THx12emGTM,8494
-reboost/hpge/utils.py,sha256=AcMS86v8s0SbeC4YA3CfH73GBZcC7Sb_RR3pRPjWAwU,2857
+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/optmap/__init__.py,sha256=imvuyld-GLw8qdwqW-lXCg2feptcTyQo3wIzPvDHwmY,93
 reboost/optmap/__main__.py,sha256=DfzkXQ7labOe53hd7jH5pAbTW491jjQYSMLyl72L4Rk,111
 reboost/optmap/cli.py,sha256=_7WBlx55eRyW_wWB-ELbFaWXin2d3xsh6Q5bFoNJaHE,8694
-reboost/optmap/convolve.py,sha256=0UCOVy6BIu5OYoSWeAb6fuROqvGu9rVswly6XvnxYoE,10763
+reboost/optmap/convolve.py,sha256=e3vTBurYSM4UmicIdDif6_cI04pV8khHEG5n7M_DWNg,16079
 reboost/optmap/create.py,sha256=GmGd0-F0eWmw7ywH8pT1lKiMb60QXCq9al8Ka_ySD1Q,14382
 reboost/optmap/evt.py,sha256=nZcB3aOPtMhySu00J23KlOTijeI5Oazde-vfM5A71e8,5664
 reboost/optmap/mapview.py,sha256=fswwXolA6au8u8gljBKy8PSXC2W7Cy_GwOV86-duYG8,6880
@@ -34,11 +34,11 @@ 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/reduction.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-reboost/spms/__init__.py,sha256=24b9GelP4ngXc7JibpWqLOtUuIaQQa0DzDQYP8lVKUc,262
-reboost/spms/pe.py,sha256=cFj-D1-pj5djRHmrfYMjNscNnQfveACMZMJq4u4th3w,5830
-reboost-0.8.4.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-reboost-0.8.4.dist-info/METADATA,sha256=YdOpd1dR4rVjYwz6dkjtxvQ6qoJQmAMR6DXgvAvH1zM,3877
-reboost-0.8.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-reboost-0.8.4.dist-info/entry_points.txt,sha256=DxhD6BidSWNot9BrejHJjQ7RRLmrMaBIl52T75oWTwM,93
-reboost-0.8.4.dist-info/top_level.txt,sha256=q-IBsDepaY_AbzbRmQoW8EZrITXRVawVnNrB-_zyXZs,8
-reboost-0.8.4.dist-info/RECORD,,
+reboost/spms/__init__.py,sha256=8I6WT8i_kUPqEDnSD0aCf6A26cjKjQQZSNrvwZ3o-Ac,415
+reboost/spms/pe.py,sha256=LwqrK1HOZWzGcNZnntaqI6r4rnDww4KW9Mao4xLFbDE,8226
+reboost-0.9.0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+reboost-0.9.0.dist-info/METADATA,sha256=wrdKu5HVcTaoKIz6mp72uJ9ehRzyF0e1xCWUe63g7DM,3877
+reboost-0.9.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+reboost-0.9.0.dist-info/entry_points.txt,sha256=DxhD6BidSWNot9BrejHJjQ7RRLmrMaBIl52T75oWTwM,93
+reboost-0.9.0.dist-info/top_level.txt,sha256=q-IBsDepaY_AbzbRmQoW8EZrITXRVawVnNrB-_zyXZs,8
+reboost-0.9.0.dist-info/RECORD,,