spikeinterface 0.102.0__tar.gz → 0.102.2__tar.gz

{spikeinterface-0.102.0/src/spikeinterface.egg-info → spikeinterface-0.102.2}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: spikeinterface
-Version: 0.102.0
+Version: 0.102.2
 Summary: Python toolkit for analysis, visualization, and comparison of spike sorting output
 Author-email: Alessio Buccino <alessiop.buccino@gmail.com>, Samuel Garcia <sam.garcia.die@gmail.com>
 Project-URL: homepage, https://github.com/SpikeInterface/spikeinterface
@@ -14,10 +14,11 @@ Classifier: Operating System :: POSIX :: Linux
 Classifier: Operating System :: Microsoft :: Windows
 Classifier: Operating System :: MacOS
 Classifier: Operating System :: OS Independent
-Requires-Python: <3.13,>=3.9
+Requires-Python: <3.14,>=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: numpy>=1.20
+Requires-Dist: numpy>=1.20; python_version < "3.13"
+Requires-Dist: numpy>=2.0.0; python_version >= "3.13"
 Requires-Dist: threadpoolctl>=3.0.0
 Requires-Dist: tqdm
 Requires-Dist: zarr<3,>=2.18
@@ -32,13 +33,11 @@ Requires-Dist: pyedflib<0.1.39,>=0.1.30; extra == "extractors"
 Requires-Dist: sonpy; python_version < "3.10" and extra == "extractors"
 Requires-Dist: lxml; extra == "extractors"
 Requires-Dist: scipy; extra == "extractors"
-Requires-Dist: ONE-api>=2.7.0; extra == "extractors"
-Requires-Dist: ibllib>=2.36.0; extra == "extractors"
+Requires-Dist: ibllib; extra == "extractors"
 Requires-Dist: pymatreader>=0.0.32; extra == "extractors"
 Requires-Dist: zugbruecke>=0.2; sys_platform != "win32" and extra == "extractors"
 Provides-Extra: streaming-extractors
-Requires-Dist: ONE-api<2.10.0,>=2.7.0; extra == "streaming-extractors"
-Requires-Dist: ibllib>=2.36.0; extra == "streaming-extractors"
+Requires-Dist: ibllib; extra == "streaming-extractors"
 Requires-Dist: pynwb>=2.6.0; extra == "streaming-extractors"
 Requires-Dist: fsspec; extra == "streaming-extractors"
 Requires-Dist: aiohttp; extra == "streaming-extractors"
@@ -78,22 +77,21 @@ Provides-Extra: test-extractors
 Requires-Dist: pooch>=1.8.2; extra == "test-extractors"
 Requires-Dist: datalad>=1.0.2; extra == "test-extractors"
 Provides-Extra: test-preprocessing
-Requires-Dist: ibllib>=2.36.0; extra == "test-preprocessing"
+Requires-Dist: ibllib; extra == "test-preprocessing"
 Requires-Dist: torch; extra == "test-preprocessing"
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-dependency; extra == "test"
 Requires-Dist: pytest-cov; extra == "test"
+Requires-Dist: psutil; extra == "test"
 Requires-Dist: huggingface_hub; extra == "test"
-Requires-Dist: ibllib>=2.36.0; extra == "test"
+Requires-Dist: ibllib; extra == "test"
 Requires-Dist: s3fs; extra == "test"
-Requires-Dist: numba; extra == "test"
+Requires-Dist: numba<0.61.0; python_version < "3.13" and extra == "test"
+Requires-Dist: numba>=0.61.0; python_version >= "3.13" and extra == "test"
 Requires-Dist: hdbscan>=0.8.33; extra == "test"
-Requires-Dist: sortingview; extra == "test"
-Requires-Dist: tridesclous>=1.6.8; extra == "test"
-Requires-Dist: pymde; extra == "test"
+Requires-Dist: sortingview>=0.12.0; extra == "test"
 Requires-Dist: torch; extra == "test"
-Requires-Dist: pynndescent; extra == "test"
 Requires-Dist: skops; extra == "test"
 Requires-Dist: huggingface_hub; extra == "test"
 Provides-Extra: docs
