S095 · Model Selection — CV, Hyperparameter Tuning, Optuna

"""model_selection_lab.py — nested CV with Optuna tuning on an imbalanced problem.
 
Pattern:
  outer StratifiedKFold(5)   — reports unbiased mean ± std PR-AUC
    for each outer fold:
      inner Optuna study      — tunes XGBoost hyperparameters via 3-fold inner CV
      refit best on outer train
      score on outer held-out
Final: train on ALL data with best params from the LAST outer fold, ship.
Held-out test set is scored ONCE at the very end.
"""
from __future__ import annotations
import numpy as np
import optuna
from optuna.samplers import TPESampler
from sklearn.datasets import make_classification
from sklearn.model_selection import StratifiedKFold, cross_val_score, train_test_split
from sklearn.metrics import average_precision_score, classification_report
from xgboost import XGBClassifier
 
optuna.logging.set_verbosity(optuna.logging.WARNING)  # keep output readable
RNG = 42
N_OUTER = 5
N_INNER = 3
N_TRIALS = 30
 
 
def make_data(pos_frac: float = 0.05, n: int = 15_000):
    return make_classification(
        n_samples=n, n_features=25, n_informative=8,
        weights=[1 - pos_frac, pos_frac], random_state=RNG,
    )
 
 
def objective(trial: optuna.Trial, X, y) -> float:
    """A single Optuna trial: sample params, run inner CV, return mean PR-AUC."""
    params = {
        "n_estimators": trial.suggest_int("n_estimators", 100, 500, step=50),
        "max_depth": trial.suggest_int("max_depth", 3, 10),
        "learning_rate": trial.suggest_float("learning_rate", 0.01, 0.3, log=True),
        "subsample": trial.suggest_float("subsample", 0.5, 1.0),
        "colsample_bytree": trial.suggest_float("colsample_bytree", 0.5, 1.0),
        "min_child_weight": trial.suggest_int("min_child_weight", 1, 10),
        "reg_lambda": trial.suggest_float("reg_lambda", 1e-3, 10.0, log=True),
        "scale_pos_weight": trial.suggest_float("scale_pos_weight", 1.0, 20.0),
    }
    clf = XGBClassifier(
        **params,
        eval_metric="aucpr",
        random_state=RNG,
        tree_method="hist",
        verbosity=0,
    )
    inner_cv = StratifiedKFold(n_splits=N_INNER, shuffle=True, random_state=RNG)
    scores = cross_val_score(clf, X, y, cv=inner_cv, scoring="average_precision", n_jobs=-1)
    return float(scores.mean())
 
 
def tune_one_outer_fold(X_train_outer, y_train_outer) -> dict:
    study = optuna.create_study(direction="maximize", sampler=TPESampler(seed=RNG))
    study.optimize(
        lambda t: objective(t, X_train_outer, y_train_outer),
        n_trials=N_TRIALS,
        show_progress_bar=False,
    )
    return study.best_params
 
 
def nested_cv(X, y) -> tuple[list[float], list[dict]]:
    """Outer 5-fold. Returns (scores_per_fold, best_params_per_fold)."""
    outer_cv = StratifiedKFold(n_splits=N_OUTER, shuffle=True, random_state=RNG)
    scores, params_list = [], []
    for i, (tr_idx, te_idx) in enumerate(outer_cv.split(X, y), 1):
        Xtr, Xte = X[tr_idx], X[te_idx]
        ytr, yte = y[tr_idx], y[te_idx]
 
        best_params = tune_one_outer_fold(Xtr, ytr)
        clf = XGBClassifier(
            **best_params, eval_metric="aucpr", random_state=RNG,
            tree_method="hist", verbosity=0,
        )
        clf.fit(Xtr, ytr)
        prob = clf.predict_proba(Xte)[:, 1]
        pr_auc = average_precision_score(yte, prob)
        print(f"  outer fold {i}/{N_OUTER}: PR-AUC = {pr_auc:.3f}  best_params={best_params}")
        scores.append(pr_auc)
        params_list.append(best_params)
    return scores, params_list
 
 
if __name__ == "__main__":
    X, y = make_data()
    print(f"Dataset: n={len(y)}, positives={int(y.sum())} ({y.mean():.2%})")
 
    # Held-out set — locked in a drawer until the very end
    X_dev, X_holdout, y_dev, y_holdout = train_test_split(
        X, y, test_size=0.15, stratify=y, random_state=RNG,
    )
    print(f"\nHeld-out size: {len(y_holdout)}  positives: {int(y_holdout.sum())}")
 
    print("\nNested CV on the dev set:")
    scores, params_list = nested_cv(X_dev, y_dev)
    print(f"\nNested CV mean PR-AUC: {np.mean(scores):.3f} ± {np.std(scores):.3f}")
 
    # Final model: tune ONE more time on ALL dev data, refit, score held-out ONCE
    print("\nFinal tuning on full dev set...")
    final_params = tune_one_outer_fold(X_dev, y_dev)
    print(f"Final params: {final_params}")
    final = XGBClassifier(
        **final_params, eval_metric="aucpr", random_state=RNG,
        tree_method="hist", verbosity=0,
    ).fit(X_dev, y_dev)
 
    prob = final.predict_proba(X_holdout)[:, 1]
    holdout_pr_auc = average_precision_score(y_holdout, prob)
    print(f"\nHELD-OUT PR-AUC (scored once): {holdout_pr_auc:.3f}")
    print("\nClassification report @ default 0.5 threshold:")
    print(classification_report(y_holdout, (prob >= 0.5).astype(int), digits=3))