Calibrating binary classifier with Venn-ABERS¶

This example shows how to calibrate a binary classifier with VennAbersCalibrator and visualize the impact on predicted probabilities.

We compare an uncalibrated model to its Venn-ABERS calibrated version using reliability diagrams and Brier scores.

from __future__ import annotations

import matplotlib.pyplot as plt
from sklearn.calibration import CalibrationDisplay
from sklearn.datasets import make_classification
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import brier_score_loss

from mapie.calibration import VennAbersCalibrator
from mapie.utils import train_conformalize_test_split

1. Build a miscalibrated binary classifier¶

We generate a toy binary dataset and fit a random forest model which is known to be miscalibrated out of the box (produces probabilities too close to 0 or 1). We use a larger dataset to ensure sufficient data for proper calibration.

X, y = make_classification(
    n_samples=5000,
    n_features=20,
    n_informative=10,
    n_redundant=2,
    class_sep=0.8,
    random_state=42,
)

# Split into train, calibration, and test sets
(X_train, X_calib, X_test, y_train, y_calib, y_test) = train_conformalize_test_split(
    X, y, train_size=0.5, conformalize_size=0.2, test_size=0.3, random_state=42
)

# Use Random Forest which tends to be miscalibrated
base_model = RandomForestClassifier(n_estimators=100, max_depth=10, random_state=42)
base_model.fit(X_train, y_train)
probs_raw = base_model.predict_proba(X_test)[:, 1]
raw_brier = brier_score_loss(y_test, probs_raw)

2. Calibrate with Venn-ABERS¶

We wrap the same base model in VennAbersCalibrator using the inductive mode (default). The calibrator uses the calibration set to learn a calibration mapping that will improve probability estimates.

va_calibrator = VennAbersCalibrator(
    estimator=RandomForestClassifier(n_estimators=100, max_depth=10, random_state=42),
    inductive=True,
    random_state=42,
)
va_calibrator.fit(X_train, y_train, X_calib=X_calib, y_calib=y_calib)
probs_va = va_calibrator.predict_proba(X_test)[:, 1]
va_brier = brier_score_loss(y_test, probs_va)

3. Reliability diagrams and Brier scores¶

Reliability diagrams show how predicted probabilities compare to observed frequencies. Perfect calibration lies on the diagonal. We also display Brier scores to quantify the improvement.

fig, axes = plt.subplots(1, 2, figsize=(12, 5))
CalibrationDisplay.from_predictions(
    y_test,
    probs_raw,
    name=f"Uncalibrated (Brier={raw_brier:.3f})",
    n_bins=10,
    ax=axes[0],
)
CalibrationDisplay.from_predictions(
    y_test,
    probs_va,
    name=f"Venn-ABERS (Brier={va_brier:.3f})",
    n_bins=10,
    ax=axes[1],
)
axes[0].set_title("Before calibration")
axes[1].set_title("After Venn-ABERS calibration")
plt.tight_layout()
plt.show()