scikit-learn · ankit-dhokariya · Aug 25, 2024 · Aug 25, 2024 · Aug 25, 2024 · Aug 25, 2024
diff --git a/doc/whats_new/v1.6.rst b/doc/whats_new/v1.6.rst
@@ -146,6 +146,13 @@ Changelog
    :pr:`123456` by :user:`Joe Bloggs <joeongithub>`.
    where 123455 is the *pull request* number, not the issue number.

+:mod:`sklearn.neighbors`
+........................
+
+- |Enhancement| :func:`sklearn.neighbors.KNeighborsRegressor.predict` now accepts `return_std` flag and returns standard deviation of the predictions.
+  :pr:`29711` by :user:`Ankit Dhokariya <ankit-dhokariya>`
+
+
 :mod:`sklearn.base`
 ...................


diff --git a/sklearn/neighbors/_regression.py b/sklearn/neighbors/_regression.py
@@ -219,7 +219,7 @@ def fit(self, X, y):
        """
        return self._fit(X, y)

-    def predict(self, X):
+    def predict(self, X, return_std=False):
        """Predict the target for the provided data.

        Parameters
@@ -228,10 +228,17 @@ def predict(self, X):
                or (n_queries, n_indexed) if metric == 'precomputed'
            Test samples.

+        return_std : bool, default=False
+            If True, the standard-deviation of the predictions are returned.
+
        Returns
        -------
        y : ndarray of shape (n_queries,) or (n_queries, n_outputs), dtype=int
            Target values.
+
+        std : ndarray of shape (n_queries,) or (n_queries, n_outputs), dtype=float
+            Standard deviation of the predictions. \
+            Only returned if `return_std` is True.
        """
        if self.weights == "uniform":
            # In that case, we do not need the distances to perform
@@ -249,17 +256,33 @@ def predict(self, X):

        if weights is None:
            y_pred = np.mean(_y[neigh_ind], axis=1)
+            if return_std:
+                y_std = np.std(_y[neigh_ind], axis=1)
        else:
            y_pred = np.empty((neigh_dist.shape[0], _y.shape[1]), dtype=np.float64)
+            if return_std:
+                y_std = np.empty((neigh_dist.shape[0], _y.shape[1]), dtype=np.float64)
            denom = np.sum(weights, axis=1)

            for j in range(_y.shape[1]):
                num = np.sum(_y[neigh_ind, j] * weights, axis=1)
                y_pred[:, j] = num / denom
+                if return_std:
+                    y_std[:, j] = np.sqrt(
+                        np.average(
+                            (_y[neigh_ind, j] - y_pred[:, j : j + 1]) ** 2,
+                            axis=1,
+                            weights=weights,
+                        )
+                    )

        if self._y.ndim == 1:
            y_pred = y_pred.ravel()
+            if return_std:
+                y_std = y_std.ravel()

+        if return_std:
+            return y_pred, y_std
        return y_pred



diff --git a/sklearn/neighbors/tests/test_neighbors.py b/sklearn/neighbors/tests/test_neighbors.py
@@ -1225,6 +1225,33 @@ def test_kneighbors_regressor(
            assert np.all(abs(y_pred - y_target) < 0.3)


+def test_kneighbors_regressor_predict_with_std(
+    n_samples=40, n_features=5, n_test_pts=10, n_neighbors=3, random_state=0
+):
+    # Test k-neighbors regression with return_std
+    rng = np.random.RandomState(random_state)
+    X = 2 * rng.rand(n_samples, n_features) - 1
+    y = np.sqrt((X**2).sum(1))
+    y /= y.max()
+
+    y_target = y[:n_test_pts]
+
+    weight_func = _weight_func
+
+    for algorithm in ALGORITHMS:
+        for weights in ["uniform", "distance", weight_func]:
+            knn = neighbors.KNeighborsRegressor(
+                n_neighbors=n_neighbors, weights=weights, algorithm=algorithm
+            )
+            knn.fit(X, y)
+            epsilon = 1e-5 * (2 * rng.rand(1, n_features) - 1)
+            y_pred, y_std = knn.predict(X[:n_test_pts] + epsilon, return_std=True)
+            assert y_pred.shape == y_target.shape
+            assert np.all(abs(y_pred - y_target) < 0.3)
+            assert y_std.shape == y_target.shape
+            assert np.all(y_std >= 0)
+
+
 def test_KNeighborsRegressor_multioutput_uniform_weight():
    # Test k-neighbors in multi-output regression with uniform weight
    rng = check_random_state(0)