Merge pull request #3697 from alejoe91/ks4-dynamic-params

Add BaseSorter._dynamic_params to be to retrieve parameters dynamically

Author: alejoe91

.github/scripts/test_kilosort4_ci.py

@@ -20,7 +20,7 @@
 
 import pytest
 import copy
-from typing import Any
+from packaging.version import parse
 from inspect import signature
 
 import numpy as np
@@ -84,8 +84,6 @@
     "duplicate_spike_ms": 0.3,
 }
 
-PARAMS_TO_TEST = list(PARAMS_TO_TEST_DICT.keys())
-
 PARAMETERS_NOT_AFFECTING_RESULTS = [
     "artifact_threshold",
     "ccg_threshold",
@@ -95,13 +93,24 @@
     "duplicate_spike_ms",  # this is because ground-truth spikes don't have violations
 ]
 
-# THIS IS A PLACEHOLDER FOR FUTURE PARAMS TO TEST
-# if parse(version("kilosort")) >= parse("4.0.X"):
-#     PARAMS_TO_TEST_DICT.update(
-#         [
-#             {"new_param": new_value},
-#         ]
-#     )
+
+# Add/Remove version specific parameters
+if parse(kilosort.__version__) >= parse("4.0.22"):
+    PARAMS_TO_TEST_DICT.update(
+        {"position_limit": 50}
+    )
+    # Position limit only affects computing spike locations after sorting
+    PARAMETERS_NOT_AFFECTING_RESULTS.append("position_limit")
+
+if parse(kilosort.__version__) >= parse("4.0.24"):
+    PARAMS_TO_TEST_DICT.update(
+        {"max_peels": 200},
+    )
+    # max_peels is not affecting the results in this short dataset
+    PARAMETERS_NOT_AFFECTING_RESULTS.append("max_peels")
+
+
+PARAMS_TO_TEST = list(PARAMS_TO_TEST_DICT.keys())
 
 
 class TestKilosort4Long:

pyproject.toml

