Estimate an Average Treatment Effect

Use a partially linear GAM to estimate a causal average treatment effect while nonparametrically controlling for confounders.

Estimating a treatment effect from observational data requires removing the influence of confounders, which are variables that affect both treatment assignment and the outcome. A partially linear causal GAM handles this by modeling the confounders nonparametrically while estimating a single linear treatment coefficient, which is the average treatment effect (ATE).

The partially linear approach is robust to misspecification of the confounder relationships because it lets smooth functions absorb complex nonlinear patterns, leaving the treatment coefficient cleanly identified.

Simulate data

We generate a dataset where treatment assignment depends on age and experience (confounding), and the true ATE is 2.5 wage units.

import numpy as np
import whittaker as wk

# Generate confounders and treatment assignment
rng = np.random.default_rng(23)
n = 500
age = rng.uniform(25, 55, n)
experience = rng.uniform(0, 20, n)
p_treat = 1 / (1 + np.exp(-(0.05 * age + 0.1 * experience - 3)))
training = rng.binomial(1, p_treat).astype(float)

# Simulate wages with true ATE of 2.5
wage = 10 + 0.3 * age + 0.5 * experience + 2.5 * training + rng.normal(0, 2, n)
data = {"wage": wage, "training": training, "age": age, "experience": experience}

Fit

CausalGAM with method="partially_linear" automatically builds a smooth term for each confounder and estimates a linear treatment coefficient. No formula argument is needed.

model = wk.CausalGAM(
    outcome="wage",
    treatment="training",
    confounders=["age", "experience"],
    method="partially_linear",
).fit(data, seed=23)

Extract the treatment effect

Calling .treatment_effect() returns a TreatmentEffect object with the ATE and its uncertainty.

te = model.treatment_effect()
te.ate
2.289889555724743

The 95% confidence interval shows the range of plausible values for the true treatment effect.

(te.ci_lower, te.ci_upper)
(2.121321339308654, 2.4584577721408323)

The p-value tests the null hypothesis that the treatment has no effect.

te.p_value
0.0

Interpret

model.summary()
'CausalGAM summary\n============================================================\nOutcome:     wage\nTreatment:   training\nConfounders: age, experience\nMethod:      partially_linear\nN folds:     5\nN obs:       500\n\nTreatment effect:\n  ATE = 2.2899 (SE = 0.0860)\n  95% CI: [2.1213, 2.4585]\n  p-value: 0.0000'

The ATE estimate near 2.5 recovers the true simulated effect. The confidence interval and p-value together confirm statistical significance: after nonparametrically removing the confounding influence of age and experience, workers who received training earn roughly 2.5 more wage units on average. A naive comparison ignoring confounders would overestimate this effect because older and more experienced workers are both more likely to be treated and to earn higher wages.