Merge branch 'main' into improve_inter_sample_shift_docstring

Author: h-mayorquin

src/spikeinterface/core/sortinganalyzer.py

@@ -2092,6 +2092,13 @@ def load_data(self):
                     import pandas as pd
 
                     ext_data = pd.read_csv(ext_data_file, index_col=0)
+                    # we need to cast the index to the unit id dtype (int or str)
+                    unit_ids = self.sorting_analyzer.unit_ids
+                    if ext_data.shape[0] == unit_ids.size:
+                        # we force dtype to be the same as unit_ids
+                        if ext_data.index.dtype != unit_ids.dtype:
+                            ext_data.index = ext_data.index.astype(unit_ids.dtype)
+
                 elif ext_data_file.suffix == ".pkl":
                     with ext_data_file.open("rb") as f:
                         ext_data = pickle.load(f)

src/spikeinterface/core/template.py

@@ -205,6 +205,7 @@ def to_sparse(self, sparsity):
             unit_ids=self.unit_ids,
             probe=self.probe,
             check_for_consistent_sparsity=self.check_for_consistent_sparsity,
+            is_scaled=self.is_scaled,
         )
 
     def get_one_template_dense(self, unit_index):

src/spikeinterface/sorters/internal/spyking_circus2.py

@@ -6,12 +6,16 @@
 import numpy as np
 
 from spikeinterface.core import NumpySorting
-from spikeinterface.core.job_tools import fix_job_kwargs
+from spikeinterface.core.job_tools import fix_job_kwargs, split_job_kwargs
 from spikeinterface.core.recording_tools import get_noise_levels
 from spikeinterface.core.template import Templates
 from spikeinterface.core.waveform_tools import estimate_templates
 from spikeinterface.preprocessing import common_reference, whiten, bandpass_filter, correct_motion
-from spikeinterface.sortingcomponents.tools import cache_preprocessing
+from spikeinterface.sortingcomponents.tools import (
+    cache_preprocessing,
+    get_prototype_and_waveforms_from_recording,
+    get_shuffled_recording_slices,
+)
 from spikeinterface.core.basesorting import minimum_spike_dtype
 from spikeinterface.core.sparsity import compute_sparsity
 from spikeinterface.core.sortinganalyzer import create_sorting_analyzer
