Patch to skip chunks without peaks or if enough peaks have been detected

src/spikeinterface/core/node_pipeline.py

@@ -117,6 +117,12 @@ def __init__(self, recording, peaks):
     def get_trace_margin(self):
         return 0
 
+    def has_peaks(self, start_frame, end_frame, segment_index):
+        sl = self.segment_slices[segment_index]
+        peaks_in_segment = self.peaks[sl]
+        i0, i1 = np.searchsorted(peaks_in_segment["sample_index"], [start_frame, end_frame])
+        return i0 < i1
+
     def get_dtype(self):
         return base_peak_dtype
 
@@ -187,6 +193,12 @@ def __init__(
     def get_trace_margin(self):
         return 0
 
+    def has_peaks(self, start_frame, end_frame, segment_index):
+        sl = self.segment_slices[segment_index]
+        peaks_in_segment = self.peaks[sl]
+        i0, i1 = np.searchsorted(peaks_in_segment["sample_index"], [start_frame, end_frame])
+        return i0 < i1
+
     def get_dtype(self):
         return base_peak_dtype
 
@@ -486,9 +498,18 @@ def _compute_peak_pipeline_chunk(segment_index, start_frame, end_frame, worker_c
     nodes = worker_ctx["nodes"]
 
     recording_segment = recording._recording_segments[segment_index]
-    traces_chunk, left_margin, right_margin = get_chunk_with_margin(
-        recording_segment, start_frame, end_frame, None, max_margin, add_zeros=True
-    )
+
+    if isinstance(nodes[0], PeakRetriever) or isinstance(nodes[0], SpikeRetriever):
+        if nodes[0].has_peaks(start_frame, end_frame, segment_index):
+            traces_chunk, left_margin, right_margin = get_chunk_with_margin(
+                recording_segment, start_frame, end_frame, None, max_margin, add_zeros=True
+            )
+        else:
+            traces_chunk, left_margin, right_margin = None, None, None
+    else:
+        traces_chunk, left_margin, right_margin = get_chunk_with_margin(
+            recording_segment, start_frame, end_frame, None, max_margin, add_zeros=True
+        )
 
     # compute the graph
     pipeline_outputs = {}