Fix remove_artifacts sparsity handling for median/average modes (#3290) (#4685)

Co-authored-by: Claude <noreply@anthropic.com>

Author: Leonard013

src/spikeinterface/preprocessing/remove_artifacts.py

@@ -376,7 +376,12 @@ def get_traces(self, start_frame, end_frame, channel_indices):
                     mask = self.sparsity[label]
                 else:
                     mask = None
-                artifact_duration = len(self.artifacts[label])
+                artifact = self.artifacts[label]
+                if mask is not None:
+                    # restrict the template to the sparse channels so that it aligns
+                    # with the sparsified traces (see issue #3290)
+                    artifact = artifact[:, mask]
+                artifact_duration = len(artifact)
                 if self.time_pad > 0:
                     jitters = np.arange(-self.time_pad, self.time_pad, 1)
                 else:
@@ -404,7 +409,7 @@ def get_traces(self, start_frame, end_frame, channel_indices):
                     if mask is not None:
                         trace_slice_values = trace_slice_values[:, mask]
 
-                    artifact_slice_values = self.artifacts[label][artifact_slice]
+                    artifact_slice_values = artifact[artifact_slice]
 
                     norm = np.linalg.norm(trace_slice_values) * np.linalg.norm(artifact_slice_values)
                     best_amplitudes[count] = (
@@ -435,11 +440,9 @@ def get_traces(self, start_frame, end_frame, channel_indices):
                     best_amp = 1
 
                 if mask is not None:
-                    traces[trace_slice][:, mask] -= (best_amp * self.artifacts[label][artifact_slice]).astype(
-                        traces.dtype
-                    )
+                    traces[trace_slice][:, mask] -= (best_amp * artifact[artifact_slice]).astype(traces.dtype)
                 else:
-                    traces[trace_slice] -= (best_amp * self.artifacts[label][artifact_slice]).astype(traces.dtype)
+                    traces[trace_slice] -= (best_amp * artifact[artifact_slice]).astype(traces.dtype)
             traces = traces[:, channel_indices]
 
         return traces

src/spikeinterface/preprocessing/tests/test_remove_artifacts.py

@@ -85,5 +85,47 @@ def test_remove_artifacts():
     )
 
 
+def test_remove_artifacts_sparsity():
+    # non-regression test for issue #3290: "median"/"average" modes with sparsity
+    # used to raise "shapes not aligned" because the template was kept full-channel
+    # while the traces were sparsified before the np.dot / subtraction.
+    rec = generate_recording(num_channels=4, durations=[10.0], seed=0)
+    rec.annotate(is_filtered=True)
+
+    ms = 10
+    ms_frames = int(ms * rec.get_sampling_frequency() / 1000)
+
+    # two artifact labels, each removed on a different subset of channels
+    triggers = [15000, 30000, 45000, 60000]
+    labels = [0, 1, 0, 1]
+    list_triggers = [triggers]
+    list_labels = [labels]
+    sparsity = {
+        0: np.array([True, False, True, False]),
+        1: np.array([False, True, False, True]),
+    }
+
+    for mode in ("median", "average"):
+        rec_rmart = remove_artifacts(
+            rec,
+            list_triggers,
+            ms_before=ms,
+            ms_after=ms,
+            mode=mode,
+            list_labels=list_labels,
+            sparsity=sparsity,
+        )
+        for trig, label in zip(triggers, labels):
+            mask = sparsity[label]
+            traces_clean = rec.get_traces(start_frame=trig - ms_frames, end_frame=trig + ms_frames)
+            # get_traces must not raise (the bug reported in issue #3290)
+            traces = rec_rmart.get_traces(start_frame=trig - ms_frames, end_frame=trig + ms_frames)
+            # channels outside the sparsity mask are left untouched
+            assert np.array_equal(traces[:, ~mask], traces_clean[:, ~mask])
+            # channels inside the sparsity mask have the artifact subtracted
+            assert not np.allclose(traces[:, mask], traces_clean[:, mask])
+
+
 if __name__ == "__main__":
     test_remove_artifacts()
+    test_remove_artifacts_sparsity()