Soft merges give unmerged units another unit's `acgs_3d` data

[Arthur031221]

SortingAnalyzer.merge_units with the default merging_mode="soft" carries the acgs_3d
extension forward rather than recomputing it. The units that took no part in the merge should
keep the 3D ACG they already had, and they do keep a 3D ACG, but the rows are read out of the
old data at the wrong index, so any such unit whose index the merge changed comes back holding
another unit's 3D ACG and firing rate quantiles. Nothing raises, and both arrays come back the
shape they should be.

merging_mode="hard" is not affected, because core/sortinganalyzer.py:1702-1703 puts hard
merges into recompute_dict instead of calling extension.merge, so _merge_extension_data
never runs.

Reproduction

Against d365714a4b48b1dfb67c2d380529f402dd9e54e1, Python 3.12.3, numpy 2.5.2, Linux.

import numpy as np
from spikeinterface.core import generate_ground_truth_recording, create_sorting_analyzer
from spikeinterface.postprocessing import ComputeACG3D  # noqa: F401  (registers acgs_3d)

rec, sort = generate_ground_truth_recording(durations=[20.0], num_units=6, seed=2205)
analyzer = create_sorting_analyzer(recording=rec, sorting=sort, sparse=False)
analyzer.compute("acgs_3d")
before_acgs, _, _ = analyzer.get_extension("acgs_3d").get_data()

merged = analyzer.merge_units(merge_unit_groups=[["1", "2"]], new_id_strategy="append")
after_acgs, _, _ = merged.get_extension("acgs_3d").get_data()
fresh_acgs, _, _ = merged.compute("acgs_3d").get_data()

print("unit ids before merge:", [str(u) for u in analyzer.unit_ids])
print("unit ids after merge :", [str(u) for u in merged.unit_ids])
print()
print("unit  carried 3D ACG equals a recompute   carried 3D ACG is actually the one of")
for i, unit_id in enumerate(merged.unit_ids):
    same = np.array_equal(after_acgs[i], fresh_acgs[i])
    source = [str(analyzer.unit_ids[j]) for j in range(len(analyzer.unit_ids))
              if np.array_equal(after_acgs[i], before_acgs[j])]
    print(f"{str(unit_id):>4}  {str(same):<38}  {source if source else 'the merged unit'}")

unit ids before merge: ['0', '1', '2', '3', '4', '5']
unit ids after merge : ['0', '3', '4', '5', '6']

unit  carried 3D ACG equals a recompute   carried 3D ACG is actually the one of
   0  True                                    ['0']
   3  False                                   ['1']
   4  False                                   ['2']
   5  False                                   ['3']
   6  True                                    the merged unit

Units 3, 4 and 5 were not in the merge group, so their 3D ACGs cannot have changed. Unit 3
comes back with the 3D ACG of unit 1, which is one of the two units that were merged away.
The merged unit itself is correct, because it is recomputed.

Running that same merge under both modes, and counting calls into
ComputeACG3D._merge_extension_data with a spy on the class:

merging_mode="soft": unit ids ['0', '3', '4', '5', '6']
    3D ACGs wrong           : ['3', '4', '5']
    firing quantiles wrong  : ['3', '4', '5']
merging_mode="hard": unit ids ['0', '3', '4', '5', '6']
    3D ACGs wrong           : []
    firing quantiles wrong  : []

--- is _merge_extension_data reached at all in each mode? ---
{'soft': 1, 'hard': 0}

The same thing happens through apply_curation with the same merge expressed as a curation
dict.

It also happens on the automatic path. auto_merge_units defaults to merging_mode="soft"
at curation/auto_merge.py:729 and reaches merge_units at :510. On
generate_ground_truth_recording(durations=[60.0], num_units=10, seed=2205) with
inject_some_split_units(sort, split_ids=sort.unit_ids[:3], num_split=2, seed=42), a sparse
analyzer, and random_spikes, waveforms, templates, noise_levels,
template_similarity, correlograms and acgs_3d computed first, a default
auto_merge_units(analyzer, presets=["similarity_correlograms"], n_jobs=1) run gives:

