From 056935bb7b5b69ac9e2759af58dae99785c8006e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20Royeth?= Date: Thu, 10 Sep 2026 13:36:54 -0300 Subject: [PATCH] Fix ComputeACG3D soft merge indexing bug that swaps kept units' data Co-authored-by: Arthur031221 <124417490+Arthur031221@users.noreply.github.com> --- .../postprocessing/correlograms.py | 10 +++-- .../tests/test_extension_merges.py | 45 +++++++++++++++++++ 2 files changed, 52 insertions(+), 3 deletions(-) diff --git a/src/spikeinterface/postprocessing/correlograms.py b/src/spikeinterface/postprocessing/correlograms.py index 13431a8dbb..fa63230a18 100644 --- a/src/spikeinterface/postprocessing/correlograms.py +++ b/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, diff --git a/src/spikeinterface/postprocessing/tests/test_extension_merges.py b/src/spikeinterface/postprocessing/tests/test_extension_merges.py index fa0310af5c..e4f858223e 100644 --- a/src/spikeinterface/postprocessing/tests/test_extension_merges.py +++ b/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)