Filter subsets of channels at a time to reduce peak memory

src/spikeinterface/preprocessing/filter.py

@@ -14,6 +14,7 @@
 
 HIGHPASS_ERROR_THRESHOLD_HZ = 100
 MARGIN_TO_CHUNK_PERCENT_WARNING = 0.2  # 20%
+NUM_CHANNELS_PER_BLOCK = 16
 
 
 _common_filter_docs = """**filter_kwargs : dict
@@ -194,34 +195,43 @@ def get_traces(self, start_frame, end_frame, channel_indices):
 
         from scipy.signal import sosfiltfilt, filtfilt, sosfilt, lfilter
 
-        if self.direction == "forward-backward":
-            if self.filter_mode == "sos":
-                filtered_traces = sosfiltfilt(self.coeff, traces_chunk, axis=0)
-            elif self.filter_mode == "ba":
-                b, a = self.coeff
-                filtered_traces = filtfilt(b, a, traces_chunk, axis=0)
-        else:
-            if self.direction == "backward":
-                traces_chunk = np.flip(traces_chunk, axis=0)
+        # Channels are filtered independently, so filtering in blocks gives the same result with a
+        # much lower peak memory (scipy makes several float copies of its input)
+        num_samples = traces_chunk.shape[0] - left_margin - right_margin
+        num_channels = traces_chunk.shape[1]
+        filtered_traces = np.empty((num_samples, num_channels), dtype=self.dtype)
+
+        for channel_index in range(0, num_channels, NUM_CHANNELS_PER_BLOCK):
+            channel_slice = slice(channel_index, channel_index + NUM_CHANNELS_PER_BLOCK)
+            block = traces_chunk[:, channel_slice]
+
+            if self.direction == "forward-backward":
+                if self.filter_mode == "sos":
+                    filtered_block = sosfiltfilt(self.coeff, block, axis=0)
+                elif self.filter_mode == "ba":
+                    b, a = self.coeff
+                    filtered_block = filtfilt(b, a, block, axis=0)
+            else:
+                if self.direction == "backward":
+                    block = np.flip(block, axis=0)
 
-            if self.filter_mode == "sos":
-                filtered_traces = sosfilt(self.coeff, traces_chunk, axis=0)
-            elif self.filter_mode == "ba":
-                b, a = self.coeff
-                filtered_traces = lfilter(b, a, traces_chunk, axis=0)
+                if self.filter_mode == "sos":
+                    filtered_block = sosfilt(self.coeff, block, axis=0)
+                elif self.filter_mode == "ba":
+                    b, a = self.coeff
+                    filtered_block = lfilter(b, a, block, axis=0)
 
-            if self.direction == "backward":
-                filtered_traces = np.flip(filtered_traces, axis=0)
+                if self.direction == "backward":
+                    filtered_block = np.flip(filtered_block, axis=0)
 
-        if right_margin > 0:
-            filtered_traces = filtered_traces[left_margin:-right_margin, :]
-        else:
-            filtered_traces = filtered_traces[left_margin:, :]
+            filtered_block = filtered_block[left_margin : filtered_block.shape[0] - right_margin, :]
+
+            if np.issubdtype(self.dtype, np.integer):
+                filtered_block = filtered_block.round()
 
-        if np.issubdtype(self.dtype, np.integer):
-            filtered_traces = filtered_traces.round()
+            filtered_traces[:, channel_slice] = filtered_block
 
-        return filtered_traces.astype(self.dtype)
+        return filtered_traces
 
 
 class BandpassFilterRecording(FilterRecording):