Merge branch 'clustering_merge_isocut' of github.com:samuelgarcia/spikeinterface into clustering_merge_isocut

samuelgarcia · samuelgarcia · commit 39c8e607f51c · 2025-09-10T11:21:53.000+02:00
diff --git a/src/spikeinterface/sortingcomponents/clustering/merge.py b/src/spikeinterface/sortingcomponents/clustering/merge.py
@@ -33,7 +33,7 @@ def merge_peak_labels_from_features(
     template_sparse_mask,
     recording,
     features_dict_or_folder,
-    radius_um=70.,
+    radius_um=70.0,
     method="project_distribution",
     method_kwargs={},
     **job_kwargs,
@@ -57,7 +57,6 @@ def merge_peak_labels_from_features(
         template_sparse_mask,
         recording,
         features_dict_or_folder,
-
         # sparse_wfs,
         # sparse_mask,
         radius_um=radius_um,
@@ -66,29 +65,23 @@ def merge_peak_labels_from_features(
         **job_kwargs,
     )
 
-    clean_labels, merge_template_array, merge_sparsity_mask, new_unit_ids = \
+    clean_labels, merge_template_array, merge_sparsity_mask, new_unit_ids = (
         _apply_pair_mask_on_labels_and_recompute_templates(
-            pair_mask,
-            peak_labels,
-            unit_ids,
-            templates_array, 
-            template_sparse_mask
+            pair_mask, peak_labels, unit_ids, templates_array, template_sparse_mask
         )
+    )
 
     return clean_labels, merge_template_array, merge_sparsity_mask, new_unit_ids
 
 
-
-
 def find_merge_pairs_from_features(
     peaks,
     peak_labels,
     unit_ids,
     templates_array,
     template_sparse_mask,
     recording,
-    features_dict_or_folder,      
-
+    features_dict_or_folder,
     # sparse_wfs,
     # sparse_mask,
     radius_um=70,
@@ -120,7 +113,6 @@ def find_merge_pairs_from_features(
 
     # compute template (no shift at this step)
 
-
     # templates = compute_template_from_sparse(
     #     peaks, peak_labels, labels_set, sparse_wfs, sparse_mask, total_channels, peak_shifts=None
     # )
@@ -131,13 +123,11 @@ def find_merge_pairs_from_features(
     # ms_after = features['ms_after']
     # svd_model = features['svd_model']
 
-
     # templates, final_sparsity_mask = get_templates_from_peaks_and_svd(
-    #     recording, peaks, peak_labels, ms_before, ms_after, svd_model, peaks_svd, sparse_mask, operator="average",        
+    #     recording, peaks, peak_labels, ms_before, ms_after, svd_model, peaks_svd, sparse_mask, operator="average",
     # )
     # dense_templates_array = templates.templates_array
 
-
     labels_set = unit_ids.tolist()
 
     max_chans = np.argmax(np.max(np.abs(templates_array), axis=1), axis=1)
@@ -319,8 +309,6 @@ def merge(
             final_shift = 0
             return is_merge, label0, label1, final_shift, merge_value
 
-
-
         inds = np.concatenate([inds0, inds1])
         labels = np.zeros(inds.size, dtype="int")
         labels[inds0.size :] = 1
@@ -332,7 +320,6 @@ def merge(
         wfs0 = wfs[:cut, :, :]
         wfs1 = wfs[cut:, :, :]
 
-
         # num_samples = template0.shape[0]
 
         # template0 = template0_[num_shift : num_samples - num_shift, :]
@@ -368,7 +355,6 @@ def merge(
         # wfs1 = wfs1_[:, best_shift : best_shift + template0.shape[0], :]
         # template1 = template1_[best_shift : best_shift + template0.shape[0], :]
 
-        
         feat0 = wfs0.reshape(wfs0.shape[0], -1)
         feat1 = wfs1.reshape(wfs1.shape[0], -1)
         feat = np.concatenate([feat0, feat1], axis=0)
@@ -377,11 +363,11 @@ def merge(
 
         if use_svd:
             from sklearn.decomposition import TruncatedSVD
+
             n_pca_features = 3
             tsvd = TruncatedSVD(n_pca_features, random_state=seed)
             feat = tsvd.fit_transform(feat)
 
-
         if isinstance(n_pca_features, float):
             assert 0 < n_pca_features < 1, "n_components should be in ]0, 1["
             nb_dimensions = min(feat.shape[0], feat.shape[1])
@@ -416,7 +402,6 @@ def merge(
         # else:
         #     feat = feat
 
-
         feat0 = feat[:cut]
         feat1 = feat[cut:]
 
@@ -447,7 +432,6 @@ def merge(
         feat0 = feat[:cut]
         feat1 = feat[cut:]
 
-
         if criteria == "isocut":
             dipscore, cutpoint = isocut(feat)
             is_merge = dipscore < isocut_threshold
@@ -484,7 +468,7 @@ def merge(
             final_shift = 0
 
         if DEBUG:
-        # if dipscore < 4:
+            # if dipscore < 4:
             import matplotlib.pyplot as plt
 
             flatten_wfs0 = wfs0.swapaxes(1, 2).reshape(wfs0.shape[0], -1)
@@ -570,33 +554,25 @@ def merge_peak_labels_from_templates(
     )
     pair_mask = similarity > similarity_thresh
 
-
-    clean_labels, merge_template_array, merge_sparsity_mask, new_unit_ids = \
+    clean_labels, merge_template_array, merge_sparsity_mask, new_unit_ids = (
         _apply_pair_mask_on_labels_and_recompute_templates(
-            pair_mask,
-            peak_labels,
-            unit_ids,
-            templates_array, 
-            template_sparse_mask
+            pair_mask, peak_labels, unit_ids, templates_array, template_sparse_mask
         )
+    )
 
     return clean_labels, merge_template_array, merge_sparsity_mask, new_unit_ids
 
+
 def _apply_pair_mask_on_labels_and_recompute_templates(
-    pair_mask,
-    peak_labels,
-    unit_ids,
-    templates_array, 
-    template_sparse_mask
+    pair_mask, peak_labels, unit_ids, templates_array, template_sparse_mask
 ):
     """
     Resolve pairs graph.
     Apply to new labels.
     Recompute templates.
     """
-    
-    from scipy.sparse.csgraph import connected_components
 
+    from scipy.sparse.csgraph import connected_components
 
     keep_template = np.ones(templates_array.shape[0], dtype="bool")
     clean_labels = peak_labels.copy()
@@ -638,4 +614,3 @@ def _apply_pair_mask_on_labels_and_recompute_templates(
     merge_sparsity_mask = merge_sparsity_mask[keep_template, :]
 
     return clean_labels, merge_template_array, merge_sparsity_mask, new_unit_ids
-
diff --git a/src/spikeinterface/sortingcomponents/clustering/tdc.py b/src/spikeinterface/sortingcomponents/clustering/tdc.py
@@ -22,7 +22,10 @@
 from spikeinterface.sortingcomponents.clustering.split import split_clusters
 
 # from spikeinterface.sortingcomponents.clustering.merge import merge_clusters
-from spikeinterface.sortingcomponents.clustering.merge import merge_peak_labels_from_templates, merge_peak_labels_from_features
+from spikeinterface.sortingcomponents.clustering.merge import (
+    merge_peak_labels_from_templates,
+    merge_peak_labels_from_features,
+)
 from spikeinterface.sortingcomponents.clustering.tools import get_templates_from_peaks_and_svd
 from spikeinterface.sortingcomponents.clustering.peak_svd import extract_peaks_svd
 
@@ -55,7 +58,6 @@ class TdcClustering:
             # "clusterer": "isosplit6",
             # "clusterer_kwargs": {},
             "clusterer": "isosplit",
-
             "clusterer_kwargs": {
                 "n_init": 50,
                 "min_cluster_size": 10,
@@ -71,11 +73,7 @@ class TdcClustering:
             "min_size_split": 10,
         },
         "do_merge_with_features": False,
-        "merge_features_kwargs": {
-            "merge_radius_um":50.,
-            "criteria": "isocut",
-            "isocut_threshold": 2.0
-        },
+        "merge_features_kwargs": {"merge_radius_um": 50.0, "criteria": "isocut", "isocut_threshold": 2.0},
         "do_merge_with_templates": True,
         "merge_template_kwargs": {
             "similarity_metric": "l1",
@@ -123,7 +121,6 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
 
         original_labels = peaks["channel_index"]
 
-
         clusterer = params["split"]["clusterer"]
         clusterer_kwargs = params["split"]["clusterer_kwargs"]
 
@@ -186,29 +183,24 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
                 radius_um=merge_radius_um,
                 method="project_distribution",
                 method_kwargs=dict(
-                    feature_name="peaks_svd",
-                    waveforms_sparse_mask=sparse_mask,
-                    **merge_features_kwargs
-                    ),
-                **job_kwargs
+                    feature_name="peaks_svd", waveforms_sparse_mask=sparse_mask, **merge_features_kwargs
+                ),
+                **job_kwargs,
             )
         else:
             post_merge_label1 = post_split_label.copy()
-        
+
         if params["do_merge_with_templates"]:
-            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_template_kwargs"],
-                )
+            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_template_kwargs"],
             )
         else:
             post_merge_label2 = post_merge_label1.copy()
-        
 
         dense_templates = Templates(
             templates_array=templates_array,
@@ -221,7 +213,6 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
             is_in_uV=False,
         )
 
-
         sparsity = ChannelSparsity(template_sparse_mask, unit_ids, recording.channel_ids)
         templates = dense_templates.to_sparse(sparsity)
 
@@ -243,9 +234,7 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
         templates = templates.select_units(labels_set)
         labels_set = templates.unit_ids
 
-
         more_outs = dict(
             templates=templates,
         )
         return labels_set, final_peak_labels, more_outs
-
diff --git a/src/spikeinterface/sortingcomponents/clustering/tools.py b/src/spikeinterface/sortingcomponents/clustering/tools.py
@@ -127,7 +127,7 @@ def aggregate_sparse_features(peaks, peak_indices, sparse_feature, sparse_target
 #     templates: numpy.array
 #         Templates shape : (len(labels_set), num_samples, total_channels)
 #     """
-    
+
 #     # NOTE SAM I think this is wrong, we should remove
 
 #     n = len(labels_set)
