S094 · Imbalanced Data — SMOTE, Class Weights, Thresholds

"""imbalance_lab.py — four fixes for imbalanced classification, honestly evaluated."""
from __future__ import annotations
import numpy as np
from sklearn.datasets import make_classification
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split, StratifiedKFold, cross_val_score
from sklearn.pipeline import Pipeline as SkPipeline
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import (
    average_precision_score,
    precision_recall_curve,
    confusion_matrix,
    classification_report,
)
from imblearn.pipeline import Pipeline as ImbPipeline
from imblearn.over_sampling import SMOTE
 
RNG = 42
 
 
def make_data(pos_frac: float = 0.01, n: int = 20_000):
    return make_classification(
        n_samples=n, n_features=20, n_informative=6,
        weights=[1 - pos_frac, pos_frac], random_state=RNG,
    )
 
 
def cv_pr_auc(pipe, X, y) -> float:
    cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=RNG)
    return cross_val_score(pipe, X, y, cv=cv, scoring="average_precision").mean()
 
 
def evaluate_at_threshold(pipe, X_test, y_test, target_recall: float = 0.90) -> None:
    prob = pipe.predict_proba(X_test)[:, 1]
    prec, rec, thr = precision_recall_curve(y_test, prob)
    mask = rec[:-1] >= target_recall
    if not mask.any():
        print(f"  Cannot reach recall {target_recall}")
        return
    idx = int(np.argmax(prec[:-1][mask]))
    chosen_thr = float(thr[mask][idx])
    pred = (prob >= chosen_thr).astype(int)
    tn, fp, fn, tp = confusion_matrix(y_test, pred).ravel()
    print(f"  threshold to hit recall≥{target_recall}: {chosen_thr:.3f}")
    print(f"  TP={tp:>4}  FP={fp:>5}  FN={fn:>3}  TN={tn:>5}")
    print(f"  precision@recall={target_recall}: {prec[:-1][mask][idx]:.3f}")
 
 
def build_default() -> SkPipeline:
    return SkPipeline([("scale", StandardScaler()), ("clf", LogisticRegression(max_iter=5000))])
 
 
def build_class_weighted() -> SkPipeline:
    return SkPipeline([
        ("scale", StandardScaler()),
        ("clf", LogisticRegression(max_iter=5000, class_weight="balanced")),
    ])
 
 
def build_smote() -> ImbPipeline:
    return ImbPipeline([
        ("scale", StandardScaler()),
        ("smote", SMOTE(random_state=RNG, k_neighbors=5)),
        ("clf", LogisticRegression(max_iter=5000)),
    ])
 
 
if __name__ == "__main__":
    X, y = make_data(pos_frac=0.01)
    print(f"Dataset: n={len(y)}, positives={int(y.sum())} ({y.mean():.2%})")
 
    Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.3, stratify=y, random_state=RNG)
 
    variants = {
        "default":         build_default(),
        "class_weighted":  build_class_weighted(),
        "smote":           build_smote(),
    }
 
    print("\n5-fold PR-AUC (train fold only):")
    for name, pipe in variants.items():
        score = cv_pr_auc(pipe, Xtr, ytr)
        print(f"  {name:<15s}: {score:.3f}")
 
    # Fit each on full train, evaluate operating point on test
    print("\nHeld-out test, threshold tuned to recall ≥ 0.90:")
    for name, pipe in variants.items():
        pipe.fit(Xtr, ytr)
        print(f"\n[{name}]")
        evaluate_at_threshold(pipe, Xte, yte, target_recall=0.90)
 
    # Full sklearn report for the SMOTE variant at 0.5 threshold as a sanity check
    print("\nSMOTE classifier @ default 0.5 threshold:")
    prob = variants["smote"].predict_proba(Xte)[:, 1]
    pred_default = (prob >= 0.5).astype(int)
    print(classification_report(yte, pred_default, digits=3))