Handle empty annulus in local common reference

src/spikeinterface/preprocessing/common_reference.py

@@ -78,6 +78,8 @@ class CommonReferenceRecording(BasePreprocessor):
     min_local_neighbors : int, default: 5
         Use in the local CAR implementation to set a minimum number of neighbors. If the number of neighbors within the
         annulus is less than this number, then the closest neighbors are used until this number is reached.
+        If a channel has no other channel beyond the exclude radius at all, it is left unreferenced and a
+        warning is raised at init.
     dtype : None or dtype, default: None
         If None the parent dtype is kept.
 
@@ -154,21 +156,33 @@ def __init__(
             # For the median operator, the neighbors are extracted from the kernel on-the-fly via nonzero.
             local_kernel = np.zeros((num_chans, num_chans))
             not_enough_channels = []
+            no_neighbor_channels = []
             for i in range(num_chans):
                 annulus_mask = (dist[i, :] > local_radius[0]) & (dist[i, :] <= local_radius[1])
                 if np.sum(annulus_mask) >= min_local_neighbors:
                     neighbors_i = closest_inds[i, annulus_mask]
                 else:
-                    # Not enough channels in the annulus — take the closest ones beyond the inner radius
-                    not_enough_channels.append(str(recording.channel_ids[i]))
+                    # Not enough channels in the annulus, take the closest ones beyond the inner radius
                     beyond_inner = dist[i, :] > local_radius[0]
                     neighbors_i = closest_inds[i, beyond_inner][:min_local_neighbors]
+                    if len(neighbors_i) == 0:
+                        # No channel at all lies beyond the inner radius, so there is nothing to build a
+                        # reference from. The kernel row is left at zero, which keeps the channel unreferenced.
+                        no_neighbor_channels.append(str(recording.channel_ids[i]))
+                        continue
+                    not_enough_channels.append(str(recording.channel_ids[i]))
                 local_kernel[i, neighbors_i] = 1 / len(neighbors_i)
             if len(not_enough_channels) > 0:
                 warnings.warn(
                     f"The following channels did not have enough neighbors in the annulus and used the closest "
                     f"{min_local_neighbors} channels beyond the inner radius instead: {', '.join(not_enough_channels)}"
                 )
+            if len(no_neighbor_channels) > 0:
+                warnings.warn(
+                    f"The following channels have no channel beyond the exclude radius of "
+                    f"{local_radius[0]} and are left unreferenced: {', '.join(no_neighbor_channels)}. "
+                    f"Consider lowering the exclude radius of 'local_radius'."
+                )
         dtype_ = fix_dtype(recording, dtype)
         BasePreprocessor.__init__(self, recording, dtype=dtype_)
 
@@ -254,6 +268,10 @@ def get_traces(self, start_frame, end_frame, channel_indices):
                     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 = np.nonzero(self.local_kernel[channel_index])[0]
+                        if channel_neighborhood.size == 0:
+                            # No neighbor was found at init time, the channel is left unreferenced
+                            re_referenced_traces[:, i] = traces[:, channel_index]
+                            continue
                         channel_shift = self.operator_func(traces[:, channel_neighborhood], axis=1)
                         re_referenced_traces[:, i] = traces[:, channel_index] - channel_shift
                 else:  # then it must be local average, use local_kernel

src/spikeinterface/preprocessing/tests/test_common_reference.py

@@ -249,6 +249,54 @@ def test_min_local_radius():
         )
 
 
+@pytest.mark.parametrize("operator", ["median", "average"])
+def test_local_reference_no_channel_beyond_exclude_radius(operator):
+    # When no channel at all lies beyond the exclude (inner) radius, there is nothing to build a
+    # local reference from. This used to raise a ZeroDivisionError at init. The channels should now
+    # be left unreferenced and a warning should be raised instead.
+    recording = generate_recording(durations=[1.0], num_channels=4, set_probe=False)
+    recording = recording.rename_channels(np.array(["a", "b", "c", "d"]))
+    # Tetrode-like geometry: every channel is within 30 um of every other one
+    recording.set_dummy_probe_from_locations(np.array([[0.0, 0.0], [0.0, 20.0], [20.0, 0.0], [20.0, 20.0]]))
+
+    with pytest.warns(UserWarning, match="left unreferenced"):
+        rec_local = common_reference(recording, reference="local", local_radius=(30.0, 55.0), operator=operator)
+
+    traces = recording.get_traces()
+    assert np.array_equal(rec_local.get_traces(), traces)
+
+
+@pytest.mark.parametrize("operator", ["median", "average"])
+def test_local_reference_partially_empty_annulus(operator):
+    # Only some channels have no channel beyond the exclude radius. Those are left unreferenced while
+    # the others are referenced normally.
+    recording = generate_recording(durations=[1.0], num_channels=5, set_probe=False)
+    recording = recording.rename_channels(np.array(["a", "b", "c", "d", "e"]))
+    # "a" sits at the center of a cross, so the farthest channel from it is 100 um away. Every other
+    # channel has the opposite arm of the cross 200 um away.
+    locations = np.array([[0.0, 0.0], [-100.0, 0.0], [100.0, 0.0], [0.0, -100.0], [0.0, 100.0]])
+    recording.set_dummy_probe_from_locations(locations)
+
+    with pytest.warns(UserWarning, match="left unreferenced"):
+        rec_local = common_reference(
+            recording,
+            reference="local",
+            local_radius=(150.0, 400.0),
+            operator=operator,
+            min_local_neighbors=1,
+        )
+
+    traces = recording.get_traces()
+    referenced_traces = rec_local.get_traces()
+
+    # "a" has no channel beyond 150 um, so it is passed through unchanged
+    assert np.array_equal(referenced_traces[:, 0], traces[:, 0])
+    # the four arms are each referenced to the single opposite arm, which is 200 um away
+    for channel_index, opposite_index in ((1, 2), (2, 1), (3, 4), (4, 3)):
+        expected = traces[:, channel_index] - traces[:, opposite_index]
+        assert np.allclose(referenced_traces[:, channel_index], expected, atol=1e-5)
+
+
 @pytest.mark.skip(reason="This test can be used to check local CAR vs local CMR performance")
 def test_local_car_vs_cmr_performance():
     import time