diff --git a/CHANGELOG.md b/CHANGELOG.md index 78b4cae..2ee2867 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,10 @@ per-individual ITEs (abduction + `do`, including the mediated `X5 → X6 → Y` path and the treatment heterogeneity) on this DGP, validated against the known truth (r ≈ 0.99, ATE within ~0.01) and contrasted with the biased naive contrast. +- `experiments/ite_train_size.py` — ITE/ATE recovery vs training-set size + (n = 500..20000) for that all-CI model, evaluated on a fixed 5k test set, with + the **train-vs-test NLL gap** as an overfitting diagnostic (the gap, and the ITE + error, shrink systematically as n grows). ## 0.3.0 (2026-06-19) diff --git a/experiments/ite_train_size.py b/experiments/ite_train_size.py new file mode 100644 index 0000000..8922b88 --- /dev/null +++ b/experiments/ite_train_size.py @@ -0,0 +1,226 @@ +"""ITE recovery vs training-set size for an all-CI TRAM-DAG. + +Companion experiment to ``notebooks/ite_observational.py``: fit the same all-CI +S-learner TRAM-DAG on increasingly large observational samples +(n = 500 .. 20000), evaluate every model on the **same fixed 5000-row test set**, +and report + + * ATE recovery (predicted vs the known true ATE), + * ITE recovery (correlation and MAE vs the per-individual true ITE), and + * the **train vs test NLL gap** — the overfitting diagnostic: a flexible + all-CI model trained on little data fits its train set better than the test + set; the gap should shrink as n grows. + +Averaged over a few seeds (data draw + weight init). Results cached to JSON; +two figures written to ``results/ite_train_size/plots/``. + +Run from ``experiments/``:: + + uv run python ite_train_size.py + uv run python ite_train_size.py --seeds 3 --sizes 500 1000 2000 5000 10000 20000 +""" + +from __future__ import annotations + +import argparse +import time + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 + +from paper_common import results_dir, save_json # noqa: E402 + +from tramdag import CausalFlowDAG, ContinuousNode, I, OrdinalNode # noqa: E402 +from tramdag.simulations import ITEObservational # noqa: E402 +from tramdag.simulations.ite_observational import COLUMNS # noqa: E402 + +SIZES = [500, 1000, 2000, 5000, 10000, 20000] +TEST_N = 5000 +EPOCHS = 600 + + +def make_spec() -> dict: + """The notebook's all-CI S-learner spec (one joint CI per child node).""" + return {"X1": ContinuousNode(), "X2": ContinuousNode(), "X3": ContinuousNode(), + "Tr": OrdinalNode(levels=2, terms=[I("X1", "X2")]), + "X5": ContinuousNode(terms=[I("Tr")]), + "X6": ContinuousNode(terms=[I("X5")]), + "Y": ContinuousNode(terms=[I("Tr", "X1", "X2", "X3", "X5", "X6")])} + + +def evaluate(flow, test_obs, ite_true, true_ate) -> dict: + """Per-individual ITE on the test set via abduction + do(Tr=0/1).""" + u = flow.abduct(test_obs, seed=0) + y1 = flow.sample(do={"Tr": 1.0}, u=u)["Y"].to_numpy() + y0 = flow.sample(do={"Tr": 0.0}, u=u)["Y"].to_numpy() + ite = y1 - y0 + return {"ite_pred": ite, + "ate_pred": float(ite.mean()), + "ate_err": float(ite.mean() - true_ate), + "corr": float(np.corrcoef(ite, ite_true)[0, 1]), + "mae": float(np.abs(ite - ite_true).mean())} + + +def run(sizes: list[int], n_seeds: int, scenario: int) -> dict: + gen = ITEObservational(seed=123, scenario=scenario) + true_ate = gen.true_ate(mc_n=400_000)["ate_true"] + + # one fixed test set, shared by every model + test = gen.with_truth(TEST_N, seed_offset=999) + test_obs = test[list(COLUMNS)] + ite_true = test["ITE_true"].to_numpy() + + rows: list[dict] = [] + scatter: dict[int, np.ndarray] = {} # seed-0 ITE predictions, for the grid + for n in sizes: + for s in range(n_seeds): + train = gen.observational(n, seed_offset=10 + s) + t0 = time.perf_counter() + flow = CausalFlowDAG(make_spec(), seed=s) + flow.fit(train, epochs=EPOCHS, learning_rate=1e-2, schedule="plateau", + plateau_patience=25, verbose=0) + m = evaluate(flow, test_obs, ite_true, true_ate) + row = {"n": n, "seed": s, "secs": time.perf_counter() - t0, + "train_nll": float(sum(flow.nll(train).values())), + "test_nll": float(sum(flow.nll(test_obs).values())), + **{k: v for k, v in m.items() if k != "ite_pred"}} + rows.append(row) + if s == 0: + scatter[n] = m["ite_pred"] + print(f"n={n:6d} seed={s} ATE={row['ate_pred']:+.3f} " + f"(err {row['ate_err']:+.3f}) r={row['corr']:.3f} " + f"MAE={row['mae']:.3f} NLL train/test " + f"{row['train_nll']:.3f}/{row['test_nll']:.3f} " + f"[{row['secs']:.0f}s]") + return {"true_ate": true_ate, "scenario": scenario, "sizes": sizes, + "n_seeds": n_seeds, "test_n": TEST_N, "epochs": EPOCHS, + "rows": rows, "ite_true": ite_true, "scatter": scatter} + + +def _per_seed(rows, sizes, key): + """List (per size) of the per-seed values — for scatter overlays.""" + return [np.array([r[key] for r in rows if r["n"] == n]) for n in sizes] + + +def _median(rows, sizes, key): + """Median over seeds per size (robust to the small-n instability blow-up).""" + return np.array([np.median([r[key] for r in rows if r["n"] == n]) + for n in sizes]) + + +def plot_scatter_grid(res, path) -> None: + sizes, scatter, ite_true = res["sizes"], res["scatter"], res["ite_true"] + rows_by = {(r["n"], r["seed"]): r for r in res["rows"]} + lim = [ite_true.min(), ite_true.max()] + ncol = 3 + nrow = int(np.ceil(len(sizes) / ncol)) + fig, axes = plt.subplots(nrow, ncol, figsize=(3.4 * ncol, 3.2 * nrow), + sharex=True, sharey=True) + for ax, n in zip(axes.ravel(), sizes): + ip = scatter[n] + r = rows_by[(n, 0)] + ax.plot(lim, lim, color="0.6", lw=1, ls="--") + ax.scatter(ite_true, ip, s=5, alpha=0.25, color="#1b9e77") + ax.set_title(f"n_train = {n}\nr={r['corr']:.3f} ATE err={r['ate_err']:+.3f}", + fontsize=9) + ax.set_xlabel("true ITE") + ax.set_ylabel("predicted ITE") + for ax in axes.ravel()[len(sizes):]: + ax.set_visible(False) + fig.suptitle("ITE recovery on the fixed 5k test set vs training size " + "(seed 0)", fontsize=11) + fig.tight_layout() + fig.savefig(path, dpi=100, bbox_inches="tight") + plt.close(fig) + + +def _scatter_seeds(ax, sizes, per_seed, color, ylim): + """Overlay per-seed points; those outside ylim are clipped to the top edge + so an off-scale blow-up stays visible (as a clipped marker) without + rescaling the axis.""" + n_out = 0 + for n, vals in zip(sizes, per_seed): + clipped = np.clip(vals, *ylim) + n_out += int((vals > ylim[1]).sum() + (vals < ylim[0]).sum()) + ax.scatter([n] * len(vals), clipped, s=18, color=color, + alpha=0.35, zorder=3) + return n_out + + +def plot_curves(res, path) -> None: + sizes, rows = res["sizes"], res["rows"] + true_ate = res["true_ate"] + # robust headline = median over seeds (the n=500 instability can blow a + # single seed's ATE/MAE up by ~1e7; mean would be meaningless there) + ate_med, corr_med, mae_med = (_median(rows, sizes, k) + for k in ("ate_pred", "corr", "mae")) + trn_med, tst_med = _median(rows, sizes, "train_nll"), _median(rows, sizes, "test_nll") + + fig, ax = plt.subplots(1, 3, figsize=(13, 4)) + # (1) ATE recovery — y-range pinned near the true ATE; blow-ups clip to edge + ate_ylim = (true_ate - 0.35, true_ate + 0.35) + ax[0].axhline(true_ate, color="0.5", ls="--", label="true ATE") + n_out = _scatter_seeds(ax[0], sizes, _per_seed(rows, sizes, "ate_pred"), + "#d95f02", ate_ylim) + ax[0].plot(sizes, np.clip(ate_med, *ate_ylim), marker="o", color="#d95f02", + label="predicted ATE (median)") + ax[0].set_ylim(*ate_ylim) + ax[0].set(xscale="log", xlabel="training size n", ylabel="ATE", + title="ATE recovery") + if n_out: + ax[0].annotate(f"{n_out} seed(s) off-scale\n(small-n instability)", + xy=(sizes[0], ate_ylim[0]), xytext=(0, 12), + textcoords="offset points", fontsize=7, color="#d95f02") + ax[0].legend(fontsize=8, loc="lower right") + # (2) ITE recovery — correlation (left) + MAE (right), median lines + corr_ylim, mae_ylim = (-0.1, 1.0), (0.0, 0.45) + _scatter_seeds(ax[1], sizes, _per_seed(rows, sizes, "corr"), "#1b9e77", corr_ylim) + ax[1].plot(sizes, corr_med, marker="o", color="#1b9e77", label="corr (median)") + ax[1].set_ylim(*corr_ylim) + ax[1].set(xscale="log", xlabel="training size n", ylabel="ITE correlation", + title="ITE recovery") + ax1b = ax[1].twinx() + _scatter_seeds(ax1b, sizes, _per_seed(rows, sizes, "mae"), "#7570b3", mae_ylim) + ax1b.plot(sizes, np.clip(mae_med, *mae_ylim), marker="s", color="#7570b3", + label="MAE (median)") + ax1b.set_ylim(*mae_ylim) + ax1b.set_ylabel("ITE MAE") + h1, l1 = ax[1].get_legend_handles_labels() + h2, l2 = ax1b.get_legend_handles_labels() + ax[1].legend(h1 + h2, l1 + l2, fontsize=8, loc="center right") + # (3) overfitting: train vs test NLL gap + ax[2].plot(sizes, trn_med, marker="o", color="#1b9e77", label="train NLL") + ax[2].plot(sizes, tst_med, marker="o", color="#d95f02", label="test NLL") + ax[2].set(xscale="log", xlabel="training size n", ylabel="total NLL", + title="overfitting: train vs test NLL") + ax[2].legend(fontsize=8) + fig.suptitle("ITE/ATE recovery vs training size (median over seeds; " + "points = seeds). Fixed 5k test set.", fontsize=10) + fig.tight_layout() + fig.savefig(path, dpi=100, bbox_inches="tight") + plt.close(fig) + + +def main(argv: list[str] | None = None) -> None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--sizes", type=int, nargs="+", default=SIZES) + p.add_argument("--seeds", type=int, default=3) + p.add_argument("--scenario", type=int, default=1, choices=(1, 2, 3, 4)) + args = p.parse_args(argv) + + out = results_dir("ite_train_size") + res = run(args.sizes, args.seeds, args.scenario) + + # JSON-friendly copy (drop the big arrays) + save_json(out / "results.json", + {k: v for k, v in res.items() if k not in ("scatter", "ite_true")}) + plot_scatter_grid(res, out / "plots" / "ite_scatter_grid.png") + plot_curves(res, out / "plots" / "ate_ite_curves.png") + print(f"\nwrote {out}/results.json and 2 plots under {out}/plots/") + + +if __name__ == "__main__": + main()