reboost 0.8.3__py3-none-any.whl

reboost/shape/cluster.py

@@ -0,0 +1,260 @@
+from __future__ import annotations
+
+import logging
+
+import awkward as ak
+import numba
+import numpy as np
+from lgdo import VectorOfVectors
+
+log = logging.getLogger(__name__)
+
+
+def apply_cluster(
+    cluster_run_lengths: VectorOfVectors | ak.Array, field: ak.Array | VectorOfVectors
+) -> VectorOfVectors:
+    """Apply clustering to a field.
+
+    Parameters
+    ----------
+    cluster_ids
+        run lengths of each cluster
+    field
+        the field to cluster
+    """
+    if isinstance(cluster_run_lengths, VectorOfVectors):
+        cluster_run_lengths = cluster_run_lengths.view_as("ak")
+
+    if isinstance(field, VectorOfVectors):
+        field = field.view_as("ak")
+
+    n_cluster = ak.num(cluster_run_lengths, axis=-1)
+    clusters = ak.unflatten(ak.flatten(field), ak.flatten(cluster_run_lengths))
+
+    # reshape into cluster oriented
+    return VectorOfVectors(ak.unflatten(clusters, n_cluster))
+
+
+def cluster_by_step_length(
+    trackid: ak.Array | VectorOfVectors,
+    pos_x: ak.Array | VectorOfVectors,
+    pos_y: ak.Array | VectorOfVectors,
+    pos_z: ak.Array | VectorOfVectors,
+    dist: ak.Array | VectorOfVectors | None = None,
+    surf_cut: float | None = None,
+    threshold: float = 0.1,
+    threshold_surf: float | None = None,
+) -> VectorOfVectors:
+    """Perform clustering based on the step length.
+
+    Steps are clustered based on distance, if either:
+     - a step is in a new track,
+     - a step moves from surface to bulk region (or visa versa),
+     - the distance between the current step and the first step of the current cluster is above a threshold.
+
+    Then a new cluster is started. The surface region is defined as the volume
+    less than surf_cut distance to the surface. This allows for a fine tuning of the
+    parameters to be different for bulk and surface.
+
+    Parameters
+    ----------
+    trackid
+        index of the track.
+    pos_x
+        x position of the step.
+    pos_y
+        y position of the step.
+    pos_z
+        z position of the step.
+    dist
+        distance to the detector surface. Can be `None` in which case all steps are treated as being in the "bulk".
+    surf_cut
+        Size of the surface region (in mm), if `None` no selection is applied (default).
+    threshold
+        Distance threshold in mm to combine steps in the bulk.
+    threshold_surf
+        Distance threshold in mm to combine steps in the surface.
+
+    Returns
+    -------
+    Array of the run lengths of each cluster within a hit.
+    """
+    # type conversions
+    if isinstance(pos_x, VectorOfVectors):
+        pos_x = pos_x.view_as("ak")
+
+    if isinstance(pos_y, VectorOfVectors):
+        pos_y = pos_y.view_as("ak")
+
+    if isinstance(pos_z, VectorOfVectors):
+        pos_z = pos_z.view_as("ak")
+
+    if isinstance(trackid, VectorOfVectors):
+        trackid = trackid.view_as("ak")
+
+    if isinstance(dist, VectorOfVectors):
+        dist = dist.view_as("ak")
+
+    pos = np.vstack(
+        [
+            ak.flatten(pos_x).to_numpy().astype(np.float64),
+            ak.flatten(pos_y).to_numpy().astype(np.float64),
+            ak.flatten(pos_z).to_numpy().astype(np.float64),
+        ]
+    ).T
+
+    indices_flat = cluster_by_distance_numba(
+        ak.flatten(ak.local_index(trackid)).to_numpy(),
+        ak.flatten(trackid).to_numpy(),
+        pos,
+        dist_to_surf=ak.flatten(dist).to_numpy() if dist is not None else dist,
+        surf_cut=surf_cut,
+        threshold=threshold,
+        threshold_surf=threshold_surf,
+    )
+
+    # reshape into being event oriented
+    indices = ak.unflatten(indices_flat, ak.num(ak.local_index(trackid)))
+
+    # number of steps per cluster
+    counts = ak.run_lengths(indices)
+
+    return VectorOfVectors(counts)
+
+
+@numba.njit
+def cluster_by_distance_numba(
+    local_index: np.ndarray,
+    trackid: np.ndarray,
+    pos: np.ndarray,
+    dist_to_surf: np.ndarray | None,
+    surf_cut: float | None = None,
+    threshold: float = 0.1,
+    threshold_surf: float | None = None,
+) -> np.ndarray:
+    """Cluster steps by the distance between points in the same track.
+
+    This function gives the basic numerical calculations for
+    :func:`cluster_by_step_length`.
+
+    Parameters
+    ----------
+    local_index
+        1D array of the local index within each hit (step group)
+    trackid
+        1D array of index of the track
+    pos
+        `(n,3)` size array of the positions
+    dist_to_surf
+        1D array of the distance to the detector surface. Can be `None` in which case all steps are treated as being in the bulk.
+    surf_cut
+        Size of the surface region (in mm), if `None` no selection is applied.
+    threshold
+        Distance threshold in mm to combine steps in the bulk.
+    threshold_surf
+        Distance threshold in mm to combine steps in the surface.
+
+    Returns
+    -------
+    np.ndarray
+        1D array of cluster indices
+    """
+
+    def _dist(a, b):
+        return np.sqrt(np.sum((a - b) ** 2))
+
+    n = len(local_index)
+    out = np.zeros((n,), dtype=numba.int32)
+
+    trackid_prev = -1
+    pos_prev = np.zeros(3, dtype=numba.float64)
+    cluster_idx = -1
+    is_surf_prev = False
+
+    for idx in range(n):
+        # consider a surface and a bulk region
+        if dist_to_surf is not None:
+            thr = threshold if dist_to_surf[idx] > surf_cut else threshold_surf
+
+            new_cluster = (
+                (trackid[idx] != trackid_prev)
+                or (is_surf_prev and (dist_to_surf[idx] > surf_cut))
+                or ((not is_surf_prev) and (dist_to_surf[idx] < surf_cut))
+                or (_dist(pos[idx, :], pos_prev) > thr)
+            )
+        # basic clustering without split into surface / bulk
+        else:
+            thr = threshold
+            new_cluster = (trackid[idx] != trackid_prev) or (_dist(pos[idx, :], pos_prev) > thr)
+
+        # New hit, reset cluster index
+        if idx == 0 or local_index[idx] == 0:
+            cluster_idx = 0
+            pos_prev = pos[idx]
+
+        # either new track, moving from surface to bulk,
+        # moving from bulk to surface, or stepping more than
+        # the threshold. Start a new cluster.
+        elif new_cluster:
+            cluster_idx += 1
+            pos_prev = pos[idx, :]
+
+        out[idx] = cluster_idx
+
+        # Update previous values
+        trackid_prev = trackid[idx]
+        if dist_to_surf is not None:
+            is_surf_prev = dist_to_surf[idx] < surf_cut
+
+    return out
+
+
+def step_lengths(
+    x_cluster: ak.Array | VectorOfVectors,
+    y_cluster: ak.Array | VectorOfVectors,
+    z_cluster: ak.Array | VectorOfVectors,
+) -> VectorOfVectors:
+    """Compute the distance between consecutive steps.
+
+    This is based on calculating the distance between consecutive steps in the same track,
+    thus the input arrays should already be clustered (have dimension 3). The output
+    will have a similar shape to the input with one less entry in the outermost dimension.
+
+    Example config (assuming that the clustered positions are obtained already):
+
+    .. code-block:: yaml
+
+        step_lengths: reboost.shape.cluster.step_lengths(HITS.cluster_x,HITS.cluster_y,HITS.cluster_z))
+
+    Parameters
+    ----------
+    x_cluster
+        The x location of each step in each cluster and event.
+    y_cluster
+        The y location of each step in each cluster and event.
+    z_cluster
+        The z location of each step in each cluster and event.
+
+    Returns
+    -------
+    a `VectorOfVectors` of the step lengths in each cluster.
+    """
+    data = [x_cluster, y_cluster, z_cluster]
+
+    for idx, var in enumerate(data):
+        if isinstance(var, VectorOfVectors):
+            data[idx] = var.view_as("ak")
+        # check shape
+        if data[idx].ndim != 3:
+            msg = f"The input array for step lengths must be 3 dimensional not {data[idx.dim]}"
+            raise ValueError(msg)
+
+    counts = ak.num(data[0], axis=-1)
+    data = np.vstack([ak.flatten(ak.flatten(var)).to_numpy() for var in data])
+    dist = np.append(np.sqrt(np.sum(np.diff(data, axis=1) ** 2, axis=0)), 0)
+
+    n_cluster = ak.num(counts, axis=-1)
+    clusters = ak.unflatten(ak.Array(dist), ak.flatten(counts))
+
+    out = ak.unflatten(clusters, n_cluster)
+    return VectorOfVectors(out[:, :, :-1])

