reboost 0.8.3__py3-none-any.whl

reboost/hpge/surface.py

@@ -0,0 +1,284 @@
+from __future__ import annotations
+
+import logging
+
+import awkward as ak
+import numba
+import numpy as np
+import pygeomhpges
+from lgdo import VectorOfVectors
+from lgdo.types import LGDO
+from numpy.typing import ArrayLike
+from scipy import stats
+
+log = logging.getLogger(__name__)
+
+
+def distance_to_surface(
+    positions_x: VectorOfVectors,
+    positions_y: VectorOfVectors,
+    positions_z: VectorOfVectors,
+    hpge: pygeomhpges.base.HPGe,
+    det_pos: ArrayLike,
+    *,
+    surface_type: str | None = None,
+    unit: str = "mm",
+    distances_precompute: VectorOfVectors | None = None,
+    precompute_cutoff: float | None = None,
+) -> VectorOfVectors:
+    """Computes the distance from each step to the detector surface.
+
+    The calculation can be performed for any surface type `nplus`, `pplus`,
+    `passive` or `None`. In order to speed up the calculation we provide
+    an option to only compute the distance for points within a certain distance
+    of any surface (as computed by remage and stored in the "distances_precompute") argument.
+
+    Parameters
+    ----------
+    positions_x
+        Global x positions for each step.
+    positions_y
+        Global y positions for each step.
+    positions_z
+        Global z positions for each step.
+    hpge
+        HPGe object.
+    det_pos
+        position of the detector origin, must be a 3 component array corresponding to `(x,y,z)`.
+    surface_type
+        string of which surface to use, can be `nplus`, `pplus` `passive` or None (in which case the distance to any surface is calculated).
+    unit
+        unit for the hit tier positions table.
+    distances_precompute
+        VectorOfVectors of distance to any surface computed by remage.
+    precompute_cutoff
+        cutoff on distances_precompute to not compute the distance for (in mm)
+
+    Returns
+    -------
+    VectorOfVectors with the same shape as `positions_x/y/z` of the distance to the surface.
+
+    Note
+    ----
+    `positions_x/positions_y/positions_z` must all have the same shape.
+    """
+    factor = np.array([1, 100, 1000])[unit == np.array(["mm", "cm", "m"])][0]
+
+    # compute local positions
+    pos = []
+    sizes = []
+
+    for idx, pos_tmp in enumerate([positions_x, positions_y, positions_z]):
+        local_pos_tmp = ak.Array(pos_tmp) * factor - det_pos[idx]
+        local_pos_flat_tmp = ak.flatten(local_pos_tmp).to_numpy()
+        pos.append(local_pos_flat_tmp)
+        sizes.append(ak.num(local_pos_tmp, axis=1))
+
+    if not ak.all(sizes[0] == sizes[1]) or not ak.all(sizes[0] == sizes[2]):
+        msg = "all position vector of vector must have the same shape"
+        raise ValueError(msg)
+
+    size = sizes[0]
+    # restructure the positions
+    local_positions = np.vstack(pos).T
+
+    # get indices
+    surface_indices = (
+        np.where(np.array(hpge.surfaces) == surface_type) if surface_type is not None else None
+    )
+
+    # distance calc itself
+    if distances_precompute is None:
+        distances = hpge.distance_to_surface(local_positions, surface_indices=surface_indices)
+    else:
+        # decide when the calculation needs to be run
+        if isinstance(distances_precompute, LGDO):
+            distances_precompute = distances_precompute.view_as("ak")
+
+        distances_precompute_flat = ak.flatten(distances_precompute)
+        distances = np.full_like(distances_precompute_flat.to_numpy(), np.nan, dtype=float)
+
+        # values to compute
+        indices = distances_precompute_flat < precompute_cutoff
+
+        # compute the distances
+        distances[indices] = hpge.distance_to_surface(
+            local_positions[indices], surface_indices=surface_indices
+        )
+
+    return VectorOfVectors(ak.unflatten(distances, size), dtype=np.float32)
+
+
+@numba.njit(cache=True)
+def _advance_diffusion(
+    charge: np.ndarray,
+    factor: float,
+    recomb: float = 0,
+    recomb_depth: float = 600,
+    delta_x: float = 10,
+):
+    """Make a step of diffusion using explicit Euler scheme.
+
+    Parameters
+    ----------
+    charge
+        charge in each space bin up to the FCCD
+    factor
+        the factor of diffusion for the Euler scheme
+    recomb
+        the recomination probability.
+    recomb_depth
+        the depth of the recombination region.
+    delta_x
+        the width of each spatial bin.
+
+    Returns
+    -------
+    a tuple of the charge distribution at the next time step and the collected charge.
+    """
+    charge_xp1 = np.append(charge[1:], [0])
+    charge_xm1 = np.append([0], charge[:-1])
+
+    # collected charge
+    collected = factor * charge[-1]
+
+    # charge at the next step
+    charge_new = charge_xp1 * factor + charge_xm1 * factor + charge * (1 - 2 * factor)
+
+    # correction for recombination
+    charge_new[0 : int(recomb_depth / delta_x)] = (1 - recomb) * charge_new[
+        0 : int(recomb_depth / delta_x)
+    ]
+
+    return charge_new, collected
+
+
+@numba.njit(cache=True)
+def _compute_diffusion_impl(
+    init_charge: np.ndarray,
+    nsteps: int,
+    factor: float,
+    recomb: float = 0,
+    recomb_depth: float = 600,
+    delta_x: float = 10,
+):
+    """Compute the charge collected as a function of time.
+
+    Parameters
+    ----------
+    init_charge
+        Initial charge distribution.
+    nsteps
+        Number of time steps to take.
+    kwargs
+        Keyword arguments to pass to :func:`_advance_diffusion`
+    """
+    charge = init_charge
+    collected_charge = np.zeros(nsteps)
+
+    for i in range(nsteps):
+        charge, collected = _advance_diffusion(
+            charge, factor=factor, recomb=recomb, recomb_depth=recomb_depth, delta_x=delta_x
+        )
+        collected_charge[i] = collected
+
+    return collected_charge
+
+
+def get_surface_library(fccd: float, dist_step_in_um: float, **kwargs):
+    """Build the surface response library by calling `reboost.hpge.surface.get_surface_response`.
+
+    Parameters
+    ----------
+    fccd
+        The value of the FCCD
+    dist_step_in_um
+        The distance steps to use in building the library
+    **kwargs
+        Other keyword arguments to `reboost.hpge.surface.get_surface_response`.
+
+    Returns
+    -------
+    2D array of the cumulative-charge arriving at the p-n junction as a function
+    of time, for each distance.
+    """
+    steps = int(fccd / dist_step_in_um)
+
+    out = np.zeros((10000, steps))
+
+    for step in range(steps):
+        out[:, step] = get_surface_response(fccd, init=step * dist_step_in_um, **kwargs)
+
+    return out
+
+
+def get_surface_response(
+    fccd: float,
+    init: float = 0,
+    *,
+    recomb_depth: float = 500,
+    recomb: float = 0.002,
+    init_size: float = 0.0,
+    factor: float = 0.29,
+    nsteps: int = 10000,
+    delta_x: float = 10,
+):
+    """Extract the surface response current pulse based on diffusion.
+
+    This extracts the amount of charge arrived (cumulative) at the p-n
+    junction as a function of time, based on diffusion. The final
+    induced waveform on the p-n contact is obtained from convolution
+    with the bulk pulse.
+
+    Parameters
+    ----------
+    fccd
+        the full charge collection depth (in um)
+    recomb_depth
+        the depth of the recombination region (in um)
+    init
+        the initial position of the charge (in um)
+    recomb
+        the recombination rate
+    init_size
+        the initial size of the charge cloud (in um)
+    factor
+        the factor for the explicit Euler scheme (the probability of charge diffusuion)
+    nsteps
+        the number of time steps.
+    delta_x
+        the width of each position bin.
+    """
+    # number of position steps
+    nx = int(fccd / delta_x)
+
+    # initial charge
+    charge = np.zeros(nx)
+
+    # generate initial conditions
+    x = (fccd / nx) * np.arange(nx)
+    x_full = (fccd / nx) * np.arange(2 * nx)
+
+    # generate initial conditions
+    if init_size != 0:
+        charge = stats.norm.pdf(x, loc=init, scale=init_size)
+        charge_full = stats.norm.pdf(x_full, loc=init, scale=init_size)
+        charge_col = [(np.sum(charge_full) - np.sum(charge)) / np.sum(charge_full)]
+        charge = charge / np.sum(charge_full)
+    elif int(init * nx / fccd) < len(charge):
+        charge[int(init * nx / fccd)] = 1
+        charge_col = np.array([])
+    else:
+        charge_col = np.array([1])
+
+    # run the simulation
+    charge_collected = _compute_diffusion_impl(
+        charge,
+        nsteps=nsteps,
+        factor=factor,
+        recomb=recomb,
+        recomb_depth=recomb_depth,
+        delta_x=delta_x,
+    )
+
+    return np.cumsum(np.concatenate((charge_col, charge_collected)))

