handle pos_label

Author: GaetandeCast

skore/src/skore/_sklearn/_estimator/metrics_accessor.py

@@ -2060,6 +2060,7 @@ def confusion_matrix(
         X: ArrayLike | None = None,
         y: ArrayLike | None = None,
         threshold: bool = False,
+        pos_label: PositiveLabel | None = _DEFAULT,
     ) -> ConfusionMatrixDisplay:
         """Plot the confusion matrix.
 
@@ -2091,6 +2092,12 @@ def confusion_matrix(
             requires the estimator to have `predict_proba` or `decision_function`
             methods.
 
+        pos_label : int, float, bool, str or None, default=_DEFAULT
+            The label to consider as the positive class when computing the metric. Use
+            this parameter to override the positive class. By default, the positive
+            class is set to the one provided when creating the report. If `None`,
+            the metric is computed considering each class as a positive class.
+
         Returns
         -------
         display : :class:`~skore._sklearn._plot.ConfusionMatrixDisplay`
@@ -2114,9 +2121,12 @@ def confusion_matrix(
         >>> display = report.metrics.confusion_matrix(threshold=True)
         >>> display.plot(threshold=0.7)
         """
+        if pos_label is _DEFAULT:
+            pos_label = self._parent.pos_label
+
         display_kwargs = {
             "display_labels": self._parent.estimator_.classes_.tolist(),
-            "pos_label": self._parent.pos_label,
+            "pos_label": pos_label,
             "threshold": threshold,
         }
 

skore/src/skore/_sklearn/_plot/metrics/confusion_matrix.py

@@ -10,7 +10,7 @@
 from skore._externals._sklearn_compat import confusion_matrix_at_thresholds
 from skore._sklearn._plot.base import DisplayMixin
 from skore._sklearn._plot.utils import _validate_style_kwargs
-from skore._sklearn.types import ReportType, YPlotData
+from skore._sklearn.types import PositiveLabel, ReportType, YPlotData
 
 
 class ConfusionMatrixDisplay(DisplayMixin):
@@ -178,6 +178,9 @@ def _plot_single_estimator(
             heatmap_kwargs or {},
         )
         normalize_by = "normalized_by_" + normalize if normalize else "count"
