Update TDC2 and SC2 with various improvement in sorting components (clustering + matching)

src/spikeinterface/postprocessing/template_similarity.py

@@ -333,12 +333,17 @@ def compute_similarity_with_templates_array(
     mask = np.ones((num_templates, other_num_templates, num_channels), dtype=bool)
 
     if sparsity is not None and other_sparsity is not None:
+
+        # make the input more flexible with either The object or the array mask
+        sparsity_mask = sparsity.mask if isinstance(sparsity, ChannelSparsity) else sparsity
+        other_sparsity_mask = other_sparsity.mask if isinstance(other_sparsity, ChannelSparsity) else other_sparsity
+
         if support == "intersection":
-            mask = np.logical_and(sparsity.mask[:, np.newaxis, :], other_sparsity.mask[np.newaxis, :, :])
+            mask = np.logical_and(sparsity_mask[:, np.newaxis, :], other_sparsity_mask[np.newaxis, :, :])
         elif support == "union":
-            mask = np.logical_and(sparsity.mask[:, np.newaxis, :], other_sparsity.mask[np.newaxis, :, :])
+            mask = np.logical_and(sparsity_mask[:, np.newaxis, :], other_sparsity_mask[np.newaxis, :, :])
             units_overlaps = np.sum(mask, axis=2) > 0
-            mask = np.logical_or(sparsity.mask[:, np.newaxis, :], other_sparsity.mask[np.newaxis, :, :])
+            mask = np.logical_or(sparsity_mask[:, np.newaxis, :], other_sparsity_mask[np.newaxis, :, :])
             mask[~units_overlaps] = False
 
     assert num_shifts < num_samples, "max_lag is too large"

src/spikeinterface/sorters/internal/spyking_circus2.py

@@ -307,7 +307,7 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
                 _, peak_labels, svd_model, svd_features, sparsity_mask = outputs
                 from spikeinterface.sortingcomponents.clustering.tools import get_templates_from_peaks_and_svd
 
-                templates = get_templates_from_peaks_and_svd(
+                templates, _ = get_templates_from_peaks_and_svd(
                     recording_w,
                     selected_peaks,
                     peak_labels,

src/spikeinterface/sortingcomponents/clustering/circus.py

@@ -13,17 +13,15 @@
     HAVE_HDBSCAN = False
 
 import random, string
-from spikeinterface.core import get_global_tmp_folder
-from spikeinterface.core.basesorting import minimum_spike_dtype
-from spikeinterface.core.waveform_tools import estimate_templates
+from spikeinterface.core import get_global_tmp_folder, Templates
 from .clustering_tools import remove_duplicates_via_matching
 from spikeinterface.core.recording_tools import get_noise_levels, get_channel_distances
 from spikeinterface.sortingcomponents.peak_selection import select_peaks
 from spikeinterface.core.template import Templates
 from spikeinterface.core.sparsity import compute_sparsity
 from spikeinterface.sortingcomponents.tools import remove_empty_templates
 from spikeinterface.sortingcomponents.clustering.peak_svd import extract_peaks_svd
-
+from spikeinterface.sortingcomponents.clustering.merge import merge_peak_labels_from_templates
 
 from spikeinterface.sortingcomponents.tools import extract_waveform_at_max_channel
 
@@ -56,11 +54,18 @@ class CircusClustering:
         "few_waveforms": None,
         "ms_before": 0.5,
         "ms_after": 0.5,
+        "remove_small_snr": False,
         "noise_threshold": 4,
         "rank": 5,
         "templates_from_svd": False,
         "noise_levels": None,
         "tmp_folder": None,
+        "do_merge": True,
+        "merge_kwargs": {
+            "similarity_metric": "l1",
+            "num_shifts": 3,
+            "similarity_thresh": 0.8,
+        },
         "verbose": True,
         "debug": False,
     }
@@ -170,10 +175,11 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
                 ms_after,
                 **job_kwargs,
             )
+            sparse_mask2 = sparse_mask
         else:
             from spikeinterface.sortingcomponents.clustering.tools import get_templates_from_peaks_and_svd
 
-            templates = get_templates_from_peaks_and_svd(
+            templates, sparse_mask2 = get_templates_from_peaks_and_svd(
                 recording,
                 peaks,
                 peak_labels,
@@ -185,44 +191,54 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
                 operator="median",
             )
 
-        templates_array = templates.templates_array
-        best_channels = np.argmax(np.abs(templates_array[:, nbefore, :]), axis=1)
-        peak_snrs = np.abs(templates_array[:, nbefore, :])
-        best_snrs_ratio = (peak_snrs / params["noise_levels"])[np.arange(len(peak_snrs)), best_channels]
-        old_unit_ids = templates.unit_ids.copy()
-        valid_templates = best_snrs_ratio > params["noise_threshold"]
-
-        mask = np.isin(peak_labels, old_unit_ids[~valid_templates])
-        peak_labels[mask] = -1
-
-        from spikeinterface.core.template import Templates
-
-        templates = Templates(
-            templates_array=templates_array[valid_templates],
-            sampling_frequency=fs,
-            nbefore=templates.nbefore,
-            sparsity_mask=None,
-            channel_ids=recording.channel_ids,
-            unit_ids=templates.unit_ids[valid_templates],
-            probe=recording.get_probe(),
-            is_scaled=False,
-        )
+        if params["do_merge"]:
+            peak_labels, merge_template_array, new_unit_ids = merge_peak_labels_from_templates(
+                peaks, peak_labels, templates.unit_ids,
+                templates.templates_array, sparse_mask2,
+                **params["merge_kwargs"]
+            )
+
+            templates  = Templates(
+                templates_array=merge_template_array,
+                sampling_frequency=fs,
+                nbefore=templates.nbefore,
+                sparsity_mask=None,
+                channel_ids=recording.channel_ids,
+                unit_ids=new_unit_ids,
+                probe=recording.get_probe(),
+                is_scaled=False
+            )
+
+        if params["remove_small_snr"] :
+            templates_array = templates.templates_array
+            best_channels = np.argmax(np.abs(templates_array[:, nbefore, :]), axis=1)
+            peak_snrs = np.abs(templates_array[:, nbefore, :])
+            best_snrs_ratio = (peak_snrs / params["noise_levels"])[np.arange(len(peak_snrs)), best_channels]
+            old_unit_ids = templates.unit_ids.copy()
+            valid_templates = best_snrs_ratio > params["noise_threshold"]
+
+            mask = np.isin(peak_labels, old_unit_ids[~valid_templates])
+            peak_labels[mask] = -1
+
+            templates = templates.select_units(templates.unit_ids[valid_templates])
+
+        labels = templates.unit_ids
 
         if params["debug"]:
             templates_folder = tmp_folder / "dense_templates"
             templates.to_zarr(folder_path=templates_folder)
 
-        sparsity = compute_sparsity(templates, noise_levels=params["noise_levels"], **params["sparsity"])
-        templates = templates.to_sparse(sparsity)
-        empty_templates = templates.sparsity_mask.sum(axis=1) == 0
-        old_unit_ids = templates.unit_ids.copy()
-        templates = remove_empty_templates(templates)
+        # sparsity = compute_sparsity(templates, noise_levels=params["noise_levels"], **params["sparsity"])
+        # templates = templates.to_sparse(sparsity)
+        # empty_templates = templates.sparsity_mask.sum(axis=1) == 0
+        # old_unit_ids = templates.unit_ids.copy()
+        # templates = remove_empty_templates(templates)
 
-        mask = np.isin(peak_labels, old_unit_ids[empty_templates])
-        peak_labels[mask] = -1
+        # mask = np.isin(peak_labels, old_unit_ids[empty_templates])
+        # peak_labels[mask] = -1
 
-        labels = np.unique(peak_labels)
-        labels = labels[labels >= 0]
+        # labels = np.unique(peak_labels)
+        # labels = labels[labels >= 0]
 
         if params["remove_mixtures"]:
             if verbose:

src/spikeinterface/sortingcomponents/clustering/graph_clustering.py

@@ -79,7 +79,7 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
             ms_after=ms_after,
             radius_um=radius_um,
             motion_aware=motion_aware,
-            motion=None,
+            motion=motion,
             **params["extract_peaks_svd_kwargs"],
             # **job_kwargs
         )

src/spikeinterface/sortingcomponents/clustering/merge.py

@@ -754,3 +754,72 @@ def merge(
     NormalizedTemplateDiff,
 ]
 find_pair_method_dict = {e.name: e for e in find_pair_method_list}
+
+
+
+def merge_peak_labels_from_templates(peaks, peak_labels, unit_ids,
+                                     templates_array, sparsity_mask,
+                                     similarity_metric="l1",
+                                     similarity_thresh=0.8,
+                                     num_shifts=3,
+                                     ):
+    assert len(unit_ids) == templates_array.shape[0]
+
+    from spikeinterface.postprocessing.template_similarity import compute_similarity_with_templates_array
+    from scipy.sparse.csgraph import connected_components
+
+
+    similarity = compute_similarity_with_templates_array(
+        templates_array,
+        templates_array,
+        method=similarity_metric,
+        num_shifts=num_shifts,
+        support="union",
+        sparsity=sparsity_mask,
+        other_sparsity=sparsity_mask,
+    )
+    pair_mask = similarity > similarity_thresh
+
+    # import matplotlib.pyplot as plt
+    # fig, ax = plt.subplots()
+    # ax.hist(similarity.flatten(), bins=np.linspace(0, 1, 50), log=True)
+    # ax.axvline(similarity_thresh)
+
+
+    keep_template = np.ones(templates_array.shape[0], dtype="bool")
+    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)
+
+    merge_template_array = templates_array.copy()
+    new_unit_ids = np.zeros(n_components, dtype=unit_ids.dtype)
+    for c in range(n_components):
+        merge_group = np.flatnonzero(group_labels == c)
+        g0 = merge_group[0]
+        new_unit_ids[c] = unit_ids[g0]
+        if len(merge_group) > 1:
+            weights = np.zeros(len(merge_group), dtype=np.float32)
+
+            # import matplotlib.pyplot as plt
+            # fig, ax = plt.subplots()
+            # for i, l in enumerate(merge_group):
+            #     temp_flat = merge_template_array[l, :, :].T.flatten()
+            #     ax.plot(temp_flat)
+            # sim = similarity[merge_group[0], merge_group[1]]
+            # ax.set_title(f"{sim} {similarity_thresh}")
+
+
+            for i, l in enumerate(merge_group):
+                label = unit_ids[l]
+                weights[i] = np.sum(peak_labels == label)
+                if i > 0:
+                    clean_labels[peak_labels == label] = unit_ids[g0]
+                    keep_template[l] = False
+            weights /= weights.sum()
+            merge_template_array[g0, :, :] = np.sum(merge_template_array[merge_group, :, :] * weights[:, np.newaxis, np.newaxis], axis=0)
+
+
+    return clean_labels, merge_template_array, new_unit_ids
+

src/spikeinterface/sortingcomponents/clustering/method_list.py

@@ -22,7 +22,7 @@
     # "sliding_nn": SlidingNNClustering,
     "position_and_features": PositionAndFeaturesClustering,
     "random_projections": RandomProjectionClustering,
-    "circus": CircusClustering,
+    "circus_clustering": CircusClustering,
     "tdc_clustering": TdcClustering,
     "graph_clustering": GraphClustering,
 }