test: add cases for different subplots functionality for PredictionErrorDisplay

Author: waridrox

skore/tests/unit/sklearn/plot/test_prediction_error.py

@@ -577,3 +577,143 @@ def test_prediction_error_display_wrong_report_type(pyplot, regression_data):
     )
     with pytest.raises(ValueError, match=err_msg):
         display.plot()
+
+
+def test_prediction_error_display_subplots_basic(pyplot, regression_data):
+    """Test that subplots=True creates multiple subplots with default parameters."""
+    estimator, X_train, X_test, y_train, y_test = regression_data
+    report = ComparisonReport(
+        reports={
+            "estimator 1": EstimatorReport(
+                estimator,
+                X_train=X_train,
+                y_train=y_train,
+                X_test=X_test,
+                y_test=y_test,
+            ),
+            "estimator 2": EstimatorReport(
+                estimator,
+                X_train=X_train,
+                y_train=y_train,
+                X_test=X_test,
+                y_test=y_test,
+            ),
+        },
+    )
+    display = report.metrics.prediction_error()
+    display.plot(subplots=True)
+
+    assert hasattr(display, "figure_")
+    assert len(display.scatter_) == 2
+
+    axes = display.figure_.get_axes()
+    assert len(axes) == 2
+    assert "Model: estimator 1" in axes[0].get_title()
+    assert "Model: estimator 2" in axes[1].get_title()
+
+
+def test_prediction_error_display_subplots_custom_layout(pyplot, regression_data):
+    """Test subplots with custom layout parameters."""
+    estimator, X_train, X_test, y_train, y_test = regression_data
+    report = ComparisonReport(
+        reports={
+            "estimator 1": EstimatorReport(
+                estimator,
+                X_train=X_train,
+                y_train=y_train,
+                X_test=X_test,
+                y_test=y_test,
+            ),
+            "estimator 2": EstimatorReport(
+                estimator,
+                X_train=X_train,
+                y_train=y_train,
+                X_test=X_test,
+                y_test=y_test,
+            ),
+            "estimator 3": EstimatorReport(
+                estimator,
+                X_train=X_train,
+                y_train=y_train,
+                X_test=X_test,
+                y_test=y_test,
+            ),
+        },
+    )
+    display = report.metrics.prediction_error()
+
+    figsize = (10, 8)
+    display.plot(subplots=True, nrows=3, ncols=1, figsize=figsize)
+
+    assert hasattr(display, "figure_")
+    assert display.figure_.get_size_inches()[0] == figsize[0]
+    assert display.figure_.get_size_inches()[1] == figsize[1]
+
+    axes = display.figure_.get_axes()
+    assert len(axes) == 3
+
+    pos1 = axes[0].get_position()
+    pos2 = axes[1].get_position()
+    pos3 = axes[2].get_position()
+
+    # Same column (similar x positions)
+    assert abs(pos1.x0 - pos2.x0) < 0.1
+    assert abs(pos2.x0 - pos3.x0) < 0.1
+
+    # Different rows (decreasing y positions)
+    assert pos1.y0 > pos2.y0
+    assert pos2.y0 > pos3.y0
+
+
+def test_prediction_error_display_subplots_cross_validation(
+    pyplot, regression_data_no_split
+):
+    """Test subplots with cross-validation data."""
+    (estimator, X, y), cv = regression_data_no_split, 3
+    report = CrossValidationReport(estimator, X=X, y=y, cv_splitter=cv)
+    display = report.metrics.prediction_error()
+    display.plot(subplots=True)
+
+    assert hasattr(display, "figure_")
+
+    # Check number of subplots matches number of CV folds
+    axes = display.figure_.get_axes()
+    assert len(axes) == cv
+
+    # Check titles for each fold
+    for i, ax in enumerate(axes):
+        assert f"Fold #{i + 1}" in ax.get_title()
+
+
+def test_prediction_error_display_ax_and_subplots_error(pyplot, regression_data):
+    """Test that an error is raised when both ax and subplots=True are specified."""
+    estimator, X_train, X_test, y_train, y_test = regression_data
+    report = EstimatorReport(
+        estimator, X_train=X_train, y_train=y_train, X_test=X_test, y_test=y_test
+    )
+    display = report.metrics.prediction_error()
+
+    # Create a figure and axis to pass
+    fig, ax = pyplot.subplots()
+
+    # Test that error is raised when both ax and subplots=True are specified
+    with pytest.raises(
+        ValueError, match="Cannot specify both 'ax' and 'subplots=True'"
+    ):
+        display.plot(ax=ax, subplots=True)
+
+
+def test_prediction_error_display_subplots_estimator_report(pyplot, regression_data):
+    """Test subplots with simple estimator report (should be a single plot)."""
+    estimator, X_train, X_test, y_train, y_test = regression_data
+    report = EstimatorReport(
+        estimator, X_train=X_train, y_train=y_train, X_test=X_test, y_test=y_test
+    )
+    display = report.metrics.prediction_error()
+    display.plot(subplots=True)
+
+    # For a single estimator, we should get a single plot
+    assert hasattr(display, "figure_")
+    axes = display.figure_.get_axes()
+    assert len(axes) == 1
+    assert "Model: LinearRegression" in axes[0].get_title()