Sparse waveforms for Spyking Circus 2

src/spikeinterface/sorters/internal/spyking_circus2.py

@@ -3,7 +3,6 @@
 import os
 import shutil
 import numpy as np
-import os
 
 from spikeinterface.core import NumpySorting, load_extractor, BaseRecording, get_noise_levels, extract_waveforms
 from spikeinterface.core.job_tools import fix_job_kwargs
@@ -21,18 +20,17 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
     sorter_name = "spykingcircus2"
 
     _default_params = {
-        "general": {"ms_before": 2, "ms_after": 2, "radius_um": 100},
-        "waveforms": {"max_spikes_per_unit": 200, "overwrite": True},
+        "general": {"ms_before": 2, "ms_after": 2, "radius_um": 75},
+        "waveforms": {"max_spikes_per_unit": 200, "overwrite": True, "sparse": True, "method": "ptp", "threshold": 1},
         "filtering": {"dtype": "float32"},
         "detection": {"peak_sign": "neg", "detect_threshold": 5},
         "selection": {"n_peaks_per_channel": 5000, "min_n_peaks": 20000},
         "localization": {},
         "clustering": {},
         "matching": {},
-        "registration": {},
         "apply_preprocessing": True,
-        "shared_memory": False,
-        "job_kwargs": {},
+        "shared_memory": True,
+        "job_kwargs": {"n_jobs": -1},
     }
 
     @classmethod
@@ -63,8 +61,6 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
         ## First, we are filtering the data
         filtering_params = params["filtering"].copy()
         if params["apply_preprocessing"]:
-            # if recording.is_filtered == True:
-            #    print('Looks like the recording is already filtered, check preprocessing!')
             recording_f = bandpass_filter(recording, **filtering_params)
             recording_f = common_reference(recording_f)
         else:
@@ -103,8 +99,11 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
         ## We launch a clustering (using hdbscan) relying on positions and features extracted on
         ## the fly from the snippets
         clustering_params = params["clustering"].copy()
-        clustering_params.update(params["waveforms"])
-        clustering_params.update(params["general"])
+        clustering_params["waveforms_kwargs"] = params["waveforms"]
+
+        for k in ["ms_before", "ms_after"]:
+            clustering_params["waveforms_kwargs"][k] = params["general"][k]
+
         clustering_params.update(dict(shared_memory=params["shared_memory"]))
         clustering_params["job_kwargs"] = job_kwargs
         clustering_params["tmp_folder"] = sorter_output_folder / "clustering"
@@ -126,6 +125,9 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
         waveforms_params = params["waveforms"].copy()
         waveforms_params.update(job_kwargs)
 
+        for k in ["ms_before", "ms_after"]:
+            waveforms_params[k] = params["general"][k]
+
         if params["shared_memory"]:
             mode = "memory"
             waveforms_folder = None
@@ -143,6 +145,10 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
         matching_params.update({"noise_levels": noise_levels})
 
         matching_job_params = job_kwargs.copy()
+        for value in ["chunk_size", "chunk_memory", "total_memory", "chunk_duration"]:
+            if value in matching_job_params:
+                matching_job_params.pop(value)
+
         matching_job_params["chunk_duration"] = "100ms"
 
         spikes = find_spikes_from_templates(

src/spikeinterface/sortingcomponents/benchmark/benchmark_matching.py

@@ -600,29 +600,38 @@ def plot_comparison_matching(
             else:
                 ax = axs[j]
             comp1, comp2 = comp_per_method[method1], comp_per_method[method2]
-            for performance, color in zip(performance_names, colors):
-                perf1 = comp1.get_performance()[performance]
-                perf2 = comp2.get_performance()[performance]
-                ax.plot(perf2, perf1, ".", label=performance, color=color)
-            ax.plot([0, 1], [0, 1], "k--", alpha=0.5)
-            ax.set_ylim(ylim)
-            ax.set_xlim(ylim)
-            ax.spines[["right", "top"]].set_visible(False)
-            ax.set_aspect("equal")
-
-            if j == 0:
-                ax.set_ylabel(f"{method1}")
-            else:
-                ax.set_yticks([])
-            if i == num_methods - 1:
-                ax.set_xlabel(f"{method2}")
+            if i <= j:
+                for performance, color in zip(performance_names, colors):
+                    perf1 = comp1.get_performance()[performance]
+                    perf2 = comp2.get_performance()[performance]
+                    ax.plot(perf2, perf1, ".", label=performance, color=color)
+
+                ax.plot([0, 1], [0, 1], "k--", alpha=0.5)
+                ax.set_ylim(ylim)
+                ax.set_xlim(ylim)
+                ax.spines[["right", "top"]].set_visible(False)
+                ax.set_aspect("equal")
+
+                if j == i:
+                    ax.set_ylabel(f"{method1}")
+                else:
+                    ax.set_yticks([])
+                if i == j:
+                    ax.set_xlabel(f"{method2}")
+                else:
+                    ax.set_xticks([])
+                if i == num_methods - 1 and j == num_methods - 1:
+                    patches = []
+                    for color, name in zip(colors, performance_names):
+                        patches.append(mpatches.Patch(color=color, label=name))
+                    ax.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc="upper left", borderaxespad=0.0)
             else:
+                ax.spines["bottom"].set_visible(False)
+                ax.spines["left"].set_visible(False)
+                ax.spines["top"].set_visible(False)
+                ax.spines["right"].set_visible(False)
                 ax.set_xticks([])
-            if i == num_methods - 1 and j == num_methods - 1:
-                patches = []
-                for color, name in zip(colors, performance_names):
-                    patches.append(mpatches.Patch(color=color, label=name))
-                ax.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc="upper left", borderaxespad=0.0)
+                ax.set_yticks([])
     plt.tight_layout(h_pad=0, w_pad=0)
     return fig, axs
 

