Causal ML: Forests, Meta-Learners, TMLE, Neural¶
Machine-learning-based heterogeneous treatment effect estimators.
Causal forests¶
r = sp.causal_forest(df, y='y', treat='d', covariates=[...],
n_estimators=2000,
honest=True,
min_samples_leaf=5)
r.cate(new_X) # conditional ATE for new units
r.variable_importance()
r.subgroup_test(by='age_bin') # test heterogeneity
Meta-learners¶
| Learner | Description | Reference |
|---|---|---|
sp.metalearner(..., learner='s') |
Single-learner on (X, D) |
Künzel et al. (2019) |
sp.metalearner(..., learner='t') |
Two-learner — separate treated/control | Künzel et al. (2019) |
sp.metalearner(..., learner='x') |
Cross-learner combining S and T | Künzel et al. (2019) |
sp.metalearner(..., learner='r') |
Residualised (Robinson-style) | Nie & Wager (2021) |
sp.metalearner(..., learner='dr') |
Doubly-robust — outcome + propensity | Kennedy (2023) |
r = sp.metalearner(df, y='y', treat='d', covariates=[...], learner='x',
outcome_model=RandomForestRegressor())
r.cate(new_X)
Plus sp.cate_eval(...) for overlap, calibration, and CATE QQ plots.
TMLE¶
r = sp.tmle(df, y='y', treat='d', covariates=[...],
outcome_library=['rf', 'lasso'], # Super Learner ensemble
propensity_library=['rf', 'logistic'])
r.ate, r.ci
Neural causal models¶
sp.tarnet(df, y='y', treat='d', covariates=[...],
epochs=100, repr_layers=[200, 100]) # Shalit, Johansson, Sontag (2017)
sp.cfrnet(df, y='y', treat='d', covariates=[...], # Counterfactual Regression Net
ipm_weight=1.0)
sp.dragonnet(df, y='y', treat='d', covariates=[...], # Shi, Blei, Veitch (2019)
targeted_reg_weight=1.0)
sp.deepiv(df, y='y', treat='d', instruments='z', # Hartford et al. (2017)
covariates=[...])
Causal discovery¶
sp.notears(df, w_threshold=0.3, lambda1=0.1) # Zheng et al. 2018
sp.pc_algorithm(df, alpha=0.05) # PC-stable; ci_test='chi-square' for categorical data
sp.lingam(df) # Shimizu 2006
sp.ges(df) # Chickering 2002
Graphs: declare, test, intervene, identify¶
g = sp.dag("Z -> X; Z -> Y; X -> M -> Y") # also 'X <- Z', quoted names, digraph { ... }
g.adjustment_sets("X", "Y"); g.frontdoor_sets("X", "Y"); g.bad_controls("X", "Y")
g.test_implications(df) # chi-square for labels, Fisher z for numbers
g.recommend_estimator("X", "Y") # every strategy the graph licenses
net = sp.bayes_net(g, df) # discrete causal Bayesian network
net.query("Y", evidence={"X": 1}) # seeing
net.query("Y", do={"X": 1}) # doing
net.counterfactual("Y", evidence={...}, do={...}) # structural models only
res = sp.identify(g, "X", "Y") # Shpitser-Pearl ID
res.estimand; res.estimate(df) # works with latent confounders
ett = sp.identify_counterfactual( # ID* / IDC*
g, [("Y", 1, {"X": 0})], given=[("X", 1)]) # P(Y[X=0] = 1 | X = 1)
ett.identifiable; ett.estimand; ett.estimate(df)
Policy learning¶
sp.policy_tree(df, y='y', treat='d', covariates=[...],
depth=3) # Athey-Wager 2021
sp.policy_value(tree, df_test)
Bayesian causal forests¶
sp.bcf(df, y='y', treat='d', covariates=[...],
n_trees_mu=200, n_trees_tau=50) # Hahn, Murray, Carvalho 2020
# since 1.41 the same as sp.bart_causal(df, 'y', 'd', [...], method='bcf')