Wrong correlograms from method="numpy" when a unit has no spikes in a segment

[Arthur031221]

_compute_correlograms_numpy sizes its accumulator from len(sorting.unit_ids) and adds each segment into it. correlogram_for_one_segment sizes the array it returns from len(np.unique(spike_unit_indices)), the units that fire in that segment. A unit with no spikes in a segment is missing from that segment's np.unique, so the two sizes always disagree, and most of the time that raises ValueError.

The case that goes wrong quietly needs a sorting of more than one unit, a segment in which exactly one of them fires, and at least one count from that unit in that segment. num_units is then 1 for the segment while the accumulator is the full size, so the segment result is (1, 1, num_bins) or (1, num_bins) and numpy broadcasts it in rather than refusing. compute_auto_correlograms builds the row for unit index 0 whatever the firing unit actually is, and a (1, num_bins) result always broadcasts, so that segment cannot raise. It goes wrong in one of two directions: if the firing unit is index 0 its counts are added to every unit, and if it is not, the row is empty and that unit's counts for the segment are dropped. compute_correlograms broadcasts the same way when the firing unit is index 0 and raises when it is not, because the unit coordinate is then out of range. Neither happens when the sorting has only one unit, because one firing unit is then the full size and there is nothing to broadcast, nor when that unit produces no counts in the segment, because broadcasting zeros changes nothing.

method="numba" returned the expected counts in every case I tried. numba is not a core dependency, so pip install "spikeinterface[full]" gets it and method="auto" resolves to numba there, but a plain pip install spikeinterface does not and auto falls back to numpy. method="numpy" is affected either way.

Reproduction

Python 3.12.3, numpy 2.5.2, released spikeinterface 0.104.8 from PyPI, numba not installed:

import numpy as np
from spikeinterface.core import NumpySorting
from spikeinterface.postprocessing import compute_correlograms, compute_auto_correlograms

# Three units, two segments. Every unit fires somewhere, so no unit is empty.
# Segment 1 has spikes from unit 0 only. Unit 2 has a single spike in the whole
# sorting, so its autocorrelogram has to be all zeros.
sorting = NumpySorting.from_samples_and_labels(
    samples_list=[np.array([0, 120, 240, 480]), np.array([0, 150, 300, 600])],
    labels_list=[np.array([0, 1, 2, 1]), np.array([0, 0, 0, 0])],
    sampling_frequency=30000.0,
    unit_ids=[0, 1, 2],
)
print("spikes per unit:", sorting.count_num_spikes_per_unit())
print("non-empty units:", sorting.get_non_empty_unit_ids())

for method in ("numpy", "numba"):
    try:
        acg, bins = compute_auto_correlograms(sorting, window_ms=40.0, bin_ms=5.0, method=method)
        print(f"{method:5s} autocorrelogram of unit 2:", acg[2])
    except Exception as exc:
        print(f"{method:5s} ->", f"{type(exc).__name__}: {exc}")

# The same shape of sorting, except that segment 1 now has two units firing
# instead of one.
sorting = NumpySorting.from_samples_and_labels(
    samples_list=[np.array([0, 120, 240, 480]), np.array([0, 150, 300, 600])],
    labels_list=[np.array([0, 1, 2, 1]), np.array([0, 2, 0, 2])],
    sampling_frequency=30000.0,
    unit_ids=[0, 1, 2],
)
for method in ("numpy", "numba"):
    try:
        ccg, bins = compute_correlograms(sorting, window_ms=40.0, bin_ms=5.0, method=method)
        print(f"{method:5s} cross-correlograms {ccg.shape} sum {int(ccg.sum())}")
    except Exception as exc:
        print(f"{method:5s} ->", f"{type(exc).__name__}: {exc}")

spikes per unit: {np.int64(0): np.int64(5), np.int64(1): np.int64(2), np.int64(2): np.int64(1)}
non-empty units: [0 1 2]
numpy autocorrelogram of unit 2: [1 1 2 2 0 2 2 1]
numba -> AssertionError: numba version of this function requires installation of numba
numpy -> ValueError: invalid entry in coordinates array
numba -> AssertionError: numba version of this function requires installation of numba