_merge_extension_data calls : 1
units before auto merge     : 13
units after auto merge      : 12
3D ACGs wrong               : ['3', '4', '5', '6', '7', '8', '9']
firing quantiles wrong      : ['3', '4', '5', '6', '7', '8', '9']

Seven of the twelve surviving units. That run repeated on a tree carrying the change below
gives [] for both lines.

Cause

src/spikeinterface/postprocessing/correlograms.py, ComputeACG3D._merge_extension_data, at
d365714a4:

1224        new_sorting = new_sorting_analyzer.sorting
...
1236        new_unit_ids_indices = new_sorting.ids_to_indices(new_unit_ids)
1237        old_unit_ids = [unit_id for unit_id in new_sorting_analyzer.unit_ids if unit_id not in new_unit_ids]
1238        old_unit_ids_indices = new_sorting.ids_to_indices(old_unit_ids)
...
1244        new_acgs_3d[old_unit_ids_indices, :, :] = self.data["acgs_3d"][old_unit_ids_indices, :, :]
1247        new_firing_quantiles[old_unit_ids_indices, :] = self.data["firing_quantiles"][old_unit_ids_indices, :]

old_unit_ids_indices is looked up in new_sorting on line 1238 and then used on both sides
of the assignment. The left side wants an index into the merged sorting, which is what it is.
The right side reads self.data, which is laid out in the order of the sorting before the
merge. A merge removes unit ids, so a kept unit generally has a lower index afterwards, and
the row that gets copied is whichever unit used to sit at that lower index.

Extensions that keep unit data across a merge look the source rows up against the sorting the
data belongs to. ComputeAutoCorrelograms._merge_extension_data at correlograms.py:296,
ComputeTemplates._merge_extension_data at core/analyzer_extension_core.py:584 and
ComputeUnitLocations._merge_extension_data at postprocessing/unit_locations.py:75 all use
self.sorting_analyzer.sorting.id_to_index(unit_id) for the units they keep.
ComputeACG3D._select_units_extension_data at correlograms.py:1212 also uses the old
sorting, which is why select_units is not affected.

Which units are affected

Every merge group of size 2 and 3 over 6 units, for both new_id_strategy values, 70
configurations, each kept unit compared against a fresh recompute on the merged analyzer:

merge configurations tested       : 70
kept units checked                : 240
kept units whose index changed    : 135
kept units with a wrong 3D ACG    : 135
cases where the rule disagrees    : 0

A kept unit gets another unit's data exactly when the merge moved it to a different index.

Cross-tabulating that against the easier-to-check "does it sit after the first merged unit in
the original list", over the same 70 configurations, 120 kept units per strategy:

new_id_strategy="append": 120 kept units checked, 85 at a changed index
                        wrong   correct
  sits after first merged     85         0   (total 85)
  sits before or at it         0        35   (total 35)

new_id_strategy="take_first": 120 kept units checked, 50 at a changed index
                        wrong   correct
  sits after first merged     50        35   (total 85)
  sits before or at it         0        35   (total 35)

For append the two statements coincide. For take_first they do not, because the merged
unit can reuse the first id of the group and hold that slot, which leaves 35 units sitting
after it at an unchanged index and correct. Nothing before the first merged unit is ever
wrong, under either strategy, so merging only the last units in the list moves nobody and
nothing is wrong in that case.

Suggested fix

diff --git a/src/spikeinterface/postprocessing/correlograms.py b/src/spikeinterface/postprocessing/correlograms.py
index 13431a8db..3ab2e590d 100644
--- a/src/spikeinterface/postprocessing/correlograms.py
+++ b/src/spikeinterface/postprocessing/correlograms.py
@@ -1234,17 +1234,20 @@ class ComputeACG3D(AnalyzerExtension):
         )
 
         new_unit_ids_indices = new_sorting.ids_to_indices(new_unit_ids)
+        # a unit which takes no part in the merge keeps its data, but it does not keep its
+        # index, so the source rows have to be looked up in the sorting they come from
         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)
