From 41094c112080079851d7e01d195eff1c556debb7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20Royeth?= Date: Tue, 6 Oct 2026 16:49:52 -0300 Subject: [PATCH] Speed up common median reference on channel-sized axes --- .../preprocessing/common_reference.py | 20 ++++++++- .../tests/test_common_reference.py | 45 +++++++++++++++++++ 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/src/spikeinterface/preprocessing/common_reference.py b/src/spikeinterface/preprocessing/common_reference.py index 6da823d0e4..b481a17a71 100644 --- a/src/spikeinterface/preprocessing/common_reference.py +++ b/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 diff --git a/src/spikeinterface/preprocessing/tests/test_common_reference.py b/src/spikeinterface/preprocessing/tests/test_common_reference.py index 139213056f..808f7517e0 100644 --- a/src/spikeinterface/preprocessing/tests/test_common_reference.py +++ b/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")