The AutoCATE class is a high-level API facilitating an AutoML framework for CATE estimation, built on top of the EconML library.
AutoCATE is experimental and may change significantly in future versions.
The CATE is defined as \(\mathbb{E}[\tau \mid \mathbf{X}]\) where \(\tau\) is the treatment effect and \(\mathbf{X}\) is the set of features leveraged for treatment effect heterogeneity.
This class is built on top of the EconML library and provides a high-level AutoML API for fitting, validating, and making predictions/inference with CATE models, with best practices built directly into the API. The class is designed to be easy to use and understand, while still providing flexibility for advanced users.
Note that first-stage models are estimated using a full AutoML framework via Flaml, whereas the second-stage models are currently estimated & selected based on a pre-specified set of models (or custom CATE models) passed, thus there is no tuning of hyperparameters. This and other AutoML features are on the roadmap for future versions of AutoCATE.
For technical details on the AutoCATE class, see here.
Note: All the standard assumptions of Causal Inference apply to this class (e.g., exogeneity/unconfoundedness, overlap, positivity, etc.). The class does not check for these assumptions and assumes that the user has already thought through these assumptions before using the class.
The str representing the column name for the outcome variable.
required
T
str
The str representing the column name(s) for the treatment variable(s).
required
X
Sequence[str] | None
The sequence of feature names representing the feature set to be utilized for estimating heterogeneity/CATE.
None
W
Sequence[str] | None
The sequence of feature names representing the confounder/control feature set to be utilized only for nuisance function estimation. When W is passed, only Orthogonal learners will be leveraged.
None
discrete_treatment
bool
A boolean indicating whether the treatment is discrete/categorical or continuous.
True
discrete_outcome
bool
A boolean indicating whether the outcome is binary or continuous.
False
model_Y
dict | BaseEstimator | None
A dictionary of FLAML kwarg overrides or a BaseEstimator instance for the outcome model - \(\mathbb{E}[Y \mid \mathbf{X},\mathbf{W}]\).
None
model_T
dict | BaseEstimator | None
A dictionary of FLAML kwarg overrides or a BaseEstimator instance for the treatment/propensity model - \(\mathbb{E}[T \mid \mathbf{X},\mathbf{W}]\).
None
model_regression
dict | BaseEstimator | None
A dictionary of FLAML kwarg overrides or a BaseEstimator instance for the regression model - \(\mathbb{E}[Y \mid \mathbf{X},\mathbf{W},T]\).
None
enable_categorical
bool
A boolean indicating whether to enable categorical encoding for the models. When set to True, pandas categorical types will be converted to ordinal encodings. For one-hot encoding, please implement it prior to passing the data to the model.
False
n_jobs
int
The number of jobs to run in parallel for model training.
1
use_ray
bool
A boolean indicating whether to use Ray for distributed computing, utilized in both AutoML for first-stage models and AutoML for CATE models. n_concurrent_tasks defaults to 4 if not specified in FLAML kwarg overrides (e.g., for model_*).
False
ray_remote_func_options_kwargs
dict | None
A dictionary of Ray remote function options for ray remote function that fits each individual CATE model. See here for options.
None
use_spark
bool
A boolean indicating whether to use Spark for distributed computing, utilized only in AutoML for first-stage models. n_concurrent_tasks defaults to 4 if not specified in FLAML kwarg overrides (e.g., for model_*).
The list of str representing the column names for the outcome variable \(Y\).
T
list[str]
The list of str representing the column names for the treatment variable \(T\).
X
list[str]
The list of str representing the confounder/control feature set \(\mathbf{X}\) to be utilized for estimating heterogeneity/CATE and nuisance function estimation where applicable.
W
list[str] | None
The list of str representing the confounder/control feature set \(\mathbf{W}\) to be utilized only for nuisance function estimation, where applicable. These will be included by default in Meta-Learners.
available_estimators
list[str]
A list of the available CATE estimators out of the box. Validity of estimator at runtime will depend on the outcome and treatment types and be automatically selected.
model_Y
BaseEstimator
The selected outcome model - \(\mathbb{E}[Y \mid \mathbf{X},\mathbf{W}]\).
model_T
BaseEstimator
The selected treatment model - \(\mathbb{E}[T \mid \mathbf{X},\mathbf{W}]\).
model_regression
BaseEstimator
The selected regression model - \(\mathbb{E}[Y \mid \mathbf{X},\mathbf{W},T]\).
rscores
dict[str, float]
The dictionary of the Rscores on the validation set for each CATE estimator fitted during model selection.
Run end-to-end fitting, validation & model selection, and testing for CATE models.
Parameters
Name
Type
Description
Default
df
PandasConvertibleDataFrame
The dataset to fit the CATE model on. Accepts a Pandas DataFrame or a compatible object with a to_pandas() or .toPandas() method.
required
cate_estimators
Sequence[str]
The out-of-the-box CATE estimators to use. Accessible via self.available_estimators.
list(available_estimators.keys())
additional_cate_estimators
Sequence[AutoCateEstimator]
Additional CATE estimators to use.
list()
use_cached_models
bool
Whether to use cached first-stage/nuisance models, if previously fitted.
False
ensemble
bool
Whether to use an ensemble of CATE estimators.
False
refit_final
bool
Whether to refit the final CATE estimator on the entire dataset after model selection.
True
bootstrap_inference
bool
Whether to use bootstrap inference for the final CATE estimator, when other inference methods are not available (e.g., metalearners). This can be computationally expensive.
False
n_bootstrap_samples
int
The number of bootstrap samples to use for bootstrap inference, if bootstrap_inference is True.
100
validation_fraction
float
The fraction of the dataset to use for validation.
Calculate the average treatment effect(s) (ATE) \(\mathbb{E}[\tau(\mathbf{X})]\).
This method can be used for group average treatment effects (GATEs) \(\mathbb{E}[\tau(\mathbf{X}) \mid G]\) by filtering the input df to a specific group.
Two effect modes are supported: discrete and marginal.
In discrete mode, the effect is calculated between two specific treatment levels T0 and T1 - \(\mathbb{E}[\tau(\mathbf{X}, T0, T1)]\).
In marginal mode, the effect is calculated for each observation as a gradient around their treatment levels T - \(\mathbb{E}[\partial_{\tau}(T,\mathbf{x})]\)
The CATE estimates as an array if return_inference is False. Otherwise, the return value is an instance of NormalInferenceResults if asymptotic inference is available for the best estimator or EmpiricalInferenceResults if not.
