chore: improve error message for ZeroDivisionError in LowessRegression

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
 
 

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)