Merge branch 'main' into MGAMZ-Improve_job_kwarg_violation_hint

Author: MGAMZ

examples/tutorials/comparison/generate_erroneous_sorting.py

@@ -44,7 +44,8 @@ def generate_erroneous_sorting():
     units_err = {}
 
     # sorting_true have 10 units
-    np.random.seed(0)
+    # np.random.seed(0)
+    rng = np.random.default_rng(seed=0)
 
     # unit 1 2 are perfect
     for u in [1, 2]:
@@ -54,29 +55,29 @@ def generate_erroneous_sorting():
     # unit 3 4 (medium) 10 (low) have medium to low agreement
     for u, score in [(3, 0.8), (4, 0.75), (10, 0.3)]:
         st = sorting_true.get_unit_spike_train(u)
-        st = np.sort(np.random.choice(st, size=int(st.size * score), replace=False))
+        st = np.sort(rng.choice(st, size=int(st.size * score), replace=False))
         units_err[u] = st
 
     # unit 5 6 are over merge
     st5 = sorting_true.get_unit_spike_train(5)
     st6 = sorting_true.get_unit_spike_train(6)
     st = np.unique(np.concatenate([st5, st6]))
-    st = np.sort(np.random.choice(st, size=int(st.size * 0.7), replace=False))
+    st = np.sort(rng.choice(st, size=int(st.size * 0.7), replace=False))
     units_err[56] = st
 
     # unit 7 is over split in 2 part
     st7 = sorting_true.get_unit_spike_train(7)
     st70 = st7[::2]
     units_err[70] = st70
     st71 = st7[1::2]
-    st71 = np.sort(np.random.choice(st71, size=int(st71.size * 0.9), replace=False))
+    st71 = np.sort(rng.choice(st71, size=int(st71.size * 0.9), replace=False))
     units_err[71] = st71
 
     # unit 8 is redundant 3 times
     st8 = sorting_true.get_unit_spike_train(8)
-    st80 = np.sort(np.random.choice(st8, size=int(st8.size * 0.65), replace=False))
-    st81 = np.sort(np.random.choice(st8, size=int(st8.size * 0.6), replace=False))
-    st82 = np.sort(np.random.choice(st8, size=int(st8.size * 0.55), replace=False))
+    st80 = np.sort(rng.choice(st8, size=int(st8.size * 0.65), replace=False))
+    st81 = np.sort(rng.choice(st8, size=int(st8.size * 0.6), replace=False))
+    st82 = np.sort(rng.choice(st8, size=int(st8.size * 0.55), replace=False))
     units_err[80] = st80
     units_err[81] = st81
     units_err[82] = st82

src/spikeinterface/benchmark/benchmark_clustering.py

@@ -192,6 +192,11 @@ def plot_performances_vs_snr(self, **kwargs):
 
         return plot_performances_vs_snr(self, **kwargs)
 
+    def plot_performances_vs_firing_rate(self, **kwargs):
+        from .benchmark_plot_tools import plot_performances_vs_firing_rate
+
+        return plot_performances_vs_firing_rate(self, **kwargs)
+
     def plot_performances_comparison(self, *args, **kwargs):
         from .benchmark_plot_tools import plot_performances_comparison
 

src/spikeinterface/benchmark/benchmark_plot_tools.py

@@ -373,7 +373,8 @@ def plot_agreement_matrix(study, ordered=True, case_keys=None, axs=None):
 
     num_axes = len(case_keys)
     if axs is None:
-        fig, axs = plt.subplots(ncols=num_axes, squeeze=True)
+        fig, axs = plt.subplots(ncols=num_axes, squeeze=False)
+        axs = axs[0, :]
     else:
         assert len(axs) == num_axes, "axs should have the same number of axes as case_keys"
         fig = axs[0].get_figure()
@@ -399,12 +400,13 @@ def plot_agreement_matrix(study, ordered=True, case_keys=None, axs=None):
     return fig
 
 
