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:

\[ P(Y, T, X) \]

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

\[ e(X) = P(T | X)\]

and our outcome model

\[\begin{split} \begin{aligned} P(Y \mid T, X) &= \frac{P(Y, T, X)}{P(T, X)} \cdot P(T, X) \\ &= \frac{P(Y \mid T, X) \cdot P(T \mid X) \cdot P(X)}{P(T \mid X) \cdot P(X)} \cdot P(T, X) \\ &= \frac{P(Y \mid T, X) \cdot e(X) \cdot P(X)}{e(X) \cdot P(X)} \cdot P(T, X) \end{aligned} \end{split}\]

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 umeasured 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 Model
- Specifying 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)

../_images/e37a1fca3731994fe876ec6e10598bbeac9322de02e9af9f546db9d1b7418861.svg

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)

../_images/155746e6fab0d426e3656fc48d590fb7aeb60950d003941920916bd175236237.svg

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

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");

../_images/673f33947f13176390258516a289ef8c54c4a276c871105523cb590450b4d6ce.png

		mean	sd	hdi94_lb	hdi94_ub	r_hat
1-stage-model	alpha_trt	0.0002	0.0216	-0.041	0.04	1.00
	beta_[trt]	2.314	0.056	2.2	2.4	1.00
	beta_[x1]	1.296	0.068	1.2	1.4	1.00
	beta_ps	-15.3	0.297	-16	-15	1.00
	alpha_outcome	8.51	0.162	8.2	8.8	1.00
2-stage-model	beta_[trt]	2.07	0.189	1.7	2.4	1.00
	beta_[x1]	1.15	0.0415	1.1	1.2	1.00
	alpha_outcome	3.071	0.061	3	3.2	1.00
	beta_ps	-4.179	0.212	-4.6	-3.8	1.00
2-stage-model_spline	beta_[trt]	1.981	0.18	1.6	2.3	1.00
	beta_[x1]	1.1726	0.0395	1.1	1.2	1.00
	alpha_outcome	0.93	0.294	0.37	1.5	1.00
	beta_ps	0.02	1	-1.8	1.9	1.00
Simple Regression	beta_[trt]	-1.308	0.085	-1.5	-1.1	1.00
	beta_[x1]	1.107	0.0432	1	1.2	1.00
	alpha_outcome	2.665	0.06	2.6	2.8	1.00

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)

../_images/f92189419d1ac75edbe5a189f13e0af480d1cf1b6bd12b644e524af4ec50753f.svg

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)

../_images/2d2104e34aad127f4b3d4a3ed8ccb7ee33b40ed55b97ba0ec4cdf187d14e6cbe.png

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");

../_images/fefb38dba7c577c431233a5efb8c53bc97c6f3b76601903d9251335f0d8dbdad.png

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.

\[ \mu = \alpha + X\beta + f(p) \]

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]

/opt/anaconda3/envs/CausalPy-pymc6-1047-f/lib/python3.14/site-packages/pytensor/tensor/type.py:684: RuntimeWarning: invalid value encountered in divide
  f"numpy allclose failed for abs_err {np.max(abs(a - b))} and rel_err {np.max(abs(a - b) / (abs(a) + abs(b)))}"
Sampling: [like]

For sake of illustration we calculate the mean value for \(Y(1)\)

idata_trt["posterior_predictive"]["like"].mean().item()

3.1563710546578654

The mean value for \(Y(0)\)

idata_ntrt["posterior_predictive"]["like"].mean().item()

0.8425251534492368

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.3138459012086283

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
)

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");

/opt/anaconda3/envs/CausalPy-pymc6-1047-f/lib/python3.14/site-packages/arviz_plots/plot_collection.py:56: 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.
  data = xr.concat(ds_list, dim="model").assign_coords(model=list(data))

../_images/e2e54f3be67b58f0f325490cb076229b48b49eb7f709675f0fde1b47a5b39b60.png

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 approximtely recover the reported -10 treatment estimate with a large degree of uncertainty. The modular 2 stage estimates pulls away from the joint model estimate.

		mean	sd	hdi94_lb	hdi94_ub	r_hat
1-stage-model	alpha_trt	-0.259	0.053	-0.36	-0.16	1.00
	beta_[trt]	-9.937	0.256	-10	-9.5	1.00
	beta_[income]	-10.097	0.218	-11	-9.7	1.00
	beta_ps	-16.24	0.65	-17	-15	1.00
	alpha_outcome	46.61	0.4	46	47	1.00
