respond to Sam

Author: chrishalcrow

src/spikeinterface/postprocessing/correlograms.py

@@ -108,17 +108,25 @@ def _merge_extension_data(
             # 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):
 
-                merged_spike_train_length = len(new_sorting_analyzer.sorting.get_unit_spike_train(new_unit_id))
+                num_segments = new_sorting_analyzer.get_num_segments()
+                for segment_index in range(num_segments):
 
-                old_spike_train_lengths = len(
-                    np.concatenate(
-                        [self.sorting_analyzer.sorting.get_unit_spike_train(unit_id) for unit_id in merge_unit_group]
+                    merged_spike_train_length = len(
+                        new_sorting_analyzer.sorting.get_unit_spike_train(new_unit_id, segment_index=segment_index)
                     )
-                )
 
-                if merged_spike_train_length != old_spike_train_lengths:
-                    can_apply_soft_method = False
-                    break
+                    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)
@@ -134,11 +142,14 @@ def _merge_extension_data(
 
                 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 the first unit index 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[1:], :, :] = 0
 
+                # Sum unit columns of the correlogram matrix: C_{l,k} = C_{l,i} + C_{l,j}
+                # and put this sum in the first unit index from the merge group
                 new_row = np.sum(correlograms[:, merge_unit_group_indices, :], axis=1)
                 correlograms[:, merge_unit_group_indices[0], :] = new_row
                 correlograms[:, merge_unit_group_indices[1:], :] = 0
@@ -147,6 +158,8 @@ def _merge_extension_data(
                     need_to_append = True
                     delete_from = 0
 
+            # When new_id_strategy='append' is used, the new summed rows/columns to the
+            # end of the correlogram matrix
             if need_to_append is True:
                 old_num_units = np.shape(correlograms)[0]
                 correlograms = np.pad(correlograms, ((0, len(new_unit_ids)), (0, len(new_unit_ids)), (0, 0)))
@@ -158,6 +171,7 @@ def _merge_extension_data(
                     correlograms[:, [old_loc, new_loc], :] = correlograms[:, [new_loc, old_loc], :]
                     correlograms[[old_loc, new_loc], :, :] = correlograms[[new_loc, old_loc], :, :]
 
+            # Delete the leftover units from the merge
             units_to_delete = np.concatenate([merge_unit_group[delete_from:] for merge_unit_group in merge_unit_groups])
             indices_to_delete = self.sorting_analyzer.sorting.ids_to_indices(units_to_delete)
 

src/spikeinterface/postprocessing/tests/test_extension_merges.py

@@ -10,7 +10,7 @@ def test_correlograms_merge():
     sum as recomputing the correlograms from scratch.
     """
 
-    rec, sort = generate_ground_truth_recording()
+    rec, sort = generate_ground_truth_recording(durations=[10, 10])
 
     sorting_analyzer = create_sorting_analyzer(recording=rec, sorting=sort)
     sorting_analyzer.compute("correlograms")