Fix and speed up to_reordered_spike_vector - #4695
Open
grahamfindlay wants to merge 5 commits into
Open
Conversation
alejoe91
added
core
…erface#4683) Also fixes the docstring, which was wrong in two places: it described the pre-SpikeInterface#4606 vanilla order, and documented a default that was not the actual default.
Adds reorder_spike_vector_by_unit_and_segment() to sorting_tools. With numba, this is a counting sort. Without, it is a stable argsort on a bucket narrowed to the smallest dtype that fits, which lets numpy radix sort it. On ~400M spikes, 342 units, 1 segment, to_reordered_spike_vector goes from 137s to: * ~12s with numba (~11x), at unchanged peak memory usage * ~85s with numpy (~1.6x), for +2 bytes/spike (the narrowed bucket) ordered_spikes, slices and order are identical to the previous implementation. Also tested on synthetic recordings of various sizes.
to_reordered_spike_vector was caching under str(lexsort), but the three readers of _cached_lexsorted_spike_vector -- get_unit_spike_train, count_num_spikes_per_unit and precompute_spike_trains -- probed it with the tuple, so those lookups never matched, and the cache was silently missing every time. The bug was possible because the same lexsort literal had to be retyped at eight sites, and they drifted apart. Using constants LEXSORT_UNIT_COMPACT / LEXSORT_SEGMENT_COMPACT is not a perfect solution (it would probably be better not to reach into the private dictionary in so many places), but it it should at least minimize the likelihood that the same mistake gets made again.
backing. reorder_spike_vector_by_unit_and_segment assumed a three-field int64 in-memory array. However, spike vectors can have extra fields (e.g. "channel_index", also int64), including arbitrary user-defined fields (e.g. "amplitude", maybe float64). They can also be lazy (e.g. ZarrSpikeVector). Updated the numba kernel to take extra int64 field(s). Anything spike vector with a non-int64 field now gets routed through the numpy fallback path (which is dtype-agnsotic). Lazy vectors are normalised with np.asarray (which np.lexsort was already doing implicitly).
grahamfindlay
force-pushed
the
perf/to-reordered-spike-vector
branch
from
August 7, 2026 23:28
7dea432 to
f1ff67c
Compare
Contributor
Author
|
I was refreshing this to keep it up-to-date with main, and learned from two of the intervening PRs that (1) spike vectors can be lazy and (2) can have extra fields like |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follows #4606 and #4683. Implements part of #4502. @alejoe91 @samuelgarcia
elif ("sample_index", "unit_index", "segment_index"):into_reordered_spike_vector.to_reordered_spike_vectorto non-vanilla orderings, and fixes some inaccuracies in the docstring.to_reordered_spike_vector. There is a numba path with a counting sort that drops the time on an ~400M spike / 342 unit / 1 segment sorting from 137s to 12s, with no additional memory usage. There is a numpy fallback path using a radix sort that drops the time from 137s to 85s, with a slightly increased memory cost (34 bytes/spike instead of 32 bytes/spike). The outputs (ordered_spikes,slices,order) are identical to the old implementation. I also tested on synthetic data of various sizes (1M, 10M, 40M spikes) to make sure both paths are strict improvements at all input sizes.