@@ -159,7 +159,8 @@ test = [
     "s3fs",
 
     # tridesclous
-    "numba",
+    "numba<0.61.0;python_version<'3.13'",
+    "numba>=0.61.0;python_version>='3.13'",
     "hdbscan>=0.8.33",  # Previous version had a broken wheel
 
     # for sortingview backend

src/spikeinterface/preprocessing/tests/test_whiten.py

@@ -56,7 +56,7 @@ def get_test_recording(self, dtype, means=None):
 
         return means, cov_mat, recording
 
-    def get_test_data_with_known_distribution(self, num_samples, dtype, means=None):
+    def get_test_data_with_known_distribution(self, num_samples, dtype, means=None, seed=0):
         """
         Create multivariate normal data with known means and covariance matrixs.
         If `dtype` is int16, scale to full range of int16 before cast.
@@ -68,7 +68,8 @@ def get_test_data_with_known_distribution(self, num_samples, dtype, means=None):
 
         cov_mat = np.array([[1, 0.5, 0], [0.5, 1, -0.25], [0, -0.25, 1]])
 
-        data = np.random.multivariate_normal(means, cov_mat, num_samples)
+        rng = np.random.RandomState(seed)
+        data = rng.multivariate_normal(means, cov_mat, num_samples)
 
         # Set the dtype, if `int16`, first scale to +/- 1 then cast to int16 range.
         if dtype == np.float32:

src/spikeinterface/sorters/basesorter.py

@@ -152,17 +152,25 @@ def initialize_folder(cls, recording, output_folder, verbose, remove_existing_fo
 
         return output_folder
 
+    @classmethod
+    def _dynamic_params(cls):
+        # optional
+        # can be implemented in subclass for dynamic default parameters
+        return cls._default_params, cls._params_description
+
     @classmethod
     def default_params(cls):
-        p = copy.deepcopy(cls._default_params)
+        default_params, _ = cls._dynamic_params()
+        p = copy.deepcopy(default_params)
         if cls.requires_binary_data:
             job_kwargs = get_global_job_kwargs()
             p.update(job_kwargs)
         return p
 
     @classmethod
     def params_description(cls):
-        p = copy.deepcopy(cls._params_description)
+        _, default_params_description = cls._dynamic_params()
+        p = copy.deepcopy(default_params_description)
         if cls.requires_binary_data:
             p.update(default_job_kwargs_description)
         return p

src/spikeinterface/sorters/external/kilosort4.py

@@ -2,6 +2,7 @@
 
 from pathlib import Path
 from typing import Union
+import warnings
 from packaging import version
 
 
@@ -22,95 +23,29 @@ class Kilosort4Sorter(BaseSorter):
     gpu_capability = "nvidia-optional"
     requires_binary_data = False
 
-    _default_params = {
-        "batch_size": 60000,
-        "nblocks": 1,
-        "Th_universal": 9,
-        "Th_learned": 8,
+    _si_default_params = {
         "do_CAR": True,
         "invert_sign": False,
-        "nt": 61,
-        "shift": None,
-        "scale": None,
-        "artifact_threshold": None,
-        "nskip": 25,
-        "whitening_range": 32,
-        "highpass_cutoff": 300,
-        "binning_depth": 5,
-        "sig_interp": 20,
-        "drift_smoothing": [0.5, 0.5, 0.5],
-        "nt0min": None,
-        "dmin": None,
-        "dminx": 32,
-        "min_template_size": 10,
-        "template_sizes": 5,
-        "nearest_chans": 10,
-        "nearest_templates": 100,
-        "max_channel_distance": None,
-        "templates_from_data": True,
-        "n_templates": 6,
-        "n_pcs": 6,
-        "Th_single_ch": 6,
-        "acg_threshold": 0.2,
-        "ccg_threshold": 0.25,
-        "cluster_downsampling": 20,
-        "cluster_pcs": 64,
-        "x_centers": None,
-        "duplicate_spike_ms": 0.25,
-        "scaleproc": None,
+        "save_extra_vars": False,
         "save_preprocessed_copy": False,
         "torch_device": "auto",
         "bad_channels": None,
         "clear_cache": False,
-        "save_extra_vars": False,
         "do_correction": True,
-        "keep_good_only": False,
         "skip_kilosort_preprocessing": False,
+        "keep_good_only": False,
         "use_binary_file": True,
         "delete_recording_dat": True,
     }
 
-    _params_description = {
-        "batch_size": "Number of samples included in each batch of data.",
-        "nblocks": "Number of non-overlapping blocks for drift correction (additional nblocks-1 blocks are created in the overlaps). Default value: 1.",
-        "Th_universal": "Spike detection threshold for universal templates. Th(1) in previous versions of Kilosort. Default value: 9.",
-        "Th_learned": "Spike detection threshold for learned templates. Th(2) in previous versions of Kilosort. Default value: 8.",
-        "do_CAR": "Whether to perform common average reference. Default value: True.",
-        "invert_sign": "Invert the sign of the data. Default value: False.",
-        "nt": "Number of samples per waveform. Also size of symmetric padding for filtering. Default value: 61.",
-        "shift": "Scalar shift to apply to data before all other operations. Default None.",
-        "scale": "Scaling factor to apply to data before all other operations. Default None.",
-        "artifact_threshold": "If a batch contains absolute values above this number, it will be zeroed out under the assumption that a recording artifact is present. By default, the threshold is infinite (so that no zeroing occurs). Default value: None.",
-        "nskip": "Batch stride for computing whitening matrix. Default value: 25.",
-        "whitening_range": "Number of nearby channels used to estimate the whitening matrix. Default value: 32.",
-        "highpass_cutoff": "High-pass filter cutoff frequency in Hz. Default value: 300.",
-        "binning_depth": "For drift correction, vertical bin size in microns used for 2D histogram. Default value: 5.",
-        "sig_interp": "For drift correction, sigma for interpolation (spatial standard deviation). Approximate smoothness scale in units of microns. Default value: 20.",
-        "drift_smoothing": "Amount of gaussian smoothing to apply to the spatiotemporal drift estimation, for x,y,time axes in units of registration blocks (for x,y axes) and batch size (for time axis). The x,y smoothing has no effect for `nblocks = 1`.",
-        "nt0min": "Sample index for aligning waveforms, so that their minimum or maximum value happens here. Default of 20. Default value: None.",
-        "dmin": "Vertical spacing of template centers used for spike detection, in microns. Determined automatically by default. Default value: None.",
-        "dminx": "Horizontal spacing of template centers used for spike detection, in microns. Default value: 32.",
-        "min_template_size": "Standard deviation of the smallest, spatial envelope Gaussian used for universal templates. Default value: 10.",
-        "template_sizes": "Number of sizes for universal spike templates (multiples of the min_template_size). Default value: 5.",
-        "nearest_chans": "Number of nearest channels to consider when finding local maxima during spike detection. Default value: 10.",
-        "nearest_templates": "Number of nearest spike template locations to consider when finding local maxima during spike detection. Default value: 100.",
-        "max_channel_distance": "Templates farther away than this from their nearest channel will not be used. Also limits distance between compared channels during clustering. Default value: None.",
-        "templates_from_data": "Indicates whether spike shapes used in universal templates should be estimated from the data or loaded from the predefined templates. Default value: True.",
-        "n_templates": "Number of single-channel templates to use for the universal templates (only used if templates_from_data is True). Default value: 6.",
-        "n_pcs": "Number of single-channel PCs to use for extracting spike features (only used if templates_from_data is True). Default value: 6.",
-        "Th_single_ch": "For single channel threshold crossings to compute universal- templates. In units of whitened data standard deviations. Default value: 6.",
-        "acg_threshold": 'Fraction of refractory period violations that are allowed in the ACG compared to baseline; used to assign "good" units. Default value: 0.2.',
-        "ccg_threshold": "Fraction of refractory period violations that are allowed in the CCG compared to baseline; used to perform splits and merges. Default value: 0.25.",
-        "cluster_downsampling": "Inverse fraction of nodes used as landmarks during clustering (can be 1, but that slows down the optimization). Default value: 20.",
-        "cluster_pcs": "Maximum number of spatiotemporal PC features used for clustering. Default value: 64.",
-        "x_centers": "Number of x-positions to use when determining center points for template groupings. If None, this will be determined automatically by finding peaks in channel density. For 2D array type probes, we recommend specifying this so that centers are placed every few hundred microns.",
-        "duplicate_spike_bins": "Number of bins for which subsequent spikes from the same cluster are assumed to be artifacts. A value of 0 disables this step. Default value: 7.",
-        "save_extra_vars": "If True, additional kwargs are saved to the output",
-        "scaleproc": "int16 scaling of whitened data, if None set to 200.",
-        "save_preprocessed_copy": "Save a pre-processed copy of the data (including drift correction) to temp_wh.dat in the results directory and format Phy output to use that copy of the data",
-        "torch_device": "Select the torch device auto/cuda/cpu",
-        "bad_channels": "A list of channel indices (rows in the binary file) that should not be included in sorting. Listing channels here is equivalent to excluding them from the probe dictionary.",
-        "clear_cache": "If True, force pytorch to free up memory reserved for its cache in between memory-intensive operations. Note that setting `clear_cache=True` is NOT recommended unless you encounter GPU out-of-memory errors, since this can result in slower sorting.",
+    _si_params_description = {
+        "do_CAR": "If True, common average reference is performed. Default is True. (run_kilosrt parameter)",
+        "invert_sign": "Invert the sign of the data. Default value: False. (run_kilosort parameter)",
+        "save_extra_vars": "If True, additional kwargs are saved to the output. Default is False. (run_kilosort parameter)",
+        "save_preprocessed_copy": "Save a pre-processed copy of the data (including drift correction) to temp_wh.dat in the results directory and format Phy output to use that copy of the data. (run_kilosort parameter)",
+        "torch_device": "Select the torch device auto/cuda/cpu. Default is 'auto'. (run_kilosort parameter)",
+        "bad_channels": "A list of channel indices (rows in the binary file) that should not be included in sorting. Listing channels here is equivalent to excluding them from the probe dictionary. (run_kilosort parameter)",
+        "clear_cache": "If True, force pytorch to free up memory reserved for its cache in between memory-intensive operations. Note that setting `clear_cache=True` is NOT recommended unless you encounter GPU out-of-memory errors, since this can result in slower sorting. (run_kilosort parameter)",
         "do_correction": "If True, drift correction is performed. Default is True. (spikeinterface parameter)",
         "skip_kilosort_preprocessing": "Can optionally skip the internal kilosort preprocessing. (spikeinterface parameter)",
         "keep_good_only": "If True, only the units labeled as 'good' by Kilosort are returned in the output. (spikeinterface parameter)",
@@ -151,6 +86,36 @@ def get_sorter_version(cls):
         """kilosort.__version__ <4.0.10 is always '4'"""
         return importlib_version("kilosort")
 
+    @classmethod
+    def _dynamic_params(cls):
+        if cls.is_installed():
+            import kilosort as ks
+
+            # 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"]
+            default_params = {}
+            default_params_descriptions = {}
+            ks_params = ks.parameters.MAIN_PARAMETERS.copy()
+            ks_params.update(ks.parameters.EXTRA_PARAMETERS)
+            for param, param_value in ks_params.items():
+                if param not in skip_main:
+                    default_params[param] = param_value["default"]
+                    desc = param_value.get("description")
+                    if desc is None:
+                        desc = ""
+                    else:
+                        # get rid of escape characters and extra spaces
+                        desc = " ".join(desc.replace("\n", "").split())
+                    default_params_descriptions[param] = desc
+            default_params.update(cls._si_default_params)
+            default_params_descriptions.update(cls._si_params_description)
+            return default_params, default_params_descriptions
+        else:
+            warnings.warn("Kilosort4 is not installed. Please install kilosort4 to get the parameters.")
+            return {}, {}
+
     @classmethod
     def initialize_folder(cls, recording, output_folder, verbose, remove_existing_folder):
         if not cls.is_installed():