Repository: mars-project/mars

Issue:
[BUG] Fix compatibility with kmeans for scikit-learn 0.24.0
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

Scikit-learn has upgraded to 0.24.0, and the result of kmeans is not consistent, we need to fix the compatibility.




Discussion:

Resolving patch:
diff --git a/mars/learn/cluster/_k_means_init.py b/mars/learn/cluster/_k_means_init.py
--- a/mars/learn/cluster/_k_means_init.py
+++ b/mars/learn/cluster/_k_means_init.py
@@ -200,9 +200,9 @@ def _kmeans_plusplus(*args, **kwargs):
             [ctx[inp.key] for inp in op.inputs], device=op.device, ret_extra=True)
 
         with device(device_id):
-            ctx[op.outputs[0].key], _ = _kmeans_plusplus(
+            ctx[op.outputs[0].key] = _kmeans_plusplus(
                 x, op.n_clusters, x_squared_norms=x_squared_norms, random_state=op.state,
-                n_local_trials=op.n_local_trials)
+                n_local_trials=op.n_local_trials)[0]
 
 
 ###############################################################################
diff --git a/mars/learn/cluster/_kmeans.py b/mars/learn/cluster/_kmeans.py
--- a/mars/learn/cluster/_kmeans.py
+++ b/mars/learn/cluster/_kmeans.py
@@ -180,19 +180,12 @@ def k_means(X, n_clusters, sample_weight=None, init='k-means||',
         return est.cluster_centers_, est.labels_, est.inertia_
 
 
-def _kmeans_single_elkan(X, sample_weight, n_clusters, max_iter=300,
-                         init='k-means++', verbose=False, x_squared_norms=None,
-                         random_state=None, tol=1e-4, oversampling_factor=2,
-                         init_iter=5, X_mean=None, session=None, run_kwargs=None):
-    random_state = check_random_state(random_state)
+def _kmeans_single_elkan(X, sample_weight, centers_init, n_clusters, max_iter=300,
+                         verbose=False, x_squared_norms=None,
+                         tol=1e-4, X_mean=None, session=None, run_kwargs=None):
     sample_weight = _check_normalize_sample_weight(sample_weight, X)
 
-    # init
-    centers = _init_centroids(X, n_clusters, init, random_state=random_state,
-                              x_squared_norms=x_squared_norms,
-                              oversampling_factor=oversampling_factor,
-                              init_iter=init_iter)
-
+    centers = centers_init
     # execute X, centers and tol first
     tol = mt.asarray(tol)
     to_run = [X, sample_weight, centers, x_squared_norms, tol]
@@ -266,19 +259,12 @@ def _kmeans_single_elkan(X, sample_weight, n_clusters, max_iter=300,
     return labels, inertia, centers, i + 1
 
 
-def _kmeans_single_lloyd(X, sample_weight, n_clusters, max_iter=300,
-                         init='k-means++', verbose=False, x_squared_norms=None,
-                         random_state=None, tol=1e-4, oversampling_factor=2,
-                         init_iter=5, X_mean=None, session=None, run_kwargs=None):
-    random_state = check_random_state(random_state)
+def _kmeans_single_lloyd(X, sample_weight, centers_init, n_clusters, max_iter=300,
+                         verbose=False, x_squared_norms=None,
+                         tol=1e-4, X_mean=None, session=None, run_kwargs=None):
     sample_weight = _check_normalize_sample_weight(sample_weight, X)
 
-    # init
-    centers = _init_centroids(X, n_clusters, init, random_state=random_state,
-                              x_squared_norms=x_squared_norms,
-                              oversampling_factor=oversampling_factor,
-                              init_iter=init_iter)
-
+    centers = centers_init
     # execute X, centers and tol first
     tol = mt.asarray(tol)
     to_run = [X, centers, x_squared_norms, tol]
@@ -622,6 +608,60 @@ def __init__(self, n_clusters=8, init='k-means||', n_init=1,
         self.oversampling_factor = oversampling_factor
         self.init_iter = init_iter
 
+    def _check_params(self, X):
+        # n_init
+        if self.n_init <= 0:
+            raise ValueError(
+                f"n_init should be > 0, got {self.n_init} instead.")
+        self._n_init = self.n_init
+
+        # max_iter
+        if self.max_iter <= 0:
+            raise ValueError(
+                f"max_iter should be > 0, got {self.max_iter} instead.")
+
+        # n_clusters
+        if X.shape[0] < self.n_clusters:
+            raise ValueError(f"n_samples={X.shape[0]} should be >= "
+                             f"n_clusters={self.n_clusters}.")
+
+        # tol
+        self._tol = _tolerance(X, self.tol)
+
+        # algorithm
+        if self.algorithm not in ("auto", "full", "elkan"):
+            raise ValueError(f"Algorithm must be 'auto', 'full' or 'elkan', "
+                             f"got {self.algorithm} instead.")
+
+        self._algorithm = self.algorithm
+        if self._algorithm == "auto":
+            # note(xuye.qin):
+            # Different from scikit-learn,
+            # for now, full seems more efficient when data is large,
+            # elkan needs to be tuned more
+            # old: algorithm = "full" if self.n_clusters == 1 else "elkan"
+            self._algorithm = "full"
+        if self._algorithm == "elkan" and self.n_clusters == 1:
+            warnings.warn("algorithm='elkan' doesn't make sense for a single "
+                          "cluster. Using 'full' instead.", RuntimeWarning)
+            self._algorithm = "full"
+
+        # init
+        if not (hasattr(self.init, '__array__') or callable(self.init)
+                or (isinstance(self.init, str)
+                    and self.init in ["k-means++", "k-means||", "random"])):
+            raise ValueError(
+                f"init should be either 'k-means++'， 'k-mean||', 'random', "
+                f"a tensor, a ndarray or a "
+                f"callable, got '{self.init}' instead.")
+
+        if hasattr(self.init, '__array__') and self._n_init != 1:
+            warnings.warn(
+                f"Explicit initial center position passed: performing only"
+                f" one init in {self.__class__.__name__} instead of "
+                f"n_init={self._n_init}.", RuntimeWarning, stacklevel=2)
+            self._n_init = 1
+
     def _check_test_data(self, X):
         X = check_array(X, accept_sparse=True, dtype=[np.float64, np.float32],
                         order='C', accept_large_sparse=False)
@@ -716,19 +756,6 @@ def fit(self, X, y=None, sample_weight=None, session=None, run_kwargs=None):
         self
             Fitted estimator.
         """
-        random_state = check_random_state(self.random_state).to_numpy()
-
-        n_init = self.n_init
-        if n_init <= 0:
-            raise ValueError("Invalid number of initializations. "
-                             f"n_init={n_init} must be bigger than zero.")
-
-        if self.max_iter <= 0:  # pragma: no cover
-            raise ValueError(
-                'Number of iterations should be a positive number, '
-                f'got {self.max_iter} instead'
-            )
-
         expect_chunk_size_on_columns = mt.tensor(X).shape[1]
         if not np.isnan(expect_chunk_size_on_columns):
             X = mt.tensor(X, chunk_size={1: expect_chunk_size_on_columns})
@@ -738,12 +765,11 @@ def fit(self, X, y=None, sample_weight=None, session=None, run_kwargs=None):
                                 order='C', copy=self.copy_x,
                                 accept_large_sparse=False)
         # verify that the number of samples given is larger than k
-        n_samples = _num_samples(X)
-        if np.isnan(n_samples):  # pragma: no cover
+        if np.isnan(_num_samples(X)):  # pragma: no cover
             X.execute(session=session, **(run_kwargs or dict()))
-            n_samples = _num_samples(X)
-        if n_samples < self.n_clusters:
-            raise ValueError(f"n_samples={n_samples} should be >= n_clusters={self.n_clusters}")
+
+        self._check_params(X)
+        random_state = check_random_state(self.random_state).to_numpy()
 
         tol = _tolerance(X, self.tol)
 
@@ -753,13 +779,6 @@ def fit(self, X, y=None, sample_weight=None, session=None, run_kwargs=None):
             init = check_array(init, dtype=X.dtype.type, copy=True, order='C')
             _validate_center_shape(X, self.n_clusters, init)
 
-            if n_init != 1:  # pragma: no cover
-                warnings.warn(
-                    'Explicit initial center position passed: '
-                    'performing only one init in k-means instead of n_init=%d'
-                    % n_init, RuntimeWarning, stacklevel=2)
-                n_init = 1
-
         # subtract of mean of x for more accurate distance computations
         X_mean = None
         if not X.issparse():
@@ -775,38 +794,24 @@ def fit(self, X, y=None, sample_weight=None, session=None, run_kwargs=None):
 
         best_labels, best_inertia, best_centers = None, None, None
 
-        algorithm = self.algorithm
-        if algorithm == "elkan" and self.n_clusters == 1:  # pragma: no cover
-            warnings.warn("algorithm='elkan' doesn't make sense for a single "
-                          "cluster. Using 'full' instead.", RuntimeWarning)
-            algorithm = "full"
-
-        if algorithm == "auto":
-            # note(xuye.qin):
-            # Different from scikit-learn,
-            # for now, full seems more efficient when data is large,
-            # elkan needs to be tuned more
-            # old: algorithm = "full" if self.n_clusters == 1 else "elkan"
-            algorithm = "full"
-
-        if algorithm == "full":
+        if self._algorithm == "full":
             kmeans_single = _kmeans_single_lloyd
-        elif algorithm == "elkan":
+        else:
             kmeans_single = _kmeans_single_elkan
-        else:  # pragma: no cover
-            raise ValueError(f"Algorithm must be 'auto', 'full' or 'elkan', got {algorithm}")
 
-        # seeds for the initializations of the kmeans runs.
-        seeds = random_state.randint(np.iinfo(np.int32).max, size=n_init)
+        for i in range(self._n_init):  # pylint: disable=unused-variable
+            # Initialize centers
+            centers_init = _init_centroids(
+                X, self.n_clusters, init, random_state=random_state,
+                x_squared_norms=x_squared_norms,
+                oversampling_factor=self.oversampling_factor,
+                init_iter=self.init_iter)
 
-        for seed in seeds:
             # run a k-means once
             labels, inertia, centers, n_iter_ = kmeans_single(
-                X, sample_weight, self.n_clusters, max_iter=self.max_iter,
-                init=init, verbose=self.verbose, tol=tol,
-                x_squared_norms=x_squared_norms, random_state=seed,
-                oversampling_factor=self.oversampling_factor,
-                init_iter=self.init_iter, X_mean=X_mean,
+                X, sample_weight, centers_init, self.n_clusters,
+                max_iter=self.max_iter, verbose=self.verbose, tol=tol,
+                x_squared_norms=x_squared_norms, X_mean=X_mean,
                 session=session, run_kwargs=run_kwargs)
             inertia = inertia.fetch(session=session)
             # determine if these results are the best so far


Test patch:
diff --git a/mars/learn/cluster/tests/test_k_means.py b/mars/learn/cluster/tests/test_k_means.py
--- a/mars/learn/cluster/tests/test_k_means.py
+++ b/mars/learn/cluster/tests/test_k_means.py
@@ -20,6 +20,7 @@
 import pytest
 import scipy.sparse as sp
 try:
+    from sklearn.cluster import KMeans as SK_KMeans
     from sklearn.datasets import make_blobs
     from sklearn.metrics.cluster import v_measure_score
     from sklearn.utils._testing import assert_raise_message, assert_warns
@@ -420,3 +421,24 @@ def testKMeansInitLargeNClusters(self):
             nbytes = c.nbytes
             if not np.isnan(nbytes):
                 self.assertLessEqual(nbytes, chunk_bytes_limit)
+
+    def testConsistentResultWithSklearn(self):
+        rnd = np.random.RandomState(0)
+        X, _ = make_blobs(random_state=rnd)
+        raw = X
+        X = mt.tensor(X, chunk_size=50)
+
+        km_elkan = KMeans(algorithm='elkan', n_clusters=5,
+                          random_state=0, n_init=1, tol=1e-4,
+                          init='k-means++')
+        sk_km_elkan = SK_KMeans(algorithm='elkan', n_clusters=5,
+                                random_state=0, n_init=1, tol=1e-4,
+                                init='k-means++')
+
+        km_elkan.fit(X)
+        sk_km_elkan.fit(raw)
+
+        np.testing.assert_allclose(km_elkan.cluster_centers_, sk_km_elkan.cluster_centers_)
+        np.testing.assert_array_equal(km_elkan.labels_, sk_km_elkan.labels_)
+
+        self.assertEqual(km_elkan.n_iter_, sk_km_elkan.n_iter_)