Heterogeneous treatment effects: the VC term¶
Most causal questions are not "does the treatment help?" but "whom does it help, and by how much?". A TRAM-DAG answers that with a varying-coefficient term, which gives one edge of the DAG a treatment effect that depends on covariates:
$$\text{shift} \;\mathrel{+}=\; \beta(\text{modifiers})\cdot x_t, \qquad \beta(x) = \beta_0 + b_\Theta(x)$$
docs/varying-coefficients.md states the
model, its penalty and its centering; this notebook runs it.
The notebook scores the recovered beta(x) against a known effect
function. It then shows what a plain flexible shift does not give: a
read-out you can plot, and a propensity-centered variant that survives
confounding.
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from tramdag import CS, LS, SI, VC, CausalFlowDAG, ContinuousNode, OrdinalNode
from tramdag.callbacks import EarlyStopping
from tramdag.plots import plot_varying_coef
plt.rcParams["figure.dpi"] = 110
1. A DGP with a known effect function¶
Three covariates, a confounded treatment (X1 and X2 push assignment),
and an outcome whose conditional is exactly a logistic-latent flow:
$$h(y) + g(x) + \beta(x)\,t = U,\qquad U\sim\text{Logistic}(0,1)$$
with $h(y) = 2y$, a nonlinear prognostic part $g(x) = \tfrac12 X_1^2 + X_2 - \tfrac12 X_3$, and the effect function
$$\beta(x) = -1 + 0.8\,X_2 - 0.6\,X_3 .$$
X2 is deliberately both a confounder and an effect modifier. In this case
a naive estimate fails hardest. The simulator writes beta as the same
expression that true_beta reports.
B0, B2, B3 = -1.0, 0.8, -0.6
def true_beta(df):
"""The effect function the model has to recover, on the latent scale."""
return B0 + B2 * df["X2"].to_numpy() + B3 * df["X3"].to_numpy()
def simulate(n, seed):
rng = np.random.default_rng(seed)
x1, x2, x3 = (rng.normal(size=n) for _ in range(3))
t = (rng.logistic(size=n) > -(0.4 * x1 + 0.4 * x2)).astype(float)
g = 0.5 * x1**2 + x2 - 0.5 * x3
beta = B0 + B2 * x2 + B3 * x3
y = (rng.logistic(size=n) - g - beta * t) / 2.0
return pd.DataFrame({"X1": x1, "X2": x2, "X3": x3, "T": t, "Y": y})
train, val, test = simulate(4500, 1), simulate(500, 2), simulate(2000, 3)
print(f"treated fraction: {train['T'].mean():.2f}")
print(
f"true effect ranges from {true_beta(test).min():+.2f} to {true_beta(test).max():+.2f}"
)
treated fraction: 0.48 true effect ranges from -4.41 to +2.29
2. Which covariates modify the effect? (effect_modifier_scan)¶
Before you choose the modifiers, ask the data. Fit the simple all-LS
model, which takes seconds and is deterministic. Then look at the per-row
scores of the treatment coefficient. With a constant effect, these scores
show no pattern along any covariate. The scan measures that pattern per
covariate and ranks the covariates.
def cheap_all_ls():
return {
"X1": ContinuousNode(),
"X2": ContinuousNode(),
"X3": ContinuousNode(),
"T": OrdinalNode(2, LS("X1") + LS("X2")),
"Y": ContinuousNode(LS("X1") + LS("X2") + LS("X3") + LS("T")),
}
screen = CausalFlowDAG(cheap_all_ls(), seed=0)
screen.fit_classical(train)
scan = screen.effect_modifier_scan(train, "Y", t="T")
print(scan)
assert scan["flag"][["X1", "X2", "X3"]].all()
stat p_value crit_5pct flag X2 4.164609 1.722864e-15 1.3581 True X1 4.141297 2.537664e-15 1.3581 True X3 3.161564 4.159703e-09 1.3581 True
X2 and X3 are the true modifiers, and both flag. X1 flags too, but it
does not enter $\beta(x)$. Here $g$ contains $\tfrac12 X_1^2$, which the
simple linear model cannot represent.
docs/scores.md explains why a misspecified prognostic
part flags too.
To show this, generate the data again with a linear prognostic $X_1$ and
change nothing else. X1 then does not flag, because the prognostic part is
linear in X1 and the simple model has the right form.
def simulate_linear_x1(n, seed):
rng = np.random.default_rng(seed)
x1, x2, x3 = (rng.normal(size=n) for _ in range(3))
t = (rng.logistic(size=n) > -(0.4 * x1 + 0.4 * x2)).astype(float)
g = 0.5 * x1 + x2 - 0.5 * x3
y = (rng.logistic(size=n) - g - (B0 + B2 * x2 + B3 * x3) * t) / 2.0
return pd.DataFrame({"X1": x1, "X2": x2, "X3": x3, "T": t, "Y": y})
well_specified = simulate_linear_x1(4500, 1)
screen2 = CausalFlowDAG(cheap_all_ls(), seed=0)
screen2.fit_classical(well_specified)
scan2 = screen2.effect_modifier_scan(well_specified, "Y", t="T")
print(scan2)
assert scan2["flag"][["X2", "X3"]].all()
assert not scan2["flag"]["X1"]
stat p_value crit_5pct flag X2 4.608240 7.174333e-19 1.3581 True X3 3.474392 6.545501e-11 1.3581 True X1 0.553054 9.197166e-01 1.3581 False
The figure draws the scaled running sum $B_j$ of the scan for each candidate
(docs/scores.md explains it). The left panel uses the
quadratic X1, the right panel the linear X1. stat is the largest
distance of the path from zero, the dot on each line. If a path leaves the
grey band, the scan flags its covariate.
CANDIDATES = {"X1": "#2a78d6", "X2": "#eb6834", "X3": "#1baf7a"}
def plot_cusum(ax, flow, df, title):
"""Draw each candidate's scaled running sum of the T scores, as the scan forms it."""
psi = flow.scores(df, "Y")["T[1]"].to_numpy()
n = len(psi)
scan = flow.effect_modifier_scan(df, "Y", t="T")
crit = scan["crit_5pct"].iloc[0]
ax.axhspan(
-crit, crit, color="0.93", lw=0, label=f"$|B_j| < {crit:.3f}$: no flag at 5%"
)
for c in (-crit, crit):
ax.axhline(c, color="0.55", lw=1, ls="--")
ax.axhline(0, color="0.75", lw=0.8)
frac = np.arange(1, n + 1) / n
for name, color in CANDIDATES.items():
b = np.cumsum(psi[np.argsort(df[name].to_numpy(), kind="stable")])
b /= psi.std() * np.sqrt(n)
j = np.abs(b).argmax()
ax.plot(frac, b, color=color, lw=2, label=name)
ax.plot(frac[j], b[j], "o", ms=8, color=color, mec="white", mew=2)
ax.annotate(
f"{name}: stat {scan.loc[name, 'stat']:.2f}",
(frac[j], b[j]),
xytext=(0, 10 if b[j] > 0 else -16),
textcoords="offset points",
ha="center",
color="0.2",
)
ax.set(
title=title, xlabel="fraction of rows, sorted by the candidate", ylim=(-6, 5)
)
fig, axes = plt.subplots(1, 2, figsize=(11, 4), sharey=True)
plot_cusum(axes[0], screen, train, "X1 quadratic in the outcome: X1 flags too")
plot_cusum(
axes[1], screen2, well_specified, "X1 linear in the outcome: only X2, X3 flag"
)
axes[0].set_ylabel("scaled running sum $B_j$ of the T scores")
fig.legend(
*axes[1].get_legend_handles_labels(), loc="lower center", ncol=4, frameon=False
)
fig.tight_layout(rect=(0, 0.07, 1, 1))
plt.show()
3. The spec: prognostic part and effect head are separate¶
CS("X1", "X2", "X3") absorbs the prognostic signal, as flexibly as
necessary. VC("X2", "X3", t="T") carries the effect. X2 and X3 appear
twice on purpose: prognostically through the shift, and as modifiers through
the head. Only the treatment named by t= owns its edge.
spec = {
"X1": ContinuousNode(),
"X2": ContinuousNode(),
"X3": ContinuousNode(),
"T": OrdinalNode(2, LS("X1") + LS("X2")),
"Y": ContinuousNode(CS("X1", "X2", "X3") + VC("X2", "X3", t="T")),
}
flow = CausalFlowDAG(spec, seed=0)
flow.fit(
train,
epochs=300,
learning_rate=1e-2,
batch_size=512,
validation_data=val,
seed=0,
callbacks=EarlyStopping(),
)
print(flow.to_matrix())
X1 X2 X3 T Y X1 LS CS['X1', 'X2', 'X3'] X2 LS CS['X1', 'X2', 'X3']+VCm X3 CS['X1', 'X2', 'X3']+VCm T VC Y
In the meta-adjacency view (docs/interpretation.md)
VC marks the treatment edge and VCm the modifiers, so X2 visibly enters
Y twice.
4. Reading the effect out¶
varying_coef gives $\beta(x)$ per row. It is deterministic and does not
look at the outcome. It is a property of the fitted model, not of the rows'
Y values.
beta_hat = flow.varying_coef(test, "Y")
beta_true = true_beta(test)
corr = float(np.corrcoef(beta_hat, beta_true)[0, 1])
beta0 = flow.nodes["Y"].shifts["T"].beta0.item()
print(f"corr(beta_hat, beta_true) = {corr:.3f}")
print(f"beta0 = {beta0:+.3f} (true constant part {B0:+.1f})")
print(f"mean |error| = {np.abs(beta_hat - beta_true).mean():.3f}")
assert corr > 0.95, f"the fitted beta(x) correlates only {corr:.3f} with the truth"
fig, axes = plt.subplots(1, 2, figsize=(10, 3.8))
axes[0].scatter(beta_true, beta_hat, s=6, alpha=0.35)
lims = [min(beta_true.min(), beta_hat.min()), max(beta_true.max(), beta_hat.max())]
axes[0].plot(lims, lims, "k--", lw=1)
axes[0].set_xlabel(r"true $\beta(x)$")
axes[0].set_ylabel(r"fitted $\hat\beta(x)$")
axes[0].set_title(f"recovery, corr = {corr:.3f}")
plot_varying_coef(flow, test, "Y", by="X2", ax=axes[1])
x3_med = test["X3"].median()
grid = np.linspace(test["X2"].min(), test["X2"].max(), 2)
axes[1].plot(grid, B0 + B2 * grid + B3 * x3_med, "k--", lw=1)
along_x2 = test.assign(X3=x3_med)
gap_x2 = np.abs(flow.varying_coef(along_x2, "Y") - true_beta(along_x2)).max()
print(f"X3 at its median: max |beta_hat - beta_true| along X2 = {gap_x2:.3f}")
assert gap_x2 < 0.75, f"the fitted beta(X2) leaves the true line by {gap_x2:.3f}"
axes[1].legend(["fitted", r"$\beta = 0$", "true"])
axes[1].set_title(r"$\beta$ along $X_2$, $X_3$ at its median")
fig.tight_layout()
plt.show()
corr(beta_hat, beta_true) = 0.997 beta0 = -0.917 (true constant part -1.0) mean |error| = 0.097 X3 at its median: max |beta_hat - beta_true| along X2 = 0.321
The right panel holds X3 at its median, because the effect also depends on
X3. Along X2, the fitted line follows the true line.
5. The read-out is an identity, not a summary¶
For a binary treatment, $\beta(x)$ equals the difference of the abducted latents between the two arms, with the outcome held fixed. The implementation must satisfy this definition to float precision:
u1 = flow.abduct(test.assign(T=1.0), seed=0)["Y"].to_numpy()
u0 = flow.abduct(test.assign(T=0.0), seed=0)["Y"].to_numpy()
gap = np.abs(beta_hat - (u1 - u0)).max()
print(f"max |beta(x) - (u(T=1) - u(T=0))| = {gap:.2e}")
assert gap < 1e-5, f"the read-out differs from the latent difference by {gap:.2e}"
max |beta(x) - (u(T=1) - u(T=0))| = 4.77e-07
6. Confounding: why propensity= exists¶
Everything above had a prognostic part flexible enough to absorb $g$. Real models are misspecified, and when the misfit correlates with treatment assignment, the effect head absorbs the confounding instead of the effect.
The configuration where this bites hardest is deliberately simple: one
covariate, a strong propensity $e(x)=\sigma(2x)$, a constant true effect
$\tau = -1$, and a quadratic prognostic part fitted with a linear term.
propensity="ps" replaces $t$ by $t - \hat e(x)$ with out-of-fold propensities
from the training-frame column ps.
docs/varying-coefficients.md says why
they must be out of fold.
TAU = -1.0
def confounded(n, seed):
rng = np.random.default_rng(seed)
x = rng.normal(size=n)
t = (rng.logistic(size=n) > -2.0 * x).astype(float)
y = (rng.logistic(size=n) - 1.2 * x**2 - TAU * t) / 2.0
return pd.DataFrame({"X": x, "T": t, "Y": y})
c_train, c_val, c_test = confounded(5400, 0), confounded(600, 7), confounded(3000, 1000)
Stage 1 of the centered design is the caller's job: the cross-fitted $P(T=1 \mid X)$. A treatment model that never saw a fold predicts that fold, as double machine learning requires.
fold_id = np.random.default_rng(0).permutation(len(c_train)) % 5
e_oof = np.empty(len(c_train))
for j in range(5):
t_spec = {
"X": ContinuousNode(SI(transform="affine")),
"T": OrdinalNode(2, LS("X")),
}
proxy = CausalFlowDAG(t_spec, seed=0)
proxy.fit_classical(c_train.iloc[fold_id != j][["X", "T"]])
e_oof[fold_id == j] = proxy.pmf(c_train.iloc[fold_id == j], "T")[:, 1]
mae = {}
for propensity in (None, "ps"):
spec_c = {
"X": ContinuousNode(SI(transform="affine")),
"T": OrdinalNode(2, LS("X")),
"Y": ContinuousNode(LS("X") + VC("X", propensity=propensity, t="T")),
}
fc = CausalFlowDAG(spec_c, seed=0)
fc.fit(
c_train.assign(ps=e_oof) if propensity else c_train,
epochs=250,
learning_rate=1e-2,
batch_size=512,
validation_data=c_val,
seed=0,
callbacks=EarlyStopping(),
)
b = fc.varying_coef(c_test, "Y")
centered = propensity is not None
mae[centered] = float(np.abs(b - TAU).mean())
print(
f"centered={centered!s:5s} mean |beta - tau| = {mae[centered]:.3f}"
f" mean beta = {b.mean():+.3f} (true tau {TAU:+.1f})"
)
ratio = mae[False] / mae[True]
print(f"centering cuts the bias by a factor of {ratio:.1f}")
assert ratio >= 2, f"centering should at least halve the bias: {ratio:.2f}"
centered=False mean |beta - tau| = 1.303 mean beta = +0.044 (true tau -1.0)
centered=True mean |beta - tau| = 0.279 mean beta = -0.722 (true tau -1.0) centering cuts the bias by a factor of 4.7
Without centering the confounding swallows most of the effect; with centering
the estimate returns near the truth. The printed factor is the measured bias
reduction, and tests/test_vc_centered.py requires at least two.
What to take away¶
| you want | use | read it with |
|---|---|---|
| one interpretable effect | LS("T") |
ls_coefficients() |
| an effect that varies with covariates | VC(*modifiers, t="T") |
varying_coef(df, node) |
| the same under confounding + a misspecified prognostic part | VC(..., propensity="ps") |
the same read-out |
| a shortlist of modifiers before you commit | effect_modifier_scan on a simple all-LS fit |
its flag column, read as screening |
Why not a joint CS over treatment and modifiers, and what the penalty and
the centering mean: docs/varying-coefficients.md.
The scan: docs/scores.md.