Baselines and the ladder

import matplotlib.pyplot as plt
import pandas as pd
from lensml import (
    FEATURES,
    GOLD_CMAP,
    RNG_SEED,
    WET_QUANTILE,
    SignRule,
    cite,
    ladder,
    pooled,
    run_folds,
    setup,
    show,
    show_fig,
)
from sklearn.base import clone
from sklearn.dummy import DummyClassifier
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import ConfusionMatrixDisplay
from sklearn.model_selection import FixedThresholdClassifier
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.tree import DecisionTreeClassifier

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}"
)

show(SignRule)  # rung 0c: wet if the autumn Niño 3.4 anomaly is positive

seed = RNG_SEED
models = {
    "always-dry": DummyClassifier(strategy="most_frequent"),
    "always-wet": DummyClassifier(strategy="constant", constant=1),
    "nino-sign": Pipeline([("rule", SignRule())]),
    "logistic": Pipeline(
        [
            ("scale", StandardScaler()),
            ("clf", LogisticRegression(random_state=seed)),
        ]
    ),
    "tree": DecisionTreeClassifier(
        max_depth=3, min_samples_leaf=20, random_state=seed
    ),
    "boosting": GradientBoostingClassifier(
        n_estimators=100, max_depth=2, learning_rate=0.05, random_state=seed
    ),
}
module = ladder(seed)  # the same six, as chapters 5-7 build them
assert list(models) == list(module)
assert all(
    repr(m.get_params()) == repr(module[k].get_params())
    for k, m in models.items()
)

results = {
    name: run_folds(m, train, seed, S.n_folds) for name, m in models.items()
}
print("trained:", ", ".join(models))
coef = clone(models["logistic"]).fit(X_tr, y_tr)[-1].coef_[0, 0]
print(
    f"logistic coefficient on the scaled Niño 3.4 anomaly (all training members): "
    f"{coef:+.3f} log-odds per SD"
)

summary = pd.DataFrame(
    {
        name: dict(
            precision_fold=s["precision"].mean(),
            precision_sd=s["precision"].std(ddof=1),
            recall_fold=s["recall"].mean(),
            recall_sd=s["recall"].std(ddof=1),
            **pooled(cm),
            called_wet=cm[:, 1].sum() / cm.sum(),
            accuracy=s["accuracy"].mean(),
            f1=s["f1"].mean(),
            ap=s["ap"].mean(),
        )
        for name, (s, cm) in results.items()
    }
).T.round(3)
summary["base_rate"] = round(float(y_tr.mean()), 3)
print(
    "nino-sign summed confusion matrix (rows = true, columns = predicted); the "
    "figure below draws the same matrix:"
)
print(
    pd.DataFrame(
        results["nino-sign"][1],
        index=["true dry", "true wet"],
        columns=["pred dry", "pred wet"],
    )
)
summary

fig, axes = plt.subplots(1, len(results), figsize=(3.0 * len(results), 3.4))
for ax, (name, (s, cm)) in zip(axes, results.items()):
    ConfusionMatrixDisplay(cm, display_labels=["dry", "wet"]).plot(
        ax=ax, colorbar=False, cmap=GOLD_CMAP
    )
    ax.set_xlim(-0.85, 1.85)
    ax.set_ylim(1.85, -0.85)  # a canvas frame around the matrix
    pr = pooled(cm)
    ax.set_title(
        f"{name}\npooled P={pr['precision_pooled']:.2f} "
        f"R={pr['recall_pooled']:.2f}",
        fontsize=9,
    )
fig.suptitle(
    f"Summed confusion matrices over {S.n_folds} member-grouped folds (rows = "
    "true, columns = predicted); titles = pooled precision/recall",
    fontsize=9,
)
cite(fig, S.access_date)
plt.tight_layout()
show_fig(
    fig,
    (
        "Six summed confusion matrices over member-grouped folds, one per"
        " rung — always-dry, always-wet, nino-sign, logistic regression, "
        "decision tree, gradient boosting — rows true, columns predicted,"
        " with pooled precision and recall in each title."
    ),
)

# sklearn's default, and the designed base rate
thresholds = [0.5, 1 - WET_QUANTILE]
rows = []
for thr in thresholds:
    m = FixedThresholdClassifier(
        clone(models["logistic"]),
        threshold=thr,
        response_method="predict_proba",
    )
    s, cm = run_folds(m, train, RNG_SEED, S.n_folds)
    rows.append(
        dict(
            model="logistic",
            threshold=round(thr, 3),
            **pooled(cm),
            f1=s["f1"].mean(),
            ap=s["ap"].mean(),
            predicted_wet=int(cm[:, 1].sum()),
            of=int(cm.sum()),
        )
    )
pd.DataFrame(rows).round(3)