refactor: improve subplot sharing for consistent axis scales

Author: waridrox

skore/src/skore/sklearn/_plot/metrics/precision_recall_curve.py

@@ -137,6 +137,7 @@ def _plot_single_estimator(
         *,
         estimator_name: str,
         pr_curve_kwargs: list[dict[str, Any]],
+        ax: Optional[Axes] = None,
     ) -> tuple[Axes, list[Line2D], Union[str, None]]:
         """Plot precision-recall curve for a single estimator.
 
@@ -150,6 +151,9 @@ def _plot_single_estimator(
             binary case, we should have a single dict. In multiclass case, we should
             have a list of dicts, one per class.
 
+        ax : matplotlib.axes.Axes, default=None
+            The axes to plot on. If None, self.ax_ is used.
+
         Returns
         -------
         ax : matplotlib.axes.Axes
@@ -165,6 +169,11 @@ def _plot_single_estimator(
         lines: list[Line2D] = []
         line_kwargs: dict[str, Any] = {"drawstyle": "steps-post"}
 
+        # Use the provided axis or self.ax_
+        plot_ax = ax if ax is not None else self.ax_
+        if plot_ax is None:
+            _, plot_ax = plt.subplots()
+
         if self.ml_task == "binary-classification":
             pos_label = cast(PositiveLabel, self.pos_label)
 
@@ -181,7 +190,7 @@ def _plot_single_estimator(
                     f"AP = {self.average_precision[pos_label][0]:0.2f}"
                 )
 
-            (line,) = self.ax_.plot(
+            (line,) = plot_ax.plot(
                 self.recall[pos_label][0],
                 self.precision[pos_label][0],
                 **line_kwargs_validated,
@@ -219,22 +228,23 @@ def _plot_single_estimator(
                         f"AP = {average_precision_class:0.2f}"
                     )
 
-                (line,) = self.ax_.plot(
+                (line,) = plot_ax.plot(
                     recall_class, precision_class, **line_kwargs_validated
                 )
                 lines.append(line)
 
             info_pos_label = None  # irrelevant for multiclass
 
-        self.ax_.legend(bbox_to_anchor=(1.02, 1), title=estimator_name)
+        plot_ax.legend(bbox_to_anchor=(1.02, 1), title=estimator_name)
 
-        return self.ax_, lines, info_pos_label
+        return plot_ax, lines, info_pos_label
 
     def _plot_cross_validated_estimator(
         self,
         *,
         estimator_name: str,
         pr_curve_kwargs: list[dict[str, Any]],
+        ax: Optional[Axes] = None,
     ) -> tuple[Axes, list[Line2D], Union[str, None]]:
         """Plot precision-recall curve for a cross-validated estimator.
 
