API Reference#

Submodules#

causalml.inference.tree module#

class causalml.inference.tree.CausalRandomForestRegressor(n_estimators: int = 100, *, control_name: int | str = 0, criterion: str = 'causal_mse', alpha: float = 0.05, max_depth: int = None, min_samples_split: int = 60, min_samples_leaf: int = 100, min_group_samples: int = 50, min_weight_fraction_leaf: float = 0.0, max_features: int | float | str = 1.0, max_leaf_nodes: int = None, min_impurity_decrease: float = -inf, bootstrap: bool = True, oob_score: bool = False, n_jobs: int = None, random_state: int = None, verbose: int = 0, warm_start: bool = False, ccp_alpha: float | str = 0.0, groups_penalty: float = 0.5, max_samples: int = None, groups_cnt: bool = True, groups_cnt_mode: str = 'nodes', honesty: bool = True, estimation_sample_size: float = 0.5, cv_folds: int = 5)[source]#

Bases: SerializableLearner, ForestRegressor

A random forest of CausalTreeRegressor estimators.

Note

Observational data needs inverse-propensity weights. The causal criteria compare raw group means with no adjustment for treatment assignment, so a varying propensity biases both the splits and the leaf estimates. Pass inverse-propensity weights as sample_weight to fit; see the note on CausalTreeRegressor for the recipe and the measured effect.

calculate_error(X_train: ndarray, X_test: ndarray, inbag: ndarray = None, calibrate: bool = True, memory_constrained: bool = False, memory_limit: int = None) → ndarray[source]#

Calculate error bars from scikit-learn RandomForest estimators Source: scikit-learn-contrib/forest-confidence-interval

Parameters:

  • X_train – (np.ndarray), training subsample of feature matrix, (n_train_sample, n_features)

  • X_test – (np.ndarray), test subsample of feature matrix, (n_train_sample, n_features)

  • inbag – (ndarray, optional), The inbag matrix that fit the data. If set to None (default) it will be inferred from the forest. However, this only works for trees for which bootstrapping was set to True. That is, if sampling was done with replacement. Otherwise, users need to provide their own inbag matrix.

  • calibrate – (boolean, optional) Whether to apply calibration to mitigate Monte Carlo noise. Some variance estimates may be negative due to Monte Carlo effects if the number of trees in the forest is too small. To use calibration, Default: True

  • memory_constrained – (boolean, optional) Whether or not there is a restriction on memory. If False, it is assumed that a ndarray of shape (n_train_sample,n_test_sample) fits in main memory. Setting to True can actually provide a speedup if memory_limit is tuned to the optimal range.

  • memory_limit – (int, optional) An upper bound for how much memory the intermediate matrices will take up in Megabytes. This must be provided if memory_constrained=True.

Returns:

(np.ndarray), An array with the unbiased sampling variance for a RandomForest object.

Note

