Protect median against nans in get_prototype_spike

SpikeInterface · Jul 29, 2024 · 983cf75 · 983cf75
1 parent 1148b88
commit 983cf75
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/spikeinterface/sortingcomponents/tools.py b/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):
+    from spikeinterface.sortingcomponents.peak_selection import select_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, few_peaks, ms_before=ms_before, ms_after=ms_after, **job_kwargs
     )
     with np.errstate(divide="ignore", invalid="ignore"):
-        prototype = np.median(waveforms[:, :, 0] / (np.abs(waveforms[:, nbefore, 0][:, np.newaxis])), axis=0)
+        prototype = np.nanmedian(waveforms[:, :, 0] / (np.abs(waveforms[:, nbefore, 0][:, np.newaxis])), axis=0)
     return prototype