Causal Forests

56.1 Trees that split on effects (40 min)

Concept. A regression tree splits customers where spend differs. A causal tree (Athey & Imbens, 2016) splits where the effect differs: it looks for the split that makes the coupon effect most different between the two sides. A causal forest (Wager & Athey 2018; Athey, Tibshirani & Wager 2019) averages many such trees. Three ingredients make it statistically valid:

• Honesty: each tree uses one half of its sample to choose splits and the other half to estimate effects in the leaves, so the estimates aren't tuned to the noise that chose the splits.
• Subsampling: each tree sees a random subset, which makes the forest's variance estimable and gives confidence intervals.
• Orthogonalisation: EconML's CausalForestDML first residualises spend and coupon on X (Day 52), so the forest only has to find effect differences, not untangle confounding.

56.2 A causal forest on the coupon data (60 min)

Create day56/causal_forest.py. Its fit_forest function is reused for the rest of the module (fitting takes about 40 seconds):

# day56/causal_forest.py   (python -m day56.causal_forest)
"""
Causal forests (Wager & Athey, 2018; Athey, Tibshirani & Wager, 2019), in EconML's CausalForestDML:
residualise spend and coupon on X with ML (Day 52), then grow 'honest' trees that split where the
EFFECT differs, not where spend differs. Each customer gets a CATE and a confidence interval.
"""
import numpy as np
from econml.dml import CausalForestDML
from sklearn.model_selection import train_test_split

from causalkit.coupons import X_COLS, coupons
from causalkit.learners import clf, reg


def fit_forest(df=None):
    """Fit on half of the customers; return the forest and the (train, test) halves as pandas."""
    full = (coupons() if df is None else df).with_row_index("row")
    train, test = train_test_split(full.to_pandas(), test_size=0.5, random_state=0)
    cf = CausalForestDML(model_y=reg(), model_t=clf(), discrete_treatment=True, cv=5, n_estimators=1000,
                         min_samples_leaf=20, max_samples=0.45, honest=True, random_state=0)
    cf.fit(train.spend, train.coupon, X=train[X_COLS])
    return cf, train, test


if __name__ == "__main__":
    cf, train, test = fit_forest()
    tau_hat = cf.effect(test[X_COLS])
    lo, hi = cf.effect_interval(test[X_COLS], alpha=0.05)
    tau = test._tau.to_numpy()
    print(f"PEHE {np.sqrt(np.mean((tau_hat - tau) ** 2)):.2f}; "
          f"correlation with true CATE {np.corrcoef(tau_hat, tau)[0, 1]:.3f}")
    print(f"Customer-level 95% intervals contain the true CATE for {np.mean((lo <= tau) & (tau <= hi)):.0%} "
          f"of customers; median width ${np.median(hi - lo):.0f}")
    ate = cf.ate_inference(X=test[X_COLS])
    print(f"ATE {ate.mean_point:.2f}   truth {tau.mean():.2f}")
    imp = sorted(zip(X_COLS, cf.feature_importances_), key=lambda kv: -kv[1])
    print("Where the forest splits on effect differences (importance):")
    for name, v in imp[:5]:
        print(f"   {name:<14} {v:.3f}")
    print(f"   noise features z1-z10 together: {sum(v for n, v in imp if n.startswith('z')):.3f}")

python -m day56.causal_forest

PEHE 8.64; correlation with true CATE 0.838
Customer-level 95% intervals contain the true CATE for 47% of customers; median width $11
ATE 15.87   truth 16.47
Where the forest splits on effect differences (importance):
   deal_share     0.887
   app_user       0.018
   visits         0.013
   pre_spend      0.009
   z9             0.006
   noise features z1-z10 together: 0.048

• The forest finds the main driver: 89% of its effect-splits are on deal_share. The noise features get under 5%.
• Its accuracy (PEHE 8.6) is between the meta-learners of Day 55: forests are not magic.
• The customer-level 95% intervals cover the true effect for only 47% of customers. The forest's intervals describe uncertainty around its own smoothed estimate, and a forest smooths: customers with extreme true effects get estimates pulled towards their neighbours'. Individual-level intervals are the least reliable output of any CATE method. Use averages over groups instead.

56.3 Group effects without shrinkage (55 min)

Concept. The business already uses groups: deal-seeker bands, app users. Their group average treatment effects (GATEs) are what marketing can act on. Two ways to compute them:

