Add post-fit confidence intervals to BaseTLearner via store_bootstraps and return_ci

causalml/inference/meta/base.py

@@ -2,6 +2,9 @@
 import logging
 import numpy as np
 import pandas as pd
+from joblib import Parallel, delayed
+from sklearn.base import clone
+from tqdm import tqdm
 
 from causalml.inference.meta.explainer import Explainer
 from causalml.inference.meta.utils import check_p_conditions, convert_pd_to_np
@@ -70,6 +73,69 @@ def bootstrap(self, X, treatment, y, p=None, size=10000, rng=None):
         self.fit(X=X_b, treatment=treatment_b, y=y_b, p=p_b)
         return self.predict(X=X, p=p)
 
+    def fit_bootstrap_ensemble(
+        self,
+        X,
+        treatment,
+        y,
+        n_bootstraps=200,
+        bootstrap_size=10000,
+        random_state=None,
+        n_jobs=1,
+    ):
+        """Train and store a bootstrap ensemble for post-fit CI estimation.
+
+        Fits n_bootstraps cloned copies of the learner on bootstrap samples and
+        stores them in self.bootstrap_models_. Used by predict(return_ci=True)
+        to compute percentile-based confidence intervals on new data without
+        refitting.
+
+        Note: storing N bootstraps of a GBM-based learner with k treatment groups
+        holds 2*N*k model objects in memory. Monitor RAM for large N or heavy
+        base learners.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+            n_bootstraps (int, optional): number of bootstrap iterations. Default: 200.
+            bootstrap_size (int, optional): number of samples per bootstrap. Default: 10000.
+            random_state (int, optional): random seed for reproducibility.
+            n_jobs (int, optional): number of parallel jobs. -1 uses all cores. Default: 1.
+        """
+        rng = np.random.RandomState(random_state)
+        seeds = rng.randint(0, np.iinfo(np.int32).max, size=n_bootstraps)
+        logger.info("Storing bootstrap ensemble ({} iterations)".format(n_bootstraps))
+
+        def _fit_one(seed):
+            local_rng = np.random.RandomState(seed)
+            idxs = local_rng.choice(np.arange(X.shape[0]), size=bootstrap_size)
+            X_b, treatment_b, y_b = X[idxs], treatment[idxs], y[idxs]
+            models_c_b = {group: clone(self.model_c) for group in self.t_groups}
+            models_t_b = {group: clone(self.model_t) for group in self.t_groups}
+            for group in self.t_groups:
+                mask = (treatment_b == group) | (treatment_b == self.control_name)
+                treatment_filt = treatment_b[mask]
+                X_filt = X_b[mask]
+                y_filt = y_b[mask]
+                w = (treatment_filt == group).astype(int)
+                if w.sum() == 0 or (w == 0).sum() == 0:
+                    logger.warning(
+                        "Bootstrap sample has no treated or no control units "
+                        "for group {}. Falling back to global model — "
+                        "CI may be underestimated.".format(group)
+                    )
+                    models_c_b[group] = self.models_c[group]
+                    models_t_b[group] = self.models_t[group]
+                    continue
+                models_c_b[group].fit(X_filt[w == 0], y_filt[w == 0])
+                models_t_b[group].fit(X_filt[w == 1], y_filt[w == 1])
+            return models_c_b, models_t_b
+
+        self.bootstrap_models_ = Parallel(n_jobs=n_jobs)(
+            delayed(_fit_one)(s) for s in tqdm(seeds)
+        )
+
     @staticmethod
     def _format_p(p, t_groups):
         """Format propensity scores into a dictionary of {treatment group: propensity scores}.

causalml/inference/meta/tlearner.py

@@ -62,20 +62,44 @@ def __init__(
 
         self.ate_alpha = ate_alpha
         self.control_name = control_name
+        self.bootstrap_models_ = None
 
     def __repr__(self):
         return "{}(model_c={}, model_t={})".format(
             self.__class__.__name__, self.model_c.__repr__(), self.model_t.__repr__()
         )
 
     @ignore_warnings(category=ConvergenceWarning)
-    def fit(self, X, treatment, y, p=None):
+    def fit(
+        self,
+        X,
+        treatment,
+        y,
+        p=None,
+        store_bootstraps=False,
+        n_bootstraps=200,
+        bootstrap_size=10000,
+        random_state=None,
+        n_jobs=1,
+    ):
         """Fit the inference model
 
         Args:
             X (np.matrix or np.array or pd.Dataframe): a feature matrix
             treatment (np.array or pd.Series): a treatment vector
             y (np.array or pd.Series): an outcome vector