This method delegates to forestci (forest-confidence-interval), which calls scikit-learn’s private _get_n_samples_bootstrap and _generate_sample_indices with their pre-1.9 signatures. Under scikit-learn >= 1.9 those helpers require an extra sample_weight argument, so forestci (<= 0.7) would raise TypeError here. As a temporary measure this call runs inside _forestci_sklearn19_compat(), which patches those helpers to pass sample_weight=None (the pre-1.9 uniform-bootstrap default, so results are unchanged), pending the upstream fix (forest-confidence-interval#122) and a forestci pin bump. The shim is not thread-safe (see its docstring). See uber/causalml#906.

fit(X: ndarray, treatment: ndarray, y: ndarray, sample_weight: ndarray = None)[source]#

Fit Causal RandomForest

Parameters:

  • X – (np.ndarray), feature matrix

  • treatment – (np.ndarray), treatment vector

  • y – (np.ndarray), outcome vector

  • sample_weight – (np.ndarray), sample weights. Pass inverse-propensity weights here on observational data — see the note on the class.

Returns:

self

predict(X: ndarray, with_outcomes: bool = False) → ndarray[source]#

Predict individual treatment effects

Parameters:

  • X (np.ndarray) – a feature matrix

  • with_outcomes (bool) – include outcomes Y_hat(X|T=0), Y_hat(X|T=1) along with individual treatment effect

Returns:

individual treatment effect (ITE), dim=(samples, groups-1)

or ITE with outcomes: [Y_hat(X|T=0), Y_hat(X|T=1),…,Y_hat(X|T=n), ITE_1, ITE_2,…,ITE_n], dim=(samples, 2*groups-1)

Return type:

(np.ndarray)

set_fit_request(*, sample_weight: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → CausalRandomForestRegressor#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • sample_weight (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for sample_weight parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, with_outcomes: bool | None | str = '$UNCHANGED$') → CausalRandomForestRegressor#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

with_outcomes (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for with_outcomes parameter in predict.

Returns:

self – The updated object.

Return type:

object

set_score_request(*, sample_weight: bool | None | str = '$UNCHANGED$') → CausalRandomForestRegressor#

Configure whether metadata should be requested to be passed to the score method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to score if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to score.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

sample_weight (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for sample_weight parameter in score.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.tree.CausalTreeRegressor(*, criterion: str = 'causal_mse', splitter: str = 'best', alpha: float = 0.05, control_name: int | str = 0, max_depth: int = None, min_samples_split: int | float = 60, min_weight_fraction_leaf: float = 0.0, max_features: int | float | str = None, max_leaf_nodes: int = None, min_impurity_decrease: float = -inf, ccp_alpha: float | str = 0.0, groups_penalty: float = 0.5, min_group_samples: int = 50, min_samples_leaf: int = 100, random_state: int = None, groups_cnt: bool = False, groups_cnt_mode: str = 'nodes', node_pvalues: bool = False, honesty: bool = True, estimation_sample_size: float = 0.5, cv_folds: int = 5)[source]#

Bases: SerializableLearner, RegressorMixin, BaseCausalDecisionTree

A Causal Tree regressor class. The Causal Tree is a decision tree regressor with a split criteria for treatment effects. Details are available at Athey and Imbens (2015).

Note

Observational data needs inverse-propensity weights. Every criterion here compares raw group means, with no adjustment for how treatment was assigned, so when the propensity varies with X both the split search and the leaf estimates inherit the confounding bias. Pass inverse-propensity weights as sample_weight to correct it:

e_hat = cross_val_predict(clf, X, treatment, cv=5, method="predict_proba")[:, 1]
e_hat = np.clip(e_hat, 0.05, 0.95)
ipw = np.where(treatment == 1, 1 / e_hat, 1 / (1 - e_hat))
model.fit(X=X, treatment=treatment, y=y, sample_weight=ipw)

Measured on Nie & Wager Setup A over 10 seeds (a confounded design whose propensity and treatment effect share the same features), for CausalRandomForestRegressor with a cross-fitted propensity:

weights

corr with true tau

ATE error

CATE RMSE

none

0.381

0.389

0.433

inverse-propensity

0.844

0.067

0.142

No correction is needed for a randomized experiment. See docs/examples/causal_tree_honesty_parity.ipynb for the full comparison.

bootstrap(X: ndarray, treatment: ndarray, y: ndarray, sample_size: int, seed: int) → ndarray[source]#

Runs a single bootstrap.

Fits on bootstrapped sample, then predicts on whole population.

Parameters:

  • X (np.ndarray) – a feature matrix

  • treatment (np.ndarray) – a treatment vector

  • y (np.ndarray) – an outcome vector

  • sample_size (int) – bootstrap sample size

  • seed – (int): bootstrap seed

Returns:

bootstrap predictions

Return type:

(np.ndarray)

bootstrap_pool(X: ndarray, treatment: ndarray, y: ndarray, n_bootstraps: int, bootstrap_size: int, n_jobs: int, verbose: bool)[source]#

Run a pool of bootstraps :param X: a feature matrix :type X: np.ndarray :param treatment: a treatment vector :type treatment: np.ndarray :param y: an outcome vector :type y: np.ndarray :param n_bootstraps: number of bootstrap iterations :type n_bootstraps: int :param bootstrap_size: number of samples per bootstrap :type bootstrap_size: int :param n_jobs: number of processes :type n_jobs: int :param verbose: whether to output progress logs :type verbose: bool

Returns:

(np.ndarray), bootstrap estimates

estimate_ate(X: ndarray, treatment: ndarray, y: ndarray, pretrain: bool = False) → tuple[source]#

Estimate the Average Treatment Effect (ATE).

Parameters:

  • X (np.ndarray) – a feature matrix

  • treatment (np.array) – a treatment vector

  • y (np.ndarray) – an outcome vector

  • pretrain (bool) – whether a model has been fit, default False. When True the fitted tree is reused instead of being refit, so these rows can be ones it never saw. The default refits on X and estimates from the same rows, which leaves the estimate carrying whatever the tree overfit.

Returns:

tuple, The mean and confidence interval (LB, UB) of the ATE estimate.

fit(X: ndarray, treatment: ndarray, y: ndarray, sample_weight: ndarray | None = None, check_input: bool = True, prepare_data: bool = True)[source]#

Fit CausalTreeRegressor

Parameters:

  • X (np.ndarray) – feature matrix

  • treatment (np.ndarray) – treatment vector, includes control group

  • y (np.ndarray) – outcome vector

  • sample_weight (np.ndarray) – sample_weight, optional. Weights the split search and the leaf estimates alike, including the honest re-estimation, so inverse-propensity weights passed here correct the confounding bias on observational data — see the note on the class.

  • check_input (bool, optional) – default=False

  • prepare_data (bool) – default=True

Returns:

self

fit_predict(X: ndarray, treatment: ndarray, y: ndarray, return_ci: bool = False, n_bootstraps: int = 1000, bootstrap_size: int = 10000, n_jobs: int = 1, verbose: bool = False) → tuple[source]#

Fit the Causal Tree model and predict treatment effects.

Parameters:

  • X (np.ndarray) – a feature matrix

  • treatment (np.ndarray) – a treatment vector

  • y (np.array) – an outcome vector

  • return_ci (bool) – whether to return confidence intervals

  • n_bootstraps (int) – number of bootstrap iterations

  • bootstrap_size (int) – number of samples per bootstrap

  • n_jobs (int) – the number of jobs for bootstrap

  • verbose (str) – whether to output progress logs

Returns:

  • te (numpy.ndarray): Predictions of treatment effects.

  • te_lower (numpy.ndarray, optional): lower bounds of treatment effects

  • te_upper (numpy.ndarray, optional): upper bounds of treatment effects

Return type:

(tuple)

predict(X: ndarray, with_outcomes: bool = False, check_input=True) → ndarray[source]#

Predict individual treatment effects

Parameters:

  • X (np.ndarray) – a feature matrix

  • with_outcomes (bool) – include outcomes Y_hat(X|T=0), Y_hat(X|T=1),…,Y_hat(X|T=n) along with individual treatment effects

  • check_input (bool) – Allow to bypass several input checking.

Returns:

individual treatment effect (ITE), dim=(samples, groups)

or ITE with outcomes: [Y_hat(X|T=0), Y_hat(X|T=1),…,Y_hat(X|T=n), ITE_1, ITE_2,…,ITE_n], dim=(samples, 2*groups-1)

Return type:

(np.ndarray)

set_fit_request(*, check_input: bool | None | str = '$UNCHANGED$', prepare_data: bool | None | str = '$UNCHANGED$', sample_weight: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → CausalTreeRegressor#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • check_input (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for check_input parameter in fit.

  • prepare_data (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for prepare_data parameter in fit.

  • sample_weight (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for sample_weight parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, check_input: bool | None | str = '$UNCHANGED$', with_outcomes: bool | None | str = '$UNCHANGED$') → CausalTreeRegressor#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • check_input (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for check_input parameter in predict.

  • with_outcomes (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for with_outcomes parameter in predict.

Returns:

self – The updated object.

Return type:

object

set_score_request(*, sample_weight: bool | None | str = '$UNCHANGED$') → CausalTreeRegressor#

Configure whether metadata should be requested to be passed to the score method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to score if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to score.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

sample_weight (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for sample_weight parameter in score.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.tree.UpliftRandomForestClassifier(control_name, n_estimators=10, max_features=10, random_state=None, max_depth=5, min_samples_leaf=100, min_samples_treatment=10, n_reg=10, early_stopping_eval_diff_scale=1, evaluationFunction='KL', normalization=True, honesty=False, estimation_sample_size=0.5, n_jobs=None, joblib_prefer: str = 'threads')[source]#

Bases: _KernelUpliftRandomForestClassifier

Uplift random forest classifier.

Kernel-backed drop-in for the historical Cython UpliftRandomForestClassifier (issue #955 switchover). The legacy constructor names and defaults are preserved: evaluationFunction selects the split criterion, max_features defaults to 10 (clamped to the feature count, as the legacy forest did), and the fitted trees are exposed as uplift_forest. predict returns the per-treatment uplift deltas (full_output=True returns the full frame). Bags kernel-backed uplift trees on scikit-learn’s ForestRegressor (see _KernelUpliftRandomForestClassifier).

early_stopping_eval_diff_scale and fit’s X_val / treatment_val / y_val are accepted for backward compatibility but ignored: validation-set early stopping is not implemented on the kernel trees.

n_jobs defaults to None (one worker). Peak memory during fit grows roughly in proportion to the number of workers, since each concurrent tree fit holds its own working set; see _KernelUpliftRandomForestClassifier for measurements. Results are unaffected by n_jobs.

fit(X: ndarray, treatment: ndarray, y: ndarray, X_val: ndarray = None, treatment_val: ndarray = None, y_val: ndarray = None, sample_weight: ndarray = None)[source]#

Fit the forest. X_val / treatment_val / y_val (legacy early stopping) are accepted for backward compatibility but ignored.

set_fit_request(*, X_val: bool | None | str = '$UNCHANGED$', sample_weight: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', treatment_val: bool | None | str = '$UNCHANGED$', y_val: bool | None | str = '$UNCHANGED$') → UpliftRandomForestClassifier#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • X_val (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for X_val parameter in fit.

  • sample_weight (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for sample_weight parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

  • treatment_val (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment_val parameter in fit.

  • y_val (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for y_val parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, full_output: bool | None | str = '$UNCHANGED$') → UpliftRandomForestClassifier#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

full_output (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for full_output parameter in predict.

Returns:

self – The updated object.

Return type:

object

set_score_request(*, sample_weight: bool | None | str = '$UNCHANGED$') → UpliftRandomForestClassifier#

Configure whether metadata should be requested to be passed to the score method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to score if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to score.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

sample_weight (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for sample_weight parameter in score.

Returns:

self – The updated object.

Return type:

object

property uplift_forest#

The fitted trees (legacy attribute name; aliases estimators_).

class causalml.inference.tree.UpliftTreeClassifier(control_name, max_features=None, max_depth=3, min_samples_leaf=100, min_samples_treatment=10, n_reg=100, early_stopping_eval_diff_scale=1, evaluationFunction='KL', normalization=True, honesty=False, estimation_sample_size=0.5, prune_fraction=None, min_gain=0.0001, prune_rule='maxAbsDiff', random_state=None)[source]#

Bases: _KernelUpliftTreeClassifier

Uplift tree classifier.

Kernel-backed drop-in for the historical Cython UpliftTreeClassifier (issue #955 switchover). The legacy constructor names, defaults, and predict semantics are preserved: evaluationFunction selects the split criterion (KL / ED / Chi / CTS / DDP / IT / CIT / IDDP); predict returns per-group P(Y=1|T=g) including the control column; fill re-annotates the fitted tree with new data. The tree is grown on the shared _tree Cython kernel (see _KernelUpliftTreeClassifier).

early_stopping_eval_diff_scale and fit’s X_val / treatment_val / y_val are accepted for backward compatibility but ignored: validation-set early stopping is not implemented on the kernel tree.

prune_fraction (default None, off) makes fit do the split-and-prune itself: that fraction of the rows is held out stratified on (treatment, outcome), the tree is grown on the rest, and prune() runs on the holdout with min_gain / prune_rule. The pruning rows are taken out before the honest split, so neither the split search nor the estimation half sees them, and n_nodes_before_pruning_ records the size pruning started from. prune() is unchanged for callers managing their own holdout. prune_fraction and estimation_sample_size are held-out fractions, so fit rejects anything outside (0, 1) — including 0.0, which would otherwise read as off.

On make_uplift_classification (n=3000, 6 seeds, max_depth=None, min_samples_leaf=20), prune_fraction=0.3 took held-out qini from -1.74 to -1.44, and to 0.83 combined with honesty=True; an unpruned tree at that depth is badly overfit. Measured on one simulated design, so treat the magnitudes as indicative.

property feature_importances_: ndarray#

Non-negative normalized feature importances (sum to 1).

The kernel’s raw importances can be signed (an uplift split gain is not a monotone impurity decrease); the legacy contract exposes non-negative importances, so magnitudes are taken and renormalized.

fill(X: ndarray, treatment: ndarray, y: ndarray)[source]#

Re-estimate leaf probabilities and node counts on new data.

Routes (X, treatment, y) through the fitted tree and overwrites each leaf’s per-group P(Y=1|T=g) and every node’s per-group counts from the supplied data (typically a validation set) using the existing group encoding – mirroring the legacy UpliftTreeClassifier.fill used to re-annotate the plotted tree. Returns self.

fit(X: ndarray, treatment: ndarray, y: ndarray, X_val: ndarray = None, treatment_val: ndarray = None, y_val: ndarray = None, sample_weight: ndarray | None = None, check_input: bool = True)[source]#

Fit the uplift tree. X_val / treatment_val / y_val (legacy early stopping) are accepted for backward compatibility but ignored.

predict(X: ndarray, check_input: bool = True) → ndarray[source]#

Per-group P(Y=1|T=g) for each row (control in column 0).

Preserves the historical UpliftTreeClassifier.predict return shape (n_samples, n_groups); a leaf with no samples of an otherwise-present group yields a NaN rate under min_samples_treatment=0, which is zero-filled to match the legacy uplift_classification_results.

set_fit_request(*, X_val: bool | None | str = '$UNCHANGED$', check_input: bool | None | str = '$UNCHANGED$', sample_weight: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', treatment_val: bool | None | str = '$UNCHANGED$', y_val: bool | None | str = '$UNCHANGED$') → UpliftTreeClassifier#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • X_val (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for X_val parameter in fit.

  • check_input (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for check_input parameter in fit.

  • sample_weight (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for sample_weight parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

  • treatment_val (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment_val parameter in fit.

  • y_val (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for y_val parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, check_input: bool | None | str = '$UNCHANGED$') → UpliftTreeClassifier#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

check_input (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for check_input parameter in predict.

Returns:

self – The updated object.

Return type:

object

causalml.inference.tree.cat_continuous(x, granularity='Medium')[source]#

Categorize (bin) continuous variable based on percentile.

Parameters:

  • x (list) – Feature values.

  • granularity (string, optional, (default = 'Medium')) – Control the granularity of the bins, optional values are: ‘High’, ‘Medium’, ‘Low’.

Returns:

res – List of percentile bins for the feature value.

Return type:

list

causalml.inference.tree.cat_group(dfx, kpix, n_group=10)[source]#

Category Reduction for Categorical Variables

Parameters:

  • dfx (dataframe) – The inputs data dataframe.

  • kpix (string) – The column of the feature.

  • n_group (int, optional (default = 10)) – The number of top category values to be remained, other category values will be put into “Other”.

Return type:

The transformed categorical feature value list.

causalml.inference.tree.cat_transform(dfx, kpix, kpi1)[source]#

Encoding string features.

Parameters:

  • dfx (dataframe) – The inputs data dataframe.

  • kpix (string) – The column of the feature.

  • kpi1 (list) – The list of feature names.

Returns:

  • dfx (DataFrame) – The updated dataframe containing the encoded data.

  • kpi1 (list) – The updated feature names containing the new dummy feature names.

causalml.inference.tree.cv_fold_index(n, i, k, random_seed=2018)[source]#

Encoding string features.

Parameters:

  • dfx (dataframe) – The inputs data dataframe.

  • kpix (string) – The column of the feature.

  • kpi1 (list) – The list of feature names.

Returns:

  • dfx (DataFrame) – The updated dataframe containing the encoded data.

  • kpi1 (list) – The updated feature names containing the new dummy feature names.

causalml.inference.tree.get_tree_leaves_mask(tree) → ndarray[source]#

Get mask array for tree leaves.

Parameters:

tree (CausalTreeRegressor) – a fitted tree object

Returns:

boolean mask, True at each leaf node

Return type:

(numpy.ndarray)

causalml.inference.tree.kpi_transform(dfx, kpi_combo, kpi_combo_new)[source]#

Feature transformation from continuous feature to binned features for a list of features

Parameters:

  • dfx (DataFrame) – DataFrame containing the features.

  • kpi_combo (list of string) – List of feature names to be transformed

  • kpi_combo_new (list of string) – List of new feature names to be assigned to the transformed features.

Returns:

dfx – Updated DataFrame containing the new features.

Return type:

DataFrame

causalml.inference.tree.plot_dist_tree_leaves_values(tree: CausalTreeRegressor, title: str = 'Leaves values distribution', figsize: tuple = (5, 5), fontsize: int = 12) → None[source]#

Create distplot for tree leaves values :param tree: (CausalTreeRegressor), Tree object :param title: (str), plot title :param figsize: (tuple), figure size :param fontsize: (int), title font size

Returns: None

causalml.inference.tree.uplift_tree_plot(decisionTree, x_names)[source]#

Convert the tree to dot graph for plots.

Parameters:

  • decisionTree (object) – object of DecisionTree class

  • x_names (list) – List of feature names

Return type:

Dot class representing the tree graph.

causalml.inference.tree.uplift_tree_string(decisionTree, x_names)[source]#

Convert the tree to string for print.

Parameters:

  • decisionTree (object) – object of DecisionTree class

  • x_names (list) – List of feature names

Return type:

A string representation of the tree.

causalml.inference.meta module#

class causalml.inference.meta.BaseDRClassifier(learner=None, control_outcome_learner=None, treatment_outcome_learner=None, treatment_effect_learner=None, ate_alpha=0.05, control_name=0)[source]#

Bases: BaseDRLearner

A parent class for DR-learner classifier classes.

predict(X, treatment=None, y=None, p=None, return_components=False, verbose=True)[source]#

Predict treatment effects (classifier variant — uses predict_proba for outcomes).

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix. A pl.LazyFrame is collected once at the start of this method.

  • treatment (np.array, pd.Series, or pl.Series, optional) – a treatment vector. Used for computing classification metrics when y is also provided.

  • y (np.array, pd.Series, or pl.Series, optional) – an outcome vector. Used for computing classification metrics when treatment is also provided.

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1). Currently not used in prediction but kept for API consistency.

  • return_components (bool, optional) – whether to return outcome probabilities for treatment and control groups separately. Defaults to False.

  • verbose (bool, optional) – whether to output progress logs. Defaults to True.

Returns:

predictions of treatment effects. If return_components is True, also returns a dict of predicted probabilities for the control group (yhat_cs) and a dict for the treatment group (yhat_ts).

Return type:

(numpy.ndarray)

set_fit_request(*, p: bool | None | str = '$UNCHANGED$', seed: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → BaseDRClassifier#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • seed (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for seed parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → BaseDRClassifier#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.BaseDRLearner(learner=None, control_outcome_learner=None, treatment_outcome_learner=None, treatment_effect_learner=None, ate_alpha=0.05, control_name=0)[source]#

Bases: BaseLearner

A parent class for DR-learner regressor classes.

A DR-learner estimates treatment effects with machine learning models.

Details of DR-learner are available at Kennedy (2020).

bootstrap(X, treatment, y, p=None, size=10000, rng=None, seed=None)[source]#

Runs a single bootstrap with optional deterministic cross-fit seed.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, or pl.DataFrame) – a feature matrix. Resampled natively via filter_index().

  • treatment (np.array) – a treatment vector (numpy)

  • y (np.array) – an outcome vector (numpy)

  • p (dict, optional) – a dict of {treatment group: propensity scores (numpy)}

  • size (int, optional) – number of samples to draw with replacement

  • rng (np.random.Generator, optional) – random number generator for deterministic resampling

  • seed (int, optional) – random seed for cross-fitting within the resampled fit() call

Returns:

Predictions of treatment effects on the full X

from a model trained on the resampled subset.

Return type:

(numpy.ndarray)

estimate_ate(X, treatment, y, p=None, bootstrap_ci=False, n_bootstraps=1000, bootstrap_size=10000, seed=None, pretrain=False)[source]#

Estimate the Average Treatment Effect (ATE).

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.

  • bootstrap_ci (bool) – whether run bootstrap for confidence intervals

  • n_bootstraps (int) – number of bootstrap iterations

  • bootstrap_size (int) – number of samples per bootstrap

  • seed (int) – random seed for cross-fitting

  • pretrain (bool) – whether a model has been fit, default False.

Returns:

The mean and confidence interval (LB, UB) of the ATE estimate.

fit(X, treatment, y, p=None, seed=None)[source]#

Fit the inference model.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix. A pl.LazyFrame is collected once at the start of this method; the feature matrix is otherwise kept in its native format throughout, including the KFold partitions (sliced via filter_index()).

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.

  • seed (int) – random seed for cross-fitting

fit_predict(X, treatment, y, p=None, return_ci=False, n_bootstraps=1000, bootstrap_size=10000, return_components=False, verbose=True, seed=None)[source]#

Fit the treatment effect and outcome models of the DR learner and predict treatment effects.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.

  • return_ci (bool) – whether to return confidence intervals

  • n_bootstraps (int) – number of bootstrap iterations

  • bootstrap_size (int) – number of samples per bootstrap

  • return_components (bool, optional) – whether to return outcome for treatment and control seperately

  • verbose (str) – whether to output progress logs

  • seed (int) – random seed for cross-fitting

Returns:

Predictions of treatment effects.

Return type:

(numpy.ndarray)

predict(X, treatment=None, y=None, p=None, return_components=False, verbose=True)[source]#

Predict treatment effects.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix. A pl.LazyFrame is collected once at the start of this method.

  • treatment (np.array, pd.Series, or pl.Series, optional) – a treatment vector

  • y (np.array, pd.Series, or pl.Series, optional) – an outcome vector

  • verbose (bool, optional) – whether to output progress logs

Returns:

Predictions of treatment effects.

Return type:

(numpy.ndarray)

set_fit_request(*, p: bool | None | str = '$UNCHANGED$', seed: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → BaseDRLearner#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • seed (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for seed parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → BaseDRLearner#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.BaseDRRegressor(learner=None, control_outcome_learner=None, treatment_outcome_learner=None, treatment_effect_learner=None, ate_alpha=0.05, control_name=0)[source]#

Bases: BaseDRLearner

A parent class for DR-learner regressor classes.

set_fit_request(*, p: bool | None | str = '$UNCHANGED$', seed: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → BaseDRRegressor#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • seed (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for seed parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → BaseDRRegressor#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.BaseRClassifier(outcome_learner=None, effect_learner=None, propensity_learner=LogisticRegressionCV(Cs=4, cv=StratifiedKFold(n_splits=4, random_state=42, shuffle=True), l1_ratios=array([0.001, 0.33366667, 0.66633333, 0.999]), penalty='elasticnet', random_state=42, scoring='neg_log_loss', solver='saga'), ate_alpha=0.05, control_name=0, n_fold=5, random_state=None)[source]#

Bases: BaseRLearner

A parent class for R-learner classifier classes.

fit(X, treatment, y, p=None, sample_weight=None, verbose=True)[source]#

Fit the R-learner classifier (uses predict_proba for outcome estimates).

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix. A pl.LazyFrame is collected once at the start of this method.

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.

  • sample_weight (np.array, pd.Series, or pl.Series, optional) – an array of sample weights indicating the weight of each observation for effect_learner. If None, it assumes equal weight.

  • verbose (bool, optional) – whether to output progress logs

predict(X, p=None, return_components=False)[source]#

Predict treatment effects.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix.

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – propensity scores.

  • return_components (bool) – whether to return nuisance components.

Returns:

numpy.ndarray or tuple

set_fit_request(*, p: bool | None | str = '$UNCHANGED$', sample_weight: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → BaseRClassifier#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • sample_weight (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for sample_weight parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$') → BaseRClassifier#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.BaseRLearner(learner=None, outcome_learner=None, effect_learner=None, propensity_learner=LogisticRegressionCV(Cs=4, cv=StratifiedKFold(n_splits=4, random_state=42, shuffle=True), l1_ratios=array([0.001, 0.33366667, 0.66633333, 0.999]), penalty='elasticnet', random_state=42, scoring='neg_log_loss', solver='saga'), ate_alpha=0.05, control_name=0, n_fold=5, random_state=None, cv_n_jobs=-1)[source]#

Bases: BaseLearner

A parent class for R-learner classes.

An R-learner estimates treatment effects with two machine learning models and the propensity score.

Details of R-learner are available at Nie and Wager (2019).

estimate_ate(X, treatment=None, y=None, p=None, sample_weight=None, bootstrap_ci=False, n_bootstraps=1000, bootstrap_size=10000, pretrain=False)[source]#

Estimate the Average Treatment Effect (ATE).

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix

  • treatment (np.array, pd.Series, or pl.Series) – only needed when pretrain=False, a treatment vector

  • y (np.array, pd.Series, or pl.Series) – only needed when pretrain=False, an outcome vector

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.

  • sample_weight (np.array, pd.Series, or pl.Series, optional) – an array of sample weights indicating the weight of each observation for effect_learner. If None, it assumes equal weight.

  • bootstrap_ci (bool) – whether run bootstrap for confidence intervals

  • n_bootstraps (int) – number of bootstrap iterations

  • bootstrap_size (int) – number of samples per bootstrap

  • pretrain (bool) – whether a model has been fit, default False.

Returns:

The mean and confidence interval (LB, UB) of the ATE estimate.

fit(X, treatment, y, p=None, sample_weight=None, verbose=True)[source]#

Fit the treatment effect and outcome models of the R learner.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix. A pl.LazyFrame is collected once at the start of this method; the feature matrix is otherwise kept in its native format throughout, including the call to cross_val_predict (scikit-learn >= 1.6 accepts pandas and Polars DataFrames natively).

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.

  • sample_weight (np.array, pd.Series, or pl.Series, optional) – an array of sample weights indicating the weight of each observation for effect_learner. If None, it assumes equal weight.

  • verbose (bool, optional) – whether to output progress logs

fit_predict(X, treatment, y, p=None, sample_weight=None, return_ci=False, n_bootstraps=1000, bootstrap_size=10000, return_components=False, verbose=True)[source]#

Fit the R learner and predict treatment effects.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.

  • sample_weight (np.array, pd.Series, or pl.Series, optional) – an array of sample weights indicating the weight of each observation for effect_learner. If None, it assumes equal weight.

  • return_ci (bool) – whether to return confidence intervals

  • n_bootstraps (int) – number of bootstrap iterations

  • bootstrap_size (int) – number of samples per bootstrap

  • return_components (bool, optional) – whether to return the nuisance

  • prediction (outcome)

  • effects. (to treatment)

  • verbose (bool) – whether to output progress logs

Returns:

Predictions of treatment effects.

Return type:

(numpy.ndarray)

predict(X, p=None, return_components=False)[source]#

Predict treatment effects.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix.

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – propensity scores.

  • return_components (bool) – whether to return nuisance components.

Returns:

numpy.ndarray or tuple

set_fit_request(*, p: bool | None | str = '$UNCHANGED$', sample_weight: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → BaseRLearner#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • sample_weight (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for sample_weight parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$') → BaseRLearner#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.BaseRRegressor(learner=None, outcome_learner=None, effect_learner=None, propensity_learner=LogisticRegressionCV(Cs=4, cv=StratifiedKFold(n_splits=4, random_state=42, shuffle=True), l1_ratios=array([0.001, 0.33366667, 0.66633333, 0.999]), penalty='elasticnet', random_state=42, scoring='neg_log_loss', solver='saga'), ate_alpha=0.05, control_name=0, n_fold=5, random_state=None)[source]#

Bases: BaseRLearner

A parent class for R-learner regressor classes.

set_fit_request(*, p: bool | None | str = '$UNCHANGED$', sample_weight: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → BaseRRegressor#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • sample_weight (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for sample_weight parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$') → BaseRRegressor#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.BaseSClassifier(learner=None, ate_alpha=0.05, control_name=0)[source]#

Bases: BaseSLearner

A parent class for S-learner classifier classes.

predict(X, treatment=None, y=None, p=None, return_components=False, verbose=True)[source]#

Predict treatment effects.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix. A pl.LazyFrame is collected once at the start of this method.

  • treatment (np.array, pd.Series, or pl.Series, optional) – a treatment vector

  • y (np.array, pd.Series, or pl.Series, optional) – an outcome vector

  • return_components (bool, optional) – whether to return outcome for treatment and control seperately

  • verbose (bool, optional) – whether to output progress logs

Returns:

Predictions of treatment effects.

Return type:

(numpy.ndarray)

set_fit_request(*, p: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → BaseSClassifier#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → BaseSClassifier#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.BaseSLearner(learner=None, ate_alpha=0.05, control_name=0)[source]#

Bases: BaseLearner

A parent class for S-learner classes.

An S-learner estimates treatment effects with one machine learning model.

Details of S-learner are available at Kunzel et al. (2018).

estimate_ate(X, treatment, y, p=None, bootstrap_ci=False, n_bootstraps=1000, bootstrap_size=10000, pretrain=False)[source]#

Estimate the Average Treatment Effect (ATE).

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • bootstrap_ci (bool) – whether to return confidence intervals

  • n_bootstraps (int) – number of bootstrap iterations

  • bootstrap_size (int) – number of samples per bootstrap

  • pretrain (bool) – whether a model has been fit, default False.

Returns:

The mean and confidence interval (LB, UB) of the ATE estimate.

fit(X, treatment, y, p=None)[source]#

Fit the inference model.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix. A pl.LazyFrame is collected once at the start of this method; the feature matrix is otherwise kept in its native format throughout.

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

fit_predict(X, treatment, y, p=None, return_ci=False, n_bootstraps=1000, bootstrap_size=10000, return_components=False, verbose=True)[source]#

Fit the inference model of the S learner and predict treatment effects.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • return_ci (bool, optional) – whether to return confidence intervals

  • n_bootstraps (int, optional) – number of bootstrap iterations

  • bootstrap_size (int, optional) – number of samples per bootstrap

  • return_components (bool, optional) – whether to return outcome for treatment and control seperately

  • verbose (bool, optional) – whether to output progress logs

Returns:

Predictions of treatment effects. Output dim: [n_samples, n_treatment].

If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment], UB [n_samples, n_treatment]

Return type:

(numpy.ndarray)

predict(X, treatment=None, y=None, p=None, return_components=False, verbose=True)[source]#

Predict treatment effects.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix. A pl.LazyFrame is collected once at the start of this method.

  • treatment (np.array, pd.Series, or pl.Series, optional) – a treatment vector

  • y (np.array, pd.Series, or pl.Series, optional) – an outcome vector

  • return_components (bool, optional) – whether to return outcome for treatment and control seperately

  • verbose (bool, optional) – whether to output progress logs

Returns:

Predictions of treatment effects.

Return type:

(numpy.ndarray)

set_fit_request(*, p: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → BaseSLearner#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → BaseSLearner#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.BaseSRegressor(learner=None, ate_alpha=0.05, control_name=0)[source]#

Bases: BaseSLearner

A parent class for S-learner regressor classes.

set_fit_request(*, p: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → BaseSRegressor#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → BaseSRegressor#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.BaseTClassifier(learner=None, control_learner=None, treatment_learner=None, ate_alpha=0.05, control_name=0)[source]#

Bases: BaseTLearner

A parent class for T-learner classifier classes.

predict(X, treatment=None, y=None, p=None, return_components=False, verbose=True, return_ci=False)[source]#

Predict treatment effects.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix. A pl.LazyFrame is collected once at the start of this method.

  • treatment (np.array, pd.Series, or pl.Series, optional) – a treatment vector

  • y (np.array, pd.Series, or pl.Series, optional) – an outcome vector

  • return_components (bool, optional) – whether to return outcome for treatment and control seperately

  • verbose (bool, optional) – whether to output progress logs

  • return_ci (bool, optional) – whether to return confidence intervals using the stored bootstrap ensemble.

Returns:

Predictions of treatment effects.

Return type:

(numpy.ndarray)

set_fit_request(*, bootstrap_size: bool | None | str = '$UNCHANGED$', n_bootstraps: bool | None | str = '$UNCHANGED$', n_jobs: bool | None | str = '$UNCHANGED$', p: bool | None | str = '$UNCHANGED$', random_state: bool | None | str = '$UNCHANGED$', store_bootstraps: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → BaseTClassifier#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • bootstrap_size (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for bootstrap_size parameter in fit.

  • n_bootstraps (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for n_bootstraps parameter in fit.

  • n_jobs (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for n_jobs parameter in fit.

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • random_state (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for random_state parameter in fit.

  • store_bootstraps (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for store_bootstraps parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_ci: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → BaseTClassifier#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_ci (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_ci parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.BaseTLearner(learner=None, control_learner=None, treatment_learner=None, ate_alpha=0.05, control_name=0)[source]#

Bases: BaseLearner

A parent class for T-learner regressor classes.

A T-learner estimates treatment effects with two machine learning models.

Details of T-learner are available at Kunzel et al. (2018).

estimate_ate(X, treatment, y, p=None, bootstrap_ci=False, n_bootstraps=1000, bootstrap_size=10000, pretrain=False)[source]#

Estimate the Average Treatment Effect (ATE).

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • bootstrap_ci (bool) – whether to return confidence intervals

  • n_bootstraps (int) – number of bootstrap iterations

  • bootstrap_size (int) – number of samples per bootstrap

  • pretrain (bool) – whether a model has been fit, default False.

Returns:

The mean and confidence interval (LB, UB) of the ATE estimate.

fit(X, treatment, y, p=None, store_bootstraps=False, n_bootstraps=200, bootstrap_size=10000, random_state=None, n_jobs=1)[source]#

Fit the inference model.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix. A pl.LazyFrame is collected once at the start of this method; the feature matrix is otherwise kept in its native format throughout.

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • p – unused, kept for API consistency

  • store_bootstraps (bool, optional) – if True, trains a bootstrap ensemble during fit and stores it in self.bootstrap_models_ for post-fit CI estimation via predict(return_ci=True). Default: False.

  • n_bootstraps (int, optional) – number of bootstrap iterations. Default: 200.

  • n_jobs (int, optional) – number of parallel jobs for bootstrap fitting. -1 uses all available cores. Default: 1.

  • bootstrap_size (int, optional) – number of samples per bootstrap. Default: 10000.

  • random_state (int, optional) – random seed for reproducible bootstrap sampling.

fit_predict(X, treatment, y, p=None, return_ci=False, n_bootstraps=1000, bootstrap_size=10000, return_components=False, verbose=True)[source]#

Fit the inference model of the T learner and predict treatment effects.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • return_ci (bool) – whether to return confidence intervals

  • n_bootstraps (int) – number of bootstrap iterations

  • bootstrap_size (int) – number of samples per bootstrap

  • return_components (bool, optional) – whether to return outcome for treatment and control seperately

  • verbose (str) – whether to output progress logs

Returns:

Predictions of treatment effects. Output dim: [n_samples, n_treatment].

If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment], UB [n_samples, n_treatment]

Return type:

(numpy.ndarray)

predict(X, treatment=None, y=None, p=None, return_components=False, verbose=True, return_ci=False)[source]#

Predict treatment effects.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix. A pl.LazyFrame is collected once at the start of this method.

  • treatment (np.array, pd.Series, or pl.Series, optional) – a treatment vector

  • y (np.array, pd.Series, or pl.Series, optional) – an outcome vector

  • return_components (bool, optional) – whether to return outcome for treatment and control seperately

  • verbose (bool, optional) – whether to output progress logs

  • return_ci (bool, optional) – whether to return confidence intervals using the stored bootstrap ensemble. Requires fit() to have been called with store_bootstraps=True.

Returns:

Predictions of treatment effects. If return_ci=True,

returns (te, te_lower, te_upper) each of shape [n_samples, n_treatment].

Return type:

(numpy.ndarray)

set_fit_request(*, bootstrap_size: bool | None | str = '$UNCHANGED$', n_bootstraps: bool | None | str = '$UNCHANGED$', n_jobs: bool | None | str = '$UNCHANGED$', p: bool | None | str = '$UNCHANGED$', random_state: bool | None | str = '$UNCHANGED$', store_bootstraps: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → BaseTLearner#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • bootstrap_size (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for bootstrap_size parameter in fit.

  • n_bootstraps (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for n_bootstraps parameter in fit.

  • n_jobs (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for n_jobs parameter in fit.

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • random_state (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for random_state parameter in fit.

  • store_bootstraps (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for store_bootstraps parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_ci: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → BaseTLearner#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_ci (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_ci parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.BaseTRegressor(learner=None, control_learner=None, treatment_learner=None, ate_alpha=0.05, control_name=0)[source]#

Bases: BaseTLearner

A parent class for T-learner regressor classes.

set_fit_request(*, bootstrap_size: bool | None | str = '$UNCHANGED$', n_bootstraps: bool | None | str = '$UNCHANGED$', n_jobs: bool | None | str = '$UNCHANGED$', p: bool | None | str = '$UNCHANGED$', random_state: bool | None | str = '$UNCHANGED$', store_bootstraps: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → BaseTRegressor#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • bootstrap_size (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for bootstrap_size parameter in fit.

  • n_bootstraps (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for n_bootstraps parameter in fit.

  • n_jobs (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for n_jobs parameter in fit.

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • random_state (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for random_state parameter in fit.

  • store_bootstraps (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for store_bootstraps parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_ci: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → BaseTRegressor#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_ci (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_ci parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.BaseXClassifier(outcome_learner=None, effect_learner=None, control_outcome_learner=None, treatment_outcome_learner=None, control_effect_learner=None, treatment_effect_learner=None, ate_alpha=0.05, control_name=0)[source]#

Bases: BaseXLearner

A parent class for X-learner classifier classes.

fit(X, treatment, y, p=None)[source]#

Fit the inference model.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix. A pl.LazyFrame is collected once at the start of this method; the feature matrix is otherwise kept in its native format throughout.

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.

predict(X, treatment=None, y=None, p=None, return_components=False, verbose=True)[source]#

Predict treatment effects (classifier variant — uses predict_proba).

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix. A pl.LazyFrame is collected once at the start of this method.

  • treatment (np.array, pd.Series, or pl.Series, optional) – a treatment vector

  • y (np.array, pd.Series, or pl.Series, optional) – an outcome vector

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.

  • return_components (bool, optional) – whether to return outcome for treatment and control seperately

  • verbose (bool, optional) – whether to output progress logs

Returns:

Predictions of treatment effects.

Return type:

(numpy.ndarray)

set_fit_request(*, p: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → BaseXClassifier#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → BaseXClassifier#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.BaseXLearner(learner=None, control_outcome_learner=None, treatment_outcome_learner=None, control_effect_learner=None, treatment_effect_learner=None, ate_alpha=0.05, control_name=0)[source]#

Bases: BaseLearner

A parent class for X-learner regressor classes.

An X-learner estimates treatment effects with four machine learning models.

Details of X-learner are available at Kunzel et al. (2018).

estimate_ate(X, treatment, y, p=None, bootstrap_ci=False, n_bootstraps=1000, bootstrap_size=10000, pretrain=False)[source]#

Estimate the Average Treatment Effect (ATE).

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.

  • bootstrap_ci (bool) – whether run bootstrap for confidence intervals

  • n_bootstraps (int) – number of bootstrap iterations

  • bootstrap_size (int) – number of samples per bootstrap

  • pretrain (bool) – whether a model has been fit, default False.

Returns:

The mean and confidence interval (LB, UB) of the ATE estimate.

fit(X, treatment, y, p=None)[source]#

Fit the inference model.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix. A pl.LazyFrame is collected once at the start of this method; the feature matrix is otherwise kept in its native format throughout.

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.

fit_predict(X, treatment, y, p=None, return_ci=False, n_bootstraps=1000, bootstrap_size=10000, return_components=False, verbose=True)[source]#

Fit the X-learner and predict treatment effects.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.

  • return_ci (bool) – whether to return confidence intervals

  • n_bootstraps (int) – number of bootstrap iterations

  • bootstrap_size (int) – number of samples per bootstrap

  • return_components (bool, optional) – whether to return outcome for treatment and control seperately

  • verbose (str) – whether to output progress logs

Returns:

Predictions of treatment effects.

Return type:

(numpy.ndarray)

predict(X, treatment=None, y=None, p=None, return_components=False, verbose=True)[source]#

Predict treatment effects.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix. A pl.LazyFrame is collected once at the start of this method.

  • treatment (np.array, pd.Series, or pl.Series, optional) – a treatment vector

  • y (np.array, pd.Series, or pl.Series, optional) – an outcome vector

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.

  • return_components (bool, optional) – whether to return outcome for treatment and control seperately

  • verbose (bool, optional) – whether to output progress logs

Returns:

Predictions of treatment effects.

Return type:

(numpy.ndarray)

set_fit_request(*, p: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → BaseXLearner#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → BaseXLearner#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.BaseXRegressor(learner=None, control_outcome_learner=None, treatment_outcome_learner=None, control_effect_learner=None, treatment_effect_learner=None, ate_alpha=0.05, control_name=0)[source]#

Bases: BaseXLearner

A parent class for X-learner regressor classes.

set_fit_request(*, p: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → BaseXRegressor#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → BaseXRegressor#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

exception causalml.inference.meta.CausalMLVersionMismatchWarning[source]#

Bases: UserWarning

Raised when a saved model was created with a different causalml version.

class causalml.inference.meta.LRSRegressor(ate_alpha=0.05, control_name=0)[source]#

Bases: BaseSRegressor

estimate_ate(X, treatment, y, p=None, pretrain=False)[source]#

Estimate the Average Treatment Effect (ATE).

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • pretrain (bool) – whether a model has been fit, default False.

Returns:

The mean and confidence interval (LB, UB) of the ATE estimate.

set_fit_request(*, p: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → LRSRegressor#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → LRSRegressor#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.MLPTRegressor(ate_alpha=0.05, control_name=0, *args, **kwargs)[source]#

Bases: BaseTRegressor

set_fit_request(*, bootstrap_size: bool | None | str = '$UNCHANGED$', n_bootstraps: bool | None | str = '$UNCHANGED$', n_jobs: bool | None | str = '$UNCHANGED$', p: bool | None | str = '$UNCHANGED$', random_state: bool | None | str = '$UNCHANGED$', store_bootstraps: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → MLPTRegressor#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • bootstrap_size (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for bootstrap_size parameter in fit.

  • n_bootstraps (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for n_bootstraps parameter in fit.

  • n_jobs (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for n_jobs parameter in fit.

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • random_state (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for random_state parameter in fit.

  • store_bootstraps (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for store_bootstraps parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_ci: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → MLPTRegressor#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_ci (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_ci parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.TMLELearner(learner, ate_alpha=0.05, control_name=0, cv=None)[source]#

Bases: object

Targeted maximum likelihood estimation.

Ref: Gruber, S., & Van Der Laan, M. J. (2009). Targeted maximum likelihood estimation: A gentle introduction.

estimate_ate(X, treatment, y, p, segment=None, return_ci=False)[source]#

Estimate the Average Treatment Effect (ATE).

Parameters:

  • X (np.matrix or np.array or pd.Dataframe) – a feature matrix

  • treatment (np.array or pd.Series) – a treatment vector

  • y (np.array or pd.Series) – an outcome vector

  • p (np.ndarray or pd.Series or dict) – an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1)

  • segment (np.array, optional) – An optional segment vector of int. If given, the ATE and its CI will be estimated for each segment.

  • return_ci (bool, optional) – Whether to return confidence intervals

Returns:

The ATE and its confidence interval (LB, UB) for each treatment, t and segment, s

Return type:

(tuple)

class causalml.inference.meta.XGBDRRegressor(ate_alpha=0.05, control_name=0, *args, **kwargs)[source]#

Bases: BaseDRRegressor

set_fit_request(*, p: bool | None | str = '$UNCHANGED$', seed: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → XGBDRRegressor#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • seed (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for seed parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → XGBDRRegressor#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.XGBRClassifier(propensity_learner=LogisticRegressionCV(Cs=4, cv=StratifiedKFold(n_splits=4, random_state=42, shuffle=True), l1_ratios=array([0.001, 0.33366667, 0.66633333, 0.999]), penalty='elasticnet', random_state=42, scoring='neg_log_loss', solver='saga'), ate_alpha=0.05, control_name=0, n_fold=5, random_state=None, outcome_xgb_kwargs=None, effect_xgb_kwargs=None)[source]#

Bases: BaseRClassifier

An R-learner classifier using XGBoost models.

The outcome model is an XGBClassifier (its predict_proba drives the outcome cross-fit) and the effect model is an XGBRegressor. Every constructor argument is stored verbatim (scikit-learn convention) so that get_params() / clone() work correctly; the XGBoost models are constructed in fit().

fit(X, treatment, y, p=None, sample_weight=None, verbose=True)[source]#

Build the XGBoost models, then fit as an R-learner classifier.

set_fit_request(*, p: bool | None | str = '$UNCHANGED$', sample_weight: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → XGBRClassifier#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • sample_weight (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for sample_weight parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$') → XGBRClassifier#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.XGBRRegressor(early_stopping=True, test_size=0.3, early_stopping_rounds=30, effect_learner_objective='reg:squarederror', effect_learner_n_estimators=500, random_state=42, ate_alpha=0.05, control_name=0, n_fold=5, xgb_kwargs=None)[source]#

Bases: BaseRRegressor

An R-learner regressor using XGBoost models.

Stores every constructor argument verbatim (scikit-learn convention) so that get_params() / clone() work correctly. All XGBRegressor construction is deferred to fit().

Additional XGBoost keyword arguments (e.g. max_depth, learning_rate) are accepted via **xgb_kwargs and stored verbatim as self.xgb_kwargs, so that get_params() surfaces them and clone() round-trips them correctly.

fit(X, treatment, y, p=None, sample_weight=None, verbose=True)[source]#

Fit using early-stopping XGBoost R-learner.

Parameters:

  • X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame) – a feature matrix. A pl.LazyFrame is collected once at the start of this method.

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector

  • y (np.array, pd.Series, or pl.Series) – an outcome vector

  • p (np.ndarray, pd.Series, pl.Series, or dict, optional) – an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.

  • sample_weight (np.array, pd.Series, or pl.Series, optional) – an array of sample weights indicating the weight of each observation for effect_learner. If None, it assumes equal weight.

  • verbose (bool, optional) – whether to output progress logs

set_fit_request(*, p: bool | None | str = '$UNCHANGED$', sample_weight: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → XGBRRegressor#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • sample_weight (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for sample_weight parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$') → XGBRRegressor#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.XGBTClassifier(ate_alpha=0.05, control_name=0, xgb_kwargs=None)[source]#

Bases: BaseTClassifier

A T-learner classifier using XGBoost models.

Stores XGBoost hyperparameters verbatim (scikit-learn convention) so that get_params() / clone() work correctly; the XGBClassifier models are constructed in fit().

fit(X, treatment, y, *args, **kwargs)[source]#

Build the XGBoost outcome model, then fit as a T-learner classifier.

set_fit_request(*, treatment: bool | None | str = '$UNCHANGED$') → XGBTClassifier#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_ci: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → XGBTClassifier#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_ci (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_ci parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

class causalml.inference.meta.XGBTRegressor(ate_alpha=0.05, control_name=0, *args, **kwargs)[source]#

Bases: BaseTRegressor

set_fit_request(*, bootstrap_size: bool | None | str = '$UNCHANGED$', n_bootstraps: bool | None | str = '$UNCHANGED$', n_jobs: bool | None | str = '$UNCHANGED$', p: bool | None | str = '$UNCHANGED$', random_state: bool | None | str = '$UNCHANGED$', store_bootstraps: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$') → XGBTRegressor#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • bootstrap_size (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for bootstrap_size parameter in fit.

  • n_bootstraps (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for n_bootstraps parameter in fit.

  • n_jobs (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for n_jobs parameter in fit.

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in fit.

  • random_state (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for random_state parameter in fit.

  • store_bootstraps (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for store_bootstraps parameter in fit.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, p: bool | None | str = '$UNCHANGED$', return_ci: bool | None | str = '$UNCHANGED$', return_components: bool | None | str = '$UNCHANGED$', treatment: bool | None | str = '$UNCHANGED$', verbose: bool | None | str = '$UNCHANGED$') → XGBTRegressor#

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • p (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for p parameter in predict.

  • return_ci (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_ci parameter in predict.

  • return_components (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for return_components parameter in predict.

  • treatment (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for treatment parameter in predict.

  • verbose (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for verbose parameter in predict.

Returns:

self – The updated object.

Return type:

object

causalml.inference.meta.load_learner(path)[source]#

Load any saved causal learner without specifying the class.

This is a convenience function that skips the class-match check, useful when you don’t know which learner type was saved.

Parameters:

path (str) – file path to the saved model.

Returns:

The restored learner instance.

causalml.inference.iv module#

class causalml.inference.iv.BaseDRIVLearner(learner=None, control_outcome_learner=None, treatment_outcome_learner=None, treatment_effect_learner=None, ate_alpha=0.05, control_name=0)[source]#

Bases: SerializableLearner

A parent class for DRIV-learner regressor classes.

A DRIV-learner estimates endogenous treatment effects for compliers with machine learning models.

Details of DR-learner are available at Kennedy (2020). The DR moment condition for LATE comes from Chernozhukov et al (2018).

bootstrap(X, assignment, treatment, y, p, pZ, size=10000, seed=None)[source]#

Runs a single bootstrap. Fits on bootstrapped sample, then predicts on whole population.

estimate_ate(X, assignment, treatment, y, p=None, pZ=None, bootstrap_ci=False, n_bootstraps=1000, bootstrap_size=10000, seed=None, calibrate=True, pretrain=False)[source]#

Estimate the Average Treatment Effect (ATE) for compliers.

Parameters:

  • X (np.matrix or np.array or pd.Dataframe) – a feature matrix

  • assignment (np.array or pd.Series) – an assignment vector. The assignment is the instrumental variable that does not depend on unknown confounders. The assignment status influences treatment in a monotonic way, i.e. one can only be more likely to take the treatment if assigned.

  • treatment (np.array or pd.Series) – a treatment vector

  • y (np.array or pd.Series) – an outcome vector

  • p (2-tuple of np.ndarray or pd.Series or dict, optional) – The first (second) element corresponds to unassigned (assigned) units. Each is an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1). If None will run ElasticNetPropensityModel() to generate the propensity scores.

  • pZ (np.array or pd.Series, optional) – an array of assignment probability of float (0,1); if None will run ElasticNetPropensityModel() to generate the assignment probability score.

  • bootstrap_ci (bool) – whether run bootstrap for confidence intervals

  • n_bootstraps (int) – number of bootstrap iterations

  • bootstrap_size (int) – number of samples per bootstrap

  • seed (int) – random seed for cross-fitting

  • pretrain (bool) – whether a model has been fit, default False. When True the fitted models are reused instead of being refit, so these rows can be ones they never saw. p is required in that case: the propensity fit() stored belongs to the rows it was fit on, and the standard error needs the propensity of the rows estimated here. assignment and pZ are only used to fit, so they go unused.

Returns:

The mean and confidence interval (LB, UB) of the ATE estimate.

fit(X, assignment, treatment, y, p=None, pZ=None, seed=None, calibrate=True)[source]#

Fit the inference model.

Parameters:

  • X (np.matrix or np.array or pd.Dataframe) – a feature matrix

  • assignment (np.array or pd.Series) – a (0,1)-valued assignment vector. The assignment is the instrumental variable that does not depend on unknown confounders. The assignment status influences treatment in a monotonic way, i.e. one can only be more likely to take the treatment if assigned.

  • treatment (np.array or pd.Series) – a treatment vector

  • y (np.array or pd.Series) – an outcome vector

  • p (2-tuple of np.ndarray or pd.Series or dict, optional) – The first (second) element corresponds to unassigned (assigned) units. Each is an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1). If None will run ElasticNetPropensityModel() to generate the propensity scores.

  • pZ (np.array or pd.Series, optional) – an array of assignment probability of float (0,1); if None will run ElasticNetPropensityModel() to generate the assignment probability score.

  • seed (int) – random seed for cross-fitting

fit_predict(X, assignment, treatment, y, p=None, pZ=None, return_ci=False, n_bootstraps=1000, bootstrap_size=10000, return_components=False, verbose=True, seed=None, calibrate=True)[source]#

Fit the treatment effect and outcome models of the R learner and predict treatment effects.

Parameters:

  • X (np.matrix or np.array or pd.Dataframe) – a feature matrix

  • assignment (np.array or pd.Series) – a (0,1)-valued assignment vector. The assignment is the instrumental variable that does not depend on unknown confounders. The assignment status influences treatment in a monotonic way, i.e. one can only be more likely to take the treatment if assigned.

  • treatment (np.array or pd.Series) – a treatment vector

  • y (np.array or pd.Series) – an outcome vector

  • p (2-tuple of np.ndarray or pd.Series or dict, optional) – The first (second) element corresponds to unassigned (assigned) units. Each is an array of propensity scores of float (0,1) in the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1). If None will run ElasticNetPropensityModel() to generate the propensity scores.

  • pZ (np.array or pd.Series, optional) – an array of assignment probability of float (0,1); if None will run ElasticNetPropensityModel() to generate the assignment probability score.

  • return_ci (bool) – whether to return confidence intervals

  • n_bootstraps (int) – number of bootstrap iterations

  • bootstrap_size (int) – number of samples per bootstrap

  • return_components (bool, optional) – whether to return outcome for treatment and control seperately

  • verbose (str) – whether to output progress logs

  • seed (int) – random seed for cross-fitting

Returns:

Predictions of treatment effects for compliers, , i.e. those individuals

who take the treatment only if they are assigned. Output dim: [n_samples, n_treatment] If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment], UB [n_samples, n_treatment]

Return type:

(numpy.ndarray)

get_importance(X=None, tau=None, model_tau_feature=None, features=None, method='auto', normalize=True, test_size=0.3, random_state=None)[source]#

Builds a model (using X to predict estimated/actual tau), and then calculates feature importances based on a specified method.

Currently supported methods are:

  • auto (calculates importance based on estimator’s default implementation of feature importance; estimator must be tree-based). Note: if none provided, it uses lightgbm’s LGBMRegressor as estimator, and “gain” as importance type.

  • permutation (calculates importance based on mean decrease in accuracy when a feature column is permuted; estimator can be any form).

Hint: for permutation, downsample data for better performance especially if X.shape[1] is large.

Parameters:

  • X (np.matrix or np.array or pd.Dataframe) – a feature matrix

  • tau (np.array) – a treatment effect vector (estimated/actual)

  • model_tau_feature (sklearn/lightgbm/xgboost model object) – an unfitted model object

  • features (np.array) – list/array of feature names. If None, an enumerated list will be used

  • method (str) – auto, permutation

  • normalize (bool) – normalize by sum of importances if method=auto (defaults to True)

  • test_size (float/int) – if float, represents the proportion of the dataset to include in the test split. If int, represents the absolute number of test samples (used for estimating permutation importance)

  • random_state (int/RandomState instance/None) – random state used in permutation importance estimation

get_shap_values(X=None, model_tau_feature=None, tau=None, features=None)[source]#

Builds a model (using X to predict estimated/actual tau), and then calculates shapley values. :param X: a feature matrix :type X: np.matrix or np.array or pd.Dataframe :param tau: a treatment effect vector (estimated/actual) :type tau: np.array :param model_tau_feature: an unfitted model object :type model_tau_feature: sklearn/lightgbm/xgboost model object :param features: list/array of feature names. If None, an enumerated list will be used. :type features: optional, np.array

plot_importance(X=None, tau=None, model_tau_feature=None, features=None, method='auto', normalize=True, test_size=0.3, random_state=None)[source]#

Builds a model (using X to predict estimated/actual tau), and then plots feature importances based on a specified method.

Currently supported methods are:

  • auto (calculates importance based on estimator’s default implementation of feature importance; estimator must be tree-based). Note: if none provided, it uses lightgbm’s LGBMRegressor as estimator, and “gain” as importance type.

  • permutation (calculates importance based on mean decrease in accuracy when a feature column is permuted; estimator can be any form).

Hint: for permutation, downsample data for better performance especially if X.shape[1] is large.

Parameters:

  • X (np.matrix or np.array or pd.Dataframe) – a feature matrix

  • tau (np.array) – a treatment effect vector (estimated/actual)

  • model_tau_feature (sklearn/lightgbm/xgboost model object) – an unfitted model object

  • features (optional, np.array) – list/array of feature names. If None, an enumerated list will be used

  • method (str) – auto, permutation

  • normalize (bool) – normalize by sum of importances if method=auto (defaults to True)

  • test_size (float/int) – if float, represents the proportion of the dataset to include in the test split. If int, represents the absolute number of test samples (used for estimating permutation importance)

  • random_state (int/RandomState instance/None) – random state used in permutation importance estimation

plot_shap_dependence(treatment_group, feature_idx, X, tau, model_tau_feature=None, features=None, shap_dict=None, interaction_idx='auto', **kwargs)[source]#

Plots dependency of shapley values for a specified feature, colored by an interaction feature.

If shapley values have been pre-computed, pass it through the shap_dict parameter. If shap_dict is not provided, this builds a new model (using X to predict estimated/actual tau), and then calculates shapley values.

This plots the value of the feature on the x-axis and the SHAP value of the same feature on the y-axis. This shows how the model depends on the given feature, and is like a richer extension of the classical partial dependence plots. Vertical dispersion of the data points represents interaction effects.

Parameters:

  • treatment_group (str or int) – name of treatment group to create dependency plot on

  • feature_idx (str or int) – feature index / name to create dependency plot on

  • X (np.matrix or np.array or pd.Dataframe) – a feature matrix

  • tau (np.array) – a treatment effect vector (estimated/actual)

  • model_tau_feature (sklearn/lightgbm/xgboost model object) – an unfitted model object

  • features (optional, np.array) – list/array of feature names. If None, an enumerated list will be used.

  • shap_dict (optional, dict) – a dict of shapley value matrices. If None, shap_dict will be computed.

  • interaction_idx (optional, str or int) – feature index / name used in coloring scheme as interaction feature. If “auto” then shap.common.approximate_interactions is used to pick what seems to be the strongest interaction (note that to find to true strongest interaction you need to compute the SHAP interaction values).

plot_shap_values(X=None, tau=None, model_tau_feature=None, features=None, shap_dict=None, **kwargs)[source]#

Plots distribution of shapley values.

If shapley values have been pre-computed, pass it through the shap_dict parameter. If shap_dict is not provided, this builds a new model (using X to predict estimated/actual tau), and then calculates shapley values.

Parameters:

  • X (np.matrix or np.array or pd.Dataframe) – a feature matrix. Required if shap_dict is None.

  • tau (np.array) – a treatment effect vector (estimated/actual)

  • model_tau_feature (sklearn/lightgbm/xgboost model object) – an unfitted model object

  • features (optional, np.array) – list/array of feature names. If None, an enumerated list will be used.

  • shap_dict (optional, dict) – a dict of shapley value matrices. If None, shap_dict will be computed.

predict(X, treatment=None, y=None, return_components=False, verbose=True)[source]#

Predict treatment effects.

Parameters:

  • X (np.matrix or np.array or pd.Dataframe) – a feature matrix

  • treatment (np.array or pd.Series, optional) – a treatment vector

  • y (np.array or pd.Series, optional) – an outcome vector

  • verbose (bool, optional) – whether to output progress logs

Returns:

Predictions of treatment effects for compliers, i.e. those individuals

who take the treatment only if they are assigned.

Return type:

(numpy.ndarray)

class causalml.inference.iv.BaseDRIVRegressor(learner=None, control_outcome_learner=None, treatment_outcome_learner=None, treatment_effect_learner=None, ate_alpha=0.05, control_name=0)[source]#

Bases: BaseDRIVLearner

A parent class for DRIV-learner regressor classes.

class causalml.inference.iv.IVRegressor[source]#

Bases: SerializableLearner

A wrapper class that uses IV2SLS from statsmodel

A linear 2SLS model that estimates the average treatment effect with endogenous treatment variable.

fit(X, treatment, y, w)[source]#

Fits the 2SLS model.

Parameters:

  • X (np.matrix or np.array or pd.Dataframe) – a feature matrix

  • treatment (np.array or pd.Series) – a treatment vector

  • y (np.array or pd.Series) – an outcome vector

  • w (np.array or pd.Series) – an instrument vector

predict()[source]#

Returns the average treatment effect and its estimated standard error

Returns:

average treatment effect (float): standard error of the estimation

Return type:

(float)

class causalml.inference.iv.XGBDRIVRegressor(ate_alpha=0.05, control_name=0, *args, **kwargs)[source]#

Bases: BaseDRIVRegressor

causalml.inference.torch module#

class causalml.inference.torch.CEVAE(outcome_dist='studentt', latent_dim=20, hidden_dim=200, num_epochs=50, num_layers=3, batch_size=100, learning_rate=0.001, learning_rate_decay=0.1, num_samples=1000, weight_decay=0.0001)[source]#

Bases: object

fit(X, treatment, y, p=None)[source]#

Fits CEVAE.

Parameters:

  • X (np.matrix or np.array or pd.Dataframe) – a feature matrix

  • treatment (np.array or pd.Series) – a treatment vector

  • y (np.array or pd.Series) – an outcome vector

fit_predict(X, treatment, y, p=None)[source]#

Fits the CEVAE model and then predicts.

Parameters:

  • X (np.matrix or np.array or pd.Dataframe) – a feature matrix

  • treatment (np.array or pd.Series) – a treatment vector

  • y (np.array or pd.Series) – an outcome vector

Returns:

Predictions of treatment effects.

Return type:

(np.ndarray)

predict(X, treatment=None, y=None, p=None)[source]#

Calls predict on fitted DragonNet.

Parameters:

X (np.matrix or np.array or pd.Dataframe) – a feature matrix

Returns:

Predictions of treatment effects.

Return type:

(np.ndarray)

causalml.inference.tf module#

class causalml.inference.tf.DragonNet(neurons_per_layer=200, targeted_reg=True, ratio=1.0, val_split=0.2, batch_size=64, epochs=100, learning_rate=1e-05, momentum=0.9, reg_l2=0.01, use_adam=True, adam_epochs=30, adam_learning_rate=0.001, loss_func=<function dragonnet_loss_binarycross>, verbose=True)[source]#

Bases: object

fit(X, treatment, y, p=None)[source]#

Fits the DragonNet model.

Parameters:

  • X (np.matrix or np.array or pd.Dataframe) – a feature matrix

  • treatment (np.array or pd.Series) – a treatment vector

  • y (np.array or pd.Series) – an outcome vector

fit_predict(X, treatment, y, p=None, return_components=False)[source]#

Fits the DragonNet model and then predicts.

Parameters:

  • X (np.matrix or np.array or pd.Dataframe) – a feature matrix

  • treatment (np.array or pd.Series) – a treatment vector

  • y (np.array or pd.Series) – an outcome vector

  • return_components (bool, optional) – whether to return

Returns:

predictions based on return_components flag

if return_components=False (default), each row is treatment effect if return_components=True, each row is (outcome do(t=0), outcome do(t=1), propensity, epsilon)

Return type:

(np.array)

load(h5_filepath, ratio=1.0, dragonnet_loss=<function dragonnet_loss_binarycross>)[source]#

Load the dragonnet model from a H5 file.

Parameters:

  • h5_filepath (H5 file path) – H5 file path

  • ratio (float) – weight assigned to the targeted regularization loss component

  • dragonnet_loss (function) – a loss function

make_dragonnet(input_dim)[source]#

Neural net predictive model. The dragon has three heads.

Parameters:

input_dim (int) – number of rows in input

Returns:

DragonNet model

Return type:

model (keras.models.Model)

predict(X, treatment=None, y=None, p=None)[source]#

Calls predict on fitted DragonNet.

Parameters:

X (np.matrix or np.array or pd.Dataframe) – a feature matrix

Returns:

a 2D array with shape (X.shape[0], 4),

where each row takes the form of (outcome do(t=0), outcome do(t=1), propensity, epsilon)

Return type:

(np.array)

predict_propensity(X)[source]#

Predicts the individual propensity scores.

Parameters:

X (np.matrix or np.array or pd.Dataframe) – a feature matrix

Returns:

propensity score vector

Return type:

(np.array)

predict_tau(X)[source]#

Predicts the individual treatment effect (tau / “ITE”).

Parameters:

X (np.matrix or np.array or pd.Dataframe) – a feature matrix

Returns:

treatment effect vector

Return type:

(np.array)

save(h5_filepath)[source]#

Save the dragonnet model as a H5 file.

Parameters:

h5_filepath (H5 file path) – H5 file path

causalml.inference.jax module#

class causalml.inference.jax.CEVAE(outcome_dist: str = 'studentt', latent_dim: int = 20, hidden_dim: int = 200, num_epochs: int = 50, num_layers: int = 3, batch_size: int = 100, learning_rate: float = 0.001, learning_rate_decay: float = 0.1, num_samples: int = 1000, weight_decay: float = 0.0001, seed: int = 0)[source]#

Bases: object

JAX/flax.nnx CEVAE for treatment-effect estimation.

Mirrors the API of causalml.inference.torch.cevae.CEVAE.

Parameters:

  • outcome_dist – Outcome distribution as one of "bernoulli", "exponential", "laplace", "normal", "studentt".

  • latent_dim – Dimension of the latent variable z.

  • hidden_dim – Width of the hidden layers of the fully-connected nets.

  • num_epochs – Number of training epochs.

  • num_layers – Number of hidden layers in the fully-connected nets.

  • batch_size – Mini-batch size.

  • learning_rate – Initial Adam learning rate.

  • learning_rate_decay – Overall LR decay across training; per-step decay is learning_rate_decay ** (1 / num_steps).

  • num_samples – Number of Monte Carlo samples used by predict().

  • weight_decay – Decoupled (AdamW-style) weight decay coefficient.

  • seed – PRNG seed for parameter initialization and mini-batch shuffling.

fit(X, treatment, y, p=None)[source]#

Fits CEVAE on (X, treatment, y).

Parameters:

  • X – Feature matrix of shape (n, feature_dim).

  • treatment – Binary treatment vector of shape (n,).

  • y – Outcome vector of shape (n,).

  • p – Ignored (kept for API compatibility with meta-learners).

fit_predict(X, treatment, y, p=None)[source]#

Fits the model then returns ITE estimates on the same X.

Parameters:

  • X – Feature matrix.

  • treatment – Binary treatment vector.

  • y – Outcome vector.

  • p – Ignored (API compatibility).

Returns:

np.ndarray of shape (n,) with ITE estimates.

load(path, feature_dim)[source]#

Restores parameters from an orbax checkpoint.

Parameters:

  • path – Directory of a previously saved checkpoint.

  • feature_dim – Number of input features (needed to rebuild the net).

predict(X, treatment=None, y=None, p=None)[source]#

Predicts individual treatment effect for each row of X.

Parameters:

  • X – Feature matrix of shape (n, feature_dim).

  • treatment – Ignored (API compatibility).

  • y – Ignored (API compatibility).

  • p – Ignored (API compatibility).

Returns:

np.ndarray of shape (n,) with ITE estimates.

save(path)[source]#

Saves parameters to an orbax checkpoint directory.

class causalml.inference.jax.DragonNet(neurons_per_layer=200, targeted_reg=True, ratio=1.0, val_split=0.2, batch_size=64, epochs=100, learning_rate=1e-05, momentum=0.9, reg_l2=0.01, use_adam=True, adam_epochs=30, adam_learning_rate=0.001, loss_func=<function dragonnet_loss_binarycross>, verbose=True, seed=0)[source]#

Bases: object

JAX/flax.nnx DragonNet for treatment effect estimation.

Ports the TF DragonNet to JAX using flax.nnx, exposing an identical sklearn-style API. Two-phase training: Adam warm-up followed by SGD with Nesterov momentum.

Parameters:

  • neurons_per_layer – Width of shared representation layers.

  • targeted_reg – Whether to apply targeted regularization.

  • ratio – Weight for the targeted regularization loss component.

  • val_split – Fraction of training data reserved for validation.

  • batch_size – Mini-batch size.

  • epochs – Maximum SGD epochs.

  • learning_rate – SGD learning rate.

  • momentum – SGD Nesterov momentum.

  • reg_l2 – L2 regularization coefficient for outcome-head kernels.

  • use_adam – Whether to run an Adam warm-up phase before SGD.

  • adam_epochs – Maximum Adam epochs.

  • adam_learning_rate – Adam learning rate.

  • loss_func – Base loss function; defaults to dragonnet_loss_binarycross.

  • verbose – Whether to print epoch summaries.

  • seed – Random seed for parameter initialization and data shuffling.

Note

Unlike the TensorFlow backend, this JAX implementation does not currently use ReduceLROnPlateau or TerminateOnNaN callbacks.

fit(X, treatment, y, p=None)[source]#

Fits the DragonNet model.

Parameters:

  • X – Feature matrix of shape (n, p). Accepts np.ndarray or pd.DataFrame.

  • treatment – Binary treatment vector of shape (n,).

  • y – Outcome vector of shape (n,).

  • p – Ignored (kept for API compatibility with meta-learners).

fit_predict(X, treatment, y, p=None, return_components=False)[source]#

Fits the model and returns treatment effect estimates.

Parameters:

  • X – Feature matrix of shape (n, p).

  • treatment – Binary treatment vector of shape (n,).

  • y – Outcome vector of shape (n,).

  • p – Ignored (API compatibility).

  • return_components – Ignored for now; always returns ITEs.

Returns:

np.ndarray of shape (n, 1) with treatment effect estimates.

load(path, input_dim=None)[source]#

Restores model parameters from an orbax checkpoint.

Parameters:

  • path – Directory path of a previously saved checkpoint.

  • input_dim – Deprecated. Architecture metadata is read from the checkpoint’s config.json.

predict(X, treatment=None, y=None, p=None)[source]#

Runs forward pass on fitted DragonNet.

Parameters:

  • X – Feature matrix of shape (n, p).

  • treatment – Ignored (API compatibility).

  • y – Ignored (API compatibility).

  • p – Ignored (API compatibility).

Returns:

columns are (y0, y1, propensity, epsilon).

Return type:

np.ndarray of shape (n, 4)

predict_propensity(X)[source]#

Predicts individual propensity scores.

Parameters:

X – Feature matrix of shape (n, p).

Returns:

np.ndarray of shape (n,) with propensity scores.

predict_tau(X)[source]#

Predicts individual treatment effects (ITE).

Parameters:

X – Feature matrix of shape (n, p).

Returns:

np.ndarray of shape (n, 1) with treatment effect estimates.

save(path)[source]#

Saves model parameters to an orbax checkpoint directory.

Parameters:

path – Directory path where the checkpoint will be written.

causalml.optimize module#

class causalml.optimize.CounterfactualUnitSelector(learner, nevertaker_payoff, alwaystaker_payoff, complier_payoff, defier_payoff, organic_conversion=None)[source]#

Bases: object

A highly experimental implementation of the counterfactual unit selection model proposed by Li and Pearl (2019).

Parameters:

  • learner (object) – The base learner used to estimate the segment probabilities.

  • nevertaker_payoff (float) – The payoff from targeting a never-taker

  • alwaystaker_payoff (float) – The payoff from targeting an always-taker

  • complier_payoff (float) – The payoff from targeting a complier

  • defier_payoff (float) – The payoff from targeting a defier

  • organic_conversion (float, optional (default=None)) –

    The organic conversion rate in the population without an intervention. If None, the organic conversion rate is obtained from tne control group.

    NB: The organic conversion in the control group is not always the same as the organic conversion rate without treatment.

  • data (DataFrame) – A pandas DataFrame containing the features, treatment assignment indicator and the outcome of interest.

  • treatment (string) – A string corresponding to the name of the treatment column. The assumed coding in the column is 1 for treatment and 0 for control.

  • outcome (string) – A string corresponding to the name of the outcome column. The assumed coding in the column is 1 for conversion and 0 for no conversion.

References

Li, Ang, and Judea Pearl. 2019. “Unit Selection Based on Counterfactual Logic.” https://ftp.cs.ucla.edu/pub/stat_ser/r488.pdf.

fit(data, treatment, outcome)[source]#

Fits the class.

predict(data, treatment, outcome)[source]#

Predicts an individual-level payoff. If gain equality is satisfied, uses the exact function; if not, uses the midpoint between bounds.

class causalml.optimize.CounterfactualValueEstimator(treatment, control_name, treatment_names, y_proba, cate, value, conversion_cost, impression_cost, *args, **kwargs)[source]#

Bases: object

Parameters:

  • treatment (array, shape = (num_samples, )) – An array of treatment group indicator values.

  • control_name (string) – The name of the control condition as a string. Must be contained in the treatment array.

  • treatment_names (list, length = cate.shape[1]) – A list of treatment group names. NB: The order of the items in the list must correspond to the order in which the conditional average treatment effect estimates are in cate_array.

  • y_proba (array, shape = (num_samples, )) – The predicted probability of conversion using the Y ~ X model across the total sample.

  • cate (array, shape = (num_samples, len(set(treatment)))) – Conditional average treatment effect estimations from any model.

  • value (array, shape = (num_samples, )) – Value of converting each unit.

  • conversion_cost (shape = (num_samples, len(set(treatment)))) – The cost of a treatment that is triggered if a unit converts after having been in the treatment, such as a promotion code.

  • impression_cost (shape = (num_samples, len(set(treatment)))) – The cost of a treatment that is the same for each unit whether or not they convert, such as a cost associated with a promotion channel.

Notes

Because we get the conditional average treatment effects from cate-learners relative to the control condition, we subtract the cate for the unit in their actual treatment group from y_proba for that unit, in order to recover the control outcome. We then add the cates to the control outcome to obtain y_proba under each condition. These outcomes are counterfactual because just one of them is actually observed.

predict_best()[source]#

Predict the best treatment group based on the highest counterfactual value for a treatment.

predict_counterfactuals()[source]#

Predict the counterfactual values for each treatment group.

class causalml.optimize.PolicyLearner(outcome_learner=GradientBoostingRegressor(), treatment_learner=GradientBoostingClassifier(), policy_learner=DecisionTreeClassifier(), clip_bounds=(0.001, 0.999), n_fold=5, random_state=None, calibration=False)[source]#

Bases: object

A Learner that learns a treatment assignment policy with observational data using doubly robust estimator of causal effect for binary treatment.

Details of the policy learner are available at Athey and Wager (2018).

fit(X, treatment, y, p=None, dhat=None)[source]#

Fit the treatment assignment policy learner.

Parameters:

  • X (np.matrix) – a feature matrix

  • treatment (np.array) – a treatment vector (1 if treated, otherwise 0)

  • y (np.array) – an outcome vector

  • p (optional, np.array) – user provided propensity score vector between 0 and 1

  • dhat (optinal, np.array) – user provided predicted treatment effect vector

Returns:

returns an instance of self.

Return type:

self

predict(X)[source]#

Predict treatment assignment that optimizes the outcome.

Parameters:

X (np.matrix) – a feature matrix

Returns:

predictions of treatment assignment.

Return type:

(numpy.ndarray)

predict_proba(X)[source]#

Predict treatment assignment score that optimizes the outcome.

Parameters:

X (np.matrix) – a feature matrix

Returns:

predictions of treatment assignment score.

Return type:

(numpy.ndarray)

causalml.optimize.get_actual_value(treatment, observed_outcome, conversion_value, conditions, conversion_cost, impression_cost)[source]#

Set the conversion and impression costs based on a dict of parameters.

Calculate the actual value of targeting a user with the actual treatment group using the above parameters.

Params#

treatmentarray, shape = (num_samples, )

Treatment array.

observed_outcomearray, shape = (num_samples, )

Observed outcome array, aka y.

conversion_valuearray, shape = (num_samples, )

The value of converting a given user.

conditionslist, len = len(set(treatment))

List of treatment conditions.

conversion_costarray, shape = (num_samples, num_treatment)

Array of conversion costs for each unit in each treatment.

impression_costarray, shape = (num_samples, num_treatment)

Array of impression costs for each unit in each treatment.

returns:

  • actual_value (array, shape = (num_samples, )) – Array of actual values of havng a user in their actual treatment group.

  • conversion_value (array, shape = (num_samples, )) – Array of payoffs from converting a user.

causalml.optimize.get_pns_bounds(data_exp, data_obs, T, Y, type='PNS')[source]#
Parameters:

  • data_exp (DataFrame) – Data from an experiment.

  • data_obs (DataFrame) – Data from an observational study

  • T (str) – Name of the binary treatment indicator

  • y (str) – Name of the binary outcome indicator

  • type (str) –

    Type of probability of causation desired. Acceptable args are:

    • PNS: Probability of necessary and sufficient causation

    • PS: Probability of sufficient causation

    • PN: Probability of necessary causation

Notes

Based on Equation (24) in Tian and Pearl (2000).

To capture the counterfactual notation, we use 1 and 0 to indicate the actual and counterfactual values of a variable, respectively, and we use do to indicate the effect of an intervention.

The experimental and observational data are either assumed to come to the same population, or from random samples of the population. If the data are from a sample, the bounds may be incorrectly calculated because the relevant quantities in the Tian-Pearl equations are defined e.g. as \(P(Y|do(T))\), not \(P(Y|do(T), S)\) where \(S\) corresponds to sample selection. Bareinboim and Pearl (2016) discuss conditions under which \(P(Y|do(T))\) can be recovered from \(P(Y|do(T), S)\).

causalml.optimize.get_treatment_costs(treatment, control_name, cc_dict, ic_dict)[source]#

Set the conversion and impression costs based on a dict of parameters.

Calculate the actual cost of targeting a user with the actual treatment group using the above parameters.

Params#

treatmentarray, shape = (num_samples, )

Treatment array.

control_name, str

Control group name as string.

cc_dictdict

Dict containing the conversion cost for each treatment.

ic_dict

Dict containing the impression cost for each treatment.

returns:

  • conversion_cost (ndarray, shape = (num_samples, num_treatments)) – An array of conversion costs for each treatment.

  • impression_cost (ndarray, shape = (num_samples, num_treatments)) – An array of impression costs for each treatment.

  • conditions (list, len = len(set(treatment))) – A list of experimental conditions.

causalml.optimize.get_uplift_best(cate, conditions)[source]#

Takes the CATE prediction from a learner, adds the control outcome array and finds the name of the argmax conditon.

Params#

catearray, shape = (num_samples, )

The conditional average treatment effect prediction.

conditions : list, len = len(set(treatment))

returns:

uplift_recomm_name – The experimental group recommended by the learner.

rtype:

array, shape = (num_samples, )

causalml.dataset module#

causalml.dataset.bar_plot_summary(synthetic_summary, k, drop_learners=[], drop_cols=[], sort_cols=['MSE', 'Abs % Error of ATE'])[source]#

Generates a bar plot comparing learner performance.

Parameters:

  • synthetic_summary (pd.DataFrame) – summary generated by get_synthetic_summary()

  • k (int) – number of simulations (used only for plot title text)

  • drop_learners (list, optional) – list of learners (str) to omit when plotting

  • drop_cols (list, optional) – list of metrics (str) to omit when plotting

  • sort_cols (list, optional) – list of metrics (str) to sort on when plotting

causalml.dataset.bar_plot_summary_holdout(train_summary, validation_summary, k, drop_learners=[], drop_cols=[])[source]#

Generates a bar plot comparing learner performance by training and validation

Parameters:

  • train_summary (pd.DataFrame) – summary for training synthetic data generated by get_synthetic_summary_holdout()

  • validation_summary (pd.DataFrame) – summary for validation synthetic data generated by get_synthetic_summary_holdout()

  • k (int) – number of simulations (used only for plot title text)

  • drop_learners (list, optional) – list of learners (str) to omit when plotting

  • drop_cols (list, optional) – list of metrics (str) to omit when plotting

causalml.dataset.clear_data_dir(data_home=None) → None[source]#

Delete the cache directory, so a bad download can be recovered from.

Parameters:

data_home (str or Path, optional) – an explicit cache directory

causalml.dataset.distr_plot_single_sim(synthetic_preds, kind='kde', drop_learners=[], bins=50, histtype='step', alpha=1, linewidth=1, bw_method=1)[source]#

Plots the distribution of each learner’s predictions (for a single simulation). Kernel Density Estimation (kde) and actual histogram plots supported.

Parameters:

  • synthetic_preds (dict) – dictionary of predictions generated by get_synthetic_preds()

  • kind (str, optional) – ‘kde’ or ‘hist’

  • drop_learners (list, optional) – list of learners (str) to omit when plotting

  • bins (int, optional) – number of bins to plot if kind set to ‘hist’

  • histtype (str, optional) – histogram type if kind set to ‘hist’

  • alpha (float, optional) – alpha (transparency) for plotting

  • linewidth (int, optional) – line width for plotting

  • bw_method (float, optional) – parameter for kde

causalml.dataset.fetch_ihdp(replication=0, split='train', data_home=None, download_if_missing=True, return_X_y_t=False)[source]#

Load one replication of the IHDP benchmark.

Covariates from the Infant Health and Development Program randomized trial, with outcomes simulated on response surface B (Hill 2011). The file holds 100 replications of the same 747-unit sample: each one simulates its own outcomes and draws its own 672 / 75 train-test split, so row i is a different unit in each replication and the number of treated rows varies with it. Only the pooled 747 covariate rows are common to all of them.

Results on IHDP are reported as a mean and standard error across replications, so replication selects one and the caller loops. A single replication is not comparable to a published IHDP number.

Both potential outcome surfaces are returned, so the individual effect tau = mu1 - mu0 is known and PEHE is defined.

Parameters:

  • replication (int) – which replication to load, 0 to 99

  • split (str) – "train" (672 rows) or "test" (75 rows)

  • data_home (str or Path, optional) – cache directory

  • download_if_missing (bool) – if False, raise rather than download

  • return_X_y_t (bool) – return (X, y, treatment) instead of a Bunch

Returns:

sklearn.utils.Bunch with data, target (the factual outcome), treatment, tau, mu0, mu1, y_cf, feature_names and DESCR; or the triple if return_X_y_t

Raises:

ValueError – if split is not train/test or replication is out of range

causalml.dataset.fetch_lalonde(data_home=None, download_if_missing=True, return_X_y_t=False)[source]#

Load the LaLonde National Supported Work experiment.

A randomized job-training experiment: 445 rows, 185 treated, outcome re78 (1978 earnings in dollars). Its experimental estimate is the yardstick observational estimators are judged against (LaLonde 1986; Dehejia and Wahba 1999). The sample here is the Dehejia-Wahba one, taken from the MIT-licensed causaldata package at a pinned revision.

There is no per-unit ground truth, so PEHE is not defined on it. The experimental difference in means, about 1794 dollars, is what an estimate is compared against.

Parameters:

  • data_home (str or Path, optional) – cache directory

  • download_if_missing (bool) – if False, raise rather than download

  • return_X_y_t (bool) – return (X, y, treatment) instead of a Bunch

Returns:

sklearn.utils.Bunch with data, target, treatment, feature_names and DESCR; or the triple if return_X_y_t

causalml.dataset.fetch_twins(data_home=None, download_if_missing=True, return_X_y_t=False, random_state=None)[source]#

Load the Twins benchmark.

Same-sex twin births from the NBER linked birth / infant death records, 11400 pairs and 30 covariates. Treatment is being the heavier twin and the outcome is one-year mortality. Because both twins are observed, both potential outcomes are measured rather than simulated — the only dataset here where the ground truth is not a modelling assumption. Introduced as a causal benchmark by Louizos et al. (2017).

The raw outcome columns hold days survived with 9999 standing for “survived the year”, so mortality is outcome < 9999. That reproduces the 17.7% mortality for the lighter twin reported in the paper; reading the column as a number instead gives a mean near 8000 and no meaning.

One twin per pair is revealed, which makes this an observational dataset. The assignment here is randomized (a fair coin per pair), so the design is an RCT with known counterfactuals. The confounded variants in the literature assign treatment from a covariate, and they differ between papers; both potential outcomes are returned so a caller can build their own and say which.

Parameters:

  • data_home (str or Path, optional) – cache directory

  • download_if_missing (bool) – if False, raise rather than download

  • return_X_y_t (bool) – return (X, y, treatment) instead of a Bunch

  • random_state (int or np.random.RandomState, optional) – seeds which twin of each pair is revealed

Returns:

sklearn.utils.Bunch with data, target (the revealed twin’s mortality), treatment, tau, y0, y1, feature_names and DESCR; or the triple if return_X_y_t

causalml.dataset.get_data_home(data_home=None) → Path[source]#

Return the directory the benchmark loaders cache datasets in.

Resolution order: the data_home argument, then the CAUSALML_DATA environment variable, then ~/causalml-data. The directory is created if it does not exist.

Parameters:

data_home (str or Path, optional) – an explicit cache directory

Returns:

pathlib.Path, the cache directory

causalml.dataset.get_synthetic_auuc(synthetic_preds, drop_learners=[], outcome_col='y', treatment_col='w', treatment_effect_col='tau', plot=True)[source]#

Get auuc values for cumulative gains of model estimates in quantiles.

For details, reference get_cumgain() and plot_gain() :param synthetic_preds: dictionary of predictions generated by get_synthetic_preds() :type synthetic_preds: dict :param or get_synthetic_preds_holdout(): :param outcome_col: the column name for the actual outcome :type outcome_col: str, optional :param treatment_col: the column name for the treatment indicator (0 or 1) :type treatment_col: str, optional :param treatment_effect_col: the column name for the true treatment effect :type treatment_effect_col: str, optional :param plot: plot the cumulative gain chart or not :type plot: boolean,optional

Returns:

auuc values by learner for cumulative gains of model estimates

Return type:

(pandas.DataFrame)

causalml.dataset.get_synthetic_preds(synthetic_data_func, n=1000, estimators={})[source]#

Generate predictions for synthetic data using specified function (single simulation)

Parameters:

  • synthetic_data_func (function) – synthetic data generation function

  • n (int, optional) – number of samples

  • estimators (dict of object) – dict of names and objects of treatment effect estimators

Returns:

dict of the actual and estimates of treatment effects

Return type:

(dict)

causalml.dataset.get_synthetic_preds_holdout(synthetic_data_func, n=1000, valid_size=0.2, estimators={})[source]#

Generate predictions for synthetic data using specified function (single simulation) for train and holdout

Parameters:

  • synthetic_data_func (function) – synthetic data generation function

  • n (int, optional) – number of samples

  • valid_size (float,optional) – validaiton/hold out data size

  • estimators (dict of object) – dict of names and objects of treatment effect estimators

Returns:

synthetic training and validation data dictionaries:

  • preds_dict_train (dict): synthetic training data dictionary

  • preds_dict_valid (dict): synthetic validation data dictionary

Return type:

(tuple)

causalml.dataset.get_synthetic_summary(synthetic_data_func, n=1000, k=1, estimators={})[source]#

Generate a summary for predictions on synthetic data using specified function

Parameters:

  • synthetic_data_func (function) – synthetic data generation function

  • n (int, optional) – number of samples per simulation

  • k (int, optional) – number of simulations

causalml.dataset.get_synthetic_summary_holdout(synthetic_data_func, n=1000, valid_size=0.2, k=1)[source]#

Generate a summary for predictions on synthetic data for train and holdout using specified function

Parameters:

  • synthetic_data_func (function) – synthetic data generation function

  • n (int, optional) – number of samples per simulation

  • valid_size (float,optional) – validation/hold out data size

  • k (int, optional) – number of simulations

Returns:

summary evaluation metrics of predictions for train and validation:

  • summary_train (pandas.DataFrame): training data evaluation summary

  • summary_train (pandas.DataFrame): validation data evaluation summary

Return type:

(tuple)

causalml.dataset.make_uplift_classification(n_samples=1000, treatment_name=['control', 'treatment1', 'treatment2', 'treatment3'], y_name='conversion', n_classification_features=10, n_classification_informative=5, n_classification_redundant=0, n_classification_repeated=0, n_uplift_increase_dict={'treatment1': 2, 'treatment2': 2, 'treatment3': 2}, n_uplift_decrease_dict={'treatment1': 0, 'treatment2': 0, 'treatment3': 0}, delta_uplift_increase_dict={'treatment1': 0.02, 'treatment2': 0.05, 'treatment3': 0.1}, delta_uplift_decrease_dict={'treatment1': 0.0, 'treatment2': 0.0, 'treatment3': 0.0}, n_uplift_increase_mix_informative_dict={'treatment1': 1, 'treatment2': 1, 'treatment3': 1}, n_uplift_decrease_mix_informative_dict={'treatment1': 0, 'treatment2': 0, 'treatment3': 0}, positive_class_proportion=0.5, random_seed=20190101)[source]#

Generate a synthetic dataset for classification uplift modeling problem.

Parameters:

  • n_samples (int, optional (default=1000)) – The number of samples to be generated for each treatment group.

  • treatment_name (list, optional (default = ['control','treatment1','treatment2','treatment3'])) – The list of treatment names.

  • y_name (string, optional (default = 'conversion')) – The name of the outcome variable to be used as a column in the output dataframe.

  • n_classification_features (int, optional (default = 10)) – Total number of features for base classification

  • n_classification_informative (int, optional (default = 5)) – Total number of informative features for base classification

  • n_classification_redundant (int, optional (default = 0)) – Total number of redundant features for base classification

  • n_classification_repeated (int, optional (default = 0)) – Total number of repeated features for base classification

  • n_uplift_increase_dict (dictionary, optional (default: {'treatment1': 2, 'treatment2': 2, 'treatment3': 2})) – Number of features for generating positive treatment effects for corresponding treatment group. Dictionary of {treatment_key: number_of_features_for_increase_uplift}.

  • n_uplift_decrease_dict (dictionary, optional (default: {'treatment1': 0, 'treatment2': 0, 'treatment3': 0})) – Number of features for generating negative treatment effects for corresponding treatment group. Dictionary of {treatment_key: number_of_features_for_increase_uplift}.

  • delta_uplift_increase_dict (dictionary, optional (default: {'treatment1': .02, 'treatment2': .05, 'treatment3': .1})) – Positive treatment effect created by the positive uplift features on the base classification label. Dictionary of {treatment_key: increase_delta}.

  • delta_uplift_decrease_dict (dictionary, optional (default: {'treatment1': 0., 'treatment2': 0., 'treatment3': 0.})) – Negative treatment effect created by the negative uplift features on the base classification label. Dictionary of {treatment_key: increase_delta}.

  • n_uplift_increase_mix_informative_dict (dictionary, optional) – Number of positive mix features for each treatment. The positive mix feature is defined as a linear combination of a randomly selected informative classification feature and a randomly selected positive uplift feature. The linear combination is made by two coefficients sampled from a uniform distribution between -1 and 1. default: {‘treatment1’: 1, ‘treatment2’: 1, ‘treatment3’: 1}

  • n_uplift_decrease_mix_informative_dict (dictionary, optional) – Number of negative mix features for each treatment. The negative mix feature is defined as a linear combination of a randomly selected informative classification feature and a randomly selected negative uplift feature. The linear combination is made by two coefficients sampled from a uniform distribution between -1 and 1. default: {‘treatment1’: 0, ‘treatment2’: 0, ‘treatment3’: 0}

  • positive_class_proportion (float, optional (default = 0.5)) – The proportion of positive label (1) in the control group.

  • random_seed (int, optional (default = 20190101)) – The random seed to be used in the data generation process.

Returns:

  • df_res (DataFrame) – A data frame containing the treatment label, features, and outcome variable.

  • x_name (list) – The list of feature names generated.

Notes

The algorithm for generating the base classification dataset is adapted from the make_classification method in the sklearn package, that uses the algorithm in Guyon [1] designed to generate the “Madelon” dataset.

References

causalml.dataset.make_uplift_classification_logistic(n_samples=10000, treatment_name=['control', 'treatment1', 'treatment2', 'treatment3'], y_name='conversion', n_classification_features=10, n_classification_informative=5, n_classification_redundant=0, n_classification_repeated=0, n_uplift_dict={'treatment1': 2, 'treatment2': 2, 'treatment3': 3}, n_mix_informative_uplift_dict={'treatment1': 1, 'treatment2': 1, 'treatment3': 0}, delta_uplift_dict={'treatment1': 0.02, 'treatment2': 0.05, 'treatment3': -0.05}, positive_class_proportion=0.1, random_seed=20200101, feature_association_list=['linear', 'quadratic', 'cubic', 'relu', 'sin', 'cos'], random_select_association=True, error_std=0.05)[source]#

Generate a synthetic dataset for classification uplift modeling problem.

Parameters:

  • n_samples (int, optional (default=1000)) – The number of samples to be generated for each treatment group.

  • treatment_name (list, optional (default = ['control','treatment1','treatment2','treatment3'])) – The list of treatment names. The first element must be ‘control’ as control group, and the rest are treated as treatment groups.

  • y_name (string, optional (default = 'conversion')) – The name of the outcome variable to be used as a column in the output dataframe.

  • n_classification_features (int, optional (default = 10)) – Total number of features for base classification

  • n_classification_informative (int, optional (default = 5)) – Total number of informative features for base classification

  • n_classification_redundant (int, optional (default = 0)) – Total number of redundant features for base classification

  • n_classification_repeated (int, optional (default = 0)) – Total number of repeated features for base classification

  • n_uplift_dict (dictionary, optional (default: {'treatment1': 2, 'treatment2': 2, 'treatment3': 3})) – Number of features for generating heterogeneous treatment effects for corresponding treatment group. Dictionary of {treatment_key: number_of_features_for_uplift}.

  • n_mix_informative_uplift_dict (dictionary, optional (default: {'treatment1': 1, 'treatment2': 1, 'treatment3': 1})) – Number of mix features for each treatment. The mix feature is defined as a linear combination of a randomly selected informative classification feature and a randomly selected uplift feature. The mixture is made by a weighted sum (p*feature1 + (1-p)*feature2), where the weight p is drawn from a uniform distribution between 0 and 1.

  • delta_uplift_dict (dictionary, optional (default: {'treatment1': .02, 'treatment2': .05, 'treatment3': -.05})) – Treatment effect (delta), can be positive or negative. Dictionary of {treatment_key: delta}.

  • positive_class_proportion (float, optional (default = 0.1)) – The proportion of positive label (1) in the control group, or the mean of outcome variable for control group.

  • random_seed (int, optional (default = 20200101)) – The random seed to be used in the data generation process.

  • feature_association_list (list, optional (default = ['linear','quadratic','cubic','relu','sin','cos'])) – List of uplift feature association patterns to the treatment effect. For example, if the feature pattern is ‘quadratic’, then the treatment effect will increase or decrease quadratically with the feature. The values in the list must be one of (‘linear’,’quadratic’,’cubic’,’relu’,’sin’,’cos’). However, the same value can appear multiple times in the list.

  • random_select_association (boolean, optional (default = True)) – How the feature patterns are selected from the feature_association_list to be applied in the data generation process. If random_select_association = True, then for every uplift feature, a random feature association pattern is selected from the list. If random_select_association = False, then the feature association pattern is selected from the list in turns to be applied to each feature one by one.

  • error_std (float, optional (default = 0.05)) – Standard deviation to be used in the error term of the logistic regression. The error is drawn from a normal distribution with mean 0 and standard deviation specified in this argument.

Returns:

  • df1 (DataFrame) – A data frame containing the treatment label, features, and outcome variable.

  • x_name (list) – The list of feature names generated.

causalml.dataset.scatter_plot_single_sim(synthetic_preds)[source]#

Creates a grid of scatter plots comparing each learner’s predictions with the truth (for a single simulation).

Parameters:

synthetic_preds (dict) – dictionary of predictions generated by get_synthetic_preds() or get_synthetic_preds_holdout()

causalml.dataset.scatter_plot_summary(synthetic_summary, k, drop_learners=[], drop_cols=[])[source]#

Generates a scatter plot comparing learner performance. Each learner’s performance is plotted as a point in the (Abs % Error of ATE, MSE) space.

Parameters:

  • synthetic_summary (pd.DataFrame) – summary generated by get_synthetic_summary()

  • k (int) – number of simulations (used only for plot title text)

  • drop_learners (list, optional) – list of learners (str) to omit when plotting

  • drop_cols (list, optional) – list of metrics (str) to omit when plotting

causalml.dataset.scatter_plot_summary_holdout(train_summary, validation_summary, k, label=['Train', 'Validation'], drop_learners=[], drop_cols=[])[source]#

Generates a scatter plot comparing learner performance by training and validation.

Parameters:

  • train_summary (pd.DataFrame) – summary for training synthetic data generated by get_synthetic_summary_holdout()

  • validation_summary (pd.DataFrame) – summary for validation synthetic data generated by get_synthetic_summary_holdout()

  • label (string, optional) – legend label for plot

  • k (int) – number of simulations (used only for plot title text)

  • drop_learners (list, optional) – list of learners (str) to omit when plotting

  • drop_cols (list, optional) – list of metrics (str) to omit when plotting

causalml.dataset.simulate_easy_propensity_difficult_baseline(n=1000, p=5, sigma=1.0, adj=0.0)[source]#

Synthetic data with easy propensity and a difficult baseline

From Setup C in Nie X. and Wager S. (2018) ‘Quasi-Oracle Estimation of Heterogeneous Treatment Effects’

Parameters:

  • n (int, optional) – number of observations

  • p (int optional) – number of covariates (>=3)

  • sigma (float) – standard deviation of the error term

  • adj (float) – no effect. added for consistency

Returns:

Synthetically generated samples with the following outputs:

  • y ((n,)-array): outcome variable.

  • X ((n,p)-ndarray): independent variables.

  • w ((n,)-array): treatment flag with value 0 or 1.

  • tau ((n,)-array): individual treatment effect.

  • b ((n,)-array): expected outcome.

  • e ((n,)-array): propensity of receiving treatment.

Return type:

(tuple)

causalml.dataset.simulate_hidden_confounder(n=10000, p=5, sigma=1.0, adj=0.0)[source]#

Synthetic dataset with a hidden confounder biasing treatment.

From Louizos et al. (2018) “Causal Effect Inference with Deep Latent-Variable Models”

Parameters:

  • n (int, optional) – number of observations

  • p (int optional) – number of covariates (>=3)

  • sigma (float) – standard deviation of the error term

  • adj (float) – no effect. added for consistency

Returns:

Synthetically generated samples with the following outputs:

  • y ((n,)-array): outcome variable.

  • X ((n,p)-ndarray): independent variables.

  • w ((n,)-array): treatment flag with value 0 or 1.

  • tau ((n,)-array): individual treatment effect.

  • b ((n,)-array): expected outcome.

  • e ((n,)-array): propensity of receiving treatment.

Return type:

(tuple)

causalml.dataset.simulate_nuisance_and_easy_treatment(n=1000, p=5, sigma=1.0, adj=0.0)[source]#

Synthetic data with a difficult nuisance components and an easy treatment effect

From Setup A in Nie X. and Wager S. (2018) ‘Quasi-Oracle Estimation of Heterogeneous Treatment Effects’

Parameters:

  • n (int, optional) – number of observations

  • p (int optional) – number of covariates (>=5)

  • sigma (float) – standard deviation of the error term

  • adj (float) – adjustment term for the distribution of propensity, e. Higher values shift the distribution to 0.

Returns:

Synthetically generated samples with the following outputs:

  • y ((n,)-array): outcome variable.

  • X ((n,p)-ndarray): independent variables.

  • w ((n,)-array): treatment flag with value 0 or 1.

  • tau ((n,)-array): individual treatment effect.

  • b ((n,)-array): expected outcome.

  • e ((n,)-array): propensity of receiving treatment.

Return type:

(tuple)

causalml.dataset.simulate_randomized_trial(n=1000, p=5, sigma=1.0, adj=0.0)[source]#

Synthetic data of a randomized trial

From Setup B in Nie X. and Wager S. (2018) ‘Quasi-Oracle Estimation of Heterogeneous Treatment Effects’

Parameters:

  • n (int, optional) – number of observations

  • p (int optional) – number of covariates (>=5)

  • sigma (float) – standard deviation of the error term

  • adj (float) – no effect. added for consistency

Returns:

Synthetically generated samples with the following outputs:

  • y ((n,)-array): outcome variable.

  • X ((n,p)-ndarray): independent variables.

  • w ((n,)-array): treatment flag with value 0 or 1.

  • tau ((n,)-array): individual treatment effect.

  • b ((n,)-array): expected outcome.

  • e ((n,)-array): propensity of receiving treatment.

Return type:

(tuple)

causalml.dataset.simulate_unrelated_treatment_control(n=1000, p=5, sigma=1.0, adj=0.0)[source]#

Synthetic data with unrelated treatment and control groups.

From Setup D in Nie X. and Wager S. (2018) ‘Quasi-Oracle Estimation of Heterogeneous Treatment Effects’

Parameters:

  • n (int, optional) – number of observations

  • p (int optional) – number of covariates (>=3)

  • sigma (float) – standard deviation of the error term

  • adj (float) – adjustment term for the distribution of propensity, e. Higher values shift the distribution to 0.

Returns:

Synthetically generated samples with the following outputs:

  • y ((n,)-array): outcome variable.

  • X ((n,p)-ndarray): independent variables.

  • w ((n,)-array): treatment flag with value 0 or 1.

  • tau ((n,)-array): individual treatment effect.

  • b ((n,)-array): expected outcome.

  • e ((n,)-array): propensity of receiving treatment.

Return type:

(tuple)

causalml.dataset.synthetic_data(mode=1, n=1000, p=5, sigma=1.0, adj=0.0)[source]#

Synthetic data in Nie X. and Wager S. (2018) ‘Quasi-Oracle Estimation of Heterogeneous Treatment Effects’

Parameters:

  • mode (int, optional) – mode of the simulation: 1 for difficult nuisance components and an easy treatment effect. 2 for a randomized trial. 3 for an easy propensity and a difficult baseline. 4 for unrelated treatment and control groups. 5 for a hidden confounder biasing treatment.

  • n (int, optional) – number of observations

  • p (int optional) – number of covariates (>=5)

  • sigma (float) – standard deviation of the error term

  • adj (float) – adjustment term for the distribution of propensity, e. Higher values shift the distribution to 0. It does not apply to mode == 2 or 3.

Returns:

Synthetically generated samples with the following outputs:

  • y ((n,)-array): outcome variable.

  • X ((n,p)-ndarray): independent variables.

  • w ((n,)-array): treatment flag with value 0 or 1.

  • tau ((n,)-array): individual treatment effect.

  • b ((n,)-array): expected outcome.

  • e ((n,)-array): propensity of receiving treatment.

Return type:

(tuple)

causalml.match module#

class causalml.match.MatchOptimizer(treatment_col='is_treatment', ps_col='pihat', user_col=None, matching_covariates=['pihat'], max_smd=0.1, max_deviation=0.1, caliper_range=(0.01, 0.5), max_pihat_range=(0.95, 0.999), max_iter_per_param=5, min_users_per_group=1000, smd_cols=['pihat'], dev_cols_transformations={'pihat': <function mean>}, dev_factor=1.0, verbose=True)[source]#

Bases: object

check_table_one(tableone, matched, score_cols, pihat_threshold, caliper)[source]#
match_and_check(score_cols, pihat_threshold, caliper)[source]#
search_best_match(df)[source]#
single_match(score_cols, pihat_threshold, caliper)[source]#
class causalml.match.NearestNeighborMatch(caliper=0.2, replace=False, ratio=1, shuffle=True, treatment_to_control=True, random_state=None, n_jobs=-1)[source]#

Bases: object

Propensity score matching based on the nearest neighbor algorithm.

caliper#

threshold to be considered as a match.

Type:

float

replace#

whether to match with replacement or not

Type:

bool

ratio#

ratio of control / treatment to be matched.

Type:

int

shuffle#

whether to shuffle the treatment group data before matching

Type:

bool

treatment_to_control#

whether to match treatment to control or control to treatment

Type:

bool

random_state#

RandomState or an int seed

Type:

numpy.random.RandomState or int

n_jobs#

The number of parallel jobs to run for neighbors search. None means 1 unless in a joblib.parallel_backend context. -1 means using all processors

Type:

int

Fitted attributes (populated after match()):
matched_indexes_ (pandas.DataFrame): two-column dataframe with the

(from, to) pairs of original data indices produced by the most recent match() call. from corresponds to the matching source group (treatment if treatment_to_control else control); to is the matched counterpart from the opposite group. Each row is one matched pair, so with ratio > 1 a single from index can appear multiple times against distinct to indices. Useful for joining matched pairs back to upstream metadata or auditing the matching outcome (see uber/causalml#621). match_by_group calls match() once per group, so the attribute reflects the last group processed.

match(data, treatment_col, score_cols)[source]#

Find matches from the control group by matching on specified columns (propensity preferred).

Parameters:

  • data (pandas.DataFrame) – total input data

  • treatment_col (str) – the column name for the treatment

  • score_cols (list) – list of column names for matching (propensity column should be included)

Returns:

The subset of data consisting of matched

treatment and control group data.

Return type:

(pandas.DataFrame)

match_by_group(data, treatment_col, score_cols, groupby_col)[source]#

Find matches from the control group stratified by groupby_col, by matching on specified columns (propensity preferred).

Parameters:

  • data (pandas.DataFrame) – total sample data

  • treatment_col (str) – the column name for the treatment

  • score_cols (list) – list of column names for matching (propensity column should be included)

  • groupby_col (str) – the column name to be used for stratification

Returns:

The subset of data consisting of matched

treatment and control group data.

Return type:

(pandas.DataFrame)

causalml.match.create_table_one(data, treatment_col, features, with_std=True, with_counts=True)[source]#

Report balance in input features between the treatment and control groups.

References

R’s tableone at CRAN: kaz-yos/tableone Python’s tableone at PyPi: tompollard/tableone

Parameters:

  • data (pandas.DataFrame) – total or matched sample data

  • treatment_col (str) – the column name for the treatment

  • features (list of str) – the column names of features

  • with_std (bool) – whether to output std together with mean values as in <mean> (<std>) format

  • with_counts (bool) – whether to include a row counting the total number of samples

Returns:

A table with the means and standard deviations in

the treatment and control groups, and the SMD between two groups for the features.

Return type:

(pandas.DataFrame)

causalml.match.smd(feature, treatment)[source]#

Calculate the standard mean difference (SMD) of a feature between the treatment and control groups.

The definition is available at https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3144483/#s11title

Parameters:

  • feature (pandas.Series) – a column of a feature to calculate SMD for

  • treatment (pandas.Series) – a column that indicate whether a row is in the treatment group or not

Returns:

The SMD of the feature

Return type:

(float)

causalml.propensity module#

class causalml.propensity.ElasticNetPropensityModel(clip_bounds=(0.001, 0.999), calibrate=True, **model_kwargs)[source]#

Bases: LogisticRegressionPropensityModel

class causalml.propensity.GradientBoostedPropensityModel(early_stop=False, clip_bounds=(0.001, 0.999), calibrate=True, **model_kwargs)[source]#

Bases: PropensityModel

Gradient boosted propensity score model with optional early stopping.

Notes

Please see the xgboost documentation for more information on gradient boosting tuning parameters: https://xgboost.readthedocs.io/en/latest/python/python_api.html

fit(X, y, stop_val_size=0.2)[source]#

Fit a propensity model.

Parameters:

  • X (numpy.ndarray, pd.DataFrame, or pl.DataFrame) – a feature matrix

  • y (numpy.ndarray, pd.Series, or pl.Series) – the binary treatment indicator, not the outcome. A propensity score is P(W = 1 | X), so this vector must be the treatment assignment.

class causalml.propensity.LogisticRegressionPropensityModel(clip_bounds=(0.001, 0.999), calibrate=True, **model_kwargs)[source]#

Bases: PropensityModel

Propensity regression model based on the LogisticRegression algorithm.

class causalml.propensity.PropensityModel(clip_bounds=(0.001, 0.999), calibrate=True, **model_kwargs)[source]#

Bases: object

fit(X, y)[source]#

Fit a propensity model.

Parameters:

  • X (numpy.ndarray, pd.DataFrame, or pl.DataFrame) – a feature matrix. scikit-learn >= 1.6 accepts pandas and Polars DataFrames natively, so no conversion is performed here.

  • y (numpy.ndarray, pd.Series, or pl.Series) – the binary treatment indicator, not the outcome. A propensity score is P(W = 1 | X), so this vector must be the treatment assignment.

fit_predict(X, y)[source]#

Fit a propensity model and predict propensity scores.

Parameters:

  • X (numpy.ndarray, pd.DataFrame, or pl.DataFrame) – a feature matrix

  • y (numpy.ndarray, pd.Series, or pl.Series) – the binary treatment indicator, not the outcome. A propensity score is P(W = 1 | X), so this vector must be the treatment assignment.

Returns:

Propensity scores between 0 and 1.

Return type:

(numpy.ndarray)

predict(X)[source]#

Predict propensity scores.

Parameters:

X (numpy.ndarray, pd.DataFrame, or pl.DataFrame) – a feature matrix

Returns:

Propensity scores between 0 and 1.

Return type:

(numpy.ndarray)

causalml.propensity.compute_propensity_score(X, treatment, p_model=None, X_pred=None, treatment_pred=None, calibrate_p=True, clip_bounds=(0.001, 0.999))[source]#

Generate propensity score if user didn’t provide and optionally calibrate.

Parameters:

  • X (np.matrix, pd.DataFrame, or pl.DataFrame) – features for training

  • treatment (np.array, pd.Series, or pl.Series) – a treatment vector for training

  • p_model (model object, optional) – a binary classifier with either a predict_proba or predict method

  • X_pred (np.matrix, pd.DataFrame, or pl.DataFrame, optional) – features for prediction

  • treatment_pred (np.array, pd.Series, or pl.Series, optional) – a treatment vector for prediction

  • calibrate_p (bool, optional) – whether calibrate the propensity score

  • clip_bounds (tuple, optional) – lower and upper bounds for clipping propensity scores. Bounds should be implemented such that: 0 < lower < upper < 1, to avoid division by zero in BaseRLearner.fit_predict() step.

Returns:

(tuple)

  • p (numpy.ndarray): propensity score

  • p_model (PropensityModel): either the original p_model or a trained ElasticNetPropensityModel

causalml.propensity.compute_r_residuals(X, treatment, y, outcome_learner, propensity_learner=None, p=None, method='predict', n_folds=5, random_state=None, n_jobs=-1, compute_w_residual=True)[source]#

Cross-fitted outcome/treatment residuals for the R-loss (Nie & Wager, 2021).

Computes out-of-fold m_hat(X) = E[Y|X] and e_hat(X) = E[W|X] via n_folds-fold cross-fitting, stratified on treatment so every fold retains both arms, and returns:

y_residual = y - m_hat(X) w_residual = w - e_hat(X)

A candidate CATE model tau_hat is scored against these via the R-loss:

R-loss(tau_hat) = mean[(y_residual - w_residual * tau_hat(X)) ** 2]

This is also the quantity BaseRLearner.fit() implicitly minimizes: fitting the per-arm effect model against target (y_residual / w_residual) with sample_weight = w_residual ** 2 is the weighted-least-squares solution to the same R-loss objective.

Parameters:

  • X (numpy.ndarray or pandas.DataFrame) – a feature matrix

  • treatment (numpy.ndarray or pandas.Series) – a binary treatment indicator (0 or 1)

  • y (numpy.ndarray or pandas.Series) – an outcome vector

  • outcome_learner (model) – a model to estimate E[Y|X]. Must implement fit/predict (or predict_proba if method=”predict_proba”)

  • propensity_learner (PropensityModel, optional) – passed through to compute_propensity_score(). Ignored if p is given. Defaults to ElasticNetPropensityModel

  • p (numpy.ndarray or pandas.Series, optional) – pre-computed propensity scores. If given, propensity is not re-estimated in-fold

  • method (str, optional) – “predict” or “predict_proba” (for classifier outcome learners, e.g. BaseRClassifier). Only the positive-class column is used for “predict_proba”. Default “predict”

  • n_folds (int, optional) – number of cross-fitting folds. Default 5

  • random_state (int or None, optional) – random seed for the fold splitter

  • n_jobs (int, optional) – parallel jobs forwarded to cross_val_predict for the outcome model. Default -1

  • compute_w_residual (bool, optional) – whether to compute and return w_residual. If False, skips propensity estimation entirely (no in-fold propensity model is fit) and returns w_residual=None. Set False when only the outcome residual is needed – e.g. BaseRLearner.fit(), which already has propensity scores from elsewhere and would otherwise pay for a redundant per-fold propensity fit whose output is discarded. Default True.

Returns:

  • y_residual (numpy.ndarray): y - m_hat(X), out-of-fold

  • w_residual (numpy.ndarray or None): w - e_hat(X), out-of-fold (or w - p directly if p was supplied), or None if compute_w_residual=False

Return type:

(tuple)

causalml.metrics module#

class causalml.metrics.Sensitivity(df, inference_features, p_col, treatment_col, outcome_col, learner, *args, **kwargs)[source]#

Bases: object

A Sensitivity Check class to support Placebo Treatment, Irrelevant Additional Confounder and Subset validation refutation methods to verify causal inference.

Reference: microsoft/dowhy

get_ate_ci(X, p, treatment, y)[source]#

Return the confidence intervals for treatment effects prediction.

Parameters:

  • X (np.matrix) – a feature matrix

  • p (np.array) – a propensity score vector between 0 and 1

  • treatment (np.array) – a treatment vector (1 if treated, otherwise 0)

  • y (np.array) – an outcome vector

Returns:

Mean and confidence interval (LB, UB) of the ATE estimate.

Return type:

(numpy.ndarray)

static get_class_object(method_name, *args, **kwargs)[source]#

Return class object based on input method :param method_name: a list of sensitivity analysis method :type method_name: list of str

Returns:

Sensitivy Class

Return type:

(class)

get_potential_outcome_predictions(X, p, treatment, y)[source]#

Return separate potential-outcome predictions mu1_hat, mu0_hat.

Only supported for S/T/DR-learner-style objects, whose fit_predict(…, return_components=True) returns the fitted outcome regressions (mu0_hat, mu1_hat) directly. X-learner and R-learner are explicitly unsupported: X-learner’s “components” are two CATE estimates from its second-stage tau models, not potential outcomes, and R-learner has no outcome-regression decomposition to extract.

Parameters:

  • X – same as get_prediction()

  • p – same as get_prediction()

  • treatment – same as get_prediction()

  • y – same as get_prediction()

Returns:

(mu1_hat, mu0_hat)

Return type:

(tuple of np.array)

Raises:

NotImplementedError – if the learner does not expose potential-outcome regressions via return_components.

get_prediction(X, p, treatment, y)[source]#

Return the treatment effects prediction.

Parameters:

  • X (np.matrix) – a feature matrix

  • p (np.array) – a propensity score vector between 0 and 1

  • treatment (np.array) – a treatment vector (1 if treated, otherwise 0)

  • y (np.array) – an outcome vector

Returns:

Predictions of treatment effects

Return type:

(numpy.ndarray)

sensitivity_analysis(methods, sample_size=None, confound='one_sided', alpha_range=None)[source]#

Return the sensitivity data by different method

Parameters:

  • method (list of str) – a list of sensitivity analysis method

  • sample_size (float, optional) – ratio for subset the original data

  • confound (string, optional) – the name of confouding function

  • alpha_range (np.array, optional) – a parameter to pass the confounding function

Returns:

a feature matrix p (np.array): a propensity score vector between 0 and 1 treatment (np.array): a treatment vector (1 if treated, otherwise 0) y (np.array): an outcome vector

Return type:

X (np.matrix)

sensitivity_estimate()[source]#
summary(method)[source]#

Summary report :param method_name: sensitivity analysis method :type method_name: str

Returns:

a summary dataframe

Return type:

(pd.DataFrame)

class causalml.metrics.SensitivityPlaceboTreatment(*args, **kwargs)[source]#

Bases: Sensitivity

Replaces the treatment variable with a new variable randomly generated.

sensitivity_estimate()[source]#

Summary report :param return_ci: sensitivity analysis method :type return_ci: str

Returns:

a summary dataframe

Return type:

(pd.DataFrame)

class causalml.metrics.SensitivityRandomCause(*args, **kwargs)[source]#

Bases: Sensitivity

Adds an irrelevant random covariate to the dataframe.

sensitivity_estimate()[source]#
class causalml.metrics.SensitivityRandomReplace(*args, **kwargs)[source]#

Bases: Sensitivity

Replaces a random covariate with an irrelevant variable.

sensitivity_estimate()[source]#

Replaces a random covariate with an irrelevant variable.

class causalml.metrics.SensitivitySelectionBias(*args, confound='one_sided', alpha_range=None, sensitivity_features=None, **kwargs)[source]#

Bases: Sensitivity

Reference:

[1] Blackwell, Matthew. “A selection bias approach to sensitivity analysis for causal effects.” Political Analysis 22.2 (2014): 169-182. https://www.mattblackwell.org/files/papers/causalsens.pdf

[2] Confouding parameter alpha_range using the same range as in: mattblackwell/causalsens

causalsens()[source]#
static partial_rsqs_confounding(sens_df, feature_name, partial_rsqs_value, range=0.01)[source]#

Check partial rsqs values of feature corresponding confounding amonunt of ATE :param sens_df: a data frame output from causalsens :type sens_df: pandas.DataFrame :param feature_name: feature name to check :type feature_name: str :param partial_rsqs_value: partial rsquare value of feature :type partial_rsqs_value: float :param range: range to search from sens_df :type range: float

Return: min and max value of confounding amount

static plot(sens_df, partial_rsqs_df=None, type='raw', ci=False, partial_rsqs=False)[source]#

Plot the results of a sensitivity analysis against unmeasured :param sens_df: a data frame output from causalsens :type sens_df: pandas.DataFrame :param partial_rsqs_d: a data frame output from causalsens including partial rsqure :type partial_rsqs_d: pandas.DataFrame :param type: the type of plot to draw, ‘raw’ or ‘r.squared’ are supported :type type: str, optional :param ci: whether plot confidence intervals :type ci: bool, optional :param partial_rsqs: whether plot partial rsquare results :type partial_rsqs: bool, optional

summary(method='Selection Bias')[source]#

Summary report for Selection Bias Method :param method_name: sensitivity analysis method :type method_name: str

Returns:

a summary dataframe

Return type:

(pd.DataFrame)

class causalml.metrics.SensitivitySubsetData(*args, **kwargs)[source]#

Bases: Sensitivity

Takes a random subset of size sample_size of the data.

sensitivity_estimate()[source]#
causalml.metrics.ape(y, p)[source]#

Absolute Percentage Error (APE). :param y: target :type y: float :param p: prediction :type p: float

Returns:

APE

Return type:

e (float)

causalml.metrics.ate_error(tau, tau_hat)[source]#

Absolute error in the average treatment effect, usually written eps_ATE.

abs(mean(tau_hat) - mean(tau)). A model can order units well and still be biased on the average, and vice versa, so this is not implied by PEHE or by AUUC.

On an experiment with no per-unit ground truth, pass the experimental estimate as tau: the comparison is then against the number the experiment licenses rather than against a per-unit truth that does not exist.

Parameters:

  • tau (np.ndarray or float) – the true individual effects, or the true ATE

  • tau_hat (np.ndarray or float) – the estimated individual effects, or the estimated ATE

Returns:

float, the absolute difference of the two averages

causalml.metrics.auuc_score(df, outcome_col='y', treatment_col='w', treatment_effect_col='tau', normalize=True, tmle=False, *args, return_ci=False, n_bootstrap=200, alpha=0.05, random_state=None, **kwarg)[source]#

Calculate the AUUC (Area Under the Uplift Curve) score.

Parameters:

  • df (pandas.DataFrame) – a data frame with model estimates and actual data as columns. Columns not matching outcome_col, treatment_col, or treatment_effect_col are treated as model prediction columns whose AUUC will be computed.

  • outcome_col (str, optional) – the column name for the actual outcome

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1)

  • treatment_effect_col (str, optional) – the column name for the true treatment effect

  • normalize (bool, optional) – whether to normalize the y-axis to 1 or not

  • return_ci (bool, optional) – whether to return standard errors and bootstrap confidence intervals. Default False, so existing callers are unaffected.

  • n_bootstrap (int, optional) – number of half-sample bootstrap iterations. Only used when return_ci=True. Default 200.

  • alpha (float, optional) – significance level for confidence intervals. Only used when return_ci=True. Default 0.05.

  • random_state (int or None, optional) – random seed for the bootstrap sampler. Pass an integer for reproducible results. Default None.

Returns:

(pandas.Series): the AUUC score for each model estimate column If return_ci=True: (pandas.DataFrame): AUUC score, standard error and confidence interval bounds for each model estimate column.

Return type:

If return_ci=False

Note

No p-value is reported for AUUC. A ranking drawn at random scores about 0.5 here rather than 0, so testing H0: AUUC = 0 would reject for essentially every model and say nothing about whether the model beats random. Use qini_score(), which is already measured against the random curve, when a test against random is what is wanted.

causalml.metrics.classification_metrics(y, p, w=None, metrics={'AUC': <function roc_auc_score>, 'Log Loss': <function logloss>})[source]#

Log metrics for classifiers.

Parameters:

  • y (numpy.array) – target

  • p (numpy.array) – prediction

  • w (numpy.array, optional) – a treatment vector (1 or True: treatment, 0 or False: control). If given, log metrics for the treatment and control group separately

  • metrics (dict, optional) – a dictionary of the metric names and functions

causalml.metrics.compute_dr_pseudo_outcomes(X, treatment, y, p=None, learner=LGBMRegressor(learning_rate=0.05, n_estimators=300, num_leaves=64, verbose=-1), control_outcome_learner=None, treatment_outcome_learner=None, n_folds=5, p_clip_bounds=(0.02, 0.98), random_state=None)[source]#

Construct cross-fitted doubly-robust (AIPW) pseudo-outcomes for CATE evaluation.

For each unit i, the pseudo-outcome is

phi_i = (w_i - e(X_i)) / (e(X_i) * (1 - e(X_i))) * (y_i - mu_w(X_i)) + mu_1(X_i) - mu_0(X_i)

where e is the propensity score and mu_0/mu_1 are the control/treatment outcome regressions. Under either correct propensity or correct outcome-model specification, E[phi_i | X_i] is an unbiased estimate of the true CATE tau(X_i) (Kennedy, 2023), which is why phi can stand in for the unobserved ground-truth treatment effect when scoring fitted CATE models.

Nuisance models (propensity and outcome regressions) are cross-fitted with n_folds-fold splitting so that phi_i is always constructed from models that did not see unit i during training. This is the same doubly-robust formula used internally by BaseDRLearner.fit().

This is a standalone helper so the pseudo-outcomes can be computed once and reused across multiple scoring calls – e.g. passed to dr_score() directly, or to rate_score(..., treatment_effect_col=...) for RATE on observational data – without re-fitting nuisance models for each.

Parameters:

  • X (numpy.ndarray or pandas.DataFrame) – a feature matrix

  • treatment (numpy.ndarray or pandas.Series) – a binary treatment indicator (0 or 1)

  • y (numpy.ndarray or pandas.Series) – an outcome vector

  • p (numpy.ndarray or pandas.Series, optional) – propensity scores. If None, they are estimated in-fold via causalml.propensity.compute_propensity_score (ElasticNetPropensityModel by default)

  • learner (model, optional) – a model used for both control and treatment outcome regressions if the group-specific learners below are not given

  • control_outcome_learner (model, optional) – a model to estimate outcomes in the control group

  • treatment_outcome_learner (model, optional) – a model to estimate outcomes in the treatment group

  • n_folds (int, optional) – number of cross-fitting folds. Default 5.

  • p_clip_bounds (tuple, optional) – lower and upper bounds for clipping propensity scores before they’re used as AIPW weights. The default ElasticNetPropensityModel clips to (1e-3, 1 - 1e-3) internally for numerical stability of the model itself, but that’s too permissive once the score is inverted here: a handful of near-boundary, cross-fitted propensities (e.g. 0.001, arising from isotonic calibration on a single fold) can produce AIPW weights in the hundreds and dominate the mean. Tighter trimming bounds the variance at the cost of some bias for units with extreme propensity; (0.02, 0.98) is a reasonable default for that trade-off. Default (0.02, 0.98).

  • random_state (int or None, optional) – random seed for the fold splitter. Default None.

Returns:

the cross-fitted DR pseudo-outcomes, one per row of X

Return type:

(numpy.ndarray)

causalml.metrics.dr_score(df, X=None, treatment_col='w', outcome_col='y', pseudo_outcome_col=None, p=None, learner=LGBMRegressor(learning_rate=0.05, n_estimators=300, num_leaves=64, verbose=-1), control_outcome_learner=None, treatment_outcome_learner=None, n_folds=5, p_clip_bounds=(0.02, 0.98), return_ci=False, n_bootstrap=200, alpha=0.05, random_state=None)[source]#

Score fitted CATE models via the doubly-robust (DR) pseudo-outcome loss.

Following Kennedy (2023), this constructs cross-fitted AIPW pseudo-outcomes phi (see compute_dr_pseudo_outcomes()) and scores each candidate CATE model by its mean squared error against phi:

DR loss(tau_hat) = mean((tau_hat(X) - phi) ** 2)

Lower is better. Unlike held-out outcome MSE, this measures accuracy of the treatment effect estimate rather than the outcome level, without requiring access to counterfactual outcomes. Mahajan et al. (2024) found DR-based metrics dominate across 78 benchmark datasets for CATE model selection.

Pseudo-outcomes can either be supplied directly (via pseudo_outcome_col, e.g. computed once with compute_dr_pseudo_outcomes() and reused across multiple scoring calls or shared with rate_score()) or computed internally from X, treatment_col, and outcome_col.

Parameters:

  • df (pandas.DataFrame) – a data frame with fitted CATE model estimates as columns, plus either pseudo_outcome_col or both outcome_col and treatment_col

  • X (numpy.ndarray or pandas.DataFrame, optional) – feature matrix used to fit the DR nuisance models. Required unless pseudo_outcome_col is already present in df

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1). Ignored if pseudo_outcome_col is provided

  • outcome_col (str, optional) – the column name for the actual outcome. Ignored if pseudo_outcome_col is provided

  • pseudo_outcome_col (str, optional) – the column name of pre-computed DR pseudo-outcomes (e.g. from compute_dr_pseudo_outcomes()). If given and present in df, nuisance models are not re-fit

  • p (numpy.ndarray or pandas.Series, optional) – propensity scores. Only used when pseudo-outcomes are computed internally

  • learner (model, optional) – a model for both control and treatment outcome regressions if the group-specific learners below are not given

  • control_outcome_learner (model, optional) – a model to estimate outcomes in the control group

  • treatment_outcome_learner (model, optional) – a model to estimate outcomes in the treatment group

  • n_folds (int, optional) – number of cross-fitting folds for nuisance estimation. Default 5

  • p_clip_bounds (tuple, optional) – bounds for clipping propensity scores used as AIPW weights when pseudo-outcomes are computed internally. See compute_dr_pseudo_outcomes() for why this is tighter than a propensity model’s own internal clipping. Ignored if pseudo_outcome_col is provided. Default (0.02, 0.98)

  • return_ci (bool, optional) – whether to return bootstrap confidence intervals. Default False

  • n_bootstrap (int, optional) – number of half-sample bootstrap iterations. Only used when return_ci=True. Default 200

  • alpha (float, optional) – significance level for confidence intervals. Only used when return_ci=True. Default 0.05

  • random_state (int or None, optional) – random seed for cross-fitting and the bootstrap sampler. Default None

Returns:

(pandas.Series): DR loss for each model column (lower is better) If return_ci=True: (pandas.DataFrame): DR loss, standard error, and confidence interval bounds for each model column

Return type:

If return_ci=False

causalml.metrics.get_cumgain(df, outcome_col='y', treatment_col='w', treatment_effect_col='tau', normalize=False, random_seed=42)[source]#

Get cumulative gains of model estimates in population.

If the true treatment effect is provided (e.g. in synthetic data), it’s calculated as the cumulative gain of the true treatment effect in each population. Otherwise, it’s calculated as the cumulative difference between the mean outcomes of the treatment and control groups in each population.

For details, see Section 4.1 of Gutierrez and G{‘e}rardy (2016), Causal Inference and Uplift Modeling: A review of the literature.

For the former, treatment_effect_col should be provided. For the latter, both outcome_col and treatment_col should be provided.

Parameters:

  • df (pandas.DataFrame) – a data frame with model estimates and actual data as columns

  • outcome_col (str, optional) – the column name for the actual outcome

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1)

  • treatment_effect_col (str, optional) – the column name for the true treatment effect

  • normalize (bool, optional) – whether to normalize the y-axis to 1 or not

  • random_seed (int, optional) – deprecated

Returns:

cumulative gains of model estimates in population

Return type:

(pandas.DataFrame)

causalml.metrics.get_cumlift(df, outcome_col='y', treatment_col='w', treatment_effect_col='tau', random_seed=42)[source]#

Get average uplifts of model estimates in cumulative population.

If the true treatment effect is provided (e.g. in synthetic data), it’s calculated as the mean of the true treatment effect in each of cumulative population. Otherwise, it’s calculated as the difference between the mean outcomes of the treatment and control groups in each of cumulative population.

For details, see Section 4.1 of Gutierrez and G{‘e}rardy (2016), Causal Inference and Uplift Modeling: A review of the literature.

For the former, treatment_effect_col should be provided. For the latter, both outcome_col and treatment_col should be provided.

Parameters:

  • df (pandas.DataFrame) – a data frame with model estimates and actual data as columns

  • outcome_col (str, optional) – the column name for the actual outcome

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1)

  • treatment_effect_col (str, optional) – the column name for the true treatment effect

  • random_seed (int, optional) – deprecated

Returns:

average uplifts of model estimates in cumulative population

Return type:

(pandas.DataFrame)

causalml.metrics.get_qini(df, outcome_col='y', treatment_col='w', treatment_effect_col='tau', normalize=False, random_seed=42)[source]#

Get Qini of model estimates in population.

If the true treatment effect is provided (e.g. in synthetic data), it’s calculated as the cumulative gain of the true treatment effect in each population. Otherwise, it’s calculated as the cumulative difference between the mean outcomes of the treatment and control groups in each population.

For details, see Radcliffe (2007), Using Control Group to Target on Predicted Lift: Building and Assessing Uplift Models

For the former, treatment_effect_col should be provided. For the latter, both outcome_col and treatment_col should be provided.

Parameters:

  • df (pandas.DataFrame) – a data frame with model estimates and actual data as columns

  • outcome_col (str, optional) – the column name for the actual outcome

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1)

  • treatment_effect_col (str, optional) – the column name for the true treatment effect

  • normalize (bool, optional) – whether to normalize the y-axis to 1 or not

  • random_seed (int, optional) – deprecated

Returns:

cumulative gains of model estimates in population

Return type:

(pandas.DataFrame)

causalml.metrics.get_tmlegain(df, inference_col, learner=LGBMRegressor(learning_rate=0.05, n_estimators=300, num_leaves=64), outcome_col='y', treatment_col='w', p_col='p', n_segment=5, cv=None, ci=False)[source]#

Get TMLE based average uplifts of model estimates of segments.

Parameters:

  • df (pandas.DataFrame) – a data frame with model estimates and actual data as columns

  • inferenece_col (list of str) – a list of columns that used in learner for inference

  • learner (optional) – a model used by TMLE to estimate the outcome

  • outcome_col (str, optional) – the column name for the actual outcome

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1)

  • p_col (str, optional) – the column name for propensity score

  • n_segment (int, optional) – number of segment that TMLE will estimated for each

  • cv (sklearn.model_selection._BaseKFold, optional) – sklearn CV object

  • ci (bool, optional) – whether return confidence intervals for ATE or not

Returns:

cumulative gains of model estimates based of TMLE

Return type:

(pandas.DataFrame)

causalml.metrics.get_tmleqini(df, inference_col, learner=LGBMRegressor(learning_rate=0.05, n_estimators=300, num_leaves=64), outcome_col='y', treatment_col='w', p_col='p', n_segment=5, cv=None, ci=False, normalize=False)[source]#

Get TMLE based Qini of model estimates by segments.

Parameters:

  • df (pandas.DataFrame) – a data frame with model estimates and actual data as columns

  • inferenece_col (list of str) – a list of columns that used in learner for inference

  • learner (optional) – a model used by TMLE to estimate the outcome

  • outcome_col (str, optional) – the column name for the actual outcome

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1)

  • p_col (str, optional) – the column name for propensity score

  • n_segment (int, optional) – number of segment that TMLE will estimated for each

  • cv (sklearn.model_selection._BaseKFold, optional) – sklearn CV object

  • ci (bool, optional) – whether return confidence intervals for ATE or not

Returns:

cumulative gains of model estimates based of TMLE

Return type:

(pandas.DataFrame)

causalml.metrics.get_toc(df, outcome_col='y', treatment_col='w', treatment_effect_col='tau', normalize=False)[source]#

Get the Targeting Operator Characteristic (TOC) of model estimates in population.

TOC(q) is the difference between the ATE among the top-q fraction of units ranked by the prioritization score and the overall ATE. A positive TOC at low q indicates the model successfully identifies units with above-average treatment benefit.

By definition, TOC(0) = 0 and TOC(1) = 0 (the subset ATE equals the overall ATE when the entire population is selected).

If the true treatment effect is provided (e.g. in synthetic data), it’s used directly to calculate TOC. Otherwise, it’s estimated as the difference between the mean outcomes of the treatment and control groups in each quantile band.

Note: when using observed outcomes (i.e. without treatment_effect_col), the subset ATE is estimated via a naive difference-in-means. This is valid for randomized experiments (RCTs) but may be biased for observational data due to confounding within quantile bands. For observational settings, compute doubly-robust (AIPW) pseudo-outcomes externally and pass them as treatment_effect_col. See Yadlowsky et al. (2021), Section 4 for details.

If a quantile band contains only treated or only control units, the code falls back to TOC(q) = 0 for that band (i.e., subset ATE is set to the overall ATE). This is a conservative approximation and is logged as a warning.

For details, see Yadlowsky et al. (2021), Evaluating Treatment Prioritization Rules via Rank-Weighted Average Treatment Effects. https://arxiv.org/abs/2111.07966

For the former, treatment_effect_col should be provided. For the latter, both outcome_col and treatment_col should be provided.

Parameters:

  • df (pandas.DataFrame) – a data frame with model estimates and actual data as columns

  • outcome_col (str, optional) – the column name for the actual outcome

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1)

  • treatment_effect_col (str, optional) – the column name for the true treatment effect

  • normalize (bool, optional) – whether to normalize the TOC curve by its maximum absolute value. Uses max(|TOC|) as the reference to avoid division by zero at q=1 where TOC is always zero by definition.

Returns:

TOC values of model estimates in population, indexed by quantile q

Return type:

(pandas.DataFrame)

causalml.metrics.gini(y, p)[source]#

Normalized Gini Coefficient.

Parameters:

  • y (numpy.array) – target

  • p (numpy.array) – prediction

Returns:

normalized Gini coefficient

Return type:

e (numpy.float64)

causalml.metrics.logloss(y, p)[source]#

Bounded log loss error. :param y: target :type y: numpy.array :param p: prediction :type p: numpy.array

Returns:

bounded log loss error

causalml.metrics.mape(y, p)[source]#

Mean Absolute Percentage Error (MAPE). :param y: target :type y: numpy.array :param p: prediction :type p: numpy.array

Returns:

MAPE

Return type:

e (numpy.float64)

causalml.metrics.pehe(tau, tau_hat, squared=True)[source]#

Precision in estimating heterogeneous effects (Hill, 2011).

The mean squared error of the individual treatment effect:

PEHE = mean((tau_hat - tau) ** 2)

Papers report both this and its square root; the root is on the same scale as the effect itself, which makes it the easier one to read.

Only defined where tau is known per unit, which means simulated or semi-synthetic data. Computing it against another model’s predictions measures agreement between two models, not accuracy.

Parameters:

  • tau (np.ndarray) – the true individual treatment effect

  • tau_hat (np.ndarray) – the estimated individual treatment effect

  • squared (bool) – if False, return the root

Returns:

float, the (root) mean squared error of the individual effect

causalml.metrics.plot(df, kind='gain', tmle=False, n=100, figsize=(8, 8), ci=False, plot_chance_level=True, chance_level_kw=None, ax: Axes | None = None, *args, **kwarg) → Axes[source]#

Plot one of the lift/gain/Qini charts of model estimates.

A factory method for plot_lift(), plot_gain(), plot_qini(), plot_tmlegain() and plot_tmleqini(). For details, pleas see docstrings of each function.

Parameters:

  • df (pandas.DataFrame) – a data frame with model estimates and actual data as columns.

  • kind (str, optional) – the kind of plot to draw. ‘lift’, ‘gain’, and ‘qini’ are supported.

  • n (int, optional) – the number of samples to be used for plotting.

  • figsize (set of float, optional) – the size of the figure to plot.

  • ci (bool, optional) – whether to plot confidence intervals or not. Only available for tmle=True. Default is False.

  • plot_chance_level (bool, optional) – whether to plot the chance level (i.e., random) line or not. Default is True.

  • chance_level_line_kw (dict, optional) – the keyword arguments for the chance level line. Default is None.

  • *args – Variable length argument list.

  • **kwargs – Arbitrary keyword arguments.

causalml.metrics.plot_gain(df, outcome_col='y', treatment_col='w', treatment_effect_col='tau', normalize=False, random_seed=42, n=100, figsize=(8, 8), ax: Axes | None = None)[source]#

Plot the cumulative gain chart (or uplift curve) of model estimates.

If the true treatment effect is provided (e.g. in synthetic data), it’s calculated as the cumulative gain of the true treatment effect in each population. Otherwise, it’s calculated as the cumulative difference between the mean outcomes of the treatment and control groups in each population.

For details, see Section 4.1 of Gutierrez and G{‘e}rardy (2016), Causal Inference and Uplift Modeling: A review of the literature.

For the former, treatment_effect_col should be provided. For the latter, both outcome_col and treatment_col should be provided.

Parameters:

  • df (pandas.DataFrame) – a data frame with model estimates and actual data as columns

  • outcome_col (str, optional) – the column name for the actual outcome

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1)

  • treatment_effect_col (str, optional) – the column name for the true treatment effect

  • normalize (bool, optional) – whether to normalize the y-axis to 1 or not

  • random_seed (int, optional) – random seed for numpy.random.rand()

  • n (int, optional) – the number of samples to be used for plotting

causalml.metrics.plot_lift(df, outcome_col='y', treatment_col='w', treatment_effect_col='tau', random_seed=42, n=100, figsize=(8, 8))[source]#

Plot the lift chart of model estimates in cumulative population.

If the true treatment effect is provided (e.g. in synthetic data), it’s calculated as the mean of the true treatment effect in each of cumulative population. Otherwise, it’s calculated as the difference between the mean outcomes of the treatment and control groups in each of cumulative population.

For details, see Section 4.1 of Gutierrez and G{‘e}rardy (2016), Causal Inference and Uplift Modeling: A review of the literature.

For the former, treatment_effect_col should be provided. For the latter, both outcome_col and treatment_col should be provided.

Parameters:

  • df (pandas.DataFrame) – a data frame with model estimates and actual data as columns

  • outcome_col (str, optional) – the column name for the actual outcome

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1)

  • treatment_effect_col (str, optional) – the column name for the true treatment effect

  • random_seed (int, optional) – deprecated

  • n (int, optional) – the number of samples to be used for plotting

causalml.metrics.plot_qini(df, outcome_col='y', treatment_col='w', treatment_effect_col='tau', normalize=False, random_seed=42, n=100, figsize=(8, 8), ax: Axes | None = None) → Axes[source]#

Plot the Qini chart (or uplift curve) of model estimates.

If the true treatment effect is provided (e.g. in synthetic data), it’s calculated as the cumulative gain of the true treatment effect in each population. Otherwise, it’s calculated as the cumulative difference between the mean outcomes of the treatment and control groups in each population.

For details, see Radcliffe (2007), Using Control Group to Target on Predicted Lift: Building and Assessing Uplift Models

For the former, treatment_effect_col should be provided. For the latter, both outcome_col and treatment_col should be provided.

Parameters:

  • df (pandas.DataFrame) – a data frame with model estimates and actual data as columns

  • outcome_col (str, optional) – the column name for the actual outcome

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1)

  • treatment_effect_col (str, optional) – the column name for the true treatment effect

  • normalize (bool, optional) – whether to normalize the y-axis to 1 or not

  • random_seed (int, optional) – deprecated

  • n (int, optional) – the number of samples to be used for plotting

  • ci (bool, optional) – whether return confidence intervals for ATE or not

causalml.metrics.plot_tmlegain(df, inference_col, learner=LGBMRegressor(learning_rate=0.05, n_estimators=300, num_leaves=64, verbose=-1), outcome_col='y', treatment_col='w', p_col='tau', n_segment=5, cv=None, ci=False, figsize=(8, 8))[source]#

Plot the lift chart based of TMLE estimation

Parameters:

  • df (pandas.DataFrame) – a data frame with model estimates and actual data as columns

  • inferenece_col (list of str) – a list of columns that used in learner for inference

  • learner (optional) – a model used by TMLE to estimate the outcome

  • outcome_col (str, optional) – the column name for the actual outcome

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1)

  • p_col (str, optional) – the column name for propensity score

  • n_segment (int, optional) – number of segment that TMLE will estimated for each

  • cv (sklearn.model_selection._BaseKFold, optional) – sklearn CV object

  • ci (bool, optional) – whether return confidence intervals for ATE or not

causalml.metrics.plot_tmleqini(df, inference_col, learner=LGBMRegressor(learning_rate=0.05, n_estimators=300, num_leaves=64), outcome_col='y', treatment_col='w', p_col='tau', n_segment=5, cv=None, ci=False, figsize=(8, 8))[source]#

Plot the qini chart based of TMLE estimation

Parameters:

  • df (pandas.DataFrame) – a data frame with model estimates and actual data as columns

  • inferenece_col (list of str) – a list of columns that used in learner for inference

  • learner (optional) – a model used by TMLE to estimate the outcome

  • outcome_col (str, optional) – the column name for the actual outcome

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1)

  • p_col (str, optional) – the column name for propensity score

  • n_segment (int, optional) – number of segment that TMLE will estimated for each

  • cv (sklearn.model_selection._BaseKFold, optional) – sklearn CV object

  • ci (bool, optional) – whether return confidence intervals for ATE or not

causalml.metrics.plot_toc(df, outcome_col='y', treatment_col='w', treatment_effect_col='tau', normalize=False, n=100, figsize=(8, 8), ax: Axes | None = None) → Axes[source]#

Plot the Targeting Operator Characteristic (TOC) curve of model estimates.

The TOC(q) shows the excess ATE when treating only the top-q fraction of units prioritized by a model score, relative to the overall ATE. A positive and steeply decreasing curve indicates the model effectively ranks high-benefit units first.

If the true treatment effect is provided (e.g. in synthetic data), it’s used directly. Otherwise, it’s estimated from observed outcomes and treatment assignments.

For details, see Yadlowsky et al. (2021), Evaluating Treatment Prioritization Rules via Rank-Weighted Average Treatment Effects. https://arxiv.org/abs/2111.07966

For the former, treatment_effect_col should be provided. For the latter, both outcome_col and treatment_col should be provided.

Parameters:

  • df (pandas.DataFrame) – a data frame with model estimates and actual data as columns

  • outcome_col (str, optional) – the column name for the actual outcome

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1)

  • treatment_effect_col (str, optional) – the column name for the true treatment effect

  • normalize (bool, optional) – whether to normalize the TOC curve by its maximum absolute value before plotting

  • n (int, optional) – the number of samples to be used for plotting

  • figsize (tuple, optional) – the size of the figure to plot

  • ax (plt.Axes, optional) – an existing axes object to draw on

Returns:

the matplotlib Axes with the TOC plot

Return type:

(plt.Axes)

causalml.metrics.plug_in_t_score(df, X, treatment_col='w', outcome_col='y', learner=LGBMRegressor(learning_rate=0.05, n_estimators=300, num_leaves=64, verbose=-1), control_outcome_learner=None, treatment_outcome_learner=None, n_folds=5, return_ci=False, n_bootstrap=200, alpha=0.05, random_state=None)[source]#

Score fitted CATE models against a cross-fitted plug-in T-learner proxy.

Fits a simple T-learner – separate control/treatment outcome regressions – with n_folds-fold cross-fitting, and uses mu_1(X) - mu_0(X) on each held-out fold as a proxy for the true CATE. Candidate models are then scored by mean squared error against this proxy:

T-loss(tau_hat) = mean((tau_hat(X) - (mu_1(X) - mu_0(X))) ** 2)

Lower is better. This is a simpler baseline than dr_score() – it isn’t doubly robust and is biased under a misspecified outcome model – but Mahajan et al. (2024) found it is never dominated across their benchmark datasets despite its simplicity, making it a useful complement to DR-based scoring rather than a replacement.

Parameters:

  • df (pandas.DataFrame) – a data frame with fitted CATE model estimates as columns

  • X (numpy.ndarray or pandas.DataFrame) – feature matrix used to fit the plug-in T-learner nuisance models

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1)

  • outcome_col (str, optional) – the column name for the actual outcome

  • learner (model, optional) – a model for both control and treatment outcome regressions if the group-specific learners below are not given

  • control_outcome_learner (model, optional) – a model to estimate outcomes in the control group

  • treatment_outcome_learner (model, optional) – a model to estimate outcomes in the treatment group

  • n_folds (int, optional) – number of cross-fitting folds. Default 5

  • return_ci (bool, optional) – whether to return bootstrap confidence intervals. Default False

  • n_bootstrap (int, optional) – number of half-sample bootstrap iterations. Only used when return_ci=True. Default 200

  • alpha (float, optional) – significance level for confidence intervals. Only used when return_ci=True. Default 0.05

  • random_state (int or None, optional) – random seed for cross-fitting and the bootstrap sampler. Default None

Returns:

(pandas.Series): plug-in T-learner loss for each model column (lower is better) If return_ci=True: (pandas.DataFrame): loss, standard error, and confidence interval bounds for each model column

Return type:

If return_ci=False

causalml.metrics.policy_risk(y, treatment, tau_hat, control_name=0)[source]#

Expected loss of the treat-if-tau_hat-is-positive policy.

Following Shalit, Johansson and Sontag (2017):

R_pol = 1 - (E[Y(1) | pi = 1] P(pi = 1) + E[Y(0) | pi = 0] P(pi = 0))

where pi treats a unit when its estimated effect is positive. Each conditional expectation is estimated from the units that were actually assigned that way, which is why the treatment has to be randomized: on observational data those subgroups differ for reasons the policy did not choose, and the number stops describing the policy.

The outcome is assumed to be a benefit in [0, 1], as employment is on Jobs. For a cost, or a scale other than [0, 1], the complement to 1 is not meaningful and the policy value itself is the quantity to report.

Parameters:

  • y (np.ndarray) – the observed outcome

  • treatment (np.ndarray) – the treatment assignment

  • tau_hat (np.ndarray) – the estimated individual treatment effect

  • control_name – the value of treatment marking a control unit

Returns:

float, one minus the value of the policy

Raises:

ValueError – if either arm of the policy has no units to estimate from

causalml.metrics.qini_score(df, outcome_col='y', treatment_col='w', treatment_effect_col='tau', normalize=True, tmle=False, *args, return_ci=False, n_bootstrap=200, alpha=0.05, random_state=None, **kwarg)[source]#

Calculate the Qini score: the area between the Qini curves of a model and random.

For details, see Radcliffe (2007), Using Control Group to Target on Predicted Lift: Building and Assessing Uplift Models

Parameters:

  • df (pandas.DataFrame) – a data frame with model estimates and actual data as columns

  • outcome_col (str, optional) – the column name for the actual outcome

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1)

  • treatment_effect_col (str, optional) – the column name for the true treatment effect

  • normalize (bool, optional) – whether to normalize the y-axis to 1 or not

  • return_ci (bool, optional) – whether to return standard errors, bootstrap confidence intervals and p-values. Default False, so existing callers are unaffected.

  • n_bootstrap (int, optional) – number of half-sample bootstrap iterations. Only used when return_ci=True. Default 200.

  • alpha (float, optional) – significance level for confidence intervals. Only used when return_ci=True. Default 0.05.

  • random_state (int or None, optional) – random seed for the bootstrap sampler. Pass an integer for reproducible results. Default None.

Returns:

(pandas.Series): the Qini score for each model estimate column If return_ci=True: (pandas.DataFrame): Qini score, standard error, confidence interval bounds and p-value for each model estimate column.

Return type:

If return_ci=False

Note

The p-value tests H0: Qini = 0, which here means the model’s ranking is no better than random at finding units that benefit. That null is meaningful because the score is already the area between the model curve and the random curve — unlike AUUC, whose random baseline is about 0.5.

causalml.metrics.rate_score(df, outcome_col='y', treatment_col='w', treatment_effect_col='tau', weighting='autoc', normalize=False, return_ci=False, n_bootstrap=200, alpha=0.05, random_state=None)[source]#

Calculate the Rank-weighted Average Treatment Effect (RATE) score.

RATE is the weighted area under the Targeting Operator Characteristic (TOC) curve:

RATE = integral_0^1 alpha(q) * TOC(q) dq

Two standard weighting schemes are supported (Yadlowsky et al., 2021):

  • "autoc": alpha(q) = 1/q. Places more weight on the highest-priority units. Most powerful when treatment effects are concentrated in a small subgroup.

  • "qini": alpha(q) = q. Uniform weighting across units; reduces to the Qini coefficient. More powerful when treatment effects are diffuse across the population.

A positive RATE indicates the prioritization rule effectively identifies units with above-average treatment benefit. A RATE near zero suggests little heterogeneity or a poor prioritization rule.

Note: the integral is approximated via a weighted mean over the discrete quantile grid using midpoint values. Weights are normalized to sum to 1 (i.e. weights / weights.sum()), so the absolute scale matches the TOC values but may differ slightly from the paper’s continuous integral definition. Model rankings are preserved.

When return_ci=True, standard errors and confidence intervals are estimated via the half-sample bootstrap (m = n // 2 draws without replacement), which gives valid coverage for the RATE functional per the Yadlowsky et al. (2021) functional CLT. The p-value tests H0: RATE = 0 (i.e. the model’s prioritization is no better than random) using a two-sided z-test. When using observed outcomes (without treatment_effect_col), the underlying TOC estimates the subset ATE via naive difference-in-means, which is valid for RCTs but biased for observational data. For observational settings, pass AIPW pseudo-outcomes as treatment_effect_col. See the get_toc() docstring for details.

For details, see Yadlowsky et al. (2021), Evaluating Treatment Prioritization Rules via Rank-Weighted Average Treatment Effects. https://arxiv.org/abs/2111.07966

For the former, treatment_effect_col should be provided. For the latter, both outcome_col and treatment_col should be provided.

Parameters:

  • df (pandas.DataFrame) – a data frame with model estimates and actual data as columns

  • outcome_col (str, optional) – the column name for the actual outcome

  • treatment_col (str, optional) – the column name for the treatment indicator (0 or 1)

  • treatment_effect_col (str, optional) – the column name for the true treatment effect

  • weighting (str, optional) – the weighting scheme for the RATE integral. One of "autoc" (default) or "qini".

  • normalize (bool, optional) – whether to normalize the TOC curve before scoring

  • return_ci (bool, optional) – whether to return bootstrap confidence intervals and p-values. Default False.

  • n_bootstrap (int, optional) – number of half-sample bootstrap iterations. Only used when return_ci=True. Default 200.

  • alpha (float, optional) – significance level for confidence intervals. Only used when return_ci=True. Default 0.05.

  • random_state (int or None, optional) – random seed for the bootstrap sampler. Pass an integer for reproducible results. Default None.

Returns:

(pandas.Series): RATE scores of model estimates If return_ci=True: (pandas.DataFrame): RATE score, standard error, CI lower bound, CI upper bound, and p-value for each model estimate column

Return type:

If return_ci=False

causalml.metrics.regression_metrics(y, p, w=None, metrics={'Gini': <function gini>, 'RMSE': <function rmse>, 'sMAPE': <function smape>})[source]#

Log metrics for regressors.

Parameters:

  • y (numpy.array) – target

  • p (numpy.array) – prediction

  • w (numpy.array, optional) – a treatment vector (1 or True: treatment, 0 or False: control). If given, log metrics for the treatment and control group separately

  • metrics (dict, optional) – a dictionary of the metric names and functions

causalml.metrics.rlearner_score(df, X=None, treatment_col='w', outcome_col='y', y_residual_col=None, w_residual_col=None, outcome_learner=None, propensity_learner=None, n_folds=5, return_ci=False, n_bootstrap=200, alpha=0.05, random_state=None)[source]#

Score fitted CATE models via the R-loss (Nie & Wager, 2021).

R-loss(tau_hat) = mean[((y - m(X)) - (w - e(X)) * tau_hat(X)) ** 2]

where m(X) = E[Y|X] and e(X) = E[W|X] are cross-fitted nuisance regressions (see causalml.propensity.compute_r_residuals()). This is the loss BaseRLearner.fit() already minimizes internally to train its own effect model; exposing it standalone gives R-loss-based comparison across arbitrary fitted CATE models – EconML RScorer parity. Lower is better.

R-score complements dr_score() and plug_in_t_score() on the CATE-accuracy axis (as opposed to rate_score()’s targeting/ranking axis); Mahajan et al. (2024) found DR-loss dominates and plug-in-T is never dominated, with R-loss not the standout of the three – useful as a third opinion, particularly for parity with EconML workflows already using RScorer.

Residuals can be supplied directly (e.g. precomputed once with compute_r_residuals() and reused across scoring calls) via y_residual_col / w_residual_col, or computed internally from X, treatment_col, and outcome_col.

Parameters:

  • df (pandas.DataFrame) – a data frame with fitted CATE model estimates as columns, plus either y_residual_col/w_residual_col or both outcome_col and treatment_col

  • X (numpy.ndarray or pandas.DataFrame, optional) – feature matrix for the R-loss nuisance models. Required unless residual columns are given

  • treatment_col (str, optional) – treatment indicator column (0 or 1). Ignored if residual columns are provided

  • outcome_col (str, optional) – outcome column. Ignored if residual columns are provided

  • y_residual_col (str, optional) – precomputed y - m_hat(X) column

  • w_residual_col (str, optional) – precomputed w - e_hat(X) column

  • outcome_learner (model, optional) – model for E[Y|X]. Required unless residual columns are provided

  • propensity_learner (PropensityModel, optional) – passed to compute_r_residuals(). Defaults to ElasticNetPropensityModel

  • n_folds (int, optional) – cross-fitting folds. Default 5

  • return_ci (bool, optional) – whether to return bootstrap CIs. Default False

  • n_bootstrap (int, optional) – half-sample bootstrap iterations. Default 200

  • alpha (float, optional) – CI significance level. Default 0.05

  • random_state (int or None, optional) – random seed. Default None

Returns:

(pandas.Series): R-loss for each model column (lower is better) If return_ci=True: (pandas.DataFrame): R-loss, se, and CI bounds per model column

Return type:

If return_ci=False

causalml.metrics.rmse(y, p)[source]#

Root Mean Squared Error (RMSE). :param y: target :type y: numpy.array :param p: prediction :type p: numpy.array

Returns:

RMSE

Return type:

e (numpy.float64)

causalml.metrics.smape(y, p)[source]#

Symmetric Mean Absolute Percentage Error (sMAPE). :param y: target :type y: numpy.array :param p: prediction :type p: numpy.array

Returns:

sMAPE

Return type:

e (numpy.float64)

causalml.feature_selection module#

class causalml.feature_selection.FilterSelect[source]#

Bases: object

A class for feature importance methods.

filter_D(data, features, y_name, n_bins=10, method='KL', control_group='control', experiment_group_column='treatment_group_key', null_impute=None)[source]#

Rank features based on the chosen divergence measure.

Parameters:

  • data (pd.Dataframe) – DataFrame containing outcome, features, and experiment group

  • treatment_indicator (string) – the column name for binary indicator of treatment (1) or control (0)

  • features (list of string) – list of feature names, that are columns in the data DataFrame

  • y_name (string) – name of the outcome variable

  • method (string, optional, default = 'KL') – taking one of the following values {‘F’, ‘LR’, ‘KL’, ‘ED’, ‘Chi’} The feature selection method to be used to rank the features. ‘F’ for F-test ‘LR’ for likelihood ratio test ‘KL’, ‘ED’, ‘Chi’ for bin-based uplift filter methods, KL divergence, Euclidean distance, Chi-Square respectively

  • experiment_group_column (string, optional, default = 'treatment_group_key') – the experiment column name in the DataFrame, which contains the treatment and control assignment label

  • control_group (string, optional, default = 'control') – name for control group, value in the experiment group column

  • n_bins (int, optional, default = 10) – number of bins to be used for bin-based uplift filter methods

  • null_impute (str, optional, default=None) – impute np.nan present in the data taking on of the followin strategy values {‘mean’, ‘median’, ‘most_frequent’, None}. If Value is None and null is present then exception will be raised

Returns:

pd.DataFrame

a data frame containing the feature importance statistics

Return type:

all_result

filter_F(data, treatment_indicator, features, y_name, order=1)[source]#

Rank features based on the F-statistics of the interaction.

Parameters:

  • data (pd.Dataframe) – DataFrame containing outcome, features, and experiment group

  • treatment_indicator (string) – the column name for binary indicator of treatment (1) or control (0)

  • features (list of string) – list of feature names, that are columns in the data DataFrame

  • y_name (string) – name of the outcome variable

  • order (int) – the order of feature to be evaluated with the treatment effect, order takes 3 values: 1,2,3. order = 1 corresponds to linear importance of the feature, order=2 corresponds to quadratic and linear importance of the feature,

  • forms. (order= 3 will calculate feature importance up to cubic)

Returns:

pd.DataFrame

a data frame containing the feature importance statistics

Return type:

all_result

filter_LR(data, treatment_indicator, features, y_name, order=1, disp=True)[source]#

Rank features based on the LRT-statistics of the interaction.

Parameters:

  • data (pd.Dataframe) – DataFrame containing outcome, features, and experiment group

  • treatment_indicator (string) – the column name for binary indicator of treatment (1) or control (0)

  • feature_name (string) – feature name, as one column in the data DataFrame

  • y_name (string) – name of the outcome variable

  • order (int) – the order of feature to be evaluated with the treatment effect, order takes 3 values: 1,2,3. order = 1 corresponds to linear importance of the feature, order=2 corresponds to quadratic and linear importance of the feature,

  • forms. (order= 3 will calculate feature importance up to cubic)

Returns:

pd.DataFrame

a data frame containing the feature importance statistics

Return type:

all_result

get_importance(data, features, y_name, method, experiment_group_column='treatment_group_key', control_group='control', treatment_group='treatment', n_bins=5, null_impute=None, order=1, disp=False)[source]#

Rank features based on the chosen statistic of the interaction.

Parameters:

  • data (pd.Dataframe) – DataFrame containing outcome, features, and experiment group

  • features (list of string) – list of feature names, that are columns in the data DataFrame

  • y_name (string) – name of the outcome variable

  • method (string, optional, default = 'KL') – taking one of the following values {‘F’, ‘LR’, ‘KL’, ‘ED’, ‘Chi’} The feature selection method to be used to rank the features. ‘F’ for F-test ‘LR’ for likelihood ratio test ‘KL’, ‘ED’, ‘Chi’ for bin-based uplift filter methods, KL divergence, Euclidean distance, Chi-Square respectively

  • experiment_group_column (string) – the experiment column name in the DataFrame, which contains the treatment and control assignment label

  • control_group (string) – name for control group, value in the experiment group column

  • treatment_group (string) – name for treatment group, value in the experiment group column

  • n_bins (int, optional) – number of bins to be used for bin-based uplift filter methods

  • null_impute (str, optional, default=None) – impute np.nan present in the data taking on of the following strategy values {‘mean’, ‘median’, ‘most_frequent’, None}. If value is None and null is present then exception will be raised

  • order (int) – the order of feature to be evaluated with the treatment effect for F filter and LR filter, order takes 3 values: 1,2,3. order = 1 corresponds to linear importance of the feature, order=2 corresponds to quadratic and linear importance of the feature,

  • forms. (order= 3 will calculate feature importance up to cubic)

  • disp (bool) – Set to True to print convergence messages for Logistic regression convergence in LR method.

Returns:

pd.DataFrame

a data frame with following columns: [‘method’, ‘feature’, ‘rank’, ‘score’, ‘p_value’, ‘misc’]

Return type:

all_result

causalml.features module#

class causalml.features.LabelEncoder(min_obs=10)[source]#

Bases: BaseEstimator

Label Encoder that groups infrequent values into one label.

Code from jeongyoonlee/Kaggler

min_obs#

minimum number of observation to assign a label.

Type:

int

label_encoders#

label encoders for columns

Type:

list of dict

label_maxes#

maximum of labels for columns

Type:

list of int

fit(X, y=None)[source]#
fit_transform(X, y=None)[source]#

Encode categorical columns into label encoded columns

Parameters:

X (pandas.DataFrame) – categorical columns to encode

Returns:

label encoded columns

Return type:

X (pandas.DataFrame)

transform(X)[source]#

Encode categorical columns into label encoded columns

Parameters:

X (pandas.DataFrame) – categorical columns to encode

Returns:

label encoded columns

Return type:

X (pandas.DataFrame)

class causalml.features.OneHotEncoder(min_obs=10)[source]#

Bases: BaseEstimator

One-Hot-Encoder that groups infrequent values into one dummy variable.

Code from jeongyoonlee/Kaggler

min_obs#

minimum number of observation to create a dummy variable

Type:

int

label_encoders#

label encoders and their maximums for columns

Type:

list of (dict, int)

fit(X, y=None)[source]#
fit_transform(X, y=None)[source]#

Encode categorical columns into sparse matrix with one-hot-encoding.

Parameters:

X (pandas.DataFrame) – categorical columns to encode

Returns:

sparse matrix encoding categorical variables into dummy variables

transform(X)[source]#

Encode categorical columns into sparse matrix with one-hot-encoding.

Parameters:

X (pandas.DataFrame) – categorical columns to encode

Returns:

sparse matrix encoding categorical

variables into dummy variables

Return type:

X_new (scipy.sparse.coo_matrix)

causalml.features.load_data(data, features, transformations={})[source]#

Load data and set the feature matrix and label vector.

Parameters:

  • data (pandas.DataFrame) – total input data

  • features (list of str) – column names to be used in the inference model

  • transformation (dict of (str, func)) – transformations to be applied to features

Returns:

a feature matrix

Return type:

X (numpy.ndarray)

Module contents#

On this page
Show Source