Merge pull request scikit-learn#2941 from perimosocordiae/patch-1

BUG: avoid NaNs throwing off class probabilities

Author: GaelVaroquaux

sklearn/neighbors/classification.py

@@ -200,6 +200,10 @@ def predict_proba(self, X):
         weights = _get_weights(neigh_dist, self.weights)
         if weights is None:
             weights = np.ones_like(neigh_ind)
+        else:
+            # Some weights may be infinite (zero distance), which can cause
+            # downstream NaN values when used for normalization.
+            weights[np.isinf(weights)] = np.finfo('f').max
 
         all_rows = np.arange(X.shape[0])
         probabilities = []

sklearn/neighbors/tests/test_neighbors.py

@@ -201,6 +201,14 @@ def test_kneighbors_classifier_predict_proba():
     cls.fit(X, y.astype(str))
     y_prob = cls.predict_proba(X)
     assert_array_equal(real_prob, y_prob)
+    # Check that it works with weights='distance'
+    cls = neighbors.KNeighborsClassifier(
+        n_neighbors=2, p=1, weights='distance')
+    cls.fit(X, y)
+    y_prob = cls.predict_proba(np.array([[0, 2, 0], [2, 2, 2]]))
+    real_prob = np.array([[0, 1, 0], [0, 0.4, 0.6]])
+    assert_array_almost_equal(real_prob, y_prob)
+
 
 
 def test_radius_neighbors_classifier(n_samples=40,