Merge pull request #3730 from samuelgarcia/plot_bench

Some benchmark plot improvement

Author: alejoe91

src/spikeinterface/benchmark/benchmark_base.py

@@ -366,6 +366,39 @@ def get_units_snr(self, key):
     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:
     """

src/spikeinterface/benchmark/benchmark_clustering.py

@@ -181,6 +181,21 @@ def plot_performances_vs_snr(self, **kwargs):
 
         return plot_performances_vs_snr(self, **kwargs)
 
+    def plot_performances_comparison(self, *args, **kwargs):
+        from .benchmark_plot_tools import plot_performances_comparison
+
+        return plot_performances_comparison(self, *args, **kwargs)
+
+    def plot_performance_losses(self, *args, **kwargs):
+        from .benchmark_plot_tools import plot_performance_losses
+
+        return plot_performance_losses(self, *args, **kwargs)
+
+    def plot_performances_vs_depth_and_snr(self, *args, **kwargs):
+        from .benchmark_plot_tools import plot_performances_vs_depth_and_snr
+
+        return plot_performances_vs_depth_and_snr(self, *args, **kwargs)
+
     def plot_error_metrics(self, metric="cosine", case_keys=None, figsize=(15, 5)):
 
         if case_keys is None:
@@ -351,104 +366,6 @@ def plot_metrics_vs_depth_and_snr(self, metric="agreement", case_keys=None, figs
 
         return fig
 
-    def plot_unit_losses(self, cases_before, cases_after, metric="agreement", figsize=None):
-
-        fig, axs = plt.subplots(ncols=len(cases_before), nrows=1, figsize=figsize)
-
-        for count, (case_before, case_after) in enumerate(zip(cases_before, cases_after)):
-
-            ax = axs[count]
-            dataset_key = self.cases[case_before]["dataset"]
-            _, gt_sorting1 = self.datasets[dataset_key]
-            positions = gt_sorting1.get_property("gt_unit_locations")
-
-            analyzer = self.get_sorting_analyzer(case_before)
-            metrics_before = analyzer.get_extension("quality_metrics").get_data()
-            x = metrics_before["snr"].values
-
-            y_before = self.get_result(case_before)["gt_comparison"].get_performance()[metric].values
-            y_after = self.get_result(case_after)["gt_comparison"].get_performance()[metric].values
-            ax.set_ylabel("depth (um)")
-            ax.set_ylabel("snr")
-            if count > 0:
-                ax.set_ylabel("")
-                ax.set_yticks([], [])
-            im = ax.scatter(positions[:, 1], x, c=(y_after - y_before), cmap="coolwarm")
-            im.set_clim(-1, 1)
-            # fig.colorbar(im, ax=ax)
-            # ax.set_title(k)
-
-        fig.subplots_adjust(right=0.85)
-        cbar_ax = fig.add_axes([0.9, 0.1, 0.025, 0.75])
-        cbar = fig.colorbar(im, cax=cbar_ax, label=metric)
-        # cbar.set_clim(-1, 1)
-
-        return fig
-
-    def plot_comparison_clustering(
-        self,
-        case_keys=None,
-        performance_names=["accuracy", "recall", "precision"],
-        colors=["g", "b", "r"],
-        ylim=(-0.1, 1.1),
-        figsize=None,
-    ):
-
-        if case_keys is None:
-            case_keys = list(self.cases.keys())
-        import pylab as plt
-
-        num_methods = len(case_keys)
-        fig, axs = plt.subplots(ncols=num_methods, nrows=num_methods, figsize=(10, 10))
-        for i, key1 in enumerate(case_keys):
-            for j, key2 in enumerate(case_keys):
-                if len(axs.shape) > 1:
-                    ax = axs[i, j]
-                else:
-                    ax = axs[j]
-                comp1 = self.get_result(key1)["gt_comparison"]
-                comp2 = self.get_result(key2)["gt_comparison"]
-                if i <= j:
-                    for performance, color in zip(performance_names, colors):
-                        perf1 = comp1.get_performance()[performance]
-                        perf2 = comp2.get_performance()[performance]
-                        ax.plot(perf2, perf1, ".", label=performance, color=color)
-
-                    ax.plot([0, 1], [0, 1], "k--", alpha=0.5)
-                    ax.set_ylim(ylim)
-                    ax.set_xlim(ylim)
-                    ax.spines[["right", "top"]].set_visible(False)
-                    ax.set_aspect("equal")
-
-                    label1 = self.cases[key1]["label"]
-                    label2 = self.cases[key2]["label"]
-                    if j == i:
-                        ax.set_ylabel(f"{label1}")
-                    else:
-                        ax.set_yticks([])
-                    if i == j:
-                        ax.set_xlabel(f"{label2}")
-                    else:
-                        ax.set_xticks([])
-                    if i == num_methods - 1 and j == num_methods - 1:
-                        patches = []
-                        import matplotlib.patches as mpatches
-
-                        for color, name in zip(colors, performance_names):
-                            patches.append(mpatches.Patch(color=color, label=name))
-                        ax.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc="upper left", borderaxespad=0.0)
-                else:
-                    ax.spines["bottom"].set_visible(False)
-                    ax.spines["left"].set_visible(False)
-                    ax.spines["top"].set_visible(False)
-                    ax.spines["right"].set_visible(False)
-                    ax.set_xticks([])
-                    ax.set_yticks([])
-
-        plt.tight_layout(h_pad=0, w_pad=0)
-
-        return fig
-
     def plot_some_over_merged(self, case_keys=None, overmerged_score=0.05, max_units=5, figsize=None):
         if case_keys is None:
             case_keys = list(self.cases.keys())

src/spikeinterface/benchmark/benchmark_matching.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import warnings
+
 from spikeinterface.sortingcomponents.matching import find_spikes_from_templates
 from spikeinterface.core import NumpySorting
 from spikeinterface.comparison import CollisionGTComparison, compare_sorter_to_ground_truth
@@ -77,6 +79,11 @@ def plot_performances_comparison(self, **kwargs):
 
         return plot_performances_comparison(self, **kwargs)
 
+    def plot_performances_vs_depth_and_snr(self, *args, **kwargs):
+        from .benchmark_plot_tools import plot_performances_vs_depth_and_snr
+
+        return plot_performances_vs_depth_and_snr(self, *args, **kwargs)
+
     def plot_collisions(self, case_keys=None, figsize=None):
         if case_keys is None:
             case_keys = list(self.cases.keys())
@@ -138,39 +145,13 @@ def plot_unit_counts(self, case_keys=None, **kwargs):
 
         return plot_unit_counts(self, case_keys, **kwargs)
 
-    def plot_unit_losses(self, before, after, metric=["accuracy"], figsize=None):
-        import matplotlib.pyplot as plt
-
-        fig, axs = plt.subplots(ncols=1, nrows=len(metric), figsize=figsize, squeeze=False)
-
-        for count, k in enumerate(metric):
+    def plot_unit_losses(self, *args, **kwargs):
+        from .benchmark_plot_tools import plot_performance_losses
 
-            ax = axs[0, count]
+        warnings.warn("plot_unit_losses() is now plot_performance_losses()")
+        return plot_performance_losses(self, *args, **kwargs)
 
-            label = self.cases[after]["label"]
+    def plot_performance_losses(self, *args, **kwargs):
+        from .benchmark_plot_tools import plot_performance_losses
 
-            positions = self.get_result(before)["gt_comparison"].sorting1.get_property("gt_unit_locations")
-
-            analyzer = self.get_sorting_analyzer(before)
-            metrics_before = analyzer.get_extension("quality_metrics").get_data()
-            x = metrics_before["snr"].values
-
-            y_before = self.get_result(before)["gt_comparison"].get_performance()[k].values
-            y_after = self.get_result(after)["gt_comparison"].get_performance()[k].values
-            # if count < 2:
-            # ax.set_xticks([], [])
-            # elif count == 2:
-            ax.set_xlabel("depth (um)")
-            im = ax.scatter(positions[:, 1], x, c=(y_after - y_before), cmap="coolwarm")
-            fig.colorbar(im, ax=ax, label=k)
-            im.set_clim(-1, 1)
-            ax.set_title(k)
-            ax.set_ylabel("snr")
-
-        # fig.subplots_adjust(right=0.85)
-        # cbar_ax = fig.add_axes([0.9, 0.1, 0.025, 0.75])
-        # cbar = fig.colorbar(im, cax=cbar_ax, label=metric)
-
-        # if count == 2:
-        #    ax.legend()
-        return fig
+        return plot_performance_losses(self, *args, **kwargs)