Reorder spike vectors with a flexible-dtype numba kernel

src/spikeinterface/core/sorting_tools.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from spikeinterface.core.base import BaseExtractor, minimum_spike_dtype, unit_period_dtype
+from spikeinterface.core.base import BaseExtractor, unit_period_dtype
 from spikeinterface.core.basesorting import BaseSorting
 from spikeinterface.core.numpyextractors import NumpySorting
 
@@ -149,26 +149,18 @@ def vector_to_list_of_spiketrain_numba(sample_indices, unit_indices, num_units):
     return vector_to_list_of_spiketrain_numba
 
 
-def _is_flat_int64_view(dtype: np.dtype) -> bool:
+def _numba_can_reorder(dtype: np.dtype) -> bool:
     """
-    Whether a spike-vector dtype can be safely viewed as a flat (num_spikes, num_fields) int64
-    matrix, which is what the numba counting sort moves rows through.
+    Whether the numba counting sort can handle this spike-vector dtype.
 
-    Requires every field to be int64, packed with no padding, and the first three fields to be
-    `minimum_spike_dtype`'s in order, because the kernel addresses unit_index and segment_index
-    positionally (columns 1 and 2) rather than by name.
+    We need integer `unit_index` and `segment_index` fields and no object fields.
 
-    The all-int64 rule is deliberately stricter than correctness demands -- the kernel copies rows
-    bitwise, so any 8-byte field would in fact round-trip through the view. Keeping it narrow means
-    the kernel only ever sees the layout it is written for, and it costs nothing in practice: every
-    spike-vector dtype spikeinterface constructs is all-int64.
+    Numba can handle any other fields (e.g. extra float fields, string fields) just fine through the record copy:
     """
     names = dtype.names
-    if names is None or names[:3] != tuple(name for name, _ in minimum_spike_dtype):
+    if names is None or dtype.hasobject:
         return False
