ENH: support arbitrary dtype in kNN classifiers

GaelVaroquaux · GaelVaroquaux · commit 70cab6d85d2b · 2012-10-09T00:01:00.000+02:00
Fixes scikit-learn#1224
diff --git a/sklearn/neighbors/classification.py b/sklearn/neighbors/classification.py
@@ -138,7 +138,7 @@ def predict(self, X):
         else:
             mode, _ = weighted_mode(pred_labels, weights, axis=1)
 
-        return mode.flatten().astype(np.int)
+        return mode.flatten().astype(self._y.dtype)
 
     def predict_proba(self, X):
         """Return probability estimates for the test data X.
@@ -312,11 +312,11 @@ def predict(self, X):
         weights = _get_weights(neigh_dist, self.weights)
 
         if weights is None:
-            mode = np.asarray([stats.mode(pl)[0] for pl in pred_labels],
-                              dtype=np.int)
+            mode = np.array([stats.mode(pl)[0] for pl in pred_labels],
+                            dtype=self._y.dtype)
         else:
-            mode = np.asarray([weighted_mode(pl, w)[0]
-                               for (pl, w) in zip(pred_labels, weights)],
-                              dtype=np.int)
+            mode = np.array([weighted_mode(pl, w)[0]
+                             for (pl, w) in zip(pred_labels, weights)],
+                            dtype=self._y.dtype)
 
-        return mode.flatten().astype(np.int)
+        return mode.flatten()
diff --git a/sklearn/neighbors/tests/test_neighbors.py b/sklearn/neighbors/tests/test_neighbors.py
@@ -175,6 +175,7 @@ def test_kneighbors_classifier(n_samples=40,
     rng = np.random.RandomState(random_state)
     X = 2 * rng.rand(n_samples, n_features) - 1
     y = ((X ** 2).sum(axis=1) < .5).astype(np.int)
+    y_str = y.astype(str)
 
     weight_func = _weight_func
 
@@ -187,6 +188,10 @@ def test_kneighbors_classifier(n_samples=40,
             epsilon = 1e-5 * (2 * rng.rand(1, n_features) - 1)
             y_pred = knn.predict(X[:n_test_pts] + epsilon)
             assert_array_equal(y_pred, y[:n_test_pts])
+            # Test prediction with y_str
+            knn.fit(X, y_str)
+            y_pred = knn.predict(X[:n_test_pts] + epsilon)
+            assert_array_equal(y_pred, y_str[:n_test_pts])
 
 
 def test_kneighbors_classifier_predict_proba():
@@ -219,6 +224,7 @@ def test_radius_neighbors_classifier(n_samples=40,
     rng = np.random.RandomState(random_state)
     X = 2 * rng.rand(n_samples, n_features) - 1
     y = ((X ** 2).sum(axis=1) < .5).astype(np.int)
+    y_str = y.astype(str)
 
     weight_func = _weight_func
 
@@ -231,6 +237,9 @@ def test_radius_neighbors_classifier(n_samples=40,
             epsilon = 1e-5 * (2 * rng.rand(1, n_features) - 1)
             y_pred = neigh.predict(X[:n_test_pts] + epsilon)
             assert_array_equal(y_pred, y[:n_test_pts])
+            neigh.fit(X, y_str)
+            y_pred = neigh.predict(X[:n_test_pts] + epsilon)
+            assert_array_equal(y_pred, y_str[:n_test_pts])
 
 
 def test_radius_neighbors_classifier_when_no_neighbors():
