Fix #1045: seed and stratify the GBM propensity early-stopping split

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(

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))