Repository: SpikeInterface/spikeinterface
Pull request #2362 (2023-12-20 12:01:19)

Fix overflow problems with CAR
This should fix https://github.com/SpikeInterface/spikeinterface/issues/2312.

The following code (reported by @JoeZiminski ) fails on main and works with this branch:

```python
import numpy as np
import spikeinterface.preprocessing as spre
from spikeinterface.core.generate import generate_recording

# Generate a recording, int16 with positive Vm (typical for NP recording)
num_channels = 384
recording = generate_recording(num_channels=num_channels, set_probe=True, durations=[0.5])
recording = spre.scale(recording, gain=100, offset=70)  # reduce offset ~ 30 to pass assert
recording = spre.astype(recording, "int16")


# Preprocess in SI
common_ref = spre.common_reference(recording, operator="average")
common_ref_data = common_ref.get_traces()

# Perform CAR with np.mean() casting to float64 under the hood.
recording_data = recording.get_traces()
means = np.mean(recording_data, axis=1).astype(np.int16)  # with dtype=np.int16 assert passes
test_common_ref_data = recording_data - means[:, np.newaxis]

assert np.allclose(common_ref_data, test_common_ref_data, rtol=0, atol=1)

# need to comment out the assert to run this block under failing conditions 
bad_idx = np.where(~np.isclose(common_ref_data, test_common_ref_data, rtol=0, atol=1))
if np.any(bad_idx):
    true_mean = means[bad_idx[0][0]]
    bad_mean = common_ref._recording_segments[0].operator_func(recording_data)[bad_idx[0][0]]
    print(true_mean)
    print(bad_mean)
```

The problem I think is some magic with the operators for the grouping part. To solve this issue I separated the implementation for the non-grouping case from the one for grouping. The former is easier to implement and avoids the implicit casting with the current operator application. I also added some tests to see that the slicing with `channel_ids` work as intended. That said, the bug is still not solved for the grouping case but I think we should solve this for the most common use case and then address that when we have more time and the need if that ever comes to be.


Patch:
diff --git a/src/spikeinterface/preprocessing/common_reference.py b/src/spikeinterface/preprocessing/common_reference.py
index 75926118d..07920b61c 100644
--- a/src/spikeinterface/preprocessing/common_reference.py
+++ b/src/spikeinterface/preprocessing/common_reference.py
@@ -51,9 +51,9 @@ class CommonReferenceRecording(BasePreprocessor):
         List of lists containing the channel ids for splitting the reference. The CMR, CAR, or referencing with respect to
         single channels are applied group-wise. However, this is not applied for the local CAR.
         It is useful when dealing with different channel groups, e.g. multiple tetrodes.
-    ref_channel_ids: list or int, default: None
+    ref_channel_ids: list or str or int, default: None
         If no "groups" are specified, all channels are referenced to "ref_channel_ids". If "groups" is provided, then a
-        list of channels to be applied to each group is expected. If "single" reference, a list of one channel  or an
+        list of channels to be applied to each group is expected. If "single" reference, a list of one channel or an
         int is expected.
     local_radius: tuple(int, int), default: (30, 55)
         Use in the local CAR implementation as the selecting annulus with the following format:
@@ -80,7 +80,7 @@ class CommonReferenceRecording(BasePreprocessor):
     def __init__(
         self,
         recording: BaseRecording,
-        reference: Literal["global", "single", "global"] = "global",
+        reference: Literal["global", "single", "local"] = "global",
         operator: Literal["median", "average"] = "median",
         groups=None,
         ref_channel_ids=None,
@@ -107,10 +107,10 @@ class CommonReferenceRecording(BasePreprocessor):
                 else:
                     assert (
                         len(ref_channel_ids) == 1
-                    ), "'ref_channel_ids' with no 'groups' must be int or a list of one element"
+                    ), "'ref_channel_ids' with no 'groups' must be a single channel id or a list of one element"
                 ref_channel_ids = np.asarray(ref_channel_ids)
                 assert np.all(
-                    [ch in recording.get_channel_ids() for ch in ref_channel_ids]
+                    [ch in recording.channel_ids for ch in ref_channel_ids]
                 ), "Some 'ref_channel_ids' are wrong!"
         elif reference == "local":
             assert groups is None, "With 'local' CAR, the group option should not be used."
