Add the option for clustering methods to return temporal shifts for peaks while clustering (#4401)

Co-authored-by: Samuel Garcia <sam.garcia.die@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

src/spikeinterface/sorters/internal/lupin.py

@@ -308,6 +308,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

@@ -348,6 +348,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

@@ -164,13 +164,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(
@@ -183,6 +185,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 (
@@ -210,6 +214,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

@@ -299,16 +299,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,
@@ -343,7 +346,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

@@ -556,13 +556,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(
@@ -580,7 +580,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()
@@ -603,10 +606,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:
@@ -617,7 +625,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)
@@ -637,4 +645,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

@@ -132,13 +132,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(
@@ -151,6 +153,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
 
@@ -160,4 +164,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)