Merge pull request #3607 from chrishalcrow/fast-correlogram-merge

Add a fast correlogram merge

Author: samuelgarcia

src/spikeinterface/postprocessing/correlograms.py

@@ -3,6 +3,7 @@
 import warnings
 import numpy as np
 from spikeinterface.core.sortinganalyzer import register_result_extension, AnalyzerExtension, SortingAnalyzer
+from copy import deepcopy
 
 from spikeinterface.core.waveforms_extractor_backwards_compatibility import MockWaveformExtractor
 
@@ -93,9 +94,98 @@ def _select_extension_data(self, unit_ids):
     def _merge_extension_data(
         self, merge_unit_groups, new_unit_ids, new_sorting_analyzer, censor_ms=None, verbose=False, **job_kwargs
     ):
-        # recomputing correlogram is fast enough and much easier in this case
-        new_ccgs, new_bins = _compute_correlograms_on_sorting(new_sorting_analyzer.sorting, **self.params)
-        new_data = dict(ccgs=new_ccgs, bins=new_bins)
+        """
+        When two units are merged, their cross-correlograms with other units become the sum
+        of the previous cross-correlograms. More precisely, if units i and j get merged into
+        unit k, then the new unit's cross-correlogram with any other unit l is:
+            C_{k,l} = C_{i,l} + C_{j,l}
+            C_{l,k} = C_{l,k} + C_{l,j}
+        Here, we apply this formula to quickly compute correlograms for merged units.
+        """
+
+        can_apply_soft_method = True
+        if censor_ms is not None:
+            # if censor_ms has no effect, can apply "soft" method. Check if any spikes have been removed
+            for new_unit_id, merge_unit_group in zip(new_unit_ids, merge_unit_groups):
+
+                num_segments = new_sorting_analyzer.get_num_segments()
+                for segment_index in range(num_segments):
+
+                    merged_spike_train_length = len(
+                        new_sorting_analyzer.sorting.get_unit_spike_train(new_unit_id, segment_index=segment_index)
+                    )
+
+                    old_spike_train_lengths = len(
+                        np.concatenate(
+                            [
+                                self.sorting_analyzer.sorting.get_unit_spike_train(unit_id, segment_index=segment_index)
+                                for unit_id in merge_unit_group
+                            ]
+                        )
+                    )
+
+                    if merged_spike_train_length != old_spike_train_lengths:
+                        can_apply_soft_method = False
+                        break
+
+        if can_apply_soft_method is False:
+            new_ccgs, new_bins = _compute_correlograms_on_sorting(new_sorting_analyzer.sorting, **self.params)
+            new_data = dict(ccgs=new_ccgs, bins=new_bins)
+        else:
+
+            # Make a transformation dict, which tells us how unit_indices from the
+            # old to the new sorter are mapped.
+            old_to_new_unit_index_map = {}
+            for old_unit in self.sorting_analyzer.unit_ids:
+                old_unit_index = self.sorting_analyzer.sorting.id_to_index(old_unit)
+                unit_involved_in_merge = False
+                for merge_unit_group, new_unit_id in zip(merge_unit_groups, new_unit_ids):
+                    new_unit_index = new_sorting_analyzer.sorting.id_to_index(new_unit_id)
+                    # check if the old_unit is involved in a merge
+                    if old_unit in merge_unit_group:
+                        # check if it is mapped to itself
+                        if old_unit == new_unit_id:
+                            old_to_new_unit_index_map[old_unit_index] = new_unit_index
+                        # or to a unit_id outwith the old ones
+                        elif new_unit_id not in self.sorting_analyzer.unit_ids:
+                            if new_unit_index not in old_to_new_unit_index_map.values():
+                                old_to_new_unit_index_map[old_unit_index] = new_unit_index
+                        unit_involved_in_merge = True
+                if unit_involved_in_merge is False:
+                    old_to_new_unit_index_map[old_unit_index] = new_sorting_analyzer.sorting.id_to_index(old_unit)
+
+            need_to_append = False
+            delete_from = 1
+
+            correlograms, new_bins = deepcopy(self.get_data())
+
+            for new_unit_id, merge_unit_group in zip(new_unit_ids, merge_unit_groups):
+
+                merge_unit_group_indices = self.sorting_analyzer.sorting.ids_to_indices(merge_unit_group)
+
+                # Sum unit rows of the correlogram matrix: C_{k,l} = C_{i,l} + C_{j,l}
+                # and place this sum in all indices from the merge group
+                new_col = np.sum(correlograms[merge_unit_group_indices, :, :], axis=0)
+                # correlograms[merge_unit_group_indices[0], :, :] = new_col
+                correlograms[merge_unit_group_indices, :, :] = new_col
+                # correlograms[merge_unit_group_indices[1:], :, :] = 0
+
+                # Sum unit columns of the correlogram matrix: C_{l,k} = C_{l,i} + C_{l,j}
+                # and put this sum in all indices from the merge group
+                new_row = np.sum(correlograms[:, merge_unit_group_indices, :], axis=1)
+
+                for merge_unit_group_index in merge_unit_group_indices:
+                    correlograms[:, merge_unit_group_index, :] = new_row
+
+            new_correlograms = np.zeros(
+                (len(new_sorting_analyzer.unit_ids), len(new_sorting_analyzer.unit_ids), correlograms.shape[2])
+            )
+            for old_index_1, new_index_1 in old_to_new_unit_index_map.items():
+                for old_index_2, new_index_2 in old_to_new_unit_index_map.items():
+                    new_correlograms[new_index_1, new_index_2, :] = correlograms[old_index_1, old_index_2, :]
+                    new_correlograms[new_index_2, new_index_1, :] = correlograms[old_index_2, old_index_1, :]
+
+            new_data = dict(ccgs=new_correlograms, bins=new_bins)
         return new_data
 
     def _run(self, verbose=False):

