diff --git a/src/spikeinterface/preprocessing/_row_median.py b/src/spikeinterface/preprocessing/_row_median.py new file mode 100644 index 0000000000..1d740830ce --- /dev/null +++ b/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 diff --git a/src/spikeinterface/preprocessing/_row_median_numba.py b/src/spikeinterface/preprocessing/_row_median_numba.py new file mode 100644 index 0000000000..73b1dd0e4d --- /dev/null +++ b/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 diff --git a/src/spikeinterface/preprocessing/common_reference.py b/src/spikeinterface/preprocessing/common_reference.py index 6da823d0e4..e601bcf2ca 100644 --- a/src/spikeinterface/preprocessing/common_reference.py +++ b/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" ) diff --git a/src/spikeinterface/preprocessing/tests/test_row_median.py b/src/spikeinterface/preprocessing/tests/test_row_median.py new file mode 100644 index 0000000000..eae228f40c --- /dev/null +++ b/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))