BUG: preserve correlogram unit axes across segments

src/spikeinterface/postprocessing/correlograms.py

@@ -521,14 +521,14 @@ def _compute_correlograms_numpy(sorting, window_size, bin_size):
         spike_times = spikes[seg_index]["sample_index"]
         spike_unit_indices = spikes[seg_index]["unit_index"]
 
-        c0 = correlogram_for_one_segment(spike_times, spike_unit_indices, window_size, bin_size)
+        c0 = correlogram_for_one_segment(spike_times, spike_unit_indices, window_size, bin_size, num_units=num_units)
 
         correlograms += c0
 
     return correlograms
 
 
-def correlogram_for_one_segment(spike_times, spike_unit_indices, window_size, bin_size):
+def correlogram_for_one_segment(spike_times, spike_unit_indices, window_size, bin_size, num_units):
     """
     A very well optimized algorithm for the cross-correlation of
     spike trains, copied from the Phy package, written by Cyrille Rossant.
@@ -545,6 +545,8 @@ def correlogram_for_one_segment(spike_times, spike_unit_indices, window_size, bi
         The window size over which to perform the cross-correlation, in samples
     bin_size : int
         The size of which to bin lags, in samples.
+    num_units : int
+        Number of units in the complete sorting.
 
     Returns
     -------
@@ -572,8 +574,6 @@ def correlogram_for_one_segment(spike_times, spike_unit_indices, window_size, bi
     match within the window size.
     """
     num_bins, num_half_bins = _compute_num_bins(window_size, bin_size)
-    num_units = len(np.unique(spike_unit_indices))
-
     correlograms = np.zeros((num_units, num_units, num_bins), dtype="int64")
 
     # At a given shift, the mask precises which spikes have matching spikes
@@ -963,14 +963,16 @@ def _compute_auto_correlograms_numpy(sorting, window_size, bin_size):
         spike_times = spikes[seg_index]["sample_index"]
         spike_unit_indices = spikes[seg_index]["unit_index"]
 
-        c0 = auto_correlogram_for_one_segment(spike_times, spike_unit_indices, window_size, bin_size)
+        c0 = auto_correlogram_for_one_segment(
+            spike_times, spike_unit_indices, window_size, bin_size, num_units=num_units
+        )
 
         correlograms += c0
 
     return correlograms
 
 
-def auto_correlogram_for_one_segment(spike_times, spike_unit_indices, window_size, bin_size):
+def auto_correlogram_for_one_segment(spike_times, spike_unit_indices, window_size, bin_size, num_units):
     """
     A very well optimized algorithm for the auto-correlation of
     spike trains, copied from the Phy package, written by Cyrille Rossant.
@@ -987,6 +989,8 @@ def auto_correlogram_for_one_segment(spike_times, spike_unit_indices, window_siz
         The window size over which to perform the cross-correlation, in samples
     bin_size : int
         The size of which to bin lags, in samples.
+    num_units : int
+        Number of units in the complete sorting.
 
     Returns
     -------
@@ -1014,8 +1018,6 @@ def auto_correlogram_for_one_segment(spike_times, spike_unit_indices, window_siz
     match within the window size.
     """
     num_bins, num_half_bins = _compute_num_bins(window_size, bin_size)
-    num_units = len(np.unique(spike_unit_indices))
-
     correlograms = np.zeros((num_units, num_bins), dtype="int64")
 
     for unit_ind in range(num_units):

src/spikeinterface/postprocessing/tests/test_correlograms.py

@@ -24,9 +24,11 @@
     _compute_correlograms_on_sorting,
     _compute_auto_correlograms_on_sorting,
     _make_bins,
+    auto_correlogram_for_one_segment,
     compute_acgs_3d,
     compute_correlograms,
     compute_auto_correlograms,
+    correlogram_for_one_segment,
 )
 from spikeinterface.postprocessing.tests.common_extension_tests import AnalyzerExtensionCommonTestSuite
 
@@ -135,6 +137,57 @@ def test_equal_results_correlograms(window_and_bin_ms):
     assert np.array_equal(result_numpy, result_numba)
 
 
+def test_segment_helpers_preserve_explicit_unit_count():
+    # Silent units must retain their positions in the complete sorting.
+    samples = np.array([0, 1000])
+    labels = np.array([0, 2])
+    for helper, shape in (
+        (correlogram_for_one_segment, (3, 3, 20)),
+        (auto_correlogram_for_one_segment, (3, 20)),
+    ):
+        result = helper(samples, labels, window_size=100, bin_size=10, num_units=3)
+        np.testing.assert_array_equal(result, np.zeros(shape, dtype="int64"))
+
+
+@pytest.mark.parametrize("num_units", [2, 3])
+def test_equal_results_when_units_are_silent_in_a_segment(num_units):
+    """Keep global unit coordinates when a segment has silent units."""
+    sorting = NumpySorting.from_samples_and_labels(
+        samples_list=[
+            np.array([0, 40]) if num_units == 2 else np.array([0, 20, 40, 60]),
+            np.array([0, 20, 40, 60]),
+            np.array([], dtype="int64"),
+        ],
+        labels_list=[
+            np.array([0, 0]) if num_units == 2 else np.array([0, 2, 0, 2]),
+            np.array([0, 1, 0, 1]),
+            np.array([], dtype="int64"),
+        ],
+        sampling_frequency=1000.0,
+        unit_ids=np.arange(num_units),
+    )
+
+    ccg_numpy, _ = compute_correlograms(sorting, window_ms=100.0, bin_ms=10.0, method="numpy")
+    acg_numpy, _ = compute_auto_correlograms(sorting, window_ms=100.0, bin_ms=10.0, method="numpy")
+
+    # Each active unit has two spikes 40 ms apart. Unit 0 occurs in both
+    # nonempty segments, and their timestamps must never be correlated.
+    expected = np.zeros((num_units, num_units, 10), dtype="int64")
+    expected[0, 0, [1, 9]] = 2
+    for unit_index in range(1, num_units):
+        expected[unit_index, unit_index, [1, 9]] = 1
+        expected[0, unit_index, [3, 7]] = [2, 1]
+        expected[unit_index, 0, [3, 7]] = [1, 2]
+    np.testing.assert_array_equal(ccg_numpy, expected)
+    np.testing.assert_array_equal(acg_numpy, expected[np.arange(num_units), np.arange(num_units)])
+
+    if HAVE_NUMBA:
+        ccg_numba, _ = compute_correlograms(sorting, window_ms=100.0, bin_ms=10.0, method="numba")
+        acg_numba, _ = compute_auto_correlograms(sorting, window_ms=100.0, bin_ms=10.0, method="numba")
+        assert np.array_equal(ccg_numpy, ccg_numba)
+        assert np.array_equal(acg_numpy, acg_numba)
+
+
 @pytest.mark.skipif(not HAVE_NUMBA, reason="Numba not available")
 @pytest.mark.parametrize("window_and_bin_ms", [(60.0, 2.0), (3.57, 1.6421)])
 def test_equal_results_fast_correlograms(window_and_bin_ms):