spikeinterface 0.102.2__tar.gz → 0.103.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spikeinterface
-Version: 0.102.2
+Version: 0.103.0
 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
@@ -22,22 +22,24 @@ 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
-Requires-Dist: neo>=0.14.0
-Requires-Dist: probeinterface>=0.2.23
+Requires-Dist: neo>=0.14.1
+Requires-Dist: probeinterface>=0.3.0
 Requires-Dist: packaging
+Requires-Dist: pydantic
+Requires-Dist: numcodecs<0.16.0
 Provides-Extra: extractors
 Requires-Dist: MEArec>=1.8; extra == "extractors"
 Requires-Dist: pynwb>=2.6.0; extra == "extractors"
 Requires-Dist: hdmf-zarr>=0.11.0; extra == "extractors"
-Requires-Dist: pyedflib<0.1.39,>=0.1.30; extra == "extractors"
+Requires-Dist: pyedflib>=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: ibllib; extra == "extractors"
+Requires-Dist: ibllib>=3.4.1; python_version >= "3.10" and 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: ibllib; extra == "streaming-extractors"
+Requires-Dist: ibllib>=3.4.1; python_version >= "3.10" and extra == "streaming-extractors"
 Requires-Dist: pynwb>=2.6.0; extra == "streaming-extractors"
 Requires-Dist: fsspec; extra == "streaming-extractors"
 Requires-Dist: aiohttp; extra == "streaming-extractors"
@@ -56,7 +58,7 @@ Requires-Dist: networkx; extra == "full"
 Requires-Dist: distinctipy; extra == "full"
 Requires-Dist: matplotlib>=3.6; extra == "full"
 Requires-Dist: cuda-python; platform_system != "Darwin" and extra == "full"
-Requires-Dist: numba; extra == "full"
+Requires-Dist: numba>=0.59; extra == "full"
 Requires-Dist: skops; extra == "full"
 Requires-Dist: huggingface_hub; extra == "full"
 Provides-Extra: widgets
@@ -70,23 +72,24 @@ Requires-Dist: scipy; extra == "qualitymetrics"
 Requires-Dist: pandas; extra == "qualitymetrics"
 Requires-Dist: numba; extra == "qualitymetrics"
 Provides-Extra: test-core
-Requires-Dist: pytest; extra == "test-core"
+Requires-Dist: pytest<8.4.0; extra == "test-core"
 Requires-Dist: pytest-dependency; extra == "test-core"
 Requires-Dist: psutil; extra == "test-core"
+Requires-Dist: pytest-mock; extra == "test-core"
 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; extra == "test-preprocessing"
+Requires-Dist: ibllib>=3.4.1; python_version >= "3.10" and extra == "test-preprocessing"
 Requires-Dist: torch; extra == "test-preprocessing"
 Provides-Extra: test
-Requires-Dist: pytest; extra == "test"
+Requires-Dist: pytest<8.4.0; 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; extra == "test"
+Requires-Dist: ibllib>=3.4.1; python_version >= "3.10" and extra == "test"
 Requires-Dist: s3fs; extra == "test"
+Requires-Dist: pynapple; 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"
@@ -94,6 +97,7 @@ Requires-Dist: sortingview>=0.12.0; extra == "test"
 Requires-Dist: torch; extra == "test"
 Requires-Dist: skops; extra == "test"
 Requires-Dist: huggingface_hub; extra == "test"
+Requires-Dist: pytest-mock; extra == "test"
 Provides-Extra: docs
 Requires-Dist: Sphinx; extra == "docs"
 Requires-Dist: sphinx_rtd_theme>=1.2; extra == "docs"
@@ -109,8 +113,7 @@ Requires-Dist: numba; extra == "docs"
 Requires-Dist: networkx; extra == "docs"
 Requires-Dist: skops; extra == "docs"
 Requires-Dist: scikit-learn; extra == "docs"
-Requires-Dist: pooch>=1.8.2; extra == "docs"
-Requires-Dist: datalad>=1.0.2; extra == "docs"
+Requires-Dist: huggingface_hub; extra == "docs"
 Provides-Extra: dev
 Requires-Dist: spikeinterface[test]; extra == "dev"
 Requires-Dist: spikeinterface[test_core]; extra == "dev"
