get_spike_prototype can have NaN sometimes (#2980)

avoid get_spike_prototype can have NaN sometimes with margin in select_peaks

Author: yger

src/spikeinterface/sorters/internal/simplesorter.py

@@ -112,9 +112,12 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
 
         ms_before = params["waveforms"]["ms_before"]
         ms_after = params["waveforms"]["ms_after"]
+        nbefore = int(ms_before * sampling_frequency / 1000.0)
+        nafter = int(ms_after * sampling_frequency / 1000.0)
 
         # SVD for time compression
-        few_peaks = select_peaks(peaks, method="uniform", n_peaks=5000)
+
+        few_peaks = select_peaks(peaks, recording=recording, method="uniform", n_peaks=5000, margin=(nbefore, nafter))
         few_wfs = extract_waveform_at_max_channel(
             recording, few_peaks, ms_before=ms_before, ms_after=ms_after, **job_kwargs
         )

src/spikeinterface/sortingcomponents/clustering/circus.py

@@ -90,7 +90,7 @@ def main_function(cls, recording, peaks, params):
         tmp_folder.mkdir(parents=True, exist_ok=True)
 
         # SVD for time compression
-        few_peaks = select_peaks(peaks, method="uniform", n_peaks=10000)
+        few_peaks = select_peaks(peaks, recording=recording, method="uniform", n_peaks=10000, margin=(nbefore, nafter))
         few_wfs = extract_waveform_at_max_channel(
             recording, few_peaks, ms_before=ms_before, ms_after=ms_after, **params["job_kwargs"]
         )

src/spikeinterface/sortingcomponents/clustering/tdc.py

@@ -73,8 +73,11 @@ def main_function(cls, recording, peaks, params):
         ms_before = params["waveforms"]["ms_before"]
         ms_after = params["waveforms"]["ms_after"]
 
+        nbefore = int(ms_before * sampling_frequency / 1000.0)
+        nafter = int(ms_after * sampling_frequency / 1000.0)
+
         # SVD for time compression
-        few_peaks = select_peaks(peaks, method="uniform", n_peaks=5000)
+        few_peaks = select_peaks(peaks, recording=recording, method="uniform", n_peaks=5000, margin=(nbefore, nafter))
         few_wfs = extract_waveform_at_max_channel(
             recording, few_peaks, ms_before=ms_before, ms_after=ms_after, **job_kwargs
         )

src/spikeinterface/sortingcomponents/peak_selection.py

@@ -6,7 +6,9 @@
 import numpy as np
 
 
-def select_peaks(peaks, method="uniform", seed=None, return_indices=False, **method_kwargs):
+def select_peaks(
+    peaks, recording=None, method="uniform", seed=None, return_indices=False, margin=None, **method_kwargs
+):
     """
     Method to select a subset of peaks from a set of peaks.
     Usually use for reducing computational foorptint of downstream methods.
@@ -28,6 +30,9 @@ def select_peaks(peaks, method="uniform", seed=None, return_indices=False, **met
         The seed for random generations
     return_indices: bool
         If True, return the indices of selection such that selected_peaks = peaks[selected_indices]
+    margin : Margin in timesteps. default: None. Otherwise should be a tuple (nbefore, nafter)
+        preventing peaks to be selected at the borders of the segments. A recording should be provided to get the duration
+        of the segments
 
     method_kwargs: dict of kwargs method
         Keyword arguments for the chosen method:
@@ -66,8 +71,27 @@ def select_peaks(peaks, method="uniform", seed=None, return_indices=False, **met
         return_indices is True.
     """
 
+    if margin is not None:
+        assert recording is not None, "recording should be provided if margin is not None"
+
     selected_indices = select_peak_indices(peaks, method=method, seed=seed, **method_kwargs)
     selected_peaks = peaks[selected_indices]
+
+    if margin is not None:
+        to_keep = np.zeros(len(selected_peaks), dtype=bool)
+        offset = 0
+        for segment_index in range(recording.get_num_segments()):
+            duration = recording.get_num_frames(segment_index)
+            i0, i1 = np.searchsorted(selected_peaks["segment_index"], [segment_index, segment_index + 1])
+            while selected_peaks["sample_index"][i0] <= margin[0] + offset:
+                i0 += 1
+            while selected_peaks["sample_index"][i1 - 1] >= (duration - margin[1]) + offset:
+                i1 -= 1
+            to_keep[i0:i1] = True
+            offset += duration
+        selected_indices = selected_indices[to_keep]
+        selected_peaks = peaks[selected_indices]
+
     if return_indices:
         return selected_peaks, selected_indices
     else:

src/spikeinterface/sortingcomponents/tests/test_peak_selection.py

@@ -45,6 +45,11 @@ def test_select_peaks():
             selected_peaks.size <= n_peaks
         ), "selected_peaks is not the right size when return_indices=False, select_per_channel=False"
 
+        selected_peaks = select_peaks(peaks, recording=recording, method=method, margin=(10, 10), **select_kwargs)
+        assert (
+            selected_peaks.size <= n_peaks
+        ), "selected_peaks is not the right size when return_indices=False, select_per_channel=False"
+
         selected_peaks = select_peaks(peaks, method=method, select_per_channel=True, **select_kwargs)
         assert selected_peaks.size <= (
             n_peaks * recording.get_num_channels()

src/spikeinterface/sortingcomponents/tools.py

@@ -70,18 +70,18 @@ def extract_waveform_at_max_channel(rec, peaks, ms_before=0.5, ms_after=1.5, **j
 
 
 def get_prototype_spike(recording, peaks, ms_before=0.5, ms_after=0.5, nb_peaks=1000, **job_kwargs):
-    if peaks.size > nb_peaks:
-        idx = np.sort(np.random.choice(len(peaks), nb_peaks, replace=False))
-        some_peaks = peaks[idx]
-    else:
-        some_peaks = peaks
-
     nbefore = int(ms_before * recording.sampling_frequency / 1000.0)
+    nafter = int(ms_after * recording.sampling_frequency / 1000.0)
+
+    from spikeinterface.sortingcomponents.peak_selection import select_peaks
+
+    few_peaks = select_peaks(peaks, recording=recording, method="uniform", n_peaks=nb_peaks, margin=(nbefore, nafter))
 
     waveforms = extract_waveform_at_max_channel(
-        recording, some_peaks, ms_before=ms_before, ms_after=ms_after, **job_kwargs
+        recording, few_peaks, ms_before=ms_before, ms_after=ms_after, **job_kwargs
     )
-    prototype = np.nanmedian(waveforms[:, :, 0] / (np.abs(waveforms[:, nbefore, 0][:, np.newaxis])), axis=0)
+    with np.errstate(divide="ignore", invalid="ignore"):
+        prototype = np.median(waveforms[:, :, 0] / (np.abs(waveforms[:, nbefore, 0][:, np.newaxis])), axis=0)
     return prototype