reboost/hpge/utils.py

@@ -0,0 +1,79 @@
+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)

reboost/iterator.py

@@ -0,0 +1,226 @@
+from __future__ import annotations
+
+import logging
+import time
+import typing
+
+import awkward as ak
+from lgdo.lh5 import LH5Store
+from lgdo.types import LGDO, Table
+
+from . import build_glm
+
+log = logging.getLogger(__name__)
+
+
+class GLMIterator:
+    """A class to iterate over the rows of an event lookup map."""
+
+    def __init__(
+        self,
+        glm_file: str | None,
+        stp_file: str,
+        lh5_group: str,
+        start_row: int,
+        n_rows: int | None,
+        *,
+        stp_field: str = "stp",
+        buffer: int = 10000,
+        time_dict: dict | None = None,
+        reshaped_files: bool = False,
+    ):
+        """Constructor for the GLMIterator.
+
+        The GLM iterator provides a way to iterate over the
+        simulated geant4 evtids, extracting the number of hits or steps for
+        each range in evtids. This ensures a single simulated event
+        is not split between two iterations and allows to specify a
+        start and an end evtid to extract.
+
+        In case the data is already reshaped and we do not need to
+        read a specific range of evtids this iterator is just loops
+        over the input stp field. Otherwise if the GLM file is not provided
+        this is created in memory.
+
+        Parameters
+        ----------
+        glm_file
+            the file containing the event lookup map, if `None` the glm will
+            be created in memory if needed.
+        stp_file
+            the file containing the steps to read.
+        lh5_group
+            the name of the lh5 group to read.
+        start_row
+            the first row to read.
+        n_rows
+            the number of rows to read, if `None` read them all.
+        stp_field
+            name of the group.
+        buffer
+            the number of rows to read at once.
+        time_dict
+            time profiling data structure.
+        reshaped_files
+            flag for whether the files are reshaped.
+        """
+        # initialise
+        self.glm_file = glm_file
+        self.stp_file = stp_file
+        self.lh5_group = lh5_group
+        self.start_row = start_row
+        self.start_row_tmp = start_row
+        self.n_rows = n_rows
+        self.buffer = buffer
+        self.current_i_entry = 0
+        self.stp_field = stp_field
+        self.reshaped_files = reshaped_files
+
+        # would be good to replace with an iterator
+        self.sto = LH5Store()
+        self.n_rows_read = 0
+        self.time_dict = time_dict
+        self.glm = None
+        self.use_glm = True
+
+        glm_n_rows = 0
+
+        # build the glm in memory if needed
+        if self.glm_file is None and (
+            (self.n_rows is not None) or (self.start_row != 0) or not reshaped_files
+        ):
+            if self.time_dict is not None:
+                time_start = time.time()
+
+            self.glm = build_glm.build_glm(
+                stp_file, None, out_table_name="glm", id_name="evtid", lh5_groups=[lh5_group]
+            )
+
+            if self.time_dict is not None:
+                self.time_dict.update_field("read/glm", time_start)
+
+            glm_n_rows = len(self.glm)
+
+        elif self.glm_file is None:
+            self.use_glm = False
+        else:
+            glm_n_rows = self.sto.read_n_rows(f"glm/{self.lh5_group}", self.glm_file)
+
+        # get the number of stp rows
+        try:
+            stp_n_rows = self.sto.read_n_rows(f"{self.stp_field}/{self.lh5_group}", self.stp_file)
+        except Exception:
+            stp_n_rows = 0
+
+        # heuristics for a good buffer length
+        if self.use_glm:
+            self.buffer = int(buffer * glm_n_rows / (1 + stp_n_rows))
+            msg = f"Number of stp rows {stp_n_rows}, number of glm rows {glm_n_rows} changing buffer from {buffer} to {self.buffer}"
+            log.debug(msg)
+
+    def __iter__(self) -> typing.Iterator:
+        self.current_i_entry = 0
+        self.n_rows_read = 0
+        self.start_row_tmp = self.start_row
+        return self
+
+    def get_n_rows(self):
+        """Get the number of rows to read."""
+        # get the number of rows to read
+        if self.time_dict is not None:
+            time_start = time.time()
+
+        if self.n_rows is not None:
+            rows_left = self.n_rows - self.n_rows_read
+            n_rows = self.buffer if (self.buffer > rows_left) else rows_left
+        else:
+            n_rows = self.buffer
+
+        glm_rows = None
+        start = 0
+        n = 0
+
+        if self.use_glm:
+            if self.glm_file is not None:
+                glm_rows = self.sto.read(
+                    f"glm/{self.lh5_group}",
+                    self.glm_file,
+                    start_row=self.start_row_tmp,
+                    n_rows=n_rows,
+                )
+                n_rows_read = len(glm_rows.view_as("ak"))
+
+            else:
+                # get the maximum row to read
+                max_row = self.start_row_tmp + n_rows
+                max_row = min(len(self.glm[self.lh5_group]), max_row)
+
+                if max_row != self.start_row_tmp:
+                    glm_rows = Table(self.glm[self.lh5_group][self.start_row_tmp : max_row])
+
+                n_rows_read = max_row - self.start_row_tmp
+
+            if self.time_dict is not None:
+                self.time_dict.update_field("read/glm", time_start)
+
+            self.n_rows_read += n_rows_read
+            self.start_row_tmp += n_rows_read
+
+            # view our glm as an awkward array
+            if glm_rows is not None:
+                glm_ak = glm_rows.view_as("ak")
+
+                # remove empty rows
+                glm_ak = glm_ak[glm_ak.n_rows > 0]
+
+                if len(glm_ak) > 0:
+                    # extract range of stp rows to read
+                    start = glm_ak.start_row[0]
+                    n = ak.sum(glm_ak.n_rows)
+
+        else:
+            start = self.start_row_tmp
+            n = n_rows
+            n_rows_read = n
+            self.start_row_tmp += n
+
+        return start, n, n_rows_read
+
+    def __next__(self) -> tuple[LGDO, int, int]:
+        """Read one chunk.
+
+        Returns
+        -------
+        a tuple of:
+            - the steps
+            - the chunk index
+            - the number of steps read
+        """
+        # read the glm rows]
+        start, n, n_rows_read = self.get_n_rows()
+
+        if self.time_dict is not None:
+            time_start = time.time()
+
+        try:
+            stp_rows = self.sto.read(
+                f"{self.stp_field}/{self.lh5_group}",
+                self.stp_file,
+                start_row=int(start),
+                n_rows=int(n),
+            )
+            n_steps = len(stp_rows.view_as("ak"))
+
+        except OverflowError:
+            raise StopIteration from None
+
+        if n_rows_read == 0 or n_steps == 0:
+            raise StopIteration
+
+        # save time
+        if self.time_dict is not None:
+            self.time_dict.update_field("read/stp", time_start)
+
+        self.current_i_entry += 1
+
+        return (stp_rows, self.current_i_entry - 1, n_steps)

reboost/log_utils.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+import logging
+
+import colorlog
+
+
+def setup_log(level: int | None = None, multiproc: bool = False) -> None:
+    """Setup a colored logger for this package.
+
+    Parameters
+    ----------
+    level
+        initial log level, or ``None`` to use the default.
+    multiproc
+        set to ``True`` to include process ID in log output (i.e. for multiprocessing setups)
+    """
+    fmt = "%(log_color)s%(name)s [%(levelname)s]"
+    if multiproc:
+        fmt += " (pid=%(process)s)"
+    fmt += " %(message)s"
+
+    handler = colorlog.StreamHandler()
+    handler.setFormatter(colorlog.ColoredFormatter(fmt))
+
+    logger = logging.getLogger("reboost")
+    logger.addHandler(handler)
+    if level is not None:
+        logger.setLevel(level)