reboost/shape/group.py

@@ -0,0 +1,189 @@
+from __future__ import annotations
+
+import logging
+
+import awkward as ak
+import numpy as np
+from dbetto import AttrsDict
+from lgdo import Table, VectorOfVectors
+from numpy.typing import ArrayLike
+
+log = logging.getLogger(__name__)
+
+
+def isin(channels: ak.Array, chan_list: list):
+    """Check if each element of the awkward array channels is in the channel list."""
+    num_channels = ak.num(channels, axis=-1)
+    channels_flat = ak.flatten(channels)
+    isin = np.isin(channels_flat, chan_list)
+
+    # unflatten
+    return ak.unflatten(isin, num_channels)
+
+
+def get_isin_group(
+    channels: ArrayLike, groups: AttrsDict, tcm_tables: dict, group: str = "off"
+) -> ak.Array:
+    """For each channel check if it is in the group.
+
+    Parameters
+    ----------
+    channels
+        Array of the channel indices.
+    groups
+        A mapping of the group for every channel name.
+    tcm_tables
+        the mapping of indices to table names
+    group
+        the group to select.
+
+    Returns
+    -------
+    an awkward array of the same shape of channels of booleans.
+    """
+    usability = {uid: groups[name] for name, uid in tcm_tables.items()}
+    group_idx = [key for key, item in usability.items() if item == group]
+
+    return isin(channels, group_idx)
+
+
+def _sort_data(obj: ak.Array, *, time_name: str = "time", evtid_name: str = "evtid") -> ak.Array:
+    """Sort the data by evtid then time.
+
+    Parameters
+    ----------
+    obj
+        array of records containing fields `time` and `evtid`.
+    time_name
+        name of the time field in `obj`.
+    evtid_name
+        name of the evtid field in `obj`.
+
+    Returns
+    -------
+    sorted awkward array
+    """
+    obj = obj[ak.argsort(obj[evtid_name])]
+    obj_unflat = ak.unflatten(obj, ak.run_lengths(obj[evtid_name]))
+
+    indices = ak.argsort(obj_unflat[time_name], axis=-1)
+    sorted_obj = obj_unflat[indices]
+
+    return ak.flatten(sorted_obj)
+
+
+def group_by_evtid(data: Table | ak.Array, *, evtid_name: str = "evtid") -> Table:
+    """Simple grouping by evtid.
+
+    Takes the input `stp` :class:`lgdo.Table` from remage and defines groupings of steps (i.e the
+    `cumulative_length` for a vector of vectors). This then defines the output table (also :class:`lgdo.Table`),
+    on which processors can add fields.
+
+    Parameters
+    ----------
+    data
+        LGDO Table which must contain the `evtid` field.
+    evtid_name
+        the name of the index field in the input table.
+
+    Returns
+    -------
+    LGDO table of :class:`VectorOfVector` for each field.
+
+    Note
+    ----
+    The input table must be sorted (by `evtid`).
+    """
+    # convert to awkward
+    obj_ak = data.view_as("ak") if isinstance(data, Table) else data
+
+    # extract cumulative lengths
+    counts = ak.run_lengths(obj_ak[evtid_name])
+    cumulative_length = np.cumsum(counts)
+
+    # convert to numpy
+    if isinstance(cumulative_length, ak.Array):
+        cumulative_length = cumulative_length.to_numpy()
+
+    # build output table
+    out_tbl = Table(size=len(cumulative_length))
+
+    for f in obj_ak.fields:
+        out_tbl.add_field(
+            f,
+            VectorOfVectors(
+                cumulative_length=cumulative_length, flattened_data=obj_ak[f].to_numpy()
+            ),
+        )
+    return out_tbl
+
+
+def group_by_time(
+    data: Table | ak.Array,
+    window: float = 10,
+    time_name: str = "time",
+    evtid_name: str = "evtid",
+    fields: list | None = None,
+) -> Table:
+    """Grouping of steps by `evtid` and `time`.
+
+    Takes the input `stp` :class:`lgdo.Table` from remage and defines groupings of steps (i.e the
+    `cumulative_length` for a vector of vectors). This then defines the output table (also :class:`lgdo.Table`),
+    on which processors can add fields.
+
+    The windowing is based on defining a new group when the `evtid` changes or when the time increases by `> window`,
+    which is in units of us.
+
+    Parameters
+    ----------
+    data
+        :class:`lgdo.Table` or `ak.Array` which must contain the time_name and evtid_name fields
+    window
+        time window in us used to search for coincident hits
+    time_name
+        name of the timing field
+    evtid_name
+        name of the evtid field
+    fields
+        names of fields to include in the output table, if None includes all
+
+    Returns
+    -------
+    LGDO table of :class:`VectorOfVector` for each field.
+
+    Note
+    ----
+    The input table must be sorted (first by `evtid` then `time`).
+    """
+    obj = data.view_as("ak") if isinstance(data, Table) else data
+    obj = _sort_data(obj, time_name=time_name, evtid_name=evtid_name)
+
+    # get difference
+    time_diffs = np.diff(obj[time_name])
+    index_diffs = np.diff(obj[evtid_name])
+
+    # index of the last element in each run
+    time_change = (time_diffs > window * 1000) & (index_diffs == 0)
+    index_change = index_diffs > 0
+
+    # cumulative length is just the index of changes plus 1
+    cumulative_length = np.array(np.where(time_change | index_change))[0] + 1
+
+    # add the las grouping
+    cumulative_length = np.append(cumulative_length, len(obj[time_name]))
+
+    # convert to numpy
+    if isinstance(cumulative_length, ak.Array):
+        cumulative_length = cumulative_length.to_numpy()
+
+    # build output table
+    out_tbl = Table(size=len(cumulative_length))
+
+    fields = obj.fields if fields is None else fields
+    for f in fields:
+        out_tbl.add_field(
+            f,
+            VectorOfVectors(cumulative_length=cumulative_length, flattened_data=obj[f].to_numpy()),
+        )
+
+    return out_tbl