@@ -173,57 +173,86 @@ class CommonReferenceRecordingSegment(BasePreprocessorSegment):
         self.neighbors = neighbors
         self.temp = None
         self.dtype = dtype
-
-        if self.operator == "median":
-            self.operator_func = lambda x: np.median(x, axis=1, out=self.temp)[:, None]
-        elif self.operator == "average":
-            self.operator_func = lambda x: np.mean(x, axis=1, out=self.temp)[:, None]
+        self.operator_func = operator = np.mean if self.operator == "average" else np.median
 
     def get_traces(self, start_frame, end_frame, channel_indices):
-        # need input trace
-        all_traces = self.parent_recording_segment.get_traces(start_frame, end_frame, slice(None))
-        all_traces = all_traces.astype(self.dtype)
-        self.temp = np.zeros((all_traces.shape[0],), dtype=all_traces.dtype)
-        _channel_indices = np.arange(all_traces.shape[1])
-        if channel_indices is not None:
-            _channel_indices = _channel_indices[channel_indices]
-
-        out_traces = np.zeros((all_traces.shape[0], _channel_indices.size), dtype=self.dtype)
-
-        if self.reference == "global":
-            for chan_inds, chan_group_inds in self._groups(_channel_indices):
-                out_inds = np.array([np.where(_channel_indices == i)[0][0] for i in chan_inds])
-                out_traces[:, out_inds] = all_traces[:, chan_inds] - self.operator_func(all_traces[:, chan_group_inds])
-
-        elif self.reference == "single":
-            for i, (chan_inds, _) in enumerate(self._groups(_channel_indices)):
-                out_inds = np.array([np.where(_channel_indices == i)[0][0] for i in chan_inds])
-                out_traces[:, out_inds] = all_traces[:, chan_inds] - self.operator_func(
-                    all_traces[:, [self.ref_channel_indices[i]]]
-                )
-
-        elif self.reference == "local":
-            for i, chan_ind in enumerate(_channel_indices):
-                out_traces[:, [i]] = all_traces[:, [chan_ind]] - self.operator_func(
-                    all_traces[:, self.neighbors[chan_ind]]
-                )
-        return out_traces
-
-    def _groups(self, channel_indices):
-        selected_groups = []
-        selected_channels = []
-        if self.group_indices:
-            for chan_inds in self.group_indices:
-                sel_inds = [ind for ind in channel_indices if ind in chan_inds]
-                # if no channels are in a group, do not return the group
-                if len(sel_inds) > 0:
-                    selected_groups.append(chan_inds)
-                    selected_channels.append(sel_inds)
+        # Let's do the case with group_indices equal None as that is easy
+        if self.group_indices is None:
+            # We need all the channels to calculate the reference
+            traces = self.parent_recording_segment.get_traces(start_frame, end_frame, slice(None))
+
+            if self.reference == "global":
+                shift = self.operator_func(traces, axis=1, keepdims=True)
+                re_referenced_traces = traces[:, channel_indices] - shift
+            elif self.reference == "single":
+                # single channel -> no need of operator
+                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:
+                    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
+
+            return re_referenced_traces.astype(self.dtype, copy=False)
+
+        # Then the old implementation for backwards compatibility that supports grouping
         else:
