# Arsenic wells: average predictive comparisons
Source: `Arsenic/arsenic_logistic_apc.Rmd`
This page ports the ROS average predictive comparison example. We fit logistic regressions for switching wells and then average counterfactual changes in predicted probability over the observed households.
## Setup and data
```{python}
from pathlib import Path
import numpy as np
import pandas as pd
import lapylace as lp
from scipy.special import expit
root = Path("../../ROS-Examples")
wells = pd.read_csv(root / "Arsenic/data/wells.csv")
wells["dist100"] = wells["dist"] / 100
wells["educ4"] = wells["educ"] / 4
wells.head()
```
```{python}
def coef_medians(fit):
return pd.Series(
[np.median(fit.alpha_draws()), *np.median(fit.beta_draws(), axis=0)],
index=["Intercept", *fit.columns],
)
def posterior_mean_probability(fit, data):
return fit.posterior_epred(data).mean(axis=0)
```
## Model with distance, arsenic, and education
The R page uses `stan_glm`; the Python port uses the same formula-first Bayesian GLM path through `lapylace`.
```{python}
fit_7 = lp.stan_glm(
"switch ~ dist100 + arsenic + educ4",
data=wells,
family=lp.bernoulli(),
prior=lp.normal(0, 2.5),
prior_intercept=lp.normal(0, 5),
chains=4,
parallel_chains=4,
iter_warmup=500,
iter_sampling=1000,
seed=15201,
refresh=100,
)
coef_7 = coef_medians(fit_7)
coef_7.round(2)
```
## Average predictive comparisons
A helper for changing one covariate while holding each household's other covariates fixed:
```{python}
def average_predictive_difference(fit, data, variable, lo, hi):
lo_data = data.copy()
hi_data = data.copy()
lo_data[variable] = lo
hi_data[variable] = hi
return float((posterior_mean_probability(fit, hi_data) - posterior_mean_probability(fit, lo_data)).mean())
apc_distance = average_predictive_difference(fit_7, wells, "dist100", 0, 1)
apc_arsenic = average_predictive_difference(fit_7, wells, "arsenic", 0.5, 1.0)
apc_education = average_predictive_difference(fit_7, wells, "educ4", 0, 3) # 0 vs 12 years
pd.Series({
"distance: 0 to 100 meters": apc_distance,
"arsenic: 0.5 to 1.0": apc_arsenic,
"education: 0 to 12 years": apc_education,
}).round(2)
```
These are average changes in predicted probability, not regression coefficients. They are on the probability scale and therefore depend on the observed distribution of the other predictors.
## Interaction model
The original also computes the distance APC after centering and adding interactions with education.
```{python}
wells["c_dist100"] = wells["dist100"] - wells["dist100"].mean()
wells["c_arsenic"] = wells["arsenic"] - wells["arsenic"].mean()
wells["c_educ4"] = wells["educ4"] - wells["educ4"].mean()
fit_8 = lp.stan_glm(
"switch ~ c_dist100 + c_arsenic + c_educ4 + c_dist100:c_educ4 + c_arsenic:c_educ4",
data=wells,
family=lp.bernoulli(),
prior=lp.normal(0, 2.5),
prior_intercept=lp.normal(0, 5),
chains=4,
parallel_chains=4,
iter_warmup=500,
iter_sampling=1000,
seed=15202,
refresh=100,
)
coef_8 = coef_medians(fit_8)
coef_8.round(2)
```
For a centered predictor, `c_dist100=0` is an average-distance household. The original comparison sets the centered distance predictor to 0 and 1, again averaging over observed arsenic and education.
```{python}
apc_distance_interaction = average_predictive_difference(fit_8, wells, "c_dist100", 0, 1)
round(apc_distance_interaction, 2)
```
## Direct formula check
The same APC can be computed by evaluating the fitted linear predictor by hand:
```{python}
b = coef_8
hi, lo = 1, 0
eta_hi = (
b["Intercept"] + b["c_dist100"] * hi + b["c_arsenic"] * wells["c_arsenic"]
+ b["c_educ4"] * wells["c_educ4"]
+ b["c_dist100:c_educ4"] * hi * wells["c_educ4"]
+ b["c_arsenic:c_educ4"] * wells["c_arsenic"] * wells["c_educ4"]
)
eta_lo = eta_hi - b["c_dist100"] - b["c_dist100:c_educ4"] * wells["c_educ4"]
round(float((expit(eta_hi) - expit(eta_lo)).mean()), 2)
```