KidIQ: k-fold cross-validation

Source: KidIQ/kidiq_kcv.Rmd

The R page compares predictive performance for several stan_glm formulas. This Python translation uses the same formulas with lapylace for the Bayesian model and scikit-learn only for deterministic fold splitting.

Setup

from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import lapylace as lp

root = Path("../../ROS-Examples")

def coef_medians(fit):
    return pd.Series(
        [np.median(fit.alpha_draws()), *np.median(fit.beta_draws(), axis=0)],
        index=["Intercept", *fit.columns],
    )

def fit_glm(formula, data, family=None, seed=1, prior_scale=2.5, intercept_scale=5, aux_scale=10, **kwargs):
    return lp.stan_glm(
        formula,
        data=data,
        family=family or lp.gaussian(),
        prior=lp.normal(0, prior_scale),
        prior_intercept=lp.normal(0, intercept_scale),
        prior_aux=lp.exponential(aux_scale),
        chains=2,
        parallel_chains=2,
        iter_warmup=300,
        iter_sampling=500,
        seed=seed,
        refresh=100,
        **kwargs,
    )
from sklearn.model_selection import KFold
from sklearn.metrics import mean_squared_error

kidiq = pd.read_csv(root / "KidIQ/data/kidiq.csv")
kidiq.head()
kid_score mom_hs mom_iq mom_work mom_age
0 65 1 121.117529 4 27
1 98 1 89.361882 4 25
2 85 1 115.443165 4 27
3 83 1 99.449639 3 25
4 115 1 92.745710 4 27

Baseline Bayesian fit

fit = fit_glm("kid_score ~ mom_hs + mom_iq", kidiq, seed=17101, prior_scale=10, intercept_scale=30, aux_scale=30)
fit.summary(["alpha", "beta", "sigma"])
Mean MCSE StdDev 5% 50% 95% N_Eff N_Eff/s R_hat
alpha 24.744100 0.332486 6.018050 15.354300 24.882700 34.03400 327.61500 1502.82000 1.005890
beta[1] 5.601030 0.078803 2.124470 1.990300 5.609650 9.01148 726.80100 3333.95000 0.999462
beta[2] 0.576259 0.003331 0.060353 0.479869 0.575611 0.67201 328.18400 1505.43000 1.004930
sigma 18.172930 0.026440 0.667380 17.122700 18.132300 19.27370 636.89239 2921.52472 0.998800

K-fold predictive comparison

For speed in the rendered book, each fold uses posterior predictive means from a short Stan run. The comparison is deliberately small but follows the same target as the R example: out-of-fold prediction error.

formulas = [
    "kid_score ~ mom_hs",
    "kid_score ~ mom_hs + mom_iq",
    "kid_score ~ mom_hs * mom_iq",
]

rows = []
kf = KFold(n_splits=5, shuffle=True, random_state=17102)
for j, formula in enumerate(formulas, start=1):
    fold_rmse = []
    for fold, (train_idx, test_idx) in enumerate(kf.split(kidiq), start=1):
        train = kidiq.iloc[train_idx]
        test = kidiq.iloc[test_idx]
        m = fit_glm(formula, train, seed=17110 + 10*j + fold, prior_scale=10, intercept_scale=30, aux_scale=30)
        pred = m.posterior_epred(test).mean(axis=0)
        fold_rmse.append(np.sqrt(mean_squared_error(test["kid_score"], pred)))
    rows.append({"formula": formula, "rmse": np.mean(fold_rmse), "fold_sd": np.std(fold_rmse)})
pd.DataFrame(rows)
formula rmse fold_sd
0 kid_score ~ mom_hs 19.852398 1.481193
1 kid_score ~ mom_hs + mom_iq 18.134966 1.367551
2 kid_score ~ mom_hs * mom_iq 18.089493 1.401176