KidIQ: PSIS-LOO

Source: KidIQ/kidiq_loo.Rmd

The original page uses loo with stan_glm objects. lapylace exposes the same pointwise log-likelihood through ArviZ, so the translation can stay in the Bayesian workflow.

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,
    )

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

fit_small = fit_glm("kid_score ~ mom_hs", kidiq, seed=17301, prior_scale=10, intercept_scale=30, aux_scale=30)
fit_rich = fit_glm("kid_score ~ mom_hs + mom_iq", kidiq, seed=17302, prior_scale=10, intercept_scale=30, aux_scale=30)

import arviz as az
loos = {"mom_hs": fit_small.loo(), "mom_hs + mom_iq": fit_rich.loo()}

def loo_row(loo):
    def value(*names):
        for name in names:
            if hasattr(loo, name):
                return getattr(loo, name)
            try:
                return loo[name]
            except Exception:
                pass
        return np.nan

    return {
        "elpd": value("elpd_loo", "loo"),
        "p": value("p_loo"),
        "se": value("se"),
    }

pd.DataFrame({
    name: loo_row(loo)
    for name, loo in loos.items()
}).T

	elpd	p	se
mom_hs	NaN	NaN	13.792585
mom_hs + mom_iq	NaN	NaN	14.248313