@@ -248,6 +258,9 @@ def _plot_cross_validated_estimator(
             precision-recall curves. The length of the list should match the number of
             curves to plot.
 
+        ax : matplotlib.axes.Axes, default=None
+            The axes to plot on. If None, self.ax_ is used.
+
         Returns
         -------
         ax : matplotlib.axes.Axes
@@ -263,6 +276,11 @@ def _plot_cross_validated_estimator(
         lines: list[Line2D] = []
         line_kwargs: dict[str, Any] = {"drawstyle": "steps-post"}
 
+        # Use the provided axis or self.ax_
+        plot_ax = ax if ax is not None else self.ax_
+        if plot_ax is None:
+            _, plot_ax = plt.subplots()
+
         if self.ml_task == "binary-classification":
             pos_label = cast(PositiveLabel, self.pos_label)
             for split_idx in range(len(self.precision[pos_label])):
@@ -278,7 +296,7 @@ def _plot_cross_validated_estimator(
                     f"(AP = {average_precision_split:0.2f})"
                 )
 
-                (line,) = self.ax_.plot(
+                (line,) = plot_ax.plot(
                     recall_split, precision_split, **line_kwargs_validated
                 )
                 lines.append(line)
@@ -319,7 +337,7 @@ def _plot_cross_validated_estimator(
                     else:
                         line_kwargs_validated["label"] = None
 
-                    (line,) = self.ax_.plot(
+                    (line,) = plot_ax.plot(
                         recall_split, precision_split, **line_kwargs_validated
                     )
                     lines.append(line)
@@ -328,15 +346,16 @@ def _plot_cross_validated_estimator(
             title = f"{estimator_name} on $\\bf{{{self.data_source}}}$ set"
         else:
             title = f"{estimator_name} on $\\bf{{external}}$ set"
-        self.ax_.legend(bbox_to_anchor=(1.02, 1), title=title)
+        plot_ax.legend(bbox_to_anchor=(1.02, 1), title=title)
 
-        return self.ax_, lines, info_pos_label
+        return plot_ax, lines, info_pos_label
 
     def _plot_comparison_estimator(
         self,
         *,
         estimator_names: list[str],
         pr_curve_kwargs: list[dict[str, Any]],
+        ax: Optional[Axes] = None,
     ) -> tuple[Axes, list[Line2D], Union[str, None]]:
         """Plot precision-recall curve of several estimators.
 
@@ -350,6 +369,9 @@ def _plot_comparison_estimator(
             precision-recall curves. The length of the list should match the number of
             curves to plot.
 
+        ax : matplotlib.axes.Axes, default=None
+            The axes to plot on. If None, self.ax_ is used.
+
         Returns
         -------
         ax : matplotlib.axes.Axes
@@ -365,6 +387,11 @@ def _plot_comparison_estimator(
         lines: list[Line2D] = []
         line_kwargs: dict[str, Any] = {"drawstyle": "steps-post"}
 
+        # Use the provided axis or self.ax_
+        plot_ax = ax if ax is not None else self.ax_
+        if plot_ax is None:
+            _, plot_ax = plt.subplots()
+
         if self.ml_task == "binary-classification":
             pos_label = cast(PositiveLabel, self.pos_label)
             for est_idx, est_name in enumerate(estimator_names):
@@ -378,7 +405,7 @@ def _plot_comparison_estimator(
                 line_kwargs_validated["label"] = (
                     f"{est_name} (AP = {average_precision_est:0.2f})"
                 )
-                (line,) = self.ax_.plot(
+                (line,) = plot_ax.plot(
                     recall_est, precision_est, **line_kwargs_validated
                 )
                 lines.append(line)
@@ -414,17 +441,17 @@ def _plot_comparison_estimator(
                         f"(AP = {average_precision_mean:0.2f})"
                     )
 
-                    (line,) = self.ax_.plot(
+                    (line,) = plot_ax.plot(
                         recall_est_class, precision_est_class, **line_kwargs_validated
                     )
                     lines.append(line)
 
-        self.ax_.legend(
+        plot_ax.legend(
             bbox_to_anchor=(1.02, 1),
             title=f"{self.ml_task.title()} on $\\bf{{{self.data_source}}}$ set",
         )
 
-        return self.ax_, lines, info_pos_label
+        return plot_ax, lines, info_pos_label
 
     @StyleDisplayMixin.style_plot
     def plot(
@@ -434,6 +461,10 @@ def plot(
         estimator_name: Optional[str] = None,
         pr_curve_kwargs: Optional[Union[dict[str, Any], list[dict[str, Any]]]] = None,
         despine: bool = True,
+        subplots: bool = False,
+        nrows: Optional[int] = None,
+        ncols: Optional[int] = None,
+        figsize: Optional[tuple[float, float]] = None,
     ) -> None:
         """Plot visualization.
 
@@ -494,8 +525,16 @@ def plot(
         >>> report = EstimatorReport(classifier, **split_data)
         >>> display = report.metrics.precision_recall()
         >>> display.plot(pr_curve_kwargs={"color": "tab:red"})
+
+        With subplots:
+
+        >>> display.plot(subplots=True)
         """
-        self.figure_, self.ax_ = (ax.figure, ax) if ax is not None else plt.subplots()
+        if ax is not None and subplots:
+            raise ValueError(
+                "Cannot specify both 'ax' and 'subplots=True'. "
+                "Either provide an axes object or use subplots, but not both."
+            )
 
         if pr_curve_kwargs is None:
             pr_curve_kwargs = self._default_pr_curve_kwargs
@@ -725,6 +764,7 @@ def plot(
                     else estimator_name
                 ),
                 pr_curve_kwargs=pr_curve_kwargs,
+                ax=self.ax_,
             )
         elif self.report_type == "cross-validation":
             self.ax_, self.lines_, info_pos_label = (
@@ -735,12 +775,14 @@ def plot(
                         else estimator_name
                     ),
                     pr_curve_kwargs=pr_curve_kwargs,
+                    ax=self.ax_,
                 )
             )
         elif self.report_type == "comparison-estimator":
             self.ax_, self.lines_, info_pos_label = self._plot_comparison_estimator(
                 estimator_names=self.estimator_names,
                 pr_curve_kwargs=pr_curve_kwargs,
+                ax=self.ax_,
             )
         else:
             raise ValueError(
@@ -765,6 +807,8 @@ def plot(
         if despine:
             _despine_matplotlib_axis(self.ax_)
 
+        return self.figure_
+
     @classmethod
     def _compute_data_for_display(
         cls,