Fix ComputeACG3D soft merge indexing bug that swaps kept units' data

src/spikeinterface/postprocessing/correlograms.py

@@ -1235,16 +1235,20 @@ def _merge_extension_data(
 
         new_unit_ids_indices = new_sorting.ids_to_indices(new_unit_ids)
         old_unit_ids = [unit_id for unit_id in new_sorting_analyzer.unit_ids if unit_id not in new_unit_ids]
-        old_unit_ids_indices = new_sorting.ids_to_indices(old_unit_ids)
+        # source indices are looked up in the sorting the data was computed with, not the resulting one
+        old_unit_ids_indices_in_new = new_sorting.ids_to_indices(old_unit_ids)
+        old_unit_ids_indices_in_old = self.sorting_analyzer.sorting.ids_to_indices(old_unit_ids)
 
         new_acgs_3d = np.zeros((len(new_sorting.unit_ids), acgs_3d.shape[1], acgs_3d.shape[2]))
         new_firing_quantiles = np.zeros((len(new_sorting.unit_ids), firing_rate_quantiles.shape[1]))
 
         new_acgs_3d[new_unit_ids_indices, :, :] = acgs_3d
-        new_acgs_3d[old_unit_ids_indices, :, :] = self.data["acgs_3d"][old_unit_ids_indices, :, :]
+        new_acgs_3d[old_unit_ids_indices_in_new, :, :] = self.data["acgs_3d"][old_unit_ids_indices_in_old, :, :]
 
         new_firing_quantiles[new_unit_ids_indices, :] = firing_rate_quantiles
-        new_firing_quantiles[old_unit_ids_indices, :] = self.data["firing_quantiles"][old_unit_ids_indices, :]
+        new_firing_quantiles[old_unit_ids_indices_in_new, :] = self.data["firing_quantiles"][
+            old_unit_ids_indices_in_old, :
+        ]
 
         new_data = dict(
             acgs_3d=new_acgs_3d,

src/spikeinterface/postprocessing/tests/test_extension_merges.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pytest
 
 from spikeinterface.core import generate_ground_truth_recording, create_sorting_analyzer
 
@@ -59,3 +60,47 @@ def test_correlograms_merge():
 
     recomputed_ccgs_not_censored = merged_sorting_analyzer_not_censored.compute("correlograms").get_data()
     assert np.all(computed_ccgs_not_censored[0] == recomputed_ccgs_not_censored[0])
+
+
+@pytest.mark.parametrize("sparse", [False, True])
+def test_acgs_3d_merge(sparse):
+    """
+    A 3D-ACG only depends on the spike train of its own unit, so a unit that takes no part in
+    a merge should keep the 3D-ACG and firing rate quantiles it had before the merge. This test
+    checks that a soft merge gives the same data, for every such unit, as recomputing the 3D-ACGs
+    from scratch on the merged analyzer -- for merge groups at the start, in the middle, and at
+    the end of the unit list, where a kept unit's index does or does not shift -- and for both a
+    sparse and a dense analyzer, on a multi-segment recording.
+    """
+
+    recording, sorting = generate_ground_truth_recording(durations=[10.0, 10.0], num_units=6, seed=2205)
+
+    sorting_analyzer = create_sorting_analyzer(recording=recording, sorting=sorting, sparse=sparse)
+    sorting_analyzer.compute("acgs_3d")
+
+    trial_merges = [
+        [["0", "1"]],
+        [["2", "3"]],
+        [["4", "5"]],
+        [["0", "1"], ["3", "4"]],
+    ]
+
+    for new_id_strategy in ["append", "take_first"]:
+        for merge_unit_groups in trial_merges:
+
+            merged_sorting_analyzer = sorting_analyzer.merge_units(
+                merge_unit_groups=merge_unit_groups, new_id_strategy=new_id_strategy
+            )
+            # bins is excluded from this comparison: it is independently wrong after a merge
+            # (tracks the pre-merge unit count, see spikeinterface/spikeinterface#4737) in a way
+            # this fix does not touch -- a separate defect in how ComputeACG3D builds "bins".
+            computed_acgs_3d, computed_quantiles, _computed_bins = merged_sorting_analyzer.get_extension(
+                "acgs_3d"
+            ).get_data()
+
+            recomputed_acgs_3d, recomputed_quantiles, _recomputed_bins = merged_sorting_analyzer.compute(
+                "acgs_3d"
+            ).get_data()
+
+            assert np.array_equal(computed_acgs_3d, recomputed_acgs_3d)
+            assert np.array_equal(computed_quantiles, recomputed_quantiles)