Earnings: height and weight – Applied Bayesian Computing in Python

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

earnings = pd.read_csv(root / "Earnings/data/earnings.csv").dropna(subset=["weight", "height", "male", "ethnicity"])
earnings["c_height"] = earnings["height"] - earnings["height"].mean()
earnings.head()

	height	weight	male	earn	earnk	ethnicity	education	mother_education	father_education	walk	exercise	smokenow	tense	angry	age	c_height
0	74	210.0	1	50000.0	50.0	White	16.0	16.0	16.0	3	3	2.0	0.0	0.0	45	7.40749
1	66	125.0	0	60000.0	60.0	White	16.0	16.0	16.0	6	5	1.0	0.0	0.0	58	-0.59251
2	64	126.0	0	30000.0	30.0	White	16.0	16.0	16.0	8	1	2.0	1.0	1.0	29	-2.59251
3	65	200.0	0	25000.0	25.0	White	17.0	17.0	NaN	8	1	2.0	0.0	0.0	57	-1.59251
4	63	110.0	0	50000.0	50.0	Other	16.0	16.0	16.0	5	6	2.0	0.0	0.0	91	-3.59251

fit_1 = fit_glm("weight ~ height", earnings, seed=17401, prior_scale=10, intercept_scale=100, aux_scale=50)
fit_2 = fit_glm("weight ~ c_height + male", earnings, seed=17402, prior_scale=10, intercept_scale=100, aux_scale=50)
pd.concat({"height": coef_medians(fit_1), "height + male": coef_medians(fit_2)}, axis=1).round(2)

	height	height + male
Intercept	-171.31	151.99
height	4.92	NaN
c_height	NaN	3.92
male	NaN	11.54

grid = pd.DataFrame({"height": np.linspace(earnings.height.min(), earnings.height.max(), 100)})
grid["c_height"] = grid["height"] - earnings["height"].mean()
grid["male"] = 0
pred = fit_2.posterior_epred(grid)
fig, ax = plt.subplots()
ax.scatter(earnings.height, earnings.weight, alpha=.25, s=12)
ax.plot(grid.height, pred.mean(axis=0), color="black")
ax.fill_between(grid.height, np.quantile(pred, .1, axis=0), np.quantile(pred, .9, axis=0), color="0.8")
ax.set_xlabel("Height")
ax.set_ylabel("Weight")

Text(0, 0.5, 'Weight')

Setup