Waveform attributes

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.njit(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

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"