2-stage-model	beta_[trt]	-10.088	0.322	-11	-9.5	1.00
	beta_[income]	-10.757	0.195	-11	-10	1.00
	alpha_outcome	41.53	0.414	41	42	1.00
	beta_ps	-5.46	0.94	-7.2	-3.7	1.00
Simple Regression	beta_[trt]	-8.534	0.338	-9.2	-7.9	1.00
	beta_[income]	-11.515	0.169	-12	-11	1.00
	alpha_outcome	37.832	0.216	37	38	1.00

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.25 , 0.311, 0.37 , 0.43 , 0.501, 0.571, 0.643, 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)

../_images/446de9f23fb59641ea8362e228c3fcc63673e80168e5336d2a1bd14f258d099d.png

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	re78	hispan	white
0	NSW1	1	37	11	black	1	1	9930.0460	False	False
1	NSW2	1	22	9	hispan	0	1	3595.8940	True	False
2	NSW3	1	30	12	black	0	0	24909.4500	False	False
3	NSW4	1	27	11	black	0	1	7506.1460	False	False
4	NSW5	1	33	8	black	0	1	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
)

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");

/opt/anaconda3/envs/CausalPy-pymc6-1047-f/lib/python3.14/site-packages/arviz_plots/plot_collection.py:56: 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.
  data = xr.concat(ds_list, dim="model").assign_coords(model=list(data))

../_images/22c315321a8348f9e579d9aa195a74b8582caa748e891196bd26372e74bdbb5d.png

		mean	sd	hdi94_lb	hdi94_ub	r_hat
1-stage-model	alpha_trt	-1.359	0.137	-1.6	-1.1	1.00
	beta_[treat]	2e+03	589	510	2700	1.00
	beta_[nodegree]	-5e+02	491	-1400	460	1.00
	beta_[married]	4e+03	5e+02	3000	4900	1.00
	beta_ps	1.1	30	-55	56	1.00
	alpha_outcome	4e+03	385	3600	5000	1.00
	sigma	6571	153	6300	6900	1.00
2-stage-model	beta_[treat]	2e+03	577	500	2600	1.00
	beta_[nodegree]	-463	504	-1400	480	1.00
	beta_[married]	3926	516	3000	4900	1.00
	beta_ps	-0	29.7	-56	55	1.00
	alpha_outcome	4322	373	3600	5000	1.00
	sigma	6569	153	6300	6900	1.00
Simple Regression	beta_[treat]	2e+03	6e+02	520	2700	1.00
	beta_[nodegree]	-5e+02	493	-1400	480	1.00
	beta_[married]	4e+03	520	3000	4900	1.00
	alpha_outcome	4e+03	368	3600	5000	1.00
	sigma	6569	153	6300	6900	1.00

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)

../_images/31a1fa8a298f8e5db3bf877db9cb60bfb2efde8414061db1e1a5714cf4c8980d.png

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	smokeintensity	smokeyrs	outcome
0	42	1	30	29	-10.093960
1	36	0	20	24	2.604970
2	56	1	20	26	9.414486
3	68	1	3	53	4.990117
4	40	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
)

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");

../_images/d5bd3f11cf8ce6e2e251b330554eaaa1e49e3d95cbefa7b1ea9afa3c7aa682dd.png

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.

		mean	sd	hdi94_lb	hdi94_ub	r_hat
1-stage-model	alpha_trt	-1.146	0.0613	-1.3	-1	1.00
	beta_[trt]	3.492	0.44	2.7	4.3	1.00
	beta_[age]	-1	0.63	-2.1	0.25	1.00
	beta_[race]	-0.329	0.275	-0.86	0.17	1.00
	beta_[sex]	-1.151	0.304	-1.7	-0.6	1.00
	beta_[smokeintensity]	-0.44	0.314	-1.1	0.13	1.00
	beta_[smokeyrs]	-0.31	0.55	-1.4	0.68	1.00
	beta_[wt71]	-1.216	0.285	-1.7	-0.66	1.00
	alpha_outcome	5.24	1.01	3.4	7.2	1.00
	beta_ps	-13.6	3.9	-21	-6.6	1.00
