From d89c89ade4c91cdf489f724647b86602e8b86a6a Mon Sep 17 00:00:00 2001 From: Graham Findlay Date: Thu, 16 Jul 2026 13:44:48 -0500 Subject: [PATCH 1/5] Fix always-true `elif` in to_reordered_spike_vector --- src/spikeinterface/core/basesorting.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/spikeinterface/core/basesorting.py b/src/spikeinterface/core/basesorting.py index 9ed30e97c7..db7ad01bb3 100644 --- a/src/spikeinterface/core/basesorting.py +++ b/src/spikeinterface/core/basesorting.py @@ -1130,7 +1130,7 @@ def to_reordered_spike_vector( s1 = seg_slices[segment_index + 1] slices[unit_index, segment_index, :] = [u0 + s0, u0 + s1] - elif ("sample_index", "unit_index", "segment_index"): + elif lexsort == ("sample_index", "unit_index", "segment_index"): slices = np.zeros((num_segments, num_units, 2), dtype=np.int64) seg_slices = np.searchsorted(ordered_spikes["segment_index"], np.arange(num_segments + 1), side="left") for segment_index in range(self.get_num_segments()): From 2245823711cfce970a1c8f3f4b3f52911156c807 Mon Sep 17 00:00:00 2001 From: Graham Findlay Date: Thu, 16 Jul 2026 14:22:02 -0500 Subject: [PATCH 2/5] Restrict to_reordered_spike_vector to non-vanilla orderings (#4683) Also fixes the docstring, which was wrong in two places: it described the pre-#4606 vanilla order, and documented a default that was not the actual default. --- doc/modules/core.rst | 9 ++-- src/spikeinterface/core/basesorting.py | 66 +++++++++++++------------- 2 files changed, 39 insertions(+), 36 deletions(-) diff --git a/doc/modules/core.rst b/doc/modules/core.rst index 6ea3d25eb6..27136b1449 100644 --- a/doc/modules/core.rst +++ b/doc/modules/core.rst @@ -175,8 +175,9 @@ the spiketrain, which are optimally organized for specific types of calculation. Computations involving combined recording-sorting information, such as fetching recording chunks and spiketrain chunks to accumulate waveforms, are often quickest when spikes are time-ordered. For -this use case, we use an internal representation called the `spike_vector`. This is a unique buffer: -a numpy.array with dtype `[("sample_index", "int64"), ("unit_index", "int64"), ("segment_index", "int64")]`. +this use case, we use an internal representation called the `spike_vector`, obtained by calling +`sorting.to_spike_vector()`. This is a unique buffer: a numpy.array with dtype +`[("sample_index", "int64"), ("unit_index", "int64"), ("segment_index", "int64")]`. For computations which are done unit-by-unit, like computing isi-violations per unit, it is better that spikes from a single unit are concurrent in memory. For these other cases, we can re-order the @@ -185,8 +186,8 @@ spikes from a single unit are concurrent in memory. For these other cases, we ca * order by unit, then segment, then sample * order by segment, then unit, then sample -This is done using `sorting.to_reordered_spike_vector()`. The first time a reordering is done, the -reordered spiketrain is cached in memory by default. Users should rarely have to worry about these +This is done using `sorting.to_reordered_spike_vector()`. The first time a reordering is done, +the reordered spiketrain is cached in memory by default. Users should rarely have to worry about these details, but developers should keep memory layout in mind when implementing new features. diff --git a/src/spikeinterface/core/basesorting.py b/src/spikeinterface/core/basesorting.py index db7ad01bb3..8fd351d9c4 100644 --- a/src/spikeinterface/core/basesorting.py +++ b/src/spikeinterface/core/basesorting.py @@ -1046,31 +1046,24 @@ def to_reordered_spike_vector( Please note that the lexsort syntax is the **reverse** of natural reading. - By default the spike_vector is lexsort-ed like this: - - ("unit_index", "sample_index", "segment_index") (segment then sample then unit) - - But particular orderings can be better for some computations: - - ("sample_index", "unit_index", "segment_index") (segment then unit_index then sample) - - ("sample_index", "segment_index", "unit_index") (unit_index then segment then sample) - - Note that the last representation makes the spiketrain per segment compact in memory. + Two reorderings are supported: + - ("sample_index", "segment_index", "unit_index"): makes each unit's + spiketrain compact in memory. This is the default, and is what + unit-by-unit computations (e.g. isi violations) want. + - ("sample_index", "unit_index", "segment_index"): makes each segment compact, + with each unit's spiketrain compact within a segment. Rarely (if ever) used, + but might be useful when iterating segment by segment. This operation is internally cached. - The order vector is also computed and can be applied to other external vectors like - spike_amplitudes, spike_locations, ... - - An array of internal slices is also precomputed to have a fast access to a compact - portion of the reordered spikes. - Theses slices are stored as a 3d array to handle start->stop and depend of the lexsort itself. - Theses slices are pre computed using nested searchsorted. - Parameters ---------- - lexsort : tuple, default: ("sample_index", "unit_index", "segment_index") - Tuple for lexsort. Please note that this is the reverse natural reading order! + lexsort : tuple, default: ("sample_index", "segment_index", "unit_index") + The requested sort order, as keys would be passed to np.lexsort. + Please note that this is the reverse natural reading order! + Must be one of the two orderings listed above. return_order: bool, default: True - Return the order, or not. See Returns. + Return the numpy array needed to sort the spike vector (given the requested sort). return_slices: bool, default: True Return the slices, or not. See Returns. @@ -1080,25 +1073,34 @@ def to_reordered_spike_vector( Structured numpy array ("sample_index", "unit_index", "segment_index") with all spikes in the desired lexsort order order : np.array - Numpy array needed to sort the spike vector given the lexsort - slices : np.array - Numpy array of size (num_units, num_segments, 2) or (num_segments, num_units, 2) given the lexsort, - where one can obtain the indices amin, amax of all the (segment,unit_index) values. + Numpy array needed to sort the spike vector given the lexsort. Can be used + to sort other external vectors like spike_amplitudes, spike_locations, ... + slices : np.array + A 3D array of internal slices for fast access to a compact portion of the reordered spikes. + Depending on the lexsort, a numpy array of size (num_units, num_segments, 2) or (num_segments, num_units, 2). + The last dimension contains the start and end indices of each segment-unit pair. + + Raises + ------ + ValueError + If `lexsort` is not one of the two supported orderings. """ lexsort = tuple(lexsort) if lexsort == ("unit_index", "sample_index", "segment_index"): - assert ( - not return_order and not return_slices - ), 'If lexsort = ("unit_index", "sample_index", "segment_index"), both `return_order` and `return_slices` must be set to `False`.' - - spikes = self.to_spike_vector(concatenated=True) - return spikes + raise ValueError( + '`lexsort` = ("unit_index", "sample_index", "segment_index") is not supported: ' + "Use `to_spike_vector()` to get the default order." + ) - assert lexsort in [ - ("sample_index", "unit_index", "segment_index"), + if lexsort not in [ ("sample_index", "segment_index", "unit_index"), - ], '`lexsort` must be equal to ("unit_index", "sample_index", "segment_index"), ("sample_index", "unit_index", "segment_index") or ("sample_index", "segment_index", "unit_index")' + ("sample_index", "unit_index", "segment_index"), + ]: + raise ValueError( + '`lexsort` must be ("sample_index", "segment_index", "unit_index") or ' + f'("sample_index", "unit_index", "segment_index"); got {lexsort}.' + ) key = str(lexsort) From fef7e3bfe1f2cdb5d51d9956fec7d4706135e651 Mon Sep 17 00:00:00 2001 From: Graham Findlay Date: Thu, 16 Jul 2026 21:47:23 -0500 Subject: [PATCH 3/5] Speed up to_reordered_spike_vector with a counting sort 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. --- src/spikeinterface/core/basesorting.py | 72 +++---- src/spikeinterface/core/sorting_tools.py | 176 +++++++++++++++++- .../core/tests/test_basesorting.py | 103 +++++++++- .../core/tests/test_sorting_tools.py | 65 +++++++ 4 files changed, 365 insertions(+), 51 deletions(-) diff --git a/src/spikeinterface/core/basesorting.py b/src/spikeinterface/core/basesorting.py index 8fd351d9c4..7b4fe31028 100644 --- a/src/spikeinterface/core/basesorting.py +++ b/src/spikeinterface/core/basesorting.py @@ -1105,59 +1105,37 @@ def to_reordered_spike_vector( key = str(lexsort) if key not in self._cached_lexsorted_spike_vector.keys(): - spikes = self.to_spike_vector() - order = np.lexsort((spikes[lexsort[0]], spikes[lexsort[1]], spikes[lexsort[2]])) - ordered_spikes = spikes[order] - self._cached_lexsorted_spike_vector[key] = {} - self._cached_lexsorted_spike_vector[key]["ordered_spikes"] = ordered_spikes - self._cached_lexsorted_spike_vector[key]["order"] = order + from .sorting_tools import reorder_spike_vector_by_unit_and_segment + spikes = self.to_spike_vector() num_units = len(self.unit_ids) num_segments = self.get_num_segments() - # precompute the slices with nested search sorted - if lexsort == ("sample_index", "segment_index", "unit_index"): - # this case make spiketrain per unit compact in memory - - slices = np.zeros((num_units, num_segments, 2), dtype=np.int64) - unit_slices = np.searchsorted(ordered_spikes["unit_index"], np.arange(num_units + 1), side="left") - for unit_index, unit_id in enumerate(self.unit_ids): - u0 = unit_slices[unit_index] - u1 = unit_slices[unit_index + 1] - seg_slices = np.searchsorted( - ordered_spikes[u0:u1]["segment_index"], np.arange(num_segments + 1), side="left" - ) - for segment_index in range(num_segments): - s0 = seg_slices[segment_index] - s1 = seg_slices[segment_index + 1] - slices[unit_index, segment_index, :] = [u0 + s0, u0 + s1] - - elif lexsort == ("sample_index", "unit_index", "segment_index"): - slices = np.zeros((num_segments, num_units, 2), dtype=np.int64) - seg_slices = np.searchsorted(ordered_spikes["segment_index"], np.arange(num_segments + 1), side="left") - for segment_index in range(self.get_num_segments()): - s0 = seg_slices[segment_index] - s1 = seg_slices[segment_index + 1] - unit_slices = np.searchsorted( - ordered_spikes[s0:s1]["unit_index"], np.arange(num_units + 1), side="left" - ) - for unit_index, unit_id in enumerate(self.unit_ids): - u0 = unit_slices[unit_index] - u1 = unit_slices[unit_index + 1] - slices[segment_index, unit_index, :] = [s0 + u0, s0 + u1] - - self._cached_lexsorted_spike_vector[key]["slices"] = slices - - ordered_spikes = self._cached_lexsorted_spike_vector[key]["ordered_spikes"] - out = (ordered_spikes,) + unit_major = lexsort == ("sample_index", "segment_index", "unit_index") + slices_shape = (num_units, num_segments) if unit_major else (num_segments, num_units) + + ordered_spikes, order, counts = reorder_spike_vector_by_unit_and_segment( + spikes, num_units, num_segments, unit_major=unit_major + ) + + counts = counts.reshape(slices_shape) + stops = np.cumsum(counts.ravel()).reshape(slices_shape) + starts = stops - counts + slices = np.stack([starts, stops], axis=-1).astype(np.int64, copy=False) + + self._cached_lexsorted_spike_vector[key] = { + "ordered_spikes": ordered_spikes, + "order": order, + "slices": slices, + } + + cached = self._cached_lexsorted_spike_vector[key] + out = [cached["ordered_spikes"]] if return_order: - out += (self._cached_lexsorted_spike_vector[key]["order"],) + out.append(cached["order"]) if return_slices: - out += (self._cached_lexsorted_spike_vector[key]["slices"],) - if len(out) == 1: - return out[0] - else: - return out + out.append(cached["slices"]) + return tuple(out) if len(out) > 1 else out[0] def to_numpy_sorting(self, propagate_cache=True): """ diff --git a/src/spikeinterface/core/sorting_tools.py b/src/spikeinterface/core/sorting_tools.py index 6761d607cd..51949ebc24 100644 --- a/src/spikeinterface/core/sorting_tools.py +++ b/src/spikeinterface/core/sorting_tools.py @@ -4,13 +4,17 @@ import numpy as np -from spikeinterface.core.base import BaseExtractor, unit_period_dtype +from spikeinterface.core.base import BaseExtractor, minimum_spike_dtype, unit_period_dtype from spikeinterface.core.basesorting import BaseSorting from spikeinterface.core.numpyextractors import NumpySorting numba_spec = importlib.util.find_spec("numba") if numba_spec is not None: HAVE_NUMBA = True + # Instead of importing numba directly at the module level + # like we do in lazily imported submodules, here in `core` + # we defer the import of numba until the very last moment, + # to keep the import time of `spikeinterface.core` low. else: HAVE_NUMBA = False @@ -36,7 +40,6 @@ def spike_vector_to_spike_trains(spike_vector: list[np.array], unit_ids: np.arra """ if HAVE_NUMBA: - # the trick here is to have a function getter vector_to_list_of_spiketrain = get_numba_vector_to_list_of_spiketrain() else: vector_to_list_of_spiketrain = vector_to_list_of_spiketrain_numpy @@ -79,7 +82,6 @@ def spike_vector_to_indices(spike_vector: list[np.array], unit_ids: np.array, ab """ if HAVE_NUMBA: - # the trick here is to have a function getter vector_to_list_of_spiketrain = get_numba_vector_to_list_of_spiketrain() else: vector_to_list_of_spiketrain = vector_to_list_of_spiketrain_numpy @@ -147,6 +149,174 @@ def vector_to_list_of_spiketrain_numba(sample_indices, unit_indices, num_units): return vector_to_list_of_spiketrain_numba +def reorder_spike_vector_by_unit_and_segment( + spike_vector: np.ndarray, + num_units: int, + num_segments: int, + unit_major: bool = True, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Stable reorder of a spike vector so that each (unit_index, segment_index) group is contiguous. + + Each spike is assigned to one of `num_units * num_segments` buckets, and the spikes are stably + sorted by bucket. `unit_major` selects which of the two nestings to use: + + * True -> bucket = unit_index * num_segments + segment_index, i.e. each unit's spiketrain is + compact in memory, and within a unit each segment is compact. + * False -> bucket = segment_index * num_units + unit_index, i.e. each segment is compact, and + within a segment each unit is compact. + + The sort is stable, so any ordering already present in `spike_vector` carries over to each + bucket. In particular, since a spike vector is sample_index-ascending within each segment, and + every bucket lies inside a single segment, every bucket of the output is sample_index-ascending. + + Internally calls numba if numba is installed, in which case this is a counting sort running in + O(num_spikes), adapted from Cormen, Leiserson, Rivest and Stein (CLRS) chapter 8.2. The numpy + fallback is a stable (radix) argsort on the bucket. + + Parameters + ---------- + spike_vector : np.ndarray + Structured array with dtype `minimum_spike_dtype`. + num_units : int + The number of units. Every `unit_index` must be in [0, num_units). + num_segments : int + The number of segments. Every `segment_index` must be in [0, num_segments). + unit_major : bool, default: True + Whether unit_index or segment_index is the major key. See above. + + Returns + ------- + ordered_spikes : np.ndarray + Structured array of `minimum_spike_dtype`, the same length as `spike_vector`, with the + spikes grouped by bucket. + order : np.ndarray + 1d int64 array such that `spike_vector[order]` equals `ordered_spikes`. + counts : np.ndarray + 1d int64 array of length `num_units * num_segments`, the number of spikes in each bucket, + in bucket order. + """ + num_units, num_segments = int(num_units), int(num_segments) + if num_units < 0 or num_segments < 0: + raise ValueError(f"`num_units` and `num_segments` must not be negative; got {num_units} and {num_segments}.") + num_buckets = num_units * num_segments + + # Both nestings are just a linear combination of unit_index and segment_index: + # bucket = unit_index * unit_stride + segment_index * segment_stride + unit_stride, segment_stride = (num_segments, 1) if unit_major else (1, num_units) + + num_spikes = spike_vector.size + if num_spikes == 0: + return ( + np.empty(0, dtype=minimum_spike_dtype), + np.empty(0, dtype=np.int64), + np.zeros(num_buckets, dtype=np.int64), + ) + + out_of_range_error = ( + f"`spike_vector` has a unit_index outside [0, {num_units}) or a segment_index outside [0, {num_segments})." + ) + + if HAVE_NUMBA: + reorder_spike_vector = get_numba_reorder_spike_vector() + + # These flat (num_spikes, 3) int64 views are zero-copy + in_flat = np.ascontiguousarray(spike_vector).view(np.int64).reshape(num_spikes, 3) + out_flat = np.empty((num_spikes, 3), dtype=np.int64) + 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) + if not in_range: + raise ValueError(out_of_range_error) + + ordered_spikes = out_flat.view(minimum_spike_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 + + # Must be checked before narrowing: a negative or oversized bucket would silently wrap. + if bucket_index.min() < 0 or bucket_index.max() >= num_buckets: + raise ValueError(out_of_range_error) + + counts = np.bincount(bucket_index, minlength=num_buckets).astype(np.int64, copy=False) + + # Narrow the bucket to the smallest dtype that fits. This saves memory but also time: + # The cost of the radix sort that follows scales with the width of the key, + bucket_index = bucket_index.astype(np.min_scalar_type(num_buckets - 1), copy=False) + + order = np.argsort(bucket_index, kind="stable") # radix sort, because of integer key + ordered_spikes = spike_vector[order] + return ordered_spikes, order, counts + + +def get_numba_reorder_spike_vector(): + if hasattr(get_numba_reorder_spike_vector, "_cached_numba_function"): + return get_numba_reorder_spike_vector._cached_numba_function + + 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): + """ + Stable counting-sort of a (N, 3) int64 spike-vector flat-buffer view by (unit, segment). + + Each spike's bucket is derived on the fly as + `unit_index * unit_stride + segment_index * segment_stride`, so no bucket array is needed. + + 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``. + + `out_flat`, `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`. + + Returns False if any spike falls outside [0, num_buckets), + in which case the outputs are meaningless; True otherwise. + """ + num_spikes = in_flat.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 + if bucket < 0 or bucket >= num_buckets: + return False + counts[bucket] += 1 + + # Exclusive prefix sum, giving the write cursor of each bucket. This is kept in a separate + # buffer so that `counts` survives as the per-bucket sizes, which the caller needs. + write_pos = np.empty(num_buckets, dtype=np.int64) + running = 0 + for b in range(num_buckets): + write_pos[b] = running + running += counts[b] + + # 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 + pos = write_pos[bucket] + out_flat[pos, 0] = in_flat[i, 0] + out_flat[pos, 1] = in_flat[i, 1] + out_flat[pos, 2] = in_flat[i, 2] + order[pos] = i + write_pos[bucket] = pos + 1 + + return True + + # Cache the compiled function + get_numba_reorder_spike_vector._cached_numba_function = reorder_spike_vector_numba + + return reorder_spike_vector_numba + + # stratified sampling (isi / amplitude / pca distance ? ) def random_spikes_selection( sorting: BaseSorting, diff --git a/src/spikeinterface/core/tests/test_basesorting.py b/src/spikeinterface/core/tests/test_basesorting.py index 6972ae6024..25331e6497 100644 --- a/src/spikeinterface/core/tests/test_basesorting.py +++ b/src/spikeinterface/core/tests/test_basesorting.py @@ -3,6 +3,7 @@ but check only for BaseRecording general methods. """ +import importlib.util import time import numpy as np import pytest @@ -21,7 +22,7 @@ generate_sorting, load, ) -from spikeinterface.core.base import BaseExtractor, unit_period_dtype +from spikeinterface.core.base import BaseExtractor, minimum_spike_dtype, unit_period_dtype from spikeinterface.core.testing import check_sorted_arrays_equal, check_sortings_equal @@ -153,6 +154,106 @@ def test_BaseSorting(create_cache_folder): assert sorting.get_annotation(annotation_name) == sorting_zarr_loaded.get_annotation(annotation_name) +def _make_sorting_with_shuffled_ties(num_units, num_segments, seed=42): + """Build a NumpySorting whose cotemporal spikes are in arbitrary unit_index order. + + A spike vector is only guaranteed to be segment-blocked and sample_index-ascending within each + segment; the unit_index order among spikes sharing a sample_index is unspecified (see #4606). + Building via `NumpySorting.from_unit_dict` happens to produce unit-ascending ties, so it + can't test the shuffled tie case. + """ + rng = np.random.default_rng(seed) + num_spikes = 2_000 + + # A sample range far smaller than num_spikes, so cotemporal spikes are abundant -- including + # repeats of the same (segment, sample, unit), the tie that np.lexsort itself cannot break. + spikes = np.empty(num_spikes, dtype=minimum_spike_dtype) + spikes["sample_index"] = rng.integers(0, 200, size=num_spikes) + spikes["unit_index"] = rng.integers(0, num_units, size=num_spikes) + spikes["segment_index"] = rng.integers(0, num_segments, size=num_spikes) + + # Order by segment then sample, breaking ties randomly rather than by unit_index. + spikes = spikes[np.lexsort((rng.random(num_spikes), spikes["sample_index"], spikes["segment_index"]))] + + sorting = NumpySorting(spikes, 30_000.0, np.arange(num_units)) + assert sorting.get_num_segments() == num_segments + return sorting + + +@pytest.mark.parametrize("use_numba", [True, False], ids=["numba", "numpy"]) +@pytest.mark.parametrize( + "lexsort", + [("sample_index", "segment_index", "unit_index"), ("sample_index", "unit_index", "segment_index")], +) +def test_to_reordered_spike_vector(lexsort, use_numba, monkeypatch): + """`to_reordered_spike_vector` should group spikes by (unit, segment) without disturbing them.""" + if use_numba and importlib.util.find_spec("numba") is None: + pytest.skip("numba not installed") + monkeypatch.setattr("spikeinterface.core.sorting_tools.HAVE_NUMBA", use_numba) + + num_units, num_segments = 6, 3 + sorting = _make_sorting_with_shuffled_ties(num_units, num_segments) + spikes = sorting.to_spike_vector() + + # Make sure the input genuinely violates the old full lexsort + assert not np.array_equal( + spikes, spikes[np.lexsort((spikes["unit_index"], spikes["sample_index"], spikes["segment_index"]))] + ) + + ordered_spikes, order, slices = sorting.to_reordered_spike_vector( + lexsort=lexsort, return_order=True, return_slices=True + ) + + # The buckets, in the order the requested lexsort puts them in. + unit_major = lexsort == ("sample_index", "segment_index", "unit_index") + if unit_major: + assert slices.shape == (num_units, num_segments, 2) + groups = [(u, s) for u in range(num_units) for s in range(num_segments)] + else: + assert slices.shape == (num_segments, num_units, 2) + groups = [(u, s) for s in range(num_segments) for u in range(num_units)] + masks = [(spikes["unit_index"] == u) & (spikes["segment_index"] == s) for u, s in groups] + + assert np.array_equal(ordered_spikes, np.concatenate([spikes[mask] for mask in masks])) + + # `order` must reproduce the reordering. + assert np.array_equal(spikes[order], ordered_spikes) + + # `slices` must delimit each bucket, and together tile the whole vector. + stops = np.cumsum([mask.sum() for mask in masks]) + expected_slices = np.stack([stops - [mask.sum() for mask in masks], stops], axis=-1) + assert np.array_equal(slices.reshape(-1, 2), expected_slices) + assert slices.reshape(-1, 2)[0, 0] == 0 and slices.reshape(-1, 2)[-1, 1] == spikes.size + + +@pytest.mark.parametrize( + "unit_dict, num_units", + [({"0": np.array([], dtype="int64")}, 1), ({}, 0)], + ids=["unit_with_no_spikes", "no_units"], +) +def test_to_reordered_spike_vector_empty(unit_dict, num_units): + """Empty sortings must round-trip. + + A sorting with *no units at all* is a valid degenerate case, not an error + (see `test_empty_sorting`). It simply has no buckets. + """ + sorting = NumpySorting.from_unit_dict(unit_dict, 30_000.0) + assert len(sorting.unit_ids) == num_units + + ordered_spikes, order, slices = sorting.to_reordered_spike_vector( + lexsort=("sample_index", "segment_index", "unit_index"), + return_order=True, + return_slices=True, + ) + assert ordered_spikes.size == 0 + assert order.size == 0 + assert np.array_equal(slices, np.zeros((num_units, 1, 2), dtype=np.int64)) + + # The methods that build the reordering internally must survive it too. + sorting.precompute_spike_trains() + assert len(sorting.count_num_spikes_per_unit(outputs="dict")) == num_units + + def test_npy_sorting(): sfreq = 10 spike_times_0 = { diff --git a/src/spikeinterface/core/tests/test_sorting_tools.py b/src/spikeinterface/core/tests/test_sorting_tools.py index b8c27eab18..0bb6d9a8f3 100644 --- a/src/spikeinterface/core/tests/test_sorting_tools.py +++ b/src/spikeinterface/core/tests/test_sorting_tools.py @@ -16,10 +16,20 @@ remap_unit_indices_in_vector, set_properties_after_merging, set_properties_after_splits, + reorder_spike_vector_by_unit_and_segment, ) from spikeinterface.core.base import minimum_spike_dtype +@pytest.fixture(params=[True, False], ids=["numba", "numpy"]) +def force_numba(request, monkeypatch): + """Run each test once with numba enabled (if installed) and once with the numpy fallback.""" + if request.param and importlib.util.find_spec("numba") is None: + pytest.skip("numba not installed") + monkeypatch.setattr("spikeinterface.core.sorting_tools.HAVE_NUMBA", request.param) + return request.param + + @pytest.mark.skipif( importlib.util.find_spec("numba") is None, reason="Testing `spike_vector_to_dict` requires Python package 'numba'." ) @@ -48,6 +58,61 @@ def test_spike_vector_to_indices(): ) +def _make_spike_vector(sample_indices, unit_indices, segment_indices): + spikes = np.empty(len(sample_indices), dtype=minimum_spike_dtype) + spikes["sample_index"] = sample_indices + spikes["unit_index"] = unit_indices + spikes["segment_index"] = segment_indices + return spikes + + +def test_reorder_spike_vector_by_unit_and_segment(force_numba): + # 3 units, 1 segment, so the output is simply grouped by unit. + spikes = _make_spike_vector( + sample_indices=[10, 10, 11, 12, 12, 13], + unit_indices=[2, 0, 1, 2, 0, 0], + segment_indices=[0, 0, 0, 0, 0, 0], + ) + + 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(spikes[order], ordered_spikes) + assert np.array_equal(ordered_spikes["unit_index"], [0, 0, 0, 1, 2, 2]) + # Stability: within each bucket, sample_index keeps its (ascending) input order. + assert np.array_equal(ordered_spikes["sample_index"], [10, 12, 13, 11, 10, 12]) + + +def test_reorder_spike_vector_by_unit_and_segment_raises(force_numba): + """Out-of-range indices must raise on both paths, rather than write out of bounds.""" + spikes = _make_spike_vector([0, 1, 2], [0, 1, 0], 0) + + with pytest.raises(ValueError, match="must not be negative"): + reorder_spike_vector_by_unit_and_segment(spikes, -1, 1) + + with pytest.raises(ValueError, match="outside"): + 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) + + +@pytest.mark.parametrize("num_units", [2, 300, 70_000], ids=["uint8", "uint16", "uint32"]) +def test_reorder_spike_vector_by_unit_and_segment_bucket_dtypes(monkeypatch, num_units): + """The numpy path narrows the bucket dtype to num_buckets; every width must stay correct.""" + monkeypatch.setattr("spikeinterface.core.sorting_tools.HAVE_NUMBA", False) + rng = np.random.default_rng(0) + num_spikes = 1_000 + spikes = _make_spike_vector( + sample_indices=np.arange(num_spikes), + unit_indices=rng.integers(0, num_units, size=num_spikes), + segment_indices=0, + ) + ordered_spikes, order, counts = reorder_spike_vector_by_unit_and_segment(spikes, num_units, 1) + assert np.array_equal(spikes[order], ordered_spikes) + assert np.array_equal(ordered_spikes, spikes[np.argsort(spikes["unit_index"], kind="stable")]) + assert counts.sum() == num_spikes + + def test_random_spikes_selection(): recording, sorting = generate_ground_truth_recording( durations=[20.0, 10.0], From 87cba89cbb67010283bac1aec802293f4b108c6b Mon Sep 17 00:00:00 2001 From: Graham Findlay Date: Thu, 16 Jul 2026 22:37:47 -0500 Subject: [PATCH 4/5] Key the reorder cache by the lexsort tuple, not its repr 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. --- src/spikeinterface/core/basesorting.py | 61 +++++++++---------- .../core/tests/test_basesorting.py | 1 + .../metrics/quality/misc_metrics.py | 5 +- 3 files changed, 31 insertions(+), 36 deletions(-) diff --git a/src/spikeinterface/core/basesorting.py b/src/spikeinterface/core/basesorting.py index 7b4fe31028..ec285c5e6b 100644 --- a/src/spikeinterface/core/basesorting.py +++ b/src/spikeinterface/core/basesorting.py @@ -7,6 +7,14 @@ from .base import BaseExtractor, BaseSegment, minimum_spike_dtype from .waveform_tools import has_exceeding_spikes +#: Makes each unit's spiketrain compact in memory (unit, then segment, then sample). +LEXSORT_UNIT_COMPACT = ("sample_index", "segment_index", "unit_index") +#: Makes each segment compact, units compact within a segment (segment, then unit, then sample). +LEXSORT_SEGMENT_COMPACT = ("sample_index", "unit_index", "segment_index") + +# The reorderings that `BaseSorting.to_reordered_spike_vector()` can produce. +_ALLOWED_LEXSORTS = (LEXSORT_UNIT_COMPACT, LEXSORT_SEGMENT_COMPACT) + class BaseSorting(BaseExtractor): """ @@ -166,8 +174,7 @@ def get_unit_spike_train( If True, returns spike times in seconds instead of frames use_cache : bool, default: True If True, then precompute (or use) the to_reordered_spike_vector using - lexsort=("sample_index", "segment_index", "unit_index"), which makes a spiketrain - per unit and per segment compact in memory. + lexsort=LEXSORT_UNIT_COMPACT, which makes each unit's spiketrain compact in memory. Using the cache makes the first call quite slow but then future calls are very fast. Note: if use_cache=False, but the lexsorted cache is already computed then it will be used anyway. @@ -195,7 +202,7 @@ def get_unit_spike_train( segment_index = self._check_segment_index(segment_index) - lexsort_key = ("sample_index", "segment_index", "unit_index") + lexsort_key = LEXSORT_UNIT_COMPACT if lexsort_key in self._cached_lexsorted_spike_vector.keys(): use_cache = True @@ -580,11 +587,11 @@ def count_num_spikes_per_unit(self, outputs="dict", unit_ids=None): """ # speed strategy by order - # 1. if _cached_lexsorted_spike_vector has ("sample_index", "segment_index", "unit_index") then use it and sum + # 1. if _cached_lexsorted_spike_vector has LEXSORT_UNIT_COMPACT then use it and sum # 2. if _cached_spike_vector not None then use it with np.unique() # 3. compute spikevector and do np.unique() - cache_key = ("sample_index", "segment_index", "unit_index") + cache_key = LEXSORT_UNIT_COMPACT if unit_ids is not None: assert outputs == "dict", "count_num_spikes_per_unit() with unit_ids not None works only for output='dict'" @@ -867,9 +874,9 @@ def sample_index_to_time( def precompute_spike_trains(self): """ Pre-computes and caches all spike trains for this sorting. - This is equivalent to cache lexsort ("sample_index", "segment_index", "unit_index"). + This is equivalent to cache lexsort LEXSORT_UNIT_COMPACT. """ - cache_key = ("sample_index", "segment_index", "unit_index") + cache_key = LEXSORT_UNIT_COMPACT if cache_key not in self._cached_lexsorted_spike_vector: self.to_reordered_spike_vector(lexsort=cache_key) @@ -1034,7 +1041,7 @@ def _get_spike_vector_segment_slices(self): def to_reordered_spike_vector( self, - lexsort=("sample_index", "segment_index", "unit_index"), + lexsort=LEXSORT_UNIT_COMPACT, return_order=True, return_slices=True, ): @@ -1047,21 +1054,19 @@ def to_reordered_spike_vector( Please note that the lexsort syntax is the **reverse** of natural reading. Two reorderings are supported: - - ("sample_index", "segment_index", "unit_index"): makes each unit's - spiketrain compact in memory. This is the default, and is what + - LEXSORT_UNIT_COMPACT: ("sample_index", "segment_index", "unit_index"). + Makes each unit's spiketrain compact in memory. This is the default, and is what unit-by-unit computations (e.g. isi violations) want. - - ("sample_index", "unit_index", "segment_index"): makes each segment compact, - with each unit's spiketrain compact within a segment. Rarely (if ever) used, - but might be useful when iterating segment by segment. + - LEXSORT_SEGMENT_COMPACT: ("sample_index", "unit_index", "segment_index"). + Makes each segment compact, with each unit's spiketrain compact within a segment. + Rarely (if ever) used, but might be useful when iterating segment by segment. This operation is internally cached. Parameters ---------- - lexsort : tuple, default: ("sample_index", "segment_index", "unit_index") - The requested sort order, as keys would be passed to np.lexsort. - Please note that this is the reverse natural reading order! - Must be one of the two orderings listed above. + lexsort : tuple, default: LEXSORT_UNIT_COMPACT + The requested sort order. Must be one of the two orderings listed above. return_order: bool, default: True Return the numpy array needed to sort the spike vector (given the requested sort). return_slices: bool, default: True @@ -1085,33 +1090,23 @@ def to_reordered_spike_vector( ValueError If `lexsort` is not one of the two supported orderings. """ - lexsort = tuple(lexsort) - if lexsort == ("unit_index", "sample_index", "segment_index"): raise ValueError( '`lexsort` = ("unit_index", "sample_index", "segment_index") is not supported: ' "Use `to_spike_vector()` to get the default order." ) - if lexsort not in [ - ("sample_index", "segment_index", "unit_index"), - ("sample_index", "unit_index", "segment_index"), - ]: - raise ValueError( - '`lexsort` must be ("sample_index", "segment_index", "unit_index") or ' - f'("sample_index", "unit_index", "segment_index"); got {lexsort}.' - ) - - key = str(lexsort) + if lexsort not in _ALLOWED_LEXSORTS: + raise ValueError(f"`lexsort` must be one of {_ALLOWED_LEXSORTS}; got {lexsort}.") - if key not in self._cached_lexsorted_spike_vector.keys(): + if lexsort not in self._cached_lexsorted_spike_vector.keys(): from .sorting_tools import reorder_spike_vector_by_unit_and_segment spikes = self.to_spike_vector() num_units = len(self.unit_ids) num_segments = self.get_num_segments() - unit_major = lexsort == ("sample_index", "segment_index", "unit_index") + unit_major = lexsort == LEXSORT_UNIT_COMPACT slices_shape = (num_units, num_segments) if unit_major else (num_segments, num_units) ordered_spikes, order, counts = reorder_spike_vector_by_unit_and_segment( @@ -1123,13 +1118,13 @@ def to_reordered_spike_vector( starts = stops - counts slices = np.stack([starts, stops], axis=-1).astype(np.int64, copy=False) - self._cached_lexsorted_spike_vector[key] = { + self._cached_lexsorted_spike_vector[lexsort] = { "ordered_spikes": ordered_spikes, "order": order, "slices": slices, } - cached = self._cached_lexsorted_spike_vector[key] + cached = self._cached_lexsorted_spike_vector[lexsort] out = [cached["ordered_spikes"]] if return_order: out.append(cached["order"]) diff --git a/src/spikeinterface/core/tests/test_basesorting.py b/src/spikeinterface/core/tests/test_basesorting.py index 25331e6497..b9b0561701 100644 --- a/src/spikeinterface/core/tests/test_basesorting.py +++ b/src/spikeinterface/core/tests/test_basesorting.py @@ -23,6 +23,7 @@ load, ) from spikeinterface.core.base import BaseExtractor, minimum_spike_dtype, unit_period_dtype +from spikeinterface.core.basesorting import LEXSORT_UNIT_COMPACT from spikeinterface.core.testing import check_sorted_arrays_equal, check_sortings_equal diff --git a/src/spikeinterface/metrics/quality/misc_metrics.py b/src/spikeinterface/metrics/quality/misc_metrics.py index 304d25bdbb..259da4eded 100644 --- a/src/spikeinterface/metrics/quality/misc_metrics.py +++ b/src/spikeinterface/metrics/quality/misc_metrics.py @@ -16,6 +16,7 @@ from spikeinterface.core.analyzer_extension_core import BaseMetric from spikeinterface.core import SortingAnalyzer, NumpySorting +from spikeinterface.core.basesorting import LEXSORT_UNIT_COMPACT from spikeinterface.core.template_tools import ( get_template_amplitude_on_main_channel, get_dense_templates_array, @@ -576,9 +577,7 @@ def compute_sliding_rp_violations( contamination = {} - spikes, slices = sorting.to_reordered_spike_vector( - ["sample_index", "segment_index", "unit_index"], return_order=False - ) + spikes, slices = sorting.to_reordered_spike_vector(LEXSORT_UNIT_COMPACT, return_order=False) for unit_id in unit_ids: unit_index = sorting.id_to_index(unit_id) From f1ff67cbeee83ee8acc739db40ff459877c3993f Mon Sep 17 00:00:00 2001 From: Graham Findlay Date: Fri, 7 Aug 2026 17:37:04 -0500 Subject: [PATCH 5/5] Handle spike vectors with extra fields, non-integer dtypes, and lazy 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). --- src/spikeinterface/core/sorting_tools.py | 63 ++++++++++++++----- .../core/tests/test_sorting_tools.py | 62 ++++++++++++++++++ 2 files changed, 110 insertions(+), 15 deletions(-) diff --git a/src/spikeinterface/core/sorting_tools.py b/src/spikeinterface/core/sorting_tools.py index 51949ebc24..020100a204 100644 --- a/src/spikeinterface/core/sorting_tools.py +++ b/src/spikeinterface/core/sorting_tools.py @@ -149,6 +149,28 @@ 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: + """ + 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. + + 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. + + 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. + """ + names = dtype.names + if names is None or names[:3] != tuple(name for name, _ in minimum_spike_dtype): + 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)) + + def reorder_spike_vector_by_unit_and_segment( spike_vector: np.ndarray, num_units: int, @@ -176,8 +198,11 @@ def reorder_spike_vector_by_unit_and_segment( Parameters ---------- - spike_vector : np.ndarray - Structured array with dtype `minimum_spike_dtype`. + spike_vector : np.ndarray or array-like + Structured array whose first three fields are `minimum_spike_dtype`. Extra trailing 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. num_units : int The number of units. Every `unit_index` must be in [0, num_units). num_segments : int @@ -188,14 +213,17 @@ def reorder_spike_vector_by_unit_and_segment( Returns ------- ordered_spikes : np.ndarray - Structured array of `minimum_spike_dtype`, the same length as `spike_vector`, with the - spikes grouped by bucket. + Structured array of the same dtype and length as `spike_vector`, with the spikes grouped + by bucket. order : np.ndarray 1d int64 array such that `spike_vector[order]` equals `ordered_spikes`. counts : np.ndarray 1d int64 array of length `num_units * num_segments`, the number of spikes in each bucket, in bucket order. """ + # Materialise anything array-like (e.g. a lazily zarr-backed spike vector). + spike_vector = np.asarray(spike_vector) + num_units, num_segments = int(num_units), int(num_segments) if num_units < 0 or num_segments < 0: raise ValueError(f"`num_units` and `num_segments` must not be negative; got {num_units} and {num_segments}.") @@ -208,7 +236,7 @@ def reorder_spike_vector_by_unit_and_segment( num_spikes = spike_vector.size if num_spikes == 0: return ( - np.empty(0, dtype=minimum_spike_dtype), + np.empty(0, dtype=spike_vector.dtype), np.empty(0, dtype=np.int64), np.zeros(num_buckets, dtype=np.int64), ) @@ -217,12 +245,16 @@ 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})." ) - if HAVE_NUMBA: + # 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): reorder_spike_vector = get_numba_reorder_spike_vector() - # These flat (num_spikes, 3) int64 views are zero-copy - in_flat = np.ascontiguousarray(spike_vector).view(np.int64).reshape(num_spikes, 3) - out_flat = np.empty((num_spikes, 3), dtype=np.int64) + 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) order = np.empty(num_spikes, dtype=np.int64) counts = np.empty(num_buckets, dtype=np.int64) @@ -230,7 +262,7 @@ def reorder_spike_vector_by_unit_and_segment( if not in_range: raise ValueError(out_of_range_error) - ordered_spikes = out_flat.view(minimum_spike_dtype).reshape(num_spikes) + 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. @@ -260,7 +292,9 @@ def get_numba_reorder_spike_vector(): @jit(nopython=True, nogil=True, cache=False) def reorder_spike_vector_numba(in_flat, unit_stride, segment_stride, num_buckets, out_flat, order, counts): """ - Stable counting-sort of a (N, 3) int64 spike-vector flat-buffer view by (unit, segment). + 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. Each spike's bucket is derived on the fly as `unit_index * unit_stride + segment_index * segment_stride`, so no bucket array is needed. @@ -280,7 +314,7 @@ def reorder_spike_vector_numba(in_flat, unit_stride, segment_stride, num_buckets Returns False if any spike falls outside [0, num_buckets), in which case the outputs are meaningless; True otherwise. """ - num_spikes = in_flat.shape[0] + num_spikes, num_fields = in_flat.shape # Pass 1: histogram the buckets and do bounds-check (free! we already have to make the pass) for b in range(num_buckets): @@ -303,9 +337,8 @@ def reorder_spike_vector_numba(in_flat, unit_stride, segment_stride, num_buckets for i in range(num_spikes): bucket = in_flat[i, 1] * unit_stride + in_flat[i, 2] * segment_stride pos = write_pos[bucket] - out_flat[pos, 0] = in_flat[i, 0] - out_flat[pos, 1] = in_flat[i, 1] - out_flat[pos, 2] = in_flat[i, 2] + for field in range(num_fields): + out_flat[pos, field] = in_flat[i, field] 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 0bb6d9a8f3..6d042b2770 100644 --- a/src/spikeinterface/core/tests/test_sorting_tools.py +++ b/src/spikeinterface/core/tests/test_sorting_tools.py @@ -113,6 +113,68 @@ def test_reorder_spike_vector_by_unit_and_segment_bucket_dtypes(monkeypatch, num assert counts.sum() == num_spikes +def _legacy_reorder(spikes, unit_major=True): + """The pre-counting-sort implementation, for parity testing.""" + if unit_major: + keys = (spikes["sample_index"], spikes["segment_index"], spikes["unit_index"]) + else: + keys = (spikes["sample_index"], spikes["unit_index"], spikes["segment_index"]) + order = np.lexsort(keys) + return spikes[order] + + +@pytest.mark.parametrize("unit_major", [True, False], ids=["unit_major", "segment_major"]) +def test_reorder_spike_vector_by_unit_and_segment_extra_fields(force_numba, unit_major): + """Fields beyond `minimum_spike_dtype` must survive the reorder, travelling with their spike. + + This is needed because `to_spike_vector(main_channel_indices=...)` appends a "channel_index" field. + """ + wide_dtype = minimum_spike_dtype + [("channel_index", "int64")] + num_units, num_segments, num_spikes = 5, 2, 500 + + rng = np.random.default_rng(0) + spikes = np.empty(num_spikes, dtype=wide_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) + # Tie channel_index to unit_index so a mis-shuffled column is detectable. + spikes["channel_index"] = spikes["unit_index"] * 7 + 3 + + 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["channel_index"], ordered_spikes["unit_index"] * 7 + 3) + assert np.array_equal(ordered_spikes, _legacy_reorder(spikes, unit_major=unit_major)) + assert counts.sum() == num_spikes + + +@pytest.mark.parametrize("extra_field", [("amplitude", "float32"), ("amplitude", "float64")]) +def test_reorder_spike_vector_by_unit_and_segment_non_uniform_dtype(force_numba, extra_field): + """`NumpySorting` stores whatever dtype its caller hands it. Make sure these weird spike + vectors still get reordered correctly. + """ + dtype = minimum_spike_dtype + [extra_field] + 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["amplitude"] = [1.5, -2.25, 3.75, -4.5, 5.125, 6.0] + + ordered_spikes, order, counts = reorder_spike_vector_by_unit_and_segment(spikes, 3, 1) + + assert ordered_spikes.dtype == spikes.dtype + assert np.array_equal(counts, [3, 1, 2]) + assert np.array_equal(ordered_spikes, spikes[order]) + assert np.array_equal(ordered_spikes["amplitude"], [-2.25, 5.125, 6.0, 3.75, 1.5, -4.5]) + assert np.array_equal(ordered_spikes, _legacy_reorder(spikes)) + + def test_random_spikes_selection(): recording, sorting = generate_ground_truth_recording( durations=[20.0, 10.0],