src/spikeinterface/sortingcomponents/clustering/clustering_tools.py

@@ -536,7 +536,6 @@ def remove_duplicates_via_matching(
     waveform_extractor,
     noise_levels,
     peak_labels,
-    sparsify_threshold=1,
     method_kwargs={},
     job_kwargs={},
     tmp_folder=None,
@@ -552,16 +551,20 @@ def remove_duplicates_via_matching(
     from pathlib import Path
 
     job_kwargs = fix_job_kwargs(job_kwargs)
+
+    if waveform_extractor.is_sparse():
+        sparsity = waveform_extractor.sparsity.mask
+
     templates = waveform_extractor.get_all_templates(mode="median").copy()
     nb_templates = len(templates)
     duration = waveform_extractor.nbefore + waveform_extractor.nafter
 
     fs = waveform_extractor.recording.get_sampling_frequency()
     num_chans = waveform_extractor.recording.get_num_channels()
 
-    for t in range(nb_templates):
-        is_silent = templates[t].ptp(0) < sparsify_threshold
-        templates[t, :, is_silent] = 0
+    if waveform_extractor.is_sparse():
+        for count, unit_id in enumerate(waveform_extractor.sorting.unit_ids):
+            templates[count][:, ~sparsity[count]] = 0
 
     zdata = templates.reshape(nb_templates, -1)
 
@@ -581,6 +584,7 @@ def remove_duplicates_via_matching(
 
     recording = BinaryRecordingExtractor(tmp_filename, num_channels=num_chans, sampling_frequency=fs, dtype="float32")
     recording.annotate(is_filtered=True)
+    recording = recording.set_probe(waveform_extractor.recording.get_probe())
 
     margin = 2 * max(waveform_extractor.nbefore, waveform_extractor.nafter)
     half_marging = margin // 2
@@ -597,7 +601,6 @@ def remove_duplicates_via_matching(
             "waveform_extractor": waveform_extractor,
             "noise_levels": noise_levels,
             "amplitudes": [0.95, 1.05],
-            "sparsify_threshold": sparsify_threshold,
             "omp_min_sps": 0.1,
             "templates": None,
             "overlaps": None,

src/spikeinterface/sortingcomponents/clustering/random_projections.py

@@ -41,7 +41,6 @@ class RandomProjectionClustering:
         "ms_before": 1.5,
         "ms_after": 1.5,
         "random_seed": 42,
-        "cleaning_method": "matching",
         "shared_memory": False,
         "min_values": {"ptp": 0, "energy": 0},
         "tmp_folder": None,
@@ -160,86 +159,60 @@ def main_function(cls, recording, peaks, params):
         spikes["segment_index"] = peaks[mask]["segment_index"]
         spikes["unit_index"] = peak_labels[mask]
 
-        cleaning_method = params["cleaning_method"]
-
         if verbose:
-            print("We found %d raw clusters, starting to clean with %s..." % (len(labels), cleaning_method))
-
-        if cleaning_method == "cosine":
-            wfs_arrays = extract_waveforms_to_buffers(
-                recording,
-                spikes,
-                labels,
-                nbefore,
-                nafter,
-                mode="shared_memory",
-                return_scaled=False,
-                folder=None,
-                dtype=recording.get_dtype(),
-                sparsity_mask=None,
-                copy=True,
-                **params["job_kwargs"],
-            )
-
-            labels, peak_labels = remove_duplicates(
-                wfs_arrays, noise_levels, peak_labels, num_samples, num_chans, **params["cleaning_kwargs"]
-            )
-
-        elif cleaning_method == "dip":
-            wfs_arrays = {}
-            for label in labels:
-                mask = label == peak_labels
-                wfs_arrays[label] = hdbscan_data[mask]
-
-            labels, peak_labels = remove_duplicates_via_dip(wfs_arrays, peak_labels, **params["cleaning_kwargs"])
-
-        elif cleaning_method == "matching":
-            # create a tmp folder
-            if params["tmp_folder"] is None:
-                name = "".join(random.choices(string.ascii_uppercase + string.digits, k=8))
-                tmp_folder = get_global_tmp_folder() / name
-            else:
-                tmp_folder = Path(params["tmp_folder"])
-
-            if params["shared_memory"]:
-                waveform_folder = None
-                mode = "memory"
-            else:
-                waveform_folder = tmp_folder / "waveforms"
-                mode = "folder"
-
-            sorting_folder = tmp_folder / "sorting"
-            sorting = NumpySorting.from_times_labels(spikes["sample_index"], spikes["unit_index"], fs)
-            sorting = sorting.save(folder=sorting_folder)
-            we = extract_waveforms(
-                recording,
-                sorting,
-                waveform_folder,
-                ms_before=params["ms_before"],
-                ms_after=params["ms_after"],
-                **params["job_kwargs"],
-                return_scaled=False,
-                mode=mode,
-            )
-
-            cleaning_matching_params = params["job_kwargs"].copy()
-            cleaning_matching_params["chunk_duration"] = "100ms"
-            cleaning_matching_params["n_jobs"] = 1
-            cleaning_matching_params["verbose"] = False
-            cleaning_matching_params["progress_bar"] = False
-
-            cleaning_params = params["cleaning_kwargs"].copy()
-            cleaning_params["tmp_folder"] = tmp_folder
-
-            labels, peak_labels = remove_duplicates_via_matching(
-                we, noise_levels, peak_labels, job_kwargs=cleaning_matching_params, **cleaning_params
-            )
-
-            if params["tmp_folder"] is None:
-                shutil.rmtree(tmp_folder)
-            else:
+            print("We found %d raw clusters, starting to clean with matching..." % (len(labels)))
+
+        # create a tmp folder
+        if params["tmp_folder"] is None:
+            name = "".join(random.choices(string.ascii_uppercase + string.digits, k=8))
+            tmp_folder = get_global_tmp_folder() / name
+        else:
+            tmp_folder = Path(params["tmp_folder"])
+
+        if params["shared_memory"]:
+            waveform_folder = None
+            mode = "memory"
+        else:
+            waveform_folder = tmp_folder / "waveforms"
+            mode = "folder"
+
+        sorting_folder = tmp_folder / "sorting"
+        unit_ids = np.arange(len(np.unique(spikes["unit_index"])))
+        sorting = NumpySorting(spikes, fs, unit_ids=unit_ids)
+        sorting = sorting.save(folder=sorting_folder)
+        we = extract_waveforms(
+            recording,
+            sorting,
+            waveform_folder,
+            ms_before=params["ms_before"],
+            ms_after=params["ms_after"],
+            **params["job_kwargs"],
+            return_scaled=False,
+            mode=mode,
+        )
+
+        cleaning_matching_params = params["job_kwargs"].copy()
+        for value in ["chunk_size", "chunk_memory", "total_memory", "chunk_duration"]:
+            if value in cleaning_matching_params:
+                cleaning_matching_params.pop(value)
+        cleaning_matching_params["chunk_duration"] = "100ms"
+        cleaning_matching_params["n_jobs"] = 1
+        cleaning_matching_params["verbose"] = False
+        cleaning_matching_params["progress_bar"] = False
+
+        cleaning_params = params["cleaning_kwargs"].copy()
+        cleaning_params["tmp_folder"] = tmp_folder
+
+        labels, peak_labels = remove_duplicates_via_matching(
+            we, noise_levels, peak_labels, job_kwargs=cleaning_matching_params, **cleaning_params
+        )
+
+        if params["tmp_folder"] is None:
+            shutil.rmtree(tmp_folder)
+        else:
+            if not params["shared_memory"]:
                 shutil.rmtree(tmp_folder / "waveforms")
-                shutil.rmtree(tmp_folder / "sorting")
+            shutil.rmtree(tmp_folder / "sorting")
 
         if verbose:
             print("We kept %d non-duplicated clusters..." % len(labels))