diff --git a/sklego/linear_model.py b/sklego/linear_model.py index 7c05262c..ebe9fc43 100644 --- a/sklego/linear_model.py +++ b/sklego/linear_model.py @@ -141,7 +141,17 @@ def predict(self, X): X = check_array(X, estimator=self, dtype=FLOAT_DTYPES) check_is_fitted(self, ["X_", "y_"]) - results = np.stack([np.average(self.y_, weights=self._calc_wts(x_i=x_i)) for x_i in X]) + try: + results = np.stack([np.average(self.y_, weights=self._calc_wts(x_i=x_i)) for x_i in X]) + except ZeroDivisionError: + msg = ( + "Weights, resulting from `np.exp(-(distances**2) / self.sigma)`, are all zero. " + "Try to increase the value of `sigma` or to normalize the input data.\n\n" + "`distances` refer to the distance between each sample `x_i` with all the" + "training samples." + ) + raise ValueError(msg) + return results diff --git a/tests/test_estimators/test_lowess.py b/tests/test_estimators/test_lowess.py index a135a0aa..a4d39f3f 100644 --- a/tests/test_estimators/test_lowess.py +++ b/tests/test_estimators/test_lowess.py @@ -1,4 +1,7 @@ +import re + import numpy as np +import pytest from sklearn.utils.estimator_checks import parametrize_with_checks from sklego.linear_model import LowessRegression @@ -15,3 +18,18 @@ def test_obvious_usecase(): y = np.ones(x.shape) y_pred = LowessRegression().fit(X, y).predict(X) assert np.isclose(y, y_pred).all() + + +def test_custom_error_for_zero_division(): + x = np.arange(0, 100) + X = x.reshape(-1, 1) + y = np.ones(x.shape) + estimator = LowessRegression(sigma=1e-10).fit(X, y) + + with pytest.raises( + ValueError, match=re.escape("Weights, resulting from `np.exp(-(distances**2) / self.sigma)`, are all zero.") + ): + # Adding an offset, otherwise X to predict would be the same as X in fit method, + # leading to weight of 1 for the corresponding value. + X_pred = X[:10] + 0.5 + estimator.predict(X_pred)