Common function between benchlark_clustering/benchmark_matching/benchmark_sorter

Author: samuelgarcia

src/spikeinterface/benchmark/benchmark_clustering.py

@@ -160,49 +160,19 @@ def get_count_units(self, case_keys=None, well_detected_score=None, redundant_sc
 
         return count_units
 
-    def plot_unit_counts(self, case_keys=None, figsize=None, **extra_kwargs):
-        from spikeinterface.widgets.widget_list import plot_study_unit_counts
+    # plotting by methods
+    def plot_unit_counts(self, **kwargs):
+        from .benchmark_plot_tools import plot_unit_counts
+        return plot_unit_counts(self, **kwargs)
 
-        plot_study_unit_counts(self, case_keys, figsize=figsize, **extra_kwargs)
+    def plot_agreement_matrix(self, **kwargs):
+        from .benchmark_plot_tools import plot_agreement_matrix
+        return plot_agreement_matrix(self, **kwargs)
 
-    def plot_agreements(self, case_keys=None, figsize=(15, 15)):
-        if case_keys is None:
-            case_keys = list(self.cases.keys())
-        import pylab as plt
-
-        fig, axs = plt.subplots(ncols=len(case_keys), nrows=1, figsize=figsize, squeeze=False)
-
-        for count, key in enumerate(case_keys):
-            ax = axs[0, count]
-            ax.set_title(self.cases[key]["label"])
-            plot_agreement_matrix(self.get_result(key)["gt_comparison"], ax=ax)
-
-        return fig
-
-    def plot_performances_vs_snr(self, case_keys=None, figsize=(15, 15)):
-        if case_keys is None:
-            case_keys = list(self.cases.keys())
-        import pylab as plt
-
-        fig, axes = plt.subplots(ncols=1, nrows=3, figsize=figsize)
-
-        for count, k in enumerate(("accuracy", "recall", "precision")):
+    def plot_performances_vs_snr(self, **kwargs):
+        from .benchmark_plot_tools import plot_performances_vs_snr
+        return plot_performances_vs_snr(self, **kwargs)
 
-            ax = axes[count]
-            for key in case_keys:
-                label = self.cases[key]["label"]
-
-                analyzer = self.get_sorting_analyzer(key)
-                metrics = analyzer.get_extension("quality_metrics").get_data()
-                x = metrics["snr"].values
-                y = self.get_result(key)["gt_comparison"].get_performance()[k].values
-                ax.scatter(x, y, marker=".", label=label)
-                ax.set_title(k)
-
-            if count == 2:
-                ax.legend()
-
-        return fig
 
     def plot_error_metrics(self, metric="cosine", case_keys=None, figsize=(15, 5)):
 

src/spikeinterface/benchmark/benchmark_matching.py

@@ -61,42 +61,14 @@ def create_benchmark(self, key):
         benchmark = MatchingBenchmark(recording, gt_sorting, params)
         return benchmark
 
-    def plot_agreements(self, case_keys=None, figsize=None):
-        if case_keys is None:
-            case_keys = list(self.cases.keys())
-        import pylab as plt
-
-        fig, axs = plt.subplots(ncols=len(case_keys), nrows=1, figsize=figsize, squeeze=False)
-
-        for count, key in enumerate(case_keys):
-            ax = axs[0, count]
-            ax.set_title(self.cases[key]["label"])
-            plot_agreement_matrix(self.get_result(key)["gt_comparison"], ax=ax)
+    def plot_agreement_matrix(self, **kwargs):
+        from .benchmark_plot_tools import plot_agreement_matrix
+        return plot_agreement_matrix(self, **kwargs)
 
-    def plot_performances_vs_snr(self, case_keys=None, figsize=None, metrics=["accuracy", "recall", "precision"]):
-        if case_keys is None:
-            case_keys = list(self.cases.keys())
-
-        import matplotlib.pyplot as plt
-        fig, axs = plt.subplots(ncols=1, nrows=len(metrics), figsize=figsize, squeeze=False)
-
-        for count, k in enumerate(metrics):
+    def plot_performances_vs_snr(self, **kwargs):
+        from .benchmark_plot_tools import plot_performances_vs_snr
+        return plot_performances_vs_snr(self, **kwargs)
 
-            ax = axs[count, 0]
-            for key in case_keys:
-                label = self.cases[key]["label"]
-
-                analyzer = self.get_sorting_analyzer(key)
-                metrics = analyzer.get_extension("quality_metrics").get_data()
-                x = metrics["snr"].values
-                y = self.get_result(key)["gt_comparison"].get_performance()[k].values
-                ax.scatter(x, y, marker=".", label=label)
-                ax.set_title(k)
-
-            if count == 2:
-                ax.legend()
-
-        return fig
 
     def plot_collisions(self, case_keys=None, figsize=None):
         if case_keys is None:

src/spikeinterface/benchmark/benchmark_plot_tools.py

@@ -217,3 +217,31 @@ def plot_agreement_matrix(study, ordered=True, case_keys=None):
             ax.set_yticks([])
         ax.set_xticks([])
 
+
+def plot_performances_vs_snr(study, case_keys=None, figsize=None, metrics=["accuracy", "recall", "precision"]):
+    import matplotlib.pyplot as plt
+
+    if case_keys is None:
+        case_keys = list(study.cases.keys())
+
+    fig, axs = plt.subplots(ncols=1, nrows=len(metrics), figsize=figsize, squeeze=False)
+
+    for count, k in enumerate(metrics):
+
+        ax = axs[count, 0]
+        for key in case_keys:
+            label = study.cases[key]["label"]
+
+            analyzer = study.get_sorting_analyzer(key)
+            metrics = analyzer.get_extension("quality_metrics").get_data()
+            x = metrics["snr"].values
+            y = study.get_result(key)["gt_comparison"].get_performance()[k].values
+            ax.scatter(x, y, marker=".", label=label)
+            ax.set_title(k)
+        
+        ax.set_ylim(0, 1.05)
+        
+        if count == 2:
+            ax.legend()
+
+    return fig

src/spikeinterface/benchmark/benchmark_sorter.py

@@ -123,3 +123,18 @@ def get_count_units(self, case_keys=None, well_detected_score=None, redundant_sc
 
         return count_units
 
+    # plotting as methods
+    def plot_unit_counts(self, **kwargs):
+        from .benchmark_plot_tools import plot_unit_counts
+        return plot_unit_counts(self, **kwargs)
+
+    def plot_performances(self, **kwargs):
+        from .benchmark_plot_tools import plot_performances
+        return plot_performances(self, **kwargs)
+
+    def plot_agreement_matrix(self, **kwargs):
+        from .benchmark_plot_tools import plot_agreement_matrix
+        return plot_agreement_matrix(self, **kwargs)
+
+
+