Fix check_array dtype in MinibatchKMeans.partial_fit (#14323)

rth · jeremiedbb · commit bf8eff3feade · 2019-07-15T11:17:11.000+02:00
diff --git a/sklearn/cluster/k_means_.py b/sklearn/cluster/k_means_.py
@@ -1422,8 +1422,8 @@ class MiniBatchKMeans(KMeans):
     >>> kmeans = kmeans.partial_fit(X[0:6,:])
     >>> kmeans = kmeans.partial_fit(X[6:12,:])
     >>> kmeans.cluster_centers_
-    array([[1, 1],
-           [3, 4]])
+    array([[2. , 1. ],
+           [3.5, 4.5]])
     >>> kmeans.predict([[0, 0], [4, 4]])
     array([0, 1], dtype=int32)
     >>> # fit on the whole data
@@ -1667,7 +1667,8 @@ def partial_fit(self, X, y=None, sample_weight=None):
 
         """
 
-        X = check_array(X, accept_sparse="csr", order="C")
+        X = check_array(X, accept_sparse="csr", order="C",
+                        dtype=[np.float64, np.float32])
         n_samples, n_features = X.shape
         if hasattr(self.init, '__array__'):
             self.init = np.ascontiguousarray(self.init, dtype=X.dtype)
diff --git a/sklearn/cluster/tests/test_k_means.py b/sklearn/cluster/tests/test_k_means.py
@@ -942,3 +942,11 @@ def test_k_means_empty_cluster_relocated():
 
     assert len(set(km.labels_)) == 2
     assert_allclose(km.cluster_centers_, [[-1], [1]])
+
+
+def test_minibatch_kmeans_partial_fit_int_data():
+    # Issue GH #14314
+    X = np.array([[-1], [1]], dtype=np.int)
+    km = MiniBatchKMeans(n_clusters=2)
+    km.partial_fit(X)
+    assert km.cluster_centers_.dtype.kind == "f"