+            p: unused, kept for API consistency
+            store_bootstraps (bool, optional): if True, trains a bootstrap ensemble
+                during fit and stores it in self.bootstrap_models_ for post-fit CI
+                estimation via predict(return_ci=True). Default: False.
+            n_bootstraps (int, optional): number of bootstrap iterations. Default: 200.
+                Note: storing N bootstraps of a GBM-based learner with k treatment
+                groups holds 2*N*k model objects in memory. Monitor RAM for large N
+                or heavy base learners.
+            n_jobs (int, optional): number of parallel jobs for bootstrap fitting.
+                -1 uses all available cores. Default: 1.
+            bootstrap_size (int, optional): number of samples per bootstrap. Default: 10000.
+            random_state (int, optional): random seed for reproducible bootstrap sampling.
         """
         X, treatment, y = convert_pd_to_np(X, treatment, y)
         check_treatment_vector(treatment, self.control_name)
@@ -95,20 +119,50 @@ def fit(self, X, treatment, y, p=None):
             self.models_c[group].fit(X_filt[w == 0], y_filt[w == 0])
             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 predict(
-        self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
+        self,
+        X,
+        treatment=None,
+        y=None,
+        p=None,
+        return_components=False,
+        verbose=True,
+        return_ci=False,
     ):
         """Predict treatment effects.
 
         Args:
             X (np.matrix or np.array or pd.Dataframe): a feature matrix
             treatment (np.array or pd.Series, optional): a treatment vector
             y (np.array or pd.Series, optional): an outcome vector
-            return_components (bool, optional): whether to return outcome for treatment and control seperately
+            return_components (bool, optional): whether to return outcome for
+                treatment and control separately
             verbose (bool, optional): whether to output progress logs
+            return_ci (bool, optional): whether to return confidence intervals
+                using the stored bootstrap ensemble. Requires fit() to have been
+                called with store_bootstraps=True. CI width is controlled by
+                self.ate_alpha set at init time.
         Returns:
-            (numpy.ndarray): Predictions of treatment effects.
+            (numpy.ndarray): Predictions of treatment effects. If return_ci=True,
+                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.")
+
         X, treatment, y = convert_pd_to_np(X, treatment, y)
         yhat_cs = {}
         yhat_ts = {}
@@ -136,6 +190,25 @@ def predict(
         for i, group in enumerate(self.t_groups):
             te[:, i] = yhat_ts[group] - yhat_cs[group]
 
+        if return_ci:
+            if self.bootstrap_models_ is None:
+                raise ValueError(
+                    "No bootstrap ensemble found. Call fit(..., store_bootstraps=True) first."
+                )
+            te_bootstraps = np.zeros(
+                (X.shape[0], self.t_groups.shape[0], len(self.bootstrap_models_))
+            )
+            for b, (models_c_b, models_t_b) in enumerate(self.bootstrap_models_):
+                for i, group in enumerate(self.t_groups):
+                    te_bootstraps[:, i, b] = models_t_b[group].predict(X) - models_c_b[
+                        group
+                    ].predict(X)
+            te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
+            te_upper = np.percentile(
+                te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
+            )
+            return te, te_lower, te_upper
+
         if not return_components:
             return te
         else:

tests/test_meta_learners.py

@@ -1,5 +1,6 @@
 import numpy as np
 import pandas as pd
+import pytest
 
 from sklearn.linear_model import LinearRegression
 from sklearn.linear_model import LogisticRegression
@@ -1220,3 +1221,55 @@ 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_predict_return_ci(generate_regression_data):
+    y, X, treatment, tau, b, e = generate_regression_data()
+
+    learner = BaseTRegressor(learner=LinearRegression(), control_name=0)
+
+    # Test 1: store_bootstraps=True then predict with return_ci=True
+    learner.fit(
+        X,
+        treatment,
+        y,
+        store_bootstraps=True,
+        n_bootstraps=50,
+        bootstrap_size=500,
+        random_state=RANDOM_SEED,
+    )
+    tau_pred, lb, ub = learner.predict(X, return_ci=True)
+
+    assert tau_pred.shape == (X.shape[0], len(learner.t_groups))
+    assert lb.shape == tau_pred.shape
+    assert ub.shape == tau_pred.shape
+    assert (lb <= ub).all()
+
+    # Test 2: ValueError without store_bootstraps
+    learner2 = BaseTRegressor(learner=LinearRegression(), control_name=0)
+    learner2.fit(X, treatment, y)
+    with pytest.raises(ValueError):
+        learner2.predict(X, return_ci=True)
+
+    # Test 3: ValueError when return_ci and return_components both True
+    with pytest.raises(ValueError):
+        learner.predict(X, return_ci=True, return_components=True)
+
+    # Test 4: old API unchanged
+    tau_plain = learner.predict(X)
+    assert tau_plain.shape == (X.shape[0], len(learner.t_groups))
+
+    # Test 5: reproducibility via random_state
+    learner3 = BaseTRegressor(learner=LinearRegression(), control_name=0)
+    learner3.fit(
+        X,
+        treatment,
+        y,
+        store_bootstraps=True,
+        n_bootstraps=50,
+        bootstrap_size=500,
+        random_state=RANDOM_SEED,
+    )
+    tau2, lb2, ub2 = learner3.predict(X, return_ci=True)
+    np.testing.assert_array_equal(lb, lb2)
+    np.testing.assert_array_equal(ub, ub2)