metacausal.adapters.GenericCATEAdapter

class metacausal.adapters.GenericCATEAdapter(fn_fit, fn_cate, fn_ate=None, name='custom_cate', supported_outcome_types=('continuous',))

Bases: object

Wrap callables as a CATE-capable causal estimator.

Reduces the barrier for prototyping custom CATE estimators to two or three functions — one for fitting, one for CATE prediction, and an optional one for ATE — without implementing the full CausalEstimator protocol.

Parameters:

• fn_fit (callable) –

  fn_fit(X, T, Y, **kwargs) -> state

  Called during fit(). Receives the training arrays and any keyword arguments forwarded by CausalEnsemble (e.g. random_state). Returns an opaque state object (fitted model, dict, namedtuple, None for stateless callables …) passed to fn_cate and fn_ate at prediction time.

  Reproducibility: if fn_fit uses randomness, it is the caller’s responsibility to consume kwargs.get("random_state"). CausalEnsemble forwards the seed but cannot enforce its use.

  Deep-copy requirement: state must survive copy.deepcopy() so that bootstrap resampling and supervised cross-fitting work correctly. Closures over PyTorch tensors, open file handles, or GPU state may fail; store a picklable specification (CPU weights, a file path, or constructor arguments) and rebuild such objects lazily inside fn_cate / fn_ate instead.

• fn_cate (callable) –

  fn_cate(state, X) -> array-like of shape (n,)

  Called during cate(). Receives the fitted state and an evaluation covariate matrix X of shape (n, p). Must return a 1-D array (or (n, 1) array, which is automatically squeezed) of CATE predictions.

• fn_ate (callable or None) –

  fn_ate(state, X) -> float | ComponentAteEstimate

  Optional. Called during ate(). Receives the fitted state and the evaluation covariate matrix X. Returns either a scalar ATE or a ComponentAteEstimate (useful when the estimator provides a native, more efficient ATE — e.g. a DR/AIPW mean rather than mean(cate(X))).

  If None (default), ate() falls back to mean(cate(X_eval)).

• name (str) – Display name used in CateEstimate.component_estimates and warning messages. Default "custom_cate".

• supported_outcome_types (tuple[str, ...])

Notes

Parallelism cooperation contract. MetaCausal automatically suppresses the internal joblib-parallel knobs (n_jobs, cv_n_jobs, …) of the estimators it wraps for EconML, CausalML, DoubleML, and stochtree whenever they run inside one of MetaCausal’s own workers (fit/bootstrap, or supervised cross-fitting) – this stops the wrapped estimator’s own default from nesting a second worker pool inside the outer one and oversubscribing cores. That automatic suppression cannot reach into an opaque fn_fit/fn_cate/fn_ate callable, so if your model parallelizes internally, you are responsible for one of:

• Constructing it with its own parallelism knob fixed at 1 up front, before passing it to fn_fit; or

• Cooperating with the same signal MetaCausal’s built-in adapters use: inside fn_fit, check os.environ.get(metacausal.adapters.INNER_WORKER_ENV) and pin your model to serial when it is set (meaning you are already running inside a MetaCausal worker, so any further internal parallelism would nest and oversubscribe).

See “Extending MetaCausal” in the project README for a worked example of the second option.

Examples

Wrap a T-learner whose native ATE is a doubly robust (AIPW) mean computed at fit time, where the observed T and Y are available. This is the reason fn_ate exists: it returns an estimate the default mean(cate(X)) fallback cannot produce, because the fallback never sees the outcomes.

>>> import numpy as np
>>> from sklearn.linear_model import LinearRegression, LogisticRegression
>>> from metacausal import CausalEnsemble
>>> from metacausal.adapters import GenericCATEAdapter
>>> from metacausal.datasets import load_lalonde
>>> X, T, Y = load_lalonde()
>>> def fit_fn(X, T, Y, **kwargs):
...     treated = T == 1
...     m1 = LinearRegression().fit(X[treated], Y[treated])
...     m0 = LinearRegression().fit(X[~treated], Y[~treated])
...     e = LogisticRegression(max_iter=1000).fit(X, T).predict_proba(X)[:, 1]
...     mu1, mu0 = m1.predict(X), m0.predict(X)
...     # AIPW / doubly robust ATE on the training sample
...     ate_dr = float(np.mean(
...         (mu1 - mu0)
...         + T * (Y - mu1) / e
...         - (1 - T) * (Y - mu0) / (1 - e)
...     ))
...     return {"m1": m1, "m0": m0, "ate_dr": ate_dr}
>>> def cate_fn(state, X):
...     return state["m1"].predict(X) - state["m0"].predict(X)
>>> def ate_fn(state, X):
...     # native DR ATE precomputed at fit time; differs from mean(cate(X))
...     return state["ate_dr"]
>>> adapter = GenericCATEAdapter(fit_fn, cate_fn, fn_ate=ate_fn,
...                              name="t_learner")
>>> ens = CausalEnsemble(methods=[adapter], aggregation="median")
>>> _ = ens.fit(X, T, Y, random_state=42)
>>> ens.ate()
AteEstimate(ate=..., n_methods=1, aggregation='Median', spread=...)

Methods

__init__
ate	Return the ATE.
cate	Return CATE predictions by calling fn_cate(state, X).
fit	Fit by calling fn_fit(X, T, Y, **kwargs) and storing the state.
validate_outcome_type	No-op: the wrapped callables are opaque, so the user's supported_outcome_types declaration is the contract.

Attributes

name
supports_cate

Details

ate(X=None)

Return the ATE.

If fn_ate was provided, calls fn_ate(state, X_eval). Otherwise returns mean(cate(X_eval)).

X_eval is X if given, otherwise the training data from fit().

Parameters:

X (ndarray | None)

Return type:

ComponentAteEstimate

cate(X)

Return CATE predictions by calling fn_cate(state, X).

Accepts (n,) or (n, 1) output; raises ValueError otherwise.

Parameters:

X (ndarray)

Return type:

ComponentCateEstimate

fit(X, T, Y, **kwargs)

Fit by calling fn_fit(X, T, Y, **kwargs) and storing the state.

Parameters:

• X (ndarray)

• T (ndarray)

• Y (ndarray)

Return type:

None

validate_outcome_type(detected)

No-op: the wrapped callables are opaque, so the user’s supported_outcome_types declaration is the contract.

Parameters:

detected (str)

Return type:

None

property name: str

property supports_cate: bool