From a8b634cce149ed1f15dd229c9cf1c44a05fb018e Mon Sep 17 00:00:00 2001 From: Galen Lynch Date: Mon, 12 Oct 2026 13:56:11 -0700 Subject: [PATCH 1/5] perf: exact numba row median for common_reference np.median partitions each row with data-dependent branches and mispredicts about once per value. row_median gives the same values and dtype: int16 rows of 4 or more channels are counted into a histogram around the previous row's median, and float32 rows of 64 or more channels are narrowed with branchless counting passes. Other inputs, and installs without numba, keep np.median. On 1 s of 384 column-major channels, as filter steps return them, the median drops from 72 to 11 ms for int16 and from 101 to 38 ms for float32. --- .../preprocessing/_row_median.py | 53 +++++++ .../preprocessing/_row_median_numba.py | 143 ++++++++++++++++++ .../preprocessing/common_reference.py | 9 +- .../preprocessing/tests/test_row_median.py | 79 ++++++++++ 4 files changed, 283 insertions(+), 1 deletion(-) create mode 100644 src/spikeinterface/preprocessing/_row_median.py create mode 100644 src/spikeinterface/preprocessing/_row_median_numba.py create mode 100644 src/spikeinterface/preprocessing/tests/test_row_median.py diff --git a/src/spikeinterface/preprocessing/_row_median.py b/src/spikeinterface/preprocessing/_row_median.py new file mode 100644 index 0000000000..c896072563 --- /dev/null +++ b/src/spikeinterface/preprocessing/_row_median.py @@ -0,0 +1,53 @@ +"""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. +Integer rows are counted into a histogram over a window around the previous row's median, since filtered +integer traces repeat few values and adjacent samples have nearly the same median. float32 rows narrow a value +interval around the two middle ranks with branchless counting passes that vectorise over the row, seeded at the +row mean and stepped by interpolating the row's empirical CDF, then sort the few values left in the interval. +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 the kernels' per-row setup costs more than np.median's partition. +_MIN_INT_CHANNELS = 4 +_MIN_FLOAT_CHANNELS = 64 + + +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. + + Integer input of at most 16 bits and at least 4 channels, and float32 input of at least 64 channels, use numba + kernels when numba is installed; other input uses ``np.median``. + """ + if HAVE_NUMBA and traces.ndim == 2 and traces.shape[0] > 0 and traces.shape[1] > 0: + if traces.dtype == np.float32 and traces.shape[1] >= _MIN_FLOAT_CHANNELS: + return _get_float_kernel()(np.ascontiguousarray(traces))[:, np.newaxis] + if traces.dtype.kind in "iu" and traces.dtype.itemsize <= 2 and traces.shape[1] >= _MIN_INT_CHANNELS: + # Reading column-major traces in place beats copying them; the float32 kernel is the other way round. + return _get_int_kernel()(traces, _INT_WINDOW)[:, np.newaxis] + return np.median(traces, axis=1, keepdims=True) + + +@cache +def _get_int_kernel(): + from ._row_median_numba import int_row_median + + return int_row_median + + +@cache +def _get_float_kernel(): + from ._row_median_numba import float_row_median + + return float_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..418b13d0fe --- /dev/null +++ b/src/spikeinterface/preprocessing/_row_median_numba.py @@ -0,0 +1,143 @@ +"""numba kernels for ``_row_median``, kept at module level so that ``cache=True`` can find them again.""" + +import numpy as np +from numba import njit + +# The interval is sorted once it holds this many values; smaller needs more counting passes than it saves. +_FINAL_WIDTH = 24 +_MAX_PASSES = 48 + + +@njit(nogil=True, cache=True) +def int_row_median(x, window): + 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 + below = 0 + vmin = np.int64(window) + 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) + 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 + 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: + 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 + + +@njit(inline="always") +def _counts(row, v): + lt = 0 + le = 0 + for c in range(row.size): + lt += row[c] < v + le += row[c] <= v + return lt, le + + +@njit(inline="always") +def _sorted_median(row, scratch, r_lo, r_hi): + scratch[:] = row + scratch.sort() # NaN sorts last + if np.isnan(scratch[-1]): + return scratch[-1] + return np.float32(0.5) * (scratch[r_lo] + scratch[r_hi]) + + +@njit(inline="always") +def _median_of_row(row, buf, scratch): + C = row.size + r_lo = (C - 1) // 2 # the two middle ranks, equal for odd C + r_hi = C // 2 + total = np.float32(0.0) + vmin = row[0] + vmax = row[0] + for c in range(C): + v = row[c] + total += v + vmin = min(vmin, v) + vmax = max(vmax, v) + if not np.isfinite(total): # NaN, inf, or a sum that overflowed + return _sorted_median(row, scratch, r_lo, r_hi) + # Invariant: count(< a) == ca <= r_lo and count(<= b) == cb >= r_hi + 1. + a = vmin + ca = 0 + b = vmax + cb = C + p = total / C + target = 0.5 * (r_lo + r_hi + 1) + for _ in range(_MAX_PASSES): + if cb - ca <= _FINAL_WIDTH: + break + lt, le = _counts(row, p) + if lt <= r_lo and le >= r_hi + 1: + return p # p is a value holding both middle ranks + if lt > r_lo and le < r_hi + 1: + # p sits strictly between the two middle values, so they are its nearest neighbours + below = vmin + above = vmax + for c in range(C): + v = row[c] + below = max(below, v if v < p else vmin) + above = min(above, v if v > p else vmax) + return np.float32(0.5) * (below + above) + if lt <= r_lo: + a = p + ca = lt + else: + b = p + cb = le + q = a + (target - ca) / max(cb - ca, 1) * (b - a) + if not (a < q < b): + q = 0.5 * (a + b) + p = np.float32(q) + if cb - ca > _FINAL_WIDTH: + return _sorted_median(row, scratch, r_lo, r_hi) + n = 0 + for c in range(C): # branchless gather of the values in [a, b] + v = row[c] + buf[n] = v + n += (v >= a) & (v <= b) + middle = buf[:n] + middle.sort() + return np.float32(0.5) * (middle[r_lo - ca] + middle[r_hi - ca]) + + +# Reassociation only reorders the sum that seeds the search, so the result stays exact. +@njit(nogil=True, cache=True, fastmath={"reassoc"}) +def float_row_median(x): + T, C = x.shape + out = np.empty(T, dtype=np.float32) + buf = np.empty(C + 1, dtype=np.float32) + scratch = np.empty(C, dtype=np.float32) + for t in range(T): + out[t] = _median_of_row(x[t], buf, scratch) + 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..ccfa05e89d --- /dev/null +++ b/src/spikeinterface/preprocessing/tests/test_row_median.py @@ -0,0 +1,79 @@ +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_float_kernel, + _get_int_kernel, + row_median, +) + +pytestmark = pytest.mark.skipif(not HAVE_NUMBA, reason="row_median falls back to np.median without numba") + + +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.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)) + + +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("num_channels", [1, 2, 3, 63, 64, 384, 385]) +def test_row_median_float32(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("float32") + assert_matches_np_median(x) + assert_matches_np_median(np.asfortranarray(x)) + # row_median sends few channels to np.median, so test the kernel there directly + np.testing.assert_array_equal(_get_float_kernel()(x), np.median(x, axis=1)) + + +def test_row_median_float32_non_finite(): + x = np.random.default_rng(0).standard_normal((5, 384)).astype("float32") + 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] = 3e38 # the row sum overflows + with np.errstate(over="ignore"): + np.testing.assert_array_equal(row_median(x), np.median(x, axis=1, keepdims=True)) + + +@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)) From ff73ac5e2694b8f1d7cf85501d04308ebadf5763 Mon Sep 17 00:00:00 2001 From: Galen Lynch Date: Tue, 13 Oct 2026 07:02:01 -0700 Subject: [PATCH 2/5] docs: explain the row median algorithms in comments The histogram and interval-narrowing kernels now state their invariants, early exits and fallbacks, so a reader can check them without reconstructing the algorithm. No code changes. --- .../preprocessing/_row_median.py | 7 +- .../preprocessing/_row_median_numba.py | 64 ++++++++++++++++--- 2 files changed, 57 insertions(+), 14 deletions(-) diff --git a/src/spikeinterface/preprocessing/_row_median.py b/src/spikeinterface/preprocessing/_row_median.py index c896072563..6de3d1d140 100644 --- a/src/spikeinterface/preprocessing/_row_median.py +++ b/src/spikeinterface/preprocessing/_row_median.py @@ -1,11 +1,8 @@ """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. -Integer rows are counted into a histogram over a window around the previous row's median, since filtered -integer traces repeat few values and adjacent samples have nearly the same median. float32 rows narrow a value -interval around the two middle ranks with branchless counting passes that vectorise over the row, seeded at the -row mean and stepped by interpolating the row's empirical CDF, then sort the few values left in the interval. -Both equal ``np.median``, including for float rows holding NaN or inf. +The numba kernels in ``_row_median_numba`` replace most of those branches with counting: integer rows go through +a histogram, float32 rows through counting passes that vectorise. That module describes both algorithms. Both equal ``np.median``, including for float rows holding NaN or inf. """ from __future__ import annotations diff --git a/src/spikeinterface/preprocessing/_row_median_numba.py b/src/spikeinterface/preprocessing/_row_median_numba.py index 418b13d0fe..94881d5534 100644 --- a/src/spikeinterface/preprocessing/_row_median_numba.py +++ b/src/spikeinterface/preprocessing/_row_median_numba.py @@ -1,15 +1,36 @@ -"""numba kernels for ``_row_median``, kept at module level so that ``cache=True`` can find them again.""" +"""numba kernels for ``_row_median``, kept at module level so that ``cache=True`` can find them again. + +Both kernels return exactly what ``np.median(x, axis=1)`` returns. 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. Neither kernel sorts the whole row in the common case; each finds those two values another way. + +``int_row_median`` counts each row into a histogram. 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. + +``float_row_median`` narrows an interval ``[a, b]`` known to hold both middle values. Each pass counts how many +values fall below a probe ``p`` and moves ``a`` or ``b`` to ``p``. Once the interval holds few values, those are +gathered and sorted. The counting loop has no data-dependent branch, so it vectorises. + +Rows these shortcuts cannot handle are sorted in full, so every path returns the exact median. +""" import numpy as np from numba import njit -# The interval is sorted once it holds this many values; smaller needs more counting passes than it saves. +# The float interval is sorted once it holds this many values; smaller needs more counting passes than it saves. _FINAL_WIDTH = 24 +# Counting passes before giving up and sorting the row; real traces need about three. _MAX_PASSES = 48 @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) @@ -19,8 +40,10 @@ def int_row_median(x, window): 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) + 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: @@ -28,6 +51,9 @@ def int_row_median(x, window): 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: @@ -39,7 +65,7 @@ def int_row_median(x, window): if cum > r_hi: hi = i break - for c in range(C): # clear only the bins this row touched + 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 @@ -47,6 +73,8 @@ def int_row_median(x, window): 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]) @@ -55,6 +83,10 @@ def int_row_median(x, window): @njit(inline="always") def _counts(row, v): + """Numbers of values in ``row`` below ``v`` and at most ``v``. + + The comparisons are summed rather than branched on, so the loop compiles to vector compares. + """ lt = 0 le = 0 for c in range(row.size): @@ -65,6 +97,7 @@ def _counts(row, v): @njit(inline="always") def _sorted_median(row, scratch, r_lo, r_hi): + """Median by sorting a copy of ``row``; NaN if the row holds one, as ``np.median`` returns.""" scratch[:] = row scratch.sort() # NaN sorts last if np.isnan(scratch[-1]): @@ -74,6 +107,10 @@ def _sorted_median(row, scratch, r_lo, r_hi): @njit(inline="always") def _median_of_row(row, buf, scratch): + """Median of one float32 row, by narrowing an interval around the middle ranks. + + ``buf`` (at least ``C + 1`` values) and ``scratch`` (``C`` values) are work space reused across rows. + """ C = row.size r_lo = (C - 1) // 2 # the two middle ranks, equal for odd C r_hi = C // 2 @@ -87,13 +124,14 @@ def _median_of_row(row, buf, scratch): vmax = max(vmax, v) if not np.isfinite(total): # NaN, inf, or a sum that overflowed return _sorted_median(row, scratch, r_lo, r_hi) - # Invariant: count(< a) == ca <= r_lo and count(<= b) == cb >= r_hi + 1. + # Invariant: ca = count(< a) <= r_lo and cb = count(<= b) >= r_hi + 1, so both middle values lie in + # [a, b], and cb - ca of the row's values lie there. a = vmin ca = 0 b = vmax cb = C - p = total / C - target = 0.5 * (r_lo + r_hi + 1) + p = total / C # the mean is usually close to the median of filtered traces + target = 0.5 * (r_lo + r_hi + 1) # the count below the median, where the probe aims for _ in range(_MAX_PASSES): if cb - ca <= _FINAL_WIDTH: break @@ -101,7 +139,8 @@ def _median_of_row(row, buf, scratch): if lt <= r_lo and le >= r_hi + 1: return p # p is a value holding both middle ranks if lt > r_lo and le < r_hi + 1: - # p sits strictly between the two middle values, so they are its nearest neighbours + # Only possible for even C, with exactly r_hi values below p and none equal to it: the middle values + # are then the largest value below p and the smallest above it. below = vmin above = vmax for c in range(C): @@ -109,31 +148,38 @@ def _median_of_row(row, buf, scratch): below = max(below, v if v < p else vmin) above = min(above, v if v > p else vmax) return np.float32(0.5) * (below + above) + # p is below both middle values or above both; move that end of the interval to p. if lt <= r_lo: a = p ca = lt else: b = p cb = le + # Next probe: interpolate linearly between (a, ca) and (b, cb) to where the count reaches target, as + # regula falsi does on the row's empirical CDF. Fall back to bisection if rounding puts it outside (a, b). q = a + (target - ca) / max(cb - ca, 1) * (b - a) if not (a < q < b): q = 0.5 * (a + b) p = np.float32(q) if cb - ca > _FINAL_WIDTH: return _sorted_median(row, scratch, r_lo, r_hi) + # Gather the values in [a, b] without branching: every value is written, and the write position advances + # only past values in the interval. buf has one spare slot for the last write. n = 0 - for c in range(C): # branchless gather of the values in [a, b] + for c in range(C): v = row[c] buf[n] = v n += (v >= a) & (v <= b) middle = buf[:n] middle.sort() + # ca values lie below a, so rank r in the row is rank r - ca among the gathered values. return np.float32(0.5) * (middle[r_lo - ca] + middle[r_hi - ca]) # Reassociation only reorders the sum that seeds the search, so the result stays exact. @njit(nogil=True, cache=True, fastmath={"reassoc"}) def float_row_median(x): + """Median of each row of float32 ``x``, as float32 like ``np.median``.""" T, C = x.shape out = np.empty(T, dtype=np.float32) buf = np.empty(C + 1, dtype=np.float32) From beec8965ab273d1583f04ed66a18804a2fe0ffd2 Mon Sep 17 00:00:00 2001 From: Galen Lynch Date: Tue, 13 Oct 2026 08:07:11 -0700 Subject: [PATCH 3/5] perf: sort float rows for the median instead of the counting kernel np.sort vectorises 32- and 64-bit sorts, and on the default chain's float32 output it beats the counting kernel (29 vs 39 ms per 1 s of 384 channels) and np.median (100 ms). This is the approach of #4836. It also covers float64 and needs no numba. The int16 histogram kernel stays, since np.sort is 2.6x slower than np.median there. --- .../preprocessing/_row_median.py | 46 +++--- .../preprocessing/_row_median_numba.py | 134 ++---------------- .../preprocessing/tests/test_row_median.py | 34 ++--- 3 files changed, 51 insertions(+), 163 deletions(-) diff --git a/src/spikeinterface/preprocessing/_row_median.py b/src/spikeinterface/preprocessing/_row_median.py index 6de3d1d140..0284bee8af 100644 --- a/src/spikeinterface/preprocessing/_row_median.py +++ b/src/spikeinterface/preprocessing/_row_median.py @@ -1,8 +1,9 @@ """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. -The numba kernels in ``_row_median_numba`` replace most of those branches with counting: integer rows go through -a histogram, float32 rows through counting passes that vectorise. That module describes both algorithms. Both equal ``np.median``, including for float rows holding NaN or inf. +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 @@ -16,35 +17,44 @@ # Histogram bins; a row whose middle values fall outside the window is sorted instead. _INT_WINDOW = 1024 -# Below these channel counts the kernels' per-row setup costs more than np.median's partition. +# 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 = 64 +_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. - Integer input of at most 16 bits and at least 4 channels, and float32 input of at least 64 channels, use numba - kernels when numba is installed; other input uses ``np.median``. + float32 and float64 input of at least 3 channels is sorted. 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 HAVE_NUMBA and traces.ndim == 2 and traces.shape[0] > 0 and traces.shape[1] > 0: - if traces.dtype == np.float32 and traces.shape[1] >= _MIN_FLOAT_CHANNELS: - return _get_float_kernel()(np.ascontiguousarray(traces))[:, np.newaxis] - if traces.dtype.kind in "iu" and traces.dtype.itemsize <= 2 and traces.shape[1] >= _MIN_INT_CHANNELS: - # Reading column-major traces in place beats copying them; the float32 kernel is the other way round. + 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 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 - - -@cache -def _get_float_kernel(): - from ._row_median_numba import float_row_median - - return float_row_median diff --git a/src/spikeinterface/preprocessing/_row_median_numba.py b/src/spikeinterface/preprocessing/_row_median_numba.py index 94881d5534..73b1dd0e4d 100644 --- a/src/spikeinterface/preprocessing/_row_median_numba.py +++ b/src/spikeinterface/preprocessing/_row_median_numba.py @@ -1,28 +1,18 @@ -"""numba kernels for ``_row_median``, kept at module level so that ``cache=True`` can find them again. +"""numba kernel for ``_row_median``, kept at module level so that ``cache=True`` can find it again. -Both kernels return exactly what ``np.median(x, axis=1)`` returns. 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. Neither kernel sorts the whole row in the common case; each finds those two values another way. +``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. -``int_row_median`` counts each row into a histogram. 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. - -``float_row_median`` narrows an interval ``[a, b]`` known to hold both middle values. Each pass counts how many -values fall below a probe ``p`` and moves ``a`` or ``b`` to ``p``. Once the interval holds few values, those are -gathered and sorted. The counting loop has no data-dependent branch, so it vectorises. - -Rows these shortcuts cannot handle are sorted in full, so every path returns the exact median. +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 -# The float interval is sorted once it holds this many values; smaller needs more counting passes than it saves. -_FINAL_WIDTH = 24 -# Counting passes before giving up and sorting the row; real traces need about three. -_MAX_PASSES = 48 - @njit(nogil=True, cache=True) def int_row_median(x, window): @@ -79,111 +69,3 @@ def int_row_median(x, window): out[t] = 0.5 * (np.float64(s[r_lo]) + np.float64(s[r_hi])) center = np.int64(s[r_lo]) return out - - -@njit(inline="always") -def _counts(row, v): - """Numbers of values in ``row`` below ``v`` and at most ``v``. - - The comparisons are summed rather than branched on, so the loop compiles to vector compares. - """ - lt = 0 - le = 0 - for c in range(row.size): - lt += row[c] < v - le += row[c] <= v - return lt, le - - -@njit(inline="always") -def _sorted_median(row, scratch, r_lo, r_hi): - """Median by sorting a copy of ``row``; NaN if the row holds one, as ``np.median`` returns.""" - scratch[:] = row - scratch.sort() # NaN sorts last - if np.isnan(scratch[-1]): - return scratch[-1] - return np.float32(0.5) * (scratch[r_lo] + scratch[r_hi]) - - -@njit(inline="always") -def _median_of_row(row, buf, scratch): - """Median of one float32 row, by narrowing an interval around the middle ranks. - - ``buf`` (at least ``C + 1`` values) and ``scratch`` (``C`` values) are work space reused across rows. - """ - C = row.size - r_lo = (C - 1) // 2 # the two middle ranks, equal for odd C - r_hi = C // 2 - total = np.float32(0.0) - vmin = row[0] - vmax = row[0] - for c in range(C): - v = row[c] - total += v - vmin = min(vmin, v) - vmax = max(vmax, v) - if not np.isfinite(total): # NaN, inf, or a sum that overflowed - return _sorted_median(row, scratch, r_lo, r_hi) - # Invariant: ca = count(< a) <= r_lo and cb = count(<= b) >= r_hi + 1, so both middle values lie in - # [a, b], and cb - ca of the row's values lie there. - a = vmin - ca = 0 - b = vmax - cb = C - p = total / C # the mean is usually close to the median of filtered traces - target = 0.5 * (r_lo + r_hi + 1) # the count below the median, where the probe aims - for _ in range(_MAX_PASSES): - if cb - ca <= _FINAL_WIDTH: - break - lt, le = _counts(row, p) - if lt <= r_lo and le >= r_hi + 1: - return p # p is a value holding both middle ranks - if lt > r_lo and le < r_hi + 1: - # Only possible for even C, with exactly r_hi values below p and none equal to it: the middle values - # are then the largest value below p and the smallest above it. - below = vmin - above = vmax - for c in range(C): - v = row[c] - below = max(below, v if v < p else vmin) - above = min(above, v if v > p else vmax) - return np.float32(0.5) * (below + above) - # p is below both middle values or above both; move that end of the interval to p. - if lt <= r_lo: - a = p - ca = lt - else: - b = p - cb = le - # Next probe: interpolate linearly between (a, ca) and (b, cb) to where the count reaches target, as - # regula falsi does on the row's empirical CDF. Fall back to bisection if rounding puts it outside (a, b). - q = a + (target - ca) / max(cb - ca, 1) * (b - a) - if not (a < q < b): - q = 0.5 * (a + b) - p = np.float32(q) - if cb - ca > _FINAL_WIDTH: - return _sorted_median(row, scratch, r_lo, r_hi) - # Gather the values in [a, b] without branching: every value is written, and the write position advances - # only past values in the interval. buf has one spare slot for the last write. - n = 0 - for c in range(C): - v = row[c] - buf[n] = v - n += (v >= a) & (v <= b) - middle = buf[:n] - middle.sort() - # ca values lie below a, so rank r in the row is rank r - ca among the gathered values. - return np.float32(0.5) * (middle[r_lo - ca] + middle[r_hi - ca]) - - -# Reassociation only reorders the sum that seeds the search, so the result stays exact. -@njit(nogil=True, cache=True, fastmath={"reassoc"}) -def float_row_median(x): - """Median of each row of float32 ``x``, as float32 like ``np.median``.""" - T, C = x.shape - out = np.empty(T, dtype=np.float32) - buf = np.empty(C + 1, dtype=np.float32) - scratch = np.empty(C, dtype=np.float32) - for t in range(T): - out[t] = _median_of_row(x[t], buf, scratch) - return out diff --git a/src/spikeinterface/preprocessing/tests/test_row_median.py b/src/spikeinterface/preprocessing/tests/test_row_median.py index ccfa05e89d..727fa1ea9d 100644 --- a/src/spikeinterface/preprocessing/tests/test_row_median.py +++ b/src/spikeinterface/preprocessing/tests/test_row_median.py @@ -3,15 +3,7 @@ from spikeinterface.core import NumpyRecording from spikeinterface.preprocessing import common_reference -from spikeinterface.preprocessing._row_median import ( - HAVE_NUMBA, - _INT_WINDOW, - _get_float_kernel, - _get_int_kernel, - row_median, -) - -pytestmark = pytest.mark.skipif(not HAVE_NUMBA, reason="row_median falls back to np.median without numba") +from spikeinterface.preprocessing._row_median import HAVE_NUMBA, _INT_WINDOW, _get_int_kernel, row_median def assert_matches_np_median(x): @@ -21,6 +13,7 @@ def assert_matches_np_median(x): 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): @@ -34,6 +27,7 @@ def test_row_median_int(dtype, num_channels): np.testing.assert_array_equal(_get_int_kernel()(x, _INT_WINDOW), np.median(x, axis=1)) +@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") @@ -43,30 +37,32 @@ def test_row_median_int_outside_window(): assert_matches_np_median(x) -@pytest.mark.parametrize("num_channels", [1, 2, 3, 63, 64, 384, 385]) -def test_row_median_float32(num_channels): +@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("float32") + x = x.astype(dtype) assert_matches_np_median(x) assert_matches_np_median(np.asfortranarray(x)) - # row_median sends few channels to np.median, so test the kernel there directly - np.testing.assert_array_equal(_get_float_kernel()(x), np.median(x, axis=1)) + assert_matches_np_median(x[::2, ::-1]) -def test_row_median_float32_non_finite(): - x = np.random.default_rng(0).standard_normal((5, 384)).astype("float32") +@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] = 3e38 # the row sum overflows - with np.errstate(over="ignore"): - np.testing.assert_array_equal(row_median(x), np.median(x, axis=1, keepdims=True)) + 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"]) From 83a186c3050e77b0a3fb7d8b4f7e36e75e5efd3e Mon Sep 17 00:00:00 2001 From: Galen Lynch Date: Wed, 14 Oct 2026 05:53:16 -0700 Subject: [PATCH 4/5] chore: credit @JESUSROYETH for sorting float rows Sorting float rows for the median, in beec8965a, follows @JESUSROYETH's approach in #4836. Co-authored-by: JESUSROYETH <42451234+JESUSROYETH@users.noreply.github.com> From d967a2ff6f71d9f7ac3342c722335eb07c5a2a39 Mon Sep 17 00:00:00 2001 From: Galen Lynch Date: Wed, 14 Oct 2026 16:37:55 -0700 Subject: [PATCH 5/5] fix: send byte-swapped integer traces to np.median in row_median numba does not accept non-native byte order, so the integer kernel raised on big-endian int16 or uint16 traces. Those now use np.median, as other unsupported dtypes already did. --- src/spikeinterface/preprocessing/_row_median.py | 5 +++-- src/spikeinterface/preprocessing/tests/test_row_median.py | 6 ++++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/preprocessing/_row_median.py b/src/spikeinterface/preprocessing/_row_median.py index 0284bee8af..1d740830ce 100644 --- a/src/spikeinterface/preprocessing/_row_median.py +++ b/src/spikeinterface/preprocessing/_row_median.py @@ -25,8 +25,8 @@ 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. 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``. + 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] @@ -36,6 +36,7 @@ def row_median(traces: np.ndarray) -> np.ndarray: 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. diff --git a/src/spikeinterface/preprocessing/tests/test_row_median.py b/src/spikeinterface/preprocessing/tests/test_row_median.py index 727fa1ea9d..eae228f40c 100644 --- a/src/spikeinterface/preprocessing/tests/test_row_median.py +++ b/src/spikeinterface/preprocessing/tests/test_row_median.py @@ -27,6 +27,12 @@ def test_row_median_int(dtype, num_channels): 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)