Repository: SpikeInterface/spikeinterface
Pull request #2649 (2024-04-01 10:03:55)

Fix get_traces for a local common reference
Fixes bug in `get_traces` for common reference recordings with a `local` reference, where you couldn't select a subset of channels. Noticed it when plotting traces.

Here's minimal code which reproduces the bug

```
from spikeinterface.core import generate_recording
from spikeinterface.preprocessing import common_reference

recording = generate_recording(durations=[5.0], num_channels=4)
recording_cmr_segment = common_reference(recording, reference="local")._recording_segments[0]

recording_cmr_segment.get_traces(start_frame=0, end_frame=10, channel_indices=[1,2])
```
This should raises an error.

Also added a test which selects a subset of channels ids, and makes sure it's selected the right one.


Patch:
diff --git a/src/spikeinterface/preprocessing/common_reference.py b/src/spikeinterface/preprocessing/common_reference.py
index 07920b61c..a4369309f 100644
--- a/src/spikeinterface/preprocessing/common_reference.py
+++ b/src/spikeinterface/preprocessing/common_reference.py
@@ -189,12 +189,12 @@ class CommonReferenceRecordingSegment(BasePreprocessorSegment):
                 shift = traces[:, self.ref_channel_indices]
                 re_referenced_traces = traces[:, channel_indices] - shift
             else:  # then it must be local
-                re_referenced_traces = np.zeros_like(traces[:, channel_indices])
                 channel_indices_array = np.arange(traces.shape[1])[channel_indices]
-                for channel_index in channel_indices_array:
+                re_referenced_traces = np.zeros((traces.shape[0], len(channel_indices_array)), dtype="float32")
+                for i, channel_index in enumerate(channel_indices_array):
                     channel_neighborhood = self.neighbors[channel_index]
                     channel_shift = self.operator_func(traces[:, channel_neighborhood], axis=1)
-                    re_referenced_traces[:, channel_index] = traces[:, channel_index] - channel_shift
+                    re_referenced_traces[:, i] = traces[:, channel_index] - channel_shift
 
             return re_referenced_traces.astype(self.dtype, copy=False)
 


Test patch:
diff --git a/src/spikeinterface/preprocessing/tests/test_common_reference.py b/src/spikeinterface/preprocessing/tests/test_common_reference.py
index e64775df2..1df9b21c8 100644
--- a/src/spikeinterface/preprocessing/tests/test_common_reference.py
+++ b/src/spikeinterface/preprocessing/tests/test_common_reference.py
@@ -10,7 +10,7 @@ from tqdm import trange
 
 
 def _generate_test_recording():
-    recording = generate_recording(durations=[5.0], num_channels=4)
+    recording = generate_recording(durations=[1.0], num_channels=4)
     recording = recording.channel_slice(recording.channel_ids, np.array(["a", "b", "c", "d"]))
     return recording
 
@@ -46,11 +46,15 @@ def test_common_reference(recording):
 def test_common_reference_channel_slicing(recording):
     recording_cmr = common_reference(recording, reference="global", operator="median")
     recording_car = common_reference(recording, reference="global", operator="average")
-    recording_single_reference = common_reference(recording, reference="single", ref_channel_ids=["a"])
+    recording_single_reference = common_reference(recording, reference="single", ref_channel_ids=["b"])
     recording_local_car = common_reference(recording, reference="local", local_radius=(20, 65), operator="median")
 
-    channel_ids = ["a", "b"]
-    indices = recording.ids_to_indices(["a", "b"])
+    channel_ids = ["b", "d"]
+    indices = recording.ids_to_indices(channel_ids)
+
+    all_channel_ids = recording.channel_ids
+    all_indices = recording.ids_to_indices(all_channel_ids)
+
     original_traces = recording.get_traces()
 
     cmr_trace = recording_cmr.get_traces(channel_ids=channel_ids)
@@ -62,13 +66,50 @@ def test_common_reference_channel_slicing(recording):
     assert np.allclose(car_trace, expected_trace, atol=0.01)
 
     single_reference_trace = recording_single_reference.get_traces(channel_ids=channel_ids)
-    single_reference_index = recording.ids_to_indices(["a"])
+    single_reference_index = recording.ids_to_indices(["b"])
     expected_trace = original_traces[:, indices] - original_traces[:, single_reference_index]
 
     assert np.allclose(single_reference_trace, expected_trace, atol=0.01)
 
-    # local car
-    local_trace = recording_local_car.get_traces(channel_ids=channel_ids)
+    local_trace = recording_local_car.get_traces(channel_ids=all_channel_ids)
+    local_trace_sub = recording_local_car.get_traces(channel_ids=channel_ids)
+
+    assert np.all(local_trace[:, indices] == local_trace_sub)
+
+    # test segment slicing
+
+    start_frame = 0
+    end_frame = 10
+
+    recording_segment_cmr = recording_cmr._recording_segments[0]
+    traces_cmr_all = recording_segment_cmr.get_traces(
+        start_frame=start_frame, end_frame=end_frame, channel_indices=all_indices
+    )
+    traces_cmr_sub = recording_segment_cmr.get_traces(
+        start_frame=start_frame, end_frame=end_frame, channel_indices=indices
+    )
+
+    assert np.all(traces_cmr_all[:, indices] == traces_cmr_sub)
+
+    recording_segment_car = recording_car._recording_segments[0]
+    traces_car_all = recording_segment_car.get_traces(
+        start_frame=start_frame, end_frame=end_frame, channel_indices=all_indices
+    )
+    traces_car_sub = recording_segment_car.get_traces(
+        start_frame=start_frame, end_frame=end_frame, channel_indices=indices
+    )
+
+    assert np.all(traces_car_all[:, indices] == traces_car_sub)
+
+    recording_segment_local = recording_local_car._recording_segments[0]
+    traces_local_all = recording_segment_local.get_traces(
+        start_frame=start_frame, end_frame=end_frame, channel_indices=all_indices
+    )
+    traces_local_sub = recording_segment_local.get_traces(
+        start_frame=start_frame, end_frame=end_frame, channel_indices=indices
+    )
+
+    assert np.all(traces_local_all[:, indices] == traces_local_sub)
 
 
 def test_common_reference_groups(recording):