-            # no groups = all channels
-            selected_groups = [slice(None)]
-            selected_channels = [channel_indices]
-        return zip(selected_channels, selected_groups)
+            # need input trace
+            traces = self.parent_recording_segment.get_traces(start_frame, end_frame, slice(None))
+
+            sliced_channel_indices = np.arange(traces.shape[1])
+            if channel_indices is not None:
+                sliced_channel_indices = sliced_channel_indices[channel_indices]
+
+            re_referenced_traces = np.zeros((traces.shape[0], sliced_channel_indices.size))
+            for group_index, selected_indices_in_group, all_group_indices in self.slice_groups(sliced_channel_indices):
+                (out_indices,) = np.nonzero(np.isin(sliced_channel_indices, selected_indices_in_group))
+                in_group_traces = traces[:, selected_indices_in_group]
+
+                if self.reference == "global":
+                    shift = self.operator_func(traces[:, all_group_indices], axis=1, keepdims=True)
+                    re_referenced_traces[:, out_indices] = in_group_traces - shift
+                else:
+                    # single (as local is not allowed for groups)
+                    shift = self.operator_func(
+                        traces[:, [self.ref_channel_indices[group_index]]], axis=1, keepdims=True
+                    )
+                    re_referenced_traces[:, out_indices] = in_group_traces - shift
+
+            return re_referenced_traces.astype(self.dtype, copy=False)
+
+    def slice_groups(self, channel_indices):
+        """
+        Slice the channel indices into groups. This is used to apply the common reference to groups of channels.
+
+        Parameters
+        ----------
+        channel_indices: array-like
+            The channel indices to be sliced
+
+        Returns
+        -------
+        zip with:
+            * group_index: The index of the group
+            * selected_channels: The selected channel indices in the group
+            * group_channels: The channels indices in the group
+        """
+        selected_channels = []
+        group_channels = []
+        group_indices = []
+
+        assert self.group_indices is not None, "No groups to slice"
+        for group_index, chanel_indices in enumerate(self.group_indices):
+            selected_indices = [ind for ind in channel_indices if ind in chanel_indices]
+            # if no channels are in a group, do not return the group
+            if len(selected_indices) > 0:
+                group_channels.append(chanel_indices)
+                selected_channels.append(selected_indices)
+                group_indices.append(group_index)
+        return zip(group_indices, selected_channels, group_channels)
 
 
 common_reference = define_function_from_class(source_class=CommonReferenceRecording, name="common_reference")


Test patch:
diff --git a/src/spikeinterface/preprocessing/tests/test_common_reference.py b/src/spikeinterface/preprocessing/tests/test_common_reference.py
index 8032dc2f1..e64775df2 100644
--- a/src/spikeinterface/preprocessing/tests/test_common_reference.py
+++ b/src/spikeinterface/preprocessing/tests/test_common_reference.py
@@ -1,38 +1,33 @@
+import operator
 import pytest
-from pathlib import Path
-import shutil
 
-from spikeinterface import set_global_tmp_folder
 from spikeinterface.core import generate_recording
 
-from spikeinterface.preprocessing import CommonReferenceRecording, common_reference
+from spikeinterface.preprocessing import common_reference
 
 import numpy as np
+from tqdm import trange
 
-if hasattr(pytest, "global_test_folder"):
-    cache_folder = pytest.global_test_folder / "preprocessing"
-else:
-    cache_folder = Path("cache_folder") / "preprocessing"
 
-set_global_tmp_folder(cache_folder)
+def _generate_test_recording():
+    recording = generate_recording(durations=[5.0], num_channels=4)
+    recording = recording.channel_slice(recording.channel_ids, np.array(["a", "b", "c", "d"]))
+    return recording
 
 
-def test_common_reference():
-    rec = generate_recording(durations=[5.0], num_channels=4)
-    rec._main_ids = np.array(["a", "b", "c", "d"])
+@pytest.fixture
+def recording():
+    return _generate_test_recording()
 
-    # no groups
-    rec_cmr = common_reference(rec, reference="global", operator="median")
-    rec_car = common_reference(rec, reference="global", operator="average")
-    rec_sin = common_reference(rec, reference="single", ref_channel_ids=["a"])
-    rec_local_car = common_reference(rec, reference="local", local_radius=(20, 65), operator="median")
 
-    rec_cmr.save(verbose=False)
-    rec_car.save(verbose=False)
-    rec_sin.save(verbose=False)
-    rec_local_car.save(verbose=False)
+def test_common_reference(recording):
+    # Test simple case
+    rec_cmr = common_reference(recording, reference="global", operator="median")
+    rec_car = common_reference(recording, reference="global", operator="average")
+    rec_sin = common_reference(recording, reference="single", ref_channel_ids=["a"])
+    rec_local_car = common_reference(recording, reference="local", local_radius=(20, 65), operator="median")
 
