Bayesian R-squared examples

Source: Rsquared/rsquared.Rmd

The calculations use posterior expected values from lapylace fits. For each model, \(R^2 = Var(\mu)/(Var(\mu) + Var(\epsilon))\) is computed draw by draw.

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

def bayes_r2_gaussian(fit, data):
    mu = fit.posterior_epred(data)
    return np.var(mu, axis=1) / (np.var(mu, axis=1) + fit.stan_variables()["sigma"]**2)

kidiq = pd.read_csv(root / "KidIQ/data/kidiq.csv")
earnings = pd.read_csv(root / "Earnings/data/earnings.csv").dropna(subset=["earn", "height", "male"])
earnings["log_earn"] = np.log(earnings.earn.clip(lower=1))

fits = {
    "KidIQ": (fit_glm("kid_score ~ mom_hs + mom_iq", kidiq, seed=19301, prior_scale=10, intercept_scale=30, aux_scale=30), kidiq),
    "log earnings": (fit_glm("log_earn ~ height + male", earnings, seed=19302, prior_scale=2.5, intercept_scale=10, aux_scale=5), earnings),
}
rows = []
for label, (fit, data) in fits.items():
    r2 = bayes_r2_gaussian(fit, data)
    rows.append({"model": label, "median": np.median(r2), "lo80": np.quantile(r2,.1), "hi80": np.quantile(r2,.9)})
pd.DataFrame(rows)

	model	median	lo80	hi80
0	KidIQ	0.217660	0.181176	0.258708
1	log earnings	0.076843	0.063485	0.092963