From 8388e60366b7e0355a0dfd1d5ce72ccf701c2aec Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Thu, 18 Jun 2026 23:48:23 +0530 Subject: [PATCH 1/2] preserve categorical dtypes in meta-learners --- causalml/inference/meta/tlearner.py | 10 +++---- causalml/inference/meta/utils.py | 9 ++++++- tests/test_meta_learners.py | 42 +++++++++++++++++++++++++++++ 3 files changed, 55 insertions(+), 6 deletions(-) diff --git a/causalml/inference/meta/tlearner.py b/causalml/inference/meta/tlearner.py index 04ca796f..8f313650 100644 --- a/causalml/inference/meta/tlearner.py +++ b/causalml/inference/meta/tlearner.py @@ -77,7 +77,7 @@ def fit(self, X, treatment, y, p=None): treatment (np.array or pd.Series): a treatment vector y (np.array or pd.Series): an outcome vector """ - X, treatment, y = convert_pd_to_np(X, treatment, y) + treatment, y = convert_pd_to_np(treatment, y) check_treatment_vector(treatment, self.control_name) self.t_groups = np.unique(treatment[treatment != self.control_name]) self.t_groups.sort() @@ -88,7 +88,7 @@ def fit(self, X, treatment, y, p=None): for group in self.t_groups: mask = (treatment == group) | (treatment == self.control_name) treatment_filt = treatment[mask] - X_filt = X[mask] + X_filt = X[mask].reset_index(drop=True) if hasattr(X, "loc") else X[mask] y_filt = y[mask] w = (treatment_filt == group).astype(int) @@ -109,7 +109,7 @@ def predict( Returns: (numpy.ndarray): Predictions of treatment effects. """ - X, treatment, y = convert_pd_to_np(X, treatment, y) + treatment, y = convert_pd_to_np(treatment, y) yhat_cs = {} yhat_ts = {} @@ -169,7 +169,7 @@ def fit_predict( If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment], UB [n_samples, n_treatment] """ - X, treatment, y = convert_pd_to_np(X, treatment, y) + treatment, y = convert_pd_to_np(treatment, y) self.fit(X, treatment, y) te = self.predict(X, treatment, y, return_components=return_components) @@ -226,7 +226,7 @@ def estimate_ate( The mean and confidence interval (LB, UB) of the ATE estimate. pretrain (bool): whether a model has been fit, default False. """ - X, treatment, y = convert_pd_to_np(X, treatment, y) + treatment, y = convert_pd_to_np(treatment, y) if pretrain: te, yhat_cs, yhat_ts = self.predict(X, treatment, y, return_components=True) else: diff --git a/causalml/inference/meta/utils.py b/causalml/inference/meta/utils.py index 157eeaf6..d94f1c80 100644 --- a/causalml/inference/meta/utils.py +++ b/causalml/inference/meta/utils.py @@ -6,7 +6,14 @@ def convert_pd_to_np(*args): - output = [obj.to_numpy() if hasattr(obj, "to_numpy") else obj for obj in args] + def _convert(obj): + if isinstance(obj, pd.DataFrame) and any( + pd.api.types.is_categorical_dtype(obj[c]) for c in obj.columns + ): + return obj # pass through so learners can handle categoricals natively + return obj.to_numpy() if hasattr(obj, "to_numpy") else obj + + output = [_convert(obj) for obj in args] return output if len(output) > 1 else output[0] diff --git a/tests/test_meta_learners.py b/tests/test_meta_learners.py index d5a60216..f19ee4a4 100644 --- a/tests/test_meta_learners.py +++ b/tests/test_meta_learners.py @@ -1220,3 +1220,45 @@ def test_BaseDRClassifier(generate_classification_data): te_separate = learner_separate.fit_predict(X=X, treatment=treatment, y=y) assert te_separate.shape == te.shape + + +def test_BaseTLearner_with_categorical_features(): + np.random.seed(RANDOM_SEED) + n = 200 + + X = pd.DataFrame( + { + "num1": np.random.randn(n), + "num2": np.random.randn(n), + "cat1": pd.Categorical(np.random.choice([0, 1, 2], size=n)), + } + ) + treatment = np.random.binomial(1, 0.5, n) + y = X["num1"].values + (treatment * 0.5) + np.random.randn(n) * 0.1 + + learner = BaseTRegressor(learner=XGBRegressor(enable_categorical=True)) + learner.fit(X=X, treatment=treatment, y=y) + te = learner.predict(X=X) + + assert te.shape == (n, 1) + + +def test_BaseRLearner_with_categorical_features(): + np.random.seed(RANDOM_SEED) + n = 200 + + X = pd.DataFrame( + { + "num1": np.random.randn(n), + "num2": np.random.randn(n), + "cat1": pd.Categorical(np.random.choice([0, 1, 2], size=n)), + } + ) + treatment = np.random.binomial(1, 0.5, n) + y = X["num1"].values + (treatment * 0.5) + np.random.randn(n) * 0.1 + + learner = BaseRRegressor(learner=XGBRegressor(enable_categorical=True)) + learner.fit(X=X, treatment=treatment, y=y) + te = learner.predict(X=X) + + assert te.shape == (n, 1) From 1f044c346f30204476667d670dcc72e35be19a04 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Sat, 20 Jun 2026 12:19:55 +0530 Subject: [PATCH 2/2] fix the build errors --- causalml/inference/meta/tlearner.py | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/causalml/inference/meta/tlearner.py b/causalml/inference/meta/tlearner.py index 7ff0a78e..6cdff84f 100644 --- a/causalml/inference/meta/tlearner.py +++ b/causalml/inference/meta/tlearner.py @@ -129,7 +129,12 @@ def fit( # re-calling fit() always starts from a clean state (safe with warm_start). control_mask = treatment == self.control_name self.model_c = deepcopy(self._model_c_template) - self.model_c.fit(X[control_mask], y[control_mask]) + X_control = ( + X[control_mask].reset_index(drop=True) + if hasattr(X, "loc") + else X[control_mask] + ) + self.model_c.fit(X_control, y[control_mask]) # Expose as a shared-reference dict to preserve the public models_c API. self.models_c = {group: self.model_c for group in self.t_groups} @@ -140,6 +145,21 @@ def fit( y_filt = y[mask] w = (treatment_filt == group).astype(int) + self.models_t[group].fit(X_filt[w == 1], y_filt[w == 1]) + + if store_bootstraps: + self.fit_bootstrap_ensemble( + X=X, + treatment=treatment, + y=y, + n_bootstraps=n_bootstraps, + bootstrap_size=bootstrap_size, + random_state=random_state, + n_jobs=n_jobs, + ) + else: + self.bootstrap_models_ = None + def _compute_bootstrap_ci(self, X): """Compute bootstrap CI using stored ensemble. @@ -189,8 +209,10 @@ def predict( returns (te, te_lower, te_upper) each of shape [n_samples, n_treatment]. return_ci=True and return_components=True cannot be used together. """ + if return_ci and return_components: + raise ValueError("return_ci and return_components cannot both be True.") + treatment, y = convert_pd_to_np(treatment, y) - yhat_cs = {} yhat_ts = {} yhat_c = self.model_c.predict(X)