optmap: split up convolving processors

docs/source/manual/optical.md

@@ -145,7 +145,10 @@ map-based production of the optical response is divided into multiple steps:
    arbitrary user-supplied scaling factor for the channel {math}`c`.
 
    This is implemented in the processor
-   {func}`reboost.spms.pe.detected_photoelectrons`.
+   {func}`reboost.spms.pe.number_of_detected_photoelectrons`.
+
+3. (optional) sampling of the photon arrival times. This is implemented in the
+   processor {func}`reboost.spms.pe.photoelectron_times`.
 
 The used statistical model are developed and described in more detail in M.
 Huber's master thesis.

src/reboost/optmap/convolve.py

@@ -155,14 +155,65 @@ def iterate_stepwise_depositions_scintillate(
         msg = "the pe processors only support already reshaped output"
         raise ValueError(msg)
 
+    builder = ak.ArrayBuilder()
     rng = np.random.default_rng() if rng is None else rng
-    output_list = _iterate_stepwise_depositions_scintillate(edep_hits, rng, scint_mat_params, mode)
+    _iterate_stepwise_depositions_scintillate(edep_hits, rng, scint_mat_params, mode, builder)
+
+    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)
 
-    # convert the numba result back into an awkward array.
     builder = ak.ArrayBuilder()
-    for r in output_list:
-        with builder.list():
-            builder.extend(r)
+    rng = np.random.default_rng() if rng is None else rng
+    res = _iterate_stepwise_depositions_numdet(
+        edep_hits,
+        rng,
+        np.where(optmap.dets == det)[0][0],
+        map_scaling,
+        map_scaling_sigma,
+        optmap.edges,
+        optmap.weights,
+        builder,
+    )
+
+    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 builder.snapshot()
+
+
+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)
+
+    builder = ak.ArrayBuilder()
+    rng = np.random.default_rng() if rng is None else rng
+    _iterate_stepwise_depositions_times(edep_hits, rng, scint_mat_params, builder)
 
     return builder.snapshot()
 
@@ -270,14 +321,13 @@ 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, builder
 ):
     pdgid_map = {}
-    output_list = []
 
     for rowid in range(len(edep_hits)):  # iterate hits
         hit = edep_hits[rowid]
-        hit_output = []
+        builder.begin_list()
 
         # iterate steps inside the hit
         for si in range(len(hit.particle)):
@@ -295,12 +345,101 @@ def _iterate_stepwise_depositions_scintillate(
                 rng,
                 emission_term_model=("poisson" if mode == "no-fano" else "normal_fano"),
             )
-            hit_output.append(num_phot)
+            builder.integer(num_phot)
 
-        assert len(hit_output) == len(hit.particle)
-        output_list.append(hit_output)
+        # assert len(hit_output) == len(hit.particle)
+        builder.end_list()
+
+
+# - 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,
+    builder,
+):
+    oob = ib = det_no_stats = 0
+
+    for rowid in range(len(edep_hits)):  # iterate hits
+        hit = edep_hits[rowid]
+        builder.begin_list()
+
+        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)
+            builder.integer(pois_cnt)
+
+        builder.end_list()
+
+    return {"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, builder
+):
+    pdgid_map = {}
+
+    for rowid in range(len(edep_hits)):  # iterate hits
+        hit = edep_hits[rowid]
+        builder.begin_list()
+
+        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]
+            for ti in range(len(scint_times)):
+                builder.real(scint_times[ti])
 
-    return output_list
+        builder.end_list()
 
 
 def _get_scint_params(material: str):

src/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",
 ]

src/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"})

tests/hit/configs/spms.yaml

@@ -15,6 +15,7 @@ processing_groups:
       - edep
       - pe_times
       - pe_times_lar
+      - pe_times_lar_full
 
     pre_operations:
       num_scint_ph:
@@ -25,7 +26,13 @@ processing_groups:
       optmap_lar: reboost.spms.load_optmap(ARGS.optmap_path, DETECTOR)
 
     operations:
+      num_det_ph: reboost.spms.number_of_detected_photoelectrons(
+        HITS.xloc, HITS.yloc, HITS.zloc, HITS.num_scint_ph,
+        DETECTOR_OBJECTS.optmap_lar, DETECTOR)
       pe_times_lar:
+        reboost.spms.photoelectron_times(HITS.num_det_ph, HITS.particle,
+        HITS.time, "lar")
+      pe_times_lar_full:
         reboost.spms.detected_photoelectrons( HITS.num_scint_ph, HITS.particle,
         HITS.time, HITS.xloc, HITS.yloc, HITS.zloc, DETECTOR_OBJECTS.optmap_lar,
         "lar", DETECTOR )

tests/hit/test_build_hit.py

@@ -6,6 +6,7 @@
 import awkward as ak
 import dbetto
 import h5py
+import numpy as np
 import pytest
 from lgdo import Array, Struct, Table, VectorOfVectors, lh5
 
@@ -350,3 +351,12 @@ def test_spms(test_gen_lh5_scint, tmptestdir):
     assert "S002" in hits
 
     assert isinstance(hits["S001"]["pe_times_lar"], VectorOfVectors)
+    assert isinstance(hits["S001"]["pe_times_lar_full"], VectorOfVectors)
+
+    pe_split = hits["S001"]["pe_times_lar"].view_as("ak")
+    pe_full = hits["S001"]["pe_times_lar_full"].view_as("ak")
+    pe_ratio = ak.count(pe_split) / ak.count(pe_full)
+
+    assert np.abs(pe_ratio - 1) < 0.1
+
+    # TODO: test they are statistically close