Improvements for internal sorters (#4279)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Samuel Garcia <[email protected]>
Co-authored-by: Alessio Buccino <[email protected]>

Author: 4 people

src/spikeinterface/sorters/internal/lupin.py

@@ -35,6 +35,7 @@ class LupinSorter(ComponentsBasedSorter):
 
     _default_params = {
         "apply_preprocessing": True,
+        "preprocessing_dict": None,
         "apply_motion_correction": False,
         "motion_correction_preset": "dredge_fast",
         "clustering_ms_before": 0.3,
@@ -67,7 +68,8 @@ class LupinSorter(ComponentsBasedSorter):
 
     _params_description = {
         "apply_preprocessing": "Apply internal preprocessing or not",
-        "apply_motion_correction": "Apply motion correction or not",
+        "preprocessing_dict": "Inject customized preprocessing chain via a dict, instead of the internal one",
+        "apply_motion_correction": "Apply motion correction or not (only used when apply_preprocessing=True)",
         "motion_correction_preset": "Motion correction preset",
         "clustering_ms_before": "Milliseconds before the spike peak for clustering",
         "clustering_ms_after": "Milliseconds after the spike peak  for clustering",
@@ -111,6 +113,7 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
         from spikeinterface.preprocessing import correct_motion
         from spikeinterface.sortingcomponents.motion import InterpolateMotionRecording
         from spikeinterface.sortingcomponents.tools import clean_templates, compute_sparsity_from_peaks_and_label
+        from spikeinterface.preprocessing import apply_preprocessing_pipeline
 
         job_kwargs = params["job_kwargs"].copy()
         job_kwargs = fix_job_kwargs(job_kwargs)
@@ -156,36 +159,42 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
         # preprocessing
         if params["apply_preprocessing"]:
             if params["apply_motion_correction"]:
+
                 rec_for_motion = recording_raw
-                if params["apply_preprocessing"]:
+                if params["preprocessing_dict"] is None:
                     rec_for_motion = bandpass_filter(
                         rec_for_motion, freq_min=300.0, freq_max=6000.0, ftype="bessel", dtype="float32"
                     )
                     if apply_cmr:
                         rec_for_motion = common_reference(rec_for_motion)
-                    if verbose:
-                        print("Start correct_motion()")
-                    _, motion_info = correct_motion(
-                        rec_for_motion,
-                        folder=sorter_output_folder / "motion",
-                        output_motion_info=True,
-                        preset=params["motion_correction_preset"],
-                    )
-                    if verbose:
-                        print("Done correct_motion()")
-
-            recording = bandpass_filter(
-                recording_raw,
-                freq_min=params["freq_min"],
-                freq_max=params["freq_max"],
-                ftype="bessel",
-                filter_order=2,
-                margin_ms=20.0,
-                dtype="float32",
-            )
+                else:
+                    rec_for_motion = apply_preprocessing_pipeline(rec_for_motion, params["preprocessing_dict"])
+
+                if verbose:
+                    print("Start correct_motion()")
+                _, motion_info = correct_motion(
+                    rec_for_motion,
+                    folder=sorter_output_folder / "motion",
+                    output_motion_info=True,
+                    preset=params["motion_correction_preset"],
+                )
+                if verbose:
+                    print("Done correct_motion()")
+
+            if params["preprocessing_dict"] is None:
+                recording = bandpass_filter(
+                    recording_raw,
+                    freq_min=params["freq_min"],
+                    freq_max=params["freq_max"],
+                    ftype="bessel",
+                    filter_order=2,
+                    dtype="float32",
+                )
 
-            if apply_cmr:
-                recording = common_reference(recording)
+                if apply_cmr:
+                    recording = common_reference(recording)
+            else:
+                recording = apply_preprocessing_pipeline(recording, params["preprocessing_dict"])
 
             recording = whiten(
                 recording,
@@ -219,7 +228,7 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
 
             noise_levels = get_noise_levels(recording, return_in_uV=False)
         else:
-            recording = recording_raw
+            recording = recording_raw.astype("float32")
             noise_levels = get_noise_levels(recording, return_in_uV=False)
             cache_info = None
 
@@ -366,6 +375,8 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
                 recording,
                 sorting,
                 templates,
+                amplitude_scalings=spikes["amplitude"],
+                noise_levels=noise_levels,
                 similarity_kwargs={"method": "l1", "support": "union", "max_lag_ms": 0.1},
                 sparsity_overlap=0.5,
                 censor_ms=3.0,

src/spikeinterface/sorters/internal/spyking_circus2.py

@@ -24,7 +24,7 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
 
     _default_params = {
         "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},
+        "filtering": {"freq_min": 150, "freq_max": 7000, "ftype": "bessel", "filter_order": 2},
         "whitening": {"mode": "local", "regularize": False},
         "detection": {
             "method": "matched_filtering",
@@ -68,6 +68,7 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
         "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 + whitening",
+        "apply_whitening": "Boolean to specify whether circus 2 should whiten the recording or not",
         "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), \
@@ -145,7 +146,7 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
                     print("Skipping preprocessing (whitening only)")
                 else:
                     print("Skipping preprocessing (no whitening)")
-            recording_f = recording
+            recording_f = recording.astype("float32")
             recording_f.annotate(is_filtered=True)
 
         if apply_whitening:
@@ -189,9 +190,10 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
         )
 
         if recording_w.check_serializability("json"):
-            recording_w.dump(sorter_output_folder / "preprocessed_recording.json", relative_to=None)
+            recording_dump_file = sorter_output_folder / "preprocessed_recording.json"
         elif recording_w.check_serializability("pickle"):
-            recording_w.dump(sorter_output_folder / "preprocessed_recording.pickle", relative_to=None)
+            recording_dump_file = sorter_output_folder / "preprocessed_recording.pickle"
+        recording_w.dump(recording_dump_file, relative_to=None)
 
         recording_w, cache_info = cache_preprocessing(
             recording_w, job_kwargs=job_kwargs, **params["cache_preprocessing"]
@@ -468,10 +470,15 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
                     # np.save(fitting_folder / "amplitudes", guessed_amplitudes)
 
                 if sorting.get_non_empty_unit_ids().size > 0:
+                    from spikeinterface.core import load
+
+                    recording_ww = load(recording_dump_file)
+
                     final_analyzer = final_cleaning_circus(
-                        recording_w,
+                        recording_ww,
                         sorting,
                         templates,
+                        amplitude_scalings=spikes["amplitude"],
                         noise_levels=noise_levels,
                         job_kwargs=job_kwargs,
                         **merging_params,
@@ -495,6 +502,7 @@ def final_cleaning_circus(
     recording,
     sorting,
     templates,
+    amplitude_scalings=None,
     similarity_kwargs={"method": "l1", "support": "union", "max_lag_ms": 0.1},
     sparsity_overlap=0.5,
     censor_ms=3.0,
@@ -509,7 +517,9 @@ def final_cleaning_circus(
     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, noise_levels=noise_levels)
+    analyzer = create_sorting_analyzer_with_existing_templates(
+        sorting, recording, templates, noise_levels=noise_levels, amplitude_scalings=amplitude_scalings
+    )
     analyzer.compute("unit_locations", method="center_of_mass", **job_kwargs)
     analyzer.compute("template_similarity", **similarity_kwargs)
 

src/spikeinterface/sorters/internal/tridesclous2.py

@@ -31,6 +31,7 @@ class Tridesclous2Sorter(ComponentsBasedSorter):
 
     _default_params = {
         "apply_preprocessing": True,
+        "preprocessing_dict": None,
         "apply_motion_correction": False,
         "motion_correction_preset": "dredge_fast",
         "clustering_ms_before": 0.5,
@@ -62,7 +63,8 @@ class Tridesclous2Sorter(ComponentsBasedSorter):
 
     _params_description = {
         "apply_preprocessing": "Apply internal preprocessing or not",
-        "apply_motion_correction": "Apply motion correction or not",
+        "preprocessing_dict": "Inject customized preprocessing chain via a dict, instead of the internal one",
+        "apply_motion_correction": "Apply motion correction or not (only used when apply_preprocessing=True)",
         "motion_correction_preset": "Motion correction preset",
         "clustering_ms_before": "Milliseconds before the spike peak for clustering",
         "clustering_ms_after": "Milliseconds after the spike peak  for clustering",
@@ -104,6 +106,7 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
         from spikeinterface.preprocessing import correct_motion
         from spikeinterface.sortingcomponents.motion import InterpolateMotionRecording
         from spikeinterface.sortingcomponents.tools import clean_templates, compute_sparsity_from_peaks_and_label
+        from spikeinterface.preprocessing import apply_preprocessing_pipeline
 
         job_kwargs = params["job_kwargs"].copy()
         job_kwargs = fix_job_kwargs(job_kwargs)
@@ -121,38 +124,42 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
         # preprocessing
         if params["apply_preprocessing"]:
             if params["apply_motion_correction"]:
+
                 rec_for_motion = recording_raw
-                if params["apply_preprocessing"]:
+                if params["preprocessing_dict"] is None:
                     rec_for_motion = bandpass_filter(
                         rec_for_motion, freq_min=300.0, freq_max=6000.0, ftype="bessel", dtype="float32"
                     )
                     if apply_cmr:
                         rec_for_motion = common_reference(rec_for_motion)
-                    if verbose:
-                        print("Start correct_motion()")
-                    _, motion_info = correct_motion(
-                        rec_for_motion,
-                        folder=sorter_output_folder / "motion",
-                        output_motion_info=True,
-                        preset=params["motion_correction_preset"],
-                        # **params["motion_correction"],
-                    )
-                    if verbose:
-                        print("Done correct_motion()")
-
-            # recording = bandpass_filter(recording_raw, **params["filtering"], margin_ms=20.0, dtype="float32")
-            recording = bandpass_filter(
-                recording_raw,
-                freq_min=params["freq_min"],
-                freq_max=params["freq_max"],
-                ftype="bessel",
-                filter_order=2,
-                margin_ms=20.0,
-                dtype="float32",
-            )
+                else:
+                    rec_for_motion = apply_preprocessing_pipeline(rec_for_motion, params["preprocessing_dict"])
+
+                if verbose:
+                    print("Start correct_motion()")
+                _, motion_info = correct_motion(
+                    rec_for_motion,
+                    folder=sorter_output_folder / "motion",
+                    output_motion_info=True,
+                    preset=params["motion_correction_preset"],
+                )
+                if verbose:
+                    print("Done correct_motion()")
+
+            if params["preprocessing_dict"] is None:
+                recording = bandpass_filter(
+                    recording_raw,
+                    freq_min=params["freq_min"],
+                    freq_max=params["freq_max"],
+                    ftype="bessel",
+                    filter_order=2,
+                    dtype="float32",
+                )
 
-            if apply_cmr:
-                recording = common_reference(recording)
+                if apply_cmr:
+                    recording = common_reference(recording)
+            else:
+                recording = apply_preprocessing_pipeline(recording, params["preprocessing_dict"])
 
             if params["apply_motion_correction"]:
                 interpolate_motion_kwargs = dict(
@@ -183,7 +190,7 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
 
             noise_levels = np.ones(num_chans, dtype="float32")
         else:
-            recording = recording_raw
+            recording = recording_raw.astype("float32")
             noise_levels = get_noise_levels(recording, return_in_uV=False)
             cache_info = None
 
@@ -332,6 +339,8 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
                 recording_for_peeler,
                 sorting,
                 templates,
+                amplitude_scalings=spikes["amplitude"],
+                noise_levels=noise_levels,
                 similarity_kwargs={"method": "l1", "support": "union", "max_lag_ms": 0.1},
                 sparsity_overlap=0.5,
                 censor_ms=3.0,

src/spikeinterface/sortingcomponents/tools.py

@@ -12,12 +12,11 @@
 
 from spikeinterface.core.sparsity import ChannelSparsity
 from spikeinterface.core.waveform_tools import extract_waveforms_to_single_buffer
-from spikeinterface.core.job_tools import split_job_kwargs, fix_job_kwargs
+from spikeinterface.core.job_tools import fix_job_kwargs
 from spikeinterface.core.sortinganalyzer import create_sorting_analyzer
 from spikeinterface.core.sparsity import ChannelSparsity
 from spikeinterface.core.sparsity import compute_sparsity
-from spikeinterface.core.analyzer_extension_core import ComputeTemplates, ComputeNoiseLevels
-from spikeinterface.core.template_tools import get_template_extremum_channel_peak_shift, get_template_extremum_channel
+from spikeinterface.core.template_tools import get_template_extremum_channel_peak_shift
 from spikeinterface.core.recording_tools import get_noise_levels
 from spikeinterface.core.sorting_tools import get_numba_vector_to_list_of_spiketrain
 
@@ -480,7 +479,7 @@ def remove_empty_templates(templates):
 
 
 def create_sorting_analyzer_with_existing_templates(
-    sorting, recording, templates, remove_empty=True, noise_levels=None
+    sorting, recording, templates, remove_empty=True, noise_levels=None, amplitude_scalings=None
 ):
     sparsity = templates.sparsity
     templates_array = templates.get_dense_templates().copy()
@@ -495,6 +494,8 @@ def create_sorting_analyzer_with_existing_templates(
     else:
         non_empty_sorting = sorting
 
+    from spikeinterface.core.analyzer_extension_core import ComputeTemplates
+
     sa = create_sorting_analyzer(non_empty_sorting, recording, format="memory", sparsity=sparsity)
     sa.compute("random_spikes")
     sa.extensions["templates"] = ComputeTemplates(sa)
@@ -509,12 +510,30 @@ def create_sorting_analyzer_with_existing_templates(
     sa.extensions["templates"].run_info["runtime_s"] = 0
 
     if noise_levels is not None:
+        from spikeinterface.core.analyzer_extension_core import ComputeNoiseLevels
+
         sa.extensions["noise_levels"] = ComputeNoiseLevels(sa)
         sa.extensions["noise_levels"].params = {}
         sa.extensions["noise_levels"].data["noise_levels"] = noise_levels
         sa.extensions["noise_levels"].run_info["run_completed"] = True
         sa.extensions["noise_levels"].run_info["runtime_s"] = 0
 
+    if amplitude_scalings is not None:
+        from spikeinterface.postprocessing.amplitude_scalings import ComputeAmplitudeScalings
+
+        sa.extensions["amplitude_scalings"] = ComputeAmplitudeScalings(sa)
+        sa.extensions["amplitude_scalings"].params = dict(
+            sparsity=None,
+            max_dense_channels=16,
+            ms_before=templates.ms_before,
+            ms_after=templates.ms_after,
+            handle_collisions=False,
+            delta_collision_ms=2,
+        )
+        sa.extensions["amplitude_scalings"].data["amplitude_scalings"] = amplitude_scalings
+        sa.extensions["amplitude_scalings"].run_info["run_completed"] = True
+        sa.extensions["amplitude_scalings"].run_info["runtime_s"] = 0
+
     return sa