From 20e82ffe932f4654ef416bb19de703d2560aee4c Mon Sep 17 00:00:00 2001 From: Jeong-Yoon Lee Date: Tue, 14 Jul 2026 10:33:03 -0700 Subject: [PATCH] Add XGBTClassifier and XGBRClassifier meta-learners (#824) Add XGBoost-backed classifier variants of the T- and R-learners, mirroring the existing XGBTRegressor / XGBRRegressor. Both follow the post-#912 scikit-learn BaseEstimator convention: constructor arguments are stored verbatim and the XGBoost models are constructed in fit(), so get_params() / clone() round-trip correctly (verified by clone tests). XGBTClassifier builds XGBClassifier outcome models; XGBRClassifier uses an XGBClassifier outcome model (predict_proba drives the outcome cross-fit) and an XGBRegressor effect model, initializing through the grandparent BaseRLearner to bypass BaseRClassifier's both-learners-None guard. Both are re-exported from causalml.inference.meta. Reimplements the work of @jbbqqf in #899 (CLA-blocked). Closes #824. --- causalml/inference/meta/__init__.py | 9 +- causalml/inference/meta/rlearner.py | 68 ++++++++++++++- causalml/inference/meta/tlearner.py | 33 ++++++- tests/test_meta_learners.py | 130 ++++++++++++++++++++++++++++ 4 files changed, 237 insertions(+), 3 deletions(-) diff --git a/causalml/inference/meta/__init__.py b/causalml/inference/meta/__init__.py index 562f74eb..fdaea33a 100644 --- a/causalml/inference/meta/__init__.py +++ b/causalml/inference/meta/__init__.py @@ -5,9 +5,16 @@ BaseTLearner, BaseTRegressor, BaseTClassifier, + XGBTClassifier, ) from .xlearner import BaseXLearner, BaseXRegressor, BaseXClassifier -from .rlearner import BaseRLearner, BaseRRegressor, BaseRClassifier, XGBRRegressor +from .rlearner import ( + BaseRLearner, + BaseRRegressor, + BaseRClassifier, + XGBRRegressor, + XGBRClassifier, +) from .tmle import TMLELearner from .drlearner import BaseDRLearner, BaseDRRegressor, BaseDRClassifier, XGBDRRegressor from .serialization import load_learner, CausalMLVersionMismatchWarning diff --git a/causalml/inference/meta/rlearner.py b/causalml/inference/meta/rlearner.py index de9e3863..300bf6e1 100644 --- a/causalml/inference/meta/rlearner.py +++ b/causalml/inference/meta/rlearner.py @@ -4,7 +4,7 @@ from tqdm import tqdm from scipy.stats import norm from sklearn.model_selection import cross_val_predict, KFold, train_test_split -from xgboost import XGBRegressor +from xgboost import XGBRegressor, XGBClassifier from causalml.inference.meta.base import BaseLearner from causalml.inference.meta.utils import ( @@ -847,3 +847,69 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True): self.vars_c[group] = get_weighted_variance(diff_c, sample_weight_filt_c) self.vars_t[group] = get_weighted_variance(diff_t, sample_weight_filt_t) return self + + +class XGBRClassifier(BaseRClassifier): + """An R-learner classifier using XGBoost models. + + The outcome model is an ``XGBClassifier`` (its ``predict_proba`` drives the + outcome cross-fit) and the effect model is an ``XGBRegressor``. Every + constructor argument is stored verbatim (scikit-learn convention) so that + ``get_params()`` / ``clone()`` work correctly; the XGBoost models are + constructed in ``fit()``. + """ + + def __init__( + self, + propensity_learner=ElasticNetPropensityModel(), + ate_alpha=0.05, + control_name=0, + n_fold=5, + random_state=None, + outcome_xgb_kwargs=None, + effect_xgb_kwargs=None, + ): + """Initialize an R-learner classifier with XGBoost models. + + Args: + propensity_learner (optional): a propensity model. Defaults to + ``ElasticNetPropensityModel()``. + ate_alpha (float, optional): the confidence level alpha of the ATE estimate + control_name (str or int, optional): name of control group + n_fold (int, optional): CV folds for the outcome learner + random_state (int or RandomState, optional): random seed + outcome_xgb_kwargs (dict, optional): keyword arguments forwarded verbatim + to the ``XGBClassifier`` outcome model, e.g. ``{'max_depth': 4}``. + effect_xgb_kwargs (dict, optional): keyword arguments forwarded verbatim + to the ``XGBRegressor`` effect model. + + Note: all arguments are stored verbatim (scikit-learn convention) so that + ``get_params`` / ``clone`` work correctly. XGBoost model construction is + deferred to ``fit()``. + """ + # Store verbatim — no XGBoost construction here. + self.outcome_xgb_kwargs = outcome_xgb_kwargs + self.effect_xgb_kwargs = effect_xgb_kwargs + + # BaseRClassifier.__init__ rejects both learners being None; bypass it by + # initializing through the grandparent, since the learners are built in + # fit() rather than passed in. + BaseRLearner.__init__( + self, + learner=None, + outcome_learner=None, + effect_learner=None, + propensity_learner=propensity_learner, + ate_alpha=ate_alpha, + control_name=control_name, + n_fold=n_fold, + random_state=random_state, + ) + + def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True): + """Build the XGBoost models, then fit as an R-learner classifier.""" + self.outcome_learner = XGBClassifier(**(self.outcome_xgb_kwargs or {})) + self.effect_learner = XGBRegressor(**(self.effect_xgb_kwargs or {})) + return BaseRClassifier.fit( + self, X, treatment, y, p=p, sample_weight=sample_weight, verbose=verbose + ) diff --git a/causalml/inference/meta/tlearner.py b/causalml/inference/meta/tlearner.py index d9d555da..c4214905 100644 --- a/causalml/inference/meta/tlearner.py +++ b/causalml/inference/meta/tlearner.py @@ -12,7 +12,7 @@ else: from sklearn.utils.testing import ignore_warnings from tqdm import tqdm -from xgboost import XGBRegressor +from xgboost import XGBRegressor, XGBClassifier from causalml.inference.meta.base import BaseLearner from causalml.inference.meta.utils import ( @@ -547,3 +547,34 @@ def __init__(self, ate_alpha=0.05, control_name=0, *args, **kwargs): ate_alpha=ate_alpha, control_name=control_name, ) + + +class XGBTClassifier(BaseTClassifier): + """A T-learner classifier using XGBoost models. + + Stores XGBoost hyperparameters verbatim (scikit-learn convention) so that + ``get_params()`` / ``clone()`` work correctly; the ``XGBClassifier`` models + are constructed in ``fit()``. + """ + + def __init__(self, ate_alpha=0.05, control_name=0, xgb_kwargs=None): + """Initialize a T-learner classifier with two XGBoost models. + + Args: + ate_alpha (float, optional): the confidence level alpha of the ATE estimate + control_name (str or int, optional): name of control group + xgb_kwargs (dict, optional): keyword arguments forwarded verbatim to each + ``XGBClassifier``, e.g. ``xgb_kwargs={'max_depth': 4, 'learning_rate': 0.05}``. + + Note: all arguments are stored verbatim (scikit-learn convention) so that + ``get_params`` / ``clone`` work correctly. ``XGBClassifier`` construction + is deferred to ``fit()``. + """ + # Store verbatim — no XGBClassifier construction here. + self.xgb_kwargs = xgb_kwargs + super().__init__(ate_alpha=ate_alpha, control_name=control_name) + + def fit(self, X, treatment, y, *args, **kwargs): + """Build the XGBoost outcome model, then fit as a T-learner classifier.""" + self.learner = XGBClassifier(**(self.xgb_kwargs or {})) + return super().fit(X, treatment, y, *args, **kwargs) diff --git a/tests/test_meta_learners.py b/tests/test_meta_learners.py index 0a74a2dd..efbf551c 100644 --- a/tests/test_meta_learners.py +++ b/tests/test_meta_learners.py @@ -24,6 +24,7 @@ BaseTClassifier, XGBTRegressor, MLPTRegressor, + XGBTClassifier, ) from causalml.inference.meta import BaseXLearner, BaseXClassifier, BaseXRegressor from causalml.inference.meta import ( @@ -31,6 +32,7 @@ BaseRClassifier, BaseRRegressor, XGBRRegressor, + XGBRClassifier, ) from causalml.inference.meta import TMLELearner from causalml.inference.meta import BaseDRLearner @@ -1988,6 +1990,8 @@ def _params_repr_equal(p1, p2): (BaseSClassifier, {"learner": _DT_CLF}), (BaseTClassifier, {"learner": _DT_CLF}), (BaseXClassifier, {"outcome_learner": _DT_CLF, "effect_learner": _DT_REG}), + (XGBTClassifier, {"xgb_kwargs": {"n_estimators": 5}}), + (XGBRClassifier, {"outcome_xgb_kwargs": {"n_estimators": 5}}), ] @@ -2050,6 +2054,132 @@ def test_xgb_rregressor_clone_with_kwargs(): assert _params_repr_equal(r2.get_params(), r.get_params()) +def test_XGBTClassifier(generate_classification_data): + np.random.seed(RANDOM_SEED) + + df, x_names = generate_classification_data() + df["treatment_group_key"] = np.where( + df["treatment_group_key"] == CONTROL_NAME, 0, 1 + ) + + df_train, df_test = train_test_split(df, test_size=0.2, random_state=RANDOM_SEED) + + uplift_model = XGBTClassifier(xgb_kwargs={"n_estimators": 50}) + uplift_model.fit( + X=df_train[x_names].values, + treatment=df_train["treatment_group_key"].values, + y=df_train[CONVERSION].values, + ) + + # The outcome models are XGBClassifiers built in fit(), with xgb_kwargs forwarded. + assert isinstance(uplift_model.model_c, XGBClassifier) + assert uplift_model.model_c.n_estimators == 50 + for g in uplift_model.t_groups: + assert isinstance(uplift_model.models_t[g], XGBClassifier) + + tau_pred = uplift_model.predict( + X=df_test[x_names].values, treatment=df_test["treatment_group_key"].values + ) + + auuc_metrics = pd.DataFrame( + { + "tau_pred": tau_pred.flatten(), + "W": df_test["treatment_group_key"].values, + CONVERSION: df_test[CONVERSION].values, + "treatment_effect_col": df_test["treatment_effect"].values, + } + ) + auuc = auuc_score( + auuc_metrics, + outcome_col=CONVERSION, + treatment_col="W", + treatment_effect_col="treatment_effect_col", + normalize=True, + ) + assert auuc["tau_pred"] > 0.5 + + +def test_XGBRClassifier(generate_classification_data): + np.random.seed(RANDOM_SEED) + + df, x_names = generate_classification_data() + df["treatment_group_key"] = np.where( + df["treatment_group_key"] == CONTROL_NAME, 0, 1 + ) + + propensity_model = LogisticRegression() + propensity_model.fit(X=df[x_names].values, y=df["treatment_group_key"].values) + df["propensity_score"] = propensity_model.predict_proba(df[x_names].values)[:, 1] + + df_train, df_test = train_test_split(df, test_size=0.2, random_state=RANDOM_SEED) + + uplift_model = XGBRClassifier( + outcome_xgb_kwargs={"n_estimators": 50}, + effect_xgb_kwargs={"n_estimators": 50}, + ) + uplift_model.fit( + X=df_train[x_names].values, + p=df_train["propensity_score"].values, + treatment=df_train["treatment_group_key"].values, + y=df_train[CONVERSION].values, + ) + + # Outcome model is an XGBClassifier, effect model an XGBRegressor; kwargs forwarded. + assert isinstance(uplift_model.model_mu, XGBClassifier) + assert isinstance(uplift_model.model_tau, XGBRegressor) + assert uplift_model.model_mu.n_estimators == 50 + assert uplift_model.model_tau.n_estimators == 50 + + tau_pred = uplift_model.predict( + X=df_test[x_names].values, p=df_test["propensity_score"].values + ) + + auuc_metrics = pd.DataFrame( + { + "tau_pred": tau_pred.flatten(), + "W": df_test["treatment_group_key"].values, + CONVERSION: df_test[CONVERSION].values, + "treatment_effect_col": df_test["treatment_effect"].values, + } + ) + auuc = auuc_score( + auuc_metrics, + outcome_col=CONVERSION, + treatment_col="W", + treatment_effect_col="treatment_effect_col", + normalize=True, + ) + assert auuc["tau_pred"] > 0.5 + + +def test_xgbt_classifier_clone(): + """XGBTClassifier round-trips through clone() / get_params() (unfitted).""" + m = XGBTClassifier(xgb_kwargs={"max_depth": 3, "n_estimators": 5}) + assert m.get_params()["xgb_kwargs"] == {"max_depth": 3, "n_estimators": 5} + m2 = clone(m) + assert type(m2) is XGBTClassifier + assert m2.xgb_kwargs == {"max_depth": 3, "n_estimators": 5} + assert _params_repr_equal(m2.get_params(), m.get_params()) + assert not hasattr(m2, "t_groups") + + +def test_xgbr_classifier_clone(): + """XGBRClassifier round-trips through clone() / get_params() (unfitted).""" + m = XGBRClassifier( + outcome_xgb_kwargs={"max_depth": 3}, + effect_xgb_kwargs={"n_estimators": 5}, + ) + params = m.get_params() + assert params["outcome_xgb_kwargs"] == {"max_depth": 3} + assert params["effect_xgb_kwargs"] == {"n_estimators": 5} + m2 = clone(m) + assert type(m2) is XGBRClassifier + assert m2.outcome_xgb_kwargs == {"max_depth": 3} + assert m2.effect_xgb_kwargs == {"n_estimators": 5} + assert _params_repr_equal(m2.get_params(), m.get_params()) + assert not hasattr(m2, "t_groups") + + def test_xgb_rregressor_fit_predict_return_ci(): """XGBRRegressor.fit_predict(return_ci=True) exercises clone(self) in bootstrap.