2-stage-model	beta_[trt]	3.389	0.438	2.6	4.2	1.00
	beta_[age]	-1.76	0.48	-2.6	-0.85	1.00
	beta_[race]	0.037	0.217	-0.37	0.45	1.00
	beta_[sex]	-0.808	0.229	-1.2	-0.38	1.00
	beta_[smokeintensity]	0.04	0.227	-0.39	0.47	1.00
	beta_[smokeyrs]	0.15	0.424	-0.65	0.94	1.00
	beta_[wt71]	-1.478	0.211	-1.9	-1.1	1.00
	alpha_outcome	0.1	3	-5.5	5.7	1.00
	beta_ps	-0.2	10	-19	18	1.00
Simple Regression	beta_[trt]	3.27	0.435	2.5	4.1	1.00
	beta_[age]	-2.427	0.387	-3.2	-1.7	1.00
	beta_[race]	0.204	0.193	-0.15	0.56	1.00
	beta_[sex]	-0.685	0.218	-1.1	-0.28	1.00
	beta_[smokeintensity]	0.27	0.199	-0.11	0.65	1.00
	beta_[smokeyrs]	0.55	0.393	-0.2	1.3	1.00
	beta_[wt71]	-1.543	0.208	-1.9	-1.1	1.00
	alpha_outcome	1.795	0.216	1.4	2.2	1.00

compare_propensity_dists(idata_treatment_2s_nhefs, idata_nhefs)

../_images/6cf9de388c2b9a75add7a724ad0bbdb37aadac2f5af8bca672be28d4cf90de27.png

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

result

Comparing Inverse Propensity Score Weighting and Covariate Adjustment#

The two step procedure doesn’t jusst 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);

/opt/anaconda3/envs/CausalPy-pymc6-1047-f/lib/python3.14/site-packages/IPython/core/events.py:100: UserWarning: Creating legend with loc="best" can be slow with large amounts of data.
  func(*args, **kwargs)

/opt/anaconda3/envs/CausalPy-pymc6-1047-f/lib/python3.14/site-packages/IPython/core/pylabtools.py:170: UserWarning: Creating legend with loc="best" can be slow with large amounts of data.
  fig.canvas.print_figure(bytes_io, **kw)

../_images/92fd4dfb8ef2f420e410e81fd422ef775f5b51fed768309c9d2586563f8820ee.png

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

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]	3.22	1.37	0.66	5.8	9045	13076	1.00	0.014	0.01
beta_[age]	-1.959	0.535	-3	-0.95	12486	17710	1.00	0.0048	0.0034
beta_[race]	-0.041	0.305	-0.61	0.53	12572	18058	1.00	0.0027	0.0019
beta_[sex]	-0.98	0.337	-1.6	-0.35	12586	18319	1.00	0.003	0.0021
beta_[smokeintensity]	-0.139	0.411	-0.92	0.63	10449	14785	1.00	0.004	0.0029
beta_[smokeyrs]	0.362	0.439	-0.47	1.2	17205	23253	1.00	0.0033	0.0024
beta_[wt71]	-1.267	0.314	-1.9	-0.68	13068	18908	1.00	0.0027	0.002
beta_[active_1]	-0.51	0.21	-0.9	-0.11	28891	26837	1.00	0.0012	0.00088
beta_[active_2]	0.016	0.243	-0.44	0.48	18295	23616	1.00	0.0018	0.0013
beta_[education_2]	0.455	0.253	-0.022	0.93	19747	24896	1.00	0.0018	0.0013
beta_[education_3]	0.395	0.276	-0.13	0.91	19001	23531	1.00	0.002	0.0014
beta_[education_4]	0.409	0.225	-0.016	0.83	23444	26629	1.00	0.0015	0.001
beta_[education_5]	0.176	0.298	-0.38	0.74	13673	19340	1.00	0.0025	0.0018
beta_[exercise_1]	0.419	0.342	-0.23	1.1	13804	19746	1.00	0.0029	0.0021
beta_[exercise_2]	0.413	0.325	-0.2	1	15533	20901	1.00	0.0026	0.0018
beta_[trt]	3.332	0.436	2.5	4.2	37206	29968	1.00	0.0023	0.0016

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"
);

../_images/e6d8179915b1fbfab5e8b49e8f8c4d739346f79b6880c187faaf7cfe960d58e9.png

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#

[1]

P Aronow and B Miller. Foundations of Agnostic Statistics. Cambridge University Press, 2019.

[2] (1,2)

Li Fan, Ding Peng, and Mealli Fabrizia. Bayesian causal inference: a critical review. Philosophical Transactions of the Royal Society, 2023.

[3]

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.

[4]

Miguel A. Hernan. Causal Inference: What If. CRC Press TBC, Boca Raton, 2024.