diff --git a/causalml/inference/meta/rlearner.py b/causalml/inference/meta/rlearner.py index a7c45182..000d826e 100644 --- a/causalml/inference/meta/rlearner.py +++ b/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) | ( diff --git a/tests/test_meta_learners.py b/tests/test_meta_learners.py index 0886c69b..85cccb28 100644 --- a/tests/test_meta_learners.py +++ b/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()