API Documentation

DRLearner

DRLearner(
    is_classification: bool,
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel]
    | None = None,
    nuisance_model_params: Params
    | dict[str, Params]
    | None = None,
    treatment_model_params: Params
    | dict[str, Params]
    | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[
        str, list[CrossFitEstimator]
    ]
    | None = None,
    fitted_propensity_model: CrossFitEstimator
    | None = None,
    feature_set: Features
    | dict[str, Features]
    | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
    adaptive_clipping: bool = False,
)

Bases: _ConditionalAverageOutcomeMetaLearner

DR-Learner for CATE estimation as described by Kennedy (2020).

Importantly, the current DR-Learner implementation only supports:

• binary classes in case of a classification outcome

The DR-Learner contains the following nuisance models:

• a "propensity_model" estimating \(\Pr[W=k|X]\)
• one "variant_outcome_model" for each treatment variant (including control) estimating \(\mathbb{E}[Y|X, W=k]\)

and one treatment model for each treatment variant (without control):

• "treatment_model" which estimates \(\mathbb{E}[Y(k) - Y(0) | X]\)

If adaptive_clipping is set to True, then the pseudo outcomes are computed using adaptive propensity clipping described in section 4.1, equation DR-Switch of Mahajan et al. (2024).

Source code in metalearners/drlearner.py

def __init__(
    self,
    is_classification: bool,
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel] | None = None,
    nuisance_model_params: Params | dict[str, Params] | None = None,
    treatment_model_params: Params | dict[str, Params] | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[str, list[CrossFitEstimator]] | None = None,
    fitted_propensity_model: CrossFitEstimator | None = None,
    feature_set: Features | dict[str, Features] | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
    adaptive_clipping: bool = False,
):
    super().__init__(
        nuisance_model_factory=nuisance_model_factory,
        is_classification=is_classification,
        n_variants=n_variants,
        treatment_model_factory=treatment_model_factory,
        propensity_model_factory=propensity_model_factory,
        nuisance_model_params=nuisance_model_params,
        treatment_model_params=treatment_model_params,
        propensity_model_params=propensity_model_params,
        fitted_nuisance_models=fitted_nuisance_models,
        fitted_propensity_model=fitted_propensity_model,
        feature_set=feature_set,
        n_folds=n_folds,
        random_state=random_state,
    )
    self.adaptive_clipping = adaptive_clipping

init_args property

init_args: dict[str, Any]

Create initialization parameters for a new MetaLearner.

Importantly, this does not copy further internal state, such as the weights or parameters of trained base models.

average_treatment_effect

average_treatment_effect(
    X: Matrix, y: Vector, w: Vector, is_oos: bool
) -> tuple[np.ndarray, np.ndarray]

Compute Average Treatment Effect (ATE) for each treatment variant using the Augmented IPW estimator (Robins et al 1994). Does not require fitting a second- stage treatment model: it uses the pseudo-outcome alone and computes the point estimate and standard error. Can be used following the fit_all_nuisance method.

Parameters:

Name	Type	Description	Default
X	Matrix	Covariate matrix	required
y	Vector	Outcome vector	required
w	Vector	Treatment vector	required
is_oos	bool	indicator whether data is out of sample	required

Returns:

Type	Description
ndarray	np.ndarray: Treatment effect for each treatment variant.
ndarray	np.ndarray: Standard error for each treatment variant.

Source code in metalearners/drlearner.py

def average_treatment_effect(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
) -> tuple[np.ndarray, np.ndarray]:
    """Compute Average Treatment Effect (ATE) for each treatment variant using the
    Augmented IPW estimator (Robins et al 1994). Does not require fitting a second-
    stage treatment model: it uses the pseudo-outcome alone and computes the point
    estimate and standard error. Can be used following the
    [`fit_all_nuisance`][metalearners.drlearner.DRLearner.fit_all_nuisance] method.

    Args:
        X (Matrix): Covariate matrix
        y (Vector): Outcome vector
        w (Vector): Treatment vector
        is_oos (bool): indicator whether data is out of sample

    Returns:
        np.ndarray: Treatment effect for each treatment variant.
        np.ndarray: Standard error for each treatment variant.
    """
    if not self._nuisance_models_fit:
        raise ValueError(
            "The nuisance models need to be fitted before computing the treatment effect."
        )
    gamma_matrix = np.zeros((safe_len(X), self.n_variants - 1))
    for treatment_variant in range(1, self.n_variants):
        gamma_matrix[:, treatment_variant - 1] = self._pseudo_outcome(
            X=X,
            w=w,
            y=y,
            treatment_variant=treatment_variant,
            is_oos=is_oos,
        )
    treatment_effect = gamma_matrix.mean(axis=0)
    standard_error = gamma_matrix.std(axis=0) / np.sqrt(safe_len(X))
    return treatment_effect, standard_error

evaluate

evaluate(
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]

Evaluate the MetaLearner.

The keys in scoring which are not a name of a model contained in the MetaLearner will be ignored, for information about this names check nuisance_model_specifications and treatment_model_specifications. The values must be a list of:

• string representing a sklearn scoring method. Check here for the possible values.
• Callable with signature scorer(estimator, X, y_true, **kwargs). We recommend using sklearn.metrics.make_scorer to create such a Callable.

If some model name is not present in the keys of scoring then the default used metrics will be neg_log_loss if it is a classifier and neg_root_mean_squared_error if it is a regressor.

The returned dictionary keys have the following structure:

• For nuisance models:

  • If the cardinality is one: f"{model_kind}_{scorer}"
  • If there is one model for each treatment variant (including control): f"{model_kind}_{treatment_variant}_{scorer}"

• For treatment models: f"{model_kind}_{treatment_variant}_vs_0_{scorer}"

Where scorer is the name of the scorer if it is a string and "custom_scorer_{idx}" if it is a callable where idx is the index in the scorers list.

Source code in metalearners/drlearner.py

def evaluate(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]:
    safe_scoring = self._scoring(scoring)

    variant_outcome_evaluation = _evaluate_model_kind(
        cfes=self._nuisance_models[VARIANT_OUTCOME_MODEL],
        Xs=[index_matrix(X, w == tv) for tv in range(self.n_variants)],
        ys=[index_vector(y, w == tv) for tv in range(self.n_variants)],
        scorers=safe_scoring[VARIANT_OUTCOME_MODEL],
        model_kind=VARIANT_OUTCOME_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=False,
        feature_set=self.feature_set[VARIANT_OUTCOME_MODEL],
    )

    propensity_evaluation = _evaluate_model_kind(
        cfes=self._nuisance_models[PROPENSITY_MODEL],
        Xs=[X],
        ys=[w],
        scorers=safe_scoring[PROPENSITY_MODEL],
        model_kind=PROPENSITY_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=False,
        feature_set=self.feature_set[PROPENSITY_MODEL],
    )

    pseudo_outcome: list[np.ndarray] = []
    for treatment_variant in range(1, self.n_variants):
        tv_pseudo_outcome = self._pseudo_outcome(
            X=X,
            y=y,
            w=w,
            treatment_variant=treatment_variant,
            is_oos=is_oos,
            oos_method=oos_method,
        )
        pseudo_outcome.append(tv_pseudo_outcome)

    treatment_evaluation = _evaluate_model_kind(
        self._treatment_models[TREATMENT_MODEL],
        Xs=[X for _ in range(1, self.n_variants)],
        ys=pseudo_outcome,
        scorers=safe_scoring[TREATMENT_MODEL],
        model_kind=TREATMENT_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=True,
        feature_set=self.feature_set[TREATMENT_MODEL],
    )

    return variant_outcome_evaluation | propensity_evaluation | treatment_evaluation

explainer

