diff --git a/src/reboost/hpge/psd.py b/src/reboost/hpge/psd.py index 1254fc5..1e37615 100644 --- a/src/reboost/hpge/psd.py +++ b/src/reboost/hpge/psd.py @@ -13,7 +13,7 @@ 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. @@ -513,6 +513,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 @@ -546,11 +547,12 @@ def _get_waveform_value( ------- Value of the current waveform """ - n = len(template) out = 0 time = start + dt * idx for i in range(len(edep)): + n = len(template) + E = edep[i] mu = drift_time[i] @@ -559,6 +561,83 @@ def _get_waveform_value( 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. + + 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. + + Returns + ------- + Value of the current waveform + """ + 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( + 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) + 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) + 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 +677,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 +691,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 +707,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 +740,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, @@ -684,7 +765,8 @@ def _estimate_current_impl( 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 + array of the bulk pulse template, in the case of a full pulse shape library, 3 arrays can be passed corresponding to the + "r" and "z" coordinates of the library and the waveforms for each point. times time-stamps for the bulk pulse template """ @@ -693,7 +775,7 @@ def _estimate_current_impl( energy = np.zeros(len(dt)) time_step = 1 - n = len(template) + n = len(times) start = times[0] if include_surface_effects: @@ -707,6 +789,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 +805,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 +821,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 +835,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,8 +919,12 @@ 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 @@ -796,40 +943,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, diff --git a/src/reboost/hpge/utils.py b/src/reboost/hpge/utils.py index 370cad7..a8738bf 100644 --- a/src/reboost/hpge/utils.py +++ b/src/reboost/hpge/utils.py @@ -11,23 +11,105 @@ 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{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) diff --git a/test.lh5 b/test.lh5 new file mode 100644 index 0000000..9452b9a Binary files /dev/null and b/test.lh5 differ diff --git a/tests/conftest.py b/tests/conftest.py index c4c6666..7fe6249 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,8 +7,12 @@ 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 njit_patched(*args, **kwargs): 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() diff --git a/tests/hpge/test_current.py b/tests/hpge/test_current.py index a877b6d..a1a19ff 100644 --- a/tests/hpge/test_current.py +++ b/tests/hpge/test_current.py @@ -6,6 +6,8 @@ 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) diff --git a/tests/hpge/test_dt_heuristic.py b/tests/hpge/test_dt_heuristic.py index f70bc44..27b064b 100644 --- a/tests/hpge/test_dt_heuristic.py +++ b/tests/hpge/test_dt_heuristic.py @@ -80,7 +80,7 @@ @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): diff --git a/tests/hpge/test_hpge_map.py b/tests/hpge/test_hpge_map.py index 50a1520..6cda060 100644 --- a/tests/hpge/test_hpge_map.py +++ b/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)