Splitting by ensemble member

import itertools
import re

import numpy as np
from lensml import (
    ANOMALIES,
    BOOSTING_SETTINGS,
    FEATURES,
    N_FOLDS_MAX,
    N_REPEATS,
    RNG_SEED,
    WET_QUANTILE,
    WINDOW,
    apply_transforms,
    fit_transforms,
    member_folds,
    score_splits,
    setup,
)
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import (
    accuracy_score,
    average_precision_score,
    confusion_matrix,
    precision_score,
    recall_score,
)
from sklearn.model_selection import (
    FixedThresholdClassifier,
    StratifiedGroupKFold,
    StratifiedKFold,
)

S = setup()  # the cache, the test members, the transforms (chapters 1-2)
train, test = S.train, S.test
X_tr, y_tr = train[FEATURES].values, train["wet"].values
g_tr = train["member_id"].values
print(
    f"train: {S.n_train_members} members / {len(train)} winters   "
    f"test: {test.member_id.nunique()} members / {len(test)} winters   "
    f"N_FOLDS = {S.n_folds}"
)

N_FOLDS = S.n_folds


def grouped_folds(seed, n_folds):
    # member-grouped, class-balanced folds over the training members
    return StratifiedGroupKFold(
        n_splits=n_folds, shuffle=True, random_state=seed
    )


# the same folds as lensml.member_folds, which chapters 4-7 call
assert all(
    np.array_equal(a1, a2) and np.array_equal(b1, b2)
    for (a1, b1), (a2, b2) in zip(
        grouped_folds(RNG_SEED, N_FOLDS).split(X_tr, y_tr, g_tr),
        member_folds(RNG_SEED, N_FOLDS).split(X_tr, y_tr, g_tr),
        strict=True,
    )
)

print(
    f"base rate train {train.wet.mean():.3f}   (the test members' base rate is "
    "printed once, in chapter 6; nothing here reads their labels)"
)
print(
    f"N_FOLDS = min(N_FOLDS_MAX={N_FOLDS_MAX}, {S.n_train_members} training "
    f"members) = {N_FOLDS}"
)
for k, (a, b) in enumerate(
    grouped_folds(RNG_SEED, N_FOLDS).split(X_tr, y_tr, g_tr)
):
    assert not set(g_tr[a]) & set(g_tr[b]), f"fold {k} leaks members"
print(
    f"{N_FOLDS} member-grouped folds verified: no member appears in both the "
    "fitting and the scoring side of any fold"
)
if S.n_repeats_effective == 1:
    print(
        f"NOTE: leave-one-member-out; the trials use 1 fold draw, not {N_REPEATS}"
    )


def lag1(s):
    v = np.asarray(s, dtype=float)
    return float(np.corrcoef(v[:-1], v[1:])[0, 1])


dependence = (
    train.groupby("member_id")
    .agg(lag1_nino34_anom=("nino34_anom", lag1), lag1_wet=("wet", lag1))
    .round(3)
)
print(
    "\nlag-1 autocorrelation within training members (what consecutive winters "
    "of one member share):"
)
print(dependence.to_string())
print("mean:", dependence.mean().round(3).to_dict())

MICRO_STARTS = {"1231", "1251", "1281", "1301"}
start_year = {
    m: re.search(r"i(\d{4})p", m).group(1) for m in train.member_id.unique()
}
print(
    f"\ncorrelation between training members over {WINDOW[0]}–{WINDOW[1]} "
    "(micro-siblings share a start year; the memory check):"
)
for anom in ANOMALIES:
    wide = train.pivot(index="winter", columns="member_id", values=anom)
    corr = wide.corr()
    sib, unrel = [], []
    for a, b in itertools.combinations(wide.columns, 2):
        same = start_year[a] == start_year[b] and start_year[a] in MICRO_STARTS
        (sib if same else unrel).append(corr.loc[a, b])
    sib, unrel = np.array(sib), np.array(unrel)
    if len(sib) < 2 or len(unrel) < 2:
        print(f"{anom:12s} too few member pairs of one kind to compare")
        continue
    excess = sib.mean() - unrel.mean()
    se = np.sqrt(sib.var(ddof=1) / len(sib) + unrel.var(ddof=1) / len(unrel))
    if abs(excess) <= 2 * se:
        reading = (
            "PASS: no sibling excess detectable (|excess| ≤ 2 SE, SE a lower "
            "bound) — the member is taken as the unit"
        )
    else:
        reading = (
            "FAIL: siblings still move together — group the folds by start "
            "year, or start WINDOW later"
        )
    print(
        f"{anom:12s} micro-siblings {sib.mean():+.3f} ({len(sib)} pairs, SD "
        f"{sib.std(ddof=1):.3f}) · unrelated {unrel.mean():+.3f} ({len(unrel)} pairs, "
        f"SD {unrel.std(ddof=1):.3f}) · siblings − unrelated {excess:+.3f} (±{se:.3f}, "
        "a lower bound: the pairs share members, so they are not independent "
        f"draws) → {reading}"
    )