@@ -119,6 +117,7 @@ Requires-Dist: spikeinterface[test_core]; extra == "dev"
 Requires-Dist: spikeinterface[docs]; extra == "dev"
 Requires-Dist: black; extra == "dev"
 Requires-Dist: pre-commit; extra == "dev"
+Dynamic: license-file
 
 # SpikeInterface: a unified framework for spike sorting
 

{spikeinterface-0.102.0 → spikeinterface-0.102.2}/pyproject.toml

@@ -1,13 +1,13 @@
 [project]
 name = "spikeinterface"
-version = "0.102.0"
+version = "0.102.2"
 authors = [
   { name="Alessio Buccino", email="alessiop.buccino@gmail.com" },
   { name="Samuel Garcia", email="sam.garcia.die@gmail.com" },
 ]
 description = "Python toolkit for analysis, visualization, and comparison of spike sorting output"
 readme = "README.md"
-requires-python = ">=3.9,<3.13"
+requires-python = ">=3.9,<3.14"
 classifiers = [
     "Programming Language :: Python :: 3 :: Only",
     "License :: OSI Approved :: MIT License",
@@ -20,7 +20,8 @@ classifiers = [
 
 
 dependencies = [
-    "numpy>=1.20",
+    "numpy>=1.20;python_version<'3.13'",
+    "numpy>=2.0.0;python_version>='3.13'",
     "threadpoolctl>=3.0.0",
     "tqdm",
     "zarr>=2.18,<3",
@@ -30,7 +31,7 @@ dependencies = [
 ]
 
 [build-system]
-requires = ["setuptools>=62.0"]
+requires = ["setuptools>=78.0.2"]
 build-backend = "setuptools.build_meta"
 
 [tool.setuptools]
@@ -66,15 +67,13 @@ extractors = [
     "sonpy;python_version<'3.10'",
     "lxml", # lxml for neuroscope
     "scipy",
-    "ONE-api>=2.7.0", # alf sorter and streaming IBL
-    "ibllib>=2.36.0", # streaming IBL
+    "ibllib", # streaming IBL
     "pymatreader>=0.0.32", # For cell explorer matlab files
     "zugbruecke>=0.2; sys_platform!='win32'", # For plexon2
 ]
 
 streaming_extractors = [
-    "ONE-api>=2.7.0,<2.10.0", # alf sorter and streaming IBL
-    "ibllib>=2.36.0", # streaming IBL
+    "ibllib", # streaming IBL
     # Following dependencies are for streaming with nwb files
     "pynwb>=2.6.0",
     "fsspec",
@@ -134,12 +133,13 @@ test_extractors = [
     # Functions to download data in neo test suite
     "pooch>=1.8.2",
     "datalad>=1.0.2",
+    # Commenting out for release
     # "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git",
     # "neo @ git+https://github.com/NeuralEnsemble/python-neo.git",
 ]
 
 test_preprocessing = [
-    "ibllib>=2.36.0", # for IBL
+    "ibllib", # for IBL
     "torch",
 ]
 
@@ -148,30 +148,26 @@ test = [
     "pytest",
     "pytest-dependency",
     "pytest-cov",
+    "psutil",
 
     "huggingface_hub",
 
     # preprocessing
-    "ibllib>=2.36.0", # for IBL
+    "ibllib", # for IBL
 
     # streaming templates
     "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
-    "sortingview",
-
-
-    ## install tridesclous for testing ##
-    "tridesclous>=1.6.8",
+    "sortingview>=0.12.0",
 
-    ## sliding_nn
-    "pymde",
+    # for motion and sortingcomponents
     "torch",
-    "pynndescent",
 
     # curation
     "skops",
@@ -207,7 +203,6 @@ docs = [
     # for release we need pypi, so this needs to be commented
     # "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git",  # We always build from the latest version
     # "neo @ git+https://github.com/NeuralEnsemble/python-neo.git",  # We always build from the latest version
-
 ]
 
 dev = [

{spikeinterface-0.102.0 → spikeinterface-0.102.2}/src/spikeinterface/benchmark/benchmark_base.py

@@ -10,7 +10,7 @@ import time
 
 
 from spikeinterface.core import SortingAnalyzer
-
+from spikeinterface.core.job_tools import fix_job_kwargs, split_job_kwargs
 from spikeinterface import load, create_sorting_analyzer, load_sorting_analyzer
 from spikeinterface.widgets import get_some_colors
 
@@ -43,11 +43,37 @@ class BenchmarkStudy:
         self.analyzers = {}
         self.cases = {}
         self.benchmarks = {}
+        self.levels = None
+        self.colors_by_case = None
+        self.colors_by_levels = {}
         self.scan_folder()
-        self.colors = None
 
     @classmethod
     def create(cls, study_folder, datasets={}, cases={}, levels=None):
+        """
+        Create a BenchmarkStudy from a dict of datasets and cases.
+
+        Parameters
+        ----------
+        study_folder : str | Path
+            The folder where the study will be saved.
+        datasets : dict
+            A dict of datasets. The keys are the dataset names and the values are `SortingAnalyzer` objects.
+            Values can also be tuples with (recording, gt_sorting), but this is deprecated.
+        cases : dict
+            A dict of cases. The keys are the cases (str, or tuples) and the values are dictionaries containing:
+
+                * dataset
+                * label
+                * params
+        levels : list | None
+            If the keys of the cases are tuples, this is the list of levels names.
+
+        Returns
+        -------
+        study : BenchmarkStudy
+            The created study.
+        """
         # check that cases keys are homogeneous
         key0 = list(cases.keys())[0]
         if isinstance(key0, str):
@@ -132,6 +158,9 @@ class BenchmarkStudy:
         return cls(study_folder)
 
     def create_benchmark(self, key):
+        """
+        Create a benchmark for a given key.
+        """
         raise NotImplementedError
 
     def scan_folder(self):
@@ -167,7 +196,7 @@ class BenchmarkStudy:
             result_folder = self.folder / "results" / self.key_to_str(key)
             if result_folder.exists():
                 result = self.benchmark_class.load_folder(result_folder)
-                benchmark = self.create_benchmark(key)
+                benchmark = self.create_benchmark(key=key)
                 benchmark.result.update(result)
                 self.benchmarks[key] = benchmark
             else:
@@ -219,7 +248,7 @@ class BenchmarkStudy:
         for key in job_keys:
             benchmark = self.create_benchmark(key)
             t0 = time.perf_counter()
-            benchmark.run()
+            benchmark.run(**job_kwargs)
             t1 = time.perf_counter()
             self.benchmarks[key] = benchmark
             bench_folder = self.folder / "results" / self.key_to_str(key)
@@ -228,19 +257,49 @@ class BenchmarkStudy:
             benchmark.result["run_time"] = float(t1 - t0)
             benchmark.save_main(bench_folder)
 
-    def set_colors(self, colors=None, map_name="tab20"):
+    def set_colors(self, colors=None, map_name="tab10", levels_to_group_by=None):
+        """
+        Set colors for the study cases or for a given levels_to_group_by.
+
+        Parmeters
+        ---------
+        colors : dict | None, default: None
+            A user-defined dictionary with the case keys as keys and the colors as values.
+            Note that the case keys depend on the levels_to_group_by.
+        map_name : str, default: 'tab10'
+            The name of the colormap to use.
+        levels_to_group_by : list | None, default: None
+            The levels to group by. If None, the colors are set for the cases.
+        """
+        case_keys, _ = self.get_grouped_keys_mapping(levels_to_group_by)
+
         if colors is None:
-            case_keys = list(self.cases.keys())
-            self.colors = get_some_colors(
-                case_keys, map_name=map_name, color_engine="matplotlib", shuffle=False, margin=0
+            colors = get_some_colors(
+                case_keys, map_name=map_name, color_engine="matplotlib", shuffle=False, margin=0, resample=False
             )
+            if levels_to_group_by is None:
+                self.colors_by_case = colors
+            else:
+                level_key = tuple(levels_to_group_by) if len(levels_to_group_by) > 1 else levels_to_group_by[0]
+                self.colors_by_levels[level_key] = colors
         else:
-            self.colors = colors
-
-    def get_colors(self):
-        if self.colors is None:
-            self.set_colors()
-        return self.colors
+            assert all([key in colors for key in case_keys]), f"You must provide colors for all cases keys: {case_keys}"
+            if levels_to_group_by is None:
+                self.colors_by_case = colors
+            else:
+                level_key = tuple(levels_to_group_by) if len(levels_to_group_by) > 1 else levels_to_group_by[0]
+                self.colors_by_levels[level_key] = colors
+
+    def get_colors(self, levels_to_group_by=None):
+        if levels_to_group_by is None:
+            if self.colors_by_case is None:
+                self.set_colors()
+            return self.colors_by_case
+        else:
+            level_key = tuple(levels_to_group_by) if len(levels_to_group_by) > 1 else levels_to_group_by[0]
+            if level_key not in self.colors_by_levels:
+                self.set_colors(levels_to_group_by=levels_to_group_by)
+            return self.colors_by_levels[level_key]
 
     def get_run_times(self, case_keys=None):
         if case_keys is None:
@@ -258,12 +317,60 @@ class BenchmarkStudy:
             df.index.names = self.levels
         return df
 
-    def plot_run_times(self, case_keys=None):
+    def get_grouped_keys_mapping(self, levels_to_group_by=None):
+        """
+        Return a dictionary of grouped keys.
+
+        Parameters
+        ----------
+        levels_to_group_by : list
+            A list of levels to group by.
+
+        Returns
+        -------
+        grouped_keys : dict
+            A dictionary of grouped keys, with the new keys as keys and the list of cases
+            associated to new keys as values.
+        labels : dict
+            A dictionary of labels, with the new keys as keys and the labels as values.
+        """
+        cases = list(self.cases.keys())
+        if levels_to_group_by is None or self.levels is None:
+            keys_mapping = {key: [key] for key in cases}
+        elif len(self.levels) == 1:
+            keys_mapping = {key: [key] for key in cases}
+        else:
+            study_levels = self.levels
+            assert np.all(
+                [l in study_levels for l in levels_to_group_by]
+            ), f"levels_to_group_by must be in {study_levels}, got {levels_to_group_by}"
+            keys_mapping = {}
+            for key in cases:
+                new_key = tuple(key[list(study_levels).index(level)] for level in levels_to_group_by)
+                if len(new_key) == 1:
+                    new_key = new_key[0]
+                if new_key not in keys_mapping:
+                    keys_mapping[new_key] = []
+                keys_mapping[new_key].append(key)
+
+        if levels_to_group_by is None:
+            labels = {key: self.cases[key]["label"] for key in cases}
+        else:
+            key0 = list(keys_mapping.keys())[0]
+            if isinstance(key0, tuple):
+                labels = {key: "-".join(key) for key in keys_mapping}
+            else:
+                labels = {key: key for key in keys_mapping}
+
+        return keys_mapping, labels
+
+    def plot_run_times(self, case_keys=None, **kwargs):
         from .benchmark_plot_tools import plot_run_times
 
-        return plot_run_times(self, case_keys=case_keys)
+        return plot_run_times(self, case_keys=case_keys, **kwargs)
 
     def compute_results(self, case_keys=None, verbose=False, **result_params):
+
         if case_keys is None:
             case_keys = list(self.cases.keys())
 
@@ -304,12 +411,32 @@ class BenchmarkStudy:
         # sorting_analyzer = load_sorting_analyzer(folder)
         # return sorting_analyzer
 
+    def compute_analyzer_extension(self, extensions, dataset_keys=None, **extension_kwargs):
+        if dataset_keys is None:
+            dataset_keys = list(self.datasets.keys())
+        if not isinstance(dataset_keys, list):
+            dataset_keys = [dataset_keys]
+        for dataset_key in dataset_keys:
+            sorting_analyzer = self.get_sorting_analyzer(dataset_key=dataset_key)
+            sorting_analyzer.compute(extensions, **extension_kwargs)
+
+    def get_gt_unit_locations(self, case_key):
+        dataset_key = self.cases[case_key]["dataset"]
+        sorting_analyzer = self.get_sorting_analyzer(dataset_key=dataset_key)
+        if "gt_unit_locations" in sorting_analyzer.sorting.get_property_keys():
+            return sorting_analyzer.get_sorting_property("gt_unit_locations")
+        else:
+            if not sorting_analyzer.has_extension("unit_locations"):
+                self.compute_analyzer_extension(["unit_locations"], dataset_keys=dataset_key)
+            unit_locations_ext = sorting_analyzer.get_extension("unit_locations")
+            return unit_locations_ext.get_data()
+
     def get_templates(self, key, operator="average"):
         sorting_analyzer = self.get_sorting_analyzer(case_key=key)
         templates = sorting_analyzer.get_extenson("templates").get_data(operator=operator)
         return templates
 
-    def compute_metrics(self, case_keys=None, metric_names=["snr", "firing_rate"], force=False):
+    def compute_metrics(self, case_keys=None, metric_names=["snr", "firing_rate"], force=False, **job_kwargs):
         if case_keys is None:
             case_keys = self.cases.keys()
 
@@ -329,17 +456,13 @@ class BenchmarkStudy:
             sorting_analyzer = self.get_sorting_analyzer(key)
             qm_ext = sorting_analyzer.get_extension("quality_metrics")
             if qm_ext is None or force:
-                qm_ext = sorting_analyzer.compute("quality_metrics", metric_names=metric_names)
+                qm_ext = sorting_analyzer.compute("quality_metrics", metric_names=metric_names, **job_kwargs)
 
             # TODO remove this metics CSV file!!!!
             metrics = qm_ext.get_data()
             # metrics.to_csv(filename, sep="\t", index=True)
 
     def get_metrics(self, key):
-        import pandas as pd
-
-        dataset_key = self.cases[key]["dataset"]
-
         analyzer = self.get_sorting_analyzer(key)
         ext = analyzer.get_extension("quality_metrics")
         if ext is None:
@@ -347,16 +470,32 @@ class BenchmarkStudy:
             return None
 
         metrics = ext.get_data()
+        # add GT unit id column
+        gt_unit_ids = analyzer.sorting.unit_ids
+        metrics.loc[:, "gt_unit_id"] = gt_unit_ids
         return metrics
 
-        # filename = self.folder / "metrics" / f"{self.key_to_str(dataset_key)}.csv"
-        # if not filename.exists():
-        #     return
-        # metrics = pd.read_csv(filename, sep="\t", index_col=0)
-        # dataset_key = self.cases[key]["dataset"]
-        # recording, gt_sorting = self.datasets[dataset_key]
-        # metrics.index = gt_sorting.unit_ids
-        # return metrics
+    def get_all_metrics(self, case_keys=None):
+        """
+        Return a DataFrame with concatented metrics for multiple cases.
+        """
+        import pandas as pd
+
+        if case_keys is None:
+            case_keys = list(self.cases.keys())
+        assert all(key in self.cases for key in case_keys), "Some case keys are not in cases"
+        metrics = []
+        indices = []
+        for key in case_keys:
+            metrics.append(self.get_metrics(key))
+            indices.extend([key] * len(metrics[-1]))
+        if isinstance(case_keys[0], str):
+            index = pd.Index(indices, name=self.levels)
+        else:
+            index = pd.MultiIndex.from_tuples(indices, names=self.levels)
+        metrics = pd.concat(metrics)
+        metrics.index = index
+        return metrics
 
     def get_units_snr(self, key):
         """ """
@@ -365,6 +504,39 @@ class BenchmarkStudy:
     def get_result(self, key):
         return self.benchmarks[key].result
 
+    def get_pairs_by_level(self, level):
+        """
+        usefull for function like plot_performance_losses() where you need to plot one pair of results
+        This generate list of pairs for a given level.
+        """
+
+        level_index = self.levels.index(level)
+
+        possible_values = []
+        for key in self.cases.keys():
+            assert isinstance(key, tuple), "get_pairs_by_level need tuple keys"
+            level_value = key[level_index]
+            if level_value not in possible_values:
+                possible_values.append(level_value)
+        assert len(possible_values) == 2, "get_pairs_by_level() : you need exactly 2 value for this levels"
+
+        pairs = []
+        for key in self.cases.keys():
+
+            case0 = list(key)
+            case1 = list(key)
+            case0[level_index] = possible_values[0]
+            case1[level_index] = possible_values[1]
+            case0 = tuple(case0)
+            case1 = tuple(case1)
+
+            pair = (case0, case1)
+
+            if pair not in pairs:
+                pairs.append(pair)
+
+        return pairs
+
 
 class Benchmark:
     """