"""Experimental orchestrators shared by the driver scripts."""
from __future__ import annotations
from typing import Any
import numpy as np
import pandas as pd
from sklearn.metrics import (
accuracy_score,
cohen_kappa_score,
f1_score,
matthews_corrcoef,
roc_auc_score,
)
from sklearn.model_selection import StratifiedKFold, train_test_split
from xpectrass import get_all_models
from .calibration import CalibratedModel
from .config import TrustConfig
from .data import loso_splits
from .metrics import expected_calibration_error, multiclass_brier
[docs]
def get_estimator(model_name: str):
"""One fresh unfitted estimator from the xpectrass zoo."""
models = get_all_models()
if model_name not in models:
raise KeyError(f"Model '{model_name}' not in xpectrass zoo. "
f"Available: {sorted(models)[:5]}...")
return models[model_name]
def _proba_metrics(y_true, proba, n_classes, ece_bins) -> dict[str, float]:
y_pred = proba.argmax(axis=1)
labels = list(range(n_classes))
ece, _ = expected_calibration_error(y_true, proba, n_bins=ece_bins)
out = {
"accuracy": float(accuracy_score(y_true, y_pred)),
"macro_f1": float(f1_score(y_true, y_pred, average="macro", zero_division=0)),
"mcc": float(matthews_corrcoef(y_true, y_pred)),
"kappa": float(cohen_kappa_score(y_true, y_pred)),
"ece": ece,
"brier": multiclass_brier(y_true, proba),
}
try:
out["roc_auc_ovr"] = float(
roc_auc_score(y_true, proba, multi_class="ovr", average="macro", labels=labels)
)
except ValueError:
out["roc_auc_ovr"] = float("nan")
return out
[docs]
def run_calibration_eval(
model_name: str, route: str, X: np.ndarray, y: np.ndarray,
n_classes: int, trust: TrustConfig,
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
"""CV metrics per calibration method + a holdout reliability table."""
estimator = get_estimator(model_name)
rows: list[dict[str, Any]] = []
reliability: dict[str, Any] = {}
X_tr, X_te, y_tr, y_te = train_test_split(
X, y, test_size=trust.test_size, stratify=y, random_state=trust.random_state
)
for method in trust.calibration_methods:
skf = StratifiedKFold(trust.cv_folds, shuffle=True, random_state=trust.random_state)
fold_metrics = []
for tr_idx, va_idx in skf.split(X_tr, y_tr):
cal = CalibratedModel(estimator, method=method, cv=3,
random_state=trust.random_state)
cal.fit(X_tr[tr_idx], y_tr[tr_idx])
proba = cal.predict_proba(X_tr[va_idx])
fold_metrics.append(_proba_metrics(y_tr[va_idx], proba, n_classes, trust.ece_bins))
agg = {k: float(np.mean([m[k] for m in fold_metrics])) for k in fold_metrics[0]}
rows.append({"route": route, "model": model_name, "calibration": method, **agg})
cal = CalibratedModel(estimator, method=method, cv=trust.cv_folds,
random_state=trust.random_state).fit(X_tr, y_tr)
proba_te = cal.predict_proba(X_te)
ece_te, table = expected_calibration_error(y_te, proba_te, n_bins=trust.ece_bins)
reliability[method] = {"table": table, "ece": ece_te}
return rows, reliability
[docs]
def run_generalization(
model_name: str, route: str, X: np.ndarray, y: np.ndarray, study: np.ndarray,
trust: TrustConfig,
) -> dict[str, Any]:
"""Random 80/20 vs leave-one-study-out macro-F1 for one model."""
estimator = get_estimator(model_name)
X_tr, X_te, y_tr, y_te = train_test_split(
X, y, test_size=trust.test_size, stratify=y, random_state=trust.random_state
)
rand = CalibratedModel(estimator, method="none").fit(X_tr, y_tr)
random_f1 = f1_score(y_te, rand.predict(X_te), average="macro", zero_division=0)
n_classes = int(np.max(y)) + 1
all_labels = list(range(n_classes))
fold_rows = []
for held, tr_idx, te_idx in loso_splits(study):
model = CalibratedModel(estimator, method="none").fit(X[tr_idx], y[tr_idx])
y_te = y[te_idx]
pred = model.predict(X[te_idx])
present = sorted(set(y_te.tolist()))
fold_rows.append({
"fold": held,
"n_test": int(len(te_idx)),
"n_classes_present": len(present),
"n_classes_missing": n_classes - len(present),
# default (present-in-truth-or-prediction labels) — kept for continuity
"macro_f1": f1_score(y_te, pred, average="macro", zero_division=0),
# only classes actually present in this study: artifact-free
"macro_f1_present": f1_score(y_te, pred, labels=present,
average="macro", zero_division=0),
# all classes forced in: absent classes score 0 (most conservative)
"macro_f1_fixed": f1_score(y_te, pred, labels=all_labels,
average="macro", zero_division=0),
"accuracy": accuracy_score(y_te, pred),
})
n_test = np.array([r["n_test"] for r in fold_rows], dtype=float)
weights = n_test / n_test.sum()
def _mean(key: str) -> float:
return float(np.mean([r[key] for r in fold_rows]))
def _wmean(key: str) -> float:
return float(np.sum([w * r[key] for w, r in zip(weights, fold_rows)]))
loso_f1 = _mean("macro_f1") # unchanged headline for backward compatibility
return {
"route": route, "model": model_name,
"random_macro_f1": float(random_f1),
"loso_macro_f1": loso_f1,
"gap": float(random_f1) - loso_f1,
# artifact-free and support-weighted variants
"loso_macro_f1_present": _mean("macro_f1_present"),
"loso_macro_f1_fixed": _mean("macro_f1_fixed"),
"loso_macro_f1_weighted": _wmean("macro_f1_present"),
"loso_accuracy": _wmean("accuracy"),
"gap_present": float(random_f1) - _mean("macro_f1_present"),
"folds": fold_rows,
}