Reorder spike vectors with a flexible-dtype numba kernel by grahamfindlay · Pull Request #4784 · SpikeInterface/spikeinterface

grahamfindlay · 2026-09-15T20:26:35Z

Follows @samuelgarcia 's comment on #4695 and #4781.

The numba 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, 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)

grahamfindlay · 2026-09-17T01:23:48Z

The 2 failing tests don't appear to have anything to do with this PR. They're tolerance tests failing on test_decimate_antialias_by_chunks which funnily enough passed in its own PR (#4621) with the same scipy and numpy versions, so they must be nondeterministic. Maybe create_sinusoidal_traces just needs a seed. It looks like they're still tiny tiny single-sample discrepancies. It must be a rounding thing on Apple Silicon.

grahamfindlay · 2026-09-17T01:32:16Z

~~Yes, I know it was kinda naughty of me to push a fix for that to this PR... forgive me.~~

EDIT: No longer applicable, see below.

grahamfindlay · 2026-09-17T17:44:53Z

Removed the fix for the unrelated failing tests since another one was already just merged in #4792

grahamfindlay · 2026-10-08T21:21:23Z

@samuelgarcia the failing Python 3.14 tests are unrelated and should be fixed in #4845 so this is still ready for your review anytime

The counting-sort kernel from SpikeInterface#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)

for more information, see https://pre-commit.ci

grahamfindlay · 2026-10-30T16:26:09Z

@samuelgarcia Actually, nothing about this had to change, so it is all ready for your review.

JoeZiminski

cool @grahamfindlay this is very nice! Took me a while to look back and understand the provenance of these changes, in general a cool program of work. I expected this change to require much more than it did, very elegant.

No real comments feel free to ignore renaming suggestions, the only one to really consider is excluding some other dtype than numba can't handle e.g. float16

JoeZiminski · 2026-11-12T14:59:37Z

+        any width. Any other fields (e.g. the "channel_index" that
+        `to_spike_vector(main_channel_indices=...)` adds) are carried through to the output.
+        Anything array-like is accepted and materialised with `np.asarray`, so lazily-backed spike
+        vectors work, at the cost of being read into memory.


worth mentioning here that if any other fields are objects it forces the slow path?

JoeZiminski · 2026-11-12T15:03:21Z

-        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)


could this be called ordered_spike_vector?

JoeZiminski · 2026-11-12T15:08:45Z

    @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


could spikes be spike_vector and out be output_spike_vector? (not strong on this, but renaming out would be nice)

JoeZiminski · 2026-11-12T15:10:08Z

+        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.


Could mention here 'They can be of arbitrary integer dtype, and so are widened…"

JoeZiminski · 2026-11-12T15:11:45Z

+
+        `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.


I am curious why this doesn't work if the extra fields are objects now

JoeZiminski · 2026-11-12T15:43:10Z



+NARROW_INDEX_DTYPES = {
+    "int64/int32/int8": np.dtype([("sample_index", "int64"), ("unit_index", "int32"), ("segment_index", "int8")]),


maybe this is done elsewhere, is it also worth testing the all int64 array too?

JoeZiminski · 2026-11-12T15:47:03Z

    """
    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:


we could also check:

dtype is not float16 which numba cannot support, as well as some extended precision float and complex dtype

dtype.isnative ensures endianness is same as machine in use

alejoe91 requested a review from samuelgarcia September 16, 2026 09:11

alejoe91 added core Changes to core module performance Performance issues/improvements labels Sep 16, 2026

grahamfindlay force-pushed the perf/reorder-record-kernel branch from 62a512f to 22cd36f Compare September 17, 2026 16:42

JoeZiminski self-requested a review September 23, 2026 20:05

grahamfindlay force-pushed the perf/reorder-record-kernel branch from 22cd36f to 89ee24f Compare October 8, 2026 17:10

grahamfindlay force-pushed the perf/reorder-record-kernel branch from 89ee24f to 98f3b8d Compare October 14, 2026 16:45

grahamfindlay and others added 3 commits October 29, 2026 15:40

[pre-commit.ci] auto fixes from pre-commit.com hooks

ca3bb1b

for more information, see https://pre-commit.ci

Update reorder docstrings for flexible spike-vector dtypes

99ab7c5

grahamfindlay force-pushed the perf/reorder-record-kernel branch from 98f3b8d to 99ab7c5 Compare October 29, 2026 20:57

JoeZiminski approved these changes Nov 12, 2026

View reviewed changes



		NARROW_INDEX_DTYPES = {
		"int64/int32/int8": np.dtype([("sample_index", "int64"), ("unit_index", "int32"), ("segment_index", "int8")]),

Conversation

grahamfindlay commented Sep 15, 2026

Uh oh!

grahamfindlay commented Sep 17, 2026

Uh oh!

grahamfindlay commented Sep 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

grahamfindlay commented Sep 17, 2026

Uh oh!

grahamfindlay commented Oct 8, 2026

Uh oh!

grahamfindlay commented Oct 30, 2026

Uh oh!

JoeZiminski left a comment

Choose a reason for hiding this comment

Uh oh!

JoeZiminski Nov 12, 2026

Choose a reason for hiding this comment

Uh oh!

JoeZiminski Nov 12, 2026

Choose a reason for hiding this comment

Uh oh!

JoeZiminski Nov 12, 2026

Choose a reason for hiding this comment

Uh oh!

JoeZiminski Nov 12, 2026

Choose a reason for hiding this comment

Uh oh!

JoeZiminski Nov 12, 2026

Choose a reason for hiding this comment

Uh oh!

JoeZiminski Nov 12, 2026

Choose a reason for hiding this comment

Uh oh!

JoeZiminski Nov 12, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

grahamfindlay commented Sep 17, 2026 •

edited

Loading