@@ -26,7 +30,7 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
     _default_params = {
         "general": {"ms_before": 2, "ms_after": 2, "radius_um": 75},
         "sparsity": {"method": "snr", "amplitude_mode": "peak_to_peak", "threshold": 0.25},
-        "filtering": {"freq_min": 150, "freq_max": 7000, "ftype": "bessel", "filter_order": 2},
+        "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},
         "selection": {
@@ -53,6 +57,7 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
         "cache_preprocessing": {"mode": "memory", "memory_limit": 0.5, "delete_cache": True},
         "multi_units_only": False,
         "job_kwargs": {"n_jobs": 0.5},
+        "seed": 42,
         "debug": False,
     }
 
@@ -74,18 +79,21 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
         "merging": "A dictionary to specify the final merging param to group cells after template matching (get_potential_auto_merge)",
         "motion_correction": "A dictionary to be provided if motion correction has to be performed (dense probe only)",
         "apply_preprocessing": "Boolean to specify whether circus 2 should preprocess the recording or not. If yes, then high_pass filtering + common\
-                                                    median reference + zscore",
+                                                    median reference + whitening",
+        "apply_motion_correction": "Boolean to specify whether circus 2 should apply motion correction to the recording or not",
+        "matched_filtering": "Boolean to specify whether circus 2 should detect peaks via matched filtering (slightly slower)",
         "cache_preprocessing": "How to cache the preprocessed recording. Mode can be memory, file, zarr, with extra arguments. In case of memory (default), \
                          memory_limit will control how much RAM can be used. In case of folder or zarr, delete_cache controls if cache is cleaned after sorting",
         "multi_units_only": "Boolean to get only multi units activity (i.e. one template per electrode)",
         "job_kwargs": "A dictionary to specify how many jobs and which parameters they should used",
+        "seed": "An int to control how chunks are shuffled while detecting peaks",
         "debug": "Boolean to specify if internal data structures made during the sorting should be kept for debugging",
     }
 
     sorter_description = """Spyking Circus 2 is a rewriting of Spyking Circus, within the SpikeInterface framework
     It uses a more conservative clustering algorithm (compared to Spyking Circus), which is less prone to hallucinate units and/or find noise.
     In addition, it also uses a full Orthogonal Matching Pursuit engine to reconstruct the traces, leading to more spikes
-    being discovered."""
+    being discovered. The code is much faster and memory efficient, inheriting from all the preprocessing possibilities of spikeinterface"""
 
     @classmethod
     def get_sorter_version(cls):
@@ -114,7 +122,7 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
         from spikeinterface.sortingcomponents.clustering import find_cluster_from_peaks
         from spikeinterface.sortingcomponents.matching import find_spikes_from_templates
         from spikeinterface.sortingcomponents.tools import remove_empty_templates
-        from spikeinterface.sortingcomponents.tools import get_prototype_spike, check_probe_for_drift_correction
+        from spikeinterface.sortingcomponents.tools import check_probe_for_drift_correction
 
         job_kwargs = fix_job_kwargs(params["job_kwargs"])
         job_kwargs.update({"progress_bar": verbose})
@@ -131,10 +139,14 @@ 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 verbose:
+                print("Preprocessing the recording (bandpass filtering + CMR + whitening)")
             recording_f = bandpass_filter(recording, **filtering_params, dtype="float32")
             if num_channels > 1:
                 recording_f = common_reference(recording_f)
         else:
+            if verbose:
+                print("Skipping preprocessing (whitening only)")
             recording_f = recording
             recording_f.annotate(is_filtered=True)
 
@@ -157,12 +169,14 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
         # TODO add , regularize=True chen ready
         whitening_kwargs = params["whitening"].copy()
         whitening_kwargs["dtype"] = "float32"
-        whitening_kwargs["radius_um"] = radius_um
+        whitening_kwargs["regularize"] = whitening_kwargs.get("regularize", False)
         if num_channels == 1:
             whitening_kwargs["regularize"] = False
+        if whitening_kwargs["regularize"]:
+            whitening_kwargs["regularize_kwargs"] = {"method": "LedoitWolf"}
 
         recording_w = whiten(recording_f, **whitening_kwargs)
-        noise_levels = get_noise_levels(recording_w, return_scaled=False)
+        noise_levels = get_noise_levels(recording_w, return_scaled=False, **job_kwargs)
 
         if recording_w.check_serializability("json"):
             recording_w.dump(sorter_output_folder / "preprocessed_recording.json", relative_to=None)
@@ -173,42 +187,69 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
 
         ## Then, we are detecting peaks with a locally_exclusive method
         detection_params = params["detection"].copy()
-        detection_params.update(job_kwargs)
-
-        detection_params["radius_um"] = detection_params.get("radius_um", 50)
+        selection_params = params["selection"].copy()
+        detection_params["radius_um"] = radius_um
         detection_params["exclude_sweep_ms"] = exclude_sweep_ms
         detection_params["noise_levels"] = noise_levels
 
         fs = recording_w.get_sampling_frequency()
         nbefore = int(ms_before * fs / 1000.0)
         nafter = int(ms_after * fs / 1000.0)
 
+        skip_peaks = not params["multi_units_only"] and selection_params.get("method", "uniform") == "uniform"
+        max_n_peaks = selection_params["n_peaks_per_channel"] * num_channels
+        n_peaks = max(selection_params["min_n_peaks"], max_n_peaks)
+
+        if params["debug"]:
+            clustering_folder = sorter_output_folder / "clustering"
+            clustering_folder.mkdir(parents=True, exist_ok=True)
+            np.save(clustering_folder / "noise_levels.npy", noise_levels)
+
         if params["matched_filtering"]:
-            peaks = detect_peaks(recording_w, "locally_exclusive", **detection_params, skip_after_n_peaks=5000)
-            prototype = get_prototype_spike(recording_w, peaks, ms_before, ms_after, **job_kwargs)
+            prototype, waveforms, _ = get_prototype_and_waveforms_from_recording(
+                recording_w,
+                n_peaks=10000,
+                ms_before=ms_before,
+                ms_after=ms_after,
+                seed=params["seed"],
+                **detection_params,
+                **job_kwargs,
+            )
             detection_params["prototype"] = prototype
             detection_params["ms_before"] = ms_before
-            peaks = detect_peaks(recording_w, "matched_filtering", **detection_params)
+            if params["debug"]:
+                np.save(clustering_folder / "waveforms.npy", waveforms)
+                np.save(clustering_folder / "prototype.npy", prototype)
+            if skip_peaks:
+                detection_params["skip_after_n_peaks"] = n_peaks
+                detection_params["recording_slices"] = get_shuffled_recording_slices(
+                    recording_w, seed=params["seed"], **job_kwargs
+                )
+            peaks = detect_peaks(recording_w, "matched_filtering", **detection_params, **job_kwargs)
         else:
-            peaks = detect_peaks(recording_w, "locally_exclusive", **detection_params)
+            waveforms = None
+            if skip_peaks:
+                detection_params["skip_after_n_peaks"] = n_peaks
+                detection_params["recording_slices"] = get_shuffled_recording_slices(
+                    recording_w, seed=params["seed"], **job_kwargs
+                )
+            peaks = detect_peaks(recording_w, "locally_exclusive", **detection_params, **job_kwargs)
 
-        if verbose:
-            print("We found %d peaks in total" % len(peaks))
+        if not skip_peaks and verbose:
+            print("Found %d peaks in total" % len(peaks))
 
         if params["multi_units_only"]:
             sorting = NumpySorting.from_peaks(peaks, sampling_frequency, unit_ids=recording_w.unit_ids)
         else:
             ## We subselect a subset of all the peaks, by making the distributions os SNRs over all
             ## channels as flat as possible
             selection_params = params["selection"]
-            selection_params["n_peaks"] = min(len(peaks), selection_params["n_peaks_per_channel"] * num_channels)
-            selection_params["n_peaks"] = max(selection_params["min_n_peaks"], selection_params["n_peaks"])
-
+            selection_params["n_peaks"] = n_peaks
             selection_params.update({"noise_levels": noise_levels})
             selected_peaks = select_peaks(peaks, **selection_params)
 
             if verbose:
-                print("We kept %d peaks for clustering" % len(selected_peaks))
+                print("Kept %d peaks for clustering" % len(selected_peaks))
 
             ## We launch a clustering (using hdbscan) relying on positions and features extracted on
             ## the fly from the snippets
@@ -218,10 +259,13 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
             clustering_params["radius_um"] = radius_um
             clustering_params["waveforms"]["ms_before"] = ms_before
             clustering_params["waveforms"]["ms_after"] = ms_after
+            clustering_params["few_waveforms"] = waveforms
             clustering_params["noise_levels"] = noise_levels
-            clustering_params["ms_before"] = exclude_sweep_ms
-            clustering_params["ms_after"] = exclude_sweep_ms
+            clustering_params["ms_before"] = ms_before
+            clustering_params["ms_after"] = ms_after
+            clustering_params["verbose"] = verbose
             clustering_params["tmp_folder"] = sorter_output_folder / "clustering"
+            clustering_params["noise_threshold"] = detection_params.get("detect_threshold", 4)
 
             legacy = clustering_params.get("legacy", True)
 
@@ -246,12 +290,8 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
             unit_ids = np.arange(len(np.unique(labeled_peaks["unit_index"])))
             sorting = NumpySorting(labeled_peaks, sampling_frequency, unit_ids=unit_ids)
 
-            clustering_folder = sorter_output_folder / "clustering"
-            clustering_folder.mkdir(parents=True, exist_ok=True)
-
-            if not params["debug"]:
-                shutil.rmtree(clustering_folder)
-            else:
+            if params["debug"]:
+                np.save(clustering_folder / "peak_labels", peak_labels)
                 np.save(clustering_folder / "labels", labels)
                 np.save(clustering_folder / "peaks", selected_peaks)
 
@@ -294,7 +334,7 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
                     np.save(fitting_folder / "spikes", spikes)
 
                 if verbose:
-                    print("We found %d spikes" % len(spikes))
+                    print("Found %d spikes" % len(spikes))
 
                 ## And this is it! We have a spyking circus
                 sorting = np.zeros(spikes.size, dtype=minimum_spike_dtype)
@@ -334,10 +374,10 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
                 sorting.save(folder=curation_folder)
                 # np.save(fitting_folder / "amplitudes", guessed_amplitudes)
 
-            sorting = final_cleaning_circus(recording_w, sorting, templates, **merging_params)
+            sorting = final_cleaning_circus(recording_w, sorting, templates, merging_params, **job_kwargs)
 
             if verbose:
-                print(f"Final merging, keeping {len(sorting.unit_ids)} units")
+                print(f"Kept {len(sorting.unit_ids)} units after final merging")
 
         folder_to_delete = None
         cache_mode = params["cache_preprocessing"].get("mode", "memory")
@@ -376,17 +416,18 @@ def create_sorting_analyzer_with_templates(sorting, recording, templates, remove
     return sa
 
 
-def final_cleaning_circus(recording, sorting, templates, **merging_kwargs):
+def final_cleaning_circus(recording, sorting, templates, merging_kwargs, **job_kwargs):
 
     from spikeinterface.core.sorting_tools import apply_merges_to_sorting
 
     sa = create_sorting_analyzer_with_templates(sorting, recording, templates)
 
-    sa.compute("unit_locations", method="monopolar_triangulation")
+    sa.compute("unit_locations", method="monopolar_triangulation", **job_kwargs)
     similarity_kwargs = merging_kwargs.pop("similarity_kwargs", {})
-    sa.compute("template_similarity", **similarity_kwargs)
+    sa.compute("template_similarity", **similarity_kwargs, **job_kwargs)
     correlograms_kwargs = merging_kwargs.pop("correlograms_kwargs", {})
-    sa.compute("correlograms", **correlograms_kwargs)
+    sa.compute("correlograms", **correlograms_kwargs, **job_kwargs)
+
     auto_merge_kwargs = merging_kwargs.pop("auto_merge", {})
     merges = get_potential_auto_merge(sa, resolve_graph=True, **auto_merge_kwargs)
     sorting = apply_merges_to_sorting(sa.sorting, merges)