diff --git a/strax/processing/peak_properties.py b/strax/processing/peak_properties.py index ece9ebd14..6ca39b417 100644 --- a/strax/processing/peak_properties.py +++ b/strax/processing/peak_properties.py @@ -100,3 +100,52 @@ def compute_widths(peaks, select_peaks_indices=None): i = len(desired_fr) // 2 peaks['width'][select_peaks_indices] = fr_times[:, i:] - fr_times[:, ::-1][:, i:] peaks['area_decile_from_midpoint'][select_peaks_indices] = fr_times[:, ::2] - fr_times[:, i].reshape(-1,1) + +@export +@numba.jit(nopython=True, cache=True) +def compute_wf_attributes(data, sample_length, n_samples: int, downsample_wf=False): + """ + Compute waveform attribures + Quantiles: represent the amount of time elapsed for + a given fraction of the total waveform area to be observed in n_samples + i.e. n_samples = 10, then quantiles are equivalent deciles + Waveforms: downsampled waveform to n_samples + :param data: waveform e.g. peaks or peaklets + :param n_samples: compute quantiles for a given number of samples + :return: waveforms and quantiles of size n_samples + """ + assert data.shape[0] == len(sample_length), "ararys must have same size" + + num_samples = data.shape[1] + + waveforms = np.zeros((len(data), n_samples), dtype=np.float64) + quantiles = np.zeros((len(data), n_samples), dtype=np.float64) + + # Cannot compute with with more samples than actual waveform sample + assert num_samples > n_samples, "cannot compute with more samples than the actual waveform" + + + step_size = int(num_samples / n_samples) + steps = np.arange(0, num_samples + 1, step_size) + inter_points = np.linspace(0., 1. - (1. / n_samples), n_samples) + cumsum_steps = np.zeros(n_samples + 1, dtype=np.float64) + frac_of_cumsum = np.zeros(num_samples + 1) + sample_number_div_dt = np.arange(0, num_samples + 1, 1) + for i, (samples, dt) in enumerate(zip(data, sample_length)): + if np.sum(samples) == 0: + continue + # reset buffers + frac_of_cumsum[:] = 0 + cumsum_steps[:] = 0 + frac_of_cumsum[1:] = np.cumsum(samples) + frac_of_cumsum[1:] = frac_of_cumsum[1:] / frac_of_cumsum[-1] + cumsum_steps[:-1] = np.interp(inter_points, frac_of_cumsum, sample_number_div_dt * dt) + cumsum_steps[-1] = sample_number_div_dt[-1] * dt + quantiles[i] = cumsum_steps[1:] - cumsum_steps[:-1] + + if downsample_wf: + + for j in range(n_samples): + waveforms[i][j] = np.sum(samples[steps[j]:steps[j + 1]]) + waveforms[i] /= (step_size * dt) + return quantiles, waveforms diff --git a/tests/test_peak_properties.py b/tests/test_peak_properties.py index c368481f6..c81bf7f25 100644 --- a/tests/test_peak_properties.py +++ b/tests/test_peak_properties.py @@ -3,6 +3,7 @@ from hypothesis import given, strategies, example, settings import tempfile from math import ceil, floor +import pytest def get_filled_peaks(peak_length, data_length, n_widths): dtype = [(('Start time since unix epoch [ns]', 'time'), np.int64), @@ -79,3 +80,30 @@ def test_compute_widths(peak_length, data_length, n_widths): mess = ("Highly unlikely that from at least 10 positive area peaks " "none were able to compute the width") assert np.any(peaks['width'] != pre_peaks['width']), mess + +@settings(max_examples=100, deadline=None) +@given( + # number of peaks + strategies.integers(min_value=0, max_value=20), + # length of the data field in the peaks + strategies.integers(min_value=2, max_value=20), + # Number of widths to compute + strategies.integers(min_value=2, max_value=10), +) +def test_compute_wf_attributes(peak_length, data_length, n_widths): + """ + Test strax.compute_wf_attribute + """ + peaks = get_filled_peaks(peak_length, data_length, n_widths) + wf = np.zeros((len(peaks), 10), dtype=np.float64) + q = np.zeros((len(peaks), 10), dtype=np.float64) + + try: + q, wf = strax.compute_wf_attributes(peaks['data'], peaks['dt'], 10, True) + except AssertionError as e: + if "zero waveform" in str(e): + print("cannot compute with a zero waveform") + elif "more samples than the actual waveform" in str(e): + print("cannot compute with more samples than the actual waveform") + + assert np.all(~np.isnan(q)) and np.all(~np.isnan(wf)), "attributes contains NaN values"