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=None):
     """
     A very well optimized algorithm for the cross-correlation of
     spike trains, copied from the Phy package, written by Cyrille Rossant.
@@ -545,6 +545,9 @@ 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 or None
+        Number of units in the complete sorting. If ``None``, infer it from
+        the largest unit index in ``spike_unit_indices``.
 
     Returns
     -------
@@ -572,7 +575,8 @@ 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))
+    if num_units is None:
+        num_units = int(np.max(spike_unit_indices)) + 1 if spike_unit_indices.size else 0
 
     correlograms = np.zeros((num_units, num_units, num_bins), dtype="int64")
 
@@ -963,14 +967,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=None):
     """
     A very well optimized algorithm for the auto-correlation of
     spike trains, copied from the Phy package, written by Cyrille Rossant.
@@ -987,6 +993,9 @@ 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 or None
+        Number of units in the complete sorting. If ``None``, infer it from
+        the largest unit index in ``spike_unit_indices``.
 
     Returns
     -------
@@ -1014,7 +1023,8 @@ 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))
+    if num_units is None:
+        num_units = int(np.max(spike_unit_indices)) + 1 if spike_unit_indices.size else 0
 
     correlograms = np.zeros((num_units, num_bins), dtype="int64")
 

src/spikeinterface/postprocessing/tests/test_correlograms.py

@@ -135,6 +135,45 @@ def test_equal_results_correlograms(window_and_bin_ms):
     assert np.array_equal(result_numpy, result_numba)
 
 
+@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):