1. Average the forest's predictions within each group. Simple, but forest predictions are shrunk towards the overall average, and so are their averages.
2. Average the doubly robust (AIPW) scores within each group (Chernozhukov et al.'s GATES). Each customer's score is a noisy but unbiased estimate of their effect; averaging over thousands removes the noise without shrinking anything. DoubleML's IRM.gate() does this, with confidence intervals.

Create day56/gates.py:

# day56/gates.py   (python -m day56.gates)
"""
Group average treatment effects (GATEs) for groups the business already uses. Two ways:
  1. average the causal forest's CATE predictions within each group;
  2. average the doubly robust (AIPW) scores within each group: DoubleML's IRM.gate().
Forest predictions are SMOOTHED towards the overall average (that's how forests reduce noise), so
group averages of them inherit that shrinkage. Averages of doubly robust scores don't.
"""
import doubleml as dml
import numpy as np
import pandas as pd
import polars as pl

from causalkit.coupons import X_COLS, coupons
from causalkit.learners import clf, reg
from day56.causal_forest import fit_forest

np.random.seed(56)  # DoubleML draws its cross-fitting folds from numpy's global RNG

df = coupons().to_pandas()
cf, train, test = fit_forest()
irm = dml.DoubleMLIRM(dml.DoubleMLData(df[X_COLS + ["coupon", "spend"]], "spend", "coupon", X_COLS),
                      ml_g=reg(), ml_m=clf(), n_folds=5, trimming_threshold=0.02).fit()

partitions = [   # each partition's groups must not overlap: one gate() call per partition
    {"deal share < 15%": lambda d: d.deal_share < 0.15,
     "deal share 15-25%": lambda d: d.deal_share.between(0.15, 0.25, inclusive="left"),
     "deal share 25-40%": lambda d: d.deal_share.between(0.25, 0.40, inclusive="left"),
     "deal share 40%+": lambda d: d.deal_share >= 0.40},
    {"no app": lambda d: d.app_user == 0, "app user": lambda d: d.app_user == 1},
]
rows = []
for splits in partitions:
    gate = irm.gate(groups=pd.DataFrame({k: f(df) for k, f in splits.items()})).confint()
    for name, f in splits.items():
        rows.append({"group": name, "customers": int(f(df).sum()),
                     "forest GATE": round(cf.effect(test[f(test)][X_COLS]).mean(), 1),
                     "DR-score GATE": round(gate.loc[name, "effect"], 1),
                     "DR 95% CI": f"[{gate.loc[name, '2.5 %']:.1f}, {gate.loc[name, '97.5 %']:.1f}]",
                     "true average": round(df._tau[f(df)].mean(), 1)})
with pl.Config(tbl_width_chars=110, fmt_str_lengths=30):
    print(pl.DataFrame(rows))

python -m day56.gates

│ group             ┆ customers ┆ forest GATE ┆ DR-score GATE ┆ DR 95% CI    ┆ true average │
│ deal share < 15%  ┆ 4421      ┆ 9.6         ┆ 4.1           ┆ [1.4, 6.7]   ┆ 2.9          │
│ deal share 15-25% ┆ 4759      ┆ 10.5        ┆ 9.5           ┆ [7.1, 11.8]  ┆ 10.6         │
│ deal share 25-40% ┆ 6066      ┆ 15.4        ┆ 18.4          ┆ [16.4, 20.5] ┆ 18.7         │
│ deal share 40%+   ┆ 4754      ┆ 27.5        ┆ 30.9          ┆ [28.2, 33.5] ┆ 32.3         │
│ no app            ┆ 11016     ┆ 14.9        ┆ 11.3          ┆ [9.7, 12.9]  ┆ 12.0         │
│ app user          ┆ 8984      ┆ 17.1        ┆ 21.9          ┆ [20.1, 23.8] ┆ 22.1         │

The forest says the least deal-prone customers respond by $9.6; the truth is $2.9, and the DR-score GATE gets $4.1. At the top, the forest says $27.5 against $32.3 true. The forest's ranking is right, but its spread is compressed by about half. The DR-score GATEs are within $1.4 of the truth everywhere, with honest intervals. Note the app-user gap: the forest says $2.2 between app users and others, the truth is $10.1.

Rule: use a CATE model to form groups; use doubly robust scores to estimate the groups' effects.

gate() needs non-overlapping groups, so the script calls it once per partition (deal bands, then app use).

56.4 What noise features cost (30 min)

Try it. Refit the forest without the ten noise columns z1–z10. Does it improve? Create day56/no_noise_features.py.

# day56/no_noise_features.py   (python -m day56.no_noise_features)
import numpy as np
from econml.dml import CausalForestDML

from causalkit.coupons import BEHAVIOUR
from causalkit.learners import clf, reg
from day56.causal_forest import fit_forest

_, train, test = fit_forest()
cols = BEHAVIOUR + ["lapsing"]
cf = CausalForestDML(model_y=reg(), model_t=clf(), discrete_treatment=True, cv=5, n_estimators=1000,
                     min_samples_leaf=20, max_samples=0.45, honest=True, random_state=0)
cf.fit(train.spend, train.coupon, X=train[cols])
pred = cf.effect(test[cols])
print(f"Without z1-z10: PEHE {np.sqrt(np.mean((pred - test._tau) ** 2)):.2f}, "
      f"correlation {np.corrcoef(pred, test._tau)[0, 1]:.3f}")

Without z1-z10: PEHE 7.77, correlation 0.871

56.5 Day 56 mini-project: rules marketing can use (65 min)

Marketing can't act on 20,000 individual estimates, but it can act on three or four rules. Fit a depth-2 tree to the forest's predictions (EconML's SingleTreeCateInterpreter), turn its leaves into segments, and estimate each segment's effect with DR scores. Create day56/segment_rules.py:

# day56/segment_rules.py   (python -m day56.segment_rules)
"""
Mini-project: marketing can't act on 20,000 individual CATEs, but it can act on a few rules.
Fit a depth-2 tree to the causal forest's CATEs (EconML's SingleTreeCateInterpreter), read off the
segments, and check each segment's effect against the truth with doubly robust GATEs.
"""
import doubleml as dml
import numpy as np
import pandas as pd
import polars as pl
from econml.cate_interpreter import SingleTreeCateInterpreter
from sklearn.tree import export_text

from causalkit.coupons import X_COLS, coupons
from causalkit.learners import clf, reg
from day56.causal_forest import fit_forest

np.random.seed(56)  # DoubleML draws its cross-fitting folds from numpy's global RNG

cf, train, test = fit_forest()
interp = SingleTreeCateInterpreter(max_depth=2, min_samples_leaf=500, include_model_uncertainty=False)
interp.interpret(cf, test[X_COLS])
tree = interp.tree_model_
print(export_text(tree, feature_names=X_COLS, decimals=2))



def leaf_rules(tree, names):
    """Plain-language rule for each leaf of a fitted sklearn tree."""
    t, out = tree.tree_, {}

    def walk(node, conds):
        if t.children_left[node] == -1:
            out[node] = " and ".join(conds)
            return
        f, thr = names[t.feature[node]], t.threshold[node]
        walk(t.children_left[node], conds + [f"{f} <= {thr:.2f}"])
        walk(t.children_right[node], conds + [f"{f} > {thr:.2f}"])

    walk(0, [])
    return out


rules = leaf_rules(tree, X_COLS)
df = coupons().to_pandas()
leaf = tree.apply(df[X_COLS].to_numpy())
irm = dml.DoubleMLIRM(dml.DoubleMLData(df[X_COLS + ["coupon", "spend"]], "spend", "coupon", X_COLS),
                      ml_g=reg(), ml_m=clf(), n_folds=5, trimming_threshold=0.02).fit()
groups = pd.DataFrame({rules[k]: leaf == k for k in np.unique(leaf)})
gate = irm.gate(groups=groups).confint()
rows = [{"segment": name, "customers": int(groups[name].sum()), "DR-score GATE": round(gate.loc[name, "effect"], 1),
         "95% CI": f"[{gate.loc[name, '2.5 %']:.1f}, {gate.loc[name, '97.5 %']:.1f}]",
         "true average": round(df._tau[groups[name]].mean(), 1)} for name in groups]
with pl.Config(tbl_width_chars=120, fmt_str_lengths=50):
    print(pl.DataFrame(rows))

python -m day56.segment_rules

|--- deal_share <= 0.37
|   |--- deal_share <= 0.29
|   |   |--- value: [10.36]
|   |--- deal_share >  0.29
|   |   |--- value: [16.03]
|--- deal_share >  0.37
|   |--- deal_share <= 0.44
|   |   |--- value: [22.61]
|   |--- deal_share >  0.44
|   |   |--- value: [28.47]

│ segment                                   ┆ customers ┆ DR-score GATE ┆ 95% CI       ┆ true average │
│ deal_share <= 0.37 and deal_share <= 0.29 ┆ 10810     ┆ 7.8           ┆ [6.2, 9.4]   ┆ 8.0          │
│ deal_share <= 0.37 and deal_share > 0.29  ┆ 3646      ┆ 19.0          ┆ [16.3, 21.6] ┆ 19.6         │
│ deal_share > 0.37 and deal_share <= 0.44  ┆ 1881      ┆ 26.2          ┆ [22.7, 29.7] ┆ 24.8         │
│ deal_share > 0.37 and deal_share > 0.44   ┆ 3663      ┆ 32.4          ┆ [29.2, 35.6] ┆ 34.2         │

The rules are simple: four deal_share bands, with cut points at 29%, 37% and 44%. The tree's own leaf values (10, 16, 23, 28) are the forest's shrunk estimates again; the DR-score GATEs (8, 19, 26, 32) are within $1.8 of the truth. A brief for marketing reads: "customers who bought more than 29% of their items on promotion respond to a $10 coupon with $19–32 of extra spend; the rest, about $8".

One caution to write in day56/NOTES.md: the cut points were chosen by looking at the data, so testing "is segment 4 different from segment 1?" on the same data overstates the evidence. Day 57 shows how to validate on data the model never saw.

56.6 Save your work with Git (20 min)

git add .
git commit -m "Day 56: causal forest, DR-score GATEs, interpretable segment rules"
git push

Day 56 checklist

• I can explain honesty, subsampling and orthogonalisation in a causal forest
• I know that individual-level CATE intervals are unreliable, and why
• I estimate group effects with doubly robust scores, not by averaging forest predictions
• I can turn a forest into a few rules with a shallow interpreter tree