metacausal.aggregation.CausalStacking

class metacausal.aggregation.CausalStacking(split=<factory>, propensity_model=None, outcome_model=None, propensity_trim=0.01, fit_nuisance_fn=None, pseudo_outcome_fn=None)

Bases: SupervisedStrategy

Causal Stacking (Han & Wu, 2022).

Finds a convex combination of component CATEs that minimises the mean squared error of the ensemble prediction against the DR/AIPW pseudo-outcome:

min_{θ ≥ 0, Σθ=1} (1/m) Σ_i (Γ_i - θ·τ(x_i))²

This is the special case of Q-Aggregation with ν=0 and β=0.

Parameters:

• split (CrossFitSplit or TrainAvgSplit) – Data splitting strategy. Default: 5-fold cross-fitting.

• propensity_model (sklearn classifier or None) – Model for P(T=1|X). None selects HistGradientBoostingClassifier.

• outcome_model (sklearn regressor / classifier or None) – Model for E[Y|X, T]. For continuous Y, must be a regressor; for binary Y, must be a classifier (predict_proba). None selects an outcome-type-appropriate default (HistGradientBoostingRegressor or HistGradientBoostingClassifier).

• propensity_trim (float) – Clip propensity scores to [trim, 1-trim].

• fit_nuisance_fn (callable or None) – Optional replacement for the default nuisance fitting procedure. See SupervisedStrategy.fit_nuisance_fn for the required signature.

• pseudo_outcome_fn (callable or None) –

  Optional replacement for dr_pseudo_outcome. Must have signature:

  pseudo_outcome_fn(Y, T, nuisance) -> ndarray of shape (m,)
  

  Use this to inject targeted-learning influence functions or other semiparametric pseudo-outcome constructions. None (default) uses the standard AIPW form.

Examples

>>> from sklearn.linear_model import LinearRegression
>>> from sklearn.ensemble import HistGradientBoostingRegressor as HGBR
>>> from metacausal import CausalEnsemble
>>> from metacausal.adapters import GenericCATEAdapter
>>> from metacausal.aggregation import CausalStacking
>>> from metacausal.datasets import load_lalonde
>>> X, T, Y = load_lalonde()
>>> def fit_linear(X, T, Y, **kwargs):
...     treated = T == 1
...     m1 = LinearRegression().fit(X[treated], Y[treated])
...     m0 = LinearRegression().fit(X[~treated], Y[~treated])
...     return (m1, m0)
>>> def fit_hgb(X, T, Y, **kwargs):
...     treated = T == 1
...     m1 = HGBR(max_iter=20).fit(X[treated], Y[treated])
...     m0 = HGBR(max_iter=20).fit(X[~treated], Y[~treated])
...     return (m1, m0)
>>> def cate_fn(state, X):
...     m1, m0 = state
...     return m1.predict(X) - m0.predict(X)
>>> methods = [
...     GenericCATEAdapter(fit_linear, cate_fn, name="linear"),
...     GenericCATEAdapter(fit_hgb, cate_fn, name="hgb"),
... ]
>>> ens = CausalEnsemble(methods=methods, aggregation=CausalStacking())
>>> _ = ens.fit(X, T, Y, random_state=42)
>>> ens.ate()
AteEstimate(ate=..., n_methods=2, aggregation='CausalStacking', spread=...)

Methods

__init__
aggregate	Aggregate component CATEs using stored weights.
fit_weights	Compute Causal Stacking weights from OOF predictions and nuisance estimates.

Attributes

ensemble_weights	Ensemble weights computed during fit().
fit_nuisance_fn
outcome_model
propensity_model
propensity_trim
pseudo_outcome_fn
strategy_family
split

Details

fit_weights(cate_predictions, Y, T, X, nuisance)

Compute Causal Stacking weights from OOF predictions and nuisance estimates.

Parameters:

• cate_predictions (array of shape (K, m))

• Y (arrays of shape (m,))

• T (arrays of shape (m,))

• X (array of shape (m, p))

• nuisance (NuisanceEstimates)

Returns:

EnsembleWeights (model_names populated by _fit_supervised).

Return type:

EnsembleWeights

pseudo_outcome_fn: Any = None