Improve Split PCA

Author: yger

src/spikeinterface/sorters/internal/spyking_circus2.py

@@ -32,7 +32,7 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
         "sparsity": {"method": "snr", "amplitude_mode": "peak_to_peak", "threshold": 0.25},
         "filtering": {"freq_min": 150, "freq_max": 7000, "ftype": "bessel", "filter_order": 2, "margin_ms": 10},
         "whitening": {"mode": "local", "regularize": False},
-        "detection": {"peak_sign": "neg", "detect_threshold": 4},
+        "detection": {"peak_sign": "neg", "detect_threshold": 5},
         "selection": {
             "method": "uniform",
             "n_peaks_per_channel": 5000,
@@ -134,7 +134,7 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
         ms_before = params["general"].get("ms_before", 2)
         ms_after = params["general"].get("ms_after", 2)
         radius_um = params["general"].get("radius_um", 75)
-        exclude_sweep_ms = params["detection"].get("exclude_sweep_ms", max(ms_before, ms_after) / 2)
+        exclude_sweep_ms = params["detection"].get("exclude_sweep_ms", max(ms_before, ms_after))
 
         ## First, we are filtering the data
         filtering_params = params["filtering"].copy()
@@ -265,6 +265,7 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
             clustering_params["ms_after"] = ms_after
             clustering_params["verbose"] = verbose
             clustering_params["tmp_folder"] = sorter_output_folder / "clustering"
+            clustering_params["debug"] = params["debug"]
             clustering_params["noise_threshold"] = detection_params.get("detect_threshold", 4)
 
             legacy = clustering_params.get("legacy", True)

src/spikeinterface/sortingcomponents/clustering/circus.py

@@ -41,7 +41,12 @@ class CircusClustering:
     """
 
     _default_params = {
-        "hdbscan_kwargs": {"min_cluster_size": 10, "allow_single_cluster": True, "min_samples": 5},
+        "hdbscan_kwargs": {
+            "min_cluster_size": 20,
+            "cluster_selection_epsilon": 0.5,
+            "cluster_selection_method": "leaf",
+            "allow_single_cluster": True,
+        },
         "cleaning_kwargs": {},
         "waveforms": {"ms_before": 2, "ms_after": 2},
         "sparsity": {"method": "snr", "amplitude_mode": "peak_to_peak", "threshold": 0.25},
@@ -51,7 +56,7 @@ class CircusClustering:
             "returns_split_count": True,
         },
         "radius_um": 100,
-        "n_svd": [5, 2],
+        "n_svd": 5,
         "few_waveforms": None,
         "ms_before": 0.5,
         "ms_after": 0.5,
@@ -60,6 +65,7 @@ class CircusClustering:
         "noise_levels": None,
         "tmp_folder": None,
         "verbose": True,
+        "debug": False,
     }
 
     @classmethod
@@ -110,7 +116,7 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
 
         from sklearn.decomposition import TruncatedSVD
 
-        tsvd = TruncatedSVD(params["n_svd"][0])
+        tsvd = TruncatedSVD(params["n_svd"])
         tsvd.fit(wfs)
 
         model_folder = tmp_folder / "tsvd_model"
@@ -166,8 +172,8 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
                 sub_data = all_pc_data[mask]
                 sub_data = sub_data.reshape(len(sub_data), -1)
 
-                if all_pc_data.shape[1] > params["n_svd"][1]:
-                    tsvd = PCA(params["n_svd"][1], whiten=True)
+                if all_pc_data.shape[1] > params["n_svd"]:
+                    tsvd = PCA(params["n_svd"], whiten=True)
                 else:
                     tsvd = PCA(all_pc_data.shape[1], whiten=True)
 
@@ -207,7 +213,12 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
             original_labels = peaks["channel_index"]
             from spikeinterface.sortingcomponents.clustering.split import split_clusters
 
-            min_size = 2 * params["hdbscan_kwargs"].get("min_cluster_size", 10)
+            min_size = 2 * params["hdbscan_kwargs"].get("min_cluster_size", 20)
+
+            if params["debug"]:
+                debug_folder = tmp_folder / "split"
+            else:
+                debug_folder = None
 
             peak_labels, _ = split_clusters(
                 original_labels,
@@ -221,9 +232,9 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
                     waveforms_sparse_mask=sparse_mask,
                     min_size_split=min_size,
                     clusterer_kwargs=d["hdbscan_kwargs"],
-                    n_pca_features=params["n_svd"][1],
-                    scale_n_pca_by_depth=True,
+                    n_pca_features=5,
                 ),
+                debug_folder=debug_folder,
                 **params["recursive_kwargs"],
                 **job_kwargs,
             )

src/spikeinterface/sortingcomponents/clustering/method_list.py

@@ -3,7 +3,8 @@
 from .dummy import DummyClustering
 from .position import PositionClustering
 from .sliding_hdbscan import SlidingHdbscanClustering
-from .sliding_nn import SlidingNNClustering
+
+# from .sliding_nn import SlidingNNClustering
 from .position_and_pca import PositionAndPCAClustering
 from .position_ptp_scaled import PositionPTPScaledClustering
 from .position_and_features import PositionAndFeaturesClustering
@@ -17,7 +18,7 @@
     "position_ptp_scaled": PositionPTPScaledClustering,
     "position_and_pca": PositionAndPCAClustering,
     "sliding_hdbscan": SlidingHdbscanClustering,
-    "sliding_nn": SlidingNNClustering,
+    # "sliding_nn": SlidingNNClustering,
     "position_and_features": PositionAndFeaturesClustering,
     "random_projections": RandomProjectionClustering,
     "circus": CircusClustering,

src/spikeinterface/sortingcomponents/clustering/random_projections.py

@@ -56,6 +56,7 @@ class RandomProjectionClustering:
         "noise_threshold": 4,
         "tmp_folder": None,
         "verbose": True,
+        "debug": False,
     }
 
     @classmethod

src/spikeinterface/sortingcomponents/clustering/split.py

@@ -29,6 +29,7 @@ def split_clusters(
     recursive=False,
     recursive_depth=None,
     returns_split_count=False,
+    debug_folder=None,
     **job_kwargs,
 ):
     """
@@ -42,7 +43,7 @@ def split_clusters(
         Recording object
     features_dict_or_folder: dict or folder
         A dictionary of features precomputed with peak_pipeline or a folder containing npz file for features
-    method: str, default: "hdbscan_on_local_pca"
+    method: str, default: "local_feature_clustering"
         The method name
     method_kwargs: dict, default: dict()
         The method option
@@ -70,7 +71,7 @@ def split_clusters(
     original_labels = peak_labels
     peak_labels = peak_labels.copy()
     split_count = np.zeros(peak_labels.size, dtype=int)
-
+    recursion_level = 1
     Executor = get_poolexecutor(n_jobs)
 
     with Executor(
@@ -82,19 +83,33 @@ def split_clusters(
         labels_set = np.setdiff1d(peak_labels, [-1])
         current_max_label = np.max(labels_set) + 1
         jobs = []
+
+        if debug_folder is not None:
+            if debug_folder.exists():
+                import shutil
+
+                shutil.rmtree(debug_folder)
+            debug_folder.mkdir(parents=True, exist_ok=True)
+
         for label in labels_set:
             peak_indices = np.flatnonzero(peak_labels == label)
+            if debug_folder is not None:
+                sub_folder = str(debug_folder / f"split_{label}")
+
+            else:
+                sub_folder = None
             if peak_indices.size > 0:
-                jobs.append(pool.submit(split_function_wrapper, peak_indices, 1))
+                jobs.append(pool.submit(split_function_wrapper, peak_indices, recursion_level, sub_folder))
 
         if progress_bar:
-            iterator = tqdm(jobs, desc=f"split_clusters with {method}", total=len(labels_set))
-        else:
-            iterator = jobs
+            pbar = tqdm(desc=f"split_clusters with {method}", total=len(labels_set))
+
+        for res in jobs:
+            is_split, local_labels, peak_indices, sub_folder = res.result()
+
+            if progress_bar:
+                pbar.update(1)
 
-        for res in iterator:
-            is_split, local_labels, peak_indices = res.result()
-            # print(is_split, local_labels, peak_indices)
             if not is_split:
                 continue
 
@@ -117,11 +132,21 @@ def split_clusters(
                     new_labels_set = np.setdiff1d(peak_labels[peak_indices], [-1])
                     for label in new_labels_set:
                         peak_indices = np.flatnonzero(peak_labels == label)
+                        if sub_folder is not None:
+                            new_sub_folder = sub_folder + f"_{label}"
+                        else:
+                            new_sub_folder = None
                         if peak_indices.size > 0:
                             # print('Relaunched', label, len(peak_indices), recursion_level)
-                            jobs.append(pool.submit(split_function_wrapper, peak_indices, recursion_level))
+                            jobs.append(
+                                pool.submit(split_function_wrapper, peak_indices, recursion_level, new_sub_folder)
+                            )
                             if progress_bar:
-                                iterator.total += 1
+                                pbar.total += 1
+
+        if progress_bar:
+            pbar.close()
+            del pbar
 
     if returns_split_count:
         return peak_labels, split_count
@@ -149,13 +174,13 @@ def split_worker_init(
     _ctx["peaks"] = _ctx["features"]["peaks"]
 
 
-def split_function_wrapper(peak_indices, recursion_level):
+def split_function_wrapper(peak_indices, recursion_level, debug_folder):
     global _ctx
     with threadpool_limits(limits=_ctx["max_threads_per_worker"]):
         is_split, local_labels = _ctx["method_class"].split(
-            peak_indices, _ctx["peaks"], _ctx["features"], recursion_level, **_ctx["method_kwargs"]
+            peak_indices, _ctx["peaks"], _ctx["features"], recursion_level, debug_folder, **_ctx["method_kwargs"]
         )
-    return is_split, local_labels, peak_indices
+    return is_split, local_labels, peak_indices, debug_folder
 
 
 class LocalFeatureClustering:
@@ -178,14 +203,14 @@ def split(
         peaks,
         features,
         recursion_level=1,
+        debug_folder=None,
         clusterer="hdbscan",
         feature_name="sparse_tsvd",
         neighbours_mask=None,
         waveforms_sparse_mask=None,
         clusterer_kwargs={"min_cluster_size": 25},
         min_size_split=25,
         n_pca_features=2,
-        scale_n_pca_by_depth=False,
         minimum_overlap_ratio=0.25,
     ):
         local_labels = np.zeros(peak_indices.size, dtype=np.int64)
@@ -213,26 +238,23 @@ def split(
 
         local_labels[dont_have_channels] = -2
         kept = np.flatnonzero(~dont_have_channels)
-        # print(recursion_level, kept.size, min_size_split)
+
         if kept.size < min_size_split:
             return False, None
 
         aligned_wfs = aligned_wfs[kept, :, :]
-
         flatten_features = aligned_wfs.reshape(aligned_wfs.shape[0], -1)
 
+        is_split = False
+
         if flatten_features.shape[1] > n_pca_features:
             from sklearn.decomposition import PCA
 
             # from sklearn.decomposition import TruncatedSVD
-
-            if scale_n_pca_by_depth:
-                # tsvd = TruncatedSVD(n_pca_features * recursion_level)
-                tsvd = PCA(n_pca_features * recursion_level, whiten=True)
-            else:
-                # tsvd = TruncatedSVD(n_pca_features)
-                tsvd = PCA(n_pca_features, whiten=True)
+            # tsvd = TruncatedSVD(n_pca_features)
+            tsvd = PCA(n_pca_features, whiten=True)
             final_features = tsvd.fit_transform(flatten_features)
+            del tsvd
         else:
             final_features = flatten_features
 
@@ -243,6 +265,7 @@ def split(
             clust.fit(final_features)
             possible_labels = clust.labels_
             is_split = np.setdiff1d(possible_labels, [-1]).size > 1
+            del clust
         elif clusterer == "isocut5":
             min_cluster_size = clusterer_kwargs["min_cluster_size"]
             dipscore, cutpoint = isocut5(final_features[:, 0])
@@ -255,32 +278,43 @@ def split(
             else:
                 is_split = False
         else:
-            raise ValueError(f"wrong clusterer {clusterer}")
+            raise ValueError(f"wrong clusterer {clusterer}. Possible options are 'hdbscan' or 'isocut5'.")
+
+        DEBUG = False  # only for Sam or dirty hacking
 
-        # DEBUG = True
-        DEBUG = False
-        if DEBUG:
+        if debug_folder is not None or DEBUG:
             import matplotlib.pyplot as plt
 
             labels_set = np.setdiff1d(possible_labels, [-1])
             colors = plt.colormaps["tab10"].resampled(len(labels_set))
             colors = {k: colors(i) for i, k in enumerate(labels_set)}
             colors[-1] = "k"
-            fix, axs = plt.subplots(nrows=2)
+            fig, axs = plt.subplots(nrows=2)
 
             flatten_wfs = aligned_wfs.swapaxes(1, 2).reshape(aligned_wfs.shape[0], -1)
 
-            sl = slice(None, None, 10)
+            sl = slice(None, None, 100)
             for k in np.unique(possible_labels):
                 mask = possible_labels == k
                 ax = axs[0]
-                ax.scatter(final_features[:, 0][mask][sl], final_features[:, 1][mask][sl], s=5, color=colors[k])
-
+                ax.scatter(final_features[:, 0][mask], final_features[:, 1][mask], s=5, color=colors[k])
+                if k > -1:
+                    centroid = final_features[:, :2][mask].mean(axis=0)
+                    ax.text(centroid[0], centroid[1], f"Label {k}", fontsize=10, color="k")
                 ax = axs[1]
                 ax.plot(flatten_wfs[mask][sl].T, color=colors[k], alpha=0.5)
-
-            axs[0].set_title(f"{clusterer} {is_split} {peak_indices[0]} {np.unique(possible_labels)}")
-            plt.show()
+                if k > -1:
+                    ax.plot(np.median(flatten_wfs[mask].T, axis=1), color=colors[k], lw=2)
+                ax.set_xlabel("PCA features")
+            ymin, ymax = ax.get_ylim()
+            ax.plot([n_pca_features, n_pca_features], [ymin, ymax], "k--")
+
+            axs[0].set_title(f"{clusterer} level={recursion_level}")
+            if not DEBUG:
+                fig.savefig(str(debug_folder) + ".png")
+                plt.close(fig)
+            else:
+                plt.show()
 
         if not is_split:
             return is_split, None

src/spikeinterface/sortingcomponents/clustering/tdc.py

@@ -162,7 +162,6 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
                 waveforms_sparse_mask=sparse_mask,
                 min_size_split=min_cluster_size,
                 n_pca_features=3,
-                scale_n_pca_by_depth=True,
             ),
             recursive=True,
             recursive_depth=3,