Speed up common median reference on channel-sized axes

src/spikeinterface/preprocessing/common_reference.py

@@ -11,6 +11,24 @@
 from .filter import fix_dtype
 
 
+def _median_reference(traces, axis=1, keepdims=False):
+    """Compute a median along a short channel axis using NumPy's faster sort path."""
+    num_channels = traces.shape[axis]
+    if traces.dtype.kind in "fiu" and num_channels in (1, 2):
+        return np.mean(traces, axis=axis, keepdims=keepdims)
+    if num_channels < 4 or traces.dtype.kind not in "fiu" or (traces.dtype.kind == "f" and np.isnan(traces).any()):
+        return np.median(traces, axis=axis, keepdims=keepdims)
+
+    sorted_traces = np.sort(traces, axis=axis)
+    midpoint = num_channels // 2
+    middle = [slice(None)] * traces.ndim
+    if num_channels % 2:
+        middle[axis] = slice(midpoint, midpoint + 1)
+    else:
+        middle[axis] = slice(midpoint - 1, midpoint + 1)
+    return np.mean(sorted_traces[tuple(middle)], axis=axis, keepdims=keepdims)
+
+
 class CommonReferenceRecording(BasePreprocessor):
     """
     Re-references the recording extractor traces. That is, the value of the traces are
@@ -244,7 +262,7 @@ def __init__(
         self.local_kernel = local_kernel
         self.temp = None
         self.dtype = dtype
-        self.operator_func = np.mean if self.operator == "average" else np.median
+        self.operator_func = np.mean if self.operator == "average" else _median_reference
 
     def get_traces(self, start_frame, end_frame, channel_indices):
         # Let's do the case with group_indices equal None as that is easy

src/spikeinterface/preprocessing/tests/test_common_reference.py

@@ -20,6 +20,51 @@ def recording():
     return _generate_test_recording()
 
 
+def test_median_reference_matches_numpy():
+    from spikeinterface.preprocessing.common_reference import _median_reference
+
+    rng = np.random.default_rng(0)
+    for dtype in ("int16", "float32", "float64"):
+        for num_channels in (1, 2, 3, 4, 5):
+            traces = rng.integers(-100, 100, size=(101, num_channels)).astype(dtype)
+            for keepdims in (False, True):
+                expected = np.median(traces, axis=1, keepdims=keepdims)
+                result = _median_reference(traces, axis=1, keepdims=keepdims)
+
+                assert result.dtype == expected.dtype
+                assert np.array_equal(result, expected)
+
+
+@pytest.mark.parametrize("num_channels", (2, 3, 4, 5))
+def test_median_reference_special_values(num_channels):
+    from spikeinterface.preprocessing.common_reference import _median_reference
+
+    nan_traces = np.array([[1.0, np.nan, 3.0, 4.0, 5.0], [6.0, 7.0, 8.0, 9.0, 10.0]], dtype="float32")
+    sliced = nan_traces[:, :num_channels]
+    expected = np.median(sliced, axis=1, keepdims=True)
+    result = _median_reference(sliced, axis=1, keepdims=True)
+    assert np.array_equal(result, expected, equal_nan=True)
+
+    special_traces = np.array([[-0.0, 0.0, -np.inf, np.inf, 1.0], [-np.inf, -0.0, 0.0, np.inf, -1.0]], dtype="float32")
+    sliced = special_traces[:, :num_channels]
+    with np.errstate(invalid="ignore"):
+        expected = np.median(sliced, axis=1, keepdims=True)
+        result = _median_reference(sliced, axis=1, keepdims=True)
+    if num_channels % 2:
+        assert np.array_equal(result, expected)
+    else:
+        assert result.tobytes() == expected.tobytes()
+
+    integer_limits = np.array(
+        [[np.iinfo(np.int64).max, np.iinfo(np.int64).max, np.iinfo(np.int64).min, np.iinfo(np.int64).min, 0]],
+        dtype="int64",
+    )
+    sliced = integer_limits[:, :num_channels]
+    expected = np.median(sliced, axis=1, keepdims=True)
+    result = _median_reference(sliced, axis=1, keepdims=True)
+    assert result.tobytes() == expected.tobytes()
+
+
 def test_common_reference(recording, create_cache_folder):
     # Test simple case
     rec_cmr = common_reference(recording, reference="global", operator="median")