diff --git a/src/spikeinterface/sortingcomponents/matching/tests/test_wobble.py b/src/spikeinterface/sortingcomponents/matching/tests/test_wobble.py index 0d46b790ad..36d5452933 100644 --- a/src/spikeinterface/sortingcomponents/matching/tests/test_wobble.py +++ b/src/spikeinterface/sortingcomponents/matching/tests/test_wobble.py @@ -1,3 +1,5 @@ +from types import SimpleNamespace + import pytest import numpy as np @@ -261,6 +263,63 @@ def test_compute_objective(): assert np.allclose(objective, expected_objective) +def test_find_strict_peaks_matches_argrelmax(): + from scipy.signal import argrelmax + + rng = np.random.default_rng(0) + nan_objective = rng.random(200).astype(np.float32) + nan_objective[50] = np.nan + cases = [ + (rng.random(2000).astype(np.float32), 5, 0.3), + (rng.random(50), 60, -1.0), + (rng.random(50), 61, -1.0), + (np.array([0.0, 1.0, 1.0, 1.0, 0.5, 2.0, 2.0, 0.0], dtype=np.float32), 2, -1.0), + (np.array([0.0, 3.0, 0.0, -np.inf, -np.inf, 0.0, 2.0, 0.0], dtype=np.float32), 2, -1.0), + (nan_objective, 4, 0.2), + (rng.random(2).astype(np.float32), 1, -1.0), + (rng.random(500).astype(np.float32), 3, 10.0), + ] + + for objective, order, threshold in cases: + expected = argrelmax(objective, order=order)[0] + expected = expected[objective[expected] > threshold] + actual = wobble._find_strict_peaks(objective, order, threshold) + assert np.array_equal(actual, expected) + + +def test_find_peaks_matches_argrelmax_windowing(monkeypatch): + from scipy.signal import argrelmax + + cases = [ + ([0.0, 0.0, 1.0, 4.0, 1.0, 3.0, 1.0, 0.0, 0.0], 3, 2.0), + ([0.0, 0.0, 0.0, 5.0, 1.0, 2.0, 1.0, 4.0, 0.0, 0.0], 4, 1.5), + ([0.0, 0.0, 2.0, 1.0, 2.0, 0.0, 0.0], 3, 3.0), + ] + + def no_high_res_shift(cls, spike_time_indices, *args): + zeros = np.zeros(spike_time_indices.size, dtype=np.int64) + return zeros, zeros, np.zeros(0, dtype=np.int64), np.zeros(0) + + monkeypatch.setattr(wobble.WobbleMatch, "calculate_high_res_shift", classmethod(no_high_res_shift)) + for objective_maximum, num_samples, threshold in cases: + objective_normalized = np.asarray([objective_maximum], dtype=np.float32) + window = objective_normalized[0, num_samples - 1 : -num_samples] + expected_indices = argrelmax(window, order=num_samples - 1)[0] + expected_indices = expected_indices[window[expected_indices] > threshold] + spike_train, scalings, distance_metric = wobble.WobbleMatch.find_peaks( + objective=objective_normalized, + objective_normalized=objective_normalized, + spike_trains=np.zeros((0, 2), dtype=np.int64), + params=SimpleNamespace(threshold=threshold, jitter_factor=1), + template_data=None, + template_meta=SimpleNamespace(num_samples=num_samples), + ) + + assert np.array_equal(spike_train[:, 0], expected_indices) + assert np.array_equal(scalings, np.ones(expected_indices.size, dtype=np.float32)) + assert np.array_equal(distance_metric, window[expected_indices]) + + def test_compute_scale_amplitudes(): # Arrange: Generate random 'data' seed = 0 diff --git a/src/spikeinterface/sortingcomponents/matching/wobble.py b/src/spikeinterface/sortingcomponents/matching/wobble.py index 92bf814f72..e9ee961c3e 100644 --- a/src/spikeinterface/sortingcomponents/matching/wobble.py +++ b/src/spikeinterface/sortingcomponents/matching/wobble.py @@ -326,6 +326,30 @@ def __post_init__(self): self.temporal, self.singular, self.spatial, self.temporal_jittered = self.compressed_templates +def _find_strict_peaks(objective, order, threshold): + """Find values strictly greater than their neighbors within ``order`` samples.""" + from scipy.ndimage import maximum_filter1d + + if objective.size < 3: + return np.zeros(0, dtype=np.intp) + if np.isnan(objective).any(): + from scipy.signal import argrelmax + + peak_indices = argrelmax(objective, order=order)[0] + return peak_indices[objective[peak_indices] > threshold] + + trailing_maximum = maximum_filter1d(objective, size=order, origin=(order - 1) // 2, mode="nearest") + left_maximum = np.empty_like(objective) + left_maximum[0] = objective[0] + left_maximum[1:] = trailing_maximum[:-1] + + trailing_maximum = maximum_filter1d(objective[::-1], size=order, origin=(order - 1) // 2, mode="nearest") + right_maximum = np.empty_like(objective) + right_maximum[-1] = objective[-1] + right_maximum[:-1] = trailing_maximum[-2::-1] + return np.flatnonzero((objective > threshold) & (objective > left_maximum) & (objective > right_maximum)) + + class WobbleMatch(BaseTemplateMatching): """Template matching method from the Paninski lab. @@ -618,16 +642,16 @@ def find_peaks( Finally, it generates a new spike train from the spike times, and returns it along with additional metrics about each spike. """ - from scipy import signal - # Get spike times (indices) using peaks in the objective objective_template_max = np.max(objective_normalized, axis=0) spike_window = (template_meta.num_samples - 1, objective_normalized.shape[1] - template_meta.num_samples) objective_windowed = objective_template_max[spike_window[0] : spike_window[1]] - spike_time_indices = signal.argrelmax(objective_windowed, order=template_meta.num_samples - 1)[0] + spike_time_indices = _find_strict_peaks( + objective_windowed, + order=template_meta.num_samples - 1, + threshold=params.threshold, + ) spike_time_indices += template_meta.num_samples - 1 - objective_spikes = objective_template_max[spike_time_indices] - spike_time_indices = spike_time_indices[objective_spikes > params.threshold] if len(spike_time_indices) == 0: # No new spikes found return np.zeros((0, 2), dtype=np.int32), np.zeros(0), np.zeros(0)