Merge branch 'main' into prepare_release

Author: alejoe91

.github/scripts/test_kilosort4_ci.py

@@ -76,7 +76,7 @@
     "x_centers": 5,
     "binning_depth": 1,
     "drift_smoothing": [250, 250, 250],
-    "artifact_threshold": 200,
+    "artifact_threshold": 500,
     "ccg_threshold": 1e12,
     "acg_threshold": 1e12,
     "cluster_downsampling": 2,

src/spikeinterface/benchmark/benchmark_plot_tools.py

@@ -478,8 +478,9 @@ def _plot_performances_vs_metric(
                     .get_performance()[performance_name]
                     .to_numpy(dtype="float64")
                 )
-                all_xs.append(x)
-                all_ys.append(y)
+                mask = ~np.isnan(x) & ~np.isnan(y)
+                all_xs.append(x[mask])
+                all_ys.append(y[mask])
 
             if with_sigmoid_fit:
                 max_snr = max(np.max(x) for x in all_xs)

src/spikeinterface/core/sortinganalyzer.py

@@ -1634,11 +1634,17 @@ def compute(self, input, save=True, extension_params=None, verbose=False, **kwar
             return self.compute_one_extension(extension_name=input, save=save, verbose=verbose, **kwargs)
         elif isinstance(input, dict):
             params_, job_kwargs = split_job_kwargs(kwargs)
-            assert len(params_) == 0, "Too many arguments for SortingAnalyzer.compute_several_extensions()"
+            assert len(params_) == 0, (
+                "Too many arguments for SortingAnalyzer.compute_several_extensions(), "
+                f"please remove the arguments {set(params_)} from the compute function."
+            )
             self.compute_several_extensions(extensions=input, save=save, verbose=verbose, **job_kwargs)
         elif isinstance(input, list):
             params_, job_kwargs = split_job_kwargs(kwargs)
-            assert len(params_) == 0, "Too many arguments for SortingAnalyzer.compute_several_extensions()"
+            assert len(params_) == 0, (
+                "Too many arguments for SortingAnalyzer.compute_several_extensions(), "
+                f"please remove the arguments {set(params_)} from the compute function."
+            )
             extensions = {k: {} for k in input}
             if extension_params is not None:
                 for ext_name, ext_params in extension_params.items():

src/spikeinterface/postprocessing/template_similarity.py

@@ -94,7 +94,7 @@ def _merge_extension_data(
                 new_sorting_analyzer.sparsity.mask[keep, :], new_unit_ids, new_sorting_analyzer.channel_ids
             )
 
-        new_similarity = compute_similarity_with_templates_array(
+        new_similarity, _ = compute_similarity_with_templates_array(
             new_templates_array,
             all_templates_array,
             method=self.params["method"],
@@ -146,7 +146,7 @@ def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer,
                 new_sorting_analyzer.sparsity.mask[keep, :], new_unit_ids_f, new_sorting_analyzer.channel_ids
             )
 
-        new_similarity = compute_similarity_with_templates_array(
+        new_similarity, _ = compute_similarity_with_templates_array(
             new_templates_array,
             all_templates_array,
             method=self.params["method"],
@@ -188,7 +188,7 @@ def _run(self, verbose=False):
             self.sorting_analyzer, return_in_uV=self.sorting_analyzer.return_in_uV
         )
         sparsity = self.sorting_analyzer.sparsity
-        similarity = compute_similarity_with_templates_array(
+        similarity, _ = compute_similarity_with_templates_array(
             templates_array,
             templates_array,
             method=self.params["method"],
@@ -393,7 +393,13 @@ def get_overlapping_mask_for_one_template(template_index, sparsity, other_sparsi
 
 
 def compute_similarity_with_templates_array(
-    templates_array, other_templates_array, method, support="union", num_shifts=0, sparsity=None, other_sparsity=None
+    templates_array,
+    other_templates_array,
+    method,
+    support="union",
+    num_shifts=0,
+    sparsity=None,
+    other_sparsity=None,
 ):
 
     if method == "cosine_similarity":
@@ -432,10 +438,11 @@ def compute_similarity_with_templates_array(
         templates_array, other_templates_array, num_shifts, method, sparsity_mask, other_sparsity_mask, support=support
     )
 
+    lags = np.argmin(distances, axis=0) - num_shifts
     distances = np.min(distances, axis=0)
     similarity = 1 - distances
 
-    return similarity
+    return similarity, lags
 
 
 def compute_template_similarity_by_pair(
@@ -445,7 +452,7 @@ def compute_template_similarity_by_pair(
     templates_array_2 = get_dense_templates_array(sorting_analyzer_2, return_in_uV=True)
     sparsity_1 = sorting_analyzer_1.sparsity
     sparsity_2 = sorting_analyzer_2.sparsity
-    similarity = compute_similarity_with_templates_array(
+    similarity, _ = compute_similarity_with_templates_array(
         templates_array_1,
         templates_array_2,
         method=method,

src/spikeinterface/postprocessing/tests/test_template_similarity.py

@@ -82,8 +82,9 @@ def test_compute_similarity_with_templates_array(params):
     templates_array = rng.random(size=(2, 20, 5))
     other_templates_array = rng.random(size=(4, 20, 5))
 
-    similarity = compute_similarity_with_templates_array(templates_array, other_templates_array, **params)
+    similarity, lags = compute_similarity_with_templates_array(templates_array, other_templates_array, **params)
     print(similarity.shape)
+    print(lags)
 
 
 pytest.mark.skipif(not HAVE_NUMBA, reason="Numba not available")
@@ -141,5 +142,5 @@ def test_equal_results_numba(params):
     test.cache_folder = Path("./cache_folder")
     test.test_extension(params=dict(method="l2"))
 
-    # params = dict(method="cosine", num_shifts=8)
-    # test_compute_similarity_with_templates_array(params)
+    params = dict(method="cosine", num_shifts=8)
+    test_compute_similarity_with_templates_array(params)

src/spikeinterface/sorters/external/kilosort4.py

@@ -95,7 +95,7 @@ def _dynamic_params(cls):
             # we skip some parameters that are not relevant for the user
             # n_chan_bin/sampling_frequency: retrieved from the recording
             # tmin/tmax: same ase time/frame_slice in SpikeInterface
-            skip_main = ["n_chan_bin", "sampling_frequency", "tmin", "tmax"]
+            skip_main = ["fs", "n_chan_bin", "tmin", "tmax"]
             default_params = {}
             default_params_descriptions = {}
             ks_params = ks.parameters.MAIN_PARAMETERS.copy()

src/spikeinterface/sorters/internal/spyking_circus2.py

@@ -15,15 +15,15 @@
     _set_optimal_chunk_size,
 )
 from spikeinterface.core.basesorting import minimum_spike_dtype
+from spikeinterface.core import compute_sparsity
 
 
 class Spykingcircus2Sorter(ComponentsBasedSorter):
     sorter_name = "spykingcircus2"
 
     _default_params = {
-        "general": {"ms_before": 0.5, "ms_after": 1.5, "radius_um": 100},
-        "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},
+        "general": {"ms_before": 0.5, "ms_after": 1.5, "radius_um": 100.0},
+        "filtering": {"freq_min": 150, "freq_max": 7000, "ftype": "bessel", "filter_order": 2, "margin_ms": 20},
         "whitening": {"mode": "local", "regularize": False},
         "detection": {
             "method": "matched_filtering",
@@ -38,8 +38,10 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
         "motion_correction": {"preset": "dredge_fast"},
         "merging": {"max_distance_um": 50},
         "clustering": {"method": "iterative-hdbscan", "method_kwargs": dict()},
+        "cleaning": {"min_snr": 5, "max_jitter_ms": 0.1, "sparsify_threshold": None},
         "matching": {"method": "circus-omp", "method_kwargs": dict(), "pipeline_kwargs": dict()},
         "apply_preprocessing": True,
+        "apply_whitening": True,
         "cache_preprocessing": {"mode": "memory", "memory_limit": 0.5, "delete_cache": True},
         "chunk_preprocessing": {"memory_limit": None},
         "multi_units_only": False,
@@ -85,7 +87,7 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
 
     @classmethod
     def get_sorter_version(cls):
-        return "2025.09"
+        return "2025.10"
 
     @classmethod
     def _run_from_folder(cls, sorter_output_folder, params, verbose):
@@ -114,30 +116,50 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
         num_channels = recording.get_num_channels()
         ms_before = params["general"].get("ms_before", 0.5)
         ms_after = params["general"].get("ms_after", 1.5)
-        radius_um = params["general"].get("radius_um", 100)
+        radius_um = params["general"].get("radius_um", 100.0)
         detect_threshold = params["detection"]["method_kwargs"].get("detect_threshold", 5)
         peak_sign = params["detection"].get("peak_sign", "neg")
         deterministic = params["deterministic_peaks_detection"]
         debug = params["debug"]
         seed = params["seed"]
         apply_preprocessing = params["apply_preprocessing"]
+        apply_whitening = params["apply_whitening"]
         apply_motion_correction = params["apply_motion_correction"]
         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()
         if apply_preprocessing:
             if verbose:
-                print("Preprocessing the recording (bandpass filtering + CMR + whitening)")
+                if apply_whitening:
+                    print("Preprocessing the recording (bandpass filtering + CMR + whitening)")
+                else:
+                    print("Preprocessing the recording (bandpass filtering + CMR)")
             recording_f = bandpass_filter(recording, **filtering_params, dtype="float32")
-            if num_channels > 1:
+            if num_channels >= 32:
                 recording_f = common_reference(recording_f)
         else:
             if verbose:
                 print("Skipping preprocessing (whitening only)")
             recording_f = recording
             recording_f.annotate(is_filtered=True)
 
+        if apply_whitening:
+            ## We need to whiten before the template matching step, to boost the results
+            # TODO add , regularize=True chen ready
+            whitening_kwargs = params["whitening"].copy()
+            whitening_kwargs["dtype"] = "float32"
+            whitening_kwargs["seed"] = params["seed"]
+            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"}
+                whitening_kwargs["apply_mean"] = True
+            recording_w = whiten(recording_f, **whitening_kwargs)
+        else:
+            recording_w = recording_f
+
         valid_geometry = check_probe_for_drift_correction(recording_f)
         if apply_motion_correction:
             if not valid_geometry:
@@ -151,27 +173,13 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
                 motion_correction_kwargs = params["motion_correction"].copy()
                 motion_correction_kwargs.update({"folder": motion_folder})
                 noise_levels = get_noise_levels(
-                    recording_f, return_in_uV=False, random_slices_kwargs={"seed": seed}, **job_kwargs
+                    recording_w, return_in_uV=False, random_slices_kwargs={"seed": seed}, **job_kwargs
                 )
                 motion_correction_kwargs["detect_kwargs"] = {"noise_levels": noise_levels}
-                recording_f = correct_motion(recording_f, **motion_correction_kwargs, **job_kwargs)
+                recording_w = correct_motion(recording_w, **motion_correction_kwargs, **job_kwargs)
         else:
             motion_folder = None
 
-        ## We need to whiten before the template matching step, to boost the results
-        # TODO add , regularize=True chen ready
-        whitening_kwargs = params["whitening"].copy()
-        whitening_kwargs["dtype"] = "float32"
-        whitening_kwargs["seed"] = params["seed"]
-        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"}
-            whitening_kwargs["apply_mean"] = True
-
-        recording_w = whiten(recording_f, **whitening_kwargs)
-
         noise_levels = get_noise_levels(
             recording_w, return_in_uV=False, random_slices_kwargs={"seed": seed}, **job_kwargs
         )
@@ -325,18 +333,33 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
             if not clustering_from_svd:
                 from spikeinterface.sortingcomponents.clustering.tools import get_templates_from_peaks_and_recording
 
-                templates = get_templates_from_peaks_and_recording(
+                dense_templates = get_templates_from_peaks_and_recording(
                     recording_w,
                     selected_peaks,
                     peak_labels,
                     ms_before,
                     ms_after,
                     job_kwargs=job_kwargs,
                 )
+
+                sparsity = compute_sparsity(dense_templates, method="radius", radius_um=radius_um)
+                threshold = params["cleaning"].get("sparsify_threshold", None)
+                if threshold is not None:
+                    sparsity_snr = compute_sparsity(
+                        dense_templates,
+                        method="snr",
+                        amplitude_mode="peak_to_peak",
+                        noise_levels=noise_levels,
+                        threshold=threshold,
+                    )
+                    sparsity.mask = sparsity.mask & sparsity_snr.mask
+
+                templates = dense_templates.to_sparse(sparsity)
+
             else:
                 from spikeinterface.sortingcomponents.clustering.tools import get_templates_from_peaks_and_svd
 
-                templates, _ = get_templates_from_peaks_and_svd(
+                dense_templates, new_sparse_mask = get_templates_from_peaks_and_svd(
                     recording_w,
                     selected_peaks,
                     peak_labels,
@@ -348,16 +371,14 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
                     operator="median",
                 )
                 # this release the peak_svd memmap file
+                templates = dense_templates.to_sparse(new_sparse_mask)
 
             del more_outs
 
-            templates = clean_templates(
-                templates,
-                noise_levels=noise_levels,
-                min_snr=detect_threshold,
-                max_jitter_ms=0.1,
-                remove_empty=True,
-            )
+            cleaning_kwargs = params.get("cleaning", {}).copy()
+            cleaning_kwargs["noise_levels"] = noise_levels
+            cleaning_kwargs["remove_empty"] = True
+            templates = clean_templates(templates, **cleaning_kwargs)
 
             if verbose:
                 print("Kept %d clean clusters" % len(templates.unit_ids))
@@ -416,7 +437,12 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
 
                 if sorting.get_non_empty_unit_ids().size > 0:
                     final_analyzer = final_cleaning_circus(
-                        recording_w, sorting, templates, job_kwargs=job_kwargs, **merging_params
+                        recording_w,
+                        sorting,
+                        templates,
+                        noise_levels=noise_levels,
+                        job_kwargs=job_kwargs,
+                        **merging_params,
                     )
                     final_analyzer.save_as(format="binary_folder", folder=sorter_output_folder / "final_analyzer")
 
@@ -451,14 +477,15 @@ def final_cleaning_circus(
     max_distance_um=50,
     template_diff_thresh=np.arange(0.05, 0.5, 0.05),
     debug_folder=None,
-    job_kwargs=None,
+    noise_levels=None,
+    job_kwargs=dict(),
 ):
 
     from spikeinterface.sortingcomponents.tools import create_sorting_analyzer_with_existing_templates
     from spikeinterface.curation.auto_merge import auto_merge_units
 
     # First we compute the needed extensions
-    analyzer = create_sorting_analyzer_with_existing_templates(sorting, recording, templates)
+    analyzer = create_sorting_analyzer_with_existing_templates(sorting, recording, templates, noise_levels=noise_levels)
     analyzer.compute("unit_locations", method="center_of_mass", **job_kwargs)
     analyzer.compute("template_similarity", **similarity_kwargs)