From 1a37e3ec417f1b24efcd79e782f3d2f944e5bf01 Mon Sep 17 00:00:00 2001 From: Leonardo Scappatura <95079472+Leonard013@users.noreply.github.com> Date: Sun, 12 Jul 2026 15:16:03 +0200 Subject: [PATCH] Fix remove_artifacts sparsity handling for median/average modes When `sparsity` is provided with mode "median" or "average", `RemoveArtifactsRecordingSegment.get_traces` sparsified the traces to the masked channels but kept the artifact template full-channel. The np.dot and the subtraction then operated on mismatched channel dimensions and raised `ValueError: shapes ... not aligned`. Restrict the template to the sparse channels at use time so it aligns with the sparsified traces. The stored template (and `_kwargs`) stay full-channel, so serialization round-trips remain correct. The no-sparsity path is unchanged, and both artifact-provisioning paths (precomputed `artifacts` and templates computed via `estimate_templates`) are covered. Add a regression test exercising median and average modes with sparsity and multiple labels. Fixes #3290 This change was authored with the assistance of Claude (Anthropic). Co-authored-by: Claude --- .../preprocessing/remove_artifacts.py | 15 ++++--- .../tests/test_remove_artifacts.py | 42 +++++++++++++++++++ 2 files changed, 51 insertions(+), 6 deletions(-) diff --git a/src/spikeinterface/preprocessing/remove_artifacts.py b/src/spikeinterface/preprocessing/remove_artifacts.py index ded1bb6158..b20f1351b0 100644 --- a/src/spikeinterface/preprocessing/remove_artifacts.py +++ b/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 diff --git a/src/spikeinterface/preprocessing/tests/test_remove_artifacts.py b/src/spikeinterface/preprocessing/tests/test_remove_artifacts.py index 3461bf9b1b..85ba61cdd4 100644 --- a/src/spikeinterface/preprocessing/tests/test_remove_artifacts.py +++ b/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()