Merge pull request #3830 from samuelgarcia/fix_plot_benchmark

benchmark plot : go back to previous with no level in plot_run_times and plot_unit_ counts

Author: yger

src/spikeinterface/benchmark/benchmark_base.py

@@ -619,3 +619,68 @@ def run(self):
     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.loc[key, "num_gt"] = len(gt_sorting.get_unit_ids())
+            count_units.loc[key, "num_sorter"] = len(sorting.get_unit_ids())
+            count_units.loc[key, "num_well_detected"] = comp.count_well_detected_units(well_detected_score)
+
+            if comp.exhaustive_gt:
+                count_units.loc[key, "num_redundant"] = comp.count_redundant_units(redundant_score)
+                count_units.loc[key, "num_overmerged"] = comp.count_overmerged_units(overmerged_score)
+                count_units.loc[key, "num_false_positive"] = comp.count_false_positive_units(redundant_score)
+                count_units.loc[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

src/spikeinterface/benchmark/benchmark_clustering.py

@@ -11,7 +11,7 @@
 
 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
 
@@ -93,7 +93,7 @@ def compute_result(self, **result_params):
     ]
 
 
-class ClusteringStudy(BenchmarkStudy):
+class ClusteringStudy(BenchmarkStudy, MixinStudyUnitCount):
 
     benchmark_class = ClusteringBenchmark
 
@@ -196,6 +196,11 @@ def plot_performances_vs_depth_and_snr(self, *args, **kwargs):
 
         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:

src/spikeinterface/benchmark/benchmark_matching.py

@@ -11,7 +11,7 @@
 )
 
 import numpy as np
-from .benchmark_base import Benchmark, BenchmarkStudy
+from .benchmark_base import Benchmark, BenchmarkStudy, MixinStudyUnitCount
 from spikeinterface.core.basesorting import minimum_spike_dtype
 
 
@@ -53,7 +53,7 @@ def compute_result(self, with_collision=False, **result_params):
     _result_key_saved = [("gt_collision", "pickle"), ("gt_comparison", "pickle")]
 
 
-class MatchingStudy(BenchmarkStudy):
+class MatchingStudy(BenchmarkStudy, MixinStudyUnitCount):
 
     benchmark_class = MatchingBenchmark
 
@@ -84,6 +84,11 @@ def plot_performances_vs_depth_and_snr(self, *args, **kwargs):
 
         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_collisions(self, case_keys=None, figsize=None):
         if case_keys is None:
             case_keys = list(self.cases.keys())
@@ -104,42 +109,6 @@ def plot_collisions(self, case_keys=None, figsize=None):
 
         return fig
 
-    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"]
-        comp = self.get_result(case_keys[0])["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"]
-            assert comp is not None, "You need to do study.run_comparisons() first"
-
-            gt_sorting = comp.sorting1
-            sorting = comp.sorting2
-
-            count_units.loc[key, "num_gt"] = len(gt_sorting.get_unit_ids())
-            count_units.loc[key, "num_sorter"] = len(sorting.get_unit_ids())
-            count_units.loc[key, "num_well_detected"] = comp.count_well_detected_units(well_detected_score)
-
-            if comp.exhaustive_gt:
-                count_units.loc[key, "num_redundant"] = comp.count_redundant_units(redundant_score)
-                count_units.loc[key, "num_overmerged"] = comp.count_overmerged_units(overmerged_score)
-                count_units.loc[key, "num_false_positive"] = comp.count_false_positive_units(redundant_score)
-                count_units.loc[key, "num_bad"] = comp.count_bad_units()
-
-        return count_units
-
     def plot_unit_counts(self, case_keys=None, **kwargs):
         from .benchmark_plot_tools import plot_unit_counts