Add the option for clustering methods to return temporal shifts for peaks while clustering

src/spikeinterface/sorters/internal/lupin.py

@@ -310,6 +310,10 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
             job_kwargs=job_kwargs,
         )
 
+        if more_outs["time_shifts"] is not None:
+            time_shifts = more_outs["time_shifts"]
+            peaks["sample_index"] += time_shifts
+
         mask = clustering_label >= 0
         kept_peaks = peaks[mask]
         kept_labels = clustering_label[mask]

src/spikeinterface/sorters/internal/tridesclous2.py

@@ -350,6 +350,7 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
         sorting = NumpySorting(final_spikes, sampling_frequency, templates.unit_ids)
 
         auto_merge = True
+
         analyzer_final = None
         if auto_merge:
             from spikeinterface.sorters.internal.spyking_circus2 import final_cleaning_circus

src/spikeinterface/sortingcomponents/clustering/iterative_hdbscan.py

@@ -166,13 +166,15 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
             print("Kept %d raw clusters" % len(labels))
 
         if params["merge_from_templates"] is not None:
-            peak_labels, merge_template_array, new_sparse_mask, new_unit_ids = merge_peak_labels_from_templates(
-                peaks,
-                peak_labels,
-                templates.unit_ids,
-                templates.templates_array,
-                new_sparse_mask,
-                **params["merge_from_templates"],
+            peak_labels, merge_template_array, new_sparse_mask, new_unit_ids, time_shifts = (
+                merge_peak_labels_from_templates(
+                    peaks,
+                    peak_labels,
+                    templates.unit_ids,
+                    templates.templates_array,
+                    new_sparse_mask,
+                    **params["merge_from_templates"],
+                )
             )
 
             templates = Templates(
@@ -185,6 +187,8 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
                 probe=recording.get_probe(),
                 is_in_uV=False,
             )
+        else:
+            time_shifts = None
 
         # clean very small cluster before peeler
         if (
@@ -212,6 +216,7 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
         more_outs = dict(
             svd_model=svd_model,
             peaks_svd=peaks_svd,
+            time_shifts=time_shifts,
             peak_svd_sparse_mask=sparse_mask,
         )
         return labels, peak_labels, more_outs

src/spikeinterface/sortingcomponents/clustering/iterative_isosplit.py

@@ -293,16 +293,19 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
             num_shifts = params_merge_from_templates["num_shifts"]
             num_shifts = min((num_shifts, nbefore, nafter))
             params_merge_from_templates["num_shifts"] = num_shifts
-            post_merge_label2, templates_array, template_sparse_mask, unit_ids = merge_peak_labels_from_templates(
-                peaks,
-                post_merge_label1,
-                unit_ids,
-                templates_array,
-                template_sparse_mask,
-                **params_merge_from_templates,
+            post_merge_label2, templates_array, template_sparse_mask, unit_ids, time_shifts = (
+                merge_peak_labels_from_templates(
+                    peaks,
+                    post_merge_label1,
+                    unit_ids,
+                    templates_array,
+                    template_sparse_mask,
+                    **params_merge_from_templates,
+                )
             )
         else:
             post_merge_label2 = post_merge_label1.copy()
+            time_shifts = None
 
         dense_templates = Templates(
             templates_array=templates_array,
@@ -337,7 +340,5 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
 
         labels_set = templates.unit_ids
 
-        more_outs = dict(
-            templates=templates,
-        )
+        more_outs = dict(templates=templates, time_shifts=time_shifts)
         return labels_set, final_peak_labels, more_outs

src/spikeinterface/sortingcomponents/clustering/merging_tools.py

@@ -558,13 +558,13 @@ def merge_peak_labels_from_templates(
     if not use_lags:
         lags = None
 
-    clean_labels, merge_template_array, merge_sparsity_mask, new_unit_ids = (
+    clean_labels, merge_template_array, merge_sparsity_mask, new_unit_ids, time_shifts = (
         _apply_pair_mask_on_labels_and_recompute_templates(
             pair_mask, peak_labels, unit_ids, templates_array, template_sparse_mask, lags
         )
     )
 
-    return clean_labels, merge_template_array, merge_sparsity_mask, new_unit_ids
+    return clean_labels, merge_template_array, merge_sparsity_mask, new_unit_ids, time_shifts
 
 
 def _apply_pair_mask_on_labels_and_recompute_templates(
@@ -582,7 +582,10 @@ def _apply_pair_mask_on_labels_and_recompute_templates(
     clean_labels = peak_labels.copy()
     n_components, group_labels = connected_components(pair_mask, directed=False, return_labels=True)
 
-    # print("merges", templates_array.shape[0], "to", n_components)
+    if lags is not None:
+        time_shifts = np.zeros(len(peak_labels), dtype=np.int32)
+    else:
+        time_shifts = None
 
     merge_template_array = templates_array.copy()
     merge_sparsity_mask = template_sparse_mask.copy()
@@ -605,10 +608,15 @@ def _apply_pair_mask_on_labels_and_recompute_templates(
 
             for i, l in enumerate(merge_group):
                 label = unit_ids[l]
-                weights[i] = np.sum(peak_labels == label)
+                mask = peak_labels == label
+                weights[i] = np.sum(mask)
                 if i > 0:
-                    clean_labels[peak_labels == label] = unit_ids[g0]
+                    clean_labels[mask] = unit_ids[g0]
                     keep_template[l] = False
+                    if lags is not None:
+                        shift = lags[l, g0]  # which is the same as  -lags[g0, l]
+                        time_shifts[mask] += shift
+
             weights /= weights.sum()
 
             if lags is None:
@@ -619,7 +627,7 @@ def _apply_pair_mask_on_labels_and_recompute_templates(
                 # with shifts
                 accumulated_template = np.zeros_like(merge_template_array[g0, :, :])
                 for i, l in enumerate(merge_group):
-                    shift = -lags[g0, l]
+                    shift = lags[l, g0]  # which is the same as  -lags[g0, l]
                     if shift > 0:
                         # template is shifted to right
                         temp = np.zeros_like(accumulated_template)
@@ -639,4 +647,4 @@ def _apply_pair_mask_on_labels_and_recompute_templates(
     merge_template_array = merge_template_array[keep_template, :, :]
     merge_sparsity_mask = merge_sparsity_mask[keep_template, :]
 
-    return clean_labels, merge_template_array, merge_sparsity_mask, new_unit_ids
+    return clean_labels, merge_template_array, merge_sparsity_mask, new_unit_ids, time_shifts

src/spikeinterface/sortingcomponents/clustering/random_projections.py

@@ -134,13 +134,15 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
             print("Kept %d raw clusters" % len(labels))
 
         if params["merge_from_templates"] is not None:
-            peak_labels, merge_template_array, new_sparse_mask, new_unit_ids = merge_peak_labels_from_templates(
-                peaks,
-                peak_labels,
-                unit_ids,
-                templates_array,
-                np.ones((len(unit_ids), num_chans), dtype=bool),
-                **params["merge_from_templates"],
+            peak_labels, merge_template_array, new_sparse_mask, new_unit_ids, time_shifts = (
+                merge_peak_labels_from_templates(
+                    peaks,
+                    peak_labels,
+                    unit_ids,
+                    templates_array,
+                    np.ones((len(unit_ids), num_chans), dtype=bool),
+                    **params["merge_from_templates"],
+                )
             )
 
             templates = Templates(
@@ -153,6 +155,8 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
                 probe=recording.get_probe(),
                 is_in_uV=False,
             )
+        else:
+            time_shifts = None
 
         labels = templates.unit_ids
 
@@ -162,4 +166,4 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
         if verbose:
             print("Kept %d non-duplicated clusters" % len(labels))
 
-        return labels, peak_labels, dict()
+        return labels, peak_labels, dict(time_shifts=time_shifts, templates=templates)