Unit 2 has one spike in the whole sorting, so its autocorrelogram cannot contain a count at all. It comes back holding 11. What it holds is segment 1's result, which is unit 0's autocorrelogram in a (1, num_bins) array, added to every row.

The same script on main at d365714a4 with numba 0.67.0 installed, so the two methods can be compared:

spikes per unit: {np.int64(0): np.int64(5), np.int64(1): np.int64(2), np.int64(2): np.int64(1)}
non-empty units: [0 1 2]
numpy autocorrelogram of unit 2: [1 1 2 2 0 2 2 1]
numba autocorrelogram of unit 2: [0 0 0 0 0 0 0 0]
numpy -> ValueError: invalid entry in coordinates array
numba cross-correlograms (3, 3, 8) sum 23

remove_empty_units() is not a way around this. Its docstring says a unit counts as empty only when it has no spikes in all segments, so a unit silent in one segment out of several survives it, which is why get_non_empty_unit_ids() prints [0 1 2] above.

It also reaches the quality metrics

compute_correlograms is not the only caller. slidingRP_violations picks the method itself at metrics/quality/misc_metrics.py:1715 and calls _compute_correlograms_numpy directly at :1723, on the one-unit sorting that compute_sliding_rp_violations builds per unit at :597. NumpySorting takes its segment count from the last spike (core/numpyextractors.py:261-264), so a unit that is silent in a segment before one in which it fires gives a one-unit sorting with an empty segment, and the two sizes disagree there as well. Silence only in trailing segments is truncated away and is not affected. Three units, two segments, unit 2 firing in segment 1 only, with numba absent:

compute_sliding_rp_violations -> ValueError: non-broadcastable output operand with shape (1,1,8570) doesn't match the broadcast shape (0,0,8570)

With the change it returns {'0': nan, '1': 0.35, '2': nan}. #4409 was computing the correlograms extension rather than this metric, so this is a second place the same size disagreement surfaces, not the traceback in that report.

Cause

_compute_correlograms_numpy (src/spikeinterface/postprocessing/correlograms.py:513-526):

num_units = len(sorting.unit_ids)
...
correlograms = np.zeros((num_units, num_units, num_bins), dtype="int64")

for seg_index in range(num_seg):
    ...
    c0 = correlogram_for_one_segment(spike_times, spike_unit_indices, window_size, bin_size)

    correlograms += c0

correlogram_for_one_segment (:575):

num_units = len(np.unique(spike_unit_indices))

correlograms = np.zeros((num_units, num_units, num_bins), dtype="int64")

One row per unit of the sorting is what everything around it assumes. The function's own docstring at :551-553 describes the return as a "(num_units, num_units, num_bins) array of correlograms between all units", _compute_correlograms_numba allocates from len(sorting.unit_ids) at :667 and has its kernel fill that array in place, and ComputeCorrelograms._select_units_extension_data at :102 does self.data["ccgs"][unit_indices][:, unit_indices], so row i has to be unit index i.

auto_correlogram_for_one_segment carries the same num_units line at :1017, and there it also picks the wrong rows: the loop is for unit_ind in range(num_units) with unit_mask = spike_unit_indices == unit_ind at :1021-1022, so counting num_units from the spikes present makes it look for unit indices 0..k-1 whatever the firing units actually are. If only unit 2 fires, it computes one row for unit 0, finds no spikes, and returns zeros.

To find out which outcome happens when, I enumerated every choice of silent units for sortings of 2, 3, 4 and 5 units, one segment, both functions, 104 configurations, and compared numpy against numba on each. At d365714a4:

RAISE-broadcast  44
RAISE-coords     42
WRONG            18

