From 2855ccfe9bb6dd38272fc31c521ab817505bda60 Mon Sep 17 00:00:00 2001 From: Hari20032005 Date: Fri, 4 Sep 2026 12:30:48 +0530 Subject: [PATCH] Fix #1045: seed and stratify the GBM propensity early-stopping split GradientBoostedPropensityModel.fit() drew its early-stopping validation set with neither random_state nor stratify, so a model built with an explicit seed still fit differently on every call -- propensity scores moved by up to 0.55 between identical calls. Pass the same random_state expression _model uses for the XGBClassifier, matching the R-learner's own early-stopping split (rlearner.py:830), and stratify on the treatment indicator as every other treatment split in the library does. Add regression tests covering both properties. --- causalml/propensity.py | 11 ++++++++++- tests/test_propensity.py | 41 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/causalml/propensity.py b/causalml/propensity.py index 5e22d106..197bccc4 100644 --- a/causalml/propensity.py +++ b/causalml/propensity.py @@ -179,8 +179,17 @@ def fit(self, X, y, stop_val_size=0.2): so this vector must be the treatment assignment. """ if self.early_stop: + # Seed the split with the same random_state the underlying + # XGBClassifier resolves to, so that early stopping -- and + # therefore the fitted model -- is reproducible. Stratify on the + # treatment indicator so the validation set retains both arms, + # matching how the rest of the library splits on treatment. X_train, X_val, y_train, y_val = train_test_split( - X, y, test_size=stop_val_size + X, + y, + test_size=stop_val_size, + random_state=self.model_kwargs.get("random_state", 42), + stratify=y, ) self.model.fit( diff --git a/tests/test_propensity.py b/tests/test_propensity.py index ee20a5a5..b7568eea 100644 --- a/tests/test_propensity.py +++ b/tests/test_propensity.py @@ -1,5 +1,7 @@ import numpy as np +import pytest +from causalml import propensity from causalml.propensity import ( ElasticNetPropensityModel, GradientBoostedPropensityModel, @@ -71,6 +73,45 @@ def test_gradientboosted_propensity_model_earlystopping(generate_regression_data assert roc_auc_score(treatment, ps) > 0.5 +def test_gradientboosted_propensity_model_earlystopping_reproducible( + generate_regression_data, +): + """Early stopping is reproducible: the validation split is seeded (#1045).""" + y, X, treatment, tau, b, e = generate_regression_data() + + def fit_predict(random_state): + pm = GradientBoostedPropensityModel(random_state=random_state, early_stop=True) + return pm.fit_predict(X, treatment) + + np.testing.assert_array_equal(fit_predict(RANDOM_SEED), fit_predict(RANDOM_SEED)) + # A different seed must still move the split; otherwise the seed would be + # ignored in a different way (e.g. a hard-coded constant). + assert not np.array_equal(fit_predict(RANDOM_SEED), fit_predict(RANDOM_SEED + 1)) + + +def test_gradientboosted_propensity_model_earlystopping_stratified(monkeypatch): + """The early-stopping validation split keeps both treatment arms (#1045).""" + rng = np.random.RandomState(RANDOM_SEED) + X = rng.normal(size=(400, 10)) + treatment = (rng.uniform(size=400) < 0.1).astype(int) + + captured = {} + train_test_split = propensity.train_test_split + + def spy(*args, **kwargs): + split = train_test_split(*args, **kwargs) + captured["y_val"] = split[3] + return split + + monkeypatch.setattr(propensity, "train_test_split", spy) + + pm = GradientBoostedPropensityModel(random_state=RANDOM_SEED, early_stop=True) + pm.fit(X, treatment) + + # Stratification preserves the treatment rate up to rounding. + assert captured["y_val"].mean() == pytest.approx(treatment.mean(), abs=0.01) + + def test_propensity_models_imbalanced_1027(): rng = np.random.RandomState(RANDOM_SEED) X = rng.normal(size=(400, 25))