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 51ec90bf37..3372705cb0 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 @@ -581,11 +588,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'" @@ -868,9 +875,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, ): @@ -1046,31 +1053,22 @@ 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: + - 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. + - 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. - 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: LEXSORT_UNIT_COMPACT + The requested sort 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,82 +1078,59 @@ 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. - """ - lexsort = tuple(lexsort) + 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. + """ 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"), - ("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")' + if lexsort not in _ALLOWED_LEXSORTS: + raise ValueError(f"`lexsort` must be one of {_ALLOWED_LEXSORTS}; got {lexsort}.") - key = str(lexsort) + if lexsort not in self._cached_lexsorted_spike_vector.keys(): + from .sorting_tools import reorder_spike_vector_by_unit_and_segment - 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 - 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 ("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 == 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( + 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[lexsort] = { + "ordered_spikes": ordered_spikes, + "order": order, + "slices": slices, + } + + cached = self._cached_lexsorted_spike_vector[lexsort] + 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 127eae9246..6ba1b21dbd 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..b9b0561701 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,8 @@ 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.basesorting import LEXSORT_UNIT_COMPACT from spikeinterface.core.testing import check_sorted_arrays_equal, check_sortings_equal @@ -153,6 +155,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 4194f459b3..cdef4446c1 100644 --- a/src/spikeinterface/core/tests/test_sorting_tools.py +++ b/src/spikeinterface/core/tests/test_sorting_tools.py @@ -13,10 +13,20 @@ _get_ids_after_merging, generate_unit_ids_for_merge_group, remap_unit_indices_in_vector, + 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'." ) @@ -45,6 +55,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], diff --git a/src/spikeinterface/metrics/quality/misc_metrics.py b/src/spikeinterface/metrics/quality/misc_metrics.py index d1f0a78d61..d072893ab1 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)