def score_folds(model, rows, splits, refit_transforms=True):
    out = {m: [] for m in ['ap', 'precision', 'recall', 'accuracy']}
    cm = np.zeros((2, 2), dtype=int)  # rows = true, columns = predicted
    for a, b in splits:
        fit_rows, held_rows = rows.iloc[a], rows.iloc[b]
        if refit_transforms:  # curve and threshold from this fold only
            tf = fit_transforms(fit_rows)
            fit_rows = apply_transforms(fit_rows, tf)
            held_rows = apply_transforms(held_rows, tf)
        Xa, ya = fit_rows[FEATURES].values, fit_rows['wet'].values
        Xb, yb = held_rows[FEATURES].values, held_rows['wet'].values
        model.fit(Xa, ya)
        yhat = model.predict(Xb)
        p_wet = model.predict_proba(Xb)[:, 1]
        out['precision'].append(precision_score(yb, yhat, zero_division=0))
        out['recall'].append(recall_score(yb, yhat, zero_division=0))
        out['accuracy'].append(accuracy_score(yb, yhat))
        out['ap'].append(average_precision_score(yb, p_wet))
        cm += confusion_matrix(yb, yhat, labels=[0, 1])
    return {m: np.array(v) for m, v in out.items()}, cm

CONTRAST_THRESHOLD = 1 - WET_QUANTILE  # the designed base rate
MIN_WET_CALL_SHARE = (
    0.01  # below this, precision and recall gaps are a NON-TEST
)


def contrast_model(seed):
    return FixedThresholdClassifier(
        GradientBoostingClassifier(**BOOSTING_SETTINGS, random_state=seed),
        threshold=CONTRAST_THRESHOLD,
        response_method="predict_proba",
    )


arms = {
    "A member split, transforms refit in each fold   ": dict(
        runs=[], cm=np.zeros((2, 2), dtype=int)
    ),
    "B member split, transforms fit before the split ": dict(
        runs=[], cm=np.zeros((2, 2), dtype=int)
    ),
    "C member-blind split, transforms refit in fold  ": dict(
        runs=[], cm=np.zeros((2, 2), dtype=int)
    ),
}
_folds = list(grouped_folds(RNG_SEED, N_FOLDS).split(X_tr, y_tr, g_tr))
_s05, _cm05 = score_folds(
    GradientBoostingClassifier(**BOOSTING_SETTINGS, random_state=RNG_SEED),
    train,
    _folds,
)
# the same scores as lensml.score_splits, which chapters 4-7 call
_s, _cm = score_folds(contrast_model(RNG_SEED), train, _folds)
_m, _mcm = score_splits(contrast_model(RNG_SEED), train, _folds)
assert all(np.allclose(_s[m], _m[m]) for m in _s)
assert np.array_equal(_cm, _mcm)
print(
    f"for reference — the same model at sklearn's default threshold 0.5 (arm A's "
    f"split, one fold draw): wet calls {int(_cm05[:, 1].sum())} of {int(_cm05.sum())} "
    f"scored winters ({_cm05[:, 1].sum() / _cm05.sum():.2%}), recall "
    f"{_s05['recall'].mean():.3f} — too few calls to read a gap in"
)
for rep in range(S.n_repeats_effective):
    seed = RNG_SEED + rep
    by_member = list(grouped_folds(seed, N_FOLDS).split(X_tr, y_tr, g_tr))
    shuffled = StratifiedKFold(
        n_splits=N_FOLDS, shuffle=True, random_state=seed
    )
    by_year = list(
        shuffled.split(X_tr, y_tr)
    )  # member-blind: winters shuffled
    plan = [(by_member, True), (by_member, False), (by_year, True)]
    for (label, store), (splits, refit) in zip(arms.items(), plan):
        s, cm = score_folds(
            contrast_model(seed), train, splits, refit_transforms=refit
        )
        store["runs"].append(s)
        store["cm"] += cm
METRICS = ["ap", "precision", "recall", "accuracy"]
pooled_arm = {
    label: {m: np.concatenate([s[m] for s in st["runs"]]) for m in METRICS}
    for label, st in arms.items()
}
ref = pooled_arm[next(iter(pooled_arm))]


def sd(v):
    return v.std(ddof=1) if len(v) > 1 else 0.0


print(
    f"boosting at threshold {CONTRAST_THRESHOLD:.3f} (the base rate; 0.5 is sklearn's "
    f"default), {S.n_repeats_effective} fold draw(s) × {N_FOLDS} folds per arm. Per arm: "
    "scores (mean ± per-fold SD) and wet calls; then gap = arm − A (mean ± SD of the "
    "per-fold gaps)"
)
for label, s in pooled_arm.items():
    cm = arms[label]["cm"]
    wet_calls, n_scored = int(cm[:, 1].sum()), int(cm.sum())
    share = wet_calls / n_scored
    print(
        f"{label} AP {s['ap'].mean():.3f}   precision {s['precision'].mean():.3f} ± "
        f"{sd(s['precision']):.3f}   recall {s['recall'].mean():.3f} ± "
        f"{sd(s['recall']):.3f}   accuracy {s['accuracy'].mean():.3f}   wet calls "
        f"{wet_calls} of {n_scored} scored winters ({share:.1%})"
    )
    gaps = {m: s[m] - ref[m] for m in ["ap", "precision", "recall"]}
    if share >= MIN_WET_CALL_SHARE:
        reading = (
            f"gaps readable (wet calls ≥ {MIN_WET_CALL_SHARE:.0%} of winters)"
        )
    else:
        reading = "precision/recall gaps a NON-TEST: read AP only"
    print(
        f"{'':{len(label)}} gap vs A: AP {gaps['ap'].mean():+.4f} ± {sd(gaps['ap']):.4f}"
        f"   precision {gaps['precision'].mean():+.4f} ± {sd(gaps['precision']):.4f}   "
        f"recall {gaps['recall'].mean():+.4f} ± {sd(gaps['recall']):.4f}   → {reading}"
    )