+        cm_pivot = cm.pivot(
+            index="True label", columns="Predicted label", values=normalize_by
+        ).reindex(index=self.display_labels, columns=self.display_labels)
         sns.heatmap(
             cm.pivot(
                 index="true_label", columns="predicted_label", values=normalize_by
@@ -205,8 +208,8 @@ def _compute_data_for_display(
         *,
         report_type: ReportType,
         display_labels: list[str],
-        pos_label: str,
         threshold: bool = False,
+        pos_label: PositiveLabel | None,
         **kwargs,
     ) -> "ConfusionMatrixDisplay":
         """Compute the confusion matrix for display.
@@ -228,8 +231,9 @@ def _compute_data_for_display(
         display_labels : list of str
             Display labels for plot.
 
-        pos_label : str
-            The positive label.
+        pos_label : int, float, bool, str or None
+            The class considered as the positive class when computing the
+            precision and recall metrics.
 
         threshold : bool, default=False
             Whether to compute the confusion matrix at different thresholds.
@@ -246,6 +250,9 @@ def _compute_data_for_display(
         y_true_values = y_true[0].y
         y_pred_values = y_pred[0].y
         cms = []
+        if isinstance(pos_label, str):
+            neg_label = next(label for label in display_labels if label != pos_label)
+            display_labels = [neg_label, pos_label]
 
         if threshold:
             tns, fps, fns, tps, thresholds = confusion_matrix_at_thresholds(
@@ -260,7 +267,8 @@ def _compute_data_for_display(
                 sklearn_confusion_matrix(
                     y_true=y_true_values,
                     y_pred=y_pred_values,
-            normalize=None,  # we will normalize later
+                    normalize=None,  # we will normalize later
+                    labels=display_labels,
                 )
             )
             thresholds = [None]
@@ -386,4 +394,4 @@ def frame(
                 cm = cm[cm["threshold"] == threshold]
             return cm.pivot(
                 index="true_label", columns="predicted_label", values=normalize_by
-            )
+            ).reindex(index=self.display_labels, columns=self.display_labels)

skore/tests/unit/displays/confusion_matrix/test_estimator.py

@@ -316,11 +316,11 @@ def test_not_implemented_error_for_non_estimator_report(
 
 
 def test_threshold_display_creation(
-    pyplot, logistic_binary_classification_with_train_test
+    pyplot, forest_binary_classification_with_train_test
 ):
     """Check that we can create a confusion matrix display with threshold support."""
     estimator, X_train, X_test, y_train, y_test = (
-        logistic_binary_classification_with_train_test
+        forest_binary_classification_with_train_test
     )
     report = EstimatorReport(
         estimator,
@@ -339,12 +339,12 @@ def test_threshold_display_creation(
 
 
 def test_threshold_display_without_threshold(
-    pyplot, logistic_binary_classification_with_train_test
+    pyplot, forest_binary_classification_with_train_test
 ):
     """Check that do_threshold is False when threshold=False and that we raise an error
     when frame or plot is called with threshold."""
     estimator, X_train, X_test, y_train, y_test = (
-        logistic_binary_classification_with_train_test
+        forest_binary_classification_with_train_test
     )
     report = EstimatorReport(
         estimator,
@@ -371,10 +371,10 @@ def test_threshold_display_without_threshold(
         display.plot(threshold=0.5)
 
 
-def test_plot_with_threshold(pyplot, logistic_binary_classification_with_train_test):
+def test_plot_with_threshold(pyplot, forest_binary_classification_with_train_test):
     """Check that we can plot with a specific threshold."""
     estimator, X_train, X_test, y_train, y_test = (
-        logistic_binary_classification_with_train_test
+        forest_binary_classification_with_train_test
     )
     report = EstimatorReport(
         estimator,
@@ -390,11 +390,11 @@ def test_plot_with_threshold(pyplot, logistic_binary_classification_with_train_t
 
 
 def test_plot_with_default_threshold(
-    pyplot, logistic_binary_classification_with_train_test
+    pyplot, forest_binary_classification_with_train_test
 ):
     """Check that the default threshold (0.5) is used when not specified."""
     estimator, X_train, X_test, y_train, y_test = (
-        logistic_binary_classification_with_train_test
+        forest_binary_classification_with_train_test
     )
     report = EstimatorReport(
         estimator,
@@ -415,10 +415,10 @@ def test_plot_with_default_threshold(
     )
 
 
-def test_frame_with_threshold(logistic_binary_classification_with_train_test):
+def test_frame_with_threshold(forest_binary_classification_with_train_test):
     """Check that we can get a frame at a specific threshold."""
     estimator, X_train, X_test, y_train, y_test = (
-        logistic_binary_classification_with_train_test
+        forest_binary_classification_with_train_test
     )
     report = EstimatorReport(
         estimator,
@@ -434,10 +434,10 @@ def test_frame_with_threshold(logistic_binary_classification_with_train_test):
     assert frame.shape == (2, 2)
 
 
-def test_frame_all_thresholds(logistic_binary_classification_with_train_test):
+def test_frame_all_thresholds(forest_binary_classification_with_train_test):
     """Check that we get all thresholds when threshold=None."""
     estimator, X_train, X_test, y_train, y_test = (
-        logistic_binary_classification_with_train_test
+        forest_binary_classification_with_train_test
     )
     report = EstimatorReport(
         estimator,
@@ -454,12 +454,10 @@ def test_frame_all_thresholds(logistic_binary_classification_with_train_test):
     assert len(frame) == len(display.thresholds_)
 
 
-def test_threshold_normalization(
-    pyplot, logistic_binary_classification_with_train_test
-):
+def test_threshold_normalization(pyplot, forest_binary_classification_with_train_test):
     """Check that normalization works with threshold support."""
     estimator, X_train, X_test, y_train, y_test = (
-        logistic_binary_classification_with_train_test
+        forest_binary_classification_with_train_test
     )
     report = EstimatorReport(
         estimator,
@@ -483,12 +481,10 @@ def test_threshold_normalization(
     assert np.isclose(frame.sum().sum(), 1.0)
 
 
-def test_threshold_closest_match(
-    pyplot, logistic_binary_classification_with_train_test
-):
+def test_threshold_closest_match(pyplot, forest_binary_classification_with_train_test):
     """Check that the closest threshold is selected."""
     estimator, X_train, X_test, y_train, y_test = (
-        logistic_binary_classification_with_train_test
+        forest_binary_classification_with_train_test
     )
     report = EstimatorReport(
         estimator,
@@ -514,12 +510,12 @@ def test_threshold_closest_match(
 
 
 def test_frame_plot_coincidence_with_threshold(
-    pyplot, logistic_binary_classification_with_train_test
+    pyplot, forest_binary_classification_with_train_test
 ):
     """Check that the values in the frame and plot coincide when threshold is
     provided."""
     estimator, X_train, X_test, y_train, y_test = (
-        logistic_binary_classification_with_train_test
+        forest_binary_classification_with_train_test
     )
     report = EstimatorReport(
         estimator,
@@ -533,3 +529,30 @@ def test_frame_plot_coincidence_with_threshold(
     frame_values = frame.values.flatten()
     display.plot(threshold=0.5)
     assert np.allclose(frame_values, display.ax_.collections[0].get_array().flatten())
+
+
+def test_pos_label(pyplot, forest_binary_classification_with_train_test):
+    """Check that the pos_label parameter works correctly."""
+    estimator, X_train, X_test, y_train, y_test = (
+        forest_binary_classification_with_train_test
+    )
+    labels = np.array(["A", "B"], dtype=object)
+    y_train = labels[y_train]
+    y_test = labels[y_test]
+    estimator.fit(X_train, y_train)
+    report = EstimatorReport(
+        estimator,
+        X_train=X_train,
+        y_train=y_train,
+        X_test=X_test,
+        y_test=y_test,
+    )
+    display = report.metrics.confusion_matrix(pos_label="A")
+    display.plot()
+    assert display.ax_.get_xticklabels()[1].get_text() == "A"
+    assert display.ax_.get_yticklabels()[1].get_text() == "A"
+
+    display = report.metrics.confusion_matrix(pos_label="B")
+    display.plot()
+    assert display.ax_.get_xticklabels()[1].get_text() == "B"
+    assert display.ax_.get_yticklabels()[1].get_text() == "B"