From 1ea92254a0effd5f27a714ede73b5f5bf25ec177 Mon Sep 17 00:00:00 2001 From: Graham Findlay Date: Tue, 15 Sep 2026 15:23:08 -0500 Subject: [PATCH 1/2] Reorder spike vectors with a record-dtype numba kernel The counting-sort kernel from #4695 moved rows through a flat (N, num_fields) int64 view and addressed unit_index and segment_index as columns 1 and 2, essentially assuming minimum_spike_dtype. Samuel suggest that we might want to change the spike vector dtype in the future, so this is a more flexible variant that can take narrow index dtypes (e.g. int8). Also updated the numpy fallback to work for narrow index fields. Measured on ~400M spikes, 342 units: - 11.6 -> 12.0 s (3x int64 fields) - 13.1 -> 13.0 s (4x int64 fields) - 9.1 s (1 each of int64/int32/int8 fields) --- src/spikeinterface/core/sorting_tools.py | 86 ++++++++++--------- .../core/tests/test_sorting_tools.py | 78 +++++++++++++++++ 2 files changed, 123 insertions(+), 41 deletions(-) diff --git a/src/spikeinterface/core/sorting_tools.py b/src/spikeinterface/core/sorting_tools.py index f37b24046d..dde71bbdb6 100644 --- a/src/spikeinterface/core/sorting_tools.py +++ b/src/spikeinterface/core/sorting_tools.py @@ -4,7 +4,7 @@ import numpy as np -from spikeinterface.core.base import BaseExtractor, minimum_spike_dtype, unit_period_dtype +from spikeinterface.core.base import BaseExtractor, unit_period_dtype from spikeinterface.core.basesorting import BaseSorting from spikeinterface.core.numpyextractors import NumpySorting @@ -149,26 +149,18 @@ def vector_to_list_of_spiketrain_numba(sample_indices, unit_indices, num_units): return vector_to_list_of_spiketrain_numba -def _is_flat_int64_view(dtype: np.dtype) -> bool: +def _numba_can_reorder(dtype: np.dtype) -> bool: """ - Whether a spike-vector dtype can be safely viewed as a flat (num_spikes, num_fields) int64 - matrix, which is what the numba counting sort moves rows through. + Whether the numba counting sort can handle this spike-vector dtype. - Requires every field to be int64, packed with no padding, and the first three fields to be - `minimum_spike_dtype`'s in order, because the kernel addresses unit_index and segment_index - positionally (columns 1 and 2) rather than by name. + We need integer `unit_index` and `segment_index` fields and no object fields. - The all-int64 rule is deliberately stricter than correctness demands -- the kernel copies rows - bitwise, so any 8-byte field would in fact round-trip through the view. Keeping it narrow means - the kernel only ever sees the layout it is written for, and it costs nothing in practice: every - spike-vector dtype spikeinterface constructs is all-int64. + Numba can handle any other fields (e.g. extra float fields, string fields) just fine through the record copy: """ names = dtype.names - if names is None or names[:3] != tuple(name for name, _ in minimum_spike_dtype): + if names is None or dtype.hasobject: return False - if dtype.itemsize != 8 * len(names): - return False - return all(dtype.fields[name][0] == np.int64 and dtype.fields[name][1] == 8 * i for i, name in enumerate(names)) + return all(name in names and dtype.fields[name][0].kind in "iu" for name in ("unit_index", "segment_index")) def reorder_spike_vector_by_unit_and_segment( @@ -245,28 +237,37 @@ def reorder_spike_vector_by_unit_and_segment( f"`spike_vector` has a unit_index outside [0, {num_units}) or a segment_index outside [0, {num_segments})." ) - # The numba kernel expects an all-int64 unpadded dtype (e.g. `minimum_spike_dtype`), but it is - # possible that a spike vector has extra fields with other dtypes (`NumpySorting` allows that). - # So we check taht the numba path is safe, and anything else takes the dtype-agnostic numpy path. - if HAVE_NUMBA and _is_flat_int64_view(spike_vector.dtype): + # If the spike-vector's dtype has object fields and/or non-integer + # unit_index and segment_index fields, we have to use numpy. + # Otherwise, we can use the numba kernel. + if HAVE_NUMBA and _numba_can_reorder(spike_vector.dtype): reorder_spike_vector = get_numba_reorder_spike_vector() - num_fields = len(spike_vector.dtype.names) - # These flat (num_spikes, num_fields) int64 views are zero-copy - in_flat = np.ascontiguousarray(spike_vector).view(np.int64).reshape(num_spikes, num_fields) - out_flat = np.empty((num_spikes, num_fields), dtype=np.int64) + spike_vector = np.ascontiguousarray(spike_vector) + ordered_spikes = np.empty(num_spikes, dtype=spike_vector.dtype) order = np.empty(num_spikes, dtype=np.int64) counts = np.empty(num_buckets, dtype=np.int64) - in_range = reorder_spike_vector(in_flat, unit_stride, segment_stride, num_buckets, out_flat, order, counts) + in_range = reorder_spike_vector( + spike_vector, + spike_vector["unit_index"], + spike_vector["segment_index"], + unit_stride, + segment_stride, + num_buckets, + ordered_spikes, + order, + counts, + ) if not in_range: raise ValueError(out_of_range_error) - - ordered_spikes = out_flat.view(spike_vector.dtype).reshape(num_spikes) return ordered_spikes, order, counts # numpy fallback: a stable argsort by bucket is equivalent to the counting sort above. - bucket_index = spike_vector["unit_index"] * unit_stride + spike_vector["segment_index"] * segment_stride + # Widen to int64 first so narrow index fields (e.g. int8) don't overflow when multiplied by the stride. + unit_index = spike_vector["unit_index"].astype(np.int64, copy=False) + segment_index = spike_vector["segment_index"].astype(np.int64, copy=False) + bucket_index = unit_index * unit_stride + segment_index * segment_stride # Must be checked before narrowing: a negative or oversized bucket would silently wrap. if bucket_index.min() < 0 or bucket_index.max() >= num_buckets: @@ -290,11 +291,15 @@ def get_numba_reorder_spike_vector(): from numba import jit @jit(nopython=True, nogil=True, cache=False) - def reorder_spike_vector_numba(in_flat, unit_stride, segment_stride, num_buckets, out_flat, order, counts): + def reorder_spike_vector_numba( + spikes, unit_index, segment_index, unit_stride, segment_stride, num_buckets, out, order, counts + ): """ - Stable counting-sort of a (N, num_fields) int64 spike-vector flat-buffer view by - (unit, segment). `num_fields` is 3 for `minimum_spike_dtype`, more when the spike vector - carries extra int64 fields; the extra columns are copied along with their spike. + Stable counting-sort of a structured spike vector by (unit, segment). + + `unit_index` and `segment_index` are the field views (i.e., zero-copy, strided) of `spikes`. + they are widened to int64 so narrow fields can't overflow the bucket arithmetic. + Rows move as whole records (`out[pos] = spikes[i]`), so any extra fields travel with their spike. Each spike's bucket is derived on the fly as `unit_index * unit_stride + segment_index * segment_stride`, so no bucket array is needed. @@ -302,25 +307,25 @@ def reorder_spike_vector_numba(in_flat, unit_stride, segment_stride, num_buckets Two O(N) passes: 1. histogram the buckets into `counts`, 2. cumulative-sum to per-bucket write positions, then scatter each - row of `in_flat` to its destination in `out_flat` and record the - source index in `order` so that ``in[order] == out``. + record of `spikes` to its destination in `out` and record the + source index in `order` so that ``spikes[order] == out``. - `out_flat`, `order` and `counts` are filled in place. + `out`, `order` and `counts` are filled in place. Stability: within each bucket, rows keep their input order, so any - ordering already present in `in_flat` (e.g. ascending sample_index - within a (segment, unit) group) carries over to `out_flat`. + ordering already present in `spikes` (e.g. ascending sample_index + within a (segment, unit) group) carries over to `out`. Returns False if any spike falls outside [0, num_buckets), in which case the outputs are meaningless; True otherwise. """ - num_spikes, num_fields = in_flat.shape + num_spikes = spikes.shape[0] # Pass 1: histogram the buckets and do bounds-check (free! we already have to make the pass) for b in range(num_buckets): counts[b] = 0 for i in range(num_spikes): - bucket = in_flat[i, 1] * unit_stride + in_flat[i, 2] * segment_stride + bucket = np.int64(unit_index[i]) * unit_stride + np.int64(segment_index[i]) * segment_stride if bucket < 0 or bucket >= num_buckets: return False counts[bucket] += 1 @@ -335,10 +340,9 @@ def reorder_spike_vector_numba(in_flat, unit_stride, segment_stride, num_buckets # Pass 2: scatter each spike into its bucket, recording where it came from. for i in range(num_spikes): - bucket = in_flat[i, 1] * unit_stride + in_flat[i, 2] * segment_stride + bucket = np.int64(unit_index[i]) * unit_stride + np.int64(segment_index[i]) * segment_stride pos = write_pos[bucket] - for field in range(num_fields): - out_flat[pos, field] = in_flat[i, field] + out[pos] = spikes[i] order[pos] = i write_pos[bucket] = pos + 1 diff --git a/src/spikeinterface/core/tests/test_sorting_tools.py b/src/spikeinterface/core/tests/test_sorting_tools.py index 7ad8291073..3e4391ee3e 100644 --- a/src/spikeinterface/core/tests/test_sorting_tools.py +++ b/src/spikeinterface/core/tests/test_sorting_tools.py @@ -94,6 +94,9 @@ def test_reorder_spike_vector_by_unit_and_segment_raises(force_numba): reorder_spike_vector_by_unit_and_segment(spikes, 1, 1) # unit_index 1 >= num_units with pytest.raises(ValueError, match="outside"): reorder_spike_vector_by_unit_and_segment(_make_spike_vector([0], [0], [5]), 1, 1) + # numba wraps a negative index, so check the lower bound + with pytest.raises(ValueError, match="outside"): + reorder_spike_vector_by_unit_and_segment(_make_spike_vector([0], [-1], [0]), 1, 1) @pytest.mark.parametrize("num_units", [2, 300, 70_000], ids=["uint8", "uint16", "uint32"]) @@ -175,6 +178,81 @@ def test_reorder_spike_vector_by_unit_and_segment_non_uniform_dtype(force_numba, assert np.array_equal(ordered_spikes, _legacy_reorder(spikes)) +NARROW_INDEX_DTYPES = { + "int64/int32/int8": np.dtype([("sample_index", "int64"), ("unit_index", "int32"), ("segment_index", "int8")]), + "uint16/int8": np.dtype([("sample_index", "int64"), ("unit_index", "uint16"), ("segment_index", "int8")]), +} + + +def _make_random_spikes(dtype, num_units, num_segments, num_spikes, seed=0): + """Spikes in valid spike-vector order (segment-blocked, sample-ascending) with random units.""" + rng = np.random.default_rng(seed) + spikes = np.empty(num_spikes, dtype=dtype) + segment_indices = np.sort(rng.integers(0, num_segments, size=num_spikes)) + spikes["segment_index"] = segment_indices + for segment_index in range(num_segments): + in_segment = segment_indices == segment_index + spikes["sample_index"][in_segment] = np.sort(rng.integers(0, 1_000, size=in_segment.sum())) + spikes["unit_index"] = rng.integers(0, num_units, size=num_spikes) + return spikes + + +@pytest.mark.parametrize("unit_major", [True, False], ids=["unit_major", "segment_major"]) +@pytest.mark.parametrize("dtype", list(NARROW_INDEX_DTYPES.values()), ids=list(NARROW_INDEX_DTYPES.keys())) +def test_reorder_spike_vector_by_unit_and_segment_narrow_index_dtypes(force_numba, dtype, unit_major): + """Make sure that narrow unit_index / segment_index fields use numba. + + """ + # 200 units x 2 segments puts the bucket index outside int8 range, + # so the kernel has to widen the fields before the bucket math. + num_units, num_segments = 200, 2 + + spikes = _make_random_spikes(dtype, num_units, num_segments, num_spikes=2_000) + if force_numba: + from spikeinterface.core.sorting_tools import _numba_can_reorder + + assert _numba_can_reorder(spikes.dtype) + + ordered_spikes, order, counts = reorder_spike_vector_by_unit_and_segment( + spikes, num_units, num_segments, unit_major=unit_major + ) + + assert ordered_spikes.dtype == spikes.dtype + assert np.array_equal(ordered_spikes, spikes[order]) + assert np.array_equal(ordered_spikes, _legacy_reorder(spikes, unit_major=unit_major)) + assert counts.sum() == spikes.size + + +def test_reorder_spike_vector_by_unit_and_segment_object_field_falls_back(): + """Make sure spike vectors with object fields use numpy""" + dtype = minimum_spike_dtype + [("label", "O")] + spikes = np.empty(6, dtype=dtype) + spikes["sample_index"] = [10, 10, 11, 12, 12, 13] + spikes["unit_index"] = [2, 0, 1, 2, 0, 0] + spikes["segment_index"] = 0 + spikes["label"] = list("abcdef") + + ordered_spikes, order, counts = reorder_spike_vector_by_unit_and_segment(spikes, 3, 1) + + assert np.array_equal(counts, [3, 1, 2]) + assert np.array_equal(ordered_spikes, spikes[order]) + assert list(ordered_spikes["label"]) == ["b", "e", "f", "c", "a", "d"] + + +def test_numba_can_reorder(): + from spikeinterface.core.sorting_tools import _numba_can_reorder + + assert _numba_can_reorder(np.dtype(minimum_spike_dtype)) + assert _numba_can_reorder(np.dtype([("sample_index", "int64"), ("unit_index", "int32"), ("segment_index", "int8")])) + + assert not _numba_can_reorder(np.dtype(minimum_spike_dtype + [("label", "O")])) + assert not _numba_can_reorder(np.dtype([("sample_index", "int64"), ("unit_index", "int64")])) + assert not _numba_can_reorder( + np.dtype([("sample_index", "int64"), ("unit_index", "float64"), ("segment_index", "int64")]) + ) + assert not _numba_can_reorder(np.dtype("int64")) + + def test_random_spikes_selection(): recording, sorting = generate_ground_truth_recording( durations=[20.0, 10.0], From 26fbf0d01629be6318bf0fb3d2b4226b05444e01 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:27:03 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/spikeinterface/core/tests/test_sorting_tools.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/spikeinterface/core/tests/test_sorting_tools.py b/src/spikeinterface/core/tests/test_sorting_tools.py index 3e4391ee3e..ad684b1f31 100644 --- a/src/spikeinterface/core/tests/test_sorting_tools.py +++ b/src/spikeinterface/core/tests/test_sorting_tools.py @@ -200,9 +200,7 @@ def _make_random_spikes(dtype, num_units, num_segments, num_spikes, seed=0): @pytest.mark.parametrize("unit_major", [True, False], ids=["unit_major", "segment_major"]) @pytest.mark.parametrize("dtype", list(NARROW_INDEX_DTYPES.values()), ids=list(NARROW_INDEX_DTYPES.keys())) def test_reorder_spike_vector_by_unit_and_segment_narrow_index_dtypes(force_numba, dtype, unit_major): - """Make sure that narrow unit_index / segment_index fields use numba. - - """ + """Make sure that narrow unit_index / segment_index fields use numba.""" # 200 units x 2 segments puts the bucket index outside int8 range, # so the kernel has to widen the fields before the bucket math. num_units, num_segments = 200, 2