Lazy-fit model_mu in R-learner so it's only trained when return_components is used

causalml/inference/meta/rlearner.py

@@ -149,9 +149,11 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
         yhat = cross_val_predict(
             self.model_mu, X, y_np, cv=self.cv, n_jobs=self.cv_n_jobs
         )
-        # Fit the nuisance outcome model on the full data so it can be
-        # reused by predict(return_components=True).
-        self.model_mu.fit(X, y_np)
+        # Defer fitting the nuisance outcome model on the full data until
+        # predict(return_components=True) actually needs it.
+        self._model_mu_fitted = False
+        self._mu_fit_X = X
+        self._mu_fit_y = y_np
 
         for group in self.t_groups:
             mask = (treatment_np == group) | (treatment_np == self.control_name)
@@ -189,6 +191,25 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
             )
         return self
 
+    def __getstate__(self):
+        """Exclude the cached training data used only for the lazy
+        model_mu fit from pickling — keeping it would balloon save() size
+        (e.g. 2.7 KiB -> 4.8 MiB at n=20k/p=30). See PR #936 review."""
+        state = self.__dict__.copy()
+        state.pop("_mu_fit_X", None)
+        state.pop("_mu_fit_y", None)
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        # If model_mu was never lazily fit before saving, there's no cached
+        # training data to fit it with after loading, predict()'s
+        # return_components path will raise a clear error rather than an
+        # AttributeError if this case is hit.
+        if not getattr(self, "_model_mu_fitted", True):
+            self._mu_fit_X = None
+            self._mu_fit_y = None
+
     def predict(
         self,
         X,
@@ -216,6 +237,17 @@ def predict(
         if not return_components:
             return te
 
+        if not self._model_mu_fitted:
+            if self._mu_fit_X is None:
+                raise ValueError(
+                    "model_mu was not fit before this learner was saved, so "
+                    "return_components=True is unavailable after loading. "
+                    "Call predict(..., return_components=True) once before "
+                    "saving, or refit the learner."
+                )
+            self.model_mu.fit(self._mu_fit_X, self._mu_fit_y)
+            self._model_mu_fitted = True
+
         if p is None:
             if not hasattr(self, "propensity_model"):
                 raise ValueError(
@@ -292,6 +324,9 @@ def fit_predict(
             _classes_global = self._classes
             model_mu_global = deepcopy(self.model_mu)
             models_tau_global = deepcopy(self.models_tau)
+            model_mu_fitted_global = self._model_mu_fitted
+            mu_fit_X_global = self._mu_fit_X
+            mu_fit_y_global = self._mu_fit_y
             te_bootstraps = np.zeros(
                 shape=(n_rows(X), self.t_groups.shape[0], n_bootstraps)
             )
@@ -314,6 +349,9 @@ def fit_predict(
             self._classes = _classes_global
             self.model_mu = deepcopy(model_mu_global)
             self.models_tau = deepcopy(models_tau_global)
+            self._model_mu_fitted = model_mu_fitted_global
+            self._mu_fit_X = mu_fit_X_global
+            self._mu_fit_y = mu_fit_y_global
 
             return (te, te_lower, te_upper)
 
@@ -395,8 +433,12 @@ def estimate_ate(
             _classes_global = self._classes
             model_mu_global = deepcopy(self.model_mu)
             models_tau_global = deepcopy(self.models_tau)
+            model_mu_fitted_global = self._model_mu_fitted
+            mu_fit_X_global = self._mu_fit_X
+            mu_fit_y_global = self._mu_fit_y
 
             logger.info("Bootstrap Confidence Intervals for ATE")
+
             ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))
 
             for n in tqdm(range(n_bootstraps)):
@@ -418,6 +460,9 @@ def estimate_ate(
             self._classes = _classes_global
             self.model_mu = deepcopy(model_mu_global)
             self.models_tau = deepcopy(models_tau_global)
+            self._model_mu_fitted = model_mu_fitted_global
+            self._mu_fit_X = mu_fit_X_global
+            self._mu_fit_y = mu_fit_y_global
             return ate, ate_lower, ate_upper
 
 
@@ -541,7 +586,11 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
         yhat = cross_val_predict(
             self.model_mu, X, y_np, cv=self.cv, method="predict_proba", n_jobs=-1
         )[:, 1]
-        self.model_mu.fit(X, y_np)
+        # Defer fitting the nuisance outcome model on the full data until
+        # predict(return_components=True) actually needs it.
+        self._model_mu_fitted = False
+        self._mu_fit_X = X
+        self._mu_fit_y = y_np
 
         for group in self.t_groups:
             mask = (treatment_np == group) | (treatment_np == self.control_name)
@@ -596,6 +645,17 @@ def predict(
         if not return_components:
             return te
 
+        if not self._model_mu_fitted:
+            if self._mu_fit_X is None:
+                raise ValueError(
+                    "model_mu was not fit before this learner was saved, so "
+                    "return_components=True is unavailable after loading. "
+                    "Call predict(..., return_components=True) once before "
+                    "saving, or refit the learner."
+                )
+            self.model_mu.fit(self._mu_fit_X, self._mu_fit_y)
+            self._model_mu_fitted = True
+
         if p is None:
             if not hasattr(self, "propensity_model"):
                 raise ValueError(
@@ -755,7 +815,11 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
         if verbose:
             logger.info("generating out-of-fold CV outcome estimates")
         yhat = cross_val_predict(self.model_mu, X, y_np, cv=self.cv, n_jobs=-1)
-        self.model_mu.fit(X, y_np)
+        # Defer fitting the nuisance outcome model on the full data until
+        # predict(return_components=True) actually needs it.
+        self._model_mu_fitted = False
+        self._mu_fit_X = X
+        self._mu_fit_y = y_np
 
         for group in self.t_groups:
             treatment_mask = (treatment_np == group) | (

tests/test_meta_learners.py

@@ -42,6 +42,18 @@
 from .const import RANDOM_SEED, N_SAMPLE, ERROR_THRESHOLD, CONTROL_NAME, CONVERSION
 
 
+class _CountingRegressor(LinearRegression):
+    """LinearRegression that counts calls to fit(), to verify laziness."""
+
+    def __init__(self):
+        super().__init__()
+        self.fit_calls = 0
+
+    def fit(self, X, y, sample_weight=None):
+        self.fit_calls += 1
+        return super().fit(X, y, sample_weight=sample_weight)
+
+
 class ReadOnlyLinearRegression:
     """Minimal regressor that marks input arrays read-only like CatBoost."""
 
@@ -758,6 +770,193 @@ def test_BaseRLearner_predict_without_propensity_model_raises(
         )
 
 
+def test_BaseRLearner_model_mu_lazy_fit(generate_regression_data):
+    """model_mu must not be fit on the full data during fit(); only the first
+    predict(return_components=True) call should trigger it, and it should be
+    cached (not re-fit) on subsequent calls."""
+    y, X, treatment, tau, b, e = generate_regression_data()
+
+    outcome_learner = _CountingRegressor()
+    learner = BaseRLearner(
+        outcome_learner=outcome_learner,
+        effect_learner=LinearRegression(),
+    )
+    learner.fit(X=X, treatment=treatment, y=y, p=e, verbose=False)
+
+    # fit() alone must not eagerly fit model_mu on the full data.
+    assert outcome_learner.fit_calls == 0
+
+    te, yhat, p = learner.predict(X=X, p=e, return_components=True)
+    assert outcome_learner.fit_calls == 1
+
+    # A second return_components=True call must reuse the cached fit.
+    te2, yhat2, p2 = learner.predict(X=X, p=e, return_components=True)
+    assert outcome_learner.fit_calls == 1
+    np.testing.assert_array_equal(yhat, yhat2)
+
+    # Plain predict() (no components requested) must never trigger the fit.
+    fresh_outcome_learner = _CountingRegressor()
+    fresh_learner = BaseRLearner(
+        outcome_learner=fresh_outcome_learner,
+        effect_learner=LinearRegression(),
+    )
+    fresh_learner.fit(X=X, treatment=treatment, y=y, p=e, verbose=False)
+    fresh_learner.predict(X=X, p=e)
+    assert fresh_outcome_learner.fit_calls == 0
+
+
+def test_BaseRLearner_fit_predict_ci_then_predict_components(generate_regression_data):
+    """Regression test: fit_predict(return_ci=True) runs a bootstrap loop that
+    repeatedly re-fits model_mu on resampled data. Once it's done, the cached
+    training data used by predict(return_components=True) must be restored to
+    the ORIGINAL training set, not left pointing at the last bootstrap resample."""
+    y, X, treatment, tau, b, e = generate_regression_data()
+
+    learner = BaseRLearner(learner=LinearRegression())
+
+    learner.fit_predict(
+        X=X,
+        treatment=treatment,
+        y=y,
+        p=e,
+        return_ci=True,
+        n_bootstraps=5,
+        bootstrap_size=200,
+        verbose=False,
+    )
+
+    te, yhat, p = learner.predict(X=X, p=e, return_components=True)
+
+    expected_yhat = LinearRegression().fit(X, y).predict(X)
+
+    assert yhat.shape == (X.shape[0],)
+    np.testing.assert_allclose(yhat, expected_yhat, rtol=1e-6, atol=1e-8)
+
+
+def test_BaseRLearner_estimate_ate_bootstrap_then_predict_components(
+    generate_regression_data,
+):
+    """Same regression as above, but for the bootstrap_ci=True path in
+    estimate_ate() rather than fit_predict(return_ci=True)."""
+    y, X, treatment, tau, b, e = generate_regression_data()
+
+    learner = BaseRLearner(learner=LinearRegression())
+
+    learner.estimate_ate(
+        X=X,
+        treatment=treatment,
+        y=y,
+        p=e,
+        bootstrap_ci=True,
+        n_bootstraps=5,
+        bootstrap_size=200,
+    )
+
+    te, yhat, p = learner.predict(X=X, p=e, return_components=True)
+
+    expected_yhat = LinearRegression().fit(X, y).predict(X)
+
+    assert yhat.shape == (X.shape[0],)
+    np.testing.assert_allclose(yhat, expected_yhat, rtol=1e-6, atol=1e-8)
+
+
+def test_BaseRClassifier_model_mu_lazy_fit(generate_classification_data):
+    """Same laziness contract as BaseRLearner, but for BaseRClassifier
+    (predict_proba-based outcome model)."""
+    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
+    )
+
+    class _CountingClassifier(LogisticRegression):
+        def __init__(self):
+            super().__init__()
+            self.fit_calls = 0
+
+        def fit(self, X, y, sample_weight=None):
+            self.fit_calls += 1
+            return super().fit(X, y, sample_weight=sample_weight)
+
+    outcome_learner = _CountingClassifier()
+    learner = BaseRClassifier(
+        outcome_learner=outcome_learner, effect_learner=XGBRegressor()
+    )
+    learner.fit(
+        X=df[x_names].values,
+        treatment=df["treatment_group_key"].values,
+        y=df[CONVERSION].values,
+        verbose=False,
+    )
+    assert outcome_learner.fit_calls == 0
+
+    te, yhat, p = learner.predict(
+        X=df[x_names].values,
+        return_components=True,
+        p=np.full(len(df), 0.5),
+    )
+    assert outcome_learner.fit_calls == 1
+
+    learner.predict(
+        X=df[x_names].values,
+        return_components=True,
+        p=np.full(len(df), 0.5),
+    )
+    assert outcome_learner.fit_calls == 1
+
+
+def test_XGBRRegressor_model_mu_lazy_fit(generate_regression_data):
+    """Same laziness contract as BaseRLearner, but for XGBRRegressor."""
+    y, X, treatment, tau, b, e = generate_regression_data()
+
+    learner = XGBRRegressor(effect_learner_n_estimators=20, random_state=0)
+    learner.fit(X=X, treatment=treatment, y=y, p=e, verbose=False)
+
+    assert learner._model_mu_fitted is False
+
+    te, yhat, p = learner.predict(X=X, p=e, return_components=True)
+    assert learner._model_mu_fitted is True
+    assert yhat.shape == (X.shape[0],)
+
+
+def test_BaseRLearner_save_load_excludes_training_data(
+    generate_regression_data, tmp_path
+):
+    """save()/load() must not balloon in size by carrying the cached
+    training data used only for the lazy model_mu fit."""
+    y, X, treatment, tau, b, e = generate_regression_data()
+
+    learner = BaseRLearner(learner=LinearRegression())
+    learner.fit(X=X, treatment=treatment, y=y, p=e, verbose=False)
+
+    path = tmp_path / "rlearner.causalml"
+    learner.save(str(path))
+
+    loaded = BaseRLearner.load(str(path))
+    assert not hasattr(loaded, "_mu_fit_X") or loaded._mu_fit_X is None
+    assert not hasattr(loaded, "_mu_fit_y") or loaded._mu_fit_y is None
+
+    # Loaded model still predicts CATE fine (doesn't need model_mu).
+    te = loaded.predict(X=X)
+    assert te.shape == (X.shape[0], len(loaded.t_groups))
+
+    # But return_components=True raises a clear error post-load, since
+    # model_mu was never lazily fit before saving.
+    with pytest.raises(ValueError, match="not fit before this learner was saved"):
+        loaded.predict(X=X, p=e, return_components=True)
+
+    # If return_components=True was called BEFORE saving, model_mu is
+    # already fitted, and reload works fine for components too.
+    learner2 = BaseRLearner(learner=LinearRegression())
+    learner2.fit(X=X, treatment=treatment, y=y, p=e, verbose=False)
+    learner2.predict(X=X, p=e, return_components=True)  # triggers lazy fit
+    path2 = tmp_path / "rlearner_prefit.causalml"
+    learner2.save(str(path2))
+    loaded2 = BaseRLearner.load(str(path2))
+    te2, yhat2, p2 = loaded2.predict(X=X, p=e, return_components=True)
+    assert yhat2.shape == (X.shape[0],)
+
+
 def test_BaseRRegressor_without_p(generate_regression_data):
     y, X, treatment, tau, b, e = generate_regression_data()