@@ -255,4 +258,4 @@ If you find SpikeInterface useful in your research, please cite:
 ```
 
 Please also cite other relevant papers for the specific components you use.
-For a ful list of references, please check the [references](https://spikeinterface.readthedocs.io/en/latest/references.html) page.
+For a full list of references, please check the [references](https://spikeinterface.readthedocs.io/en/latest/references.html) page.

{spikeinterface-0.102.2 → spikeinterface-0.103.0}/README.md

@@ -134,4 +134,4 @@ If you find SpikeInterface useful in your research, please cite:
 ```
 
 Please also cite other relevant papers for the specific components you use.
-For a ful list of references, please check the [references](https://spikeinterface.readthedocs.io/en/latest/references.html) page.
+For a full list of references, please check the [references](https://spikeinterface.readthedocs.io/en/latest/references.html) page.

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

@@ -1,6 +1,6 @@
 [project]
 name = "spikeinterface"
-version = "0.102.2"
+version = "0.103.0"
 authors = [
   { name="Alessio Buccino", email="alessiop.buccino@gmail.com" },
   { name="Samuel Garcia", email="sam.garcia.die@gmail.com" },
@@ -25,9 +25,11 @@ dependencies = [
     "threadpoolctl>=3.0.0",
     "tqdm",
     "zarr>=2.18,<3",
-    "neo>=0.14.0",
-    "probeinterface>=0.2.23",
+    "neo>=0.14.1",
+    "probeinterface>=0.3.0",
     "packaging",
+    "pydantic",
+    "numcodecs<0.16.0", # For supporting zarr < 3
 ]
 
 [build-system]
@@ -63,17 +65,17 @@ extractors = [
     "MEArec>=1.8",
     "pynwb>=2.6.0",
     "hdmf-zarr>=0.11.0",
-    "pyedflib>=0.1.30,<0.1.39",
+    "pyedflib>=0.1.30",
     "sonpy;python_version<'3.10'",
     "lxml", # lxml for neuroscope
     "scipy",
-    "ibllib", # streaming IBL
+    "ibllib>=3.4.1;python_version>='3.10'", # streaming IBL
     "pymatreader>=0.0.32", # For cell explorer matlab files
     "zugbruecke>=0.2; sys_platform!='win32'", # For plexon2
 ]
 
 streaming_extractors = [
-    "ibllib", # streaming IBL
+    "ibllib>=3.4.1;python_version>='3.10'", # streaming IBL
     # Following dependencies are for streaming with nwb files
     "pynwb>=2.6.0",
     "fsspec",
@@ -99,7 +101,7 @@ full = [
     "distinctipy",
     "matplotlib>=3.6", # matplotlib.colormaps
     "cuda-python; platform_system != 'Darwin'",
-    "numba",
+    "numba>=0.59",
     "skops",
     "huggingface_hub"
 ]
@@ -119,7 +121,7 @@ qualitymetrics = [
 ]
 
 test_core = [
-    "pytest",
+    "pytest<8.4.0",
     "pytest-dependency",
     "psutil",
 
@@ -127,6 +129,9 @@ test_core = [
     # for release we need pypi, so this need to be commented
     # "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git",
     # "neo @ git+https://github.com/NeuralEnsemble/python-neo.git",
+
+    # for slurm jobs,
+    "pytest-mock"
 ]
 
 test_extractors = [
@@ -139,30 +144,35 @@ test_extractors = [
 ]
 
 test_preprocessing = [
-    "ibllib", # for IBL
+    "ibllib>=3.4.1;python_version>='3.10'", # streaming IBL
     "torch",
 ]
 
 
 test = [
-    "pytest",
+    "pytest<8.4.0",
     "pytest-dependency",
     "pytest-cov",
     "psutil",
 
-    "huggingface_hub",
-
     # preprocessing
-    "ibllib", # for IBL
+    "ibllib>=3.4.1;python_version>='3.10'",
 
     # streaming templates
     "s3fs",
 
-    # tridesclous
+    # exporters
+    "pynapple",
+
+    # tridesclous2
     "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
 
+    # isosplit is needed for trideclous2 noramaly but isosplit is only build until python3.11
+    # so lets wait a new build of isosplit6
+    # "isosplit6",
+
     # for sortingview backend
     "sortingview>=0.12.0",
 
@@ -177,6 +187,9 @@ test = [
     # for release we need pypi, so this need to be commented
     # "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git",
     # "neo @ git+https://github.com/NeuralEnsemble/python-neo.git",
+
+    # for slurm jobs
+    "pytest-mock",
 ]
 
 docs = [
@@ -194,11 +207,9 @@ docs = [
     "hdbscan>=0.8.33",   # For sorters spykingcircus2 + tridesclous
     "numba", # For many postprocessing functions
     "networkx",
-    "skops", # For auotmated curation
-    "scikit-learn", # For auotmated curation
-    # Download data
-    "pooch>=1.8.2",
-    "datalad>=1.0.2",
+    "skops", # For automated curation
+    "scikit-learn", # For automated curation
+    "huggingface_hub", # For automated curation
 
     # 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

spikeinterface-0.103.0/src/spikeinterface/benchmark/__init__.py

@@ -0,0 +1,9 @@
+"""
+Module to benchmark:
+  * sorters with or without ground-truth
+  * some sorting components (clustering, motion, template matching)
+"""
+
+from .residual_analysis import analyse_residual, make_residual_recording
+from .benchmark_sorter import SorterStudy
+from .benchmark_sorter_without_gt import SorterStudyWithoutGroundTruth

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

@@ -9,7 +9,8 @@ import numpy as np
 import time
 
 
-from spikeinterface.core import SortingAnalyzer
+from spikeinterface.core import SortingAnalyzer, ChannelSparsity, NumpySorting
+from spikeinterface.core import SortingAnalyzer, NumpySorting
 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
@@ -46,6 +47,7 @@ class BenchmarkStudy:
         self.levels = None
         self.colors_by_case = None
         self.colors_by_levels = {}
+        self.labels_by_levels = {}
         self.scan_folder()
 
     @classmethod
@@ -118,12 +120,57 @@ class BenchmarkStudy:
             if isinstance(data, tuple):
                 # old case : rec + sorting
                 rec, gt_sorting = data
-                analyzer = create_sorting_analyzer(
-                    gt_sorting, rec, sparse=True, format="binary_folder", folder=local_analyzer_folder
-                )
-                analyzer.compute("random_spikes")
-                analyzer.compute("templates")
-                analyzer.compute("noise_levels")
+
+                if gt_sorting is not None:
+                    if "gt_unit_locations" in gt_sorting.get_property_keys():
+                        # if real units locations is present then use it for a better sparsity
+                        # then the real max channel is used
+                        radius_um = 100.0
+                        channel_ids = rec.channel_ids
+                        unit_ids = gt_sorting.unit_ids
+                        gt_unit_locations = gt_sorting.get_property("gt_unit_locations")
+                        channel_locations = rec.get_channel_locations()
+                        max_channel_indices = np.argmin(
+                            np.linalg.norm(
+                                gt_unit_locations[:, np.newaxis, :2] - channel_locations[np.newaxis, :], axis=2
+                            ),
+                            axis=1,
+                        )
+                        mask = np.zeros((unit_ids.size, channel_ids.size), dtype="bool")
+                        distances = np.linalg.norm(
+                            channel_locations[:, np.newaxis] - channel_locations[np.newaxis, :], axis=2
+                        )
+                        for unit_ind, unit_id in enumerate(unit_ids):
+                            chan_ind = max_channel_indices[unit_ind]
+                            (chan_inds,) = np.nonzero(distances[chan_ind, :] <= radius_um)
+                            mask[unit_ind, chan_inds] = True
+                        sparsity = ChannelSparsity(mask, unit_ids, channel_ids)
+                        sparse = False
+                    else:
+                        sparse = True
+                        sparsity = None
+
+                    analyzer = create_sorting_analyzer(
+                        gt_sorting,
+                        rec,
+                        sparse=sparse,
+                        sparsity=sparsity,
+                        format="binary_folder",
+                        folder=local_analyzer_folder,
+                    )
+                    analyzer.compute("random_spikes")
+                    analyzer.compute("templates")
+                    analyzer.compute("noise_levels")
+                else:
+                    # some study/benchmark has no GT sorting
+                    # in that case we still need an analyzer for internal API
+                    gt_sorting = NumpySorting.from_samples_and_labels(
+                        [np.array([])], [np.array([])], rec.sampling_frequency, unit_ids=None
+                    )
+                    analyzer = create_sorting_analyzer(
+                        gt_sorting, rec, sparse=False, format="binary_folder", folder=local_analyzer_folder
+                    )
+
             else:
                 # new case : analzyer
                 assert isinstance(data, SortingAnalyzer)
@@ -134,8 +181,6 @@ class BenchmarkStudy:
                 else:
                     analyzer = data
 
-                rec, gt_sorting = analyzer.recording, analyzer.sorting
-
             analyzers_path[key] = str(analyzer.folder.resolve())
 
             # recordings are pickled
@@ -176,17 +221,15 @@ class BenchmarkStudy:
         self.levels = self.info["levels"]
 
         for key, folder in self.analyzers_path.items():
-            analyzer = load_sorting_analyzer(folder)
+            analyzer = load_sorting_analyzer(folder, load_extensions=False)
             self.analyzers[key] = analyzer
             # the sorting is in memory here we take the saved one because comparisons need to pickle it later
             sorting = load(analyzer.folder / "sorting")
-            self.datasets[key] = analyzer.recording, sorting
-
-        # for rec_file in (self.folder / "datasets" / "recordings").glob("*.pickle"):
-        #     key = rec_file.stem
-        #     rec = load(rec_file)
-        #     gt_sorting = load(self.folder / f"datasets" / "gt_sortings" / key)
-        #     self.datasets[key] = (rec, gt_sorting)
+            if analyzer.has_recording():
+                recording = analyzer.recording
+            else:
+                recording = None
+            self.datasets[key] = recording, sorting
 
         with open(self.folder / "cases.pickle", "rb") as f:
             self.cases = pickle.load(f)
@@ -247,6 +290,8 @@ class BenchmarkStudy:
 
         for key in job_keys:
             benchmark = self.create_benchmark(key)
+            if verbose:
+                print("### Run benchmark", key, "###")
             t0 = time.perf_counter()
             benchmark.run(**job_kwargs)
             t1 = time.perf_counter()
@@ -317,7 +362,7 @@ class BenchmarkStudy:
             df.index.names = self.levels
         return df
 
-    def get_grouped_keys_mapping(self, levels_to_group_by=None):
+    def get_grouped_keys_mapping(self, levels_to_group_by=None, case_keys=None):
         """
         Return a dictionary of grouped keys.
 
@@ -325,6 +370,8 @@ class BenchmarkStudy:
         ----------
         levels_to_group_by : list
             A list of levels to group by.
+        case_keys : list
+            Optionaly a sub list of case_keys to consider
 
         Returns
         -------
@@ -334,18 +381,19 @@ class BenchmarkStudy:
         labels : dict
             A dictionary of labels, with the new keys as keys and the labels as values.
         """
-        cases = list(self.cases.keys())
+        if case_keys is None:
+            case_keys = list(self.cases.keys())
         if levels_to_group_by is None or self.levels is None:
-            keys_mapping = {key: [key] for key in cases}
+            keys_mapping = {key: [key] for key in case_keys}
         elif len(self.levels) == 1:
-            keys_mapping = {key: [key] for key in cases}
+            keys_mapping = {key: [key] for key in case_keys}
         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:
+            for key in case_keys:
                 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]
@@ -354,13 +402,17 @@ class BenchmarkStudy:
                 keys_mapping[new_key].append(key)
 
         if levels_to_group_by is None:
-            labels = {key: self.cases[key]["label"] for key in cases}
+            labels = {key: self.cases[key]["label"] for key in case_keys}
         else:
-            key0 = list(keys_mapping.keys())[0]
-            if isinstance(key0, tuple):
-                labels = {key: "-".join(key) for key in keys_mapping}
+            level_key = tuple(levels_to_group_by) if len(levels_to_group_by) > 1 else levels_to_group_by[0]
+            if level_key in self.labels_by_levels:
+                labels = self.labels_by_levels[level_key]
             else:
-                labels = {key: key for key in keys_mapping}
+                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
 
@@ -376,12 +428,14 @@ class BenchmarkStudy:
 
         job_keys = []
         for key in case_keys:
+            if verbose:
+                print("### Compute result", key, "###")
             benchmark = self.benchmarks[key]
             assert benchmark is not None
             benchmark.compute_result(**result_params)
             benchmark.save_result(self.folder / "results" / self.key_to_str(key))
 
-    def create_sorting_analyzer_gt(self, case_keys=None, return_scaled=True, random_params={}, **job_kwargs):
+    def create_sorting_analyzer_gt(self, case_keys=None, return_in_uV=True, random_params={}, **job_kwargs):
         print("###### Study.create_sorting_analyzer_gt() is not used anymore!!!!!!")
         # if case_keys is None:
         #     case_keys = self.cases.keys()
@@ -396,7 +450,7 @@ class BenchmarkStudy:
         #     folder = base_folder / self.key_to_str(dataset_key)
         #     recording, gt_sorting = self.datasets[dataset_key]
         #     sorting_analyzer = create_sorting_analyzer(
-        #         gt_sorting, recording, format="binary_folder", folder=folder, return_scaled=return_scaled
+        #         gt_sorting, recording, format="binary_folder", folder=folder, return_in_uV=return_in_uV
         #     )
         #     sorting_analyzer.compute("random_spikes", **random_params)
         #     sorting_analyzer.compute("templates", **job_kwargs)
@@ -431,9 +485,9 @@ class BenchmarkStudy:
             unit_locations_ext = sorting_analyzer.get_extension("unit_locations")
             return unit_locations_ext.get_data()
 
-    def get_templates(self, key, operator="average"):
+    def get_templates(self, key, operator="average", outputs="numpy"):
         sorting_analyzer = self.get_sorting_analyzer(case_key=key)
-        templates = sorting_analyzer.get_extenson("templates").get_data(operator=operator)
+        templates = sorting_analyzer.get_extension("templates").get_data(operator=operator, outputs=outputs)
         return templates
 
     def compute_metrics(self, case_keys=None, metric_names=["snr", "firing_rate"], force=False, **job_kwargs):
@@ -570,7 +624,10 @@ class Benchmark:
             elif format == "zarr_templates":
                 self.result[k].to_zarr(folder / k)
             elif format == "sorting_analyzer":
-                pass
+                analyzer_folder = folder / k
+                if analyzer_folder.exists():
+                    shutil.rmtree(analyzer_folder)
+                self.result[k].save_as(format="binary_folder", folder=analyzer_folder)
             else:
                 raise ValueError(f"Save error {k} {format}")
 
@@ -600,15 +657,26 @@ class Benchmark:
             elif format == "sorting":
                 from spikeinterface.core import load_extractor
 
-                result[k] = load(folder / k)
+                sorting_folder = folder / k
+                if sorting_folder.exists():
+                    result[k] = load(sorting_folder)
             elif format == "Motion":
                 from spikeinterface.core.motion import Motion
 
-                result[k] = Motion.load(folder / k)
+                motion_folder = folder / k
+                if motion_folder.exists():
+                    result[k] = Motion.load(motion_folder)
             elif format == "zarr_templates":
                 from spikeinterface.core.template import Templates
 
-                result[k] = Templates.from_zarr(folder / k)
+                zarr_folder = folder / k
+                if zarr_folder.exists():
+
+                    result[k] = Templates.from_zarr(zarr_folder)
+            elif format == "sorting_analyzer":
+                analyzer_folder = folder / k
+                if analyzer_folder.exists():
+                    result[k] = load_sorting_analyzer(analyzer_folder)
 
         return result
 
@@ -619,3 +687,68 @@ class Benchmark:
     def compute_result(self):
         # run becnhmark result
         raise NotImplementedError
+
+
+# Common feature accross some benchmark : sorter + matching
+class MixinStudyUnitCount:
+    def get_count_units(self, case_keys=None, well_detected_score=None, redundant_score=None, overmerged_score=None):
+        import pandas as pd
+
+        if case_keys is None:
+            case_keys = list(self.cases.keys())
+
+        if isinstance(case_keys[0], str):
+            index = pd.Index(case_keys, name=self.levels)
+        else:
+            index = pd.MultiIndex.from_tuples(case_keys, names=self.levels)
+
+        columns = ["num_gt", "num_sorter", "num_well_detected"]
+        key0 = case_keys[0]
+        comp = self.get_result(key0)["gt_comparison"]
+        if comp.exhaustive_gt:
+            columns.extend(["num_false_positive", "num_redundant", "num_overmerged", "num_bad"])
+        count_units = pd.DataFrame(index=index, columns=columns, dtype=int)
+
+        for key in case_keys:
+            comp = self.get_result(key)["gt_comparison"]
+
+            gt_sorting = comp.sorting1
+            sorting = comp.sorting2
+
+            count_units.at[key, "num_gt"] = len(gt_sorting.get_unit_ids())
+            count_units.at[key, "num_sorter"] = len(sorting.get_unit_ids())
+            count_units.at[key, "num_well_detected"] = comp.count_well_detected_units(well_detected_score)
+
+            if comp.exhaustive_gt:
+                count_units.at[key, "num_redundant"] = comp.count_redundant_units(redundant_score)
+                count_units.at[key, "num_overmerged"] = comp.count_overmerged_units(overmerged_score)
+                count_units.at[key, "num_false_positive"] = comp.count_false_positive_units(redundant_score)
+                count_units.at[key, "num_bad"] = comp.count_bad_units()
+
+        return count_units
+
+    def get_performance_by_unit(self, case_keys=None):
+        import pandas as pd
+
+        if case_keys is None:
+            case_keys = self.cases.keys()
+
+        perf_by_unit = []
+        for key in case_keys:
+            comp = self.get_result(key)["gt_comparison"]
+
+            perf = comp.get_performance(method="by_unit", output="pandas")
+
+            if isinstance(key, str):
+                perf[self.levels] = key
+            elif isinstance(key, tuple):
+                for col, k in zip(self.levels, key):
+                    perf[col] = k
+
+            perf = perf.reset_index()
+            perf_by_unit.append(perf)
+
+        perf_by_unit = pd.concat(perf_by_unit)
+        perf_by_unit = perf_by_unit.set_index(self.levels)
+        perf_by_unit = perf_by_unit.sort_index()
+        return perf_by_unit

{spikeinterface-0.102.2 → spikeinterface-0.103.0}/src/spikeinterface/benchmark/benchmark_clustering.py

@@ -11,7 +11,7 @@ from spikeinterface.widgets import (
 
 import numpy as np
 from spikeinterface.core.job_tools import fix_job_kwargs, split_job_kwargs
-from .benchmark_base import Benchmark, BenchmarkStudy
+from .benchmark_base import Benchmark, BenchmarkStudy, MixinStudyUnitCount
 from spikeinterface.core.sortinganalyzer import create_sorting_analyzer
 from spikeinterface.core.template_tools import get_template_extremum_channel
 
@@ -35,8 +35,8 @@ class ClusteringBenchmark(Benchmark):
         )
         self.result["peak_labels"] = peak_labels
 
-    def compute_result(self, **result_params):
-        result_params, job_kwargs = split_job_kwargs(result_params)
+    def compute_result(self, with_template=False, **job_kwargs):
+        # result_params, job_kwargs = split_job_kwargs(result_params)
         job_kwargs = fix_job_kwargs(job_kwargs)
         self.noise = self.result["peak_labels"] < 0
         spikes = self.gt_sorting.to_spike_vector()
@@ -68,19 +68,20 @@ class ClusteringBenchmark(Benchmark):
             self.result["sliced_gt_sorting"], self.result["clustering"], exhaustive_gt=self.exhaustive_gt
         )
 
-        sorting_analyzer = create_sorting_analyzer(
-            self.result["sliced_gt_sorting"], self.recording, format="memory", sparse=False, **job_kwargs
-        )
-        sorting_analyzer.compute("random_spikes")
-        ext = sorting_analyzer.compute("templates", **job_kwargs)
-        self.result["sliced_gt_templates"] = ext.get_data(outputs="Templates")
+        if with_template:
+            sorting_analyzer = create_sorting_analyzer(
+                self.result["sliced_gt_sorting"], self.recording, format="memory", sparse=False, **job_kwargs
+            )
+            sorting_analyzer.compute("random_spikes")
+            ext = sorting_analyzer.compute("templates", **job_kwargs)
+            self.result["sliced_gt_templates"] = ext.get_data(outputs="Templates")
 
-        sorting_analyzer = create_sorting_analyzer(
-            self.result["clustering"], self.recording, format="memory", sparse=False, **job_kwargs
-        )
-        sorting_analyzer.compute("random_spikes")
-        ext = sorting_analyzer.compute("templates", **job_kwargs)
-        self.result["clustering_templates"] = ext.get_data(outputs="Templates")
+            sorting_analyzer = create_sorting_analyzer(
+                self.result["clustering"], self.recording, format="memory", sparse=False, **job_kwargs
+            )
+            sorting_analyzer.compute("random_spikes")
+            ext = sorting_analyzer.compute("templates", **job_kwargs)
+            self.result["clustering_templates"] = ext.get_data(outputs="Templates")
 
     _run_key_saved = [("peak_labels", "npy")]
 
@@ -93,7 +94,7 @@ class ClusteringBenchmark(Benchmark):
     ]
 
 
-class ClusteringStudy(BenchmarkStudy):
+class ClusteringStudy(BenchmarkStudy, MixinStudyUnitCount):
 
     benchmark_class = ClusteringBenchmark
 
@@ -196,6 +197,11 @@ class ClusteringStudy(BenchmarkStudy):
 
         return plot_performances_vs_depth_and_snr(self, *args, **kwargs)
 
+    def plot_performances_ordered(self, *args, **kwargs):
+        from .benchmark_plot_tools import plot_performances_ordered
+
+        return plot_performances_ordered(self, *args, **kwargs)
+
     def plot_error_metrics(self, metric="cosine", case_keys=None, figsize=(15, 5)):
 
         if case_keys is None: