import arviz as az
import arviz_plots as azp
import numpy as np
import pandas as pd
import pymc as pm
from matplotlib import pyplot as plt
from patsy import dmatrix
import causalpy as cp
# Setting for Mac OS spawning multi-process defaults on M1 chip
sampler_kwargs = {
"tune": 2000,
"draws": 2000,
"target_accept": 0.95,
"mp_ctx": "spawn",
"random_seed": 1040,
}
# PyTensor's numba StudentT sampler divides by zero for some prior RNG streams.
PRIOR_PREDICTIVE_SEED = 2
The Paradox of Propensity Scores in Bayesian Inference#
In causal inference the role of the propensity score is often seen to be central. For instance, we’ve seen how the propensity score can be used with the cp.InversePropensityWeighting class to correct for a species of selection bias by re-weighting our outcome variable and calculating a causal contrast on the re-weighted scale. Additionally we can use the propensity score to visualise and diagnose problems of overlap or covariate balancing across treatment and control groups.
These properties give the propensity score a large role to play in design based approaches to causal inference. The focus there is on assessing aspects of the treatment allocation to ensure we have identifiability assurances for estimands of interest. What then is their role in model-based or analysis focused Bayesian methods?
When we use cp.InversePropensityWeighting to apply various re-weighting techniques we perform a two-step manoeuvre: 1. we estimate the propensity score and 2. apply the inverse-weighting of the score to transform our outcome variable and assess causal contrasts. But being good Bayesians, we might wonder why go to all this trouble? Can we not simply estimate a full-bayesian model of treatment and outcome simultaneously?
In this notebook we’ll show why we should be careful attempting to model the joint-distribution of the propensity score and the outcome variable, but still make good use of the propensity score.
Brief Digression on the Mathematics#
Consider that we have the following three variables:
where \(Y\) is our outcome variable, \(T\) is our treatment variable, and \(X\) stands in for all other control variables in scope. Now define the propensity score
and our outcome model
but now it’s clearer to see how the propensity score just cancels out. When we’re already conditioning on \(X, T\) the information in the propensity score is technically redundant in the Bayesian setting. Add the assumption of unconfoundedness or ignorability used in causal inference. We are arguing that there is no unmeasured confounding so conditioning on \(X, T\) should be sufficient to identify the causal contrast of interest.
Structure of the Presentation#
We will first consider a simple simulated data set where we know the true values of the data generating process and then we’ll demonstrate how to fit the joint distribution of the propensity score and the outcome in a single joint distribution using PyMC. To contrast this we’ll show how to fit a two-stage version of the same model. In both models we assume independent priors between the regression for the treatment and the regression for the outcome variable. Nevertheless, we will show that the joint model exhibits a bias due to feedback when there is non-measured confounding i.e. when there is misspecified outcome model.
Model Specification
Specifying the Joint ModelSpecifying the 2 Stage Model
Application to Simulated Outcome
Application to Mosquito Net Data
Application to Lalonde Data
Application to NHEFS data.
Note the presentation here owes a debt to the work of Fan Li in [Fan et al., 2023] and her presentation here. Additionally we drew on the work and data of Andrew Heiss here and Jordan Nafa and Andrew Heiss here.
Generate Some Data#
N = 4000
np.random.seed(1043)
def inv_logit(z):
"""Compute the inverse logit (sigmoid) of z."""
return 1 / (1 + np.exp(-z))
df1 = pd.DataFrame(
{
"x1": np.random.normal(0, 1, N),
"x2": np.random.normal(0, 1, N),
"x3": np.random.normal(0, 1, N),
}
)
TREATMENT_EFFECT = 2
df1["trt"] = np.random.binomial(1, inv_logit(df1["x1"] + 6 * df1["x2"] + 7 * df1["x3"]))
df1["outcome"] = (
1
+ TREATMENT_EFFECT * df1["trt"]
+ df1["x1"]
+ 0.2 * df1["x2"]
+ -3 * df1["x3"]
+ np.random.normal(0, 1, N)
)
df1.head()
| x1 | x2 | x3 | trt | outcome | |
|---|---|---|---|---|---|
| 0 | 0.570753 | -1.372992 | 0.247824 | 0 | 1.114661 |
| 1 | 0.063973 | 0.904728 | 0.322222 | 1 | 2.952142 |
| 2 | -0.356653 | -0.188922 | -1.369041 | 0 | 5.809954 |
| 3 | -1.372285 | -0.611158 | -1.755242 | 0 | 5.509733 |
| 4 | 0.029955 | 1.773686 | -0.455371 | 1 | 6.058045 |
Specifying the Joint Model#
Now we define a model context that fits our data simultaneously for treatment and outcome. We allow that the propensity score estimated in the treatment model is used in a non-parametric spline to flexibly inform the outcome variable. Note that we will allow for a richer model specification of the treatment than we do for the outcome model. This is to try and push the models to use the propensity score information to adjust for unmeasured confounding in the outcome model.
coords = {
"betas": ["trt", "x1"],
"betas_trt": ["x1", "x2", "x3"],
"obs": range(df1.shape[0]),
}
N = df1.shape[0]
X_trt = df1[["x1", "x2", "x3"]].values
X_outcome = df1[["trt", "x1"]].values
T_data = df1["trt"].values
Y_data = df1["outcome"].values
def make_joint_model(
X_trt,
X_outcome,
T_data,
Y_data,
coords,
priors=None,
noncentred=True,
normal_outcome=True,
):
if priors is None:
priors = {
"beta_": [0, 1],
"beta_trt": [0, 1],
"alpha_trt": [0, 1],
"alpha_outcome": [0, 1],
"sigma": 1,
"beta_ps": [0, 1],
}
with pm.Model(coords=coords) as model:
X_data_trt = pm.Data("X", X_trt, dims=("obs", "beta_trt"))
X_data_outcome = pm.Data("X_outcome", X_outcome, dims=("obs", "betas"))
T_data_ = pm.Data("T", T_data, dims="obs")
Y_data_ = pm.Data("Y", Y_data, dims="obs")
if noncentred:
mu_beta_trt, sigma_beta_trt = priors["beta_trt"]
beta_trt_std = pm.Normal("beta_trt_std", 0, 1, dims="betas_trt")
beta_trt = pm.Deterministic(
"beta_trt_",
mu_beta_trt + sigma_beta_trt * beta_trt_std,
dims="betas_trt",
)
mu_beta, sigma_beta = priors["beta_"]
beta_std = pm.Normal("beta_std", 0, 1, dims="betas")
beta = pm.Deterministic(
"beta_", mu_beta + sigma_beta * beta_std, dims="betas"
)
else:
beta_trt = pm.Normal(
"beta_trt_",
priors["beta_trt"][0],
priors["beta_trt"][1],
dims="betas_trt",
)
beta = pm.Normal(
"beta_", priors["beta_"][0], priors["beta_"][1], dims="betas"
)
beta_ps = pm.Normal("beta_ps", priors["beta_ps"][0], priors["beta_ps"][1])
alpha_trt = pm.Normal(
"alpha_trt", priors["alpha_trt"][0], priors["alpha_trt"][1]
)
mu_trt = alpha_trt + pm.math.dot(X_data_trt, beta_trt)
p = pm.Deterministic("p", pm.math.invlogit(mu_trt), dims="obs")
pm.Bernoulli("t_pred", p=p, observed=T_data_, dims="obs")
alpha_outcome = pm.Normal(
"alpha_outcome", priors["alpha_outcome"][0], priors["alpha_outcome"][1]
)
mu_outcome = alpha_outcome + pm.math.dot(X_data_outcome, beta) + beta_ps * p
sigma = pm.HalfNormal("sigma", priors["sigma"])
if normal_outcome:
_ = pm.Normal("like", mu_outcome, sigma, observed=Y_data_, dims="obs")
else:
nu = pm.Exponential("nu", lam=1 / 10)
_ = pm.StudentT(
"like", nu=nu, mu=mu_outcome, sigma=sigma, observed=Y_data_, dims="obs"
)
return model
model = make_joint_model(X_trt, X_outcome, T_data, Y_data, coords)
pm.model_to_graphviz(model)
Note how the two likelihood terms are fit simultaneously.
Note
We are specifying the models in raw PyMC code for clarity, but we will bundle these methods into the neater CausalPy API below.
Specifying the 2 Stage Model#
Here we allow for a function that takes the same inputs but fits two separate models. First we fit the treatment model then store the idata_treatment this xarray object stores the posterior estimates for the propensity score. We pass this through to a second outcome model where we proceed to take a random draw from the posterior and pass it through to the outcome regression via a spline component. This allows us to express any non-linearity in the treatment effect. Additionally it can be seen as a way to augment the outcome model.
While theoretically the propensity score contains no extra information if we are already conditioning on \(X\), practically the literature reports that the propensity improves the stability of the causal estimates achievable in Bayesian causal modelling. Additionally we might want to separate covariates for predicting the outcome and the treatment. In this case, there may be extra information derived in the treatment model that be used to inform the outcome model.
def make_treatment_model(
X_trt,
T_data,
coords,
priors=None,
noncentred=True,
):
if priors is None:
priors = {
"beta_": [0, 1],
"beta_trt": [0, 1],
"alpha_trt": [0, 1],
"alpha_outcome": [0, 1],
}
with pm.Model(coords=coords) as model_trt:
X_data_trt = pm.Data("X", X_trt, dims=("obs", "betas_trt"))
T_data_ = pm.Data("T", T_data, dims="obs")
if noncentred:
mu_beta_trt, sigma_beta_trt = priors["beta_trt"]
beta_trt_std = pm.Normal("beta_trt_std", 0, 1, dims="betas_trt")
beta_trt = pm.Deterministic(
"beta_trt_",
mu_beta_trt + sigma_beta_trt * beta_trt_std,
dims="betas_trt",
)
else:
beta_trt = pm.Normal(
"beta_trt_",
priors["beta_trt"][0],
priors["beta_trt"][1],
dims="betas_trt",
)
alpha_trt = pm.Normal(
"alpha_trt", priors["alpha_trt"][0], priors["alpha_trt"][1]
)
mu_trt = alpha_trt + pm.math.dot(X_data_trt, beta_trt)
p = pm.Deterministic("p", pm.math.invlogit(mu_trt), dims="obs")
pm.Bernoulli("t_pred", p=p, observed=T_data_, dims="obs")
return model_trt
def make_outcome_model(
X_outcome,
Y_data,
coords,
priors=None,
noncentred=True,
spline_component=False,
propensity_score_idata=None,
normal_outcome=True,
winsorize_boundary=0.0,
):
if priors is None:
priors = {
"beta_": [0, 1],
"beta_trt": [0, 1],
"alpha_trt": [0, 1],
"alpha_outcome": [0, 1],
"sigma": 1,
"beta_ps": [0, 1],
}
propensity_scores = az.extract(propensity_score_idata)["p"]
with pm.Model(coords=coords) as model_outcome:
X_data_outcome = pm.Data("X_outcome", X_outcome, dims=("obs", "betas"))
Y_data_ = pm.Data("Y", Y_data, dims="obs")
if noncentred:
mu_beta, sigma_beta = priors["beta_"]
beta_std = pm.Normal("beta_std", 0, 1, dims="betas")
beta = pm.Deterministic(
"beta_", mu_beta + sigma_beta * beta_std, dims="betas"
)
else:
beta = pm.Normal(
"beta_", priors["beta_"][0], priors["beta_"][1], dims="betas"
)
beta_ps = pm.Normal("beta_ps", priors["beta_ps"][0], priors["beta_ps"][1])
chosen = np.random.choice(range(propensity_scores.shape[1]))
p = propensity_scores[:, chosen].values
p = np.clip(p, winsorize_boundary, 1 - winsorize_boundary)
alpha_outcome = pm.Normal(
"alpha_outcome", priors["alpha_outcome"][0], priors["alpha_outcome"][1]
)
mu_outcome = alpha_outcome + pm.math.dot(X_data_outcome, beta) + beta_ps * p
if spline_component:
beta_ps_spline = pm.Normal(
"beta_ps_spline", priors["beta_ps"][0], priors["beta_ps"][1], size=14
)
B = dmatrix(
"bs(ps, knots=knots, degree=3, include_intercept=True, lower_bound=0, upper_bound=1) - 1",
{"ps": p, "knots": np.linspace(0, 1, 10)},
)
B_f = np.asarray(B, order="F")
splines_summed = pm.Deterministic(
"spline_features", pm.math.dot(B_f, beta_ps_spline.T), dims="obs"
)
mu_outcome = (
alpha_outcome + pm.math.dot(X_data_outcome, beta) + splines_summed
)
sigma = pm.HalfNormal("sigma", priors["sigma"])
if normal_outcome:
_ = pm.Normal("like", mu_outcome, sigma, observed=Y_data_, dims="obs")
else:
nu = pm.Exponential("nu", lam=1 / 10)
_ = pm.StudentT(
"like", nu=nu, mu=mu_outcome, sigma=sigma, observed=Y_data_, dims="obs"
)
return model_outcome
def make_2step_model(
X_trt,
X_outcome,
T_data,
Y_data,
coords,
priors,
spline_component=False,
normal_outcome=True,
winsorize_boundary=0.0,
):
treatment_model = make_treatment_model(X_trt, T_data, coords, priors)
with treatment_model:
idata_treatment = pm.sample_prior_predictive(random_seed=PRIOR_PREDICTIVE_SEED)
idata_treatment.update(pm.sample(**sampler_kwargs))
outcome_model = make_outcome_model(
X_outcome,
Y_data,
coords,
priors,
spline_component=spline_component,
propensity_score_idata=idata_treatment,
normal_outcome=normal_outcome,
winsorize_boundary=winsorize_boundary,
)
with outcome_model:
idata_outcome = pm.sample_prior_predictive(random_seed=PRIOR_PREDICTIVE_SEED)
idata_outcome.update(pm.sample(**sampler_kwargs))
return idata_treatment, idata_outcome, treatment_model, outcome_model
model_treatment = make_treatment_model(X_trt, T_data, coords)
pm.model_to_graphviz(model_treatment)
Specifying a Simple Regression Model without Propensity Scores#
Now we specify a simple regression model which does not make use of the propensity score information. This model will be used to assess how much extra information is gleaned from the presence of the propensity score as covariate in our outcome model.
def make_reg_model(
X_outcome,
Y_data,
coords,
priors=None,
noncentred=True,
):
if priors is None:
priors = {"beta_": [0, 1], "alpha_outcome": [0, 1], "sigma": 1}
with pm.Model(coords=coords) as reg_model:
X_data_outcome = pm.Data("X_outcome", X_outcome, dims=("obs", "betas"))
Y_data_ = pm.Data("Y", Y_data, dims="obs")
if noncentred:
mu_beta, sigma_beta = priors["beta_"]
beta_std = pm.Normal("beta_std", 0, 1, dims="betas")
beta = pm.Deterministic(
"beta_", mu_beta + sigma_beta * beta_std, dims="betas"
)
else:
beta = pm.Normal(
"beta_", priors["beta_"][0], priors["beta_"][1], dims="betas"
)
alpha_outcome = pm.Normal(
"alpha_outcome", priors["alpha_outcome"][0], priors["alpha_outcome"][1]
)
mu_outcome = alpha_outcome + pm.math.dot(X_data_outcome, beta)
sigma = pm.HalfNormal("sigma", priors["sigma"])
_ = pm.Normal("like", mu_outcome, sigma, observed=Y_data_, dims="obs")
idata = pm.sample(**sampler_kwargs)
return reg_model, idata
Putting it all Together#
We are now in a position to fit both joint and modular (2-stage) models to our simulated data. We are seeking to assess how the different approaches to incorporating propensity score information impacts the accuracy of the treatment effect estimate.
Note
We are using this snippet of code to pass the posterior distribution of the propensity score through to the outcome model and have it be sampled in the MCMC process of the outcome model.
chosen = np.random.choice(range(propensity_scores.shape[1]))
p = propensity_scores[:, chosen].values
This allows us to modularise the fitting process but retain the useful information stored in the propensity score.
priors = {
"beta_": [0, 3],
"beta_trt": [0, 1],
"alpha_trt": [0, 1],
"alpha_outcome": [0, 1],
"sigma": 1,
"beta_ps": [0, 1],
}
joint_model = make_joint_model(X_trt, X_outcome, T_data, Y_data, coords, priors=priors)
with joint_model:
idata_joint = pm.sample(**sampler_kwargs)
(
idata_treatment_2s_joint,
idata_outcome_2s_joint,
treatment_model_joint,
outcome_model_joint,
) = make_2step_model(
X_trt, X_outcome, T_data, Y_data, coords, priors, spline_component=False
)
(
idata_treatment_2s_joint_spline,
idata_outcome_2s_joint_spline,
treatment_model_joint_spline,
outcome_model_joint_spline,
) = make_2step_model(
X_trt, X_outcome, T_data, Y_data, coords, priors, spline_component=True
)
reg_model, idata_outcome_simple_reg = make_reg_model(X_outcome, Y_data, coords)
Show code cell output
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_trt_std, beta_std, beta_ps, alpha_trt, alpha_outcome, sigma]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 71 seconds.
Sampling: [alpha_trt, beta_trt_std, t_pred]
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_trt_std, alpha_trt]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 13 seconds.
Sampling: [alpha_outcome, beta_ps, beta_std, like, sigma]
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_std, beta_ps, alpha_outcome, sigma]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 8 seconds.
Sampling: [alpha_trt, beta_trt_std, t_pred]
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_trt_std, alpha_trt]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 13 seconds.
Sampling: [alpha_outcome, beta_ps, beta_ps_spline, beta_std, like, sigma]
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_std, beta_ps, alpha_outcome, beta_ps_spline, sigma]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 23 seconds.
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_std, alpha_outcome, sigma]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 7 seconds.
forest_plot = azp.plot_forest(
{
"2 Stage": idata_outcome_2s_joint,
"2 Stage + Spline": idata_outcome_2s_joint_spline,
"1 Stage": idata_joint,
"Simple Regression": idata_outcome_simple_reg,
},
var_names=["beta_"],
combined=True,
ci_kind="hdi",
ci_probs=(0.5, 0.94),
figure_kwargs={"figsize": (10, 5)},
)
forest_plot.add_legend("model")
forest_axis = forest_plot.get_target(None, {"column": "forest"})
forest_axis.axvline(2, label="True Treatment Value", color="k")
forest_axis.legend()
forest_plot.add_title("Comparing Joint and 2 Stage Propensity Score Parameter Fits");
Show code cell source
compare_estimate = pd.concat(
{
"1-stage-model": az.summary(
idata_joint,
var_names=["alpha_trt", "beta_", "beta_ps", "alpha_outcome"],
ci_kind="hdi",
ci_prob=0.94,
),
"2-stage-model": az.summary(
idata_outcome_2s_joint,
var_names=["beta_", "alpha_outcome", "beta_ps"],
ci_kind="hdi",
ci_prob=0.94,
),
"2-stage-model_spline": az.summary(
idata_outcome_2s_joint_spline,
var_names=["beta_", "alpha_outcome", "beta_ps"],
ci_kind="hdi",
ci_prob=0.94,
),
"Simple Regression": az.summary(
idata_outcome_simple_reg,
var_names=["beta_", "alpha_outcome"],
ci_kind="hdi",
ci_prob=0.94,
),
}
)
compare_estimate[["mean", "sd", "hdi94_lb", "hdi94_ub", "r_hat"]]
| mean | sd | hdi94_lb | hdi94_ub | r_hat | ||
|---|---|---|---|---|---|---|
| 1-stage-model | alpha_trt | 0.000262 | 0.021438 | -0.040069 | 0.040534 | 1.000686 |
| beta_[trt] | 2.315177 | 0.056498 | 2.208553 | 2.418890 | 1.000179 | |
| beta_[x1] | 1.295097 | 0.069427 | 1.163836 | 1.424174 | 1.000823 | |
| beta_ps | -15.301841 | 0.293221 | -15.863013 | -14.771210 | 1.001305 | |
| alpha_outcome | 8.509280 | 0.161998 | 8.212963 | 8.822284 | 1.000688 | |
| 2-stage-model | beta_[trt] | 2.057667 | 0.188575 | 1.708356 | 2.413761 | 1.001568 |
| beta_[x1] | 1.155191 | 0.041141 | 1.077516 | 1.233993 | 1.001096 | |
| alpha_outcome | 3.066975 | 0.059720 | 2.956578 | 3.180192 | 1.000936 | |
| beta_ps | -4.147031 | 0.209948 | -4.536531 | -3.753973 | 1.000610 | |
| 2-stage-model_spline | beta_[trt] | 1.825330 | 0.182630 | 1.483859 | 2.167799 | 1.001031 |
| beta_[x1] | 1.162270 | 0.039138 | 1.088645 | 1.234812 | 1.000638 | |
| alpha_outcome | 1.006808 | 0.296949 | 0.443748 | 1.545549 | 1.002082 | |
| beta_ps | 0.008287 | 0.980856 | -1.858976 | 1.847617 | 1.000116 | |
| Simple Regression | beta_[trt] | -1.308161 | 0.085164 | -1.467956 | -1.146511 | 1.001107 |
| beta_[x1] | 1.106904 | 0.043156 | 1.026289 | 1.187316 | 1.001067 | |
| alpha_outcome | 2.664807 | 0.059579 | 2.552167 | 2.777695 | 1.000228 |
Here the models fail to recover substantially similar and correct results, in particular the simple regression model is widely off. However the 2-stage models seem to perform better than the joint model specification. This is interesting and demonstrates a key property of propensity scores in the Bayesian setting. Propensity scores are useful correctives within a regression context but we need to be careful how the model is specified.
The Problem of Feedback#
The issue here is sometimes called Bayesian feedback or “collider bias via the likelihood” [Griffith et al., 2020], and it’s a key issue when trying to build joint models for causal inference in the Bayesian paradigm. Because we have fit the outcome and the treatment models simultaneously, and this means that the outcome can influence the posterior distribution of the parameters \(\beta\) in the treatment model and it violates the idea of design-before-analysis. We have here an apparent example of a slight bias due to this effect. The two stage modular approach seems to better recover the treatment effect reported in the literature and avoids the risk of collider bias i.e. in the modular implementation we are able to use the propensity score to adjust for accuracy and compensate for the missing variables x2 and x1.
💡 Key Take-away: With an underspecified outcome model, we may use a well specified propensity score for adjusting the model to retrieve accurate treatment effect estimates. However, this tends to breakdown if we have estimated both propensity score and outcome in a joint bayesian model due to feedback effects. The solution is to use the propensity score in a 2 stage fashion.
pm.model_to_graphviz(outcome_model_joint_spline)
Propensity Score Quantiles - Joint Model#
We can see how the different model specifications have yielded distinct propensity score estimates - as the joint specification seems to compensate for missing covariates in the outcome model by adjusting the propensity score latent in the treatment model too. We can see the differences in the quantiles to highlight a numeric difference in the propensity score distributions between the two models.
idata_joint["posterior"]["p"].quantile(
[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
).round(3).values
array([0.225, 0.308, 0.38 , 0.443, 0.502, 0.563, 0.624, 0.691, 0.769])
Propensity Score Quantiles - 2 Stage Modular Model#
The range an position of these quantiles is vastly different than the joint model specification.
idata_treatment_2s_joint["posterior"]["p"].quantile(
[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
).round(3).values
array([0. , 0.001, 0.014, 0.11 , 0.506, 0.886, 0.986, 0.999, 1. ])
Comparing Propensity Score Skew#
We have seen how the treatment effect reported by both models differ and that the joint model exhibits a bias away from the true treatment effect. But we might want to see how this bias manifests in the propensity score distribution.
def compare_propensity_dists(idata_2s, idata_1s):
fig, ax = plt.subplots(figsize=(10, 4))
for i in range(100):
s2 = idata_2s["posterior"].ds.stack(z=("chain", "draw"))["p"][:, i]
s1 = idata_1s["posterior"].ds.stack(z=("chain", "draw"))["p"][:, i]
if i == 0:
ax.hist(
s2,
alpha=0.1,
color="orange",
label="Propensity Scores 2-stage Estimation",
)
# Pivoted to compare shape
ax.hist(
1 - s1,
alpha=0.1,
color="blue",
label="Propensity Scores 1-stage Estimation",
)
else:
ax.hist(s2, alpha=0.01, color="orange")
# Pivoted to compare shape
ax.hist(1 - s1, alpha=0.01, color="blue")
ax.legend()
ax.set_title(
"Comparing Propensity Distributions \n Two-stage and Joint Estimation routines"
)
compare_propensity_dists(idata_treatment_2s_joint, idata_joint)
These are radically different models due to the feedback mechanism. How does this phenomena play out in some real world examples?
spline_features = az.extract(idata_outcome_2s_joint_spline)["spline_features"]
propensity_scores = az.extract(idata_treatment_2s_joint_spline)["p"]
fig, ax = plt.subplots(figsize=(10, 5))
for i in range(100):
temp = pd.DataFrame(
{"prop_score": propensity_scores[:, i], "spline": spline_features[:, i]}
)
temp.sort_values("prop_score", inplace=True)
ax.plot(temp["prop_score"], temp["spline"], alpha=0.01, color="blue")
temp = pd.DataFrame(
{
"prop_score": propensity_scores.mean(axis=1),
"spline": spline_features.mean(axis=1),
}
)
temp.sort_values("prop_score", inplace=True)
ax.plot(temp["prop_score"], temp["spline"], color="k", label="Expected Value")
ax.set_title("Additive Spline Effect on the Propensity Score")
ax.legend()
ax.set_ylabel("Spline Feature Contribution")
ax.set_xlabel("Propensity Score");
Note
The idea of allowing a spline component to contribute to the mu_outcome or expectation of our model, is that in some cases there might be a non-linear relationship between the propensity for treatment and the outcome interest.
Incorporating the spline function of p grants our model this kind of flexibility to account for different response categories across levels of the propensity score.
Causal Estimate with Do-Operator#
We can also confirm the model implications through counterfactual imputation. This ties the Bayesian setting back to the potential outcome framework. The fundamental problem of causal inference, when seen as a missing data problem allows us to derive causal estimands through imputation of the potential outcomes. Here we “push forward” the posterior predictive distribution for \(Y\) under different treatment settings.
X_outcome_trt = X_outcome.copy()
X_outcome_trt[:, 0] = 1
X_outcome_ntrt = X_outcome.copy()
X_outcome_ntrt[:, 0] = 0
First we specify our counterfactual input data. Then we push them through the joint model distribution using the do-operator in PyMC to sample from the posterior predictive distribution giving us sample of the potential outcomes \(Y(1), Y(0)\)
with pm.do(
joint_model,
{"T": np.ones(len(df1), dtype=np.int32), "X_outcome": X_outcome_trt},
prune_vars=True,
) as treatment_model:
idata_trt = pm.sample_posterior_predictive(idata_joint, var_names=["like", "p"])
with pm.do(
joint_model,
{"T": np.zeros(len(df1), dtype=np.int32), "X_outcome": X_outcome_ntrt},
prune_vars=True,
) as ntreatment_model:
idata_ntrt = pm.sample_posterior_predictive(idata_joint, var_names=["like", "p"])
Sampling: [like]
/Users/benjamv/mambaforge/envs/CausalPy/lib/python3.14/site-packages/pytensor/tensor/type.py:684: RuntimeWarning: invalid value encountered in divide
Sampling: [like]
For sake of illustration we calculate the mean value for \(Y(1)\)
idata_trt["posterior_predictive"]["like"].mean().item()
3.156807245082858
The mean value for \(Y(0)\)
idata_ntrt["posterior_predictive"]["like"].mean().item()
0.8409102488278395
and their difference, which is the causal estimand of interest.
(
idata_trt["posterior_predictive"]["like"].mean().item()
- idata_ntrt["posterior_predictive"]["like"].mean().item()
)
2.3158969962550184
In this way we can impute potential outcomes gain insight into any variety of causal estimands we care to calculate. In this case, we have re-derived the insight that the joint model yields a biased treatment effect estimate.
Nets Example#
Next we’ll asses a data set used by Andrew Heiss to demonstrate propensity score methods with brms.
nets_df = cp.load_data("nets")
nets_df["trt"] = nets_df["net_num"]
nets_df["outcome"] = nets_df["malaria_risk"]
nets_df.head()
| id | net | net_num | malaria_risk | income | health | household | eligible | temperature | resistance | trt | outcome | |
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| 0 | 1 | True | 1 | 33 | 781 | 56 | 2 | False | 21.1 | 59 | 1 | 33 |
| 1 | 2 | False | 0 | 42 | 974 | 57 | 4 | False | 26.5 | 73 | 0 | 42 |
| 2 | 3 | False | 0 | 80 | 502 | 15 | 3 | False | 25.6 | 65 | 0 | 80 |
| 3 | 4 | True | 1 | 34 | 671 | 20 | 5 | True | 21.3 | 46 | 1 | 34 |
| 4 | 5 | False | 0 | 44 | 728 | 17 | 5 | False | 19.2 | 54 | 0 | 44 |
coords = {
"betas": ["trt", "income"],
"betas_trt": ["income", "temperature", "health"],
"obs": range(nets_df.shape[0]),
}
# Process and Standardise Inputs
N = nets_df.shape[0]
X_trt = nets_df[["income", "temperature", "health"]].values
X_trt = (X_trt - X_trt.mean(axis=0)) / X_trt.std(axis=0)
X_outcome = nets_df[["trt", "income"]].values
X_outcome = (X_outcome - X_outcome.mean(axis=0)) / X_outcome.std(axis=0)
T_data = nets_df["trt"].values
X_outcome[:, 0] = T_data
Y_data = nets_df["outcome"].values
priors = {
"beta_": [0, 1],
"beta_trt": [0, 1],
"alpha_trt": [0, 1],
"alpha_outcome": [40, 30],
"sigma": 15,
"beta_ps": [0, 1],
}
net_model = make_joint_model(
X_trt, X_outcome, T_data, Y_data, coords, priors=priors, normal_outcome=False
)
with net_model:
idata_net = pm.sample(tune=2000, target_accept=0.98)
idata_treatment_2s_net, idata_outcome_2s_net, treatment_model_net, outcome_model_net = (
make_2step_model(
X_trt, X_outcome, T_data, Y_data, coords, priors=priors, normal_outcome=False
)
)
reg_model_nets, idata_outcome_simple_reg_nets = make_reg_model(
X_outcome, Y_data, coords
)
Show code cell output
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_trt_std, beta_std, beta_ps, alpha_trt, alpha_outcome, sigma, nu]
Sampling 4 chains for 2_000 tune and 1_000 draw iterations (8_000 + 4_000 draws total) took 18 seconds.
Sampling: [alpha_trt, beta_trt_std, t_pred]
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_trt_std, alpha_trt]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 8 seconds.
Sampling: [alpha_outcome, beta_ps, beta_std, like, nu, sigma]
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_std, beta_ps, alpha_outcome, sigma, nu]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 9 seconds.
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_std, alpha_outcome, sigma]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 6 seconds.
forest_plot = azp.plot_forest(
{
"2 Stage": idata_outcome_2s_net,
"1 Stage": idata_net,
"Simple Regression": idata_outcome_simple_reg_nets,
},
var_names=["beta_"],
combined=True,
ci_kind="hdi",
ci_probs=(0.5, 0.94),
figure_kwargs={"figsize": (10, 4)},
)
forest_plot.add_legend("model")
forest_axis = forest_plot.get_target(None, {"column": "forest"})
forest_axis.axvline(-10, label="True Treatment Value", color="k")
forest_axis.legend()
forest_plot.add_title("Comparing Joint and 2 Stage Propensity Score Parameter Fits");
/Users/benjamv/mambaforge/envs/CausalPy/lib/python3.14/site-packages/arviz_plots/plot_collection.py:57: FutureWarning: In a future version of xarray the default value for join will change from join='outer' to join='exact'. This change will result in the following ValueError: cannot be aligned with join='exact' because index/labels/sizes are not equal along these coordinates (dimensions): 'draw' ('draw',) The recommendation is to set join explicitly for this case.
As before we have used both model specifications to derive estimates for the treatment effect. In this case we have allowed the outcome model access to only a single income predictor. And while both models seem to approximately recover the reported -10 treatment estimate with a large degree of uncertainty. The modular 2 stage estimates pulls away from the joint model estimate.
Show code cell source
compare_estimate = pd.concat(
{
"1-stage-model": az.summary(
idata_net,
var_names=["alpha_trt", "beta_", "beta_ps", "alpha_outcome"],
ci_kind="hdi",
ci_prob=0.94,
),
"2-stage-model": az.summary(
idata_outcome_2s_net,
var_names=["beta_", "alpha_outcome", "beta_ps"],
ci_kind="hdi",
ci_prob=0.94,
),
"Simple Regression": az.summary(
idata_outcome_simple_reg_nets,
var_names=["beta_", "alpha_outcome"],
ci_kind="hdi",
ci_prob=0.94,
),
}
)
compare_estimate[["mean", "sd", "hdi94_lb", "hdi94_ub", "r_hat"]]
| mean | sd | hdi94_lb | hdi94_ub | r_hat | ||
|---|---|---|---|---|---|---|
| 1-stage-model | alpha_trt | -0.261509 | 0.049503 | -0.353789 | -0.167656 | 1.000751 |
| beta_[trt] | -9.932804 | 0.271137 | -10.437929 | -9.425880 | 0.999811 | |
| beta_[income] | -10.101978 | 0.217114 | -10.509079 | -9.688758 | 1.002894 | |
| beta_ps | -16.203019 | 0.620142 | -17.366534 | -15.025392 | 1.000781 | |
| alpha_outcome | 46.583068 | 0.385070 | 45.849041 | 47.295480 | 1.000988 | |
| 2-stage-model | beta_[trt] | -10.055106 | 0.322326 | -10.647972 | -9.450657 | 1.001647 |
| beta_[income] | -10.596990 | 0.191757 | -10.956112 | -10.231711 | 1.000925 | |
| alpha_outcome | 42.192712 | 0.413656 | 41.404905 | 42.963735 | 1.001020 | |
| beta_ps | -7.119886 | 0.944632 | -8.869909 | -5.313062 | 1.001396 | |
| Simple Regression | beta_[trt] | -8.534387 | 0.337863 | -9.182214 | -7.898830 | 1.000496 |
| beta_[income] | -11.515141 | 0.168633 | -11.837117 | -11.197324 | 1.000059 | |
| alpha_outcome | 37.831886 | 0.215691 | 37.424936 | 38.242209 | 1.000618 |
This kind of difference need not be very concerning but we should check if the bias stems from a difference in latent propensity scores as before.
idata_net["posterior"]["p"].quantile(
[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
).round(3).values
array([0.181, 0.249, 0.31 , 0.369, 0.43 , 0.5 , 0.571, 0.642, 0.727])
The quantiles of the propensity distributions across both models seem different.
idata_treatment_2s_net["posterior"]["p"].quantile(
[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
).round(3).values
array([0.239, 0.283, 0.317, 0.349, 0.379, 0.412, 0.449, 0.492, 0.556])
compare_propensity_dists(idata_treatment_2s_net, idata_net)
Here we see some divergence between the propensity score estimates indicating that the regression adjustment approach in the 2 stage outcome model is correcting for bias in the joint-distribution. However both credible intervals contain the reported treatment effect of -10 so the skew is perhaps less concerning. This is despite the fact that we removed useful predictors temperature, health from the outcome model specification. The model leans on the information contained in the propensity score and weights the beta_ps information appropriately.
LaLonde Example#
The Lalonde Data set is famous because it highlights a problem with naive causal contrasts. It is discussed by Angrist and Pischke in their Mostly Harmless Econometrics as an example of how regression controls can tolerably address selection effects in a way similar to propensity score weighting. So we should hope the a well specified outcome model can identify the treatment effects plausibly here too.
lalonde = df = cp.load_data("lalonde")
lalonde[["hispan", "white"]] = pd.get_dummies(lalonde["race"], drop_first=True)
lalonde.dropna(inplace=True)
lalonde.head()
| rownames | treat | age | educ | race | married | nodegree | re74 | re75 | re78 | hispan | white | |
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| 0 | NSW1 | 1 | 37 | 11 | black | 1 | 1 | 0.0 | 0.0 | 9930.0460 | False | False |
| 1 | NSW2 | 1 | 22 | 9 | hispan | 0 | 1 | 0.0 | 0.0 | 3595.8940 | True | False |
| 2 | NSW3 | 1 | 30 | 12 | black | 0 | 0 | 0.0 | 0.0 | 24909.4500 | False | False |
| 3 | NSW4 | 1 | 27 | 11 | black | 0 | 1 | 0.0 | 0.0 | 7506.1460 | False | False |
| 4 | NSW5 | 1 | 33 | 8 | black | 0 | 1 | 0.0 | 0.0 | 289.7899 | False | False |
lalonde.groupby("treat")["re78"].mean().diff()
treat
0 NaN
1 -635.026212
Name: re78, dtype: float64
The Naive group difference suggests a negative effect. Lets see how our two models work when we remove predictors from the propensity model?
coords = {
"betas": ["treat", "nodegree", "married"],
"betas_trt": ["age", "educ", "hispan", "white", "married", "nodegree"],
"obs": range(lalonde.shape[0]),
}
N = lalonde.shape[0]
X_trt = (
lalonde[["age", "educ", "hispan", "white", "married", "nodegree"]]
.astype(np.int32)
.values
)
X_trt = (X_trt - X_trt.mean(axis=0)) / X_trt.std(axis=0)
X_outcome = lalonde[["treat", "nodegree", "married"]].astype(np.int32).values
# X_outcome = (X_outcome - X_outcome.mean()) / X_outcome.std()
T_data = lalonde["treat"].values
Y_data = lalonde["re78"].values
priors = {
"beta_": [0, 4000],
"beta_trt": [0, 1],
"alpha_trt": [0, 1],
"alpha_outcome": [2000, 500],
"sigma": 500,
"beta_ps": [0, 30],
}
lalonde_model = make_joint_model(
X_trt, X_outcome, T_data, Y_data, coords, priors, noncentred=False
)
with lalonde_model:
idata_lalonde = pm.sample(tune=5000)
(
idata_treatment_2s_lalonde,
idata_outcome_2s_lalonde,
treatment_model_lalonde,
outcome_model_lalonde,
) = make_2step_model(
X_trt,
X_outcome,
T_data,
Y_data,
coords,
priors=priors,
spline_component=True,
winsorize_boundary=0.1,
)
reg_model_lalonde, idata_outcome_simple_reg_lalonde = make_reg_model(
X_outcome, Y_data, coords, priors=priors
)
Show code cell output
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_trt_, beta_, beta_ps, alpha_trt, alpha_outcome, sigma]
Sampling 4 chains for 5_000 tune and 1_000 draw iterations (20_000 + 4_000 draws total) took 4 seconds.
Sampling: [alpha_trt, beta_trt_std, t_pred]
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_trt_std, alpha_trt]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 6 seconds.
Sampling: [alpha_outcome, beta_ps, beta_ps_spline, beta_std, like, sigma]
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_std, beta_ps, alpha_outcome, beta_ps_spline, sigma]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 7 seconds.
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_std, alpha_outcome, sigma]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 6 seconds.
forest_plot = azp.plot_forest(
{
"2 Stage": idata_outcome_2s_lalonde,
"1 Stage": idata_lalonde,
"Simple Regression": idata_outcome_simple_reg_lalonde,
},
var_names=["beta_", "sigma"],
combined=True,
ci_kind="hdi",
ci_probs=(0.5, 0.94),
figure_kwargs={"figsize": (10, 4)},
)
forest_plot.add_legend("model")
# Experimental benchmark between 1600 - 1800
forest_plot.get_target(None, {"column": "forest"}).axvline(1700, color="black")
forest_plot.add_title("Comparing Joint and 2 Stage Propensity Score Parameter Fits");
/Users/benjamv/mambaforge/envs/CausalPy/lib/python3.14/site-packages/arviz_plots/plot_collection.py:57: FutureWarning: In a future version of xarray the default value for join will change from join='outer' to join='exact'. This change will result in the following ValueError: cannot be aligned with join='exact' because index/labels/sizes are not equal along these coordinates (dimensions): 'draw' ('draw',) The recommendation is to set join explicitly for this case.
Show code cell source
compare_estimate = pd.concat(
{
"1-stage-model": az.summary(
idata_lalonde,
var_names=["alpha_trt", "beta_", "beta_ps", "alpha_outcome", "sigma"],
ci_kind="hdi",
ci_prob=0.94,
),
"2-stage-model": az.summary(
idata_outcome_2s_lalonde,
var_names=["beta_", "beta_ps", "alpha_outcome", "sigma"],
ci_kind="hdi",
ci_prob=0.94,
),
"Simple Regression": az.summary(
idata_outcome_simple_reg_lalonde,
var_names=["beta_", "alpha_outcome", "sigma"],
ci_kind="hdi",
ci_prob=0.94,
),
}
)
compare_estimate[["mean", "sd", "hdi94_lb", "hdi94_ub", "r_hat"]]
| mean | sd | hdi94_lb | hdi94_ub | r_hat | ||
|---|---|---|---|---|---|---|
| 1-stage-model | alpha_trt | -1.361970 | 0.138276 | -1.623596 | -1.115152 | 1.000000 |
| beta_[treat] | 1582.723882 | 575.510028 | 504.187842 | 2658.043056 | 1.000257 | |
| beta_[nodegree] | -468.942482 | 476.936864 | -1366.074220 | 437.684294 | 1.001207 | |
| beta_[married] | 3934.706128 | 519.311118 | 2968.317116 | 4917.281288 | 0.999965 | |
| beta_ps | 1.231694 | 30.980685 | -56.131323 | 58.727656 | 1.000404 | |
| alpha_outcome | 4327.019576 | 369.926382 | 3614.093676 | 5007.742769 | 1.000996 | |
| sigma | 6565.731355 | 156.117088 | 6284.099885 | 6863.057429 | 0.999980 | |
| 2-stage-model | beta_[treat] | 1588.615623 | 574.307543 | 524.480268 | 2679.197799 | 1.000771 |
| beta_[nodegree] | -460.634709 | 494.618035 | -1390.619855 | 466.314008 | 0.999807 | |
| beta_[married] | 3934.651215 | 520.193798 | 2949.372121 | 4903.329691 | 1.000124 | |
| beta_ps | 0.503108 | 29.986962 | -55.886896 | 56.746119 | 1.000899 | |
| alpha_outcome | 4315.586212 | 368.780215 | 3616.275074 | 5011.505222 | 1.000713 | |
| sigma | 6568.595118 | 151.036246 | 6288.224841 | 6861.433304 | 1.000467 | |
| Simple Regression | beta_[treat] | 1585.697270 | 576.800610 | 519.033031 | 2670.443099 | 1.000652 |
| beta_[nodegree] | -455.116270 | 493.441233 | -1361.728888 | 468.018926 | 1.000346 | |
| beta_[married] | 3929.656708 | 520.079653 | 2946.711539 | 4908.038304 | 1.000760 | |
| alpha_outcome | 4315.300286 | 370.375109 | 3619.773379 | 5006.659307 | 1.001019 | |
| sigma | 6570.553871 | 152.634211 | 6286.608367 | 6861.984880 | 1.001572 |
The model estimates are basically identical. We should expect mirrored propensity score distributions indicating that feedback wasn’t an issue. The outcome model leaned primarily on the covariate profile \(X\) to derive the treatment effect estimate.
compare_propensity_dists(idata_treatment_2s_lalonde, idata_lalonde)
Both estimates of the treatment effects accord well with values in the literature. We can be happy that both models are picking up on the treatment effects reasonably well, despite having a minimalist outcome model. This is because in the Bayesian setting the conditional distribution of \(X, T\) are generally sufficient when there is no unmeasured confounders. Here the model somewhat ignores the propensity score coefficients beta_ps.
NHEFS#
Finally we turn to the NHEFS data as discussed by Hernan in [Hernan, 2024]. This data is known to be have a complex covariate profile for measuring aspects smokers health. We might suspect that there is some unmeasured confounding in this data set that would be hard to pick up on with simple regression controls.
df = cp.load_data("nhefs")
df[["age", "race", "trt", "smokeintensity", "smokeyrs", "outcome"]].head()
| age | race | trt | smokeintensity | smokeyrs | outcome | |
|---|---|---|---|---|---|---|
| 0 | 42 | 1 | 0 | 30 | 29 | -10.093960 |
| 1 | 36 | 0 | 0 | 20 | 24 | 2.604970 |
| 2 | 56 | 1 | 0 | 20 | 26 | 9.414486 |
| 3 | 68 | 1 | 0 | 3 | 53 | 4.990117 |
| 4 | 40 | 0 | 0 | 20 | 19 | 4.989251 |
coords = {
"betas": ["trt", "age", "race", "sex", "smokeintensity", "smokeyrs", "wt71"],
"betas_trt": [
"age",
"race",
"sex",
"smokeintensity",
"smokeyrs",
"wt71",
"active_1",
"active_2",
"education_2",
"education_3",
"education_4",
"education_5",
"exercise_1",
"exercise_2",
"age^2",
"wt71^2",
"smokeintensity^2",
"smokeyrs^2",
],
"obs": range(df.shape[0]),
}
N = df.shape[0]
X_trt = df[
[
"age",
"race",
"sex",
"smokeintensity",
"smokeyrs",
"wt71",
"active_1",
"active_2",
"education_2",
"education_3",
"education_4",
"education_5",
"exercise_1",
"exercise_2",
"age^2",
"wt71^2",
"smokeintensity^2",
"smokeyrs^2",
]
]
X_trt = (X_trt - X_trt.mean(axis=0)) / X_trt.std(axis=0)
# Note the significantly reduced outcome model specification.
X_outcome = df[["trt", "age", "race", "sex", "smokeintensity", "smokeyrs", "wt71"]]
X_outcome = (X_outcome - X_outcome.mean(axis=0)) / X_outcome.std(axis=0)
T_data = df["trt"].values
X_outcome["trt"] = T_data
Y_data = df["outcome"].values
priors = {
"beta_": [0, 3],
"beta_trt": [0, 1],
"alpha_trt": [0, 1],
"alpha_outcome": [0, 5],
"sigma": 10,
"beta_ps": [0, 10],
}
nhefs_model = make_joint_model(X_trt, X_outcome, T_data, Y_data, coords, priors)
with nhefs_model:
idata_nhefs = pm.sample(**sampler_kwargs)
(
idata_treatment_2s_nhefs,
idata_outcome_2s_nhefs,
treatment_model_nhefs,
outcome_model_nhefs,
) = make_2step_model(
X_trt,
X_outcome,
T_data,
Y_data,
coords,
priors=priors,
spline_component=True,
)
reg_model_nhefs, idata_outcome_simple_reg_nhefs = make_reg_model(
X_outcome, Y_data, coords, priors=priors
)
Show code cell output
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_trt_std, beta_std, beta_ps, alpha_trt, alpha_outcome, sigma]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 43 seconds.
Sampling: [alpha_trt, beta_trt_std, t_pred]
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_trt_std, alpha_trt]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 31 seconds.
Sampling: [alpha_outcome, beta_ps, beta_ps_spline, beta_std, like, sigma]
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_std, beta_ps, alpha_outcome, beta_ps_spline, sigma]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 17 seconds.
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_std, alpha_outcome, sigma]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 7 seconds.
forest_plot = azp.plot_forest(
{
"2 Stage": idata_outcome_2s_nhefs,
"1 Stage": idata_nhefs,
"Simple Regression": idata_outcome_simple_reg_nhefs,
},
var_names=["beta_"],
combined=True,
ci_kind="hdi",
ci_probs=(0.5, 0.94),
figure_kwargs={"figsize": (10, 4)},
)
forest_plot.add_legend("model")
forest_axis = forest_plot.get_target(None, {"column": "forest"})
forest_axis.axvline(3.4, label="True Treatment Value", color="k")
forest_axis.legend()
forest_plot.add_title("Comparing Joint and 2 Stage Propensity Score Parameter Fits");
Here we see the model specifications start to come apart slightly, although the effect of feedback seems to primarily influence the impact of age on the outcome rather than the treatment estimate.
Show code cell source
compare_estimate = pd.concat(
{
"1-stage-model": az.summary(
idata_nhefs,
var_names=["alpha_trt", "beta_", "alpha_outcome", "beta_ps"],
ci_kind="hdi",
ci_prob=0.94,
),
"2-stage-model": az.summary(
idata_outcome_2s_nhefs,
var_names=["beta_", "alpha_outcome", "beta_ps"],
ci_kind="hdi",
ci_prob=0.94,
),
"Simple Regression": az.summary(
idata_outcome_simple_reg_nhefs,
var_names=["beta_", "alpha_outcome"],
ci_kind="hdi",
ci_prob=0.94,
),
}
)
compare_estimate[["mean", "sd", "hdi94_lb", "hdi94_ub", "r_hat"]]
| mean | sd | hdi94_lb | hdi94_ub | r_hat | ||
|---|---|---|---|---|---|---|
| 1-stage-model | alpha_trt | -1.146525 | 0.061601 | -1.262861 | -1.032347 | 1.000319 |
| beta_[trt] | 3.479889 | 0.440981 | 2.655539 | 4.307511 | 1.000408 | |
| beta_[age] | -1.023342 | 0.614837 | -2.133300 | 0.192524 | 1.002765 | |
| beta_[race] | -0.318180 | 0.271364 | -0.852243 | 0.178749 | 1.000248 | |
| beta_[sex] | -1.143594 | 0.295892 | -1.716658 | -0.600974 | 1.000646 | |
| beta_[smokeintensity] | -0.427208 | 0.316977 | -1.045601 | 0.138837 | 0.999941 | |
| beta_[smokeyrs] | -0.295940 | 0.537017 | -1.343602 | 0.673263 | 1.002598 | |
| beta_[wt71] | -1.217295 | 0.283161 | -1.723912 | -0.653215 | 1.000259 | |
| alpha_outcome | 5.197875 | 0.998728 | 3.377120 | 7.148959 | 1.001430 | |
| beta_ps | -13.442767 | 3.816180 | -20.893096 | -6.653904 | 1.001511 | |
| 2-stage-model | beta_[trt] | 3.345376 | 0.436564 | 2.527828 | 4.155537 | 0.999864 |
| beta_[age] | -2.027026 | 0.466302 | -2.888303 | -1.124701 | 1.000097 | |
| beta_[race] | 0.033121 | 0.241960 | -0.415187 | 0.481910 | 1.000402 | |
| beta_[sex] | -0.853482 | 0.248358 | -1.321867 | -0.394427 | 1.000099 | |
| beta_[smokeintensity] | 0.090196 | 0.231352 | -0.340429 | 0.527929 | 1.000340 | |
| beta_[smokeyrs] | 0.328529 | 0.431626 | -0.487111 | 1.123210 | 1.000061 | |
| beta_[wt71] | -1.501871 | 0.209419 | -1.899669 | -1.111249 | 1.000142 | |
| alpha_outcome | 1.469497 | 2.822421 | -3.763378 | 6.759324 | 1.001615 | |
| beta_ps | 0.009556 | 10.026250 | -18.980521 | 18.738393 | 0.999986 | |
| Simple Regression | beta_[trt] | 3.270098 | 0.434570 | 2.452095 | 4.090006 | 1.000919 |
| beta_[age] | -2.427394 | 0.386754 | -3.161754 | -1.702797 | 1.000361 | |
| beta_[race] | 0.203665 | 0.193293 | -0.154184 | 0.563943 | 1.000382 | |
| beta_[sex] | -0.684776 | 0.217852 | -1.096770 | -0.278612 | 1.000466 | |
| beta_[smokeintensity] | 0.270480 | 0.199169 | -0.110397 | 0.647811 | 1.001083 | |
| beta_[smokeyrs] | 0.547460 | 0.392719 | -0.204243 | 1.285268 | 1.000299 | |
| beta_[wt71] | -1.543410 | 0.208128 | -1.928755 | -1.144848 | 0.999988 | |
| alpha_outcome | 1.795025 | 0.216185 | 1.385984 | 2.199773 | 1.000747 |
compare_propensity_dists(idata_treatment_2s_nhefs, idata_nhefs)
Two-Stage Outcome Modelling with CausalPy#
Next we show how to achieve these steps with the simpler CausalPy experiment API.
formula = """trt ~ 1 + age + race + sex + smokeintensity + smokeyrs + wt71 + active_1 + active_2 +
education_2 + education_3 + education_4 + education_5 + exercise_1 + exercise_2"""
df_standardised = (df - df.mean(axis=0)) / df.std(axis=0)
df_standardised["trt"] = df["trt"]
df_standardised["outcome"] = df["outcome"]
result = cp.InversePropensityWeighting(
df_standardised,
formula=formula,
outcome_variable="outcome",
weighting_scheme="robust", ## Will be used by plots after estimation if no other scheme is specified.
model=cp.pymc_models.PropensityScore(
sample_kwargs={
"chains": 4,
"tune": 2000,
"draws": 10000,
"target_accept": 0.95,
"random_seed": 18,
"progressbar": False,
"mp_ctx": "spawn",
},
),
).fit()
result
Show code cell output
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [b]
Sampling 4 chains for 2_000 tune and 10_000 draw iterations (8_000 + 40_000 draws total) took 26 seconds.
Sampling: [t_pred]
Sampling: [b, t_pred]
<causalpy.experiments.inverse_propensity_weighting.InversePropensityWeighting at 0x3e2a14590>
Comparing Inverse Propensity Score Weighting and Covariate Adjustment#
The two step procedure doesn’t just apply for regression adjustment methods as we’ve seen here, but can be used to apply inverse weighting techniques too.
result.plot_ate(result.idata);
/Users/benjamv/mambaforge/envs/CausalPy/lib/python3.14/site-packages/IPython/core/events.py:101: UserWarning: Creating legend with loc="best" can be slow with large amounts of data.
/Users/benjamv/mambaforge/envs/CausalPy/lib/python3.14/site-packages/IPython/core/pylabtools.py:158: UserWarning: Creating legend with loc="best" can be slow with large amounts of data.
which can be compared against the two-step regression adjustment here.
idata_outcome_cp, model_outcome_cp = result.model.fit_outcome_model(
result.X_outcome,
result.y,
result.coords,
priors={"b_outcome": [0, 3], "sigma": 10, "beta_ps": [0, 10]},
noncentred=True,
normal_outcome=True,
)
Show code cell output
Sampling: [beta_ps, beta_std, like, sigma]
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_std, beta_ps, sigma]
Sampling 4 chains for 2_000 tune and 10_000 draw iterations (8_000 + 40_000 draws total) took 31 seconds.
Yielding similar, but not identical results.
az.summary(
idata_outcome_cp,
var_names=["beta_"],
ci_kind="hdi",
ci_prob=0.94,
)
| mean | sd | hdi94_lb | hdi94_ub | ess_bulk | ess_tail | r_hat | mcse_mean | mcse_sd | |
|---|---|---|---|---|---|---|---|---|---|
| beta_[Intercept] | 2.09 | 1.35 | -0.45 | 4.6 | 9423 | 13410 | 1.00 | 0.014 | 0.0099 |
| beta_[age] | -2.21 | 0.73 | -3.6 | -0.83 | 10582 | 15565 | 1.00 | 0.0071 | 0.005 |
| beta_[race] | 0.151 | 0.311 | -0.43 | 0.74 | 12213 | 16402 | 1.00 | 0.0028 | 0.002 |
| beta_[sex] | -0.796 | 0.411 | -1.6 | -0.026 | 10948 | 15890 | 1.00 | 0.0039 | 0.0028 |
| beta_[smokeintensity] | 0.179 | 0.336 | -0.45 | 0.81 | 12061 | 17478 | 1.00 | 0.0031 | 0.0022 |
| beta_[smokeyrs] | 0.46 | 0.606 | -0.67 | 1.6 | 11703 | 16382 | 1.00 | 0.0056 | 0.004 |
| beta_[wt71] | -1.51 | 0.214 | -1.9 | -1.1 | 31311 | 28523 | 1.00 | 0.0012 | 0.00085 |
| beta_[active_1] | -0.537 | 0.214 | -0.94 | -0.14 | 28702 | 29521 | 1.00 | 0.0013 | 0.00091 |
| beta_[active_2] | -0.115 | 0.207 | -0.5 | 0.27 | 30634 | 28811 | 1.00 | 0.0012 | 0.00083 |
| beta_[education_2] | 0.414 | 0.253 | -0.063 | 0.89 | 24731 | 27390 | 1.00 | 0.0016 | 0.0011 |
| beta_[education_3] | 0.396 | 0.276 | -0.12 | 0.92 | 22897 | 25511 | 1.00 | 0.0018 | 0.0013 |
| beta_[education_4] | 0.359 | 0.25 | -0.11 | 0.83 | 18949 | 23754 | 1.00 | 0.0018 | 0.0013 |
| beta_[education_5] | 0.028 | 0.302 | -0.54 | 0.6 | 15574 | 22414 | 1.00 | 0.0024 | 0.0017 |
| beta_[exercise_1] | 0.208 | 0.279 | -0.31 | 0.74 | 20521 | 24364 | 1.00 | 0.002 | 0.0014 |
| beta_[exercise_2] | 0.27 | 0.331 | -0.35 | 0.89 | 14525 | 20525 | 1.00 | 0.0027 | 0.0019 |
| beta_[trt] | 3.317 | 0.434 | 2.5 | 4.1 | 43133 | 29709 | 1.00 | 0.0021 | 0.0015 |
posterior_plot = az.plot_dist(
idata_outcome_cp,
var_names=["beta_"],
coords={"outcome_coeffs": ["trt"]},
ci_kind="hdi",
ci_prob=0.94,
figure_kwargs={"figsize": (10, 5)},
)
treatment_effect = idata_outcome_cp.posterior["beta_"].sel(outcome_coeffs="trt")
probability_below_reference = float((treatment_effect < 3.5).mean().values)
treatment_effect_axis = posterior_plot.get_target("beta_", {})
treatment_effect_axis.axvline(
3.5,
color="black",
linestyle="--",
label=f"P(effect < 3.5) = {probability_below_reference:.1%}; P(effect >= 3.5) = {1 - probability_below_reference:.1%}",
)
treatment_effect_axis.legend()
posterior_plot.add_title(
"Treatment Effect Estimated using 2-stage Regression Adjustment Approach with reduced covariates"
);
Conclusion: Modularity as Causal Discipline#
When attempting to estimate treatment effects using Bayesian inference, a natural but risky strategy is to fit a joint model for both the treatment assignment and the outcome. That is, to specify a full model and infer the parameters of both components simultaneously.
However, this joint approach introduces a feedback loop: the outcome \(Y\) can influence the estimation of the treatment mechanism \(P(T | X)\). This violates the original logic of design-based inference, where treatment assignment should be modeled independently of the observed outcomes. This phenomenon is often subtle but can lead to biased treatment effect estimates.
Across several examples, we have shown that fitting a full joint model distorts the treatment effect estimate relative to a two-step (modular) approach.
In other cases, joint and modular approaches yield nearly identical estimates — usually when the outcome is well-identified from covariates alone. With these observations in scope, we recommend that practitioners generally follow a two-step or modular approach. Either two-stage inverse propensity score weighting or regression adjustment with the propensity score as an additional covariate. Both methods are available now in CausalPy.
Framed this way we can see that joint model violates the temporal precedence of the treatment assignment and outcome process. The 2-stage Bayesian procedures ensure that the causal ordering encoded in the actual data generating process is respected in the estimation process. The confounding adjustment achieved with propensity score must occur without access to information about the outcome. A well-specified propensity score model can substantially improve causal estimates (as we’ve seen), especially when the outcome model is weak or mis-specified. This allows for a form of doubly-robust inference as discussed in [Aronow and Miller, 2019] and [Fan et al., 2023]. Propensity scores do not only serve to reduce dimensionality; they formalize the treatment mechanism and encode information that the outcome model might fail to recover. This explains their continued prominence in modern causal inference and usefulness in the Bayesian setting.
We can get nicely formatted tables from our integration with the maketables package.
from maketables import ETable
ETable(result, coef_fmt="b:.3f")
| outcome | |
|---|---|
| (1) | |
| coef | |
| age | 0.548 |
| race | -0.266 |
| sex | -0.248 |
| smokeintensity | -0.314 |
| smokeyrs | -0.334 |
| wt71 | 0.094 |
| active_1 | 0.022 |
| active_2 | 0.046 |
| education_2 | -0.029 |
| education_3 | 0.027 |
| education_4 | 0.027 |
| education_5 | 0.148 |
| exercise_1 | 0.177 |
| exercise_2 | 0.205 |
| Intercept | -1.144 |
| stats | |
| N | 1,566 |
| Format of coefficient cell: Coefficient | |
References#
P Aronow and B Miller. Foundations of Agnostic Statistics. Cambridge University Press, 2019.
Li Fan, Ding Peng, and Mealli Fabrizia. Bayesian causal inference: a critical review. Philosophical Transactions of the Royal Society, 2023.
G. J. Griffith, T. T. Morris, M. J. Tudball, A. Herbert, G. Mancano, L. Pike, G. C. Sharp, J. Sterne, T. M. Palmer, G. Davey Smith, K. Tilling, L. Zuccolo, N. M. Davies, and G. & Hemani. Collider bias undermines our understanding of covid-19 disease risk and severity. Nature communications, 2020.
Miguel A. Hernan. Causal Inference: What If. CRC Press TBC, Boca Raton, 2024.