"""Leakage-safe probability calibration wrappers.
The StandardScaler always lives inside the sklearn Pipeline that gets
cross-validated, so no calibration or CV fold ever sees test statistics.
"""
from __future__ import annotations
from typing import Any
import numpy as np
from scipy.optimize import minimize_scalar
from scipy.special import softmax
from sklearn.base import BaseEstimator, ClassifierMixin, clone
from sklearn.calibration import CalibratedClassifierCV
from sklearn.model_selection import StratifiedKFold, train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC
from .metrics import expected_calibration_error
[docs]
def make_base_pipeline(estimator: BaseEstimator, needs_proba: bool) -> Pipeline:
est = clone(estimator)
if isinstance(est, SVC):
# decision_function is enough for CV calibration; Platt inside SVC is
# only needed when we expose the raw model's predict_proba directly.
est.set_params(probability=needs_proba)
return Pipeline([("scaler", StandardScaler()), ("clf", est)])
[docs]
class TemperatureScaler:
"""Single-parameter temperature scaling fitted by NLL minimization."""
[docs]
def __init__(self) -> None:
self.temperature_: float = 1.0
[docs]
def fit(self, logits: np.ndarray, y: np.ndarray) -> "TemperatureScaler":
y = np.asarray(y)
def nll(t: float) -> float:
proba = softmax(logits / t, axis=1)
picked = np.clip(proba[np.arange(len(y)), y], 1e-12, 1.0)
return -float(np.mean(np.log(picked)))
result = minimize_scalar(nll, bounds=(0.05, 20.0), method="bounded")
self.temperature_ = float(result.x)
return self
[docs]
class CalibratedModel(ClassifierMixin, BaseEstimator):
"""One estimator + one calibration method behind a uniform proba API.
method: 'none' (as-shipped probabilities), 'sigmoid'/'isotonic'
(CalibratedClassifierCV over the scaler+clf pipeline), or 'temperature'
(fit on 80%, scale logits on the held-out 20%).
"""
[docs]
def __init__(
self,
estimator: BaseEstimator,
method: str = "sigmoid",
cv: int = 5,
random_state: int = 42,
) -> None:
self.estimator = estimator
self.method = method
self.cv = cv
self.random_state = random_state
[docs]
def fit(self, X: np.ndarray, y: np.ndarray) -> "CalibratedModel":
if self.method == "none":
self.model_ = make_base_pipeline(self.estimator, needs_proba=True).fit(X, y)
elif self.method in ("sigmoid", "isotonic"):
base = make_base_pipeline(self.estimator, needs_proba=False)
splitter = StratifiedKFold(self.cv, shuffle=True, random_state=self.random_state)
self.model_ = CalibratedClassifierCV(base, method=self.method, cv=splitter).fit(X, y)
elif self.method == "temperature":
X_fit, X_cal, y_fit, y_cal = train_test_split(
X, y, test_size=0.2, stratify=y, random_state=self.random_state
)
self.model_ = make_base_pipeline(self.estimator, needs_proba=True).fit(X_fit, y_fit)
self.scaler_t_ = TemperatureScaler().fit(self._logits(X_cal), y_cal)
else:
raise ValueError(f"Unknown calibration method '{self.method}'.")
self.classes_ = self.model_.classes_
return self
def _logits(self, X: np.ndarray) -> np.ndarray:
model: Any = self.model_
if hasattr(model, "decision_function"):
scores = model.decision_function(X)
if scores.ndim == 2:
return scores
return np.log(np.clip(model.predict_proba(X), 1e-12, 1.0))
[docs]
def predict_proba(self, X: np.ndarray) -> np.ndarray:
if self.method == "temperature":
return self.scaler_t_.transform(self._logits(X))
return self.model_.predict_proba(X)
[docs]
def predict(self, X: np.ndarray) -> np.ndarray:
return self.classes_[np.argmax(self.predict_proba(X), axis=1)]
[docs]
def reliability_curve(y_true: np.ndarray, proba: np.ndarray, n_bins: int = 15):
"""Per-bin confidence/accuracy table for reliability diagrams."""
_, table = expected_calibration_error(y_true, proba, n_bins=n_bins)
return table