reboost/spms/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from .pe import detected_photoelectrons, emitted_scintillation_photons, load_optmap
+
+__all__ = ["detected_photoelectrons", "emitted_scintillation_photons", "load_optmap"]

reboost/spms/pe.py

@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+import logging
+
+import awkward as ak
+import numpy as np
+from lgdo import VectorOfVectors
+
+from ..optmap import convolve
+from ..units import units_conv_ak
+
+log = logging.getLogger(__name__)
+
+
+def load_optmap_all(map_file: str) -> convolve.OptmapForConvolve:
+    """Load an optical map file for later use with :py:func:`detected_photoelectrons`."""
+    return convolve.open_optmap(map_file)
+
+
+def load_optmap(map_file: str, spm_det_uid: int) -> convolve.OptmapForConvolve:
+    """Load an optical map file for later use with :py:func:`detected_photoelectrons`."""
+    return convolve.open_optmap_single(map_file, spm_det_uid)
+
+
+def _nested_unflatten(data: ak.Array, lengths: ak.Array):
+    return ak.unflatten(ak.unflatten(ak.flatten(data), ak.flatten(lengths)), ak.num(lengths))
+
+
+def corrected_photoelectrons(
+    simulated_pe: ak.Array,
+    simulated_uids: ak.Array,
+    data_pe: ak.Array,
+    data_uids: ak.Array,
+    *,
+    seed: int | None = None,
+) -> tuple[ak.Array, ak.Array]:
+    r"""Add a correction to the observed number of photoelectrons (p.e.) using forced trigger data.
+
+    For every simulated event a corresponding forced trigger event in data is chosen
+    and the resulting number of p.e. for each channel (i) is:
+
+     .. math::
+
+        n_i = n_{\text{sim},i} + n_{\text{data},i}
+
+    .. warning::
+       The number of supplied forced trigger events in data should ideally be
+       more than that in the simulations. If this is not the case and "allow_data_reuse"
+       is True then some data events will be used multiple times. This introduces
+       a small amount of correlation between the simulated events, but is probably acceptable
+       in most circumstances.
+
+    Parameters
+    ----------
+    simulated_pe
+        The number of number of detected pe per sipm channel.
+    simulated_uids
+        The unique identifier (uid) for each sipm hit.
+    data_pe
+        The collection of forced trigger pe.
+    data_uids
+        The uids for each forced trigger event.
+    seed
+        Seed for random number generator
+
+    Returns
+    -------
+    a tuple of the corrected pe and sipm uids.
+    """
+    rand = np.random.default_rng(seed=seed)
+    rand_ints = rand.integers(0, len(data_pe), size=len(simulated_pe))
+
+    selected_data_pe = data_pe[rand_ints]
+    selected_data_uids = data_uids[rand_ints]
+
+    # combine sims with data
+    pe_tot = ak.concatenate([simulated_pe, selected_data_pe], axis=1)
+    uid_tot = ak.concatenate([simulated_uids, selected_data_uids], axis=1)
+
+    # sort by uid
+    order = ak.argsort(uid_tot)
+    pe_tot = pe_tot[order]
+    uid_tot = uid_tot[order]
+
+    # add an extra axis
+    n = ak.run_lengths(uid_tot)
+
+    # add another dimension
+    pe_tot = _nested_unflatten(pe_tot, n)
+    uid_tot = _nested_unflatten(uid_tot, n)
+
+    # sum pe and take the first uid (should all be the same)
+    corrected_pe = ak.sum(pe_tot, axis=-1)
+    uid_tot = ak.fill_none(ak.firsts(uid_tot, axis=-1), np.nan)
+
+    return corrected_pe, uid_tot
+
+
+def detected_photoelectrons(
+    num_scint_ph: ak.Array,
+    particle: ak.Array,
+    time: ak.Array,
+    xloc: ak.Array,
+    yloc: ak.Array,
+    zloc: ak.Array,
+    optmap: convolve.OptmapForConvolve,
+    material: str,
+    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
+    ----------
+    num_scint_ph
+        array of emitted 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.
+    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.
+    optmap
+        the optical map loaded via py:func:`load_optmap`.
+    material
+        scintillating material name.
+    spm_detector
+        SiPM detector name as used in the optical map.
+    map_scaling
+        scale the detection probability in the map for this detector by this factor.
+    map_scaling_sigma
+        if larger than zero, sample the used scaling factor for each (reshaped) event
+        from a normal distribution with this standard deviation.
+    """
+    hits = ak.Array(
+        {
+            "num_scint_ph": num_scint_ph,
+            "particle": particle,
+            "time": units_conv_ak(time, "ns"),
+            "xloc": units_conv_ak(xloc, "m"),
+            "yloc": units_conv_ak(yloc, "m"),
+            "zloc": units_conv_ak(zloc, "m"),
+        }
+    )
+
+    scint_mat_params = convolve._get_scint_params(material)
+    pe = convolve.iterate_stepwise_depositions_pois(
+        hits, optmap, scint_mat_params, spm_detector, map_scaling, map_scaling_sigma
+    )
+
+    return VectorOfVectors(pe, attrs={"units": "ns"})
+
+
+def emitted_scintillation_photons(
+    edep: ak.Array, particle: ak.Array, material: str
+) -> VectorOfVectors:
+    """Derive the number of emitted scintillation photons from scintillator hits.
+
+    Parameters
+    ----------
+    edep
+        array of deposited energy in scintillation events.
+    particle
+        array of particle PDG IDs of scintillation events.
+    material
+        scintillating material name.
+    """
+    hits = ak.Array({"edep": units_conv_ak(edep, "keV"), "particle": particle})
+
+    scint_mat_params = convolve._get_scint_params(material)
+    ph = convolve.iterate_stepwise_depositions_scintillate(hits, scint_mat_params)
+    return VectorOfVectors(ph)