Every one of the 18 has exactly one firing unit: 14 are compute_auto_correlograms and 4 are compute_correlograms where the firing unit is index 0. In every configuration in this sweep that one firing unit produces at least one count, which is why none of the 18 is a harmless zero broadcast. Split by direction, 8 report counts that no unit produced and 10 lose a unit's counts entirely, and all 10 of those are compute_auto_correlograms with the firing unit at an index other than 0, where the numpy result for that unit is all zeros against 26 counts from numba. That is 104 constructed configurations of one segment, not a proof about every input.

Where the line came from

76bc7674c, merged as #1197, moved the per-segment body of compute_correlograms_numpy into a new correlogram_for_one_segment, which derives its own num_units from len(np.unique(spike_labels)) while the caller kept len(sorting.unit_ids). Before that commit the array was allocated once from len(sorting.unit_ids) and filled in place across segments, which is the shape the numba path still has. f11a45448, merged as #4307, later copied the same line into auto_correlogram_for_one_segment.

#4409 reports ValueError: invalid entry in coordinates array with a traceback that ends inside correlogram_for_one_segment, while computing the correlograms extension on a sorting produced by the Kilosort GUI. The reporter closed it ten minutes after opening it and there are no comments on it. I do not have that data set, so I cannot say the cause there was the same one.

Suggested fix

Pass the size the caller already knows:

def correlogram_for_one_segment(spike_times, spike_unit_indices, window_size, bin_size, num_units=None):
    ...
    num_bins, num_half_bins = _compute_num_bins(window_size, bin_size)
    if num_units is None:
        num_units = int(np.max(spike_unit_indices)) + 1 if len(spike_unit_indices) > 0 else 0

c0 = correlogram_for_one_segment(spike_times, spike_unit_indices, window_size, bin_size, num_units=num_units)

auto_correlogram_for_one_segment and _compute_auto_correlograms_numpy take the same two changes, and each new parameter gets a docstring entry. That is 20 lines added and 6 removed in correlograms.py, plus 181 added in test_correlograms.py for five regression tests, each parametrised over numpy and numba.

Both segment functions are in spikeinterface.postprocessing's __init__, so this is a keyword with a default rather than a new positional argument. The default is not a no-op for anyone calling them directly: today they return len(np.unique(labels)) rows, and with this change they return max(labels) + 1 rows, which differ whenever the labels are not exactly 0..k-1. In the cases I tried that turns a raise, or a result with too few rows to index by label, into one with a row per label index, but it is still a change in what a direct caller gets back.

Measured on Linux, Python 3.12.3, numpy 2.5.2, numba 0.67.0, in one clone in one session:

• pytest src/spikeinterface/postprocessing/tests/test_correlograms.py: 40 passed, 40 skipped at d365714a4, and 50 passed, 40 skipped with the change and the five tests. All 40 skips are npyx not installed.
• reverting correlograms.py alone and keeping the tests: 5 failed, 45 passed, 40 skipped. The five failures are the numpy parametrisations. Their numba parametrisations pass either way.
• pytest src/spikeinterface/postprocessing: 97 passed, 40 skipped at d365714a4 and 107 passed, 40 skipped with the change. The ten extra passes are the five tests times two methods, and nothing fails either way.
• pytest src/spikeinterface/core/tests: 323 passed, 5 skipped before and after, and the same 15 tests in comparison, exporters and curation fail or error before and after.

For sortings where every unit fires in every segment the numpy output does not change. A sha256 over every array returned by five such sortings crossed with three window and bin size pairs is d8b8f06557f7278779589d3502aa6306ac73e9efe56259c7745b324bed6fcb2d at d365714a4 and the same with the change applied.

I am not sure the keyword is the shape you want. The alternative is for _compute_correlograms_numpy to pass its accumulator down and have the segment function fill it in place, which is what the numba path does at :682-690. Before #1197 there was no segment function and the loop filled one accumulator in place, which is the same shape. That leaves both public signatures alone and adds no allocation, at the cost of turning two pure functions into ones that write through an argument. I went with the keyword because it is the smaller change, but the in-place version is the one that matches the numba path.