Honor max_iter=0 in ordinary and staged predictions

ngboost/ngboost.py

@@ -181,7 +181,7 @@ def pred_param(self, X, max_iter=None):
         for i, (models, s, col_idx) in enumerate(
             zip(self.base_models, self.scalings, self.col_idxs)
         ):
-            if max_iter and i == max_iter:
+            if max_iter is not None and i == max_iter:
                 break
             resids = np.array([model.predict(X[:, col_idx]) for model in models]).T
             params -= self.learning_rate * resids * s
@@ -745,6 +745,8 @@ def staged_pred_dist(self, X, max_iter=None):
             A list of NGBoost distribution objects, one per boosting stage up to max_iter
         """
         predictions = []
+        if max_iter == 0:
+            return predictions
         m, n = X.shape
         params = np.ones((m, self.Dist.n_params)) * self.init_params
         for i, (models, s, col_idx) in enumerate(

tests/test_basic.py

@@ -13,6 +13,7 @@
 from sklearn.pipeline import make_pipeline
 from sklearn.preprocessing import StandardScaler
 from sklearn.tree import DecisionTreeRegressor
+from sklearn.utils import check_random_state
 from sklearn.utils.validation import check_is_fitted
 
 from ngboost import NGBClassifier, NGBRegressor, NGBSurvival
@@ -288,6 +289,64 @@ def test_regression(california_housing_data):
     assert score <= 15
 
 
+@pytest.mark.parametrize("classifier", [False, True])
+@pytest.mark.parametrize("max_iter", [0, 1, None])
+def test_prediction_iteration_limits(classifier, max_iter):
+    X = check_random_state(12).normal(size=(30, 3))
+    y = np.where(X[:, 0] > 0, "yes", "no") if classifier else X[:, 0]
+    model = (NGBClassifier if classifier else NGBRegressor)(
+        n_estimators=3, verbose=False, random_state=7
+    ).fit(X, y)
+    count = len(model.base_models) if max_iter is None else max_iter
+    expected_params = np.tile(model.init_params, (len(X), 1))
+    for models, scale, columns in list(
+        zip(model.base_models, model.scalings, model.col_idxs)
+    )[:count]:
+        expected_params -= (
+            model.learning_rate
+            * scale
+            * np.array([base.predict(X[:, columns]) for base in models]).T
+        )
+    expected_dist = model.Dist(expected_params.T)
+    expected_prediction = expected_dist.predict()
+    if classifier:
+        expected_prediction = model.classes_[expected_prediction]
+
+    np.testing.assert_allclose(model.pred_param(X, max_iter), expected_params)
+    np.testing.assert_equal(model.pred_dist(X, max_iter).params, expected_dist.params)
+    np.testing.assert_array_equal(model.predict(X, max_iter), expected_prediction)
+    staged_dist = model.staged_pred_dist(X, max_iter)
+    staged_predictions = model.staged_predict(X, max_iter)
+    assert len(staged_dist) == count
+    assert len(staged_predictions) == count
+    if count:
+        np.testing.assert_equal(staged_dist[-1].params, expected_dist.params)
+        np.testing.assert_array_equal(staged_predictions[-1], expected_prediction)
+    if classifier:
+        np.testing.assert_allclose(
+            model.predict_proba(X, max_iter), expected_dist.class_probs()
+        )
+        staged_probabilities = model.staged_predict_proba(X, max_iter)
+        assert len(staged_probabilities) == count
+        if count:
+            np.testing.assert_allclose(
+                staged_probabilities[-1], expected_dist.class_probs()
+            )
+
+
+@pytest.mark.parametrize("classifier", [False, True])
+def test_zero_staged_predictions_are_empty(classifier):
+    X = np.arange(30).reshape(10, 3)
+    y = np.arange(10) % 2
+    model = (NGBClassifier if classifier else NGBRegressor)(
+        n_estimators=2, verbose=False, random_state=7
+    ).fit(X, y)
+    assert not model.staged_pred_dist(X, max_iter=0)
+    assert not model.staged_predict(X, max_iter=0)
+    if classifier:
+        assert not model.staged_predict_proba(X, max_iter=0)
+
+
 def test_score_is_mean_log_likelihood():
     rng = np.random.default_rng(0)
     X = rng.normal(size=(200, 4))