Faster exact median for common_reference

src/spikeinterface/preprocessing/_row_median.py

@@ -0,0 +1,61 @@
+"""Exact per-row median of a ``(T, C)`` array, for common median referencing.
+
+``np.median(axis=1)`` partitions each row with data-dependent branches, mispredicting about one per value.
+Float rows are sorted instead, since ``np.sort`` vectorises 32- and 64-bit sorts on most current CPUs. Integer
+rows of at most 16 bits, whose sorts it vectorises only on some, go through the histogram kernel in
+``_row_median_numba``. Both equal ``np.median``, including for float rows holding NaN or inf.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+from functools import cache
+
+import numpy as np
+
+HAVE_NUMBA = importlib.util.find_spec("numba") is not None
+
+# Histogram bins; a row whose middle values fall outside the window is sorted instead.
+_INT_WINDOW = 1024
+# Below these channel counts, np.median's partition costs less than the kernel's per-row setup or the sort.
+_MIN_INT_CHANNELS = 4
+_MIN_FLOAT_CHANNELS = 3
+
+
+def row_median(traces: np.ndarray) -> np.ndarray:
+    """Median of each row of a 2D array, as ``(T, 1)`` in ``np.median``'s output dtype.
+
+    float32 and float64 input of at least 3 channels is sorted. Native-endian integer input of at most 16 bits
+    and at least 4 channels uses a numba kernel when numba is installed. Other input uses ``np.median``.
+    """
+    if traces.ndim == 2 and traces.shape[0] > 0:
+        num_channels = traces.shape[1]
+        if traces.dtype in (np.float32, np.float64) and num_channels >= _MIN_FLOAT_CHANNELS:
+            return _sorted_row_median(traces)
+        if (
+            HAVE_NUMBA
+            and traces.dtype.kind in "iu"
+            and traces.dtype.itemsize <= 2
+            and traces.dtype.isnative  # numba rejects byte-swapped arrays
+            and num_channels >= _MIN_INT_CHANNELS
+        ):
+            # The kernel reads column-major traces in place, which beats copying them.
+            return _get_int_kernel()(traces, _INT_WINDOW)[:, np.newaxis]
+    return np.median(traces, axis=1, keepdims=True)
+
+
+def _sorted_row_median(traces):
+    num_channels = traces.shape[1]
+    sorted_traces = np.sort(traces, axis=1)  # NaN sorts last
+    lo = sorted_traces[:, (num_channels - 1) // 2]
+    hi = sorted_traces[:, num_channels // 2]
+    median = 0.5 * (lo + hi)  # the scalar stays weak, so the sum and result keep the input dtype
+    median[np.isnan(sorted_traces[:, -1])] = np.nan  # np.median returns NaN for a row holding one
+    return median[:, np.newaxis]
+
+
+@cache
+def _get_int_kernel():
+    from ._row_median_numba import int_row_median
+
+    return int_row_median

src/spikeinterface/preprocessing/_row_median_numba.py

@@ -0,0 +1,71 @@
+"""numba kernel for ``_row_median``, kept at module level so that ``cache=True`` can find it again.
+
+``int_row_median`` returns exactly what ``np.median(x, axis=1)`` returns for integer ``x``. For a row of ``C``
+values, sorted, the median is the mean of the values at ranks ``r_lo = (C - 1) // 2`` and ``r_hi = C // 2``,
+which are the same rank when ``C`` is odd.
+
+The kernel counts each row into a histogram instead of sorting it. Filtered integer traces take few distinct
+values, and the median moves little from one sample to the next, so a window of bins centred on the previous
+row's median almost always holds both middle values. A cumulative sum over the bins then gives them directly.
+Rows whose middle values fall outside the window are sorted in full, so every path returns the exact median.
+"""
+
+import numpy as np
+from numba import njit
+
+
+@njit(nogil=True, cache=True)
+def int_row_median(x, window):
+    """Median of each row of integer ``x`` (at most 16 bits), as float64 like ``np.median``.
+
+    ``window`` is the number of histogram bins. Bin ``i`` counts the value ``base + i``, and ``base`` is chosen
+    per row so that the previous row's median sits in the middle bin.
+    """
+    T, C = x.shape
+    out = np.empty(T, dtype=np.float64)
+    counts = np.zeros(window, dtype=np.int32)
+    r_lo = (C - 1) // 2  # the two middle ranks, equal for odd C
+    r_hi = C // 2
+    center = np.int64(np.sort(x[0])[r_lo])
+    for t in range(T):
+        row = x[t]
+        base = center - window // 2
+        # Count every value into its bin. Values below the window are only counted, since the rank of anything
+        # inside the window depends on how many lie below it; values above it can be ignored.
+        below = 0
+        vmin = np.int64(window)  # lowest occupied bin, where the cumulative scan starts
+        for c in range(C):
+            i = np.int64(row[c]) - base
+            if i < 0:
+                below += 1
+            elif i < window:
+                counts[i] += 1
+                vmin = min(vmin, i)
+        # Walk the bins in value order, keeping a running count of values seen. The value at rank r is the
+        # first bin where that count exceeds r. If more than r_lo values lie below the window, the middle
+        # values are below it too and the scan is skipped; if the scan ends without reaching r_hi, they are above.
+        lo = -1
+        hi = -1
+        if below <= r_lo:
+            cum = below
+            for i in range(vmin, window):
+                cum += counts[i]
+                if lo < 0 and cum > r_lo:
+                    lo = i
+                if cum > r_hi:
+                    hi = i
+                    break
+        for c in range(C):  # clear only the bins this row touched, which is cheaper than zeroing the window
+            i = np.int64(row[c]) - base
+            if 0 <= i < window:
+                counts[i] = 0
+        if hi >= 0:
+            out[t] = 0.5 * (np.float64(lo + base) + np.float64(hi + base))
+            center = lo + base
+        else:
+            # A middle value fell outside the window, as after a large jump in the median: sort, and recentre
+            # the window on this row's median for the next row.
+            s = np.sort(row)
+            out[t] = 0.5 * (np.float64(s[r_lo]) + np.float64(s[r_hi]))
+            center = np.int64(s[r_lo])
+    return out

src/spikeinterface/preprocessing/common_reference.py

@@ -9,6 +9,7 @@
 from spikeinterface.core.baserecording import BaseRecording
 
 from .filter import fix_dtype
+from ._row_median import row_median
 
 
 class CommonReferenceRecording(BasePreprocessor):
@@ -244,7 +245,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
 
     def get_traces(self, start_frame, end_frame, channel_indices):
         # Let's do the case with group_indices equal None as that is easy
@@ -343,6 +344,12 @@ def slice_groups(self, channel_indices):
         return zip(group_indices, selected_channels, group_channels)
 
 
+def _median(traces, axis, keepdims=False):
+    assert axis == 1
+    median = row_median(traces)
+    return median if keepdims else median[:, 0]
+
+
 common_reference = define_function_handling_dict_from_class(
     source_class=CommonReferenceRecording, name="common_reference"
 )

src/spikeinterface/preprocessing/tests/test_row_median.py

@@ -0,0 +1,81 @@
+import numpy as np
+import pytest
+
+from spikeinterface.core import NumpyRecording
+from spikeinterface.preprocessing import common_reference
+from spikeinterface.preprocessing._row_median import HAVE_NUMBA, _INT_WINDOW, _get_int_kernel, row_median
+
+
+def assert_matches_np_median(x):
+    median = row_median(x)
+    expected = np.median(x, axis=1, keepdims=True)
+    assert median.dtype == expected.dtype
+    np.testing.assert_array_equal(median, expected)
+
+
+@pytest.mark.skipif(not HAVE_NUMBA, reason="integer rows fall back to np.median without numba")
+@pytest.mark.parametrize("dtype", ["int16", "uint16", "int8", "uint8"])
+@pytest.mark.parametrize("num_channels", [1, 2, 3, 384])
+def test_row_median_int(dtype, num_channels):
+    rng = np.random.default_rng(0)
+    info = np.iinfo(dtype)
+    full_range = rng.integers(info.min, info.max, (300, num_channels), endpoint=True).astype(dtype)
+    ties = rng.integers(max(info.min, -3), 4, (300, num_channels)).astype(dtype)
+    for x in (full_range, ties, full_range[:, ::-1], full_range[::2], np.asfortranarray(full_range)):
+        assert_matches_np_median(x)
+        # row_median sends few channels to np.median, so test the kernel there directly
+        np.testing.assert_array_equal(_get_int_kernel()(x, _INT_WINDOW), np.median(x, axis=1))
+
+
+@pytest.mark.parametrize("dtype", [">i2", ">u2", ">f4", "float16"])
+def test_row_median_other_dtypes(dtype):
+    x = np.random.default_rng(0).integers(0, 200, (300, 384)).astype(dtype)
+    assert_matches_np_median(x)
+
+
+@pytest.mark.skipif(not HAVE_NUMBA, reason="integer rows fall back to np.median without numba")
+def test_row_median_int_outside_window():
+    rng = np.random.default_rng(0)
+    x = rng.normal(0, 10, (1000, 384)).astype("int16")
+    x[:, 10] = np.iinfo("int16").max  # a saturated channel
+    x[:, 11] = np.iinfo("int16").min
+    x[500:] += 5000  # the median jumps beyond the histogram window
+    assert_matches_np_median(x)
+
+
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+@pytest.mark.parametrize("num_channels", [1, 2, 3, 4, 64, 384, 385])
+def test_row_median_float(dtype, num_channels):
+    rng = np.random.default_rng(0)
+    for x in (
+        rng.standard_normal((300, num_channels)),
+        rng.integers(-3, 4, (300, num_channels)),
+        np.zeros((300, num_channels)),
+    ):
+        x = x.astype(dtype)
+        assert_matches_np_median(x)
+        assert_matches_np_median(np.asfortranarray(x))
+        assert_matches_np_median(x[::2, ::-1])
+
+
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+def test_row_median_float_non_finite(dtype):
+    x = np.random.default_rng(0).standard_normal((6, 384)).astype(dtype)
+    x[0, 5] = np.nan
+    x[1, 7] = np.inf
+    x[2, 3] = -np.inf
+    x[3, [1, 2]] = [np.inf, -np.inf]
+    x[4] = np.finfo(dtype).max  # the sum of the middle values overflows
+    x[5, [0, 1]] = [np.nan, np.inf]
+    with np.errstate(over="ignore", invalid="ignore"):
+        assert_matches_np_median(x)
+
+
+@pytest.mark.parametrize("dtype", ["int16", "float32"])
+def test_common_reference_median_matches_np_median(dtype):
+    x = np.random.default_rng(0).normal(0, 20, (3000, 64)).astype(dtype)
+    recording = NumpyRecording(x, 30000.0)
+    expected = x - np.median(x, axis=1, keepdims=True)
+    if dtype == "int16":
+        expected = np.round(expected)
+    np.testing.assert_array_equal(common_reference(recording).get_traces(), expected.astype(dtype))