src/spikeinterface/postprocessing/tests/test_extension_merges.py

@@ -0,0 +1,61 @@
+import numpy as np
+
+from spikeinterface.core import generate_ground_truth_recording, create_sorting_analyzer
+
+
+def test_correlograms_merge():
+    """
+    When merging in `soft` mode, correlograms sum and we take advantage of this to make
+    a fast computation. This test checks that we get the same result using this fast
+    sum as recomputing the correlograms from scratch.
+    """
+
+    rec, sort = generate_ground_truth_recording(durations=[10, 10])
+
+    sorting_analyzer = create_sorting_analyzer(recording=rec, sorting=sort)
+    sorting_analyzer.compute("correlograms")
+
+    trial_merges = [
+        [["1", "2"]],
+        [["2", "4", "6", "8"]],
+        [["1", "4", "7"], ["2", "8"]],
+        [["4", "1", "8"], ["2", "7", "0"], ["3", "9"], ["5", "6"]],
+    ]
+
+    new_unit_ids = [["2"], ["4"], ["4", "2"], ["1", "2", "9", "100"]]
+
+    for new_id_strategy in ["append", "take_first", "user"]:
+        for merge_unit_groups, new_unit_id in zip(trial_merges, new_unit_ids):
+
+            # first, compute the correlograms of the merged units using the merge method
+            if new_id_strategy == "user":
+                merged_sorting_analyzer = sorting_analyzer.merge_units(
+                    merge_unit_groups=merge_unit_groups, new_unit_ids=new_unit_id
+                )
+            else:
+                merged_sorting_analyzer = sorting_analyzer.merge_units(
+                    merge_unit_groups=merge_unit_groups, new_id_strategy=new_id_strategy
+                )
+            computed_correlograms = merged_sorting_analyzer.get_extension("correlograms").get_data()
+
+            # Then re-compute, and compare
+            recomputed_correlograms = merged_sorting_analyzer.compute("correlograms").get_data()
+            assert np.all(computed_correlograms[0] == recomputed_correlograms[0])
+
+    # test when `censor_ms` is not None. This merge does remove some spikes.
+    merged_sorting_analyzer_censored = sorting_analyzer.merge_units(
+        merge_unit_groups=trial_merges[2], new_id_strategy="take_first", censor_ms=5
+    )
+    computed_ccgs_censored = merged_sorting_analyzer_censored.get_extension("correlograms").get_data()
+
+    recomputed_ccgs_censored = merged_sorting_analyzer_censored.compute("correlograms").get_data()
+    assert np.all(computed_ccgs_censored[0] == recomputed_ccgs_censored[0])
+
+    # This `censor_ms` does not remove spikes, so can use the soft method
+    merged_sorting_analyzer_not_censored = sorting_analyzer.merge_units(
+        merge_unit_groups=trial_merges[0], new_id_strategy="take_first", censor_ms=0
+    )
+    computed_ccgs_not_censored = merged_sorting_analyzer_not_censored.get_extension("correlograms").get_data()
+
+    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])