-    if dtype.itemsize != 8 * len(names):
-        return False
-    return all(dtype.fields[name][0] == np.int64 and dtype.fields[name][1] == 8 * i for i, name in enumerate(names))
+    return all(name in names and dtype.fields[name][0].kind in "iu" for name in ("unit_index", "segment_index"))
 
 
 def reorder_spike_vector_by_unit_and_segment(
@@ -245,28 +237,37 @@ def reorder_spike_vector_by_unit_and_segment(
         f"`spike_vector` has a unit_index outside [0, {num_units}) or a segment_index outside [0, {num_segments})."
     )
 
-    # The numba kernel expects an all-int64 unpadded dtype (e.g. `minimum_spike_dtype`), but it is
-    # possible that a spike vector has extra fields with other dtypes (`NumpySorting` allows that).
-    # So we check taht the numba path is safe, and anything else takes the dtype-agnostic numpy path.
-    if HAVE_NUMBA and _is_flat_int64_view(spike_vector.dtype):
+    # If the spike-vector's dtype has object fields and/or non-integer
+    # unit_index and segment_index fields, we have to use numpy.
+    # Otherwise, we can use the numba kernel.
+    if HAVE_NUMBA and _numba_can_reorder(spike_vector.dtype):
         reorder_spike_vector = get_numba_reorder_spike_vector()
 
-        num_fields = len(spike_vector.dtype.names)
-        # These flat (num_spikes, num_fields) int64 views are zero-copy
-        in_flat = np.ascontiguousarray(spike_vector).view(np.int64).reshape(num_spikes, num_fields)
-        out_flat = np.empty((num_spikes, num_fields), dtype=np.int64)
+        spike_vector = np.ascontiguousarray(spike_vector)
+        ordered_spikes = np.empty(num_spikes, dtype=spike_vector.dtype)
         order = np.empty(num_spikes, dtype=np.int64)
         counts = np.empty(num_buckets, dtype=np.int64)
 
-        in_range = reorder_spike_vector(in_flat, unit_stride, segment_stride, num_buckets, out_flat, order, counts)
+        in_range = reorder_spike_vector(
+            spike_vector,
+            spike_vector["unit_index"],
+            spike_vector["segment_index"],
+            unit_stride,
+            segment_stride,
+            num_buckets,
+            ordered_spikes,
+            order,
+            counts,
+        )
         if not in_range:
             raise ValueError(out_of_range_error)
-
-        ordered_spikes = out_flat.view(spike_vector.dtype).reshape(num_spikes)
         return ordered_spikes, order, counts
 
     # numpy fallback: a stable argsort by bucket is equivalent to the counting sort above.
-    bucket_index = spike_vector["unit_index"] * unit_stride + spike_vector["segment_index"] * segment_stride
+    # Widen to int64 first so narrow index fields (e.g. int8) don't overflow when multiplied by the stride.
+    unit_index = spike_vector["unit_index"].astype(np.int64, copy=False)
+    segment_index = spike_vector["segment_index"].astype(np.int64, copy=False)
+    bucket_index = unit_index * unit_stride + segment_index * segment_stride
 
     # Must be checked before narrowing: a negative or oversized bucket would silently wrap.
     if bucket_index.min() < 0 or bucket_index.max() >= num_buckets:
@@ -290,37 +291,41 @@ def get_numba_reorder_spike_vector():
     from numba import jit
 
     @jit(nopython=True, nogil=True, cache=False)
-    def reorder_spike_vector_numba(in_flat, unit_stride, segment_stride, num_buckets, out_flat, order, counts):
+    def reorder_spike_vector_numba(
+        spikes, unit_index, segment_index, unit_stride, segment_stride, num_buckets, out, order, counts
+    ):
         """
-        Stable counting-sort of a (N, num_fields) int64 spike-vector flat-buffer view by
-        (unit, segment). `num_fields` is 3 for `minimum_spike_dtype`, more when the spike vector
-        carries extra int64 fields; the extra columns are copied along with their spike.
+        Stable counting-sort of a structured spike vector by (unit, segment).
+
+        `unit_index` and `segment_index` are the field views (i.e., zero-copy, strided) of `spikes`.
+        they are widened to int64 so narrow fields can't overflow the bucket arithmetic.
+        Rows move as whole records (`out[pos] = spikes[i]`), so any extra fields travel with their spike.
 
         Each spike's bucket is derived on the fly as
         `unit_index * unit_stride + segment_index * segment_stride`, so no bucket array is needed.
 
         Two O(N) passes:
           1. histogram the buckets into `counts`,
           2. cumulative-sum to per-bucket write positions, then scatter each
-             row of `in_flat` to its destination in `out_flat` and record the
-             source index in `order` so that ``in[order] == out``.
+             record of `spikes` to its destination in `out` and record the
+             source index in `order` so that ``spikes[order] == out``.
 
-        `out_flat`, `order` and `counts` are filled in place.
+        `out`, `order` and `counts` are filled in place.
 
         Stability: within each bucket, rows keep their input order, so any
-        ordering already present in `in_flat` (e.g. ascending sample_index
-        within a (segment, unit) group) carries over to `out_flat`.
+        ordering already present in `spikes` (e.g. ascending sample_index
+        within a (segment, unit) group) carries over to `out`.
 
         Returns False if any spike falls outside [0, num_buckets),
         in which case the outputs are meaningless; True otherwise.
         """
-        num_spikes, num_fields = in_flat.shape
+        num_spikes = spikes.shape[0]
 
         # Pass 1: histogram the buckets and do bounds-check (free! we already have to make the pass)
         for b in range(num_buckets):
             counts[b] = 0
         for i in range(num_spikes):
-            bucket = in_flat[i, 1] * unit_stride + in_flat[i, 2] * segment_stride
+            bucket = np.int64(unit_index[i]) * unit_stride + np.int64(segment_index[i]) * segment_stride
             if bucket < 0 or bucket >= num_buckets:
                 return False
             counts[bucket] += 1
@@ -335,10 +340,9 @@ def reorder_spike_vector_numba(in_flat, unit_stride, segment_stride, num_buckets
 
         # Pass 2: scatter each spike into its bucket, recording where it came from.
         for i in range(num_spikes):
-            bucket = in_flat[i, 1] * unit_stride + in_flat[i, 2] * segment_stride
+            bucket = np.int64(unit_index[i]) * unit_stride + np.int64(segment_index[i]) * segment_stride
             pos = write_pos[bucket]
-            for field in range(num_fields):
-                out_flat[pos, field] = in_flat[i, field]
+            out[pos] = spikes[i]
             order[pos] = i
             write_pos[bucket] = pos + 1
 

src/spikeinterface/core/tests/test_sorting_tools.py

@@ -94,6 +94,9 @@ def test_reorder_spike_vector_by_unit_and_segment_raises(force_numba):
         reorder_spike_vector_by_unit_and_segment(spikes, 1, 1)  # unit_index 1 >= num_units
     with pytest.raises(ValueError, match="outside"):
         reorder_spike_vector_by_unit_and_segment(_make_spike_vector([0], [0], [5]), 1, 1)
+    # numba wraps a negative index, so check the lower bound
+    with pytest.raises(ValueError, match="outside"):
+        reorder_spike_vector_by_unit_and_segment(_make_spike_vector([0], [-1], [0]), 1, 1)
 
 
 @pytest.mark.parametrize("num_units", [2, 300, 70_000], ids=["uint8", "uint16", "uint32"])
@@ -175,6 +178,79 @@ def test_reorder_spike_vector_by_unit_and_segment_non_uniform_dtype(force_numba,
     assert np.array_equal(ordered_spikes, _legacy_reorder(spikes))
 
 
+NARROW_INDEX_DTYPES = {
+    "int64/int32/int8": np.dtype([("sample_index", "int64"), ("unit_index", "int32"), ("segment_index", "int8")]),
+    "uint16/int8": np.dtype([("sample_index", "int64"), ("unit_index", "uint16"), ("segment_index", "int8")]),
+}
+
+
+def _make_random_spikes(dtype, num_units, num_segments, num_spikes, seed=0):
+    """Spikes in valid spike-vector order (segment-blocked, sample-ascending) with random units."""
+    rng = np.random.default_rng(seed)
+    spikes = np.empty(num_spikes, dtype=dtype)
+    segment_indices = np.sort(rng.integers(0, num_segments, size=num_spikes))
+    spikes["segment_index"] = segment_indices
+    for segment_index in range(num_segments):
+        in_segment = segment_indices == segment_index
+        spikes["sample_index"][in_segment] = np.sort(rng.integers(0, 1_000, size=in_segment.sum()))
+    spikes["unit_index"] = rng.integers(0, num_units, size=num_spikes)
+    return spikes
+
+
+@pytest.mark.parametrize("unit_major", [True, False], ids=["unit_major", "segment_major"])
+@pytest.mark.parametrize("dtype", list(NARROW_INDEX_DTYPES.values()), ids=list(NARROW_INDEX_DTYPES.keys()))
+def test_reorder_spike_vector_by_unit_and_segment_narrow_index_dtypes(force_numba, dtype, unit_major):
+    """Make sure that narrow unit_index / segment_index fields use numba."""
+    # 200 units x 2 segments puts the bucket index outside int8 range,
+    # so the kernel has to widen the fields before the bucket math.
+    num_units, num_segments = 200, 2
+
+    spikes = _make_random_spikes(dtype, num_units, num_segments, num_spikes=2_000)
+    if force_numba:
+        from spikeinterface.core.sorting_tools import _numba_can_reorder
+
+        assert _numba_can_reorder(spikes.dtype)
+
+    ordered_spikes, order, counts = reorder_spike_vector_by_unit_and_segment(
+        spikes, num_units, num_segments, unit_major=unit_major
+    )
+
+    assert ordered_spikes.dtype == spikes.dtype
+    assert np.array_equal(ordered_spikes, spikes[order])
+    assert np.array_equal(ordered_spikes, _legacy_reorder(spikes, unit_major=unit_major))
+    assert counts.sum() == spikes.size
+
+
+def test_reorder_spike_vector_by_unit_and_segment_object_field_falls_back():
+    """Make sure spike vectors with object fields use numpy"""
+    dtype = minimum_spike_dtype + [("label", "O")]
+    spikes = np.empty(6, dtype=dtype)
+    spikes["sample_index"] = [10, 10, 11, 12, 12, 13]
+    spikes["unit_index"] = [2, 0, 1, 2, 0, 0]
+    spikes["segment_index"] = 0
+    spikes["label"] = list("abcdef")
+
+    ordered_spikes, order, counts = reorder_spike_vector_by_unit_and_segment(spikes, 3, 1)
+
+    assert np.array_equal(counts, [3, 1, 2])
+    assert np.array_equal(ordered_spikes, spikes[order])
+    assert list(ordered_spikes["label"]) == ["b", "e", "f", "c", "a", "d"]
+
+
+def test_numba_can_reorder():
+    from spikeinterface.core.sorting_tools import _numba_can_reorder
+
+    assert _numba_can_reorder(np.dtype(minimum_spike_dtype))
+    assert _numba_can_reorder(np.dtype([("sample_index", "int64"), ("unit_index", "int32"), ("segment_index", "int8")]))
+
+    assert not _numba_can_reorder(np.dtype(minimum_spike_dtype + [("label", "O")]))
+    assert not _numba_can_reorder(np.dtype([("sample_index", "int64"), ("unit_index", "int64")]))
+    assert not _numba_can_reorder(
+        np.dtype([("sample_index", "int64"), ("unit_index", "float64"), ("segment_index", "int64")])
+    )
+    assert not _numba_can_reorder(np.dtype("int64"))
+
+
 def test_random_spikes_selection():
     recording, sorting = generate_ground_truth_recording(
         durations=[20.0, 10.0],