-    traces = rec.get_traces()
+    traces = recording.get_traces()
     assert np.allclose(traces, rec_cmr.get_traces() + np.median(traces, axis=1, keepdims=True), atol=0.01)
     assert np.allclose(traces, rec_car.get_traces() + np.mean(traces, axis=1, keepdims=True), atol=0.01)
     assert not np.all(rec_sin.get_traces()[0])
@@ -41,6 +36,86 @@ def test_common_reference():
     assert np.allclose(traces[:, 0], rec_local_car.get_traces()[:, 0] + np.median(traces[:, [2, 3]], axis=1), atol=0.01)
     assert np.allclose(traces[:, 1], rec_local_car.get_traces()[:, 1] + np.median(traces[:, [3]], axis=1), atol=0.01)
 
+    # Saving tests
+    rec_cmr.save(verbose=False)
+    rec_car.save(verbose=False)
+    rec_sin.save(verbose=False)
+    rec_local_car.save(verbose=False)
+
+
+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_local_car = common_reference(recording, reference="local", local_radius=(20, 65), operator="median")
+
+    channel_ids = ["a", "b"]
+    indices = recording.ids_to_indices(["a", "b"])
+    original_traces = recording.get_traces()
+
+    cmr_trace = recording_cmr.get_traces(channel_ids=channel_ids)
+    expected_trace = original_traces[:, indices] - np.median(original_traces, axis=1, keepdims=True)
+    assert np.allclose(cmr_trace, expected_trace, atol=0.01)
+
+    car_trace = recording_car.get_traces(channel_ids=channel_ids)
+    expected_trace = original_traces[:, indices] - np.mean(original_traces, axis=1, keepdims=True)
+    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"])
+    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)
+
+
+def test_common_reference_groups(recording):
+    original_traces = recording.get_traces()
+    groups = [["a", "c", "d"], ["b"]]
+
+    # median - global
+    recording_cmr = common_reference(recording, reference="global", operator="median", groups=groups)
+    traces = recording_cmr.get_traces(channel_ids=["c", "b", "d"])
+    # c/d
+    assert np.allclose(
+        traces[:, 0], original_traces[:, 2] - np.median(original_traces[:, [0, 2, 3]], axis=1), atol=0.01
+    )
+    assert np.allclose(
+        traces[:, 2], original_traces[:, 3] - np.median(original_traces[:, [0, 2, 3]], axis=1), atol=0.01
+    )
+    # b will be all zeros
+    assert np.allclose(traces[:, 1], 0)
+
+    traces = recording_cmr.get_traces(channel_ids=["a"])
+    assert np.allclose(
+        traces[:, 0], original_traces[:, 0] - np.median(original_traces[:, [0, 2, 3]], axis=1), atol=0.01
+    )
+
+    # single
+    groups = [["a", "c"], ["b", "d"]]
+    recording_single = common_reference(recording, reference="single", ref_channel_ids=["a", "b"], groups=groups)
+    traces = recording_single.get_traces(channel_ids=["c", "b", "a", "d"])
+    # c - referenced to a
+    assert np.allclose(traces[:, 0], original_traces[:, 2] - original_traces[:, 0], atol=0.01)
+    # b - all zeros
+    assert np.allclose(traces[:, 1], 0)
+    # a - all zeros
+    assert np.allclose(traces[:, 2], 0)
+    # d - referenced to b
+    assert np.allclose(traces[:, 3], original_traces[:, 3] - original_traces[:, 1], atol=0.01)
+
+    # TODO: fix this!
+    traces = recording_single.get_traces(channel_ids=["d", "b"])
+    # d - referenced to b
+    assert np.allclose(traces[:, 0], original_traces[:, 3] - original_traces[:, 1], atol=0.01)
+    # b - all zeros
+    assert np.allclose(traces[:, 1], 0)
+
 
 if __name__ == "__main__":
-    test_common_reference()
+    recording = _generate_test_recording()
+    test_common_reference(recording)
+    test_common_reference_channel_slicing(recording)
+    test_common_reference_groups(recording)