diff --git a/ngboost/ngboost.py b/ngboost/ngboost.py index 3d3e6702..3970b03c 100644 --- a/ngboost/ngboost.py +++ b/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( diff --git a/tests/test_basic.py b/tests/test_basic.py index 298d1c48..c8373a27 100644 --- a/tests/test_basic.py +++ b/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))