fix validation fraction api

RELEASE_NOTES.md

@@ -1,5 +1,11 @@
 # RELEASE NOTES
 
+## Version 0.5.11
+
+* Fix `NGBClassifier` and `NGBSurvival` API parity by adding `validation_fraction` and `early_stopping_rounds`, with regression coverage for their validation-split paths (issue #402)
+* Replace the deprecated `friedman_mse` criterion with `squared_error` for the default tree learner and distribution tests (issue #408, PR #409)
+* Support a distinct base learner for each distribution parameter, including scikit-learn nested parameters and feature importances for all-tree configurations (issue #338, PR #411)
+
 ## Version 0.5.10
 
 * Add `load_ngboost_model` compatibility loader for models saved with scikit-learn < 1.3 and loaded under newer scikit-learn versions (issue #389)

ngboost/api.py

@@ -154,6 +154,12 @@ class NGBClassifier(NGBoost, BaseEstimator):
         tol               : numerical tolerance to be used in optimization
         random_state      : seed for reproducibility. See
                             https://stackoverflow.com/questions/28064634/random-state-pseudo-random-number-in-scikit-learn
+        validation_fraction: Proportion of training data to set
+                             aside as validation data for early stopping.
+        early_stopping_rounds:      The number of consecutive boosting iterations during which the
+                                    loss has to increase before the algorithm stops early.
+                                    Set to None to disable early stopping and validation.
+                                    None enables running over the full data set.
     Output:
         An NGBClassifier object that can be fit.
     """
@@ -173,6 +179,8 @@ def __init__(
         verbose_eval=100,
         tol=1e-4,
         random_state=None,
+        validation_fraction=0.1,
+        early_stopping_rounds=None,
     ):
         assert issubclass(
             Dist, ClassificationDistn
@@ -190,6 +198,8 @@ def __init__(
             verbose_eval,
             tol,
             random_state,
+            validation_fraction,
+            early_stopping_rounds,
         )
         self._estimator_type = "classifier"
 
@@ -255,11 +265,17 @@ class NGBSurvival(NGBoost, BaseEstimator):
         tol               : numerical tolerance to be used in optimization
         random_state      : seed for reproducibility. See
                             https://stackoverflow.com/questions/28064634/random-state-pseudo-random-number-in-scikit-learn
+        validation_fraction: Proportion of training data to set
+                             aside as validation data for early stopping.
+        early_stopping_rounds:      The number of consecutive boosting iterations during which the
+                                    loss has to increase before the algorithm stops early.
+                                    Set to None to disable early stopping and validation.
+                                    None enables running over the full data set.
     Output:
         An NGBSurvival object that can be fit.
     """
 
-    # pylint: disable=too-many-positional-arguments
+    # pylint: disable=too-many-positional-arguments,too-many-locals
     def __init__(
         self,
         Dist=LogNormal,
@@ -274,6 +290,8 @@ def __init__(
         verbose_eval=100,
         tol=1e-4,
         random_state=None,
+        validation_fraction=0.1,
+        early_stopping_rounds=None,
     ):
 
         assert issubclass(
@@ -300,6 +318,8 @@ def __init__(
             verbose_eval,
             tol,
             random_state,
+            validation_fraction,
+            early_stopping_rounds,
         )
 
     def __getstate__(self):

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "ngboost"
-version = "0.5.10dev"
+version = "0.5.11dev"
 description = "Library for probabilistic predictions via gradient boosting."
 authors = ["Stanford ML Group <avati@cs.stanford.edu>"]
 readme = "README.md"

tests/test_basic.py

@@ -5,7 +5,7 @@
 from sklearn.linear_model import Ridge
 from sklearn.tree import DecisionTreeRegressor
 
-from ngboost import NGBClassifier, NGBRegressor
+from ngboost import NGBClassifier, NGBRegressor, NGBSurvival
 from ngboost.distns import Bernoulli, Normal, k_categorical
 
 
@@ -74,6 +74,45 @@ def test_regression(california_housing_data):
     assert score <= 15
 
 
+def test_classifier_validation_fraction_is_supported(breast_cancer_data):
+    x_train, x_test, y_train, _ = breast_cancer_data
+    ngb = NGBClassifier(
+        Dist=Bernoulli,
+        n_estimators=25,
+        verbose=False,
+        random_state=1,
+        validation_fraction=0.2,
+        early_stopping_rounds=2,
+    )
+
+    assert ngb.get_params()["validation_fraction"] == 0.2
+    assert ngb.get_params()["early_stopping_rounds"] == 2
+
+    ngb.fit(x_train, y_train)
+    preds = ngb.predict(x_test)
+    assert preds.shape[0] == x_test.shape[0]
+
+
+def test_survival_validation_fraction_is_supported(
+    california_housing_survival_data,
+):
+    x_train, x_test, t_train, e_train, _ = california_housing_survival_data
+    ngb = NGBSurvival(
+        n_estimators=5,
+        verbose=False,
+        random_state=1,
+        validation_fraction=0.2,
+        early_stopping_rounds=2,
+    )
+
+    assert ngb.get_params()["validation_fraction"] == 0.2
+    assert ngb.get_params()["early_stopping_rounds"] == 2
+
+    ngb.fit(x_train, t_train, e_train)
+    preds = ngb.predict(x_test)
+    assert preds.shape[0] == x_test.shape[0]
+
+
 def test_regression_accepts_base_learner_per_distribution_parameter():
     rng = np.random.default_rng(0)
     X = rng.normal(size=(80, 4))