explainer(
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer

Create an Explainer

which can be used in feature_importances.

This function can be used in two distinct manners based on the provided parameters:

• When parameters X, cate_estimates, and cate_model_factory are all set to None, the function creates an Explainer using the pre-existing treatment models. If these models do not exist, however, it triggers a ValueError.
• On the contrary, if X, cate_estimates, and cate_model_factory are not None, the function initiates an instance of the Explainer class using these parameters. This instance then fits new models for each treatment variant, and these models are employed to calculate the importance of features.

Source code in metalearners/metalearner.py

def explainer(
    self,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer:
    r"""Create an [`Explainer`][metalearners.explainer.Explainer]

    which can be used in [`feature_importances`][metalearners.metalearner.MetaLearner.feature_importances].

    This function can be used in two distinct manners based on the provided parameters:

    *   When parameters `X`, `cate_estimates`, and `cate_model_factory` are all
        set to `None`, the function creates an [`Explainer`][metalearners.explainer.Explainer]
        using the pre-existing treatment models. If these models do not exist, however,
        it triggers a `ValueError`.
    *   On the contrary, if `X`, `cate_estimates`, and `cate_model_factory` are
        not `None`, the function initiates an instance of the [`Explainer`][metalearners.explainer.Explainer]
        class using these parameters. This instance then fits new models for each
        treatment variant, and these models are employed to calculate the importance
        of features.
    """
    if X is None and cate_estimates is None and cate_model_factory is None:
        try:
            cate_cf_models = self._treatment_models[TREATMENT_MODEL]
            cate_models = [cf._overall_estimator for cf in cate_cf_models]
        except KeyError:
            raise ValueError(
                "The metalearner does not have treatment models; hence X, cate_estimates, "
                "and cate_model_factory need to be defined."
            )
        return Explainer(cate_models=cate_models)  # type: ignore
    elif (
        X is not None
        and cate_estimates is not None
        and cate_model_factory is not None
    ):
        return Explainer.from_estimates(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    else:
        raise ValueError(
            "Either all of [X, cate_estimates, cate_model_factory] are None or all "
            "of them must be defined."
        )

feature_importances

feature_importances(
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances

Calculates the feature importance for each treatment group.

If explainer is None, a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters X, cate_estimates, cate_model_factory and cate_model_params are ignored.

If normalization = True, for each treatment variant the feature importances are normalized so that they sum to 1.

feature_names is optional but in the case it’s not passed the names of the features will default to f"Feature {i}" where i is the corresponding feature index.

The returned list contains the feature importances for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def feature_importances(
    self,
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances:
    r"""Calculates the feature importance for each treatment group.

    If `explainer` is `None`, a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `X`, `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    If `normalization = True`, for each treatment variant the feature importances
    are normalized so that they sum to 1.

    `feature_names` is optional but in the case it's not passed the names of the
    features will default to `f"Feature {i}"` where `i` is the corresponding
    feature index.

    The returned list contains the feature importances for each treatment variant in
    ascending order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.feature_importances(
        normalize=normalize, feature_names=feature_names, sort_values=sort_values
    )

fit

fit(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

n_jobs_cross_fitting will be used at the cross-fitting level and n_jobs_base_learners will be used at the stage level. None means 1 unless in a joblib.parallel_backend context. -1 means using all processors. For more information about parallelism check parallelism.

fit_params is an optional dict to be forwarded to base estimator fit calls. It supports two usages patterns:

fit_params={"parameter_of_interest": value_of_interest}

fit_params={
    "nuisance": {
        "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
        "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
    },
    "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
}

In the former approach, the parameter and value of interest are passed to all base models. In the the latter approach, the explicitly qualified parameter-value pairs are passed to respective base models and no fitting parameters are passed to base models not explicitly listed. Note that in this pattern, propensity models are considered a nuisance model.

synchronize_cross_fitting indicates whether the learning of different base models should use exactly the same data splits where possible. Note that if there are several models to be synchronized which are classifiers, these cannot be split via stratification.

Source code in metalearners/metalearner.py

def fit(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    """Fit all models of the MetaLearner.

    If pre-fitted models were passed at instantiation, these are never refitted.

    `n_jobs_cross_fitting` will be used at the cross-fitting level and
    `n_jobs_base_learners` will be used at the stage level. `None` means 1 unless in a
    [`joblib.parallel_backend`](https://joblib.readthedocs.io/en/latest/generated/joblib.parallel_backend.html#joblib.parallel_backend)
    context. `-1` means using all processors.
    For more information about parallelism check [parallelism](parallelism.md).


    `fit_params` is an optional `dict` to be forwarded to base estimator `fit` calls. It supports
    two usages patterns:

    ```python
    fit_params={"parameter_of_interest": value_of_interest}
    ```

    ```python
    fit_params={
        "nuisance": {
            "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
            "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
        },
        "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
    }
    ```

    In the former approach, the parameter and value of interest are passed to all base models. In the
    the latter approach, the explicitly qualified parameter-value pairs are passed to respective base
    models and no fitting parameters are passed to base models not explicitly listed. Note that in this
    pattern, propensity models are considered a nuisance model.

    `synchronize_cross_fitting` indicates whether the learning of different base models should use exactly
    the same data splits where possible. Note that if there are several models to be synchronized which are
    classifiers, these cannot be split via stratification.
    """
    self.fit_all_nuisance(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    self.fit_all_treatment(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    return self

fit_all_nuisance

fit_all_nuisance(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all nuisance models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the nuisance models, and in the case they should also be used by the treatment models, these should also be passed in the following call to fit_all_treatment.

This method, combined with fit_all_treatment, facilitates the segmentation of the metalearner fitting process into two distinct parts. This division allows for interventions between the two stages, such as performing feature selection for the treatment models or conducting hyperparameter optimization within the nuisance models.

Source code in metalearners/drlearner.py

def fit_all_nuisance(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    self._validate_treatment(w)
    self._validate_outcome(y, w)

    self._treatment_variants_mask = []

    qualified_fit_params = self._qualified_fit_params(fit_params)

    for treatment_variant in range(self.n_variants):
        self._treatment_variants_mask.append(w == treatment_variant)

    self._cv_split_indices: SplitIndices | None

    if synchronize_cross_fitting:
        self._cv_split_indices = self._split(X)
    else:
        self._cv_split_indices = None

    nuisance_jobs: list[_ParallelJoblibSpecification | None] = []
    for treatment_variant in range(self.n_variants):
        mask = self._treatment_variants_mask[treatment_variant]
        X_masked = index_matrix(X, mask)
        y_masked = index_vector(y, mask)
        nuisance_jobs.append(
            self._nuisance_joblib_specifications(
                X=X_masked,
                y=y_masked,
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=treatment_variant,
                n_jobs_cross_fitting=n_jobs_cross_fitting,
                fit_params=qualified_fit_params[NUISANCE][VARIANT_OUTCOME_MODEL],
            )
        )

    nuisance_jobs.append(
        self._nuisance_joblib_specifications(
            X=X,
            y=w,
            model_kind=PROPENSITY_MODEL,
            model_ord=0,
            n_jobs_cross_fitting=n_jobs_cross_fitting,
            fit_params=qualified_fit_params[NUISANCE][PROPENSITY_MODEL],
            cv=self._cv_split_indices,
        )
    )

    parallel = Parallel(n_jobs=n_jobs_base_learners)
    results = parallel(
        delayed(_fit_cross_fit_estimator_joblib)(job)
        for job in nuisance_jobs
        if job is not None
    )

    self._assign_joblib_nuisance_results(results)
    self._nuisance_models_fit = True
    return self

fit_all_treatment

fit_all_treatment(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all treatment models of the MetaLearner.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the treatment models, as the nuisance models should already be fitted.

Source code in metalearners/drlearner.py

def fit_all_treatment(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    if not hasattr(self, "_cv_split_indices"):
        raise ValueError(
            "The nuisance models need to be fitted before fitting the treatment models."
            "In particular, the MetaLearner's attribute _cv_split_indices, "
            "typically set during nuisance fitting, does not exist."
        )
    qualified_fit_params = self._qualified_fit_params(fit_params)
    treatment_jobs: list[_ParallelJoblibSpecification] = []
    for treatment_variant in range(1, self.n_variants):
        pseudo_outcomes = self._pseudo_outcome(
            X=X,
            w=w,
            y=y,
            treatment_variant=treatment_variant,
            is_oos=False,
        )

        treatment_jobs.append(
            self._treatment_joblib_specifications(
                X=X,
                y=pseudo_outcomes,
                model_kind=TREATMENT_MODEL,
                model_ord=treatment_variant - 1,
                n_jobs_cross_fitting=n_jobs_cross_fitting,
                fit_params=qualified_fit_params[TREATMENT][TREATMENT_MODEL],
                cv=self._cv_split_indices,
            )
        )
    parallel = Parallel(n_jobs=n_jobs_base_learners)
    results = parallel(
        delayed(_fit_cross_fit_estimator_joblib)(job) for job in treatment_jobs
    )
    self._assign_joblib_treatment_results(results)
    return self

fit_nuisance

fit_nuisance(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit a given nuisance model of a MetaLearner.

y represents the objective of the given nuisance model, not necessarily the outcome of the experiment. If pre-fitted models were passed at instantiation, these are never refitted.

Source code in metalearners/metalearner.py

def fit_nuisance(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit a given nuisance model of a MetaLearner.

    `y` represents the objective of the given nuisance model, not necessarily the outcome of the experiment.
    If pre-fitted models were passed at instantiation, these are never refitted.
    """
    if model_kind in self._prefitted_nuisance_models:
        return self
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._nuisance_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

fit_treatment

fit_treatment(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit the treatment model of a MetaLearner.

y represents the objective of the given treatment model, not necessarily the outcome of the experiment.

Source code in metalearners/metalearner.py

def fit_treatment(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit the treatment model of a MetaLearner.

    `y` represents the objective of the given treatment model, not necessarily the outcome of the experiment.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._treatment_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

nuisance_model_specifications classmethod

nuisance_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all first-stage models.

Source code in metalearners/drlearner.py

@classmethod
def nuisance_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return {
        PROPENSITY_MODEL: _ModelSpecifications(
            cardinality=get_one,
            predict_method=get_predict_proba,
        ),
        VARIANT_OUTCOME_MODEL: _ModelSpecifications(
            cardinality=MetaLearner._get_n_variants,
            predict_method=MetaLearner._outcome_predict_method,
        ),
    }

predict

predict(
    X, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Estimate the CATE.

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants} - 1, 1)\) if the outcome is a scalar, i.e. in case of a regression problem.

• \((n_{obs}, n_{variants} - 1, n_{classes})\) if the outcome is a class, i.e. in case of a classification problem.

In the case of multiple treatment variants, the second dimension represents the CATE of the corresponding variant vs the control (variant 0).

Source code in metalearners/drlearner.py

def predict(
    self,
    X,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    n_outputs = 2 if self.is_classification else 1
    estimates = np.zeros((safe_len(X), self.n_variants - 1, n_outputs))
    for treatment_variant in range(1, self.n_variants):
        estimates_variant = self.predict_treatment(
            X,
            is_oos=is_oos,
            oos_method=oos_method,
            model_kind=TREATMENT_MODEL,
            model_ord=treatment_variant - 1,
        )
        if self.is_classification:
            # This is to be consistent with other MetaLearners (e.g. S and T) that automatically
            # work with multiclass outcomes and return the CATE estimate for each class. As the DR-Learner only
            # works with binary classes (the pseudo outcome formula does not make sense with
            # multiple classes unless some adaptation is done) we can manually infer the
            # CATE estimate for the complementary class  -- returning a matrix of shape (N, 2).
            estimates_variant = np.stack(
                [-estimates_variant, estimates_variant], axis=1
            )
        else:
            estimates_variant = np.expand_dims(estimates_variant, 1)

        estimates[:, treatment_variant - 1] = estimates_variant
    return estimates

predict_conditional_average_outcomes

predict_conditional_average_outcomes(
    X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Predict the vectors of conditional average outcomes.

These are defined as \(\mathbb{E}[Y_i(w) | X]\) for each treatment variant \(w\).

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants}, 1)\) if the outcome is a scalar, i.e., in case of a regression problem.

• \((n_{obs}, n_{variants}, n_{classes})\) if the outcome is a class, i.e., in case of a classification problem.

Source code in metalearners/metalearner.py

def predict_conditional_average_outcomes(
    self, X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray:
    if self._treatment_variants_mask is None:
        raise ValueError(
            "The metalearner needs to be fitted before predicting."
            "In particular, the MetaLearner's attribute _treatment_variant_mask, "
            "typically set during fitting, is None."
        )
    # TODO: Consider multiprocessing
    n_obs = safe_len(X)
    nuisance_tensors = self._nuisance_tensors(n_obs)
    conditional_average_outcomes_list = nuisance_tensors[VARIANT_OUTCOME_MODEL]

    for tv in range(self.n_variants):
        if is_oos:
            conditional_average_outcomes_list[tv] = self.predict_nuisance(
                X=X,
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=True,
                oos_method=oos_method,
            )
        else:
            conditional_average_outcomes_list[tv][
                self._treatment_variants_mask[tv]
            ] = self.predict_nuisance(
                X=index_matrix(X, self._treatment_variants_mask[tv]),
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=False,
            )
            conditional_average_outcomes_list[tv][
                ~self._treatment_variants_mask[tv]
            ] = self.predict_nuisance(
                X=index_matrix(X, ~self._treatment_variants_mask[tv]),
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=True,
                oos_method=oos_method,
            )
    return np.stack(conditional_average_outcomes_list, axis=1).reshape(
        n_obs, self.n_variants, -1
    )

predict_nuisance

predict_nuisance(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given nuisance model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_nuisance(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given nuisance model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    predict_method_name = self.nuisance_model_specifications()[model_kind][
        "predict_method"
    ](self)
    predict_method = getattr(
        self._nuisance_models[model_kind][model_ord], predict_method_name
    )
    return predict_method(X_filtered, is_oos, oos_method)

predict_treatment

predict_treatment(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given treatment model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_treatment(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given treatment model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    return self._treatment_models[model_kind][model_ord].predict(
        X_filtered, is_oos, oos_method
    )

shap_values

shap_values(
    X: Matrix,
    shap_explainer_factory: type[Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]

Calculates the shap values for each treatment group.

If explainer is None a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters cate_estimates, cate_model_factory and cate_model_params are ignored.

The parameter shap_explainer_factory can be used to specify the type of shap explainer, for the different options see here.

The returned list contains the shap values for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def shap_values(
    self,
    X: Matrix,
    shap_explainer_factory: type[shap.Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]:
    """Calculates the shap values for each treatment group.

    If `explainer` is `None` a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    The parameter `shap_explainer_factory` can be used to specify the type of shap
    explainer, for the different options see
    [here](https://shap.readthedocs.io/en/latest/api.html#explainers).

    The returned list contains the shap values for each treatment variant in ascending
    order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=None if cate_estimates is None else X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.shap_values(
        X=X,
        shap_explainer_factory=shap_explainer_factory,
        shap_explainer_params=shap_explainer_params,
    )

treatment_model_specifications classmethod

treatment_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all second-stage models.

Source code in metalearners/drlearner.py

@classmethod
def treatment_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return {
        TREATMENT_MODEL: _ModelSpecifications(
            cardinality=MetaLearner._get_n_variants_minus_one,
            predict_method=get_predict,
        )
    }

RLearner

RLearner(
    is_classification: bool,
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel]
    | None = None,
    nuisance_model_params: Params
    | dict[str, Params]
    | None = None,
    treatment_model_params: Params
    | dict[str, Params]
    | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[
        str, list[CrossFitEstimator]
    ]
    | None = None,
    fitted_propensity_model: CrossFitEstimator
    | None = None,
    feature_set: Features
    | dict[str, Features]
    | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
)

Bases: MetaLearner

R-Learner for CATE estimation as described by Nie et al. (2017).

Importantly, the current R-Learner implementation only supports:

• binary classes in case of a classification outcome

The R-Learner contains two nuisance models

• a "propensity_model" estimating \(\Pr[W=k|X]\)
• an "outcome_model" estimating \(\mathbb{E}[Y|X]\)

and one treatment model per treatment variant which isn’t control

• "treatment_model" which estimates \(\mathbb{E}[Y(k) - Y(0) | X]\)

The treatment_model_factory provided needs to support the argument sample_weight in its fit method.

Source code in metalearners/metalearner.py

def __init__(
    self,
    is_classification: bool,
    # TODO: Consider whether we can make this not a state of the MetaLearner
    # but rather just a parameter of a predict call.
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel] | None = None,
    nuisance_model_params: Params | dict[str, Params] | None = None,
    treatment_model_params: Params | dict[str, Params] | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[str, list[CrossFitEstimator]] | None = None,
    fitted_propensity_model: CrossFitEstimator | None = None,
    feature_set: Features | dict[str, Features] | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
):
    nuisance_model_specifications = self.nuisance_model_specifications()
    treatment_model_specifications = self.treatment_model_specifications()

    if PROPENSITY_MODEL in treatment_model_specifications:
        raise ValueError(
            f"{PROPENSITY_MODEL} can't be used as a treatment model name"
        )
    if (
        isinstance(nuisance_model_factory, dict)
        and PROPENSITY_MODEL in nuisance_model_factory.keys()
    ):
        raise ValueError(
            "Propensity model factory should be defined using propensity_model_factory "
            "and not nuisance_model_factory."
        )
    if (
        isinstance(nuisance_model_params, dict)
        and PROPENSITY_MODEL in nuisance_model_params.keys()
    ):
        raise ValueError(
            "Propensity model params should be defined using propensity_model_params "
            "and not nuisance_model_params."
        )
    if (
        PROPENSITY_MODEL in nuisance_model_specifications
        and propensity_model_factory is None
        and fitted_propensity_model is None
    ):
        raise ValueError(
            "propensity_model_factory or fitted_propensity_model needs to be defined "
            f"as the {self.__class__.__name__} has a propensity model."
        )

    self._validate_n_variants(n_variants)
    self.is_classification = is_classification
    self.n_variants = n_variants

    self.nuisance_model_factory = _combine_propensity_and_nuisance_specs(
        propensity_model_factory,
        nuisance_model_factory,
        set(nuisance_model_specifications.keys()),
    )
    if nuisance_model_params is None:
        nuisance_model_params = {}  # type: ignore
    if propensity_model_params is None:
        propensity_model_params = {}
    self.nuisance_model_params = _combine_propensity_and_nuisance_specs(
        propensity_model_params,
        nuisance_model_params,
        set(nuisance_model_specifications.keys()),
    )

    self.treatment_model_factory = _initialize_model_dict(
        treatment_model_factory, set(treatment_model_specifications.keys())
    )
    if treatment_model_params is None:
        self.treatment_model_params = _initialize_model_dict(
            {}, set(treatment_model_specifications.keys())
        )
    else:
        self.treatment_model_params = _initialize_model_dict(
            treatment_model_params, set(treatment_model_specifications.keys())
        )

    self.n_folds = _initialize_model_dict(
        n_folds,
        set(nuisance_model_specifications.keys())
        | set(treatment_model_specifications.keys()),
    )
    for model_kind, n_folds_model_kind in self.n_folds.items():
        validate_number_positive(n_folds_model_kind, f"{model_kind} n_folds", True)
    self.random_state = random_state

    self.feature_set = _initialize_model_dict(
        feature_set,
        set(nuisance_model_specifications.keys())
        | set(treatment_model_specifications.keys()),
    )

    self._nuisance_models: dict[str, list[CrossFitEstimator]] = {}
    not_fitted_nuisance_models = set(nuisance_model_specifications.keys())
    self._prefitted_nuisance_models: set[str] = set()

    if fitted_nuisance_models is not None:
        if not set(fitted_nuisance_models.keys()) <= set(
            nuisance_model_specifications.keys()
        ) - {PROPENSITY_MODEL}:
            raise ValueError(
                "The keys present in fitted_nuisance_models should be a subset of "
                f"{set(nuisance_model_specifications.keys()) - {PROPENSITY_MODEL}}"
            )
        self._nuisance_models |= deepcopy(fitted_nuisance_models)
        not_fitted_nuisance_models -= set(fitted_nuisance_models.keys())
        self._prefitted_nuisance_models |= set(fitted_nuisance_models.keys())

    if (
        PROPENSITY_MODEL in nuisance_model_specifications.keys()
        and fitted_propensity_model is not None
    ):
        self._nuisance_models |= {PROPENSITY_MODEL: [fitted_propensity_model]}
        not_fitted_nuisance_models -= {PROPENSITY_MODEL}
        self._prefitted_nuisance_models |= {PROPENSITY_MODEL}

    for name in not_fitted_nuisance_models:
        if self.nuisance_model_factory[name] is None:
            if name == PROPENSITY_MODEL:
                raise ValueError(
                    f"A model for the nuisance model {name} needs to be defined. Either "
                    "in propensity_model_factory or in fitted_propensity_model."
                )
            else:
                raise ValueError(
                    f"A model for the nuisance model {name} needs to be defined. Either "
                    "in nuisance_model_factory or in fitted_nuisance_models."
                )

    self._nuisance_models |= {
        name: [
            CrossFitEstimator(
                n_folds=self.n_folds[name],
                estimator_factory=self.nuisance_model_factory[name],
                estimator_params=self.nuisance_model_params[name],
                random_state=self.random_state,
            )
            for _ in range(nuisance_model_specifications[name]["cardinality"](self))
        ]
        for name in not_fitted_nuisance_models
    }
    self._treatment_models: dict[str, list[CrossFitEstimator]] = {
        name: [
            CrossFitEstimator(
                n_folds=self.n_folds[name],
                estimator_factory=self.treatment_model_factory[name],
                estimator_params=self.treatment_model_params[name],
                random_state=self.random_state,
            )
            for _ in range(
                treatment_model_specifications[name]["cardinality"](self)
            )
        ]
        for name in set(treatment_model_specifications.keys())
    }

    self._validate_models()

init_args property

init_args: dict[str, Any]

Create initialization parameters for a new MetaLearner.

Importantly, this does not copy further internal state, such as the weights or parameters of trained base models.

evaluate

evaluate(
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]

Evaluate the MetaLearner.

The keys in scoring which are not a name of a model contained in the MetaLearner will be ignored, for information about this names check nuisance_model_specifications and treatment_model_specifications. The values must be a list of:

• string representing a sklearn scoring method. Check here for the possible values.
• Callable with signature scorer(estimator, X, y_true, **kwargs). We recommend using sklearn.metrics.make_scorer to create such a Callable.

If some model name is not present in the keys of scoring then the default used metrics will be neg_log_loss if it is a classifier and neg_root_mean_squared_error if it is a regressor.

The returned dictionary keys have the following structure:

• For nuisance models:

  • If the cardinality is one: f"{model_kind}_{scorer}"
  • If there is one model for each treatment variant (including control): f"{model_kind}_{treatment_variant}_{scorer}"

• For treatment models: f"{model_kind}_{treatment_variant}_vs_0_{scorer}"

Where scorer is the name of the scorer if it is a string and "custom_scorer_{idx}" if it is a callable where idx is the index in the scorers list.

In the RLearner case, the "treatment_model" is always evaluated with the r_loss besides the scorers in scoring["treatment_model"], which should support passing the sample_weight keyword argument.

Source code in metalearners/rlearner.py

def evaluate(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]:
    """In the RLearner case, the `"treatment_model"` is always evaluated with the
    [`r_loss`][metalearners.rlearner.r_loss] besides the scorers in
    `scoring["treatment_model"]`, which should support passing the `sample_weight`
    keyword argument."""
    safe_scoring = self._scoring(scoring)

    propensity_evaluation = _evaluate_model_kind(
        cfes=self._nuisance_models[PROPENSITY_MODEL],
        Xs=[X],
        ys=[w],
        scorers=safe_scoring[PROPENSITY_MODEL],
        model_kind=PROPENSITY_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=False,
        feature_set=self.feature_set[PROPENSITY_MODEL],
    )

    outcome_evaluation = _evaluate_model_kind(
        cfes=self._nuisance_models[OUTCOME_MODEL],
        Xs=[X],
        ys=[y],
        scorers=safe_scoring[OUTCOME_MODEL],
        model_kind=OUTCOME_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=False,
        feature_set=self.feature_set[OUTCOME_MODEL],
    )

    # TODO: improve this? generalize it to other metalearners?
    w_hat = self.predict_nuisance(
        X=X,
        is_oos=is_oos,
        oos_method=oos_method,
        model_kind=PROPENSITY_MODEL,
        model_ord=0,
    )

    y_hat = self.predict_nuisance(
        X=X,
        is_oos=is_oos,
        oos_method=oos_method,
        model_kind=OUTCOME_MODEL,
        model_ord=0,
    )
    if self.is_classification:
        y_hat = y_hat[:, 1]

    pseudo_outcome: list[np.ndarray] = []
    sample_weights: list[np.ndarray] = []
    masks: list[Vector] = []
    is_control = w == 0
    for treatment_variant in range(1, self.n_variants):
        is_treatment = w == treatment_variant
        mask = is_treatment | is_control
        tv_pseudo_outcome, tv_sample_weights = self._pseudo_outcome_and_weights(
            X=X,
            y=y,
            w=w,
            treatment_variant=treatment_variant,
            is_oos=is_oos,
            oos_method=oos_method,
            mask=mask,
        )
        pseudo_outcome.append(tv_pseudo_outcome)
        sample_weights.append(tv_sample_weights)
        masks.append(mask)

    treatment_evaluation = _evaluate_model_kind(
        self._treatment_models[TREATMENT_MODEL],
        Xs=[index_matrix(X, masks[tv - 1]) for tv in range(1, self.n_variants)],
        ys=pseudo_outcome,
        scorers=safe_scoring[TREATMENT_MODEL],
        model_kind=TREATMENT_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=True,
        sample_weights=sample_weights,
        feature_set=self.feature_set[TREATMENT_MODEL],
    )

    rloss_evaluation = {}
    tau_hat = self.predict(X=X, is_oos=is_oos, oos_method=oos_method)
    is_control = w == 0
    for treatment_variant in range(1, self.n_variants):
        is_treatment = w == treatment_variant
        mask = is_treatment | is_control

        propensity_estimates = w_hat[:, treatment_variant] / (
            w_hat[:, 0] + w_hat[:, treatment_variant]
        )
        cate_estimates = (
            tau_hat[:, treatment_variant - 1, 1]
            if self.is_classification
            else tau_hat[:, treatment_variant - 1, 0]
        )
        rloss_evaluation[f"r_loss_{treatment_variant}_vs_0"] = r_loss(
            cate_estimates=cate_estimates[mask],
            outcome_estimates=y_hat[mask],
            propensity_scores=propensity_estimates[mask],
            outcomes=index_vector(y, mask),
            treatments=index_vector(w, mask) == treatment_variant,
        )
    return (
        propensity_evaluation
        | outcome_evaluation
        | rloss_evaluation
        | treatment_evaluation
    )

explainer

explainer(
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer

Create an Explainer

which can be used in feature_importances.

This function can be used in two distinct manners based on the provided parameters:

• When parameters X, cate_estimates, and cate_model_factory are all set to None, the function creates an Explainer using the pre-existing treatment models. If these models do not exist, however, it triggers a ValueError.
• On the contrary, if X, cate_estimates, and cate_model_factory are not None, the function initiates an instance of the Explainer class using these parameters. This instance then fits new models for each treatment variant, and these models are employed to calculate the importance of features.

Source code in metalearners/metalearner.py

def explainer(
    self,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer:
    r"""Create an [`Explainer`][metalearners.explainer.Explainer]

    which can be used in [`feature_importances`][metalearners.metalearner.MetaLearner.feature_importances].

    This function can be used in two distinct manners based on the provided parameters:

    *   When parameters `X`, `cate_estimates`, and `cate_model_factory` are all
        set to `None`, the function creates an [`Explainer`][metalearners.explainer.Explainer]
        using the pre-existing treatment models. If these models do not exist, however,
        it triggers a `ValueError`.
    *   On the contrary, if `X`, `cate_estimates`, and `cate_model_factory` are
        not `None`, the function initiates an instance of the [`Explainer`][metalearners.explainer.Explainer]
        class using these parameters. This instance then fits new models for each
        treatment variant, and these models are employed to calculate the importance
        of features.
    """
    if X is None and cate_estimates is None and cate_model_factory is None:
        try:
            cate_cf_models = self._treatment_models[TREATMENT_MODEL]
            cate_models = [cf._overall_estimator for cf in cate_cf_models]
        except KeyError:
            raise ValueError(
                "The metalearner does not have treatment models; hence X, cate_estimates, "
                "and cate_model_factory need to be defined."
            )
        return Explainer(cate_models=cate_models)  # type: ignore
    elif (
        X is not None
        and cate_estimates is not None
        and cate_model_factory is not None
    ):
        return Explainer.from_estimates(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    else:
        raise ValueError(
            "Either all of [X, cate_estimates, cate_model_factory] are None or all "
            "of them must be defined."
        )

feature_importances

feature_importances(
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances

Calculates the feature importance for each treatment group.

If explainer is None, a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters X, cate_estimates, cate_model_factory and cate_model_params are ignored.

If normalization = True, for each treatment variant the feature importances are normalized so that they sum to 1.

feature_names is optional but in the case it’s not passed the names of the features will default to f"Feature {i}" where i is the corresponding feature index.

The returned list contains the feature importances for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def feature_importances(
    self,
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances:
    r"""Calculates the feature importance for each treatment group.

    If `explainer` is `None`, a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `X`, `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    If `normalization = True`, for each treatment variant the feature importances
    are normalized so that they sum to 1.

    `feature_names` is optional but in the case it's not passed the names of the
    features will default to `f"Feature {i}"` where `i` is the corresponding
    feature index.

    The returned list contains the feature importances for each treatment variant in
    ascending order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.feature_importances(
        normalize=normalize, feature_names=feature_names, sort_values=sort_values
    )

fit

fit(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

n_jobs_cross_fitting will be used at the cross-fitting level and n_jobs_base_learners will be used at the stage level. None means 1 unless in a joblib.parallel_backend context. -1 means using all processors. For more information about parallelism check parallelism.

fit_params is an optional dict to be forwarded to base estimator fit calls. It supports two usages patterns:

fit_params={"parameter_of_interest": value_of_interest}

fit_params={
    "nuisance": {
        "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
        "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
    },
    "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
}

In the former approach, the parameter and value of interest are passed to all base models. In the the latter approach, the explicitly qualified parameter-value pairs are passed to respective base models and no fitting parameters are passed to base models not explicitly listed. Note that in this pattern, propensity models are considered a nuisance model.

synchronize_cross_fitting indicates whether the learning of different base models should use exactly the same data splits where possible. Note that if there are several models to be synchronized which are classifiers, these cannot be split via stratification.

Source code in metalearners/metalearner.py

def fit(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    """Fit all models of the MetaLearner.

    If pre-fitted models were passed at instantiation, these are never refitted.

    `n_jobs_cross_fitting` will be used at the cross-fitting level and
    `n_jobs_base_learners` will be used at the stage level. `None` means 1 unless in a
    [`joblib.parallel_backend`](https://joblib.readthedocs.io/en/latest/generated/joblib.parallel_backend.html#joblib.parallel_backend)
    context. `-1` means using all processors.
    For more information about parallelism check [parallelism](parallelism.md).


    `fit_params` is an optional `dict` to be forwarded to base estimator `fit` calls. It supports
    two usages patterns:

    ```python
    fit_params={"parameter_of_interest": value_of_interest}
    ```

    ```python
    fit_params={
        "nuisance": {
            "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
            "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
        },
        "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
    }
    ```

    In the former approach, the parameter and value of interest are passed to all base models. In the
    the latter approach, the explicitly qualified parameter-value pairs are passed to respective base
    models and no fitting parameters are passed to base models not explicitly listed. Note that in this
    pattern, propensity models are considered a nuisance model.

    `synchronize_cross_fitting` indicates whether the learning of different base models should use exactly
    the same data splits where possible. Note that if there are several models to be synchronized which are
    classifiers, these cannot be split via stratification.
    """
    self.fit_all_nuisance(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    self.fit_all_treatment(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    return self

fit_all_nuisance

fit_all_nuisance(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all nuisance models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the nuisance models, and in the case they should also be used by the treatment models, these should also be passed in the following call to fit_all_treatment.

This method, combined with fit_all_treatment, facilitates the segmentation of the metalearner fitting process into two distinct parts. This division allows for interventions between the two stages, such as performing feature selection for the treatment models or conducting hyperparameter optimization within the nuisance models.

Source code in metalearners/rlearner.py

def fit_all_nuisance(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    self._validate_treatment(w)
    self._validate_outcome(y, w)

    qualified_fit_params = self._qualified_fit_params(fit_params)
    self._validate_fit_params(qualified_fit_params)

    if synchronize_cross_fitting:
        cv_split_indices = self._split(X)
    else:
        cv_split_indices = None

    nuisance_jobs: list[_ParallelJoblibSpecification | None] = []

    nuisance_jobs.append(
        self._nuisance_joblib_specifications(
            X=X,
            y=w,
            model_kind=PROPENSITY_MODEL,
            model_ord=0,
            n_jobs_cross_fitting=n_jobs_cross_fitting,
            fit_params=qualified_fit_params[NUISANCE][PROPENSITY_MODEL],
            cv=cv_split_indices,
        )
    )
    nuisance_jobs.append(
        self._nuisance_joblib_specifications(
            X=X,
            y=y,
            model_kind=OUTCOME_MODEL,
            model_ord=0,
            n_jobs_cross_fitting=n_jobs_cross_fitting,
            fit_params=qualified_fit_params[NUISANCE][OUTCOME_MODEL],
            cv=cv_split_indices,
        )
    )

    parallel = Parallel(n_jobs=n_jobs_base_learners)
    results = parallel(
        delayed(_fit_cross_fit_estimator_joblib)(job)
        for job in nuisance_jobs
        if job is not None
    )
    self._assign_joblib_nuisance_results(results)

    return self

fit_all_treatment

fit_all_treatment(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
    epsilon: float = _EPSILON,
) -> Self

Fit all treatment models of the MetaLearner.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the treatment models, as the nuisance models should already be fitted.

Source code in metalearners/rlearner.py

def fit_all_treatment(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
    epsilon: float = _EPSILON,
) -> Self:
    qualified_fit_params = self._qualified_fit_params(fit_params)
    treatment_jobs: list[_ParallelJoblibSpecification] = []
    self._variants_indices = []
    for treatment_variant in range(1, self.n_variants):

        is_treatment = w == treatment_variant
        is_control = w == 0
        mask = is_treatment | is_control

        self._variants_indices.append(mask)

        pseudo_outcomes, weights = self._pseudo_outcome_and_weights(
            X=X,
            w=w,
            y=y,
            treatment_variant=treatment_variant,
            mask=mask,
            epsilon=epsilon,
            is_oos=False,
        )

        X_filtered = index_matrix(X, mask)

        treatment_jobs.append(
            self._treatment_joblib_specifications(
                X=X_filtered,
                y=pseudo_outcomes,
                model_kind=TREATMENT_MODEL,
                model_ord=treatment_variant - 1,
                fit_params=qualified_fit_params[TREATMENT][TREATMENT_MODEL]
                | {_SAMPLE_WEIGHT: weights},
                n_jobs_cross_fitting=n_jobs_cross_fitting,
            )
        )
    parallel = Parallel(n_jobs=n_jobs_base_learners)
    results = parallel(
        delayed(_fit_cross_fit_estimator_joblib)(job) for job in treatment_jobs
    )
    self._assign_joblib_treatment_results(results)
    return self

fit_nuisance

fit_nuisance(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit a given nuisance model of a MetaLearner.

y represents the objective of the given nuisance model, not necessarily the outcome of the experiment. If pre-fitted models were passed at instantiation, these are never refitted.

Source code in metalearners/metalearner.py

def fit_nuisance(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit a given nuisance model of a MetaLearner.

    `y` represents the objective of the given nuisance model, not necessarily the outcome of the experiment.
    If pre-fitted models were passed at instantiation, these are never refitted.
    """
    if model_kind in self._prefitted_nuisance_models:
        return self
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._nuisance_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

fit_treatment

fit_treatment(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit the treatment model of a MetaLearner.

y represents the objective of the given treatment model, not necessarily the outcome of the experiment.

Source code in metalearners/metalearner.py

def fit_treatment(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit the treatment model of a MetaLearner.

    `y` represents the objective of the given treatment model, not necessarily the outcome of the experiment.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._treatment_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

nuisance_model_specifications classmethod

nuisance_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all first-stage models.

Source code in metalearners/rlearner.py

@classmethod
def nuisance_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return {
        PROPENSITY_MODEL: _ModelSpecifications(
            cardinality=get_one,
            predict_method=get_predict_proba,
        ),
        OUTCOME_MODEL: _ModelSpecifications(
            cardinality=get_one,
            predict_method=MetaLearner._outcome_predict_method,
        ),
    }

predict

predict(
    X, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Estimate the CATE.

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants} - 1, 1)\) if the outcome is a scalar, i.e. in case of a regression problem.

• \((n_{obs}, n_{variants} - 1, n_{classes})\) if the outcome is a class, i.e. in case of a classification problem.

In the case of multiple treatment variants, the second dimension represents the CATE of the corresponding variant vs the control (variant 0).

Source code in metalearners/rlearner.py

def predict(
    self,
    X,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    n_outputs = 2 if self.is_classification else 1
    tau_hat = np.zeros((safe_len(X), self.n_variants - 1, n_outputs))

    if is_oos:

        for treatment_variant in range(1, self.n_variants):
            variant_estimates = self.predict_treatment(
                X,
                is_oos=is_oos,
                oos_method=oos_method,
                model_kind=TREATMENT_MODEL,
                model_ord=treatment_variant - 1,
            )
            if self.is_classification:
                # This is to be consistent with other MetaLearners (e.g. S and T) that automatically
                # work with multiclass outcomes and return the CATE estimate for each class. As the R-Learner only
                # works with binary classes (the pseudo outcome formula does not make sense with
                # multiple classes unless some adaptation is done) we can manually infer the
                # CATE estimate for the complementary class  -- returning a matrix of shape (N, 2).
                variant_estimates = np.stack(
                    [-variant_estimates, variant_estimates], axis=-1
                )
            variant_estimates = variant_estimates.reshape(safe_len(X), n_outputs)
            tau_hat[:, treatment_variant - 1, :] = variant_estimates

        return tau_hat

    for treatment_variant in range(1, self.n_variants):
        variant_indices = self._variants_indices[treatment_variant - 1]

        variant_estimates = self.predict_treatment(
            index_matrix(X, variant_indices),
            is_oos=False,
            model_kind=TREATMENT_MODEL,
            model_ord=treatment_variant - 1,
        )
        if sum(~variant_indices) > 0:
            non_variant_estimates = self.predict_treatment(
                index_matrix(X, ~variant_indices),
                is_oos=True,
                oos_method=oos_method,
                model_kind=TREATMENT_MODEL,
                model_ord=treatment_variant - 1,
            )
        if self.is_classification:
            # This is to be consistent with other MetaLearners (e.g. S and T) that automatically
            # work with multiclass outcomes and return the CATE estimate for each class. As the R-Learner only
            # works with binary classes (the pseudo outcome formula does not make sense with
            # multiple classes unless some adaptation is done) we can manually infer the
            # CATE estimate for the complementary class  -- returning a matrix of shape (N, 2).
            variant_estimates = np.stack(
                [-variant_estimates, variant_estimates], axis=-1
            )
            if sum(~variant_indices) > 0:
                non_variant_estimates = np.stack(
                    [-non_variant_estimates, non_variant_estimates], axis=-1
                )
        variant_estimates = variant_estimates.reshape(
            (sum(variant_indices), n_outputs)
        )
        if sum(~variant_indices) > 0:
            non_variant_estimates = non_variant_estimates.reshape(
                (sum(~variant_indices), n_outputs)
            )
            tau_hat[~variant_indices, treatment_variant - 1] = non_variant_estimates

        tau_hat[variant_indices, treatment_variant - 1] = variant_estimates
    return tau_hat

predict_conditional_average_outcomes

predict_conditional_average_outcomes(
    X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Predict the vectors of conditional average outcomes.

These are defined as \(\mathbb{E}[Y_i(w) | X]\) for each treatment variant \(w\).

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants}, 1)\) if the outcome is a scalar, i.e., in case of a regression problem.

• \((n_{obs}, n_{variants}, n_{classes})\) if the outcome is a class, i.e., in case of a classification problem.

The conditional average outcomes are estimated as follows:

• \(Y_i(0) = \hat{\mu}(X_i) - \sum_{k=1}^{K} \hat{e}_k(X_i) \hat{\tau_k}(X_i)\)
• \(Y_i(k) = Y_i(0) + \hat{\tau_k}(X_i)\) for \(k \in \{1, \dots, K\}\)

where \(K\) is the number of treatment variants.

Source code in metalearners/rlearner.py

def predict_conditional_average_outcomes(
    self, X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray:
    r"""The conditional average outcomes are estimated as follows:

    * $Y_i(0) = \hat{\mu}(X_i) - \sum_{k=1}^{K} \hat{e}_k(X_i) \hat{\tau_k}(X_i)$
    * $Y_i(k) = Y_i(0) + \hat{\tau_k}(X_i)$ for $k \in \{1, \dots, K\}$

    where $K$ is the number of treatment variants.
    """
    n_obs = safe_len(X)

    cate_estimates = self.predict(
        X=X,
        is_oos=is_oos,
        oos_method=oos_method,
    )
    propensity_estimates = self.predict_nuisance(
        X=X,
        model_kind=PROPENSITY_MODEL,
        model_ord=0,
        is_oos=is_oos,
        oos_method=oos_method,
    )
    outcome_estimates = self.predict_nuisance(
        X=X,
        model_kind=OUTCOME_MODEL,
        model_ord=0,
        is_oos=is_oos,
        oos_method=oos_method,
    )

    conditional_average_outcomes_list = []

    control_outcomes = outcome_estimates

    # TODO: Consider whether the readability vs efficiency trade-off should be dealt with differently here.
    # One could use matrix/tensor operations instead.
    for treatment_variant in range(1, self.n_variants):
        if (n_outputs := cate_estimates.shape[2]) > 1:
            for outcome_channel in range(0, n_outputs):
                control_outcomes[:, outcome_channel] -= (
                    propensity_estimates[:, treatment_variant]
                    * cate_estimates[:, treatment_variant - 1, outcome_channel]
                )
        else:
            control_outcomes -= (
                propensity_estimates[:, treatment_variant]
                * cate_estimates[:, treatment_variant - 1, 0]
            )

    conditional_average_outcomes_list.append(control_outcomes)

    for treatment_variant in range(1, self.n_variants):
        conditional_average_outcomes_list.append(
            control_outcomes
            + np.reshape(
                cate_estimates[:, treatment_variant - 1, :],
                (control_outcomes.shape),
            )
        )

    return np.stack(conditional_average_outcomes_list, axis=1).reshape(
        n_obs, self.n_variants, -1
    )

predict_nuisance

predict_nuisance(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given nuisance model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_nuisance(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given nuisance model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    predict_method_name = self.nuisance_model_specifications()[model_kind][
        "predict_method"
    ](self)
    predict_method = getattr(
        self._nuisance_models[model_kind][model_ord], predict_method_name
    )
    return predict_method(X_filtered, is_oos, oos_method)

predict_treatment

predict_treatment(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given treatment model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_treatment(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given treatment model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    return self._treatment_models[model_kind][model_ord].predict(
        X_filtered, is_oos, oos_method
    )

shap_values

shap_values(
    X: Matrix,
    shap_explainer_factory: type[Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]

Calculates the shap values for each treatment group.

If explainer is None a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters cate_estimates, cate_model_factory and cate_model_params are ignored.

The parameter shap_explainer_factory can be used to specify the type of shap explainer, for the different options see here.

The returned list contains the shap values for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def shap_values(
    self,
    X: Matrix,
    shap_explainer_factory: type[shap.Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]:
    """Calculates the shap values for each treatment group.

    If `explainer` is `None` a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    The parameter `shap_explainer_factory` can be used to specify the type of shap
    explainer, for the different options see
    [here](https://shap.readthedocs.io/en/latest/api.html#explainers).

    The returned list contains the shap values for each treatment variant in ascending
    order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=None if cate_estimates is None else X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.shap_values(
        X=X,
        shap_explainer_factory=shap_explainer_factory,
        shap_explainer_params=shap_explainer_params,
    )

treatment_model_specifications classmethod

treatment_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all second-stage models.

Source code in metalearners/rlearner.py

@classmethod
def treatment_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return {
        TREATMENT_MODEL: _ModelSpecifications(
            cardinality=MetaLearner._get_n_variants_minus_one,
            predict_method=get_predict,
        )
    }

SLearner

SLearner(
    is_classification: bool,
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel]
    | None = None,
    nuisance_model_params: Params
    | dict[str, Params]
    | None = None,
    treatment_model_params: Params
    | dict[str, Params]
    | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[
        str, list[CrossFitEstimator]
    ]
    | None = None,
    fitted_propensity_model: CrossFitEstimator
    | None = None,
    feature_set: Features
    | dict[str, Features]
    | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
)

Bases: MetaLearner

S-Learner for CATE estimation as described by Kuenzel et al (2019).

Source code in metalearners/slearner.py

def __init__(
    self,
    is_classification: bool,
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel] | None = None,
    nuisance_model_params: Params | dict[str, Params] | None = None,
    treatment_model_params: Params | dict[str, Params] | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[str, list[CrossFitEstimator]] | None = None,
    fitted_propensity_model: CrossFitEstimator | None = None,
    feature_set: Features | dict[str, Features] | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
):
    if feature_set is not None:
        # For SLearner it does not make sense to allow feature set as we only have one model
        # and having it would bring problems when using fit_nuisance and predict_nuisance
        # as we need to add the treatment column.
        warnings.warn(
            "Base-model specific feature_sets were provided to S-Learner. "
            "These will be ignored and all available features will be used instead."
        )
    super().__init__(
        is_classification=is_classification,
        n_variants=n_variants,
        nuisance_model_factory=nuisance_model_factory,
        treatment_model_factory=treatment_model_factory,
        propensity_model_factory=propensity_model_factory,
        nuisance_model_params=nuisance_model_params,
        treatment_model_params=treatment_model_params,
        propensity_model_params=propensity_model_params,
        fitted_nuisance_models=fitted_nuisance_models,
        fitted_propensity_model=fitted_propensity_model,
        feature_set=None,
        n_folds=n_folds,
        random_state=random_state,
    )

init_args property

init_args: dict[str, Any]

Create initialization parameters for a new MetaLearner.

Importantly, this does not copy further internal state, such as the weights or parameters of trained base models.

evaluate

evaluate(
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]

Evaluate the MetaLearner.

The keys in scoring which are not a name of a model contained in the MetaLearner will be ignored, for information about this names check nuisance_model_specifications and treatment_model_specifications. The values must be a list of:

• string representing a sklearn scoring method. Check here for the possible values.
• Callable with signature scorer(estimator, X, y_true, **kwargs). We recommend using sklearn.metrics.make_scorer to create such a Callable.

If some model name is not present in the keys of scoring then the default used metrics will be neg_log_loss if it is a classifier and neg_root_mean_squared_error if it is a regressor.

The returned dictionary keys have the following structure:

• For nuisance models:

  • If the cardinality is one: f"{model_kind}_{scorer}"
  • If there is one model for each treatment variant (including control): f"{model_kind}_{treatment_variant}_{scorer}"

• For treatment models: f"{model_kind}_{treatment_variant}_vs_0_{scorer}"

Where scorer is the name of the scorer if it is a string and "custom_scorer_{idx}" if it is a callable where idx is the index in the scorers list.

Source code in metalearners/slearner.py

def evaluate(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]:
    safe_scoring = self._scoring(scoring)

    X_with_w = _append_treatment_to_covariates(
        X, w, self._supports_categoricals, self.n_variants
    )
    return _evaluate_model_kind(
        cfes=self._nuisance_models[_BASE_MODEL],
        Xs=[X_with_w],
        ys=[y],
        scorers=safe_scoring[_BASE_MODEL],
        model_kind=_BASE_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=False,
        feature_set=self.feature_set[_BASE_MODEL],
    )

explainer

explainer(
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer

Create an Explainer

which can be used in feature_importances.

This function can be used in two distinct manners based on the provided parameters:

• When parameters X, cate_estimates, and cate_model_factory are all set to None, the function creates an Explainer using the pre-existing treatment models. If these models do not exist, however, it triggers a ValueError.
• On the contrary, if X, cate_estimates, and cate_model_factory are not None, the function initiates an instance of the Explainer class using these parameters. This instance then fits new models for each treatment variant, and these models are employed to calculate the importance of features.

Source code in metalearners/metalearner.py

def explainer(
    self,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer:
    r"""Create an [`Explainer`][metalearners.explainer.Explainer]

    which can be used in [`feature_importances`][metalearners.metalearner.MetaLearner.feature_importances].

    This function can be used in two distinct manners based on the provided parameters:

    *   When parameters `X`, `cate_estimates`, and `cate_model_factory` are all
        set to `None`, the function creates an [`Explainer`][metalearners.explainer.Explainer]
        using the pre-existing treatment models. If these models do not exist, however,
        it triggers a `ValueError`.
    *   On the contrary, if `X`, `cate_estimates`, and `cate_model_factory` are
        not `None`, the function initiates an instance of the [`Explainer`][metalearners.explainer.Explainer]
        class using these parameters. This instance then fits new models for each
        treatment variant, and these models are employed to calculate the importance
        of features.
    """
    if X is None and cate_estimates is None and cate_model_factory is None:
        try:
            cate_cf_models = self._treatment_models[TREATMENT_MODEL]
            cate_models = [cf._overall_estimator for cf in cate_cf_models]
        except KeyError:
            raise ValueError(
                "The metalearner does not have treatment models; hence X, cate_estimates, "
                "and cate_model_factory need to be defined."
            )
        return Explainer(cate_models=cate_models)  # type: ignore
    elif (
        X is not None
        and cate_estimates is not None
        and cate_model_factory is not None
    ):
        return Explainer.from_estimates(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    else:
        raise ValueError(
            "Either all of [X, cate_estimates, cate_model_factory] are None or all "
            "of them must be defined."
        )

feature_importances

feature_importances(
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances

Calculates the feature importance for each treatment group.

If explainer is None, a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters X, cate_estimates, cate_model_factory and cate_model_params are ignored.

If normalization = True, for each treatment variant the feature importances are normalized so that they sum to 1.

feature_names is optional but in the case it’s not passed the names of the features will default to f"Feature {i}" where i is the corresponding feature index.

The returned list contains the feature importances for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def feature_importances(
    self,
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances:
    r"""Calculates the feature importance for each treatment group.

    If `explainer` is `None`, a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `X`, `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    If `normalization = True`, for each treatment variant the feature importances
    are normalized so that they sum to 1.

    `feature_names` is optional but in the case it's not passed the names of the
    features will default to `f"Feature {i}"` where `i` is the corresponding
    feature index.

    The returned list contains the feature importances for each treatment variant in
    ascending order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.feature_importances(
        normalize=normalize, feature_names=feature_names, sort_values=sort_values
    )

fit

fit(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

n_jobs_cross_fitting will be used at the cross-fitting level and n_jobs_base_learners will be used at the stage level. None means 1 unless in a joblib.parallel_backend context. -1 means using all processors. For more information about parallelism check parallelism.

fit_params is an optional dict to be forwarded to base estimator fit calls. It supports two usages patterns:

fit_params={"parameter_of_interest": value_of_interest}

fit_params={
    "nuisance": {
        "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
        "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
    },
    "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
}

In the former approach, the parameter and value of interest are passed to all base models. In the the latter approach, the explicitly qualified parameter-value pairs are passed to respective base models and no fitting parameters are passed to base models not explicitly listed. Note that in this pattern, propensity models are considered a nuisance model.

synchronize_cross_fitting indicates whether the learning of different base models should use exactly the same data splits where possible. Note that if there are several models to be synchronized which are classifiers, these cannot be split via stratification.

Source code in metalearners/metalearner.py

def fit(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    """Fit all models of the MetaLearner.

    If pre-fitted models were passed at instantiation, these are never refitted.

    `n_jobs_cross_fitting` will be used at the cross-fitting level and
    `n_jobs_base_learners` will be used at the stage level. `None` means 1 unless in a
    [`joblib.parallel_backend`](https://joblib.readthedocs.io/en/latest/generated/joblib.parallel_backend.html#joblib.parallel_backend)
    context. `-1` means using all processors.
    For more information about parallelism check [parallelism](parallelism.md).


    `fit_params` is an optional `dict` to be forwarded to base estimator `fit` calls. It supports
    two usages patterns:

    ```python
    fit_params={"parameter_of_interest": value_of_interest}
    ```

    ```python
    fit_params={
        "nuisance": {
            "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
            "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
        },
        "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
    }
    ```

    In the former approach, the parameter and value of interest are passed to all base models. In the
    the latter approach, the explicitly qualified parameter-value pairs are passed to respective base
    models and no fitting parameters are passed to base models not explicitly listed. Note that in this
    pattern, propensity models are considered a nuisance model.

    `synchronize_cross_fitting` indicates whether the learning of different base models should use exactly
    the same data splits where possible. Note that if there are several models to be synchronized which are
    classifiers, these cannot be split via stratification.
    """
    self.fit_all_nuisance(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    self.fit_all_treatment(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    return self

fit_all_nuisance

fit_all_nuisance(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all nuisance models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the nuisance models, and in the case they should also be used by the treatment models, these should also be passed in the following call to fit_all_treatment.

This method, combined with fit_all_treatment, facilitates the segmentation of the metalearner fitting process into two distinct parts. This division allows for interventions between the two stages, such as performing feature selection for the treatment models or conducting hyperparameter optimization within the nuisance models.

Source code in metalearners/slearner.py

def fit_all_nuisance(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    self._validate_treatment(w)
    self._validate_outcome(y, w)
    self._fitted_treatments = adapt_treatment_dtypes(w)

    mock_model = self.nuisance_model_factory[_BASE_MODEL](
        **self.nuisance_model_params[_BASE_MODEL]
    )
    self._supports_categoricals = supports_categoricals(mock_model)
    X_with_w = _append_treatment_to_covariates(
        X, w, self._supports_categoricals, self.n_variants
    )

    qualified_fit_params = self._qualified_fit_params(fit_params)

    self.fit_nuisance(
        X=X_with_w,
        y=y,
        model_kind=_BASE_MODEL,
        model_ord=0,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=qualified_fit_params[NUISANCE][_BASE_MODEL],
    )
    return self

fit_all_treatment

fit_all_treatment(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all treatment models of the MetaLearner.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the treatment models, as the nuisance models should already be fitted.

Source code in metalearners/slearner.py

def fit_all_treatment(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    return self

fit_nuisance

fit_nuisance(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit a given nuisance model of a MetaLearner.

y represents the objective of the given nuisance model, not necessarily the outcome of the experiment. If pre-fitted models were passed at instantiation, these are never refitted.

Source code in metalearners/metalearner.py

def fit_nuisance(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit a given nuisance model of a MetaLearner.

    `y` represents the objective of the given nuisance model, not necessarily the outcome of the experiment.
    If pre-fitted models were passed at instantiation, these are never refitted.
    """
    if model_kind in self._prefitted_nuisance_models:
        return self
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._nuisance_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

fit_treatment

fit_treatment(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit the treatment model of a MetaLearner.

y represents the objective of the given treatment model, not necessarily the outcome of the experiment.

Source code in metalearners/metalearner.py

def fit_treatment(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit the treatment model of a MetaLearner.

    `y` represents the objective of the given treatment model, not necessarily the outcome of the experiment.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._treatment_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

nuisance_model_specifications classmethod

nuisance_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all first-stage models.

Source code in metalearners/slearner.py

@classmethod
def nuisance_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return {
        _BASE_MODEL: _ModelSpecifications(
            cardinality=get_one,
            predict_method=MetaLearner._outcome_predict_method,
        )
    }

predict

predict(
    X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Estimate the CATE.

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants} - 1, 1)\) if the outcome is a scalar, i.e. in case of a regression problem.

• \((n_{obs}, n_{variants} - 1, n_{classes})\) if the outcome is a class, i.e. in case of a classification problem.

In the case of multiple treatment variants, the second dimension represents the CATE of the corresponding variant vs the control (variant 0).

Source code in metalearners/slearner.py

def predict(
    self,
    X: Matrix,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    conditional_average_outcomes = self.predict_conditional_average_outcomes(
        X=X, is_oos=is_oos, oos_method=oos_method
    )

    return conditional_average_outcomes[:, 1:] - (
        conditional_average_outcomes[:, [0]]
    )

predict_conditional_average_outcomes

predict_conditional_average_outcomes(
    X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Predict the vectors of conditional average outcomes.

These are defined as \(\mathbb{E}[Y_i(w) | X]\) for each treatment variant \(w\).

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants}, 1)\) if the outcome is a scalar, i.e., in case of a regression problem.

• \((n_{obs}, n_{variants}, n_{classes})\) if the outcome is a class, i.e., in case of a classification problem.

Source code in metalearners/slearner.py

def predict_conditional_average_outcomes(
    self, X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray:
    n_obs = safe_len(X)
    conditional_average_outcomes_list = []

    for treatment_variant in range(self.n_variants):
        w = np.array([treatment_variant] * n_obs)
        X_with_w = _append_treatment_to_covariates(
            X, w, self._supports_categoricals, self.n_variants
        )
        variant_predictions = self.predict_nuisance(
            X=X_with_w,
            model_kind=_BASE_MODEL,
            model_ord=0,
            is_oos=is_oos,
            oos_method=oos_method,
        )

        conditional_average_outcomes_list.append(variant_predictions)

    return np.stack(conditional_average_outcomes_list, axis=1).reshape(
        n_obs, self.n_variants, -1
    )

predict_nuisance

predict_nuisance(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given nuisance model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_nuisance(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given nuisance model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    predict_method_name = self.nuisance_model_specifications()[model_kind][
        "predict_method"
    ](self)
    predict_method = getattr(
        self._nuisance_models[model_kind][model_ord], predict_method_name
    )
    return predict_method(X_filtered, is_oos, oos_method)

predict_treatment

predict_treatment(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given treatment model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_treatment(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given treatment model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    return self._treatment_models[model_kind][model_ord].predict(
        X_filtered, is_oos, oos_method
    )

shap_values

shap_values(
    X: Matrix,
    shap_explainer_factory: type[Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]

Calculates the shap values for each treatment group.

If explainer is None a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters cate_estimates, cate_model_factory and cate_model_params are ignored.

The parameter shap_explainer_factory can be used to specify the type of shap explainer, for the different options see here.

The returned list contains the shap values for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def shap_values(
    self,
    X: Matrix,
    shap_explainer_factory: type[shap.Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]:
    """Calculates the shap values for each treatment group.

    If `explainer` is `None` a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    The parameter `shap_explainer_factory` can be used to specify the type of shap
    explainer, for the different options see
    [here](https://shap.readthedocs.io/en/latest/api.html#explainers).

    The returned list contains the shap values for each treatment variant in ascending
    order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=None if cate_estimates is None else X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.shap_values(
        X=X,
        shap_explainer_factory=shap_explainer_factory,
        shap_explainer_params=shap_explainer_params,
    )

treatment_model_specifications classmethod

treatment_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all second-stage models.

Source code in metalearners/slearner.py

@classmethod
def treatment_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return dict()

TLearner

TLearner(
    is_classification: bool,
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel]
    | None = None,
    nuisance_model_params: Params
    | dict[str, Params]
    | None = None,
    treatment_model_params: Params
    | dict[str, Params]
    | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[
        str, list[CrossFitEstimator]
    ]
    | None = None,
    fitted_propensity_model: CrossFitEstimator
    | None = None,
    feature_set: Features
    | dict[str, Features]
    | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
)

Bases: _ConditionalAverageOutcomeMetaLearner

T-Learner for CATE estimation as described by Kuenzel et al (2019).

Source code in metalearners/metalearner.py

def __init__(
    self,
    is_classification: bool,
    # TODO: Consider whether we can make this not a state of the MetaLearner
    # but rather just a parameter of a predict call.
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel] | None = None,
    nuisance_model_params: Params | dict[str, Params] | None = None,
    treatment_model_params: Params | dict[str, Params] | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[str, list[CrossFitEstimator]] | None = None,
    fitted_propensity_model: CrossFitEstimator | None = None,
    feature_set: Features | dict[str, Features] | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
):
    super().__init__(
        nuisance_model_factory=nuisance_model_factory,
        is_classification=is_classification,
        n_variants=n_variants,
        treatment_model_factory=treatment_model_factory,
        propensity_model_factory=propensity_model_factory,
        nuisance_model_params=nuisance_model_params,
        treatment_model_params=treatment_model_params,
        propensity_model_params=propensity_model_params,
        fitted_nuisance_models=fitted_nuisance_models,
        fitted_propensity_model=fitted_propensity_model,
        feature_set=feature_set,
        n_folds=n_folds,
        random_state=random_state,
    )
    self._treatment_variants_mask: list[np.ndarray] | None = None

init_args property

init_args: dict[str, Any]

Create initialization parameters for a new MetaLearner.

Importantly, this does not copy further internal state, such as the weights or parameters of trained base models.

evaluate

evaluate(
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]

Evaluate the MetaLearner.

The keys in scoring which are not a name of a model contained in the MetaLearner will be ignored, for information about this names check nuisance_model_specifications and treatment_model_specifications. The values must be a list of:

• string representing a sklearn scoring method. Check here for the possible values.
• Callable with signature scorer(estimator, X, y_true, **kwargs). We recommend using sklearn.metrics.make_scorer to create such a Callable.

If some model name is not present in the keys of scoring then the default used metrics will be neg_log_loss if it is a classifier and neg_root_mean_squared_error if it is a regressor.

The returned dictionary keys have the following structure:

• For nuisance models:

  • If the cardinality is one: f"{model_kind}_{scorer}"
  • If there is one model for each treatment variant (including control): f"{model_kind}_{treatment_variant}_{scorer}"

• For treatment models: f"{model_kind}_{treatment_variant}_vs_0_{scorer}"

Where scorer is the name of the scorer if it is a string and "custom_scorer_{idx}" if it is a callable where idx is the index in the scorers list.

Source code in metalearners/tlearner.py

def evaluate(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]:
    safe_scoring = self._scoring(scoring)

    return _evaluate_model_kind(
        cfes=self._nuisance_models[VARIANT_OUTCOME_MODEL],
        Xs=[index_matrix(X, w == tv) for tv in range(self.n_variants)],
        ys=[index_vector(y, w == tv) for tv in range(self.n_variants)],
        scorers=safe_scoring[VARIANT_OUTCOME_MODEL],
        model_kind=VARIANT_OUTCOME_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=False,
        feature_set=self.feature_set[VARIANT_OUTCOME_MODEL],
    )

explainer

explainer(
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer

Create an Explainer

which can be used in feature_importances.

This function can be used in two distinct manners based on the provided parameters:

• When parameters X, cate_estimates, and cate_model_factory are all set to None, the function creates an Explainer using the pre-existing treatment models. If these models do not exist, however, it triggers a ValueError.
• On the contrary, if X, cate_estimates, and cate_model_factory are not None, the function initiates an instance of the Explainer class using these parameters. This instance then fits new models for each treatment variant, and these models are employed to calculate the importance of features.

Source code in metalearners/metalearner.py

def explainer(
    self,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer:
    r"""Create an [`Explainer`][metalearners.explainer.Explainer]

    which can be used in [`feature_importances`][metalearners.metalearner.MetaLearner.feature_importances].

    This function can be used in two distinct manners based on the provided parameters:

    *   When parameters `X`, `cate_estimates`, and `cate_model_factory` are all
        set to `None`, the function creates an [`Explainer`][metalearners.explainer.Explainer]
        using the pre-existing treatment models. If these models do not exist, however,
        it triggers a `ValueError`.
    *   On the contrary, if `X`, `cate_estimates`, and `cate_model_factory` are
        not `None`, the function initiates an instance of the [`Explainer`][metalearners.explainer.Explainer]
        class using these parameters. This instance then fits new models for each
        treatment variant, and these models are employed to calculate the importance
        of features.
    """
    if X is None and cate_estimates is None and cate_model_factory is None:
        try:
            cate_cf_models = self._treatment_models[TREATMENT_MODEL]
            cate_models = [cf._overall_estimator for cf in cate_cf_models]
        except KeyError:
            raise ValueError(
                "The metalearner does not have treatment models; hence X, cate_estimates, "
                "and cate_model_factory need to be defined."
            )
        return Explainer(cate_models=cate_models)  # type: ignore
    elif (
        X is not None
        and cate_estimates is not None
        and cate_model_factory is not None
    ):
        return Explainer.from_estimates(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    else:
        raise ValueError(
            "Either all of [X, cate_estimates, cate_model_factory] are None or all "
            "of them must be defined."
        )

feature_importances

feature_importances(
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances

Calculates the feature importance for each treatment group.

If explainer is None, a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters X, cate_estimates, cate_model_factory and cate_model_params are ignored.

If normalization = True, for each treatment variant the feature importances are normalized so that they sum to 1.

feature_names is optional but in the case it’s not passed the names of the features will default to f"Feature {i}" where i is the corresponding feature index.

The returned list contains the feature importances for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def feature_importances(
    self,
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances:
    r"""Calculates the feature importance for each treatment group.

    If `explainer` is `None`, a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `X`, `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    If `normalization = True`, for each treatment variant the feature importances
    are normalized so that they sum to 1.

    `feature_names` is optional but in the case it's not passed the names of the
    features will default to `f"Feature {i}"` where `i` is the corresponding
    feature index.

    The returned list contains the feature importances for each treatment variant in
    ascending order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.feature_importances(
        normalize=normalize, feature_names=feature_names, sort_values=sort_values
    )

fit

fit(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

n_jobs_cross_fitting will be used at the cross-fitting level and n_jobs_base_learners will be used at the stage level. None means 1 unless in a joblib.parallel_backend context. -1 means using all processors. For more information about parallelism check parallelism.

fit_params is an optional dict to be forwarded to base estimator fit calls. It supports two usages patterns:

fit_params={"parameter_of_interest": value_of_interest}

fit_params={
    "nuisance": {
        "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
        "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
    },
    "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
}

In the former approach, the parameter and value of interest are passed to all base models. In the the latter approach, the explicitly qualified parameter-value pairs are passed to respective base models and no fitting parameters are passed to base models not explicitly listed. Note that in this pattern, propensity models are considered a nuisance model.

synchronize_cross_fitting indicates whether the learning of different base models should use exactly the same data splits where possible. Note that if there are several models to be synchronized which are classifiers, these cannot be split via stratification.

Source code in metalearners/metalearner.py

def fit(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    """Fit all models of the MetaLearner.

    If pre-fitted models were passed at instantiation, these are never refitted.

    `n_jobs_cross_fitting` will be used at the cross-fitting level and
    `n_jobs_base_learners` will be used at the stage level. `None` means 1 unless in a
    [`joblib.parallel_backend`](https://joblib.readthedocs.io/en/latest/generated/joblib.parallel_backend.html#joblib.parallel_backend)
    context. `-1` means using all processors.
    For more information about parallelism check [parallelism](parallelism.md).


    `fit_params` is an optional `dict` to be forwarded to base estimator `fit` calls. It supports
    two usages patterns:

    ```python
    fit_params={"parameter_of_interest": value_of_interest}
    ```

    ```python
    fit_params={
        "nuisance": {
            "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
            "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
        },
        "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
    }
    ```

    In the former approach, the parameter and value of interest are passed to all base models. In the
    the latter approach, the explicitly qualified parameter-value pairs are passed to respective base
    models and no fitting parameters are passed to base models not explicitly listed. Note that in this
    pattern, propensity models are considered a nuisance model.

    `synchronize_cross_fitting` indicates whether the learning of different base models should use exactly
    the same data splits where possible. Note that if there are several models to be synchronized which are
    classifiers, these cannot be split via stratification.
    """
    self.fit_all_nuisance(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    self.fit_all_treatment(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    return self

fit_all_nuisance

fit_all_nuisance(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all nuisance models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the nuisance models, and in the case they should also be used by the treatment models, these should also be passed in the following call to fit_all_treatment.

This method, combined with fit_all_treatment, facilitates the segmentation of the metalearner fitting process into two distinct parts. This division allows for interventions between the two stages, such as performing feature selection for the treatment models or conducting hyperparameter optimization within the nuisance models.

Source code in metalearners/tlearner.py

def fit_all_nuisance(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    self._validate_treatment(w)
    self._validate_outcome(y, w)

    self._treatment_variants_mask = []

    for v in range(self.n_variants):
        self._treatment_variants_mask.append(w == v)

    qualified_fit_params = self._qualified_fit_params(fit_params)

    nuisance_jobs: list[_ParallelJoblibSpecification | None] = []
    for treatment_variant in range(self.n_variants):
        mask = self._treatment_variants_mask[treatment_variant]
        X_variant = index_matrix(X, mask)
        y_variant = index_vector(y, mask)
        nuisance_jobs.append(
            self._nuisance_joblib_specifications(
                X=X_variant,
                y=y_variant,
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=treatment_variant,
                n_jobs_cross_fitting=n_jobs_cross_fitting,
                fit_params=qualified_fit_params[NUISANCE][VARIANT_OUTCOME_MODEL],
            )
        )

    parallel = Parallel(n_jobs=n_jobs_base_learners)
    results = parallel(
        delayed(_fit_cross_fit_estimator_joblib)(job)
        for job in nuisance_jobs
        if job is not None
    )
    self._assign_joblib_nuisance_results(results)
    return self

fit_all_treatment

fit_all_treatment(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all treatment models of the MetaLearner.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the treatment models, as the nuisance models should already be fitted.

Source code in metalearners/tlearner.py

def fit_all_treatment(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    return self

fit_nuisance

fit_nuisance(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit a given nuisance model of a MetaLearner.

y represents the objective of the given nuisance model, not necessarily the outcome of the experiment. If pre-fitted models were passed at instantiation, these are never refitted.

Source code in metalearners/metalearner.py

def fit_nuisance(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit a given nuisance model of a MetaLearner.

    `y` represents the objective of the given nuisance model, not necessarily the outcome of the experiment.
    If pre-fitted models were passed at instantiation, these are never refitted.
    """
    if model_kind in self._prefitted_nuisance_models:
        return self
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._nuisance_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

fit_treatment

fit_treatment(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit the treatment model of a MetaLearner.

y represents the objective of the given treatment model, not necessarily the outcome of the experiment.

Source code in metalearners/metalearner.py

def fit_treatment(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit the treatment model of a MetaLearner.

    `y` represents the objective of the given treatment model, not necessarily the outcome of the experiment.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._treatment_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

nuisance_model_specifications classmethod

nuisance_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all first-stage models.

Source code in metalearners/tlearner.py

@classmethod
def nuisance_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return {
        VARIANT_OUTCOME_MODEL: _ModelSpecifications(
            cardinality=MetaLearner._get_n_variants,
            predict_method=MetaLearner._outcome_predict_method,
        ),
    }

predict

predict(
    X, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Estimate the CATE.

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants} - 1, 1)\) if the outcome is a scalar, i.e. in case of a regression problem.

• \((n_{obs}, n_{variants} - 1, n_{classes})\) if the outcome is a class, i.e. in case of a classification problem.

In the case of multiple treatment variants, the second dimension represents the CATE of the corresponding variant vs the control (variant 0).

Source code in metalearners/tlearner.py

def predict(
    self,
    X,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    conditional_average_outcomes = self.predict_conditional_average_outcomes(
        X=X, is_oos=is_oos, oos_method=oos_method
    )

    return conditional_average_outcomes[:, 1:] - (
        conditional_average_outcomes[:, [0]]
    )

predict_conditional_average_outcomes

predict_conditional_average_outcomes(
    X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Predict the vectors of conditional average outcomes.

These are defined as \(\mathbb{E}[Y_i(w) | X]\) for each treatment variant \(w\).

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants}, 1)\) if the outcome is a scalar, i.e., in case of a regression problem.

• \((n_{obs}, n_{variants}, n_{classes})\) if the outcome is a class, i.e., in case of a classification problem.

Source code in metalearners/metalearner.py

def predict_conditional_average_outcomes(
    self, X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray:
    if self._treatment_variants_mask is None:
        raise ValueError(
            "The metalearner needs to be fitted before predicting."
            "In particular, the MetaLearner's attribute _treatment_variant_mask, "
            "typically set during fitting, is None."
        )
    # TODO: Consider multiprocessing
    n_obs = safe_len(X)
    nuisance_tensors = self._nuisance_tensors(n_obs)
    conditional_average_outcomes_list = nuisance_tensors[VARIANT_OUTCOME_MODEL]

    for tv in range(self.n_variants):
        if is_oos:
            conditional_average_outcomes_list[tv] = self.predict_nuisance(
                X=X,
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=True,
                oos_method=oos_method,
            )
        else:
            conditional_average_outcomes_list[tv][
                self._treatment_variants_mask[tv]
            ] = self.predict_nuisance(
                X=index_matrix(X, self._treatment_variants_mask[tv]),
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=False,
            )
            conditional_average_outcomes_list[tv][
                ~self._treatment_variants_mask[tv]
            ] = self.predict_nuisance(
                X=index_matrix(X, ~self._treatment_variants_mask[tv]),
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=True,
                oos_method=oos_method,
            )
    return np.stack(conditional_average_outcomes_list, axis=1).reshape(
        n_obs, self.n_variants, -1
    )

predict_nuisance

predict_nuisance(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given nuisance model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_nuisance(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given nuisance model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    predict_method_name = self.nuisance_model_specifications()[model_kind][
        "predict_method"
    ](self)
    predict_method = getattr(
        self._nuisance_models[model_kind][model_ord], predict_method_name
    )
    return predict_method(X_filtered, is_oos, oos_method)

predict_treatment

predict_treatment(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given treatment model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_treatment(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given treatment model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    return self._treatment_models[model_kind][model_ord].predict(
        X_filtered, is_oos, oos_method
    )

shap_values

shap_values(
    X: Matrix,
    shap_explainer_factory: type[Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]

Calculates the shap values for each treatment group.

If explainer is None a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters cate_estimates, cate_model_factory and cate_model_params are ignored.

The parameter shap_explainer_factory can be used to specify the type of shap explainer, for the different options see here.

The returned list contains the shap values for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def shap_values(
    self,
    X: Matrix,
    shap_explainer_factory: type[shap.Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]:
    """Calculates the shap values for each treatment group.

    If `explainer` is `None` a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    The parameter `shap_explainer_factory` can be used to specify the type of shap
    explainer, for the different options see
    [here](https://shap.readthedocs.io/en/latest/api.html#explainers).

    The returned list contains the shap values for each treatment variant in ascending
    order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=None if cate_estimates is None else X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.shap_values(
        X=X,
        shap_explainer_factory=shap_explainer_factory,
        shap_explainer_params=shap_explainer_params,
    )

treatment_model_specifications classmethod

treatment_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all second-stage models.

Source code in metalearners/tlearner.py

@classmethod
def treatment_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return dict()

XLearner

XLearner(
    is_classification: bool,
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel]
    | None = None,
    nuisance_model_params: Params
    | dict[str, Params]
    | None = None,
    treatment_model_params: Params
    | dict[str, Params]
    | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[
        str, list[CrossFitEstimator]
    ]
    | None = None,
    fitted_propensity_model: CrossFitEstimator
    | None = None,
    feature_set: Features
    | dict[str, Features]
    | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
)

Bases: _ConditionalAverageOutcomeMetaLearner

X-Learner for CATE estimation as described by Kuenzel et al (2019).

Importantly, the current X-Learner implementation only supports:

• binary classes in case of a classification outcome

Source code in metalearners/metalearner.py

def __init__(
    self,
    is_classification: bool,
    # TODO: Consider whether we can make this not a state of the MetaLearner
    # but rather just a parameter of a predict call.
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel] | None = None,
    nuisance_model_params: Params | dict[str, Params] | None = None,
    treatment_model_params: Params | dict[str, Params] | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[str, list[CrossFitEstimator]] | None = None,
    fitted_propensity_model: CrossFitEstimator | None = None,
    feature_set: Features | dict[str, Features] | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
):
    super().__init__(
        nuisance_model_factory=nuisance_model_factory,
        is_classification=is_classification,
        n_variants=n_variants,
        treatment_model_factory=treatment_model_factory,
        propensity_model_factory=propensity_model_factory,
        nuisance_model_params=nuisance_model_params,
        treatment_model_params=treatment_model_params,
        propensity_model_params=propensity_model_params,
        fitted_nuisance_models=fitted_nuisance_models,
        fitted_propensity_model=fitted_propensity_model,
        feature_set=feature_set,
        n_folds=n_folds,
        random_state=random_state,
    )
    self._treatment_variants_mask: list[np.ndarray] | None = None

init_args property

init_args: dict[str, Any]

Create initialization parameters for a new MetaLearner.

Importantly, this does not copy further internal state, such as the weights or parameters of trained base models.

evaluate

evaluate(
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]

Evaluate the MetaLearner.

The keys in scoring which are not a name of a model contained in the MetaLearner will be ignored, for information about this names check nuisance_model_specifications and treatment_model_specifications. The values must be a list of:

• string representing a sklearn scoring method. Check here for the possible values.
• Callable with signature scorer(estimator, X, y_true, **kwargs). We recommend using sklearn.metrics.make_scorer to create such a Callable.

If some model name is not present in the keys of scoring then the default used metrics will be neg_log_loss if it is a classifier and neg_root_mean_squared_error if it is a regressor.

The returned dictionary keys have the following structure:

• For nuisance models:

  • If the cardinality is one: f"{model_kind}_{scorer}"
  • If there is one model for each treatment variant (including control): f"{model_kind}_{treatment_variant}_{scorer}"

• For treatment models: f"{model_kind}_{treatment_variant}_vs_0_{scorer}"

Where scorer is the name of the scorer if it is a string and "custom_scorer_{idx}" if it is a callable where idx is the index in the scorers list.

Source code in metalearners/xlearner.py

def evaluate(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]:
    safe_scoring = self._scoring(scoring)

    variant_outcome_evaluation = _evaluate_model_kind(
        cfes=self._nuisance_models[VARIANT_OUTCOME_MODEL],
        Xs=[index_matrix(X, w == tv) for tv in range(self.n_variants)],
        ys=[index_vector(y, w == tv) for tv in range(self.n_variants)],
        scorers=safe_scoring[VARIANT_OUTCOME_MODEL],
        model_kind=VARIANT_OUTCOME_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=False,
        feature_set=self.feature_set[VARIANT_OUTCOME_MODEL],
    )

    propensity_evaluation = _evaluate_model_kind(
        cfes=self._nuisance_models[PROPENSITY_MODEL],
        Xs=[X],
        ys=[w],
        scorers=safe_scoring[PROPENSITY_MODEL],
        model_kind=PROPENSITY_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=False,
        feature_set=self.feature_set[PROPENSITY_MODEL],
    )

    conditional_average_outcome_estimates = (
        self.predict_conditional_average_outcomes(
            X=X,
            is_oos=is_oos,
            oos_method=oos_method,
        )
    )

    imputed_te_control: list[np.ndarray] = []
    imputed_te_treatment: list[np.ndarray] = []
    for treatment_variant in range(1, self.n_variants):
        tv_imputed_te_control, tv_imputed_te_treatment = self._pseudo_outcome(
            y, w, treatment_variant, conditional_average_outcome_estimates
        )
        imputed_te_control.append(tv_imputed_te_control)
        imputed_te_treatment.append(tv_imputed_te_treatment)

    te_treatment_evaluation = _evaluate_model_kind(
        self._treatment_models[TREATMENT_EFFECT_MODEL],
        Xs=[index_matrix(X, w == tv) for tv in range(1, self.n_variants)],
        ys=imputed_te_treatment,
        scorers=safe_scoring[TREATMENT_EFFECT_MODEL],
        model_kind=TREATMENT_EFFECT_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=True,
        feature_set=self.feature_set[TREATMENT_EFFECT_MODEL],
    )

    te_control_evaluation = _evaluate_model_kind(
        self._treatment_models[CONTROL_EFFECT_MODEL],
        Xs=[index_matrix(X, w == 0) for _ in range(1, self.n_variants)],
        ys=imputed_te_control,
        scorers=safe_scoring[CONTROL_EFFECT_MODEL],
        model_kind=CONTROL_EFFECT_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=True,
        feature_set=self.feature_set[CONTROL_EFFECT_MODEL],
    )

    return (
        variant_outcome_evaluation
        | propensity_evaluation
        | te_treatment_evaluation
        | te_control_evaluation
    )

explainer

explainer(
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer

Create an Explainer

which can be used in feature_importances.

This function can be used in two distinct manners based on the provided parameters:

• When parameters X, cate_estimates, and cate_model_factory are all set to None, the function creates an Explainer using the pre-existing treatment models. If these models do not exist, however, it triggers a ValueError.
• On the contrary, if X, cate_estimates, and cate_model_factory are not None, the function initiates an instance of the Explainer class using these parameters. This instance then fits new models for each treatment variant, and these models are employed to calculate the importance of features.

Source code in metalearners/metalearner.py

def explainer(
    self,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer:
    r"""Create an [`Explainer`][metalearners.explainer.Explainer]

    which can be used in [`feature_importances`][metalearners.metalearner.MetaLearner.feature_importances].

    This function can be used in two distinct manners based on the provided parameters:

    *   When parameters `X`, `cate_estimates`, and `cate_model_factory` are all
        set to `None`, the function creates an [`Explainer`][metalearners.explainer.Explainer]
        using the pre-existing treatment models. If these models do not exist, however,
        it triggers a `ValueError`.
    *   On the contrary, if `X`, `cate_estimates`, and `cate_model_factory` are
        not `None`, the function initiates an instance of the [`Explainer`][metalearners.explainer.Explainer]
        class using these parameters. This instance then fits new models for each
        treatment variant, and these models are employed to calculate the importance
        of features.
    """
    if X is None and cate_estimates is None and cate_model_factory is None:
        try:
            cate_cf_models = self._treatment_models[TREATMENT_MODEL]
            cate_models = [cf._overall_estimator for cf in cate_cf_models]
        except KeyError:
            raise ValueError(
                "The metalearner does not have treatment models; hence X, cate_estimates, "
                "and cate_model_factory need to be defined."
            )
        return Explainer(cate_models=cate_models)  # type: ignore
    elif (
        X is not None
        and cate_estimates is not None
        and cate_model_factory is not None
    ):
        return Explainer.from_estimates(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    else:
        raise ValueError(
            "Either all of [X, cate_estimates, cate_model_factory] are None or all "
            "of them must be defined."
        )

feature_importances

feature_importances(
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances

Calculates the feature importance for each treatment group.

If explainer is None, a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters X, cate_estimates, cate_model_factory and cate_model_params are ignored.

If normalization = True, for each treatment variant the feature importances are normalized so that they sum to 1.

feature_names is optional but in the case it’s not passed the names of the features will default to f"Feature {i}" where i is the corresponding feature index.

The returned list contains the feature importances for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def feature_importances(
    self,
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances:
    r"""Calculates the feature importance for each treatment group.

    If `explainer` is `None`, a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `X`, `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    If `normalization = True`, for each treatment variant the feature importances
    are normalized so that they sum to 1.

    `feature_names` is optional but in the case it's not passed the names of the
    features will default to `f"Feature {i}"` where `i` is the corresponding
    feature index.

    The returned list contains the feature importances for each treatment variant in
    ascending order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.feature_importances(
        normalize=normalize, feature_names=feature_names, sort_values=sort_values
    )

fit

fit(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

n_jobs_cross_fitting will be used at the cross-fitting level and n_jobs_base_learners will be used at the stage level. None means 1 unless in a joblib.parallel_backend context. -1 means using all processors. For more information about parallelism check parallelism.

fit_params is an optional dict to be forwarded to base estimator fit calls. It supports two usages patterns:

fit_params={"parameter_of_interest": value_of_interest}

fit_params={
    "nuisance": {
        "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
        "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
    },
    "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
}

In the former approach, the parameter and value of interest are passed to all base models. In the the latter approach, the explicitly qualified parameter-value pairs are passed to respective base models and no fitting parameters are passed to base models not explicitly listed. Note that in this pattern, propensity models are considered a nuisance model.

synchronize_cross_fitting indicates whether the learning of different base models should use exactly the same data splits where possible. Note that if there are several models to be synchronized which are classifiers, these cannot be split via stratification.

Source code in metalearners/metalearner.py

def fit(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    """Fit all models of the MetaLearner.

    If pre-fitted models were passed at instantiation, these are never refitted.

    `n_jobs_cross_fitting` will be used at the cross-fitting level and
    `n_jobs_base_learners` will be used at the stage level. `None` means 1 unless in a
    [`joblib.parallel_backend`](https://joblib.readthedocs.io/en/latest/generated/joblib.parallel_backend.html#joblib.parallel_backend)
    context. `-1` means using all processors.
    For more information about parallelism check [parallelism](parallelism.md).


    `fit_params` is an optional `dict` to be forwarded to base estimator `fit` calls. It supports
    two usages patterns:

    ```python
    fit_params={"parameter_of_interest": value_of_interest}
    ```

    ```python
    fit_params={
        "nuisance": {
            "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
            "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
        },
        "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
    }
    ```

    In the former approach, the parameter and value of interest are passed to all base models. In the
    the latter approach, the explicitly qualified parameter-value pairs are passed to respective base
    models and no fitting parameters are passed to base models not explicitly listed. Note that in this
    pattern, propensity models are considered a nuisance model.

    `synchronize_cross_fitting` indicates whether the learning of different base models should use exactly
    the same data splits where possible. Note that if there are several models to be synchronized which are
    classifiers, these cannot be split via stratification.
    """
    self.fit_all_nuisance(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    self.fit_all_treatment(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    return self

fit_all_nuisance

fit_all_nuisance(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all nuisance models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the nuisance models, and in the case they should also be used by the treatment models, these should also be passed in the following call to fit_all_treatment.

This method, combined with fit_all_treatment, facilitates the segmentation of the metalearner fitting process into two distinct parts. This division allows for interventions between the two stages, such as performing feature selection for the treatment models or conducting hyperparameter optimization within the nuisance models.

Source code in metalearners/xlearner.py

def fit_all_nuisance(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    self._validate_treatment(w)
    self._validate_outcome(y, w)

    self._treatment_variants_mask = []

    qualified_fit_params = self._qualified_fit_params(fit_params)

    self._cvs: list = []

    for treatment_variant in range(self.n_variants):
        self._treatment_variants_mask.append(w == treatment_variant)
        if synchronize_cross_fitting:
            cv_split_indices = self._split(
                index_matrix(X, self._treatment_variants_mask[treatment_variant])
            )
        else:
            cv_split_indices = None
        self._cvs.append(cv_split_indices)

    nuisance_jobs: list[_ParallelJoblibSpecification | None] = []
    for treatment_variant in range(self.n_variants):
        mask = self._treatment_variants_mask[treatment_variant]
        nuisance_jobs.append(
            self._nuisance_joblib_specifications(
                X=index_matrix(X, mask),
                y=index_vector(y, mask),
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=treatment_variant,
                n_jobs_cross_fitting=n_jobs_cross_fitting,
                fit_params=qualified_fit_params[NUISANCE][VARIANT_OUTCOME_MODEL],
                cv=self._cvs[treatment_variant],
            )
        )

    nuisance_jobs.append(
        self._nuisance_joblib_specifications(
            X=X,
            y=w,
            model_kind=PROPENSITY_MODEL,
            model_ord=0,
            n_jobs_cross_fitting=n_jobs_cross_fitting,
            fit_params=qualified_fit_params[NUISANCE][PROPENSITY_MODEL],
        )
    )

    parallel = Parallel(n_jobs=n_jobs_base_learners)
    results = parallel(
        delayed(_fit_cross_fit_estimator_joblib)(job)
        for job in nuisance_jobs
        if job is not None
    )
    self._assign_joblib_nuisance_results(results)

    return self

fit_all_treatment

fit_all_treatment(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all treatment models of the MetaLearner.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the treatment models, as the nuisance models should already be fitted.

Source code in metalearners/xlearner.py

def fit_all_treatment(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    if self._treatment_variants_mask is None:
        raise ValueError(
            "The nuisance models need to be fitted before fitting the treatment models."
            "In particular, the MetaLearner's attribute _treatment_variant_mask, "
            "typically set during nuisance fitting, is None."
        )
    if not hasattr(self, "_cvs"):
        raise ValueError(
            "The nuisance models need to be fitted before fitting the treatment models."
            "In particular, the MetaLearner's attribute _cvs, "
            "typically set during nuisance fitting, does not exist."
        )
    qualified_fit_params = self._qualified_fit_params(fit_params)

    treatment_jobs: list[_ParallelJoblibSpecification] = []

    conditional_average_outcome_estimates = (
        self.predict_conditional_average_outcomes(
            X=X,
            is_oos=False,
        )
    )

    for treatment_variant in range(1, self.n_variants):
        imputed_te_control, imputed_te_treatment = self._pseudo_outcome(
            y, w, treatment_variant, conditional_average_outcome_estimates
        )
        treatment_jobs.append(
            self._treatment_joblib_specifications(
                X=index_matrix(X, self._treatment_variants_mask[treatment_variant]),
                y=imputed_te_treatment,
                model_kind=TREATMENT_EFFECT_MODEL,
                model_ord=treatment_variant - 1,
                n_jobs_cross_fitting=n_jobs_cross_fitting,
                fit_params=qualified_fit_params[TREATMENT][TREATMENT_EFFECT_MODEL],
                cv=self._cvs[treatment_variant],
            )
        )

        treatment_jobs.append(
            self._treatment_joblib_specifications(
                X=index_matrix(X, self._treatment_variants_mask[0]),
                y=imputed_te_control,
                model_kind=CONTROL_EFFECT_MODEL,
                model_ord=treatment_variant - 1,
                n_jobs_cross_fitting=n_jobs_cross_fitting,
                fit_params=qualified_fit_params[TREATMENT][CONTROL_EFFECT_MODEL],
                cv=self._cvs[0],
            )
        )

    parallel = Parallel(n_jobs=n_jobs_base_learners)
    results = parallel(
        delayed(_fit_cross_fit_estimator_joblib)(job) for job in treatment_jobs
    )
    self._assign_joblib_treatment_results(results)
    return self

fit_nuisance

fit_nuisance(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit a given nuisance model of a MetaLearner.

y represents the objective of the given nuisance model, not necessarily the outcome of the experiment. If pre-fitted models were passed at instantiation, these are never refitted.

Source code in metalearners/metalearner.py

def fit_nuisance(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit a given nuisance model of a MetaLearner.

    `y` represents the objective of the given nuisance model, not necessarily the outcome of the experiment.
    If pre-fitted models were passed at instantiation, these are never refitted.
    """
    if model_kind in self._prefitted_nuisance_models:
        return self
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._nuisance_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

fit_treatment

fit_treatment(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit the treatment model of a MetaLearner.

y represents the objective of the given treatment model, not necessarily the outcome of the experiment.

Source code in metalearners/metalearner.py

def fit_treatment(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit the treatment model of a MetaLearner.

    `y` represents the objective of the given treatment model, not necessarily the outcome of the experiment.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._treatment_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

nuisance_model_specifications classmethod

nuisance_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all first-stage models.

Source code in metalearners/xlearner.py

@classmethod
def nuisance_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return {
        VARIANT_OUTCOME_MODEL: _ModelSpecifications(
            cardinality=MetaLearner._get_n_variants,
            predict_method=MetaLearner._outcome_predict_method,
        ),
        PROPENSITY_MODEL: _ModelSpecifications(
            cardinality=get_one,
            predict_method=get_predict_proba,
        ),
    }

predict

predict(
    X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Estimate the CATE.

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants} - 1, 1)\) if the outcome is a scalar, i.e. in case of a regression problem.

• \((n_{obs}, n_{variants} - 1, n_{classes})\) if the outcome is a class, i.e. in case of a classification problem.

In the case of multiple treatment variants, the second dimension represents the CATE of the corresponding variant vs the control (variant 0).

Source code in metalearners/xlearner.py

def predict(
    self,
    X: Matrix,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    if self._treatment_variants_mask is None:
        raise ValueError(
            "The MetaLearner needs to be fitted before predicting. "
            "In particular, the X-Learner's attribute _treatment_variant_mask, "
            "typically set during fitting, is None."
        )
    n_outputs = 2 if self.is_classification else 1
    tau_hat = np.zeros((safe_len(X), self.n_variants - 1, n_outputs))
    # Propensity score model is always a classifier so we can't use MEDIAN
    propensity_score_oos = OVERALL if oos_method == MEDIAN else oos_method
    propensity_score = self.predict_nuisance(
        X=X,
        model_kind=PROPENSITY_MODEL,
        model_ord=0,
        is_oos=is_oos,
        oos_method=propensity_score_oos,
    )

    control_indices = self._treatment_variants_mask[0]
    non_control_indices = ~control_indices

    for treatment_variant in range(1, self.n_variants):
        treatment_variant_mask = self._treatment_variants_mask[treatment_variant]
        non_treatment_variant_mask = ~treatment_variant_mask
        if is_oos:
            tau_hat_treatment = self.predict_treatment(
                X=X,
                model_kind=TREATMENT_EFFECT_MODEL,
                model_ord=treatment_variant - 1,
                is_oos=is_oos,
                oos_method=oos_method,
            )
            tau_hat_control = self.predict_treatment(
                X=X,
                model_kind=CONTROL_EFFECT_MODEL,
                model_ord=treatment_variant - 1,
                is_oos=is_oos,
                oos_method=oos_method,
            )
        else:
            tau_hat_treatment = np.zeros(safe_len(X))
            tau_hat_control = np.zeros(safe_len(X))

            tau_hat_treatment[non_treatment_variant_mask] = self.predict_treatment(
                X=index_matrix(X, non_treatment_variant_mask),
                model_kind=TREATMENT_EFFECT_MODEL,
                model_ord=treatment_variant - 1,
                is_oos=True,
                oos_method=oos_method,
            )

            tau_hat_treatment[treatment_variant_mask] = self.predict_treatment(
                X=index_matrix(X, treatment_variant_mask),
                model_kind=TREATMENT_EFFECT_MODEL,
                model_ord=treatment_variant - 1,
                is_oos=False,
            )
            tau_hat_control[control_indices] = self.predict_treatment(
                X=index_matrix(X, control_indices),
                model_kind=CONTROL_EFFECT_MODEL,
                model_ord=treatment_variant - 1,
                is_oos=False,
            )
            tau_hat_control[non_control_indices] = self.predict_treatment(
                X=index_matrix(X, non_control_indices),
                model_kind=CONTROL_EFFECT_MODEL,
                model_ord=treatment_variant - 1,
                is_oos=True,
                oos_method=oos_method,
            )

        propensity_score_treatment = propensity_score[:, treatment_variant] / (
            propensity_score[:, 0] + propensity_score[:, treatment_variant]
        )

        tau_hat_treatment_variant = (
            propensity_score_treatment * tau_hat_control
            + (1 - propensity_score_treatment) * tau_hat_treatment
        )

        if self.is_classification:
            # This is to be consistent with other MetaLearners (e.g. S and T) that automatically
            # work with multiclass outcomes and return the CATE estimate for each class. As the X-Learner only
            # works with binary classes (the pseudo outcome formula does not make sense with
            # multiple classes unless some adaptation is done) we can manually infer the
            # CATE estimate for the complementary class  -- returning a matrix of shape (N, 2).
            tau_hat_treatment_variant = np.stack(
                [-tau_hat_treatment_variant, tau_hat_treatment_variant], axis=1
            )
        else:
            tau_hat_treatment_variant = np.expand_dims(tau_hat_treatment_variant, 1)

        tau_hat[:, treatment_variant - 1] = tau_hat_treatment_variant

    return tau_hat

predict_conditional_average_outcomes

predict_conditional_average_outcomes(
    X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Predict the vectors of conditional average outcomes.

These are defined as \(\mathbb{E}[Y_i(w) | X]\) for each treatment variant \(w\).

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants}, 1)\) if the outcome is a scalar, i.e., in case of a regression problem.

• \((n_{obs}, n_{variants}, n_{classes})\) if the outcome is a class, i.e., in case of a classification problem.

Source code in metalearners/metalearner.py

def predict_conditional_average_outcomes(
    self, X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray:
    if self._treatment_variants_mask is None:
        raise ValueError(
            "The metalearner needs to be fitted before predicting."
            "In particular, the MetaLearner's attribute _treatment_variant_mask, "
            "typically set during fitting, is None."
        )
    # TODO: Consider multiprocessing
    n_obs = safe_len(X)
    nuisance_tensors = self._nuisance_tensors(n_obs)
    conditional_average_outcomes_list = nuisance_tensors[VARIANT_OUTCOME_MODEL]

    for tv in range(self.n_variants):
        if is_oos:
            conditional_average_outcomes_list[tv] = self.predict_nuisance(
                X=X,
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=True,
                oos_method=oos_method,
            )
        else:
            conditional_average_outcomes_list[tv][
                self._treatment_variants_mask[tv]
            ] = self.predict_nuisance(
                X=index_matrix(X, self._treatment_variants_mask[tv]),
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=False,
            )
            conditional_average_outcomes_list[tv][
                ~self._treatment_variants_mask[tv]
            ] = self.predict_nuisance(
                X=index_matrix(X, ~self._treatment_variants_mask[tv]),
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=True,
                oos_method=oos_method,
            )
    return np.stack(conditional_average_outcomes_list, axis=1).reshape(
        n_obs, self.n_variants, -1
    )

predict_nuisance

predict_nuisance(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given nuisance model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_nuisance(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given nuisance model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    predict_method_name = self.nuisance_model_specifications()[model_kind][
        "predict_method"
    ](self)
    predict_method = getattr(
        self._nuisance_models[model_kind][model_ord], predict_method_name
    )
    return predict_method(X_filtered, is_oos, oos_method)

predict_treatment

predict_treatment(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given treatment model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_treatment(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given treatment model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    return self._treatment_models[model_kind][model_ord].predict(
        X_filtered, is_oos, oos_method
    )

shap_values

shap_values(
    X: Matrix,
    shap_explainer_factory: type[Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]

Calculates the shap values for each treatment group.

If explainer is None a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters cate_estimates, cate_model_factory and cate_model_params are ignored.

The parameter shap_explainer_factory can be used to specify the type of shap explainer, for the different options see here.

The returned list contains the shap values for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def shap_values(
    self,
    X: Matrix,
    shap_explainer_factory: type[shap.Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]:
    """Calculates the shap values for each treatment group.

    If `explainer` is `None` a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    The parameter `shap_explainer_factory` can be used to specify the type of shap
    explainer, for the different options see
    [here](https://shap.readthedocs.io/en/latest/api.html#explainers).

    The returned list contains the shap values for each treatment variant in ascending
    order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=None if cate_estimates is None else X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.shap_values(
        X=X,
        shap_explainer_factory=shap_explainer_factory,
        shap_explainer_params=shap_explainer_params,
    )

treatment_model_specifications classmethod

treatment_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all second-stage models.

Source code in metalearners/xlearner.py

@classmethod
def treatment_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return {
        CONTROL_EFFECT_MODEL: _ModelSpecifications(
            cardinality=MetaLearner._get_n_variants_minus_one,
            predict_method=get_predict,
        ),
        TREATMENT_EFFECT_MODEL: _ModelSpecifications(
            cardinality=MetaLearner._get_n_variants_minus_one,
            predict_method=get_predict,
        ),
    }

cross_fit_estimator

CrossFitEstimator dataclass

CrossFitEstimator(
    n_folds: int,
    estimator_factory: type[_ScikitModel],
    estimator_params: dict = dict(),
    enable_overall: bool = True,
    random_state: int | None = None,
)

Helper class for cross-fitting estimators on data.

Conceptually, it allows for fitting n_folds or n_folds + 1 models on n_folds folds of the data.

estimator_factory is a class implementing an estimator with a scikit-learn interface. Instantiation parameters can be passed to estimator_params. An example argument for estimator_factory would be lightgbm.LGBMRegressor.

Importantly, the CrossFitEstimator can handle in-sample and out-of-sample (‘oos’) data for prediction. When doing in-sample prediction the single model will be used in which the respective data point has not been part of the training set. When doing oos prediction, different options exist. These options either rely on combining the n_folds models or using a model trained on all of the data (enable_overall).

n_folds can be set to 1 if the user desires to deactivate cross-fitting. In that case, the CrossFitEstimator would only fit one overall model which would be the one used for either in sample or out of sample predictions. Note that this is not recommended since it can lead to data leakage when doing in-sample predictions.

clone

clone() -> CrossFitEstimator

Construct a new unfitted CrossFitEstimator with the same init parameters.

Source code in metalearners/cross_fit_estimator.py

def clone(self) -> "CrossFitEstimator":
    r"""Construct a new unfitted CrossFitEstimator with the same init parameters."""
    return CrossFitEstimator(
        n_folds=self.n_folds,
        estimator_factory=self.estimator_factory,
        estimator_params=self.estimator_params,
        enable_overall=self.enable_overall,
        random_state=self.random_state,
    )

fit

fit(
    X: Matrix,
    y: Vector | Matrix,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit the underlying estimators.

One estimator is trained per n_folds.

If enable_overall is set, an additional estimator is trained on all data.

n_jobs_cross_fitting can be used to specify the number of jobs for cross-fitting. For more information see the sklearn glossary.

cv can optionally be passed. If passed, it is expected to be a list of (train_indices, test_indices) tuples indicating how to split the data at hand into train and test/estimation sets for different folds.

Source code in metalearners/cross_fit_estimator.py

def fit(
    self,
    X: Matrix,
    y: Vector | Matrix,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit the underlying estimators.

    One estimator is trained per ``n_folds``.

    If ``enable_overall`` is set, an additional estimator is trained on all data.

    ``n_jobs_cross_fitting`` can be used to specify the number of jobs for cross-fitting.
    For more information see the [sklearn glossary](https://scikit-learn.org/stable/glossary.html#term-n_jobs).

    ``cv`` can optionally be passed. If passed, it is expected to be a list of
    (train_indices, test_indices) tuples indicating how to split the data at hand
    into train and test/estimation sets for different folds.
    """
    _validate_data_match_prior_split(safe_len(X), self._test_indices)

    if fit_params is None:
        fit_params = dict()
    if self.n_folds > 1:
        if cv is None:
            if is_classifier(self):
                cv = StratifiedKFold(
                    n_splits=self.n_folds,
                    shuffle=True,
                    random_state=self.random_state,
                )
            else:
                cv = KFold(
                    n_splits=self.n_folds,
                    shuffle=True,
                    random_state=self.random_state,
                )
        cv_result = cross_validate(
            self.estimator_factory(**self.estimator_params),
            X,
            y,
            cv=cv,
            return_estimator=True,
            return_indices=True,
            params=fit_params,
            n_jobs=n_jobs_cross_fitting,
        )
        self._estimators = cv_result["estimator"]
        self._test_indices = cv_result["indices"]["test"]
    if self.enable_overall:
        self._overall_estimator = self._train_overall_estimator(X, y, fit_params)

    if is_classifier(self):
        self._n_classes = len(np.unique(y))
        self.classes_ = np.unique(y)
        for e in self._estimators:
            if set(e.classes_) != set(self.classes_):  # type: ignore
                raise ValueError(
                    "Some folds in cross-fitting had fewer classes than "
                    "the overall dataset. Please check the cv parameter. If you are "
                    "synchronizing the folds in a MetaLearner consider not doing it."
                )
    return self

predict

predict(
    X: Matrix,
    is_oos: bool,
    oos_method: OosMethod | None = None,
    **kwargs,
) -> np.ndarray

Predict from X.

If is_oos, the oos_method will be used to generate predictions on ‘out of sample’ data. ‘Out of sample’ refers to this data not having been used in the fit method. The oos_method 'overall' can only be used if the CrossFitEstimator has been initialized with enable_overall=True.

Source code in metalearners/cross_fit_estimator.py

def predict(
    self,
    X: Matrix,
    is_oos: bool,
    oos_method: OosMethod | None = None,
    **kwargs,
) -> np.ndarray:
    """Predict from ``X``.

    If ``is_oos``, the ``oos_method`` will be used to generate predictions
    on 'out of sample' data. 'Out of sample' refers to this data not having been
    used in the ``fit`` method. The ``oos_method`` ``'overall'`` can only be used
    if the ``CrossFitEstimator`` has been initialized with
    ``enable_overall=True``.
    """
    return self._predict(
        X=X,
        is_oos=is_oos,
        method="predict",
        oos_method=oos_method,
    )

predict_proba

predict_proba(
    X: Matrix,
    is_oos: bool,
    oos_method: OosMethod | None = None,
) -> np.ndarray

Predict probability from X.

If is_oos, the oos_method will be used to generate predictions on ‘out of sample’ data. ‘Out of sample’ refers to this data not having been used in the fit method. The oos_method 'overall' can only be used if the CrossFitEstimator has been initialized with enable_overall=True.

Source code in metalearners/cross_fit_estimator.py

def predict_proba(
    self,
    X: Matrix,
    is_oos: bool,
    oos_method: OosMethod | None = None,
) -> np.ndarray:
    """Predict probability from ``X``.

    If ``is_oos``, the ``oos_method`` will be used to generate predictions
    on 'out of sample' data. 'Out of sample' refers to this data not having been
    used in the ``fit`` method. The ``oos_method`` ``'overall'`` can only be used
    if the ``CrossFitEstimator`` has been initialized with
    ``enable_overall=True``.
    """
    return self._predict(
        X=X,
        is_oos=is_oos,
        method="predict_proba",
        oos_method=oos_method,
    )

score

score(
    X: Matrix,
    y: Vector,
    is_oos: bool,
    oos_method: OosMethod | None = None,
    sample_weight: Vector | None = None,
) -> float

Return the coefficient of determination of the prediction if the estimator is a regressor or the mean accuracy if it is a classifier.

Source code in metalearners/cross_fit_estimator.py

def score(
    self,
    X: Matrix,
    y: Vector,
    is_oos: bool,
    oos_method: OosMethod | None = None,
    sample_weight: Vector | None = None,
) -> float:
    """Return the coefficient of determination of the prediction if the estimator is
    a regressor or the mean accuracy if it is a classifier."""
    if is_classifier(self):
        return accuracy_score(
            y, self.predict(X, is_oos, oos_method), sample_weight=sample_weight
        )
    elif is_regressor(self):
        return r2_score(
            y, self.predict(X, is_oos, oos_method), sample_weight=sample_weight
        )
    else:
        raise NotImplementedError(
            "score is not implemented for this type of estimator."
        )

data_generation

compute_experiment_outputs

compute_experiment_outputs(
    mu: ndarray,
    treatment: Vector,
    sigma_y: float = 1,
    sigma_tau: float = 0.5,
    n_variants: int | None = None,
    is_classification: bool = False,
    positive_proportion: float = 0.5,
    return_probability_cate: bool = False,
    rng: Generator | None = None,
) -> tuple[np.ndarray, np.ndarray]

Compute the experiment’s observed outcomes y and the true CATE.

This function generates experiment outputs and the true CATE values based on the given potential outcomes function and treatments. The treatment effect for each observation is computed as the difference in potential outcomes. Normally distributed noise is added to the response variable \(Y_i(0)\) with standard deviation sigma_y and to each corresponding treatment effect to simulate real-world variance with standard deviation sigma_tau.

treatment must be a vector representing the treatment group assignment for each observation. Each element of the vector is an integer representing a treatment variant starting at 0.

mu must be a matrix of size (n_obs, n_variants) containing the potential outcomes for each observation and treatment variant without added noise.

n_variants can be passed to specify the number of treatment variants. If None, it is inferred from the maximum value in the ‘treatment’ vector plus one.

is_classification determines if the problem to be simulated is a classification problem. If True, the function simulates a classification problem where the response variable is binary and the proportion of positive outputs is controlled by the positive_proportion parameter. It is important to notice that the potential outputs are passed through a sigmoid function and therefore the domain of them can be \(\mathbb{R}\). Classification problems are only implemented for binary treatments.

In the case of a classification problem return_probability_cate specifies if the outputted CATE is the difference in probabilities between treating and not treating or if it samples from a Bernoulli distribution and the difference in samples is returned.

The function returns a tuple containing the following elements:

• y: numpy array of the experiment’s observed outcomes (response variable) after noise addition.

• true_cate: numpy array of the true CATE without any added noise.

Source code in metalearners/data_generation.py

def compute_experiment_outputs(
    mu: np.ndarray,
    treatment: Vector,
    sigma_y: float = 1,
    sigma_tau: float = 0.5,
    n_variants: int | None = None,
    is_classification: bool = False,
    positive_proportion: float = 0.5,
    return_probability_cate: bool = False,
    rng: np.random.Generator | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    r"""Compute the experiment's observed outcomes y and the true CATE.

    This function generates experiment outputs and the true CATE values based on the
    given potential outcomes function and treatments. The treatment effect for each
    observation is computed as the difference in potential outcomes. Normally
    distributed noise is added to the response variable $Y_i(0)$ with standard
    deviation ``sigma_y`` and to each corresponding treatment effect to simulate
    real-world variance with standard deviation ``sigma_tau``.

    ``treatment`` must be a vector representing the treatment group assignment for each
    observation. Each element of the vector is an integer representing a treatment variant
    starting at 0.

    ``mu`` must be a matrix of size ``(n_obs, n_variants)`` containing the potential
    outcomes for each observation and treatment variant without added noise.

    ``n_variants`` can be passed to specify the number of treatment variants. If None,
    it is inferred from the maximum value in the 'treatment' vector plus one.

    ``is_classification`` determines if the problem to be simulated is a classification problem.
    If True, the function simulates a classification problem where the response variable is binary
    and the proportion of positive outputs is controlled by the ``positive_proportion`` parameter.
    It is important to notice that the potential outputs are passed through a sigmoid function and
    therefore the domain of them can be $\mathbb{R}$. Classification problems are
    only implemented for binary treatments.

    In the case of a classification problem ``return_probability_cate`` specifies if the
    outputted CATE is the difference in probabilities between treating and not treating or
    if it samples from a Bernoulli distribution and the difference in samples is returned.

    The function returns a tuple containing the following elements:

    * ``y``: numpy array of the experiment's observed outcomes (response variable) after noise
      addition.

    * ``true_cate``: numpy array of the true CATE without any added noise.
    """
    if rng is None:
        rng = default_rng

    if isinstance(treatment, pd.Series):
        treatment = treatment.to_numpy()

    if n_variants is None:
        n_variants = np.max(treatment) + 1

    n_obs = mu.shape[0]

    if mu.shape[1] != n_variants:
        raise ValueError(
            f"mu should be a matrix where the second dimension has size n_variants. "
            f"n_variants is {n_variants} and mu is a matrix of shape {mu.shape}"
        )

    true_cate = mu[:, 1:] - mu[:, 0].reshape(-1, 1)
    true_y = mu[np.arange(n_obs), treatment]

    y_noise = rng.normal(loc=0, scale=sigma_y, size=n_obs)
    tau_noise = np.c_[
        np.zeros(n_obs),
        rng.normal(loc=0, scale=sigma_tau, size=(n_obs, n_variants - 1)),
    ]

    y = true_y + y_noise + tau_noise[np.arange(n_obs), treatment]

    if is_classification:
        if n_variants > 2:
            raise ValueError(
                "Generating classification problems is only implemented for binary treatments"
            )
        normalizer = np.quantile(true_y, 1 - positive_proportion)

        if return_probability_cate:
            true_cate = sigmoid(mu[:, 1] - normalizer) - sigmoid(mu[:, 0] - normalizer)
        else:
            true_cate = rng.binomial(
                n=1, p=sigmoid(mu[:, 1] - normalizer)
            ) - rng.binomial(n=1, p=sigmoid(mu[:, 0] - normalizer))
        true_cate = true_cate.reshape(-1, 1)
        y = rng.binomial(n=1, p=sigmoid(y - normalizer))

    return y, true_cate

generate_categoricals

generate_categoricals(
    n_obs: int,
    n_features: int,
    n_categories: int | ndarray | None = None,
    n_uniform: int | None = None,
    p_binomial: float = 0.5,
    use_strings: bool = False,
    rng: Generator | None = None,
) -> tuple[np.ndarray, np.ndarray]

Generate a dataset of categorical features.

Generates a dataset of n_obs observations and n_features categorical features. The first n_uniform features are sampled uniformly across their categories and the rest are sampled from a binomial distribution with parameters \(n = c_i\) and \(p = p\_binomial\) where \(c_i\) is the number of categories of feature \(i\).

n_categories is the number of categories of the features, it can either be an int which is used for all the features or an array of length n_features. If None, the number of categories for each feature is sampled from \(c_i \sim \mathcal{U}\{2,3,\dots,10\}\).

In case n_uniform is None, all features are sampled uniformly.

use_strings can be set to True if the wanted represantion of the variables are strings. If set to False it will return an array with dtype np.int64.

The function returns a np.ndarray with the sampled dataset and a np.ndarray with the number of categories for each feature.

Source code in metalearners/data_generation.py

def generate_categoricals(
    n_obs: int,
    n_features: int,
    n_categories: int | np.ndarray | None = None,
    n_uniform: int | None = None,
    p_binomial: float = 0.5,
    use_strings: bool = False,
    rng: np.random.Generator | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    r"""Generate a dataset of categorical features.

    Generates a dataset of ``n_obs`` observations and ``n_features`` categorical
    features. The first ``n_uniform`` features are sampled uniformly across their
    categories and the rest are sampled from a binomial distribution with parameters
    $n = c_i$ and $p = p\_binomial$ where $c_i$ is the number of
    categories of feature $i$.

    ``n_categories`` is the number of categories of the features, it can either be an int
    which is used for all the features or an array of length ``n_features``. If None,
    the number of categories for each feature is sampled from
    $c_i \sim \mathcal{U}\{2,3,\dots,10\}$.

    In case ``n_uniform`` is None, all features are sampled uniformly.

    ``use_strings`` can be set to ``True`` if the wanted represantion of the variables
    are strings. If set to ``False`` it will return an array with dtype ``np.int64``.

    The function returns a ``np.ndarray`` with the sampled dataset and a ``np.ndarray``
    with the number of categories for each feature.
    """
    if rng is None:
        rng = default_rng

    check_probability(p_binomial)

    if n_categories is None:
        n_categories = rng.integers(low=2, high=10, size=n_features, endpoint=True)
    n_categories = np.broadcast_to(n_categories, n_features)

    if n_uniform is None:
        n_uniform = n_features

    dtype = str if use_strings else np.int64
    balanced_features: np.ndarray = np.array([], dtype=dtype).reshape(n_obs, 0)
    unbalanced_features: np.ndarray = np.array([], dtype=dtype).reshape(n_obs, 0)

    if n_uniform > 0:
        balanced_features = rng.integers(
            low=0, high=n_categories[:n_uniform], size=(n_obs, n_uniform)
        ).astype(dtype)

    if (n_non_uniform := n_features - n_uniform) > 0:
        unbalanced_features = rng.binomial(
            n_categories[n_uniform:] - 1,
            p_binomial,
            size=(n_obs, n_non_uniform),
        ).astype(dtype)
    return (
        np.concatenate([balanced_features, unbalanced_features], axis=1),
        n_categories,
    )

generate_covariates

generate_covariates(
    n_obs: int,
    n_features: int,
    n_categoricals: int = 0,
    format: Literal["pandas", "numpy"] = "pandas",
    mu: float | ndarray | None = None,
    wishart_scale: float = 1,
    n_categories: int | ndarray | None = None,
    n_uniform: int | None = None,
    p_binomial: float = 0.5,
    use_strings: bool = False,
    rng: Generator | None = None,
) -> tuple[Matrix, list[int], np.ndarray]

Generates a dataset of covariates with both numerical and categorical features.

Dataset is composed of n_obs observations and n_features features, with the first n_features - n_categoricals being numerical and the rest being categorical. Numerical features are generated using the function generate_numericals and categorical features are generated using the function generate_categoricals.

By default, the generated dataset is returned as a Pandas DataFrame where categorical features are converted to pandas’ Categorical type. Optionally, the dataset can be returned as a numpy array with dtype float64 with format = "numpy". If generating categorical variables, working with pandas DataFrames is preferred as they have support for category dtype.

For mu and wishart_scale see the docstring for generate_numericals

For n_categories, n_uniform, p_binomial and use_strings see the docstring for generate_categoricals.

use_strings can only be set to True when using format = "pandas".

The function returns a tuple of three elements. The first element is the dataset generated (either a numpy array or a pandas DataFrame depending on format). The second element is a list of indices indicating the columns of categorical features in the dataset. The third element is a np.ndarray with the number of categories for each feature.

Source code in metalearners/data_generation.py

def generate_covariates(
    n_obs: int,
    n_features: int,
    n_categoricals: int = 0,
    format: Literal["pandas", "numpy"] = "pandas",
    mu: float | np.ndarray | None = None,
    wishart_scale: float = 1,
    n_categories: int | np.ndarray | None = None,
    n_uniform: int | None = None,
    p_binomial: float = 0.5,
    use_strings: bool = False,
    rng: np.random.Generator | None = None,
) -> tuple[Matrix, list[int], np.ndarray]:
    r"""Generates a dataset of covariates with both numerical and categorical features.

    Dataset is composed of ``n_obs`` observations and ``n_features`` features, with the
    first ``n_features - n_categoricals`` being numerical and the rest being categorical.
    Numerical features are generated using the function
    [`generate_numericals`][metalearners.data_generation.generate_numericals] and categorical features are
    generated using the function [`generate_categoricals`][metalearners.data_generation.generate_categoricals].

    By default, the generated dataset is returned as a Pandas DataFrame where categorical
    features are converted to ``pandas``' [Categorical](https://pandas.pydata.org/docs/reference/api/pandas.Categorical.html#pandas.Categorical)
    type. Optionally, the dataset can be returned as a numpy array with dtype ``float64``
    with ``format = "numpy"``. If generating categorical variables, working with pandas
    DataFrames is preferred as they have support for category dtype.

    For ``mu`` and ``wishart_scale`` see the docstring for
    [`generate_numericals`][metalearners.data_generation.generate_numericals]

    For ``n_categories``, ``n_uniform``, ``p_binomial`` and  ``use_strings``
    see the docstring for [`generate_categoricals`][metalearners.data_generation.generate_categoricals].

    ``use_strings`` can only be set to ``True`` when using ``format = "pandas"``.

    The function returns a tuple of three elements. The first element is the dataset
    generated (either a numpy array or a pandas DataFrame depending on ``format``). The
    second element is a list of indices indicating the columns of categorical features
    in the dataset. The third element is a ``np.ndarray`` with the number of categories
    for each feature.
    """
    if rng is None:
        rng = default_rng
    if format not in _FORMATS:
        raise ValueError(f"format needs to be one of {_FORMATS}")

    if format == "numpy" and use_strings:
        raise ValueError("if format is numpy then use_strings must be False")

    numerical_features = np.array([]).reshape(n_obs, 0)
    categorical_features = np.array([]).reshape(n_obs, 0)

    if n_features - n_categoricals > 0:
        numerical_features = generate_numericals(
            n_obs=n_obs,
            n_features=n_features - n_categoricals,
            mu=mu,
            wishart_scale=wishart_scale,
            rng=rng,
        )
    if n_categoricals > 0:
        categorical_features, n_categories = generate_categoricals(
            n_obs=n_obs,
            n_features=n_categoricals,
            n_categories=n_categories,
            n_uniform=n_uniform,
            p_binomial=p_binomial,
            use_strings=use_strings,
            rng=rng,
        )
    else:
        n_categories = np.array([])

    categorical_features_idx = list(range(n_features - n_categoricals, n_features))

    if format == "numpy":
        features = np.concatenate([numerical_features, categorical_features], axis=1)
    elif format == "pandas":
        numerical_features = pd.DataFrame(numerical_features)
        categorical_features = pd.DataFrame(categorical_features)
        features = pd.concat(
            [numerical_features, categorical_features], axis=1, ignore_index=True
        )
        features[categorical_features_idx] = features[categorical_features_idx].astype(
            "category"
        )
        for i, c in enumerate(categorical_features_idx):
            categories = list(range(n_categories[i]))
            if use_strings:
                categories = list(map(str, categories))  # type: ignore
            # We need to set the categories manually as there may be some unsampled categories,
            # and it may be possible that the user relies on having all of them when using OHE
            # for the potential outcomes function.
            features[c] = features[c].cat.set_categories(categories)
    return features, categorical_features_idx, n_categories

generate_numericals

generate_numericals(
    n_obs: int,
    n_features: int,
    mu: float | ndarray | None = None,
    wishart_scale: float = 1,
    rng: Generator | None = None,
) -> np.ndarray

Generate a dataset of numerical features.

Generates a dataset of n_obs observations and n_features numerical features. These are sampled from \(\mathcal{N}(\mu, \Sigma)\) where \(\mu \sim \mathcal{U}[-5,5]\) unless specified in mu and \(\Sigma \sim \mathcal{W}(d, \sigma_w I_d)\) where \(W\) is the Wishart distribution and \(d\) the number of features.

mu can be either a float or an array of length n_features.

wishart_scale should be \(\geq 0\) , in case it is 0 then \(\Sigma = I_d\).

Source code in metalearners/data_generation.py

def generate_numericals(
    n_obs: int,
    n_features: int,
    mu: float | np.ndarray | None = None,
    wishart_scale: float = 1,
    rng: np.random.Generator | None = None,
) -> np.ndarray:
    r"""Generate a dataset of numerical features.

    Generates a dataset of ``n_obs`` observations and ``n_features`` numerical features.
    These are sampled from $\mathcal{N}(\mu, \Sigma)$ where
    $\mu \sim \mathcal{U}[-5,5]$ unless specified in ``mu`` and
    $\Sigma \sim \mathcal{W}(d, \sigma_w I_d)$ where $W$ is the Wishart
    distribution and $d$ the number of features.


    ``mu`` can be either a float or an array of length ``n_features``.

    ``wishart_scale`` should be $\geq 0$ , in case it is 0 then $\Sigma = I_d$.
    """
    if rng is None:
        rng = default_rng
    if mu is None:
        mu = rng.uniform(-5, 5, size=n_features)
    mu = np.broadcast_to(mu, n_features)
    if wishart_scale < 0:
        raise ValueError("wishart_scale needs to be >= 0")
    if wishart_scale > 0:
        cov_matrix = wishart.rvs(
            df=n_features,
            scale=wishart_scale * np.eye(n_features),
            random_state=rng,
        ).reshape(n_features, n_features)
    else:
        cov_matrix = np.eye(n_features)
    features = rng.multivariate_normal(mean=mu, cov=cov_matrix, size=n_obs)

    return features

generate_treatment

generate_treatment(
    propensity_scores: ndarray, rng: Generator | None = None
) -> np.ndarray

Generates a treatment assignment based on the provided propensity scores.

The function first determines the number of treatment variants based on the shape of the input propensity scores. If the propensity score array has a single dimension or only one column in the second dimension, there are two treatment variants (treated vs not-treated), and the value is interpreted as the treatment probability. Otherwise, the second dimension of the propensity scores array indicates the number of treatment variants.

Each observation is assigned to a treatment group by drawing from a categorical distribution where the probability of each treatment group is given by the propensity scores.

propensity_scores should be of size (n_obs,) or (n_obs, n_variants), where n_obs is the number of observations and n_variants is the number of treatment variants.

The function return an array of shape (n_obs,) where each element indicates the treatment variant received.

Source code in metalearners/data_generation.py

def generate_treatment(
    propensity_scores: np.ndarray, rng: np.random.Generator | None = None
) -> np.ndarray:
    """Generates a treatment assignment based on the provided propensity scores.

    The function first determines the number of treatment variants based on the shape of
    the input propensity scores. If the propensity score array has a single dimension or
    only one column in the second dimension, there are two treatment variants (treated
    vs not-treated), and the value is interpreted as the treatment probability.
    Otherwise, the second dimension of the propensity scores array indicates the number
    of treatment variants.

    Each observation is assigned to a treatment group by drawing from a categorical
    distribution where the probability of each treatment group is given by the
    propensity scores.

    ``propensity_scores`` should be of size ``(n_obs,)`` or ``(n_obs, n_variants)``,
    where ``n_obs`` is the number of observations and ``n_variants`` is the number of
    treatment variants.

    The function return an array of shape ``(n_obs,)`` where each element indicates the
    treatment variant received.
    """
    if rng is None:
        rng = default_rng
    n_variants = get_n_variants(propensity_scores)
    propensity_scores = convert_and_pad_propensity_score(propensity_scores, n_variants)
    check_propensity_score(propensity_scores, n_variants=n_variants, sum_to_one=True)

    treatment = rng.multinomial(1, propensity_scores).argmax(axis=1)
    return treatment

insert_missing

insert_missing(
    X: Matrix,
    missing_probability: float = 0.1,
    rng: Generator | None = None,
) -> Matrix

Inserts missing values into the dataset.

Each element of the dataset has a missing_probability chance of being replaced with a NaN, thus simulating a dataset with missing values.

The function returns a copy of the original dataset, but with some elements replaced by NaNs.

Source code in metalearners/data_generation.py

def insert_missing(
    X: Matrix,
    missing_probability: float = 0.1,
    rng: np.random.Generator | None = None,
) -> Matrix:
    """Inserts missing values into the dataset.

    Each element of the dataset has a ``missing_probability`` chance of being replaced
    with a NaN, thus simulating a dataset with missing values.

    The function returns a copy of the original dataset, but with some elements replaced
    by NaNs.
    """
    if rng is None:
        rng = default_rng
    check_probability(missing_probability, zero_included=True)
    missing_mask = rng.binomial(1, p=missing_probability, size=X.shape).astype("bool")

    masked = X.copy()
    masked[missing_mask] = np.nan
    return masked

drlearner

DRLearner

DRLearner(
    is_classification: bool,
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel]
    | None = None,
    nuisance_model_params: Params
    | dict[str, Params]
    | None = None,
    treatment_model_params: Params
    | dict[str, Params]
    | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[
        str, list[CrossFitEstimator]
    ]
    | None = None,
    fitted_propensity_model: CrossFitEstimator
    | None = None,
    feature_set: Features
    | dict[str, Features]
    | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
    adaptive_clipping: bool = False,
)

Bases: _ConditionalAverageOutcomeMetaLearner

DR-Learner for CATE estimation as described by Kennedy (2020).

Importantly, the current DR-Learner implementation only supports:

• binary classes in case of a classification outcome

The DR-Learner contains the following nuisance models:

• a "propensity_model" estimating \(\Pr[W=k|X]\)
• one "variant_outcome_model" for each treatment variant (including control) estimating \(\mathbb{E}[Y|X, W=k]\)

and one treatment model for each treatment variant (without control):

• "treatment_model" which estimates \(\mathbb{E}[Y(k) - Y(0) | X]\)

If adaptive_clipping is set to True, then the pseudo outcomes are computed using adaptive propensity clipping described in section 4.1, equation DR-Switch of Mahajan et al. (2024).

Source code in metalearners/drlearner.py

def __init__(
    self,
    is_classification: bool,
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel] | None = None,
    nuisance_model_params: Params | dict[str, Params] | None = None,
    treatment_model_params: Params | dict[str, Params] | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[str, list[CrossFitEstimator]] | None = None,
    fitted_propensity_model: CrossFitEstimator | None = None,
    feature_set: Features | dict[str, Features] | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
    adaptive_clipping: bool = False,
):
    super().__init__(
        nuisance_model_factory=nuisance_model_factory,
        is_classification=is_classification,
        n_variants=n_variants,
        treatment_model_factory=treatment_model_factory,
        propensity_model_factory=propensity_model_factory,
        nuisance_model_params=nuisance_model_params,
        treatment_model_params=treatment_model_params,
        propensity_model_params=propensity_model_params,
        fitted_nuisance_models=fitted_nuisance_models,
        fitted_propensity_model=fitted_propensity_model,
        feature_set=feature_set,
        n_folds=n_folds,
        random_state=random_state,
    )
    self.adaptive_clipping = adaptive_clipping

init_args property

init_args: dict[str, Any]

Create initialization parameters for a new MetaLearner.

Importantly, this does not copy further internal state, such as the weights or parameters of trained base models.

average_treatment_effect

average_treatment_effect(
    X: Matrix, y: Vector, w: Vector, is_oos: bool
) -> tuple[np.ndarray, np.ndarray]

Compute Average Treatment Effect (ATE) for each treatment variant using the Augmented IPW estimator (Robins et al 1994). Does not require fitting a second- stage treatment model: it uses the pseudo-outcome alone and computes the point estimate and standard error. Can be used following the fit_all_nuisance method.

Parameters:

Name	Type	Description	Default
X	Matrix	Covariate matrix	required
y	Vector	Outcome vector	required
w	Vector	Treatment vector	required
is_oos	bool	indicator whether data is out of sample	required

Returns:

Type	Description
ndarray	np.ndarray: Treatment effect for each treatment variant.
ndarray	np.ndarray: Standard error for each treatment variant.

Source code in metalearners/drlearner.py

def average_treatment_effect(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
) -> tuple[np.ndarray, np.ndarray]:
    """Compute Average Treatment Effect (ATE) for each treatment variant using the
    Augmented IPW estimator (Robins et al 1994). Does not require fitting a second-
    stage treatment model: it uses the pseudo-outcome alone and computes the point
    estimate and standard error. Can be used following the
    [`fit_all_nuisance`][metalearners.drlearner.DRLearner.fit_all_nuisance] method.

    Args:
        X (Matrix): Covariate matrix
        y (Vector): Outcome vector
        w (Vector): Treatment vector
        is_oos (bool): indicator whether data is out of sample

    Returns:
        np.ndarray: Treatment effect for each treatment variant.
        np.ndarray: Standard error for each treatment variant.
    """
    if not self._nuisance_models_fit:
        raise ValueError(
            "The nuisance models need to be fitted before computing the treatment effect."
        )
    gamma_matrix = np.zeros((safe_len(X), self.n_variants - 1))
    for treatment_variant in range(1, self.n_variants):
        gamma_matrix[:, treatment_variant - 1] = self._pseudo_outcome(
            X=X,
            w=w,
            y=y,
            treatment_variant=treatment_variant,
            is_oos=is_oos,
        )
    treatment_effect = gamma_matrix.mean(axis=0)
    standard_error = gamma_matrix.std(axis=0) / np.sqrt(safe_len(X))
    return treatment_effect, standard_error

evaluate

evaluate(
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]

Evaluate the MetaLearner.

The keys in scoring which are not a name of a model contained in the MetaLearner will be ignored, for information about this names check nuisance_model_specifications and treatment_model_specifications. The values must be a list of:

• string representing a sklearn scoring method. Check here for the possible values.
• Callable with signature scorer(estimator, X, y_true, **kwargs). We recommend using sklearn.metrics.make_scorer to create such a Callable.

If some model name is not present in the keys of scoring then the default used metrics will be neg_log_loss if it is a classifier and neg_root_mean_squared_error if it is a regressor.

The returned dictionary keys have the following structure:

• For nuisance models:

  • If the cardinality is one: f"{model_kind}_{scorer}"
  • If there is one model for each treatment variant (including control): f"{model_kind}_{treatment_variant}_{scorer}"

• For treatment models: f"{model_kind}_{treatment_variant}_vs_0_{scorer}"

Where scorer is the name of the scorer if it is a string and "custom_scorer_{idx}" if it is a callable where idx is the index in the scorers list.

Source code in metalearners/drlearner.py

def evaluate(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]:
    safe_scoring = self._scoring(scoring)

    variant_outcome_evaluation = _evaluate_model_kind(
        cfes=self._nuisance_models[VARIANT_OUTCOME_MODEL],
        Xs=[index_matrix(X, w == tv) for tv in range(self.n_variants)],
        ys=[index_vector(y, w == tv) for tv in range(self.n_variants)],
        scorers=safe_scoring[VARIANT_OUTCOME_MODEL],
        model_kind=VARIANT_OUTCOME_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=False,
        feature_set=self.feature_set[VARIANT_OUTCOME_MODEL],
    )

    propensity_evaluation = _evaluate_model_kind(
        cfes=self._nuisance_models[PROPENSITY_MODEL],
        Xs=[X],
        ys=[w],
        scorers=safe_scoring[PROPENSITY_MODEL],
        model_kind=PROPENSITY_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=False,
        feature_set=self.feature_set[PROPENSITY_MODEL],
    )

    pseudo_outcome: list[np.ndarray] = []
    for treatment_variant in range(1, self.n_variants):
        tv_pseudo_outcome = self._pseudo_outcome(
            X=X,
            y=y,
            w=w,
            treatment_variant=treatment_variant,
            is_oos=is_oos,
            oos_method=oos_method,
        )
        pseudo_outcome.append(tv_pseudo_outcome)

    treatment_evaluation = _evaluate_model_kind(
        self._treatment_models[TREATMENT_MODEL],
        Xs=[X for _ in range(1, self.n_variants)],
        ys=pseudo_outcome,
        scorers=safe_scoring[TREATMENT_MODEL],
        model_kind=TREATMENT_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=True,
        feature_set=self.feature_set[TREATMENT_MODEL],
    )

    return variant_outcome_evaluation | propensity_evaluation | treatment_evaluation

explainer

explainer(
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer

Create an Explainer

which can be used in feature_importances.

This function can be used in two distinct manners based on the provided parameters:

• When parameters X, cate_estimates, and cate_model_factory are all set to None, the function creates an Explainer using the pre-existing treatment models. If these models do not exist, however, it triggers a ValueError.
• On the contrary, if X, cate_estimates, and cate_model_factory are not None, the function initiates an instance of the Explainer class using these parameters. This instance then fits new models for each treatment variant, and these models are employed to calculate the importance of features.

Source code in metalearners/metalearner.py

def explainer(
    self,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer:
    r"""Create an [`Explainer`][metalearners.explainer.Explainer]

    which can be used in [`feature_importances`][metalearners.metalearner.MetaLearner.feature_importances].

    This function can be used in two distinct manners based on the provided parameters:

    *   When parameters `X`, `cate_estimates`, and `cate_model_factory` are all
        set to `None`, the function creates an [`Explainer`][metalearners.explainer.Explainer]
        using the pre-existing treatment models. If these models do not exist, however,
        it triggers a `ValueError`.
    *   On the contrary, if `X`, `cate_estimates`, and `cate_model_factory` are
        not `None`, the function initiates an instance of the [`Explainer`][metalearners.explainer.Explainer]
        class using these parameters. This instance then fits new models for each
        treatment variant, and these models are employed to calculate the importance
        of features.
    """
    if X is None and cate_estimates is None and cate_model_factory is None:
        try:
            cate_cf_models = self._treatment_models[TREATMENT_MODEL]
            cate_models = [cf._overall_estimator for cf in cate_cf_models]
        except KeyError:
            raise ValueError(
                "The metalearner does not have treatment models; hence X, cate_estimates, "
                "and cate_model_factory need to be defined."
            )
        return Explainer(cate_models=cate_models)  # type: ignore
    elif (
        X is not None
        and cate_estimates is not None
        and cate_model_factory is not None
    ):
        return Explainer.from_estimates(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    else:
        raise ValueError(
            "Either all of [X, cate_estimates, cate_model_factory] are None or all "
            "of them must be defined."
        )

feature_importances

feature_importances(
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances

Calculates the feature importance for each treatment group.

If explainer is None, a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters X, cate_estimates, cate_model_factory and cate_model_params are ignored.

If normalization = True, for each treatment variant the feature importances are normalized so that they sum to 1.

feature_names is optional but in the case it’s not passed the names of the features will default to f"Feature {i}" where i is the corresponding feature index.

The returned list contains the feature importances for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def feature_importances(
    self,
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances:
    r"""Calculates the feature importance for each treatment group.

    If `explainer` is `None`, a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `X`, `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    If `normalization = True`, for each treatment variant the feature importances
    are normalized so that they sum to 1.

    `feature_names` is optional but in the case it's not passed the names of the
    features will default to `f"Feature {i}"` where `i` is the corresponding
    feature index.

    The returned list contains the feature importances for each treatment variant in
    ascending order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.feature_importances(
        normalize=normalize, feature_names=feature_names, sort_values=sort_values
    )

fit

fit(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

n_jobs_cross_fitting will be used at the cross-fitting level and n_jobs_base_learners will be used at the stage level. None means 1 unless in a joblib.parallel_backend context. -1 means using all processors. For more information about parallelism check parallelism.

fit_params is an optional dict to be forwarded to base estimator fit calls. It supports two usages patterns:

fit_params={"parameter_of_interest": value_of_interest}

fit_params={
    "nuisance": {
        "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
        "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
    },
    "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
}

In the former approach, the parameter and value of interest are passed to all base models. In the the latter approach, the explicitly qualified parameter-value pairs are passed to respective base models and no fitting parameters are passed to base models not explicitly listed. Note that in this pattern, propensity models are considered a nuisance model.

synchronize_cross_fitting indicates whether the learning of different base models should use exactly the same data splits where possible. Note that if there are several models to be synchronized which are classifiers, these cannot be split via stratification.

Source code in metalearners/metalearner.py

def fit(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    """Fit all models of the MetaLearner.

    If pre-fitted models were passed at instantiation, these are never refitted.

    `n_jobs_cross_fitting` will be used at the cross-fitting level and
    `n_jobs_base_learners` will be used at the stage level. `None` means 1 unless in a
    [`joblib.parallel_backend`](https://joblib.readthedocs.io/en/latest/generated/joblib.parallel_backend.html#joblib.parallel_backend)
    context. `-1` means using all processors.
    For more information about parallelism check [parallelism](parallelism.md).


    `fit_params` is an optional `dict` to be forwarded to base estimator `fit` calls. It supports
    two usages patterns:

    ```python
    fit_params={"parameter_of_interest": value_of_interest}
    ```

    ```python
    fit_params={
        "nuisance": {
            "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
            "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
        },
        "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
    }
    ```

    In the former approach, the parameter and value of interest are passed to all base models. In the
    the latter approach, the explicitly qualified parameter-value pairs are passed to respective base
    models and no fitting parameters are passed to base models not explicitly listed. Note that in this
    pattern, propensity models are considered a nuisance model.

    `synchronize_cross_fitting` indicates whether the learning of different base models should use exactly
    the same data splits where possible. Note that if there are several models to be synchronized which are
    classifiers, these cannot be split via stratification.
    """
    self.fit_all_nuisance(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    self.fit_all_treatment(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    return self

fit_all_nuisance

fit_all_nuisance(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all nuisance models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the nuisance models, and in the case they should also be used by the treatment models, these should also be passed in the following call to fit_all_treatment.

This method, combined with fit_all_treatment, facilitates the segmentation of the metalearner fitting process into two distinct parts. This division allows for interventions between the two stages, such as performing feature selection for the treatment models or conducting hyperparameter optimization within the nuisance models.

Source code in metalearners/drlearner.py

def fit_all_nuisance(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    self._validate_treatment(w)
    self._validate_outcome(y, w)

    self._treatment_variants_mask = []

    qualified_fit_params = self._qualified_fit_params(fit_params)

    for treatment_variant in range(self.n_variants):
        self._treatment_variants_mask.append(w == treatment_variant)

    self._cv_split_indices: SplitIndices | None

    if synchronize_cross_fitting:
        self._cv_split_indices = self._split(X)
    else:
        self._cv_split_indices = None

    nuisance_jobs: list[_ParallelJoblibSpecification | None] = []
    for treatment_variant in range(self.n_variants):
        mask = self._treatment_variants_mask[treatment_variant]
        X_masked = index_matrix(X, mask)
        y_masked = index_vector(y, mask)
        nuisance_jobs.append(
            self._nuisance_joblib_specifications(
                X=X_masked,
                y=y_masked,
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=treatment_variant,
                n_jobs_cross_fitting=n_jobs_cross_fitting,
                fit_params=qualified_fit_params[NUISANCE][VARIANT_OUTCOME_MODEL],
            )
        )

    nuisance_jobs.append(
        self._nuisance_joblib_specifications(
            X=X,
            y=w,
            model_kind=PROPENSITY_MODEL,
            model_ord=0,
            n_jobs_cross_fitting=n_jobs_cross_fitting,
            fit_params=qualified_fit_params[NUISANCE][PROPENSITY_MODEL],
            cv=self._cv_split_indices,
        )
    )

    parallel = Parallel(n_jobs=n_jobs_base_learners)
    results = parallel(
        delayed(_fit_cross_fit_estimator_joblib)(job)
        for job in nuisance_jobs
        if job is not None
    )

    self._assign_joblib_nuisance_results(results)
    self._nuisance_models_fit = True
    return self

fit_all_treatment

fit_all_treatment(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all treatment models of the MetaLearner.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the treatment models, as the nuisance models should already be fitted.

Source code in metalearners/drlearner.py

def fit_all_treatment(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    if not hasattr(self, "_cv_split_indices"):
        raise ValueError(
            "The nuisance models need to be fitted before fitting the treatment models."
            "In particular, the MetaLearner's attribute _cv_split_indices, "
            "typically set during nuisance fitting, does not exist."
        )
    qualified_fit_params = self._qualified_fit_params(fit_params)
    treatment_jobs: list[_ParallelJoblibSpecification] = []
    for treatment_variant in range(1, self.n_variants):
        pseudo_outcomes = self._pseudo_outcome(
            X=X,
            w=w,
            y=y,
            treatment_variant=treatment_variant,
            is_oos=False,
        )

        treatment_jobs.append(
            self._treatment_joblib_specifications(
                X=X,
                y=pseudo_outcomes,
                model_kind=TREATMENT_MODEL,
                model_ord=treatment_variant - 1,
                n_jobs_cross_fitting=n_jobs_cross_fitting,
                fit_params=qualified_fit_params[TREATMENT][TREATMENT_MODEL],
                cv=self._cv_split_indices,
            )
        )
    parallel = Parallel(n_jobs=n_jobs_base_learners)
    results = parallel(
        delayed(_fit_cross_fit_estimator_joblib)(job) for job in treatment_jobs
    )
    self._assign_joblib_treatment_results(results)
    return self

fit_nuisance

fit_nuisance(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit a given nuisance model of a MetaLearner.

y represents the objective of the given nuisance model, not necessarily the outcome of the experiment. If pre-fitted models were passed at instantiation, these are never refitted.

Source code in metalearners/metalearner.py

def fit_nuisance(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit a given nuisance model of a MetaLearner.

    `y` represents the objective of the given nuisance model, not necessarily the outcome of the experiment.
    If pre-fitted models were passed at instantiation, these are never refitted.
    """
    if model_kind in self._prefitted_nuisance_models:
        return self
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._nuisance_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

fit_treatment

fit_treatment(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit the treatment model of a MetaLearner.

y represents the objective of the given treatment model, not necessarily the outcome of the experiment.

Source code in metalearners/metalearner.py

def fit_treatment(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit the treatment model of a MetaLearner.

    `y` represents the objective of the given treatment model, not necessarily the outcome of the experiment.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._treatment_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

nuisance_model_specifications classmethod

nuisance_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all first-stage models.

Source code in metalearners/drlearner.py

@classmethod
def nuisance_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return {
        PROPENSITY_MODEL: _ModelSpecifications(
            cardinality=get_one,
            predict_method=get_predict_proba,
        ),
        VARIANT_OUTCOME_MODEL: _ModelSpecifications(
            cardinality=MetaLearner._get_n_variants,
            predict_method=MetaLearner._outcome_predict_method,
        ),
    }

predict

predict(
    X, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Estimate the CATE.

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants} - 1, 1)\) if the outcome is a scalar, i.e. in case of a regression problem.

• \((n_{obs}, n_{variants} - 1, n_{classes})\) if the outcome is a class, i.e. in case of a classification problem.

In the case of multiple treatment variants, the second dimension represents the CATE of the corresponding variant vs the control (variant 0).

Source code in metalearners/drlearner.py

def predict(
    self,
    X,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    n_outputs = 2 if self.is_classification else 1
    estimates = np.zeros((safe_len(X), self.n_variants - 1, n_outputs))
    for treatment_variant in range(1, self.n_variants):
        estimates_variant = self.predict_treatment(
            X,
            is_oos=is_oos,
            oos_method=oos_method,
            model_kind=TREATMENT_MODEL,
            model_ord=treatment_variant - 1,
        )
        if self.is_classification:
            # This is to be consistent with other MetaLearners (e.g. S and T) that automatically
            # work with multiclass outcomes and return the CATE estimate for each class. As the DR-Learner only
            # works with binary classes (the pseudo outcome formula does not make sense with
            # multiple classes unless some adaptation is done) we can manually infer the
            # CATE estimate for the complementary class  -- returning a matrix of shape (N, 2).
            estimates_variant = np.stack(
                [-estimates_variant, estimates_variant], axis=1
            )
        else:
            estimates_variant = np.expand_dims(estimates_variant, 1)

        estimates[:, treatment_variant - 1] = estimates_variant
    return estimates

predict_conditional_average_outcomes

predict_conditional_average_outcomes(
    X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Predict the vectors of conditional average outcomes.

These are defined as \(\mathbb{E}[Y_i(w) | X]\) for each treatment variant \(w\).

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants}, 1)\) if the outcome is a scalar, i.e., in case of a regression problem.

• \((n_{obs}, n_{variants}, n_{classes})\) if the outcome is a class, i.e., in case of a classification problem.

Source code in metalearners/metalearner.py

def predict_conditional_average_outcomes(
    self, X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray:
    if self._treatment_variants_mask is None:
        raise ValueError(
            "The metalearner needs to be fitted before predicting."
            "In particular, the MetaLearner's attribute _treatment_variant_mask, "
            "typically set during fitting, is None."
        )
    # TODO: Consider multiprocessing
    n_obs = safe_len(X)
    nuisance_tensors = self._nuisance_tensors(n_obs)
    conditional_average_outcomes_list = nuisance_tensors[VARIANT_OUTCOME_MODEL]

    for tv in range(self.n_variants):
        if is_oos:
            conditional_average_outcomes_list[tv] = self.predict_nuisance(
                X=X,
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=True,
                oos_method=oos_method,
            )
        else:
            conditional_average_outcomes_list[tv][
                self._treatment_variants_mask[tv]
            ] = self.predict_nuisance(
                X=index_matrix(X, self._treatment_variants_mask[tv]),
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=False,
            )
            conditional_average_outcomes_list[tv][
                ~self._treatment_variants_mask[tv]
            ] = self.predict_nuisance(
                X=index_matrix(X, ~self._treatment_variants_mask[tv]),
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=True,
                oos_method=oos_method,
            )
    return np.stack(conditional_average_outcomes_list, axis=1).reshape(
        n_obs, self.n_variants, -1
    )

predict_nuisance

predict_nuisance(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given nuisance model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_nuisance(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given nuisance model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    predict_method_name = self.nuisance_model_specifications()[model_kind][
        "predict_method"
    ](self)
    predict_method = getattr(
        self._nuisance_models[model_kind][model_ord], predict_method_name
    )
    return predict_method(X_filtered, is_oos, oos_method)

predict_treatment

predict_treatment(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given treatment model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_treatment(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given treatment model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    return self._treatment_models[model_kind][model_ord].predict(
        X_filtered, is_oos, oos_method
    )

shap_values

shap_values(
    X: Matrix,
    shap_explainer_factory: type[Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]

Calculates the shap values for each treatment group.

If explainer is None a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters cate_estimates, cate_model_factory and cate_model_params are ignored.

The parameter shap_explainer_factory can be used to specify the type of shap explainer, for the different options see here.

The returned list contains the shap values for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def shap_values(
    self,
    X: Matrix,
    shap_explainer_factory: type[shap.Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]:
    """Calculates the shap values for each treatment group.

    If `explainer` is `None` a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    The parameter `shap_explainer_factory` can be used to specify the type of shap
    explainer, for the different options see
    [here](https://shap.readthedocs.io/en/latest/api.html#explainers).

    The returned list contains the shap values for each treatment variant in ascending
    order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=None if cate_estimates is None else X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.shap_values(
        X=X,
        shap_explainer_factory=shap_explainer_factory,
        shap_explainer_params=shap_explainer_params,
    )

treatment_model_specifications classmethod

treatment_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all second-stage models.

Source code in metalearners/drlearner.py

@classmethod
def treatment_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return {
        TREATMENT_MODEL: _ModelSpecifications(
            cardinality=MetaLearner._get_n_variants_minus_one,
            predict_method=get_predict,
        )
    }

explainer

Explainer

Explainer(cate_models: list[_ScikitModel])

Responsible class for managing all functions related to feature explanation and interpretation.

The cate_models parameter should be a list of length \(n_{variants} -1\) containing a model for each treatment variant which estimates \(\tau_k\). The models should not be a CrossFitEstimator rather just a plain sklearn BaseEstimator. A suggested option in the case of a CrossFitEstimator would be to use their _overall_estimator. These models should already be fitted on the data.

Source code in metalearners/explainer.py

def __init__(
    self,
    cate_models: list[_ScikitModel],
):
    self.n_variants = len(cate_models) + 1
    self.cate_models = cate_models

feature_importances

feature_importances(
    normalize: bool = False,
    feature_names: Collection[str] | None = None,
    sort_values: bool = False,
) -> FeatureImportances

Calculates the feature importance for each treatment group.

If normalization = True, for each treatment variant the feature importances are normalized so that they sum to 1.

feature_names is optional but in the case it’s not passed the names of the features will default to f"Feature {i}" where i is the corresponding feature key.

Source code in metalearners/explainer.py

def feature_importances(
    self,
    normalize: bool = False,
    feature_names: Collection[str] | None = None,
    sort_values: bool = False,
) -> FeatureImportances:
    r"""Calculates the feature importance for each treatment group.

    If ``normalization = True``, for each treatment variant the feature importances
    are normalized so that they sum to 1.

    ``feature_names`` is optional but in the case it's not passed the names of the
    features will default to ``f"Feature {i}"`` where ``i`` is the corresponding
    feature key.
    """
    feature_importances: FeatureImportances = []
    for tv in range(self.n_variants - 1):
        if not hasattr(self.cate_models[tv], "feature_importances_"):
            raise ValueError(
                f"Model used for treatment variant {tv + 1} has no attribute feature_importances_. "
                "You need to use a model which computes them, e.g. LGBMRegressor."
            )

        variant_feature_importance_np = self.cate_models[tv].feature_importances_  # type: ignore
        if normalize:
            variant_feature_importance_np = variant_feature_importance_np / np.sum(
                variant_feature_importance_np
            )
        variant_feature_importance = _build_feature_importance_dict(
            variant_feature_importance_np,
            sort_values,
            feature_names,
        )
        feature_importances.append(variant_feature_importance)

    return feature_importances

from_estimates classmethod

from_estimates(
    X: Matrix,
    cate_estimates: ndarray,
    cate_model_factory: type[_ScikitModel],
    cate_model_params: Params | None = None,
) -> Explainer

Create an Explainer object from CATE estimates.

This function will fit a model for each treatment variant with X as its input and the corresponding CATE estimates as its output.

The cate_estimates should be the raw outcome of a MetaLearner with 3 dimensions and should not be simplified.

Source code in metalearners/explainer.py

@classmethod
def from_estimates(
    cls,
    X: Matrix,
    cate_estimates: np.ndarray,
    cate_model_factory: type[_ScikitModel],
    cate_model_params: Params | None = None,
) -> "Explainer":
    r"""Create an ``Explainer`` object from CATE estimates.

    This function will fit a model for each treatment variant with ``X`` as its input
    and the corresponding CATE estimates as its output.

    The ``cate_estimates`` should be the raw outcome of a MetaLearner with 3 dimensions
    and should not be simplified.
    """
    if safe_len(X) != len(cate_estimates) or safe_len(X) == 0:
        raise ValueError(
            "X and cate_estimates should contain the same number of observations "
            "and not be empty."
        )
    if np.any(np.isnan(cate_estimates)) or np.any(np.isinf(cate_estimates)):
        raise ValueError("cate_estimates can not contain any NaN or inf.")

    cate_estimates = simplify_output_2d(
        cate_estimates
    )  # TODO: This does not work for multiclass, do we want to consider it?
    if cate_model_params is None:
        cate_model_params = {}

    n_variants = cate_estimates.shape[1] + 1
    cate_models = [
        cate_model_factory(**cate_model_params).fit(X, cate_estimates[:, tv])
        for tv in range(n_variants - 1)
    ]
    return cls(cate_models)

shap_values

shap_values(
    X: Matrix,
    shap_explainer_factory: type[Explainer],
    shap_explainer_params: dict | None = None,
) -> list[np.ndarray]

Calculates the shap values for each treatment group.

The parameter shap_explainer_factory can be used to specify the type of shap explainer, for the different options see here.

Source code in metalearners/explainer.py

def shap_values(
    self,
    X: Matrix,
    shap_explainer_factory: type[shap.Explainer],
    shap_explainer_params: dict | None = None,
) -> list[np.ndarray]:
    """Calculates the shap values for each treatment group.

    The parameter ``shap_explainer_factory`` can be used to specify the type of shap
    explainer, for the different options see
    [here](https://shap.readthedocs.io/en/latest/api.html#explainers).
    """
    if shap_explainer_params is None:
        shap_explainer_params = {}
    shap_values = []
    for tv in range(self.n_variants - 1):
        shap_explainer = shap_explainer_factory(
            model=self.cate_models[tv],
            **shap_explainer_params,
        )
        variant_shap_values = shap_explainer.shap_values(X)
        shap_values.append(variant_shap_values)
    return shap_values

grid_search

GSResult dataclass

GSResult(
    metalearner: MetaLearner,
    train_scores: dict,
    test_scores: dict | None,
    fit_time: float,
    score_time: float,
)

Result from a single grid search evaluation.

MetaLearnerGridSearch

MetaLearnerGridSearch(
    metalearner_factory: type[MetaLearner],
    metalearner_params: Mapping[str, Any],
    base_learner_grid: Mapping[
        str, Sequence[type[_ScikitModel]]
    ],
    param_grid: Mapping[
        str, Mapping[str, Mapping[str, Sequence]]
    ],
    scoring: Scoring | None = None,
    n_jobs: int | None = None,
    random_state: int | None = None,
    verbose: int = 0,
    store_raw_results: bool = True,
    store_results: bool = True,
)

Exhaustive search over specified parameter values for a MetaLearner.

metalearner_params should contain the necessary params for the MetaLearner initialization such as n_variants and is_classification. If one wants to pass optional parameters to the MetaLearner initialization, such as n_folds or feature_set, this should be done by this way, too. Importantly, random_state must be passed through the random_state parameter and not through metalearner_params.

base_learner_grid keys should be the names of the needed base models contained in the MetaLearner defined by metalearner_factory, for information about this names check MetaLearner.nuisance_model_specifications and MetaLearner.treatment_model_specifications. The values should be sequences of model factories.

If base models are meant to be reused, they should be passed through metalearner_params and the corresponding keys should not be passed to base_learner_grid.

param_grid should contain the parameters grid for each type of model used by the base learners defined in base_learner_grid. The keys should be strings with the model class name. An example for optimizing over the :class:metalearners.DRLearner would be:

base_learner_grid = {
    "propensity_model": (LGBMClassifier, LogisticRegression),
    "variant_outcome_model": (LGBMRegressor, LinearRegression),
    "treatment_model": (LGBMRegressor)
}

param_grid = {
    "propensity_model": {
        "LGBMClassifier": {"n_estimators": [1, 2, 3], "verbose": [-1]}
    },
    "variant_outcome_model": {
        "LGBMRegressor": {"n_estimators": [1, 2], "verbose": [-1]},
    },
    "treatment_model": {
        "LGBMRegressor": {"n_estimators": [5, 10], "verbose": [-1]},
    },
}

If some model is not present in param_grid, the default parameters will be used.

For information on how to define scoring see MetaLearner.evaluate.

verbose will be passed to joblib.Parallel.

store_raw_results and store_results define which and how the results are saved after calling MetaLearnerGridSearch.fit depending on their values:

• Both are True (default): raw_results_ will be a list of GSResult with all the results and results_ will be a DataFrame with the processed results.
• store_raw_results=True and store_results=False: raw_results_ will be a list of GSResult with all the results and results will be None.
• store_raw_results=False and store_results=True: raw_results_ will be None and results_ will be a DataFrame with the processed results.
• Both are False: raw_results_ will be a generator which yields a GSResult for each configuration and results will be None. This configuration can be useful in the case the grid search is big and you do not want to store all MetaLearners objects rather evaluate them after fitting each one and just store one.

grid_size_ will contain the number of hyperparameter combinations after fitting. This attribute may be useful in the case store_raw_results = False and store_results = False. In that case, the generator object returned in raw_results_ doesn’t trigger the fitting of individual metalearners until explicitly requested, e.g. in a loop. This attribute can be use to track the progress, for instance, by creating a progress bar or a similar utility.

For an illustration see our example on Tuning hyperparameters of a MetaLearner with MetaLearnerGridSearch.

Source code in metalearners/grid_search.py

def __init__(
    self,
    metalearner_factory: type[MetaLearner],
    metalearner_params: Mapping[str, Any],
    base_learner_grid: Mapping[str, Sequence[type[_ScikitModel]]],
    param_grid: Mapping[str, Mapping[str, Mapping[str, Sequence]]],
    scoring: Scoring | None = None,
    n_jobs: int | None = None,
    random_state: int | None = None,
    verbose: int = 0,
    store_raw_results: bool = True,
    store_results: bool = True,
):
    self.metalearner_factory = metalearner_factory
    self.metalearner_params = metalearner_params
    self.scoring = scoring
    self.n_jobs = n_jobs
    self.random_state = random_state
    self.verbose = verbose
    self.store_raw_results = store_raw_results
    self.store_results = store_results

    all_base_models = set(
        metalearner_factory.nuisance_model_specifications().keys()
    ) | set(metalearner_factory.treatment_model_specifications().keys())

    self.fitted_models = set(
        metalearner_params.get("fitted_nuisance_models", {}).keys()
    )
    if metalearner_params.get("fitted_propensity_model", None) is not None:
        self.fitted_models |= {PROPENSITY_MODEL}

    self.models_to_fit = all_base_models - self.fitted_models

    if set(base_learner_grid.keys()) != self.models_to_fit:
        raise ValueError(
            "base_learner_grid keys don't match the expected model names. base_learner_grid "
            f"keys were expected to be {self.models_to_fit}."
        )
    self.base_learner_grid = base_learner_grid
    self.param_grid = param_grid

fit

fit(
    X: Matrix,
    y: Vector,
    w: Vector,
    X_test: Matrix | None = None,
    y_test: Vector | None = None,
    w_test: Vector | None = None,
    oos_method: OosMethod = OVERALL,
    **kwargs,
)

Run fit with all sets of parameters.

X_test, y_test and w_test are optional, in case they are passed all the fitted metalearners will be evaluated on it.

kwargs will be passed through to the MetaLearner.fit call of each individual MetaLearner.

Source code in metalearners/grid_search.py

def fit(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    X_test: Matrix | None = None,
    y_test: Vector | None = None,
    w_test: Vector | None = None,
    oos_method: OosMethod = OVERALL,
    **kwargs,
):
    """Run fit with all sets of parameters.

    ``X_test``, ``y_test`` and ``w_test`` are optional, in case they are passed all the
    fitted metalearners will be evaluated on it.

    ``kwargs`` will be passed through to the [`MetaLearner.fit`][metalearners.metalearner.MetaLearner.fit]
    call of each individual MetaLearner.
    """
    nuisance_models_wo_propensity = (
        set(self.metalearner_factory.nuisance_model_specifications().keys())
        - {PROPENSITY_MODEL}
    ) & self.models_to_fit

    # We don't need to intersect as treatment models can't be reused
    treatment_models = set(
        self.metalearner_factory.treatment_model_specifications().keys()
    )

    jobs: list[_FitAndScoreJob] = []

    for base_learners in ParameterGrid(self.base_learner_grid):
        nuisance_model_factory = {
            model_kind: base_learners[model_kind]
            for model_kind in nuisance_models_wo_propensity
        }
        treatment_model_factory = {
            model_kind: base_learners[model_kind] for model_kind in treatment_models
        }
        propensity_model_factory = base_learners.get(PROPENSITY_MODEL, None)
        base_learner_param_grids = {
            model_kind: list(
                ParameterGrid(
                    self.param_grid.get(model_kind, {}).get(
                        base_learners[model_kind].__name__, {}
                    )
                )
            )
            for model_kind in self.models_to_fit
        }
        for params in ParameterGrid(base_learner_param_grids):
            nuisance_model_params = {
                model_kind: params[model_kind]
                for model_kind in nuisance_models_wo_propensity
            }
            treatment_model_params = {
                model_kind: params[model_kind] for model_kind in treatment_models
            }
            propensity_model_params = params.get(PROPENSITY_MODEL, None)

            grid_metalearner_params = {
                "nuisance_model_factory": nuisance_model_factory,
                "treatment_model_factory": treatment_model_factory,
                "propensity_model_factory": propensity_model_factory,
                "nuisance_model_params": nuisance_model_params,
                "treatment_model_params": treatment_model_params,
                "propensity_model_params": propensity_model_params,
                "random_state": self.random_state,
            }

            if (
                len(
                    shared_keys := set(grid_metalearner_params.keys())
                    & set(self.metalearner_params.keys())
                )
                > 0
            ):
                raise ValueError(
                    f"{shared_keys} should not be specified in metalearner_params as "
                    "they are used internally. Please use the correct parameters."
                )

            jobs.append(
                _FitAndScoreJob(
                    metalearner_factory=self.metalearner_factory,
                    metalearner_params=dict(self.metalearner_params)
                    | grid_metalearner_params,
                    X_train=X,
                    y_train=y,
                    w_train=w,
                    X_test=X_test,
                    y_test=y_test,
                    w_test=w_test,
                    oos_method=oos_method,
                    scoring=self.scoring,
                    metalerner_fit_params=kwargs,
                )
            )

    self.grid_size_ = len(jobs)
    self.raw_results_: list[GSResult] | Generator[GSResult, None, None] | None
    self.results_: pd.DataFrame | None = None

    return_as = "list" if self.store_raw_results else "generator_unordered"
    parallel = Parallel(
        n_jobs=self.n_jobs, verbose=self.verbose, return_as=return_as
    )
    self.raw_results_ = parallel(delayed(_fit_and_score)(job) for job in jobs)
    if self.store_results:
        self.results_ = _format_results(results=self.raw_results_)  # type: ignore
        if not self.store_raw_results:
            # The generator will be empty so we replace it with None
            self.raw_results_ = None

metalearner

MetaLearner

MetaLearner(
    is_classification: bool,
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel]
    | None = None,
    nuisance_model_params: Params
    | dict[str, Params]
    | None = None,
    treatment_model_params: Params
    | dict[str, Params]
    | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[
        str, list[CrossFitEstimator]
    ]
    | None = None,
    fitted_propensity_model: CrossFitEstimator
    | None = None,
    feature_set: Features
    | dict[str, Features]
    | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
)

Bases: ABC

MetaLearner abstract class. All metalearner implementations should inherit from it.

All of

• nuisance_model_factory
• treatment_model_factory
• nuisance_model_params
• treatment_model_params
• feature_set

can either

• contain a single value, such that the value will be used for all relevant models of the respective MetaLearner or
• a dictionary mapping from the relevant models (model_kind, a str) to the respective value; at least all relevant models need to be present, more are allowed and ignored

The possible values for defining feature_set (either one single value for all the models or the values inside the dictionary specifying for each model) can be:

• None: All columns will be used.
• A list of strings or integers indicating which columns to use.
• [] meaning that no present column should be used for that model and the input of the model should be a vector of 1s.

To reuse already fitted models fitted_nuisance_models and fitted_propensity_model should be used. The models should be fitted on the same data the MetaLearner is going to call fit with. For an illustration, see our example on reusing models.

Source code in metalearners/metalearner.py

def __init__(
    self,
    is_classification: bool,
    # TODO: Consider whether we can make this not a state of the MetaLearner
    # but rather just a parameter of a predict call.
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel] | None = None,
    nuisance_model_params: Params | dict[str, Params] | None = None,
    treatment_model_params: Params | dict[str, Params] | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[str, list[CrossFitEstimator]] | None = None,
    fitted_propensity_model: CrossFitEstimator | None = None,
    feature_set: Features | dict[str, Features] | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
):
    nuisance_model_specifications = self.nuisance_model_specifications()
    treatment_model_specifications = self.treatment_model_specifications()

    if PROPENSITY_MODEL in treatment_model_specifications:
        raise ValueError(
            f"{PROPENSITY_MODEL} can't be used as a treatment model name"
        )
    if (
        isinstance(nuisance_model_factory, dict)
        and PROPENSITY_MODEL in nuisance_model_factory.keys()
    ):
        raise ValueError(
            "Propensity model factory should be defined using propensity_model_factory "
            "and not nuisance_model_factory."
        )
    if (
        isinstance(nuisance_model_params, dict)
        and PROPENSITY_MODEL in nuisance_model_params.keys()
    ):
        raise ValueError(
            "Propensity model params should be defined using propensity_model_params "
            "and not nuisance_model_params."
        )
    if (
        PROPENSITY_MODEL in nuisance_model_specifications
        and propensity_model_factory is None
        and fitted_propensity_model is None
    ):
        raise ValueError(
            "propensity_model_factory or fitted_propensity_model needs to be defined "
            f"as the {self.__class__.__name__} has a propensity model."
        )

    self._validate_n_variants(n_variants)
    self.is_classification = is_classification
    self.n_variants = n_variants

    self.nuisance_model_factory = _combine_propensity_and_nuisance_specs(
        propensity_model_factory,
        nuisance_model_factory,
        set(nuisance_model_specifications.keys()),
    )
    if nuisance_model_params is None:
        nuisance_model_params = {}  # type: ignore
    if propensity_model_params is None:
        propensity_model_params = {}
    self.nuisance_model_params = _combine_propensity_and_nuisance_specs(
        propensity_model_params,
        nuisance_model_params,
        set(nuisance_model_specifications.keys()),
    )

    self.treatment_model_factory = _initialize_model_dict(
        treatment_model_factory, set(treatment_model_specifications.keys())
    )
    if treatment_model_params is None:
        self.treatment_model_params = _initialize_model_dict(
            {}, set(treatment_model_specifications.keys())
        )
    else:
        self.treatment_model_params = _initialize_model_dict(
            treatment_model_params, set(treatment_model_specifications.keys())
        )

    self.n_folds = _initialize_model_dict(
        n_folds,
        set(nuisance_model_specifications.keys())
        | set(treatment_model_specifications.keys()),
    )
    for model_kind, n_folds_model_kind in self.n_folds.items():
        validate_number_positive(n_folds_model_kind, f"{model_kind} n_folds", True)
    self.random_state = random_state

    self.feature_set = _initialize_model_dict(
        feature_set,
        set(nuisance_model_specifications.keys())
        | set(treatment_model_specifications.keys()),
    )

    self._nuisance_models: dict[str, list[CrossFitEstimator]] = {}
    not_fitted_nuisance_models = set(nuisance_model_specifications.keys())
    self._prefitted_nuisance_models: set[str] = set()

    if fitted_nuisance_models is not None:
        if not set(fitted_nuisance_models.keys()) <= set(
            nuisance_model_specifications.keys()
        ) - {PROPENSITY_MODEL}:
            raise ValueError(
                "The keys present in fitted_nuisance_models should be a subset of "
                f"{set(nuisance_model_specifications.keys()) - {PROPENSITY_MODEL}}"
            )
        self._nuisance_models |= deepcopy(fitted_nuisance_models)
        not_fitted_nuisance_models -= set(fitted_nuisance_models.keys())
        self._prefitted_nuisance_models |= set(fitted_nuisance_models.keys())

    if (
        PROPENSITY_MODEL in nuisance_model_specifications.keys()
        and fitted_propensity_model is not None
    ):
        self._nuisance_models |= {PROPENSITY_MODEL: [fitted_propensity_model]}
        not_fitted_nuisance_models -= {PROPENSITY_MODEL}
        self._prefitted_nuisance_models |= {PROPENSITY_MODEL}

    for name in not_fitted_nuisance_models:
        if self.nuisance_model_factory[name] is None:
            if name == PROPENSITY_MODEL:
                raise ValueError(
                    f"A model for the nuisance model {name} needs to be defined. Either "
                    "in propensity_model_factory or in fitted_propensity_model."
                )
            else:
                raise ValueError(
                    f"A model for the nuisance model {name} needs to be defined. Either "
                    "in nuisance_model_factory or in fitted_nuisance_models."
                )

    self._nuisance_models |= {
        name: [
            CrossFitEstimator(
                n_folds=self.n_folds[name],
                estimator_factory=self.nuisance_model_factory[name],
                estimator_params=self.nuisance_model_params[name],
                random_state=self.random_state,
            )
            for _ in range(nuisance_model_specifications[name]["cardinality"](self))
        ]
        for name in not_fitted_nuisance_models
    }
    self._treatment_models: dict[str, list[CrossFitEstimator]] = {
        name: [
            CrossFitEstimator(
                n_folds=self.n_folds[name],
                estimator_factory=self.treatment_model_factory[name],
                estimator_params=self.treatment_model_params[name],
                random_state=self.random_state,
            )
            for _ in range(
                treatment_model_specifications[name]["cardinality"](self)
            )
        ]
        for name in set(treatment_model_specifications.keys())
    }

    self._validate_models()

init_args property

init_args: dict[str, Any]

Create initialization parameters for a new MetaLearner.

Importantly, this does not copy further internal state, such as the weights or parameters of trained base models.

evaluate abstractmethod

evaluate(
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]

Evaluate the MetaLearner.

The keys in scoring which are not a name of a model contained in the MetaLearner will be ignored, for information about this names check nuisance_model_specifications and treatment_model_specifications. The values must be a list of:

• string representing a sklearn scoring method. Check here for the possible values.
• Callable with signature scorer(estimator, X, y_true, **kwargs). We recommend using sklearn.metrics.make_scorer to create such a Callable.

If some model name is not present in the keys of scoring then the default used metrics will be neg_log_loss if it is a classifier and neg_root_mean_squared_error if it is a regressor.

The returned dictionary keys have the following structure:

• For nuisance models:

  • If the cardinality is one: f"{model_kind}_{scorer}"
  • If there is one model for each treatment variant (including control): f"{model_kind}_{treatment_variant}_{scorer}"

• For treatment models: f"{model_kind}_{treatment_variant}_vs_0_{scorer}"

Where scorer is the name of the scorer if it is a string and "custom_scorer_{idx}" if it is a callable where idx is the index in the scorers list.

Source code in metalearners/metalearner.py

@abstractmethod
def evaluate(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]:
    r"""Evaluate the MetaLearner.

    The keys in `scoring` which are not a name of a model contained in the MetaLearner
    will be ignored, for information about this names check
    [`nuisance_model_specifications`][metalearners.metalearner.MetaLearner.nuisance_model_specifications] and
    [`treatment_model_specifications`][metalearners.metalearner.MetaLearner.treatment_model_specifications].
    The values must be a list of:

    * `string` representing a `sklearn` scoring method. Check
      [here](https://scikit-learn.org/stable/modules/model_evaluation.html#common-cases-predefined-values)
      for the possible values.
    * `Callable` with signature `scorer(estimator, X, y_true, **kwargs)`. We recommend
      using [`sklearn.metrics.make_scorer`](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.make_scorer.html)
      to create such a `Callable`.

    If some model name is not present in the keys of `scoring` then the default used
    metrics will be `neg_log_loss` if it is a classifier and `neg_root_mean_squared_error`
    if it is a regressor.

    The returned dictionary keys have the following structure:

    * For nuisance models:

        * If the cardinality is one:  `f"{model_kind}_{scorer}"`
        * If there is one model for each treatment variant (including control):
          `f"{model_kind}_{treatment_variant}_{scorer}"`

    * For treatment models: `f"{model_kind}_{treatment_variant}_vs_0_{scorer}"`

    Where `scorer` is the name of the scorer if it is a string and `"custom_scorer_{idx}"`
    if it is a callable where `idx` is the index in the `scorers` list.
    """
    ...

explainer

explainer(
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer

Create an Explainer

which can be used in feature_importances.

This function can be used in two distinct manners based on the provided parameters:

• When parameters X, cate_estimates, and cate_model_factory are all set to None, the function creates an Explainer using the pre-existing treatment models. If these models do not exist, however, it triggers a ValueError.
• On the contrary, if X, cate_estimates, and cate_model_factory are not None, the function initiates an instance of the Explainer class using these parameters. This instance then fits new models for each treatment variant, and these models are employed to calculate the importance of features.

Source code in metalearners/metalearner.py

def explainer(
    self,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer:
    r"""Create an [`Explainer`][metalearners.explainer.Explainer]

    which can be used in [`feature_importances`][metalearners.metalearner.MetaLearner.feature_importances].

    This function can be used in two distinct manners based on the provided parameters:

    *   When parameters `X`, `cate_estimates`, and `cate_model_factory` are all
        set to `None`, the function creates an [`Explainer`][metalearners.explainer.Explainer]
        using the pre-existing treatment models. If these models do not exist, however,
        it triggers a `ValueError`.
    *   On the contrary, if `X`, `cate_estimates`, and `cate_model_factory` are
        not `None`, the function initiates an instance of the [`Explainer`][metalearners.explainer.Explainer]
        class using these parameters. This instance then fits new models for each
        treatment variant, and these models are employed to calculate the importance
        of features.
    """
    if X is None and cate_estimates is None and cate_model_factory is None:
        try:
            cate_cf_models = self._treatment_models[TREATMENT_MODEL]
            cate_models = [cf._overall_estimator for cf in cate_cf_models]
        except KeyError:
            raise ValueError(
                "The metalearner does not have treatment models; hence X, cate_estimates, "
                "and cate_model_factory need to be defined."
            )
        return Explainer(cate_models=cate_models)  # type: ignore
    elif (
        X is not None
        and cate_estimates is not None
        and cate_model_factory is not None
    ):
        return Explainer.from_estimates(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    else:
        raise ValueError(
            "Either all of [X, cate_estimates, cate_model_factory] are None or all "
            "of them must be defined."
        )

feature_importances

feature_importances(
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances

Calculates the feature importance for each treatment group.

If explainer is None, a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters X, cate_estimates, cate_model_factory and cate_model_params are ignored.

If normalization = True, for each treatment variant the feature importances are normalized so that they sum to 1.

feature_names is optional but in the case it’s not passed the names of the features will default to f"Feature {i}" where i is the corresponding feature index.

The returned list contains the feature importances for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def feature_importances(
    self,
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances:
    r"""Calculates the feature importance for each treatment group.

    If `explainer` is `None`, a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `X`, `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    If `normalization = True`, for each treatment variant the feature importances
    are normalized so that they sum to 1.

    `feature_names` is optional but in the case it's not passed the names of the
    features will default to `f"Feature {i}"` where `i` is the corresponding
    feature index.

    The returned list contains the feature importances for each treatment variant in
    ascending order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.feature_importances(
        normalize=normalize, feature_names=feature_names, sort_values=sort_values
    )

fit

fit(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

n_jobs_cross_fitting will be used at the cross-fitting level and n_jobs_base_learners will be used at the stage level. None means 1 unless in a joblib.parallel_backend context. -1 means using all processors. For more information about parallelism check parallelism.

fit_params is an optional dict to be forwarded to base estimator fit calls. It supports two usages patterns:

fit_params={"parameter_of_interest": value_of_interest}

fit_params={
    "nuisance": {
        "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
        "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
    },
    "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
}

In the former approach, the parameter and value of interest are passed to all base models. In the the latter approach, the explicitly qualified parameter-value pairs are passed to respective base models and no fitting parameters are passed to base models not explicitly listed. Note that in this pattern, propensity models are considered a nuisance model.

synchronize_cross_fitting indicates whether the learning of different base models should use exactly the same data splits where possible. Note that if there are several models to be synchronized which are classifiers, these cannot be split via stratification.

Source code in metalearners/metalearner.py

def fit(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    """Fit all models of the MetaLearner.

    If pre-fitted models were passed at instantiation, these are never refitted.

    `n_jobs_cross_fitting` will be used at the cross-fitting level and
    `n_jobs_base_learners` will be used at the stage level. `None` means 1 unless in a
    [`joblib.parallel_backend`](https://joblib.readthedocs.io/en/latest/generated/joblib.parallel_backend.html#joblib.parallel_backend)
    context. `-1` means using all processors.
    For more information about parallelism check [parallelism](parallelism.md).


    `fit_params` is an optional `dict` to be forwarded to base estimator `fit` calls. It supports
    two usages patterns:

    ```python
    fit_params={"parameter_of_interest": value_of_interest}
    ```

    ```python
    fit_params={
        "nuisance": {
            "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
            "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
        },
        "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
    }
    ```

    In the former approach, the parameter and value of interest are passed to all base models. In the
    the latter approach, the explicitly qualified parameter-value pairs are passed to respective base
    models and no fitting parameters are passed to base models not explicitly listed. Note that in this
    pattern, propensity models are considered a nuisance model.

    `synchronize_cross_fitting` indicates whether the learning of different base models should use exactly
    the same data splits where possible. Note that if there are several models to be synchronized which are
    classifiers, these cannot be split via stratification.
    """
    self.fit_all_nuisance(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    self.fit_all_treatment(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    return self

fit_all_nuisance abstractmethod

fit_all_nuisance(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all nuisance models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the nuisance models, and in the case they should also be used by the treatment models, these should also be passed in the following call to fit_all_treatment.

This method, combined with fit_all_treatment, facilitates the segmentation of the metalearner fitting process into two distinct parts. This division allows for interventions between the two stages, such as performing feature selection for the treatment models or conducting hyperparameter optimization within the nuisance models.

Source code in metalearners/metalearner.py

@abstractmethod
def fit_all_nuisance(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    """Fit all nuisance models of the MetaLearner.

    If pre-fitted models were passed at instantiation, these are never refitted.

    The only difference with [`fit`][metalearners.metalearner.MetaLearner.fit] parameters,
    is that if `fit_params` follows the first usage pattern (explained in
    [`fit`][metalearners.metalearner.MetaLearner.fit]), then the training parameters
    will only be used for the nuisance models, and in the case they should also be used
    by the treatment models, these should also be passed in the following call to
    [`fit_all_treatment`][metalearners.metalearner.MetaLearner.fit_all_treatment].

    This method, combined with [`fit_all_treatment`][metalearners.metalearner.MetaLearner.fit_all_treatment],
    facilitates the segmentation of the metalearner fitting process into two distinct parts.
    This division allows for interventions between the two stages, such as performing
    feature selection for the treatment models or conducting hyperparameter optimization
    within the nuisance models.
    """
    ...

fit_all_treatment abstractmethod

fit_all_treatment(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all treatment models of the MetaLearner.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the treatment models, as the nuisance models should already be fitted.

Source code in metalearners/metalearner.py

@abstractmethod
def fit_all_treatment(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    """Fit all treatment models of the MetaLearner.

    The only difference with [`fit`][metalearners.metalearner.MetaLearner.fit] parameters,
    is that if `fit_params` follows the first usage pattern (explained in
    [`fit`][metalearners.metalearner.MetaLearner.fit]), then the training parameters
    will only be used for the treatment models, as the nuisance models should already
    be fitted.
    """
    ...

fit_nuisance

fit_nuisance(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit a given nuisance model of a MetaLearner.

y represents the objective of the given nuisance model, not necessarily the outcome of the experiment. If pre-fitted models were passed at instantiation, these are never refitted.

Source code in metalearners/metalearner.py

def fit_nuisance(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit a given nuisance model of a MetaLearner.

    `y` represents the objective of the given nuisance model, not necessarily the outcome of the experiment.
    If pre-fitted models were passed at instantiation, these are never refitted.
    """
    if model_kind in self._prefitted_nuisance_models:
        return self
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._nuisance_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

fit_treatment

fit_treatment(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit the treatment model of a MetaLearner.

y represents the objective of the given treatment model, not necessarily the outcome of the experiment.

Source code in metalearners/metalearner.py

def fit_treatment(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit the treatment model of a MetaLearner.

    `y` represents the objective of the given treatment model, not necessarily the outcome of the experiment.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._treatment_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

nuisance_model_specifications abstractmethod classmethod

nuisance_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all first-stage models.

Source code in metalearners/metalearner.py

@classmethod
@abstractmethod
def nuisance_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    """Return the specifications of all first-stage models."""
    ...

predict abstractmethod

predict(
    X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Estimate the CATE.

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants} - 1, 1)\) if the outcome is a scalar, i.e. in case of a regression problem.

• \((n_{obs}, n_{variants} - 1, n_{classes})\) if the outcome is a class, i.e. in case of a classification problem.

In the case of multiple treatment variants, the second dimension represents the CATE of the corresponding variant vs the control (variant 0).

Source code in metalearners/metalearner.py

@abstractmethod
def predict(
    self,
    X: Matrix,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate the CATE.

    If `is_oos`, an acronym for 'is out of sample', is `False`,
    the estimates will stem from cross-fitting. Otherwise,
    various approaches exist, specified via `oos_method`.

    The returned ndarray is of shape:

    * $(n_{obs}, n_{variants} - 1, 1)$ if the outcome is a scalar, i.e. in case
      of a regression problem.

    * $(n_{obs}, n_{variants} - 1, n_{classes})$ if the outcome is a class,
      i.e. in case of a classification problem.

    In the case of multiple treatment variants, the second dimension represents the
    CATE of the corresponding variant vs the control (variant 0).
    """
    ...

predict_conditional_average_outcomes abstractmethod

predict_conditional_average_outcomes(
    X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Predict the vectors of conditional average outcomes.

These are defined as \(\mathbb{E}[Y_i(w) | X]\) for each treatment variant \(w\).

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants}, 1)\) if the outcome is a scalar, i.e., in case of a regression problem.

• \((n_{obs}, n_{variants}, n_{classes})\) if the outcome is a class, i.e., in case of a classification problem.

Source code in metalearners/metalearner.py

@abstractmethod
def predict_conditional_average_outcomes(
    self, X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray:
    r"""Predict the vectors of conditional average outcomes.

    These are defined as $\mathbb{E}[Y_i(w) | X]$ for each treatment variant $w$.

    If `is_oos`, an acronym for 'is out of sample', is `False`,
    the estimates will stem from cross-fitting. Otherwise,
    various approaches exist, specified via `oos_method`.

    The returned ndarray is of shape:

    * $(n_{obs}, n_{variants}, 1)$ if the outcome is a scalar, i.e., in case
    of a regression problem.

    * $(n_{obs}, n_{variants}, n_{classes})$ if the outcome is a class,
    i.e., in case of a classification problem.
    """
    ...

predict_nuisance

predict_nuisance(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given nuisance model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_nuisance(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given nuisance model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    predict_method_name = self.nuisance_model_specifications()[model_kind][
        "predict_method"
    ](self)
    predict_method = getattr(
        self._nuisance_models[model_kind][model_ord], predict_method_name
    )
    return predict_method(X_filtered, is_oos, oos_method)

predict_treatment

predict_treatment(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given treatment model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_treatment(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given treatment model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    return self._treatment_models[model_kind][model_ord].predict(
        X_filtered, is_oos, oos_method
    )

shap_values

shap_values(
    X: Matrix,
    shap_explainer_factory: type[Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]

Calculates the shap values for each treatment group.

If explainer is None a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters cate_estimates, cate_model_factory and cate_model_params are ignored.

The parameter shap_explainer_factory can be used to specify the type of shap explainer, for the different options see here.

The returned list contains the shap values for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def shap_values(
    self,
    X: Matrix,
    shap_explainer_factory: type[shap.Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]:
    """Calculates the shap values for each treatment group.

    If `explainer` is `None` a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    The parameter `shap_explainer_factory` can be used to specify the type of shap
    explainer, for the different options see
    [here](https://shap.readthedocs.io/en/latest/api.html#explainers).

    The returned list contains the shap values for each treatment variant in ascending
    order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=None if cate_estimates is None else X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.shap_values(
        X=X,
        shap_explainer_factory=shap_explainer_factory,
        shap_explainer_params=shap_explainer_params,
    )

treatment_model_specifications abstractmethod classmethod

treatment_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all second-stage models.

Source code in metalearners/metalearner.py

@classmethod
@abstractmethod
def treatment_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    """Return the specifications of all second-stage models."""
    ...

outcome_functions

constant_treatment_effect

constant_treatment_effect(
    dim: int,
    tau: float | ndarray,
    ulow: float = 0,
    uhigh: float = 1,
    rng: Generator | None = None,
) -> Callable

Generate a potential outcomes function with constant treatment effect.

\[ f(x_i, w_i) = x_i' \beta_{control} + \sum_{k=1}^{n_v-1} \tau_k \cdot \mathcal{I}(\{w_i = k\}) \]

where :math:x_i is a vector of features, :math:\tau a vector of treatment effects, :math:w_i the treatment indicator, :math:n_v the number of variants and

.. math:: \beta_{control} \sim \mathcal{U}[u_l, u_h]

dim indicates the dimension of :math:\beta and therefore it should be the number of numerical features plus the number of categories in all of the categorical features.

tau expects to be of size :math:n_v-1.

Source code in metalearners/outcome_functions.py

def constant_treatment_effect(
    dim: int,
    tau: float | np.ndarray,
    ulow: float = 0,
    uhigh: float = 1,
    rng: np.random.Generator | None = None,
) -> Callable:
    r"""Generate a potential outcomes function with constant treatment effect.

    $$
    f(x_i, w_i) = x_i' \beta_{control} + \sum_{k=1}^{n_v-1} \tau_k \cdot \mathcal{I}(\{w_i = k\})
    $$

    where :math:`x_i` is a vector of features, :math:`\tau` a vector of treatment effects,
    :math:`w_i` the treatment indicator, :math:`n_v` the number of variants and

    .. math::
        \beta_{control} \sim \mathcal{U}[u_l, u_h]

    ``dim`` indicates the dimension of :math:`\beta` and therefore it should be the
    number of numerical features plus the number of categories in all of the categorical
    features.

    ``tau`` expects to be of size :math:`n_v-1`.
    """
    if rng is None:
        rng = default_rng

    beta = _beta(ulow, uhigh, dim, rng)
    if isinstance(tau, int | float):
        tau = np.array([tau])
    tau = tau.reshape(1, -1)

    def f(X: Matrix) -> np.ndarray:
        ohe_encoded_features = (
            pd.get_dummies(X, dtype="float64").to_numpy()
            if isinstance(X, pd.DataFrame)
            else X
        )

        mu_0 = np.dot(ohe_encoded_features, beta)
        return np.c_[mu_0, mu_0.reshape(-1, 1) + tau]

    return f

linear_treatment_effect

linear_treatment_effect(
    dim: int,
    n_variants: int = 2,
    ulow: float = 0,
    uhigh: float = 1,
    rng: Generator | None = None,
) -> Callable

Generate a potential outcomes function with linear treatment effect.

.. math:: f(x_i, w_i) = x_i’ \beta_{control} + \sum_{k=1}^{n_v-1} \mathcal{I}({w_i = k}) \cdot x_i’ \beta^{(k)}

where :math:x_i is a vector of features, :math:w_i the treatment indicator, and

.. math:: \beta_{control} \sim \mathcal{U}[u_l, u_h] \beta_{(k)} \sim \mathcal{U}[u_l, u_h]

dim indicates the dimension of :math:\beta_{control} and :math:\beta_{(k)} therefore it should be the number of numerical features plus the number of categories in all of the categorical features.

Source code in metalearners/outcome_functions.py

def linear_treatment_effect(
    dim: int,
    n_variants: int = 2,
    ulow: float = 0,
    uhigh: float = 1,
    rng: np.random.Generator | None = None,
) -> Callable:
    r"""Generate a potential outcomes function with linear treatment effect.

    .. math::
        f(x_i, w_i) = x_i' \beta_{control} + \sum_{k=1}^{n_v-1} \mathcal{I}(\{w_i = k\}) \cdot x_i' \beta^{(k)}

    where :math:`x_i` is a vector of features, :math:`w_i` the treatment indicator, and

    .. math::
        \beta_{control} \sim \mathcal{U}[u_l, u_h]
        \beta_{(k)} \sim \mathcal{U}[u_l, u_h]

    ``dim`` indicates the dimension of :math:`\beta_{control}` and :math:`\beta_{(k)}`
    therefore it should be the number of numerical features plus the number of categories
    in all of the categorical features.
    """
    if n_variants < 2:
        raise ValueError("n_variants needs to be an integer greater or equal to 2")

    if rng is None:
        rng = default_rng

    beta_control = _beta(ulow, uhigh, dim, rng)
    beta = _beta(ulow, uhigh, (dim, n_variants - 1), rng)

    def f(X: Matrix) -> np.ndarray:
        ohe_encoded_features = (
            pd.get_dummies(X, dtype="float64").to_numpy()
            if isinstance(X, pd.DataFrame)
            else X
        )

        mu_0 = np.dot(ohe_encoded_features, beta_control)
        return np.c_[mu_0, mu_0.reshape(-1, 1) + np.matmul(ohe_encoded_features, beta)]

    return f

no_treatment_effect

no_treatment_effect(
    dim: int,
    n_variants: int = 2,
    ulow: float = 0,
    uhigh: float = 1,
    rng: Generator | None = None,
) -> Callable

Generate a potential outcomes function with no treatment effect.

.. math:: f(x_i, w_i) = x_i’ \beta

where :math:x_i is a vector of features and

.. math:: \beta \sim \mathcal{U}[u_l, u_h]

dim indicates the dimension of :math:\beta and therefore the number of numerical features plus the number of categories in all of the categorical features.

Source code in metalearners/outcome_functions.py

def no_treatment_effect(
    dim: int,
    n_variants: int = 2,
    ulow: float = 0,
    uhigh: float = 1,
    rng: np.random.Generator | None = None,
) -> Callable:
    r"""Generate a potential outcomes function with no treatment effect.

    .. math::
        f(x_i, w_i) = x_i' \beta

    where :math:`x_i` is a vector of features and

    .. math::
        \beta \sim \mathcal{U}[u_l, u_h]

    ``dim`` indicates the dimension of :math:`\beta` and therefore the number of
    numerical features plus the number of categories in all of the categorical features.
    """
    if n_variants < 2:
        raise ValueError("n_variants needs to be an integer greater or equal to 2")

    tau = np.broadcast_to(0, n_variants - 1)

    return constant_treatment_effect(dim, tau=tau, ulow=ulow, uhigh=uhigh, rng=rng)

rlearner

RLearner

RLearner(
    is_classification: bool,
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel]
    | None = None,
    nuisance_model_params: Params
    | dict[str, Params]
    | None = None,
    treatment_model_params: Params
    | dict[str, Params]
    | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[
        str, list[CrossFitEstimator]
    ]
    | None = None,
    fitted_propensity_model: CrossFitEstimator
    | None = None,
    feature_set: Features
    | dict[str, Features]
    | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
)

Bases: MetaLearner

R-Learner for CATE estimation as described by Nie et al. (2017).

Importantly, the current R-Learner implementation only supports:

• binary classes in case of a classification outcome

The R-Learner contains two nuisance models

• a "propensity_model" estimating \(\Pr[W=k|X]\)
• an "outcome_model" estimating \(\mathbb{E}[Y|X]\)

and one treatment model per treatment variant which isn’t control

• "treatment_model" which estimates \(\mathbb{E}[Y(k) - Y(0) | X]\)

The treatment_model_factory provided needs to support the argument sample_weight in its fit method.

Source code in metalearners/metalearner.py

def __init__(
    self,
    is_classification: bool,
    # TODO: Consider whether we can make this not a state of the MetaLearner
    # but rather just a parameter of a predict call.
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel] | None = None,
    nuisance_model_params: Params | dict[str, Params] | None = None,
    treatment_model_params: Params | dict[str, Params] | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[str, list[CrossFitEstimator]] | None = None,
    fitted_propensity_model: CrossFitEstimator | None = None,
    feature_set: Features | dict[str, Features] | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
):
    nuisance_model_specifications = self.nuisance_model_specifications()
    treatment_model_specifications = self.treatment_model_specifications()

    if PROPENSITY_MODEL in treatment_model_specifications:
        raise ValueError(
            f"{PROPENSITY_MODEL} can't be used as a treatment model name"
        )
    if (
        isinstance(nuisance_model_factory, dict)
        and PROPENSITY_MODEL in nuisance_model_factory.keys()
    ):
        raise ValueError(
            "Propensity model factory should be defined using propensity_model_factory "
            "and not nuisance_model_factory."
        )
    if (
        isinstance(nuisance_model_params, dict)
        and PROPENSITY_MODEL in nuisance_model_params.keys()
    ):
        raise ValueError(
            "Propensity model params should be defined using propensity_model_params "
            "and not nuisance_model_params."
        )
    if (
        PROPENSITY_MODEL in nuisance_model_specifications
        and propensity_model_factory is None
        and fitted_propensity_model is None
    ):
        raise ValueError(
            "propensity_model_factory or fitted_propensity_model needs to be defined "
            f"as the {self.__class__.__name__} has a propensity model."
        )

    self._validate_n_variants(n_variants)
    self.is_classification = is_classification
    self.n_variants = n_variants

    self.nuisance_model_factory = _combine_propensity_and_nuisance_specs(
        propensity_model_factory,
        nuisance_model_factory,
        set(nuisance_model_specifications.keys()),
    )
    if nuisance_model_params is None:
        nuisance_model_params = {}  # type: ignore
    if propensity_model_params is None:
        propensity_model_params = {}
    self.nuisance_model_params = _combine_propensity_and_nuisance_specs(
        propensity_model_params,
        nuisance_model_params,
        set(nuisance_model_specifications.keys()),
    )

    self.treatment_model_factory = _initialize_model_dict(
        treatment_model_factory, set(treatment_model_specifications.keys())
    )
    if treatment_model_params is None:
        self.treatment_model_params = _initialize_model_dict(
            {}, set(treatment_model_specifications.keys())
        )
    else:
        self.treatment_model_params = _initialize_model_dict(
            treatment_model_params, set(treatment_model_specifications.keys())
        )

    self.n_folds = _initialize_model_dict(
        n_folds,
        set(nuisance_model_specifications.keys())
        | set(treatment_model_specifications.keys()),
    )
    for model_kind, n_folds_model_kind in self.n_folds.items():
        validate_number_positive(n_folds_model_kind, f"{model_kind} n_folds", True)
    self.random_state = random_state

    self.feature_set = _initialize_model_dict(
        feature_set,
        set(nuisance_model_specifications.keys())
        | set(treatment_model_specifications.keys()),
    )

    self._nuisance_models: dict[str, list[CrossFitEstimator]] = {}
    not_fitted_nuisance_models = set(nuisance_model_specifications.keys())
    self._prefitted_nuisance_models: set[str] = set()

    if fitted_nuisance_models is not None:
        if not set(fitted_nuisance_models.keys()) <= set(
            nuisance_model_specifications.keys()
        ) - {PROPENSITY_MODEL}:
            raise ValueError(
                "The keys present in fitted_nuisance_models should be a subset of "
                f"{set(nuisance_model_specifications.keys()) - {PROPENSITY_MODEL}}"
            )
        self._nuisance_models |= deepcopy(fitted_nuisance_models)
        not_fitted_nuisance_models -= set(fitted_nuisance_models.keys())
        self._prefitted_nuisance_models |= set(fitted_nuisance_models.keys())

    if (
        PROPENSITY_MODEL in nuisance_model_specifications.keys()
        and fitted_propensity_model is not None
    ):
        self._nuisance_models |= {PROPENSITY_MODEL: [fitted_propensity_model]}
        not_fitted_nuisance_models -= {PROPENSITY_MODEL}
        self._prefitted_nuisance_models |= {PROPENSITY_MODEL}

    for name in not_fitted_nuisance_models:
        if self.nuisance_model_factory[name] is None:
            if name == PROPENSITY_MODEL:
                raise ValueError(
                    f"A model for the nuisance model {name} needs to be defined. Either "
                    "in propensity_model_factory or in fitted_propensity_model."
                )
            else:
                raise ValueError(
                    f"A model for the nuisance model {name} needs to be defined. Either "
                    "in nuisance_model_factory or in fitted_nuisance_models."
                )

    self._nuisance_models |= {
        name: [
            CrossFitEstimator(
                n_folds=self.n_folds[name],
                estimator_factory=self.nuisance_model_factory[name],
                estimator_params=self.nuisance_model_params[name],
                random_state=self.random_state,
            )
            for _ in range(nuisance_model_specifications[name]["cardinality"](self))
        ]
        for name in not_fitted_nuisance_models
    }
    self._treatment_models: dict[str, list[CrossFitEstimator]] = {
        name: [
            CrossFitEstimator(
                n_folds=self.n_folds[name],
                estimator_factory=self.treatment_model_factory[name],
                estimator_params=self.treatment_model_params[name],
                random_state=self.random_state,
            )
            for _ in range(
                treatment_model_specifications[name]["cardinality"](self)
            )
        ]
        for name in set(treatment_model_specifications.keys())
    }

    self._validate_models()

init_args property

init_args: dict[str, Any]

Create initialization parameters for a new MetaLearner.

Importantly, this does not copy further internal state, such as the weights or parameters of trained base models.

evaluate

evaluate(
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]

Evaluate the MetaLearner.

The keys in scoring which are not a name of a model contained in the MetaLearner will be ignored, for information about this names check nuisance_model_specifications and treatment_model_specifications. The values must be a list of:

• string representing a sklearn scoring method. Check here for the possible values.
• Callable with signature scorer(estimator, X, y_true, **kwargs). We recommend using sklearn.metrics.make_scorer to create such a Callable.

If some model name is not present in the keys of scoring then the default used metrics will be neg_log_loss if it is a classifier and neg_root_mean_squared_error if it is a regressor.

The returned dictionary keys have the following structure:

• For nuisance models:

  • If the cardinality is one: f"{model_kind}_{scorer}"
  • If there is one model for each treatment variant (including control): f"{model_kind}_{treatment_variant}_{scorer}"

• For treatment models: f"{model_kind}_{treatment_variant}_vs_0_{scorer}"

Where scorer is the name of the scorer if it is a string and "custom_scorer_{idx}" if it is a callable where idx is the index in the scorers list.

In the RLearner case, the "treatment_model" is always evaluated with the r_loss besides the scorers in scoring["treatment_model"], which should support passing the sample_weight keyword argument.

Source code in metalearners/rlearner.py

def evaluate(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]:
    """In the RLearner case, the `"treatment_model"` is always evaluated with the
    [`r_loss`][metalearners.rlearner.r_loss] besides the scorers in
    `scoring["treatment_model"]`, which should support passing the `sample_weight`
    keyword argument."""
    safe_scoring = self._scoring(scoring)

    propensity_evaluation = _evaluate_model_kind(
        cfes=self._nuisance_models[PROPENSITY_MODEL],
        Xs=[X],
        ys=[w],
        scorers=safe_scoring[PROPENSITY_MODEL],
        model_kind=PROPENSITY_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=False,
        feature_set=self.feature_set[PROPENSITY_MODEL],
    )

    outcome_evaluation = _evaluate_model_kind(
        cfes=self._nuisance_models[OUTCOME_MODEL],
        Xs=[X],
        ys=[y],
        scorers=safe_scoring[OUTCOME_MODEL],
        model_kind=OUTCOME_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=False,
        feature_set=self.feature_set[OUTCOME_MODEL],
    )

    # TODO: improve this? generalize it to other metalearners?
    w_hat = self.predict_nuisance(
        X=X,
        is_oos=is_oos,
        oos_method=oos_method,
        model_kind=PROPENSITY_MODEL,
        model_ord=0,
    )

    y_hat = self.predict_nuisance(
        X=X,
        is_oos=is_oos,
        oos_method=oos_method,
        model_kind=OUTCOME_MODEL,
        model_ord=0,
    )
    if self.is_classification:
        y_hat = y_hat[:, 1]

    pseudo_outcome: list[np.ndarray] = []
    sample_weights: list[np.ndarray] = []
    masks: list[Vector] = []
    is_control = w == 0
    for treatment_variant in range(1, self.n_variants):
        is_treatment = w == treatment_variant
        mask = is_treatment | is_control
        tv_pseudo_outcome, tv_sample_weights = self._pseudo_outcome_and_weights(
            X=X,
            y=y,
            w=w,
            treatment_variant=treatment_variant,
            is_oos=is_oos,
            oos_method=oos_method,
            mask=mask,
        )
        pseudo_outcome.append(tv_pseudo_outcome)
        sample_weights.append(tv_sample_weights)
        masks.append(mask)

    treatment_evaluation = _evaluate_model_kind(
        self._treatment_models[TREATMENT_MODEL],
        Xs=[index_matrix(X, masks[tv - 1]) for tv in range(1, self.n_variants)],
        ys=pseudo_outcome,
        scorers=safe_scoring[TREATMENT_MODEL],
        model_kind=TREATMENT_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=True,
        sample_weights=sample_weights,
        feature_set=self.feature_set[TREATMENT_MODEL],
    )

    rloss_evaluation = {}
    tau_hat = self.predict(X=X, is_oos=is_oos, oos_method=oos_method)
    is_control = w == 0
    for treatment_variant in range(1, self.n_variants):
        is_treatment = w == treatment_variant
        mask = is_treatment | is_control

        propensity_estimates = w_hat[:, treatment_variant] / (
            w_hat[:, 0] + w_hat[:, treatment_variant]
        )
        cate_estimates = (
            tau_hat[:, treatment_variant - 1, 1]
            if self.is_classification
            else tau_hat[:, treatment_variant - 1, 0]
        )
        rloss_evaluation[f"r_loss_{treatment_variant}_vs_0"] = r_loss(
            cate_estimates=cate_estimates[mask],
            outcome_estimates=y_hat[mask],
            propensity_scores=propensity_estimates[mask],
            outcomes=index_vector(y, mask),
            treatments=index_vector(w, mask) == treatment_variant,
        )
    return (
        propensity_evaluation
        | outcome_evaluation
        | rloss_evaluation
        | treatment_evaluation
    )

explainer

explainer(
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer

Create an Explainer

which can be used in feature_importances.

This function can be used in two distinct manners based on the provided parameters:

• When parameters X, cate_estimates, and cate_model_factory are all set to None, the function creates an Explainer using the pre-existing treatment models. If these models do not exist, however, it triggers a ValueError.
• On the contrary, if X, cate_estimates, and cate_model_factory are not None, the function initiates an instance of the Explainer class using these parameters. This instance then fits new models for each treatment variant, and these models are employed to calculate the importance of features.

Source code in metalearners/metalearner.py

def explainer(
    self,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer:
    r"""Create an [`Explainer`][metalearners.explainer.Explainer]

    which can be used in [`feature_importances`][metalearners.metalearner.MetaLearner.feature_importances].

    This function can be used in two distinct manners based on the provided parameters:

    *   When parameters `X`, `cate_estimates`, and `cate_model_factory` are all
        set to `None`, the function creates an [`Explainer`][metalearners.explainer.Explainer]
        using the pre-existing treatment models. If these models do not exist, however,
        it triggers a `ValueError`.
    *   On the contrary, if `X`, `cate_estimates`, and `cate_model_factory` are
        not `None`, the function initiates an instance of the [`Explainer`][metalearners.explainer.Explainer]
        class using these parameters. This instance then fits new models for each
        treatment variant, and these models are employed to calculate the importance
        of features.
    """
    if X is None and cate_estimates is None and cate_model_factory is None:
        try:
            cate_cf_models = self._treatment_models[TREATMENT_MODEL]
            cate_models = [cf._overall_estimator for cf in cate_cf_models]
        except KeyError:
            raise ValueError(
                "The metalearner does not have treatment models; hence X, cate_estimates, "
                "and cate_model_factory need to be defined."
            )
        return Explainer(cate_models=cate_models)  # type: ignore
    elif (
        X is not None
        and cate_estimates is not None
        and cate_model_factory is not None
    ):
        return Explainer.from_estimates(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    else:
        raise ValueError(
            "Either all of [X, cate_estimates, cate_model_factory] are None or all "
            "of them must be defined."
        )

feature_importances

feature_importances(
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances

Calculates the feature importance for each treatment group.

If explainer is None, a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters X, cate_estimates, cate_model_factory and cate_model_params are ignored.

If normalization = True, for each treatment variant the feature importances are normalized so that they sum to 1.

feature_names is optional but in the case it’s not passed the names of the features will default to f"Feature {i}" where i is the corresponding feature index.

The returned list contains the feature importances for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def feature_importances(
    self,
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances:
    r"""Calculates the feature importance for each treatment group.

    If `explainer` is `None`, a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `X`, `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    If `normalization = True`, for each treatment variant the feature importances
    are normalized so that they sum to 1.

    `feature_names` is optional but in the case it's not passed the names of the
    features will default to `f"Feature {i}"` where `i` is the corresponding
    feature index.

    The returned list contains the feature importances for each treatment variant in
    ascending order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.feature_importances(
        normalize=normalize, feature_names=feature_names, sort_values=sort_values
    )

fit

fit(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

n_jobs_cross_fitting will be used at the cross-fitting level and n_jobs_base_learners will be used at the stage level. None means 1 unless in a joblib.parallel_backend context. -1 means using all processors. For more information about parallelism check parallelism.

fit_params is an optional dict to be forwarded to base estimator fit calls. It supports two usages patterns:

fit_params={"parameter_of_interest": value_of_interest}

fit_params={
    "nuisance": {
        "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
        "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
    },
    "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
}

In the former approach, the parameter and value of interest are passed to all base models. In the the latter approach, the explicitly qualified parameter-value pairs are passed to respective base models and no fitting parameters are passed to base models not explicitly listed. Note that in this pattern, propensity models are considered a nuisance model.

synchronize_cross_fitting indicates whether the learning of different base models should use exactly the same data splits where possible. Note that if there are several models to be synchronized which are classifiers, these cannot be split via stratification.

Source code in metalearners/metalearner.py

def fit(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    """Fit all models of the MetaLearner.

    If pre-fitted models were passed at instantiation, these are never refitted.

    `n_jobs_cross_fitting` will be used at the cross-fitting level and
    `n_jobs_base_learners` will be used at the stage level. `None` means 1 unless in a
    [`joblib.parallel_backend`](https://joblib.readthedocs.io/en/latest/generated/joblib.parallel_backend.html#joblib.parallel_backend)
    context. `-1` means using all processors.
    For more information about parallelism check [parallelism](parallelism.md).


    `fit_params` is an optional `dict` to be forwarded to base estimator `fit` calls. It supports
    two usages patterns:

    ```python
    fit_params={"parameter_of_interest": value_of_interest}
    ```

    ```python
    fit_params={
        "nuisance": {
            "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
            "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
        },
        "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
    }
    ```

    In the former approach, the parameter and value of interest are passed to all base models. In the
    the latter approach, the explicitly qualified parameter-value pairs are passed to respective base
    models and no fitting parameters are passed to base models not explicitly listed. Note that in this
    pattern, propensity models are considered a nuisance model.

    `synchronize_cross_fitting` indicates whether the learning of different base models should use exactly
    the same data splits where possible. Note that if there are several models to be synchronized which are
    classifiers, these cannot be split via stratification.
    """
    self.fit_all_nuisance(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    self.fit_all_treatment(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    return self

fit_all_nuisance

fit_all_nuisance(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all nuisance models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the nuisance models, and in the case they should also be used by the treatment models, these should also be passed in the following call to fit_all_treatment.

This method, combined with fit_all_treatment, facilitates the segmentation of the metalearner fitting process into two distinct parts. This division allows for interventions between the two stages, such as performing feature selection for the treatment models or conducting hyperparameter optimization within the nuisance models.

Source code in metalearners/rlearner.py

def fit_all_nuisance(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    self._validate_treatment(w)
    self._validate_outcome(y, w)

    qualified_fit_params = self._qualified_fit_params(fit_params)
    self._validate_fit_params(qualified_fit_params)

    if synchronize_cross_fitting:
        cv_split_indices = self._split(X)
    else:
        cv_split_indices = None

    nuisance_jobs: list[_ParallelJoblibSpecification | None] = []

    nuisance_jobs.append(
        self._nuisance_joblib_specifications(
            X=X,
            y=w,
            model_kind=PROPENSITY_MODEL,
            model_ord=0,
            n_jobs_cross_fitting=n_jobs_cross_fitting,
            fit_params=qualified_fit_params[NUISANCE][PROPENSITY_MODEL],
            cv=cv_split_indices,
        )
    )
    nuisance_jobs.append(
        self._nuisance_joblib_specifications(
            X=X,
            y=y,
            model_kind=OUTCOME_MODEL,
            model_ord=0,
            n_jobs_cross_fitting=n_jobs_cross_fitting,
            fit_params=qualified_fit_params[NUISANCE][OUTCOME_MODEL],
            cv=cv_split_indices,
        )
    )

    parallel = Parallel(n_jobs=n_jobs_base_learners)
    results = parallel(
        delayed(_fit_cross_fit_estimator_joblib)(job)
        for job in nuisance_jobs
        if job is not None
    )
    self._assign_joblib_nuisance_results(results)

    return self

fit_all_treatment

fit_all_treatment(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
    epsilon: float = _EPSILON,
) -> Self

Fit all treatment models of the MetaLearner.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the treatment models, as the nuisance models should already be fitted.

Source code in metalearners/rlearner.py

def fit_all_treatment(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
    epsilon: float = _EPSILON,
) -> Self:
    qualified_fit_params = self._qualified_fit_params(fit_params)
    treatment_jobs: list[_ParallelJoblibSpecification] = []
    self._variants_indices = []
    for treatment_variant in range(1, self.n_variants):

        is_treatment = w == treatment_variant
        is_control = w == 0
        mask = is_treatment | is_control

        self._variants_indices.append(mask)

        pseudo_outcomes, weights = self._pseudo_outcome_and_weights(
            X=X,
            w=w,
            y=y,
            treatment_variant=treatment_variant,
            mask=mask,
            epsilon=epsilon,
            is_oos=False,
        )

        X_filtered = index_matrix(X, mask)

        treatment_jobs.append(
            self._treatment_joblib_specifications(
                X=X_filtered,
                y=pseudo_outcomes,
                model_kind=TREATMENT_MODEL,
                model_ord=treatment_variant - 1,
                fit_params=qualified_fit_params[TREATMENT][TREATMENT_MODEL]
                | {_SAMPLE_WEIGHT: weights},
                n_jobs_cross_fitting=n_jobs_cross_fitting,
            )
        )
    parallel = Parallel(n_jobs=n_jobs_base_learners)
    results = parallel(
        delayed(_fit_cross_fit_estimator_joblib)(job) for job in treatment_jobs
    )
    self._assign_joblib_treatment_results(results)
    return self

fit_nuisance

fit_nuisance(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit a given nuisance model of a MetaLearner.

y represents the objective of the given nuisance model, not necessarily the outcome of the experiment. If pre-fitted models were passed at instantiation, these are never refitted.

Source code in metalearners/metalearner.py

def fit_nuisance(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit a given nuisance model of a MetaLearner.

    `y` represents the objective of the given nuisance model, not necessarily the outcome of the experiment.
    If pre-fitted models were passed at instantiation, these are never refitted.
    """
    if model_kind in self._prefitted_nuisance_models:
        return self
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._nuisance_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

fit_treatment

fit_treatment(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit the treatment model of a MetaLearner.

y represents the objective of the given treatment model, not necessarily the outcome of the experiment.

Source code in metalearners/metalearner.py

def fit_treatment(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit the treatment model of a MetaLearner.

    `y` represents the objective of the given treatment model, not necessarily the outcome of the experiment.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._treatment_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

nuisance_model_specifications classmethod

nuisance_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all first-stage models.

Source code in metalearners/rlearner.py

@classmethod
def nuisance_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return {
        PROPENSITY_MODEL: _ModelSpecifications(
            cardinality=get_one,
            predict_method=get_predict_proba,
        ),
        OUTCOME_MODEL: _ModelSpecifications(
            cardinality=get_one,
            predict_method=MetaLearner._outcome_predict_method,
        ),
    }

predict

predict(
    X, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Estimate the CATE.

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants} - 1, 1)\) if the outcome is a scalar, i.e. in case of a regression problem.

• \((n_{obs}, n_{variants} - 1, n_{classes})\) if the outcome is a class, i.e. in case of a classification problem.

In the case of multiple treatment variants, the second dimension represents the CATE of the corresponding variant vs the control (variant 0).

Source code in metalearners/rlearner.py

def predict(
    self,
    X,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    n_outputs = 2 if self.is_classification else 1
    tau_hat = np.zeros((safe_len(X), self.n_variants - 1, n_outputs))

    if is_oos:

        for treatment_variant in range(1, self.n_variants):
            variant_estimates = self.predict_treatment(
                X,
                is_oos=is_oos,
                oos_method=oos_method,
                model_kind=TREATMENT_MODEL,
                model_ord=treatment_variant - 1,
            )
            if self.is_classification:
                # This is to be consistent with other MetaLearners (e.g. S and T) that automatically
                # work with multiclass outcomes and return the CATE estimate for each class. As the R-Learner only
                # works with binary classes (the pseudo outcome formula does not make sense with
                # multiple classes unless some adaptation is done) we can manually infer the
                # CATE estimate for the complementary class  -- returning a matrix of shape (N, 2).
                variant_estimates = np.stack(
                    [-variant_estimates, variant_estimates], axis=-1
                )
            variant_estimates = variant_estimates.reshape(safe_len(X), n_outputs)
            tau_hat[:, treatment_variant - 1, :] = variant_estimates

        return tau_hat

    for treatment_variant in range(1, self.n_variants):
        variant_indices = self._variants_indices[treatment_variant - 1]

        variant_estimates = self.predict_treatment(
            index_matrix(X, variant_indices),
            is_oos=False,
            model_kind=TREATMENT_MODEL,
            model_ord=treatment_variant - 1,
        )
        if sum(~variant_indices) > 0:
            non_variant_estimates = self.predict_treatment(
                index_matrix(X, ~variant_indices),
                is_oos=True,
                oos_method=oos_method,
                model_kind=TREATMENT_MODEL,
                model_ord=treatment_variant - 1,
            )
        if self.is_classification:
            # This is to be consistent with other MetaLearners (e.g. S and T) that automatically
            # work with multiclass outcomes and return the CATE estimate for each class. As the R-Learner only
            # works with binary classes (the pseudo outcome formula does not make sense with
            # multiple classes unless some adaptation is done) we can manually infer the
            # CATE estimate for the complementary class  -- returning a matrix of shape (N, 2).
            variant_estimates = np.stack(
                [-variant_estimates, variant_estimates], axis=-1
            )
            if sum(~variant_indices) > 0:
                non_variant_estimates = np.stack(
                    [-non_variant_estimates, non_variant_estimates], axis=-1
                )
        variant_estimates = variant_estimates.reshape(
            (sum(variant_indices), n_outputs)
        )
        if sum(~variant_indices) > 0:
            non_variant_estimates = non_variant_estimates.reshape(
                (sum(~variant_indices), n_outputs)
            )
            tau_hat[~variant_indices, treatment_variant - 1] = non_variant_estimates

        tau_hat[variant_indices, treatment_variant - 1] = variant_estimates
    return tau_hat

predict_conditional_average_outcomes

predict_conditional_average_outcomes(
    X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Predict the vectors of conditional average outcomes.

These are defined as \(\mathbb{E}[Y_i(w) | X]\) for each treatment variant \(w\).

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants}, 1)\) if the outcome is a scalar, i.e., in case of a regression problem.

• \((n_{obs}, n_{variants}, n_{classes})\) if the outcome is a class, i.e., in case of a classification problem.

The conditional average outcomes are estimated as follows:

• \(Y_i(0) = \hat{\mu}(X_i) - \sum_{k=1}^{K} \hat{e}_k(X_i) \hat{\tau_k}(X_i)\)
• \(Y_i(k) = Y_i(0) + \hat{\tau_k}(X_i)\) for \(k \in \{1, \dots, K\}\)

where \(K\) is the number of treatment variants.

Source code in metalearners/rlearner.py

def predict_conditional_average_outcomes(
    self, X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray:
    r"""The conditional average outcomes are estimated as follows:

    * $Y_i(0) = \hat{\mu}(X_i) - \sum_{k=1}^{K} \hat{e}_k(X_i) \hat{\tau_k}(X_i)$
    * $Y_i(k) = Y_i(0) + \hat{\tau_k}(X_i)$ for $k \in \{1, \dots, K\}$

    where $K$ is the number of treatment variants.
    """
    n_obs = safe_len(X)

    cate_estimates = self.predict(
        X=X,
        is_oos=is_oos,
        oos_method=oos_method,
    )
    propensity_estimates = self.predict_nuisance(
        X=X,
        model_kind=PROPENSITY_MODEL,
        model_ord=0,
        is_oos=is_oos,
        oos_method=oos_method,
    )
    outcome_estimates = self.predict_nuisance(
        X=X,
        model_kind=OUTCOME_MODEL,
        model_ord=0,
        is_oos=is_oos,
        oos_method=oos_method,
    )

    conditional_average_outcomes_list = []

    control_outcomes = outcome_estimates

    # TODO: Consider whether the readability vs efficiency trade-off should be dealt with differently here.
    # One could use matrix/tensor operations instead.
    for treatment_variant in range(1, self.n_variants):
        if (n_outputs := cate_estimates.shape[2]) > 1:
            for outcome_channel in range(0, n_outputs):
                control_outcomes[:, outcome_channel] -= (
                    propensity_estimates[:, treatment_variant]
                    * cate_estimates[:, treatment_variant - 1, outcome_channel]
                )
        else:
            control_outcomes -= (
                propensity_estimates[:, treatment_variant]
                * cate_estimates[:, treatment_variant - 1, 0]
            )

    conditional_average_outcomes_list.append(control_outcomes)

    for treatment_variant in range(1, self.n_variants):
        conditional_average_outcomes_list.append(
            control_outcomes
            + np.reshape(
                cate_estimates[:, treatment_variant - 1, :],
                (control_outcomes.shape),
            )
        )

    return np.stack(conditional_average_outcomes_list, axis=1).reshape(
        n_obs, self.n_variants, -1
    )

predict_nuisance

predict_nuisance(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given nuisance model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_nuisance(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given nuisance model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    predict_method_name = self.nuisance_model_specifications()[model_kind][
        "predict_method"
    ](self)
    predict_method = getattr(
        self._nuisance_models[model_kind][model_ord], predict_method_name
    )
    return predict_method(X_filtered, is_oos, oos_method)

predict_treatment

predict_treatment(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given treatment model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_treatment(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given treatment model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    return self._treatment_models[model_kind][model_ord].predict(
        X_filtered, is_oos, oos_method
    )

shap_values

shap_values(
    X: Matrix,
    shap_explainer_factory: type[Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]

Calculates the shap values for each treatment group.

If explainer is None a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters cate_estimates, cate_model_factory and cate_model_params are ignored.

The parameter shap_explainer_factory can be used to specify the type of shap explainer, for the different options see here.

The returned list contains the shap values for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def shap_values(
    self,
    X: Matrix,
    shap_explainer_factory: type[shap.Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]:
    """Calculates the shap values for each treatment group.

    If `explainer` is `None` a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    The parameter `shap_explainer_factory` can be used to specify the type of shap
    explainer, for the different options see
    [here](https://shap.readthedocs.io/en/latest/api.html#explainers).

    The returned list contains the shap values for each treatment variant in ascending
    order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=None if cate_estimates is None else X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.shap_values(
        X=X,
        shap_explainer_factory=shap_explainer_factory,
        shap_explainer_params=shap_explainer_params,
    )

treatment_model_specifications classmethod

treatment_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all second-stage models.

Source code in metalearners/rlearner.py

@classmethod
def treatment_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return {
        TREATMENT_MODEL: _ModelSpecifications(
            cardinality=MetaLearner._get_n_variants_minus_one,
            predict_method=get_predict,
        )
    }

r_loss

r_loss(
    cate_estimates: Vector,
    outcome_estimates: Vector,
    propensity_scores: Vector,
    outcomes: Vector,
    treatments: Vector,
) -> float

Compute the square-root of the R-loss as introduced by Nie et al.

This function computes:

\[ \sqrt{\frac{1}{N}\sum_{i=1}^N ((y_i - \mu(X_i)) - \hat{\tau}(X_i) (w_i - e(X_i)))^2} \]

The R-Learner proposed in Nie et al. (2017) relies on a loss function which can be used in combination with empirical risk minimization to learn a CATE model.

Independently of the R-Learner, one can use the R-loss for evaluating CATE estimates in general.

Source code in metalearners/rlearner.py

def r_loss(
    cate_estimates: Vector,
    outcome_estimates: Vector,
    propensity_scores: Vector,
    outcomes: Vector,
    treatments: Vector,
) -> float:
    r"""Compute the square-root of the R-loss as introduced by Nie et al.

    This function computes:

    $$
    \sqrt{\frac{1}{N}\sum_{i=1}^N ((y_i - \mu(X_i)) - \hat{\tau}(X_i)
    (w_i - e(X_i)))^2}
    $$

    The R-Learner proposed in [Nie et al. (2017)](https://arxiv.org/pdf/1712.04912.pdf)
    relies on a loss function which can be used in combination with empirical risk
    minimization to learn a CATE model.

    Independently of the R-Learner, one can use the R-loss for evaluating CATE estimates
    in general.
    """
    inputs = [
        cate_estimates,
        outcome_estimates,
        propensity_scores,
        outcomes,
        treatments,
    ]
    validate_all_vectors_same_index(inputs)

    residualised_outcomes = outcomes - outcome_estimates
    residualised_treatments = treatments - propensity_scores
    return root_mean_squared_error(
        residualised_outcomes, cate_estimates * residualised_treatments
    )

slearner

SLearner

SLearner(
    is_classification: bool,
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel]
    | None = None,
    nuisance_model_params: Params
    | dict[str, Params]
    | None = None,
    treatment_model_params: Params
    | dict[str, Params]
    | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[
        str, list[CrossFitEstimator]
    ]
    | None = None,
    fitted_propensity_model: CrossFitEstimator
    | None = None,
    feature_set: Features
    | dict[str, Features]
    | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
)

Bases: MetaLearner

S-Learner for CATE estimation as described by Kuenzel et al (2019).

Source code in metalearners/slearner.py

def __init__(
    self,
    is_classification: bool,
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel] | None = None,
    nuisance_model_params: Params | dict[str, Params] | None = None,
    treatment_model_params: Params | dict[str, Params] | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[str, list[CrossFitEstimator]] | None = None,
    fitted_propensity_model: CrossFitEstimator | None = None,
    feature_set: Features | dict[str, Features] | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
):
    if feature_set is not None:
        # For SLearner it does not make sense to allow feature set as we only have one model
        # and having it would bring problems when using fit_nuisance and predict_nuisance
        # as we need to add the treatment column.
        warnings.warn(
            "Base-model specific feature_sets were provided to S-Learner. "
            "These will be ignored and all available features will be used instead."
        )
    super().__init__(
        is_classification=is_classification,
        n_variants=n_variants,
        nuisance_model_factory=nuisance_model_factory,
        treatment_model_factory=treatment_model_factory,
        propensity_model_factory=propensity_model_factory,
        nuisance_model_params=nuisance_model_params,
        treatment_model_params=treatment_model_params,
        propensity_model_params=propensity_model_params,
        fitted_nuisance_models=fitted_nuisance_models,
        fitted_propensity_model=fitted_propensity_model,
        feature_set=None,
        n_folds=n_folds,
        random_state=random_state,
    )

init_args property

init_args: dict[str, Any]

Create initialization parameters for a new MetaLearner.

Importantly, this does not copy further internal state, such as the weights or parameters of trained base models.

evaluate

evaluate(
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]

Evaluate the MetaLearner.

The keys in scoring which are not a name of a model contained in the MetaLearner will be ignored, for information about this names check nuisance_model_specifications and treatment_model_specifications. The values must be a list of:

• string representing a sklearn scoring method. Check here for the possible values.
• Callable with signature scorer(estimator, X, y_true, **kwargs). We recommend using sklearn.metrics.make_scorer to create such a Callable.

If some model name is not present in the keys of scoring then the default used metrics will be neg_log_loss if it is a classifier and neg_root_mean_squared_error if it is a regressor.

The returned dictionary keys have the following structure:

• For nuisance models:

  • If the cardinality is one: f"{model_kind}_{scorer}"
  • If there is one model for each treatment variant (including control): f"{model_kind}_{treatment_variant}_{scorer}"

• For treatment models: f"{model_kind}_{treatment_variant}_vs_0_{scorer}"

Where scorer is the name of the scorer if it is a string and "custom_scorer_{idx}" if it is a callable where idx is the index in the scorers list.

Source code in metalearners/slearner.py

def evaluate(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]:
    safe_scoring = self._scoring(scoring)

    X_with_w = _append_treatment_to_covariates(
        X, w, self._supports_categoricals, self.n_variants
    )
    return _evaluate_model_kind(
        cfes=self._nuisance_models[_BASE_MODEL],
        Xs=[X_with_w],
        ys=[y],
        scorers=safe_scoring[_BASE_MODEL],
        model_kind=_BASE_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=False,
        feature_set=self.feature_set[_BASE_MODEL],
    )

explainer

explainer(
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer

Create an Explainer

which can be used in feature_importances.

This function can be used in two distinct manners based on the provided parameters:

• When parameters X, cate_estimates, and cate_model_factory are all set to None, the function creates an Explainer using the pre-existing treatment models. If these models do not exist, however, it triggers a ValueError.
• On the contrary, if X, cate_estimates, and cate_model_factory are not None, the function initiates an instance of the Explainer class using these parameters. This instance then fits new models for each treatment variant, and these models are employed to calculate the importance of features.

Source code in metalearners/metalearner.py

def explainer(
    self,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer:
    r"""Create an [`Explainer`][metalearners.explainer.Explainer]

    which can be used in [`feature_importances`][metalearners.metalearner.MetaLearner.feature_importances].

    This function can be used in two distinct manners based on the provided parameters:

    *   When parameters `X`, `cate_estimates`, and `cate_model_factory` are all
        set to `None`, the function creates an [`Explainer`][metalearners.explainer.Explainer]
        using the pre-existing treatment models. If these models do not exist, however,
        it triggers a `ValueError`.
    *   On the contrary, if `X`, `cate_estimates`, and `cate_model_factory` are
        not `None`, the function initiates an instance of the [`Explainer`][metalearners.explainer.Explainer]
        class using these parameters. This instance then fits new models for each
        treatment variant, and these models are employed to calculate the importance
        of features.
    """
    if X is None and cate_estimates is None and cate_model_factory is None:
        try:
            cate_cf_models = self._treatment_models[TREATMENT_MODEL]
            cate_models = [cf._overall_estimator for cf in cate_cf_models]
        except KeyError:
            raise ValueError(
                "The metalearner does not have treatment models; hence X, cate_estimates, "
                "and cate_model_factory need to be defined."
            )
        return Explainer(cate_models=cate_models)  # type: ignore
    elif (
        X is not None
        and cate_estimates is not None
        and cate_model_factory is not None
    ):
        return Explainer.from_estimates(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    else:
        raise ValueError(
            "Either all of [X, cate_estimates, cate_model_factory] are None or all "
            "of them must be defined."
        )

feature_importances

feature_importances(
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances

Calculates the feature importance for each treatment group.

If explainer is None, a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters X, cate_estimates, cate_model_factory and cate_model_params are ignored.

If normalization = True, for each treatment variant the feature importances are normalized so that they sum to 1.

feature_names is optional but in the case it’s not passed the names of the features will default to f"Feature {i}" where i is the corresponding feature index.

The returned list contains the feature importances for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def feature_importances(
    self,
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances:
    r"""Calculates the feature importance for each treatment group.

    If `explainer` is `None`, a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `X`, `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    If `normalization = True`, for each treatment variant the feature importances
    are normalized so that they sum to 1.

    `feature_names` is optional but in the case it's not passed the names of the
    features will default to `f"Feature {i}"` where `i` is the corresponding
    feature index.

    The returned list contains the feature importances for each treatment variant in
    ascending order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.feature_importances(
        normalize=normalize, feature_names=feature_names, sort_values=sort_values
    )

fit

fit(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

n_jobs_cross_fitting will be used at the cross-fitting level and n_jobs_base_learners will be used at the stage level. None means 1 unless in a joblib.parallel_backend context. -1 means using all processors. For more information about parallelism check parallelism.

fit_params is an optional dict to be forwarded to base estimator fit calls. It supports two usages patterns:

fit_params={"parameter_of_interest": value_of_interest}

fit_params={
    "nuisance": {
        "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
        "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
    },
    "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
}

In the former approach, the parameter and value of interest are passed to all base models. In the the latter approach, the explicitly qualified parameter-value pairs are passed to respective base models and no fitting parameters are passed to base models not explicitly listed. Note that in this pattern, propensity models are considered a nuisance model.

synchronize_cross_fitting indicates whether the learning of different base models should use exactly the same data splits where possible. Note that if there are several models to be synchronized which are classifiers, these cannot be split via stratification.

Source code in metalearners/metalearner.py

def fit(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    """Fit all models of the MetaLearner.

    If pre-fitted models were passed at instantiation, these are never refitted.

    `n_jobs_cross_fitting` will be used at the cross-fitting level and
    `n_jobs_base_learners` will be used at the stage level. `None` means 1 unless in a
    [`joblib.parallel_backend`](https://joblib.readthedocs.io/en/latest/generated/joblib.parallel_backend.html#joblib.parallel_backend)
    context. `-1` means using all processors.
    For more information about parallelism check [parallelism](parallelism.md).


    `fit_params` is an optional `dict` to be forwarded to base estimator `fit` calls. It supports
    two usages patterns:

    ```python
    fit_params={"parameter_of_interest": value_of_interest}
    ```

    ```python
    fit_params={
        "nuisance": {
            "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
            "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
        },
        "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
    }
    ```

    In the former approach, the parameter and value of interest are passed to all base models. In the
    the latter approach, the explicitly qualified parameter-value pairs are passed to respective base
    models and no fitting parameters are passed to base models not explicitly listed. Note that in this
    pattern, propensity models are considered a nuisance model.

    `synchronize_cross_fitting` indicates whether the learning of different base models should use exactly
    the same data splits where possible. Note that if there are several models to be synchronized which are
    classifiers, these cannot be split via stratification.
    """
    self.fit_all_nuisance(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    self.fit_all_treatment(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    return self

fit_all_nuisance

fit_all_nuisance(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all nuisance models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the nuisance models, and in the case they should also be used by the treatment models, these should also be passed in the following call to fit_all_treatment.

This method, combined with fit_all_treatment, facilitates the segmentation of the metalearner fitting process into two distinct parts. This division allows for interventions between the two stages, such as performing feature selection for the treatment models or conducting hyperparameter optimization within the nuisance models.

Source code in metalearners/slearner.py

def fit_all_nuisance(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    self._validate_treatment(w)
    self._validate_outcome(y, w)
    self._fitted_treatments = adapt_treatment_dtypes(w)

    mock_model = self.nuisance_model_factory[_BASE_MODEL](
        **self.nuisance_model_params[_BASE_MODEL]
    )
    self._supports_categoricals = supports_categoricals(mock_model)
    X_with_w = _append_treatment_to_covariates(
        X, w, self._supports_categoricals, self.n_variants
    )

    qualified_fit_params = self._qualified_fit_params(fit_params)

    self.fit_nuisance(
        X=X_with_w,
        y=y,
        model_kind=_BASE_MODEL,
        model_ord=0,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=qualified_fit_params[NUISANCE][_BASE_MODEL],
    )
    return self

fit_all_treatment

fit_all_treatment(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all treatment models of the MetaLearner.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the treatment models, as the nuisance models should already be fitted.

Source code in metalearners/slearner.py

def fit_all_treatment(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    return self

fit_nuisance

fit_nuisance(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit a given nuisance model of a MetaLearner.

y represents the objective of the given nuisance model, not necessarily the outcome of the experiment. If pre-fitted models were passed at instantiation, these are never refitted.

Source code in metalearners/metalearner.py

def fit_nuisance(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit a given nuisance model of a MetaLearner.

    `y` represents the objective of the given nuisance model, not necessarily the outcome of the experiment.
    If pre-fitted models were passed at instantiation, these are never refitted.
    """
    if model_kind in self._prefitted_nuisance_models:
        return self
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._nuisance_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

fit_treatment

fit_treatment(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit the treatment model of a MetaLearner.

y represents the objective of the given treatment model, not necessarily the outcome of the experiment.

Source code in metalearners/metalearner.py

def fit_treatment(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit the treatment model of a MetaLearner.

    `y` represents the objective of the given treatment model, not necessarily the outcome of the experiment.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._treatment_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

nuisance_model_specifications classmethod

nuisance_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all first-stage models.

Source code in metalearners/slearner.py

@classmethod
def nuisance_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return {
        _BASE_MODEL: _ModelSpecifications(
            cardinality=get_one,
            predict_method=MetaLearner._outcome_predict_method,
        )
    }

predict

predict(
    X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Estimate the CATE.

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants} - 1, 1)\) if the outcome is a scalar, i.e. in case of a regression problem.

• \((n_{obs}, n_{variants} - 1, n_{classes})\) if the outcome is a class, i.e. in case of a classification problem.

In the case of multiple treatment variants, the second dimension represents the CATE of the corresponding variant vs the control (variant 0).

Source code in metalearners/slearner.py

def predict(
    self,
    X: Matrix,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    conditional_average_outcomes = self.predict_conditional_average_outcomes(
        X=X, is_oos=is_oos, oos_method=oos_method
    )

    return conditional_average_outcomes[:, 1:] - (
        conditional_average_outcomes[:, [0]]
    )

predict_conditional_average_outcomes

predict_conditional_average_outcomes(
    X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Predict the vectors of conditional average outcomes.

These are defined as \(\mathbb{E}[Y_i(w) | X]\) for each treatment variant \(w\).

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants}, 1)\) if the outcome is a scalar, i.e., in case of a regression problem.

• \((n_{obs}, n_{variants}, n_{classes})\) if the outcome is a class, i.e., in case of a classification problem.

Source code in metalearners/slearner.py

def predict_conditional_average_outcomes(
    self, X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray:
    n_obs = safe_len(X)
    conditional_average_outcomes_list = []

    for treatment_variant in range(self.n_variants):
        w = np.array([treatment_variant] * n_obs)
        X_with_w = _append_treatment_to_covariates(
            X, w, self._supports_categoricals, self.n_variants
        )
        variant_predictions = self.predict_nuisance(
            X=X_with_w,
            model_kind=_BASE_MODEL,
            model_ord=0,
            is_oos=is_oos,
            oos_method=oos_method,
        )

        conditional_average_outcomes_list.append(variant_predictions)

    return np.stack(conditional_average_outcomes_list, axis=1).reshape(
        n_obs, self.n_variants, -1
    )

predict_nuisance

predict_nuisance(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given nuisance model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_nuisance(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given nuisance model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    predict_method_name = self.nuisance_model_specifications()[model_kind][
        "predict_method"
    ](self)
    predict_method = getattr(
        self._nuisance_models[model_kind][model_ord], predict_method_name
    )
    return predict_method(X_filtered, is_oos, oos_method)

predict_treatment

predict_treatment(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given treatment model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_treatment(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given treatment model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    return self._treatment_models[model_kind][model_ord].predict(
        X_filtered, is_oos, oos_method
    )

shap_values

shap_values(
    X: Matrix,
    shap_explainer_factory: type[Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]

Calculates the shap values for each treatment group.

If explainer is None a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters cate_estimates, cate_model_factory and cate_model_params are ignored.

The parameter shap_explainer_factory can be used to specify the type of shap explainer, for the different options see here.

The returned list contains the shap values for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def shap_values(
    self,
    X: Matrix,
    shap_explainer_factory: type[shap.Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]:
    """Calculates the shap values for each treatment group.

    If `explainer` is `None` a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    The parameter `shap_explainer_factory` can be used to specify the type of shap
    explainer, for the different options see
    [here](https://shap.readthedocs.io/en/latest/api.html#explainers).

    The returned list contains the shap values for each treatment variant in ascending
    order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=None if cate_estimates is None else X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.shap_values(
        X=X,
        shap_explainer_factory=shap_explainer_factory,
        shap_explainer_params=shap_explainer_params,
    )

treatment_model_specifications classmethod

treatment_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all second-stage models.

Source code in metalearners/slearner.py

@classmethod
def treatment_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return dict()

tlearner

TLearner

TLearner(
    is_classification: bool,
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel]
    | None = None,
    nuisance_model_params: Params
    | dict[str, Params]
    | None = None,
    treatment_model_params: Params
    | dict[str, Params]
    | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[
        str, list[CrossFitEstimator]
    ]
    | None = None,
    fitted_propensity_model: CrossFitEstimator
    | None = None,
    feature_set: Features
    | dict[str, Features]
    | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
)

Bases: _ConditionalAverageOutcomeMetaLearner

T-Learner for CATE estimation as described by Kuenzel et al (2019).

Source code in metalearners/metalearner.py

def __init__(
    self,
    is_classification: bool,
    # TODO: Consider whether we can make this not a state of the MetaLearner
    # but rather just a parameter of a predict call.
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel] | None = None,
    nuisance_model_params: Params | dict[str, Params] | None = None,
    treatment_model_params: Params | dict[str, Params] | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[str, list[CrossFitEstimator]] | None = None,
    fitted_propensity_model: CrossFitEstimator | None = None,
    feature_set: Features | dict[str, Features] | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
):
    super().__init__(
        nuisance_model_factory=nuisance_model_factory,
        is_classification=is_classification,
        n_variants=n_variants,
        treatment_model_factory=treatment_model_factory,
        propensity_model_factory=propensity_model_factory,
        nuisance_model_params=nuisance_model_params,
        treatment_model_params=treatment_model_params,
        propensity_model_params=propensity_model_params,
        fitted_nuisance_models=fitted_nuisance_models,
        fitted_propensity_model=fitted_propensity_model,
        feature_set=feature_set,
        n_folds=n_folds,
        random_state=random_state,
    )
    self._treatment_variants_mask: list[np.ndarray] | None = None

init_args property

init_args: dict[str, Any]

Create initialization parameters for a new MetaLearner.

Importantly, this does not copy further internal state, such as the weights or parameters of trained base models.

evaluate

evaluate(
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]

Evaluate the MetaLearner.

The keys in scoring which are not a name of a model contained in the MetaLearner will be ignored, for information about this names check nuisance_model_specifications and treatment_model_specifications. The values must be a list of:

• string representing a sklearn scoring method. Check here for the possible values.
• Callable with signature scorer(estimator, X, y_true, **kwargs). We recommend using sklearn.metrics.make_scorer to create such a Callable.

If some model name is not present in the keys of scoring then the default used metrics will be neg_log_loss if it is a classifier and neg_root_mean_squared_error if it is a regressor.

The returned dictionary keys have the following structure:

• For nuisance models:

  • If the cardinality is one: f"{model_kind}_{scorer}"
  • If there is one model for each treatment variant (including control): f"{model_kind}_{treatment_variant}_{scorer}"

• For treatment models: f"{model_kind}_{treatment_variant}_vs_0_{scorer}"

Where scorer is the name of the scorer if it is a string and "custom_scorer_{idx}" if it is a callable where idx is the index in the scorers list.

Source code in metalearners/tlearner.py

def evaluate(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]:
    safe_scoring = self._scoring(scoring)

    return _evaluate_model_kind(
        cfes=self._nuisance_models[VARIANT_OUTCOME_MODEL],
        Xs=[index_matrix(X, w == tv) for tv in range(self.n_variants)],
        ys=[index_vector(y, w == tv) for tv in range(self.n_variants)],
        scorers=safe_scoring[VARIANT_OUTCOME_MODEL],
        model_kind=VARIANT_OUTCOME_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=False,
        feature_set=self.feature_set[VARIANT_OUTCOME_MODEL],
    )

explainer

explainer(
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer

Create an Explainer

which can be used in feature_importances.

This function can be used in two distinct manners based on the provided parameters:

• When parameters X, cate_estimates, and cate_model_factory are all set to None, the function creates an Explainer using the pre-existing treatment models. If these models do not exist, however, it triggers a ValueError.
• On the contrary, if X, cate_estimates, and cate_model_factory are not None, the function initiates an instance of the Explainer class using these parameters. This instance then fits new models for each treatment variant, and these models are employed to calculate the importance of features.

Source code in metalearners/metalearner.py

def explainer(
    self,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer:
    r"""Create an [`Explainer`][metalearners.explainer.Explainer]

    which can be used in [`feature_importances`][metalearners.metalearner.MetaLearner.feature_importances].

    This function can be used in two distinct manners based on the provided parameters:

    *   When parameters `X`, `cate_estimates`, and `cate_model_factory` are all
        set to `None`, the function creates an [`Explainer`][metalearners.explainer.Explainer]
        using the pre-existing treatment models. If these models do not exist, however,
        it triggers a `ValueError`.
    *   On the contrary, if `X`, `cate_estimates`, and `cate_model_factory` are
        not `None`, the function initiates an instance of the [`Explainer`][metalearners.explainer.Explainer]
        class using these parameters. This instance then fits new models for each
        treatment variant, and these models are employed to calculate the importance
        of features.
    """
    if X is None and cate_estimates is None and cate_model_factory is None:
        try:
            cate_cf_models = self._treatment_models[TREATMENT_MODEL]
            cate_models = [cf._overall_estimator for cf in cate_cf_models]
        except KeyError:
            raise ValueError(
                "The metalearner does not have treatment models; hence X, cate_estimates, "
                "and cate_model_factory need to be defined."
            )
        return Explainer(cate_models=cate_models)  # type: ignore
    elif (
        X is not None
        and cate_estimates is not None
        and cate_model_factory is not None
    ):
        return Explainer.from_estimates(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    else:
        raise ValueError(
            "Either all of [X, cate_estimates, cate_model_factory] are None or all "
            "of them must be defined."
        )

feature_importances

feature_importances(
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances

Calculates the feature importance for each treatment group.

If explainer is None, a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters X, cate_estimates, cate_model_factory and cate_model_params are ignored.

If normalization = True, for each treatment variant the feature importances are normalized so that they sum to 1.

feature_names is optional but in the case it’s not passed the names of the features will default to f"Feature {i}" where i is the corresponding feature index.

The returned list contains the feature importances for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def feature_importances(
    self,
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances:
    r"""Calculates the feature importance for each treatment group.

    If `explainer` is `None`, a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `X`, `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    If `normalization = True`, for each treatment variant the feature importances
    are normalized so that they sum to 1.

    `feature_names` is optional but in the case it's not passed the names of the
    features will default to `f"Feature {i}"` where `i` is the corresponding
    feature index.

    The returned list contains the feature importances for each treatment variant in
    ascending order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.feature_importances(
        normalize=normalize, feature_names=feature_names, sort_values=sort_values
    )

fit

fit(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

n_jobs_cross_fitting will be used at the cross-fitting level and n_jobs_base_learners will be used at the stage level. None means 1 unless in a joblib.parallel_backend context. -1 means using all processors. For more information about parallelism check parallelism.

fit_params is an optional dict to be forwarded to base estimator fit calls. It supports two usages patterns:

fit_params={"parameter_of_interest": value_of_interest}

fit_params={
    "nuisance": {
        "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
        "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
    },
    "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
}

In the former approach, the parameter and value of interest are passed to all base models. In the the latter approach, the explicitly qualified parameter-value pairs are passed to respective base models and no fitting parameters are passed to base models not explicitly listed. Note that in this pattern, propensity models are considered a nuisance model.

synchronize_cross_fitting indicates whether the learning of different base models should use exactly the same data splits where possible. Note that if there are several models to be synchronized which are classifiers, these cannot be split via stratification.

Source code in metalearners/metalearner.py

def fit(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    """Fit all models of the MetaLearner.

    If pre-fitted models were passed at instantiation, these are never refitted.

    `n_jobs_cross_fitting` will be used at the cross-fitting level and
    `n_jobs_base_learners` will be used at the stage level. `None` means 1 unless in a
    [`joblib.parallel_backend`](https://joblib.readthedocs.io/en/latest/generated/joblib.parallel_backend.html#joblib.parallel_backend)
    context. `-1` means using all processors.
    For more information about parallelism check [parallelism](parallelism.md).


    `fit_params` is an optional `dict` to be forwarded to base estimator `fit` calls. It supports
    two usages patterns:

    ```python
    fit_params={"parameter_of_interest": value_of_interest}
    ```

    ```python
    fit_params={
        "nuisance": {
            "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
            "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
        },
        "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
    }
    ```

    In the former approach, the parameter and value of interest are passed to all base models. In the
    the latter approach, the explicitly qualified parameter-value pairs are passed to respective base
    models and no fitting parameters are passed to base models not explicitly listed. Note that in this
    pattern, propensity models are considered a nuisance model.

    `synchronize_cross_fitting` indicates whether the learning of different base models should use exactly
    the same data splits where possible. Note that if there are several models to be synchronized which are
    classifiers, these cannot be split via stratification.
    """
    self.fit_all_nuisance(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    self.fit_all_treatment(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    return self

fit_all_nuisance

fit_all_nuisance(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all nuisance models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the nuisance models, and in the case they should also be used by the treatment models, these should also be passed in the following call to fit_all_treatment.

This method, combined with fit_all_treatment, facilitates the segmentation of the metalearner fitting process into two distinct parts. This division allows for interventions between the two stages, such as performing feature selection for the treatment models or conducting hyperparameter optimization within the nuisance models.

Source code in metalearners/tlearner.py

def fit_all_nuisance(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    self._validate_treatment(w)
    self._validate_outcome(y, w)

    self._treatment_variants_mask = []

    for v in range(self.n_variants):
        self._treatment_variants_mask.append(w == v)

    qualified_fit_params = self._qualified_fit_params(fit_params)

    nuisance_jobs: list[_ParallelJoblibSpecification | None] = []
    for treatment_variant in range(self.n_variants):
        mask = self._treatment_variants_mask[treatment_variant]
        X_variant = index_matrix(X, mask)
        y_variant = index_vector(y, mask)
        nuisance_jobs.append(
            self._nuisance_joblib_specifications(
                X=X_variant,
                y=y_variant,
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=treatment_variant,
                n_jobs_cross_fitting=n_jobs_cross_fitting,
                fit_params=qualified_fit_params[NUISANCE][VARIANT_OUTCOME_MODEL],
            )
        )

    parallel = Parallel(n_jobs=n_jobs_base_learners)
    results = parallel(
        delayed(_fit_cross_fit_estimator_joblib)(job)
        for job in nuisance_jobs
        if job is not None
    )
    self._assign_joblib_nuisance_results(results)
    return self

fit_all_treatment

fit_all_treatment(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all treatment models of the MetaLearner.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the treatment models, as the nuisance models should already be fitted.

Source code in metalearners/tlearner.py

def fit_all_treatment(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    return self

fit_nuisance

fit_nuisance(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit a given nuisance model of a MetaLearner.

y represents the objective of the given nuisance model, not necessarily the outcome of the experiment. If pre-fitted models were passed at instantiation, these are never refitted.

Source code in metalearners/metalearner.py

def fit_nuisance(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit a given nuisance model of a MetaLearner.

    `y` represents the objective of the given nuisance model, not necessarily the outcome of the experiment.
    If pre-fitted models were passed at instantiation, these are never refitted.
    """
    if model_kind in self._prefitted_nuisance_models:
        return self
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._nuisance_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

fit_treatment

fit_treatment(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit the treatment model of a MetaLearner.

y represents the objective of the given treatment model, not necessarily the outcome of the experiment.

Source code in metalearners/metalearner.py

def fit_treatment(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit the treatment model of a MetaLearner.

    `y` represents the objective of the given treatment model, not necessarily the outcome of the experiment.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._treatment_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

nuisance_model_specifications classmethod

nuisance_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all first-stage models.

Source code in metalearners/tlearner.py

@classmethod
def nuisance_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return {
        VARIANT_OUTCOME_MODEL: _ModelSpecifications(
            cardinality=MetaLearner._get_n_variants,
            predict_method=MetaLearner._outcome_predict_method,
        ),
    }

predict

predict(
    X, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Estimate the CATE.

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants} - 1, 1)\) if the outcome is a scalar, i.e. in case of a regression problem.

• \((n_{obs}, n_{variants} - 1, n_{classes})\) if the outcome is a class, i.e. in case of a classification problem.

In the case of multiple treatment variants, the second dimension represents the CATE of the corresponding variant vs the control (variant 0).

Source code in metalearners/tlearner.py

def predict(
    self,
    X,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    conditional_average_outcomes = self.predict_conditional_average_outcomes(
        X=X, is_oos=is_oos, oos_method=oos_method
    )

    return conditional_average_outcomes[:, 1:] - (
        conditional_average_outcomes[:, [0]]
    )

predict_conditional_average_outcomes

predict_conditional_average_outcomes(
    X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Predict the vectors of conditional average outcomes.

These are defined as \(\mathbb{E}[Y_i(w) | X]\) for each treatment variant \(w\).

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants}, 1)\) if the outcome is a scalar, i.e., in case of a regression problem.

• \((n_{obs}, n_{variants}, n_{classes})\) if the outcome is a class, i.e., in case of a classification problem.

Source code in metalearners/metalearner.py

def predict_conditional_average_outcomes(
    self, X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray:
    if self._treatment_variants_mask is None:
        raise ValueError(
            "The metalearner needs to be fitted before predicting."
            "In particular, the MetaLearner's attribute _treatment_variant_mask, "
            "typically set during fitting, is None."
        )
    # TODO: Consider multiprocessing
    n_obs = safe_len(X)
    nuisance_tensors = self._nuisance_tensors(n_obs)
    conditional_average_outcomes_list = nuisance_tensors[VARIANT_OUTCOME_MODEL]

    for tv in range(self.n_variants):
        if is_oos:
            conditional_average_outcomes_list[tv] = self.predict_nuisance(
                X=X,
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=True,
                oos_method=oos_method,
            )
        else:
            conditional_average_outcomes_list[tv][
                self._treatment_variants_mask[tv]
            ] = self.predict_nuisance(
                X=index_matrix(X, self._treatment_variants_mask[tv]),
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=False,
            )
            conditional_average_outcomes_list[tv][
                ~self._treatment_variants_mask[tv]
            ] = self.predict_nuisance(
                X=index_matrix(X, ~self._treatment_variants_mask[tv]),
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=True,
                oos_method=oos_method,
            )
    return np.stack(conditional_average_outcomes_list, axis=1).reshape(
        n_obs, self.n_variants, -1
    )

predict_nuisance

predict_nuisance(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given nuisance model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_nuisance(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given nuisance model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    predict_method_name = self.nuisance_model_specifications()[model_kind][
        "predict_method"
    ](self)
    predict_method = getattr(
        self._nuisance_models[model_kind][model_ord], predict_method_name
    )
    return predict_method(X_filtered, is_oos, oos_method)

predict_treatment

predict_treatment(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given treatment model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_treatment(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given treatment model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    return self._treatment_models[model_kind][model_ord].predict(
        X_filtered, is_oos, oos_method
    )

shap_values

shap_values(
    X: Matrix,
    shap_explainer_factory: type[Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]

Calculates the shap values for each treatment group.

If explainer is None a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters cate_estimates, cate_model_factory and cate_model_params are ignored.

The parameter shap_explainer_factory can be used to specify the type of shap explainer, for the different options see here.

The returned list contains the shap values for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def shap_values(
    self,
    X: Matrix,
    shap_explainer_factory: type[shap.Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]:
    """Calculates the shap values for each treatment group.

    If `explainer` is `None` a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    The parameter `shap_explainer_factory` can be used to specify the type of shap
    explainer, for the different options see
    [here](https://shap.readthedocs.io/en/latest/api.html#explainers).

    The returned list contains the shap values for each treatment variant in ascending
    order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=None if cate_estimates is None else X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.shap_values(
        X=X,
        shap_explainer_factory=shap_explainer_factory,
        shap_explainer_params=shap_explainer_params,
    )

treatment_model_specifications classmethod

treatment_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all second-stage models.

Source code in metalearners/tlearner.py

@classmethod
def treatment_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return dict()

utils

FixedBinaryPropensity

FixedBinaryPropensity(propensity_score: float)

Bases: ClassifierMixin, BaseEstimator

Binary classifier propensity dummy model which outputs a fixed propensity, independently of covariates.

Source code in metalearners/utils.py

def __init__(self, propensity_score: float) -> None:
    if not 0 <= propensity_score <= 1:
        raise ValueError(
            f"Expected a propensity score between 0 and 1 but got {propensity_score}."
        )
    self.propensity_score = propensity_score

metalearner_factory

metalearner_factory(
    metalearner_prefix: str,
) -> type[MetaLearner]

Returns the MetaLearner class corresponding to the given prefix.

The accepted metalearner_prefix values are:

• "S" for SLearner
• "T" for TLearner
• "X" for XLearner
• "R" for RLearner
• "DR" for DRLearner

Source code in metalearners/utils.py

def metalearner_factory(metalearner_prefix: str) -> type[MetaLearner]:
    """Returns the MetaLearner class corresponding to the given prefix.

    The accepted ``metalearner_prefix`` values are:

    * ``"S"`` for [`SLearner`][metalearners.slearner.SLearner]
    * ``"T"`` for [`TLearner`][metalearners.tlearner.TLearner]
    * ``"X"`` for [`XLearner`][metalearners.xlearner.XLearner]
    * ``"R"`` for [`RLearner`][metalearners.rlearner.RLearner]
    * ``"DR"`` for [`DRLearner`][metalearners.drlearner.DRLearner]
    """
    match metalearner_prefix:
        case "T":
            return TLearner
        case "S":
            return SLearner
        case "X":
            return XLearner
        case "R":
            return RLearner
        case "DR":
            return DRLearner
        case _:
            raise ValueError(
                f"No MetaLearner implementation found for prefix {metalearner_prefix}."
            )

simplify_output

simplify_output(tensor: ndarray) -> np.ndarray

Reduces dimensions of a CATE estimation tensor if possible.

The returned results will be of shape

• \((n_{obs})\) if there are 2 tratment variants and and the outcome is either a regression outcome or a binary classification outcome.

• \((n_{obs}, n_{classes})\) if there are 2 treatment variants and and the outcome is a classification outcome with at least 3 classes.

• \((n_{obs}, n_{variants} - 1)\) if there are at least 3 variants and the outcome is either a regression outcome or a binary classification outcome.

• \((n_{obs}, n_{variants} - 1, n_{classes})\) if there are at least 3 variants and and the outcome is a classification outcome with at least 3 classes.

Source code in metalearners/utils.py

def simplify_output(tensor: np.ndarray) -> np.ndarray:
    """Reduces dimensions of a CATE estimation tensor if possible.

    The returned results will be of shape

    * $(n_{obs})$ if there are 2 tratment variants and and the outcome is either
      a regression outcome or a binary classification outcome.

    * $(n_{obs}, n_{classes})$ if there are 2 treatment variants and and the outcome
      is a classification outcome with at least 3 classes.

    * $(n_{obs}, n_{variants} - 1)$ if there are at least 3
      variants and the outcome is either a regression outcome or a binary classification
      outcome.

    * $(n_{obs}, n_{variants} - 1, n_{classes})$ if there are at least 3
      variants and and the outcome is a classification outcome with at least 3 classes.
    """
    if (n_dim := len(tensor.shape)) != 3:
        raise ValueError(
            f"Output needs to be 3-dimensional but is {n_dim}-dimensional."
        )
    n_obs, n_variants, n_outputs = tensor.shape
    if n_variants == 1 and n_outputs == 1:
        return tensor.reshape(n_obs)
    if n_variants == 1 and n_outputs == 2:
        return tensor[:, 0, 1].reshape(n_obs)
    if n_variants == 1:
        return tensor.reshape(n_obs, n_outputs)
    if n_outputs == 1:
        return tensor.reshape(n_obs, n_variants)
    if n_outputs == 2:
        return tensor[:, :, 1].reshape(n_obs, n_variants)
    return tensor

xlearner

XLearner

XLearner(
    is_classification: bool,
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel]
    | None = None,
    nuisance_model_params: Params
    | dict[str, Params]
    | None = None,
    treatment_model_params: Params
    | dict[str, Params]
    | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[
        str, list[CrossFitEstimator]
    ]
    | None = None,
    fitted_propensity_model: CrossFitEstimator
    | None = None,
    feature_set: Features
    | dict[str, Features]
    | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
)

Bases: _ConditionalAverageOutcomeMetaLearner

X-Learner for CATE estimation as described by Kuenzel et al (2019).

Importantly, the current X-Learner implementation only supports:

• binary classes in case of a classification outcome

Source code in metalearners/metalearner.py

def __init__(
    self,
    is_classification: bool,
    # TODO: Consider whether we can make this not a state of the MetaLearner
    # but rather just a parameter of a predict call.
    n_variants: int,
    nuisance_model_factory: ModelFactory | None = None,
    treatment_model_factory: ModelFactory | None = None,
    propensity_model_factory: type[_ScikitModel] | None = None,
    nuisance_model_params: Params | dict[str, Params] | None = None,
    treatment_model_params: Params | dict[str, Params] | None = None,
    propensity_model_params: Params | None = None,
    fitted_nuisance_models: dict[str, list[CrossFitEstimator]] | None = None,
    fitted_propensity_model: CrossFitEstimator | None = None,
    feature_set: Features | dict[str, Features] | None = None,
    n_folds: int | dict[str, int] = 10,
    random_state: int | None = None,
):
    super().__init__(
        nuisance_model_factory=nuisance_model_factory,
        is_classification=is_classification,
        n_variants=n_variants,
        treatment_model_factory=treatment_model_factory,
        propensity_model_factory=propensity_model_factory,
        nuisance_model_params=nuisance_model_params,
        treatment_model_params=treatment_model_params,
        propensity_model_params=propensity_model_params,
        fitted_nuisance_models=fitted_nuisance_models,
        fitted_propensity_model=fitted_propensity_model,
        feature_set=feature_set,
        n_folds=n_folds,
        random_state=random_state,
    )
    self._treatment_variants_mask: list[np.ndarray] | None = None

init_args property

init_args: dict[str, Any]

Create initialization parameters for a new MetaLearner.

Importantly, this does not copy further internal state, such as the weights or parameters of trained base models.

evaluate

evaluate(
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]

Evaluate the MetaLearner.

The keys in scoring which are not a name of a model contained in the MetaLearner will be ignored, for information about this names check nuisance_model_specifications and treatment_model_specifications. The values must be a list of:

• string representing a sklearn scoring method. Check here for the possible values.
• Callable with signature scorer(estimator, X, y_true, **kwargs). We recommend using sklearn.metrics.make_scorer to create such a Callable.

If some model name is not present in the keys of scoring then the default used metrics will be neg_log_loss if it is a classifier and neg_root_mean_squared_error if it is a regressor.

The returned dictionary keys have the following structure:

• For nuisance models:

  • If the cardinality is one: f"{model_kind}_{scorer}"
  • If there is one model for each treatment variant (including control): f"{model_kind}_{treatment_variant}_{scorer}"

• For treatment models: f"{model_kind}_{treatment_variant}_vs_0_{scorer}"

Where scorer is the name of the scorer if it is a string and "custom_scorer_{idx}" if it is a callable where idx is the index in the scorers list.

Source code in metalearners/xlearner.py

def evaluate(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
    scoring: Scoring | None = None,
) -> dict[str, float]:
    safe_scoring = self._scoring(scoring)

    variant_outcome_evaluation = _evaluate_model_kind(
        cfes=self._nuisance_models[VARIANT_OUTCOME_MODEL],
        Xs=[index_matrix(X, w == tv) for tv in range(self.n_variants)],
        ys=[index_vector(y, w == tv) for tv in range(self.n_variants)],
        scorers=safe_scoring[VARIANT_OUTCOME_MODEL],
        model_kind=VARIANT_OUTCOME_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=False,
        feature_set=self.feature_set[VARIANT_OUTCOME_MODEL],
    )

    propensity_evaluation = _evaluate_model_kind(
        cfes=self._nuisance_models[PROPENSITY_MODEL],
        Xs=[X],
        ys=[w],
        scorers=safe_scoring[PROPENSITY_MODEL],
        model_kind=PROPENSITY_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=False,
        feature_set=self.feature_set[PROPENSITY_MODEL],
    )

    conditional_average_outcome_estimates = (
        self.predict_conditional_average_outcomes(
            X=X,
            is_oos=is_oos,
            oos_method=oos_method,
        )
    )

    imputed_te_control: list[np.ndarray] = []
    imputed_te_treatment: list[np.ndarray] = []
    for treatment_variant in range(1, self.n_variants):
        tv_imputed_te_control, tv_imputed_te_treatment = self._pseudo_outcome(
            y, w, treatment_variant, conditional_average_outcome_estimates
        )
        imputed_te_control.append(tv_imputed_te_control)
        imputed_te_treatment.append(tv_imputed_te_treatment)

    te_treatment_evaluation = _evaluate_model_kind(
        self._treatment_models[TREATMENT_EFFECT_MODEL],
        Xs=[index_matrix(X, w == tv) for tv in range(1, self.n_variants)],
        ys=imputed_te_treatment,
        scorers=safe_scoring[TREATMENT_EFFECT_MODEL],
        model_kind=TREATMENT_EFFECT_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=True,
        feature_set=self.feature_set[TREATMENT_EFFECT_MODEL],
    )

    te_control_evaluation = _evaluate_model_kind(
        self._treatment_models[CONTROL_EFFECT_MODEL],
        Xs=[index_matrix(X, w == 0) for _ in range(1, self.n_variants)],
        ys=imputed_te_control,
        scorers=safe_scoring[CONTROL_EFFECT_MODEL],
        model_kind=CONTROL_EFFECT_MODEL,
        is_oos=is_oos,
        oos_method=oos_method,
        is_treatment_model=True,
        feature_set=self.feature_set[CONTROL_EFFECT_MODEL],
    )

    return (
        variant_outcome_evaluation
        | propensity_evaluation
        | te_treatment_evaluation
        | te_control_evaluation
    )

explainer

explainer(
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer

Create an Explainer

which can be used in feature_importances.

This function can be used in two distinct manners based on the provided parameters:

• When parameters X, cate_estimates, and cate_model_factory are all set to None, the function creates an Explainer using the pre-existing treatment models. If these models do not exist, however, it triggers a ValueError.
• On the contrary, if X, cate_estimates, and cate_model_factory are not None, the function initiates an instance of the Explainer class using these parameters. This instance then fits new models for each treatment variant, and these models are employed to calculate the importance of features.

Source code in metalearners/metalearner.py

def explainer(
    self,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> Explainer:
    r"""Create an [`Explainer`][metalearners.explainer.Explainer]

    which can be used in [`feature_importances`][metalearners.metalearner.MetaLearner.feature_importances].

    This function can be used in two distinct manners based on the provided parameters:

    *   When parameters `X`, `cate_estimates`, and `cate_model_factory` are all
        set to `None`, the function creates an [`Explainer`][metalearners.explainer.Explainer]
        using the pre-existing treatment models. If these models do not exist, however,
        it triggers a `ValueError`.
    *   On the contrary, if `X`, `cate_estimates`, and `cate_model_factory` are
        not `None`, the function initiates an instance of the [`Explainer`][metalearners.explainer.Explainer]
        class using these parameters. This instance then fits new models for each
        treatment variant, and these models are employed to calculate the importance
        of features.
    """
    if X is None and cate_estimates is None and cate_model_factory is None:
        try:
            cate_cf_models = self._treatment_models[TREATMENT_MODEL]
            cate_models = [cf._overall_estimator for cf in cate_cf_models]
        except KeyError:
            raise ValueError(
                "The metalearner does not have treatment models; hence X, cate_estimates, "
                "and cate_model_factory need to be defined."
            )
        return Explainer(cate_models=cate_models)  # type: ignore
    elif (
        X is not None
        and cate_estimates is not None
        and cate_model_factory is not None
    ):
        return Explainer.from_estimates(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    else:
        raise ValueError(
            "Either all of [X, cate_estimates, cate_model_factory] are None or all "
            "of them must be defined."
        )

feature_importances

feature_importances(
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances

Calculates the feature importance for each treatment group.

If explainer is None, a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters X, cate_estimates, cate_model_factory and cate_model_params are ignored.

If normalization = True, for each treatment variant the feature importances are normalized so that they sum to 1.

feature_names is optional but in the case it’s not passed the names of the features will default to f"Feature {i}" where i is the corresponding feature index.

The returned list contains the feature importances for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def feature_importances(
    self,
    feature_names: Collection[str] | None = None,
    normalize: bool = False,
    sort_values: bool = False,
    explainer: Explainer | None = None,
    X: Matrix | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> FeatureImportances:
    r"""Calculates the feature importance for each treatment group.

    If `explainer` is `None`, a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `X`, `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    If `normalization = True`, for each treatment variant the feature importances
    are normalized so that they sum to 1.

    `feature_names` is optional but in the case it's not passed the names of the
    features will default to `f"Feature {i}"` where `i` is the corresponding
    feature index.

    The returned list contains the feature importances for each treatment variant in
    ascending order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.feature_importances(
        normalize=normalize, feature_names=feature_names, sort_values=sort_values
    )

fit

fit(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

n_jobs_cross_fitting will be used at the cross-fitting level and n_jobs_base_learners will be used at the stage level. None means 1 unless in a joblib.parallel_backend context. -1 means using all processors. For more information about parallelism check parallelism.

fit_params is an optional dict to be forwarded to base estimator fit calls. It supports two usages patterns:

fit_params={"parameter_of_interest": value_of_interest}

fit_params={
    "nuisance": {
        "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
        "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
    },
    "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
}

In the former approach, the parameter and value of interest are passed to all base models. In the the latter approach, the explicitly qualified parameter-value pairs are passed to respective base models and no fitting parameters are passed to base models not explicitly listed. Note that in this pattern, propensity models are considered a nuisance model.

synchronize_cross_fitting indicates whether the learning of different base models should use exactly the same data splits where possible. Note that if there are several models to be synchronized which are classifiers, these cannot be split via stratification.

Source code in metalearners/metalearner.py

def fit(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    """Fit all models of the MetaLearner.

    If pre-fitted models were passed at instantiation, these are never refitted.

    `n_jobs_cross_fitting` will be used at the cross-fitting level and
    `n_jobs_base_learners` will be used at the stage level. `None` means 1 unless in a
    [`joblib.parallel_backend`](https://joblib.readthedocs.io/en/latest/generated/joblib.parallel_backend.html#joblib.parallel_backend)
    context. `-1` means using all processors.
    For more information about parallelism check [parallelism](parallelism.md).


    `fit_params` is an optional `dict` to be forwarded to base estimator `fit` calls. It supports
    two usages patterns:

    ```python
    fit_params={"parameter_of_interest": value_of_interest}
    ```

    ```python
    fit_params={
        "nuisance": {
            "nuisance_model_kind1": {"parameter_of_interest1": value_of_interest1},
            "nuisance_model_kind3": {"parameter_of_interest3": value_of_interest3},
        },
        "treatment": {"treatment_model_kind1": {"parameter_of_interest4": value_of_interest4}}
    }
    ```

    In the former approach, the parameter and value of interest are passed to all base models. In the
    the latter approach, the explicitly qualified parameter-value pairs are passed to respective base
    models and no fitting parameters are passed to base models not explicitly listed. Note that in this
    pattern, propensity models are considered a nuisance model.

    `synchronize_cross_fitting` indicates whether the learning of different base models should use exactly
    the same data splits where possible. Note that if there are several models to be synchronized which are
    classifiers, these cannot be split via stratification.
    """
    self.fit_all_nuisance(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    self.fit_all_treatment(
        X=X,
        y=y,
        w=w,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        fit_params=fit_params,
        synchronize_cross_fitting=synchronize_cross_fitting,
        n_jobs_base_learners=n_jobs_base_learners,
    )

    return self

fit_all_nuisance

fit_all_nuisance(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all nuisance models of the MetaLearner.

If pre-fitted models were passed at instantiation, these are never refitted.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the nuisance models, and in the case they should also be used by the treatment models, these should also be passed in the following call to fit_all_treatment.

This method, combined with fit_all_treatment, facilitates the segmentation of the metalearner fitting process into two distinct parts. This division allows for interventions between the two stages, such as performing feature selection for the treatment models or conducting hyperparameter optimization within the nuisance models.

Source code in metalearners/xlearner.py

def fit_all_nuisance(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    self._validate_treatment(w)
    self._validate_outcome(y, w)

    self._treatment_variants_mask = []

    qualified_fit_params = self._qualified_fit_params(fit_params)

    self._cvs: list = []

    for treatment_variant in range(self.n_variants):
        self._treatment_variants_mask.append(w == treatment_variant)
        if synchronize_cross_fitting:
            cv_split_indices = self._split(
                index_matrix(X, self._treatment_variants_mask[treatment_variant])
            )
        else:
            cv_split_indices = None
        self._cvs.append(cv_split_indices)

    nuisance_jobs: list[_ParallelJoblibSpecification | None] = []
    for treatment_variant in range(self.n_variants):
        mask = self._treatment_variants_mask[treatment_variant]
        nuisance_jobs.append(
            self._nuisance_joblib_specifications(
                X=index_matrix(X, mask),
                y=index_vector(y, mask),
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=treatment_variant,
                n_jobs_cross_fitting=n_jobs_cross_fitting,
                fit_params=qualified_fit_params[NUISANCE][VARIANT_OUTCOME_MODEL],
                cv=self._cvs[treatment_variant],
            )
        )

    nuisance_jobs.append(
        self._nuisance_joblib_specifications(
            X=X,
            y=w,
            model_kind=PROPENSITY_MODEL,
            model_ord=0,
            n_jobs_cross_fitting=n_jobs_cross_fitting,
            fit_params=qualified_fit_params[NUISANCE][PROPENSITY_MODEL],
        )
    )

    parallel = Parallel(n_jobs=n_jobs_base_learners)
    results = parallel(
        delayed(_fit_cross_fit_estimator_joblib)(job)
        for job in nuisance_jobs
        if job is not None
    )
    self._assign_joblib_nuisance_results(results)

    return self

fit_all_treatment

fit_all_treatment(
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self

Fit all treatment models of the MetaLearner.

The only difference with fit parameters, is that if fit_params follows the first usage pattern (explained in fit), then the training parameters will only be used for the treatment models, as the nuisance models should already be fitted.

Source code in metalearners/xlearner.py

def fit_all_treatment(
    self,
    X: Matrix,
    y: Vector,
    w: Vector,
    n_jobs_cross_fitting: int | None = None,
    fit_params: dict | None = None,
    synchronize_cross_fitting: bool = True,
    n_jobs_base_learners: int | None = None,
) -> Self:
    if self._treatment_variants_mask is None:
        raise ValueError(
            "The nuisance models need to be fitted before fitting the treatment models."
            "In particular, the MetaLearner's attribute _treatment_variant_mask, "
            "typically set during nuisance fitting, is None."
        )
    if not hasattr(self, "_cvs"):
        raise ValueError(
            "The nuisance models need to be fitted before fitting the treatment models."
            "In particular, the MetaLearner's attribute _cvs, "
            "typically set during nuisance fitting, does not exist."
        )
    qualified_fit_params = self._qualified_fit_params(fit_params)

    treatment_jobs: list[_ParallelJoblibSpecification] = []

    conditional_average_outcome_estimates = (
        self.predict_conditional_average_outcomes(
            X=X,
            is_oos=False,
        )
    )

    for treatment_variant in range(1, self.n_variants):
        imputed_te_control, imputed_te_treatment = self._pseudo_outcome(
            y, w, treatment_variant, conditional_average_outcome_estimates
        )
        treatment_jobs.append(
            self._treatment_joblib_specifications(
                X=index_matrix(X, self._treatment_variants_mask[treatment_variant]),
                y=imputed_te_treatment,
                model_kind=TREATMENT_EFFECT_MODEL,
                model_ord=treatment_variant - 1,
                n_jobs_cross_fitting=n_jobs_cross_fitting,
                fit_params=qualified_fit_params[TREATMENT][TREATMENT_EFFECT_MODEL],
                cv=self._cvs[treatment_variant],
            )
        )

        treatment_jobs.append(
            self._treatment_joblib_specifications(
                X=index_matrix(X, self._treatment_variants_mask[0]),
                y=imputed_te_control,
                model_kind=CONTROL_EFFECT_MODEL,
                model_ord=treatment_variant - 1,
                n_jobs_cross_fitting=n_jobs_cross_fitting,
                fit_params=qualified_fit_params[TREATMENT][CONTROL_EFFECT_MODEL],
                cv=self._cvs[0],
            )
        )

    parallel = Parallel(n_jobs=n_jobs_base_learners)
    results = parallel(
        delayed(_fit_cross_fit_estimator_joblib)(job) for job in treatment_jobs
    )
    self._assign_joblib_treatment_results(results)
    return self

fit_nuisance

fit_nuisance(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit a given nuisance model of a MetaLearner.

y represents the objective of the given nuisance model, not necessarily the outcome of the experiment. If pre-fitted models were passed at instantiation, these are never refitted.

Source code in metalearners/metalearner.py

def fit_nuisance(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit a given nuisance model of a MetaLearner.

    `y` represents the objective of the given nuisance model, not necessarily the outcome of the experiment.
    If pre-fitted models were passed at instantiation, these are never refitted.
    """
    if model_kind in self._prefitted_nuisance_models:
        return self
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._nuisance_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

fit_treatment

fit_treatment(
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self

Fit the treatment model of a MetaLearner.

y represents the objective of the given treatment model, not necessarily the outcome of the experiment.

Source code in metalearners/metalearner.py

def fit_treatment(
    self,
    X: Matrix,
    y: Vector,
    model_kind: str,
    model_ord: int,
    fit_params: dict | None = None,
    n_jobs_cross_fitting: int | None = None,
    cv: SplitIndices | None = None,
) -> Self:
    """Fit the treatment model of a MetaLearner.

    `y` represents the objective of the given treatment model, not necessarily the outcome of the experiment.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    self._treatment_models[model_kind][model_ord].fit(
        X_filtered,
        y,
        fit_params=fit_params,
        n_jobs_cross_fitting=n_jobs_cross_fitting,
        cv=cv,
    )
    return self

nuisance_model_specifications classmethod

nuisance_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all first-stage models.

Source code in metalearners/xlearner.py

@classmethod
def nuisance_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return {
        VARIANT_OUTCOME_MODEL: _ModelSpecifications(
            cardinality=MetaLearner._get_n_variants,
            predict_method=MetaLearner._outcome_predict_method,
        ),
        PROPENSITY_MODEL: _ModelSpecifications(
            cardinality=get_one,
            predict_method=get_predict_proba,
        ),
    }

predict

predict(
    X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Estimate the CATE.

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants} - 1, 1)\) if the outcome is a scalar, i.e. in case of a regression problem.

• \((n_{obs}, n_{variants} - 1, n_{classes})\) if the outcome is a class, i.e. in case of a classification problem.

In the case of multiple treatment variants, the second dimension represents the CATE of the corresponding variant vs the control (variant 0).

Source code in metalearners/xlearner.py

def predict(
    self,
    X: Matrix,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    if self._treatment_variants_mask is None:
        raise ValueError(
            "The MetaLearner needs to be fitted before predicting. "
            "In particular, the X-Learner's attribute _treatment_variant_mask, "
            "typically set during fitting, is None."
        )
    n_outputs = 2 if self.is_classification else 1
    tau_hat = np.zeros((safe_len(X), self.n_variants - 1, n_outputs))
    # Propensity score model is always a classifier so we can't use MEDIAN
    propensity_score_oos = OVERALL if oos_method == MEDIAN else oos_method
    propensity_score = self.predict_nuisance(
        X=X,
        model_kind=PROPENSITY_MODEL,
        model_ord=0,
        is_oos=is_oos,
        oos_method=propensity_score_oos,
    )

    control_indices = self._treatment_variants_mask[0]
    non_control_indices = ~control_indices

    for treatment_variant in range(1, self.n_variants):
        treatment_variant_mask = self._treatment_variants_mask[treatment_variant]
        non_treatment_variant_mask = ~treatment_variant_mask
        if is_oos:
            tau_hat_treatment = self.predict_treatment(
                X=X,
                model_kind=TREATMENT_EFFECT_MODEL,
                model_ord=treatment_variant - 1,
                is_oos=is_oos,
                oos_method=oos_method,
            )
            tau_hat_control = self.predict_treatment(
                X=X,
                model_kind=CONTROL_EFFECT_MODEL,
                model_ord=treatment_variant - 1,
                is_oos=is_oos,
                oos_method=oos_method,
            )
        else:
            tau_hat_treatment = np.zeros(safe_len(X))
            tau_hat_control = np.zeros(safe_len(X))

            tau_hat_treatment[non_treatment_variant_mask] = self.predict_treatment(
                X=index_matrix(X, non_treatment_variant_mask),
                model_kind=TREATMENT_EFFECT_MODEL,
                model_ord=treatment_variant - 1,
                is_oos=True,
                oos_method=oos_method,
            )

            tau_hat_treatment[treatment_variant_mask] = self.predict_treatment(
                X=index_matrix(X, treatment_variant_mask),
                model_kind=TREATMENT_EFFECT_MODEL,
                model_ord=treatment_variant - 1,
                is_oos=False,
            )
            tau_hat_control[control_indices] = self.predict_treatment(
                X=index_matrix(X, control_indices),
                model_kind=CONTROL_EFFECT_MODEL,
                model_ord=treatment_variant - 1,
                is_oos=False,
            )
            tau_hat_control[non_control_indices] = self.predict_treatment(
                X=index_matrix(X, non_control_indices),
                model_kind=CONTROL_EFFECT_MODEL,
                model_ord=treatment_variant - 1,
                is_oos=True,
                oos_method=oos_method,
            )

        propensity_score_treatment = propensity_score[:, treatment_variant] / (
            propensity_score[:, 0] + propensity_score[:, treatment_variant]
        )

        tau_hat_treatment_variant = (
            propensity_score_treatment * tau_hat_control
            + (1 - propensity_score_treatment) * tau_hat_treatment
        )

        if self.is_classification:
            # This is to be consistent with other MetaLearners (e.g. S and T) that automatically
            # work with multiclass outcomes and return the CATE estimate for each class. As the X-Learner only
            # works with binary classes (the pseudo outcome formula does not make sense with
            # multiple classes unless some adaptation is done) we can manually infer the
            # CATE estimate for the complementary class  -- returning a matrix of shape (N, 2).
            tau_hat_treatment_variant = np.stack(
                [-tau_hat_treatment_variant, tau_hat_treatment_variant], axis=1
            )
        else:
            tau_hat_treatment_variant = np.expand_dims(tau_hat_treatment_variant, 1)

        tau_hat[:, treatment_variant - 1] = tau_hat_treatment_variant

    return tau_hat

predict_conditional_average_outcomes

predict_conditional_average_outcomes(
    X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray

Predict the vectors of conditional average outcomes.

These are defined as \(\mathbb{E}[Y_i(w) | X]\) for each treatment variant \(w\).

If is_oos, an acronym for ‘is out of sample’, is False, the estimates will stem from cross-fitting. Otherwise, various approaches exist, specified via oos_method.

The returned ndarray is of shape:

• \((n_{obs}, n_{variants}, 1)\) if the outcome is a scalar, i.e., in case of a regression problem.

• \((n_{obs}, n_{variants}, n_{classes})\) if the outcome is a class, i.e., in case of a classification problem.

Source code in metalearners/metalearner.py

def predict_conditional_average_outcomes(
    self, X: Matrix, is_oos: bool, oos_method: OosMethod = OVERALL
) -> np.ndarray:
    if self._treatment_variants_mask is None:
        raise ValueError(
            "The metalearner needs to be fitted before predicting."
            "In particular, the MetaLearner's attribute _treatment_variant_mask, "
            "typically set during fitting, is None."
        )
    # TODO: Consider multiprocessing
    n_obs = safe_len(X)
    nuisance_tensors = self._nuisance_tensors(n_obs)
    conditional_average_outcomes_list = nuisance_tensors[VARIANT_OUTCOME_MODEL]

    for tv in range(self.n_variants):
        if is_oos:
            conditional_average_outcomes_list[tv] = self.predict_nuisance(
                X=X,
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=True,
                oos_method=oos_method,
            )
        else:
            conditional_average_outcomes_list[tv][
                self._treatment_variants_mask[tv]
            ] = self.predict_nuisance(
                X=index_matrix(X, self._treatment_variants_mask[tv]),
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=False,
            )
            conditional_average_outcomes_list[tv][
                ~self._treatment_variants_mask[tv]
            ] = self.predict_nuisance(
                X=index_matrix(X, ~self._treatment_variants_mask[tv]),
                model_kind=VARIANT_OUTCOME_MODEL,
                model_ord=tv,
                is_oos=True,
                oos_method=oos_method,
            )
    return np.stack(conditional_average_outcomes_list, axis=1).reshape(
        n_obs, self.n_variants, -1
    )

predict_nuisance

predict_nuisance(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given nuisance model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_nuisance(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given nuisance model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    predict_method_name = self.nuisance_model_specifications()[model_kind][
        "predict_method"
    ](self)
    predict_method = getattr(
        self._nuisance_models[model_kind][model_ord], predict_method_name
    )
    return predict_method(X_filtered, is_oos, oos_method)

predict_treatment

predict_treatment(
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray

Estimate based on a given treatment model.

Importantly, this method needs to implement the subselection of X based on the feature_set field of MetaLearner.

Source code in metalearners/metalearner.py

def predict_treatment(
    self,
    X: Matrix,
    model_kind: str,
    model_ord: int,
    is_oos: bool,
    oos_method: OosMethod = OVERALL,
) -> np.ndarray:
    """Estimate based on a given treatment model.

    Importantly, this method needs to implement the subselection of `X` based on
    the `feature_set` field of `MetaLearner`.
    """
    X_filtered = _filter_x_columns(X, self.feature_set[model_kind])
    return self._treatment_models[model_kind][model_ord].predict(
        X_filtered, is_oos, oos_method
    )

shap_values

shap_values(
    X: Matrix,
    shap_explainer_factory: type[Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]

Calculates the shap values for each treatment group.

If explainer is None a new Explainer is created using MetaLearner.explainer with the passed parameters. If explainer is not None, then the parameters cate_estimates, cate_model_factory and cate_model_params are ignored.

The parameter shap_explainer_factory can be used to specify the type of shap explainer, for the different options see here.

The returned list contains the shap values for each treatment variant in ascending order.

Source code in metalearners/metalearner.py

def shap_values(
    self,
    X: Matrix,
    shap_explainer_factory: type[shap.Explainer],
    shap_explainer_params: dict | None = None,
    explainer: Explainer | None = None,
    cate_estimates: np.ndarray | None = None,
    cate_model_factory: type[_ScikitModel] | None = None,
    cate_model_params: Params | None = None,
) -> list[np.ndarray]:
    """Calculates the shap values for each treatment group.

    If `explainer` is `None` a new [`Explainer`][metalearners.explainer.Explainer]
    is created using [`MetaLearner.explainer`][metalearners.metalearner.MetaLearner.explainer]
    with the passed parameters. If `explainer` is not `None`, then the parameters
    `cate_estimates`, `cate_model_factory` and `cate_model_params` are
    ignored.

    The parameter `shap_explainer_factory` can be used to specify the type of shap
    explainer, for the different options see
    [here](https://shap.readthedocs.io/en/latest/api.html#explainers).

    The returned list contains the shap values for each treatment variant in ascending
    order.
    """
    if explainer is None:
        explainer = self.explainer(
            X=None if cate_estimates is None else X,
            cate_estimates=cate_estimates,
            cate_model_factory=cate_model_factory,
            cate_model_params=cate_model_params,
        )
    return explainer.shap_values(
        X=X,
        shap_explainer_factory=shap_explainer_factory,
        shap_explainer_params=shap_explainer_params,
    )

treatment_model_specifications classmethod

treatment_model_specifications() -> dict[
    str, _ModelSpecifications
]

Return the specifications of all second-stage models.

Source code in metalearners/xlearner.py

@classmethod
def treatment_model_specifications(cls) -> dict[str, _ModelSpecifications]:
    return {
        CONTROL_EFFECT_MODEL: _ModelSpecifications(
            cardinality=MetaLearner._get_n_variants_minus_one,
            predict_method=get_predict,
        ),
        TREATMENT_EFFECT_MODEL: _ModelSpecifications(
            cardinality=MetaLearner._get_n_variants_minus_one,
            predict_method=get_predict,
        ),
    }