Add XGBTClassifier and XGBRClassifier meta-learners (#824)

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

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
+        )

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)

tests/test_meta_learners.py

@@ -24,13 +24,15 @@
     BaseTClassifier,
     XGBTRegressor,
     MLPTRegressor,
+    XGBTClassifier,
 )
 from causalml.inference.meta import BaseXLearner, BaseXClassifier, BaseXRegressor
 from causalml.inference.meta import (
     BaseRLearner,
     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.