-def plot_performances_vs_snr(
+def _plot_performances_vs_metric(
     study,
+    metric_name,
     case_keys=None,
     figsize=None,
     performance_names=("accuracy", "recall", "precision"),
-    snr_dataset_reference=None,
+    metric_dataset_reference=None,
     levels_to_group_by=None,
     orientation="vertical",
     show_legend=True,
@@ -414,43 +416,6 @@ def plot_performances_vs_snr(
     num_bin_average=20,
     axs=None,
 ):
-    """
-    Plots performance metrics against signal-to-noise ratio (SNR) for different cases in a study.
-
-    Parameters
-    ----------
-    study : object
-        The study object containing the cases and results.
-    case_keys : list | None, default: None
-        List of case keys to include in the plot. If None, all cases in the study are included.
-    figsize : tuple | None, default: None
-        Size of the figure.
-    performance_names : tuple, default: ("accuracy", "recall", "precision")
-        Names of the performance metrics to plot. Default is ("accuracy", "recall", "precision").
-    snr_dataset_reference : str | None, default: None
-        Reference dataset key to use for SNR. If None, the SNR of each dataset is used.
-    levels_to_group_by : list | None, default: None
-        Levels to group by when mapping case keys.
-    orientation : "vertical" | "horizontal", default: "vertical"
-        The orientation of the plot.
-    show_legend : bool, default True
-        Show legend or not
-    show_sigmoid_fit : bool, default True
-        Show sigmoid that fit the performances.
-    show_average_by_bin : bool, default False
-        Instead of the sigmoid an average by bins can be plotted.
-    scatter_size : int, default 4
-        scatter size
-    num_bin_average : int, default 2
-        Num bin for average
-    axs : matplotlib.axes.Axes | None, default: None
-        The axs to use for plotting. Should be the same size as len(performance_names).
-
-    Returns
-    -------
-    fig : matplotlib.figure.Figure
-        The resulting figure containing the plots.
-    """
     import matplotlib.pyplot as plt
 
     if case_keys is None:
@@ -499,15 +464,15 @@ def plot_performances_vs_snr(
             all_xs = []
             all_ys = []
             for sub_key in key_list:
-                if snr_dataset_reference is None:
+                if metric_dataset_reference is None:
                     # use the SNR of each dataset
                     analyzer = study.get_sorting_analyzer(sub_key)
                 else:
                     # use the same SNR from a reference dataset
-                    analyzer = study.get_sorting_analyzer(dataset_key=snr_dataset_reference)
+                    analyzer = study.get_sorting_analyzer(dataset_key=metric_dataset_reference)
 
                 quality_metrics = analyzer.get_extension("quality_metrics").get_data()
-                x = quality_metrics["snr"].to_numpy(dtype="float64")
+                x = quality_metrics[metric_name].to_numpy(dtype="float64")
                 y = (
                     study.get_result(sub_key)["gt_comparison"]
                     .get_performance()[performance_name]
@@ -564,6 +529,148 @@ def plot_performances_vs_snr(
     return fig
 
 
+def plot_performances_vs_snr(
+    study,
+    case_keys=None,
+    figsize=None,
+    performance_names=("accuracy", "recall", "precision"),
+    metric_dataset_reference=None,
+    levels_to_group_by=None,
+    orientation="vertical",
+    show_legend=True,
+    with_sigmoid_fit=False,
+    show_average_by_bin=True,
+    scatter_size=4,
+    num_bin_average=20,
+    axs=None,
+):
+    """
+    Plots performance metrics against signal-to-noise ratio (SNR) for different cases in a study.
+
+    Parameters
+    ----------
+    study : object
+        The study object containing the cases and results.
+    case_keys : list | None, default: None
+        List of case keys to include in the plot. If None, all cases in the study are included.
+    figsize : tuple | None, default: None
+        Size of the figure.
+    performance_names : tuple, default: ("accuracy", "recall", "precision")
+        Names of the performance metrics to plot. Default is ("accuracy", "recall", "precision").
+    metric_dataset_reference : str | None, default: None
+        Reference dataset metric key to use. If None, the SNR of each dataset is used.
+    levels_to_group_by : list | None, default: None
+        Levels to group by when mapping case keys.
+    orientation : "vertical" | "horizontal", default: "vertical"
+        The orientation of the plot.
+    show_legend : bool, default True
+        Show legend or not
+    show_sigmoid_fit : bool, default True
+        Show sigmoid that fit the performances.
+    show_average_by_bin : bool, default False
+        Instead of the sigmoid an average by bins can be plotted.
+    scatter_size : int, default 4
+        scatter size
+    num_bin_average : int, default 2
+        Num bin for average
+    axs : matplotlib.axes.Axes | None, default: None
+        The axs to use for plotting. Should be the same size as len(performance_names).
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        The resulting figure containing the plots.
+    """
+
+    return _plot_performances_vs_metric(
+        study,
+        "snr",
+        case_keys=case_keys,
+        figsize=figsize,
+        performance_names=performance_names,
+        metric_dataset_reference=metric_dataset_reference,
+        levels_to_group_by=levels_to_group_by,
+        orientation=orientation,
+        show_legend=show_legend,
+        with_sigmoid_fit=with_sigmoid_fit,
+        show_average_by_bin=show_average_by_bin,
+        scatter_size=scatter_size,
+        num_bin_average=num_bin_average,
+        axs=axs,
+    )
+
+
+def plot_performances_vs_firing_rate(
+    study,
+    case_keys=None,
+    figsize=None,
+    performance_names=("accuracy", "recall", "precision"),
+    metric_dataset_reference=None,
+    levels_to_group_by=None,
+    orientation="vertical",
+    show_legend=True,
+    with_sigmoid_fit=False,
+    show_average_by_bin=True,
+    scatter_size=4,
+    num_bin_average=20,
+    axs=None,
+):
+    """
+    Plots performance metrics against firing rate for different cases in a study.
+
+    Parameters
+    ----------
+    study : object
+        The study object containing the cases and results.
+    case_keys : list | None, default: None
+        List of case keys to include in the plot. If None, all cases in the study are included.
+    figsize : tuple | None, default: None
+        Size of the figure.
+    performance_names : tuple, default: ("accuracy", "recall", "precision")
+        Names of the performance metrics to plot. Default is ("accuracy", "recall", "precision").
+    metric_dataset_reference : str | None, default: None
+        Reference dataset metric key to use. If None, the SNR of each dataset is used.
+    levels_to_group_by : list | None, default: None
+        Levels to group by when mapping case keys.
+    orientation : "vertical" | "horizontal", default: "vertical"
+        The orientation of the plot.
+    show_legend : bool, default True
+        Show legend or not
+    show_sigmoid_fit : bool, default True
+        Show sigmoid that fit the performances.
+    show_average_by_bin : bool, default False
+        Instead of the sigmoid an average by bins can be plotted.
+    scatter_size : int, default 4
+        scatter size
+    num_bin_average : int, default 2
+        Num bin for average
+    axs : matplotlib.axes.Axes | None, default: None
+        The axs to use for plotting. Should be the same size as len(performance_names).
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        The resulting figure containing the plots.
+    """
+
+    return _plot_performances_vs_metric(
+        study,
+        "firing_rate",
+        case_keys=case_keys,
+        figsize=figsize,
+        performance_names=performance_names,
+        metric_dataset_reference=metric_dataset_reference,
+        levels_to_group_by=levels_to_group_by,
+        orientation=orientation,
+        show_legend=show_legend,
+        with_sigmoid_fit=with_sigmoid_fit,
+        show_average_by_bin=show_average_by_bin,
+        scatter_size=scatter_size,
+        num_bin_average=num_bin_average,
+        axs=axs,
+    )
+
+
 def plot_performances_ordered(
     study,
     case_keys=None,
@@ -614,9 +721,11 @@ def plot_performances_ordered(
 
     if axs is None:
         if orientation == "vertical":
-            fig, axs = plt.subplots(nrows=num_axes, figsize=figsize, squeeze=True)
+            fig, axs = plt.subplots(nrows=num_axes, figsize=figsize, squeeze=False)
+            axs = axs[:, 0]
         elif orientation == "horizontal":
-            fig, axs = plt.subplots(ncols=num_axes, figsize=figsize, squeeze=True)
+            fig, axs = plt.subplots(ncols=num_axes, figsize=figsize, squeeze=False)
+            axs = axs[0, :]
         else:
             raise ValueError("orientation must be 'vertical' or 'horizontal'")
     else:
@@ -850,7 +959,8 @@ def plot_performances_vs_depth_and_snr(
     case_keys, labels = study.get_grouped_keys_mapping(levels_to_group_by=levels_to_group_by, case_keys=case_keys)
 
     if axs is None:
-        fig, axs = plt.subplots(ncols=len(case_keys), figsize=figsize, squeeze=True)
+        fig, axs = plt.subplots(ncols=len(case_keys), figsize=figsize, squeeze=False)
+        axs = axs[0, :]
     else:
         assert len(axs) == len(case_keys), "axs should have the same number of axes as case_keys"
         fig = axs[0].get_figure()
@@ -927,7 +1037,8 @@ def plot_performance_losses(
     import matplotlib.pyplot as plt
 
     if axs is None:
-        fig, axs = plt.subplots(nrows=len(performance_names), figsize=figsize, squeeze=True)
+        fig, axs = plt.subplots(nrows=len(performance_names), figsize=figsize, squeeze=False)
+        axs = axs[:, 0]
     else:
         assert len(axs) == len(performance_names), "axs should have the same number of axes as performance_names"
         fig = axs[0].get_figure()

src/spikeinterface/benchmark/benchmark_sorter.py

@@ -89,6 +89,11 @@ def plot_performances_vs_snr(self, **kwargs):
 
         return plot_performances_vs_snr(self, **kwargs)
 
+    def plot_performances_vs_firing_rate(self, **kwargs):
+        from .benchmark_plot_tools import plot_performances_vs_firing_rate
+
+        return plot_performances_vs_firing_rate(self, **kwargs)
+
     def plot_performances_ordered(self, **kwargs):
         from .benchmark_plot_tools import plot_performances_ordered