+        old_unit_ids_indices = self.sorting_analyzer.sorting.ids_to_indices(old_unit_ids)
+        old_unit_ids_new_indices = new_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_new_indices, :, :] = self.data["acgs_3d"][old_unit_ids_indices, :, :]
 
         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_new_indices, :] = self.data["firing_quantiles"][old_unit_ids_indices, :]
 
         new_data = dict(
             acgs_3d=new_acgs_3d,

With that change all 70 configurations above come back clean, and the reproduction prints
True on every row.

On why the tests are green: common_extension_tests.py:155-158 does merge units for every
extension that goes through AnalyzerExtensionCommonTestSuite, and it merges
sorting_analyzer.unit_ids[:2] with the default new_id_strategy, which is the case where
every remaining unit shifts. It then asserts len(merged.unit_ids) == num_units_after_merge
and nothing about the data. test_multi_extensions.py:183 is
np.testing.assert_array_equal(data_original_unmerged, data_soft_unmerged), the check that
would fail here, but the extension_dict it loops over lists 14 extensions and acgs_3d is
not one of them.

That assertion is also why I do not think this falls under the docstring at
core/sortinganalyzer.py:1894-1895, which says soft merges "will be approximated, with no
reloading of the waveforms". Whatever latitude that gives for the merged units, the project's
own test requires the units that were not merged to come through a soft merge exactly equal.

A regression test in the shape of test_correlograms_merge, in the same file
postprocessing/tests/test_extension_merges.py, is what caught it:

def test_acgs_3d_merge():
    """
    A 3D ACG only depends on the spike train of its own unit, so the units which take no
    part in a merge keep the 3D ACG they had before it. This test checks that the merge
    method gives the same result as recomputing the 3D ACGs from scratch, for merge groups
    at the start, in the middle and at the end of the unit list.
    """

    rec, sort = generate_ground_truth_recording(durations=[20.0], num_units=6, seed=2205)

    sorting_analyzer = create_sorting_analyzer(recording=rec, sorting=sort, sparse=False)
    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
            )
            computed_acgs_3d, computed_quantiles, _ = merged_sorting_analyzer.get_extension("acgs_3d").get_data()

            recomputed_acgs_3d, recomputed_quantiles, _ = 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)

It fails on d365714a4 and passes with the change.
python -m pytest src/spikeinterface/postprocessing -q goes from 97 passed, 40 skipped to 98
passed, 40 skipped, and python -m pytest src/spikeinterface/core -q is 323 passed, 5 skipped
either way.

Age, and two open pull requests that rewrite these lines

The three lines have been as they are since d78573c78, merged in #3860 on 2025-06-12, and
they are unchanged in the 0.103.0, 0.104.0 and 0.104.8 tags.

#4713 and #4715 both replace the right-hand side with
slice_rows(self.data["acgs_3d"], old_unit_ids_indices) and keep old_unit_ids_indices on
both sides. I ran the reproduction against both heads, 5bcee5dd8 and 0325c14a5, and units
3, 4 and 5 come back wrong there too, so it survives whichever of those lands.

One thing I am not sure about: whether you would rather fix this by dropping the carry-forward
and recomputing acgs_3d on merge the way ComputeCorrelograms._split_extension_data does for
splits. That is a bigger behaviour change than the index fix and it is your call.

Two things next to this that the change above does not touch, in case they matter to whoever
picks it up.

bins, the third value get_data() returns, is carried through unchanged.
correlograms.py:1340 builds it as np.repeat(bin_times_ms, num_units, axis=0), so its
length tracks the unit count: after the 6 unit to 5 unit merge above the carried bins has
306 entries where a recompute gives 255. select_units carries it the same way. The change
above does not alter that either way, and I have not worked out whether the repeat is
deliberate.

ComputeACG3D has no _split_extension_data, so SortingAnalyzer.split_units on an analyzer
with acgs_3d computed raises a bare NotImplementedError from
AnalyzerExtension._split_extension_data at core/sortinganalyzer.py:2952. That is the soft
path, which is the one you get by default. Passing splitting_mode="hard" through
**job_kwargs returns normally with acgs_3d recomputed.