Repository: scikit-learn/scikit-learn
Pull request: ENH KMeans initialization account for sample weights
Fixes #25527

Problem statement:
KMeans initialization does not use sample weights
### Describe the bug

Clustering by KMeans does not weight the input data.

### Steps/Code to Reproduce

```py
import numpy as np
from sklearn.cluster import KMeans
x = np.array([1, 1, 5, 5, 100, 100])
w = 10**np.array([8.,8,8,8,-8,-8]) # large weights for 1 and 5, small weights for 100
x=x.reshape(-1,1)# reshape to a 2-dimensional array requested for KMeans
centers_with_weight = KMeans(n_clusters=2, random_state=0,n_init=10).fit(x,sample_weight=w).cluster_centers_
centers_no_weight = KMeans(n_clusters=2, random_state=0,n_init=10).fit(x).cluster_centers_
```

### Expected Results

centers_with_weight=[[1.],[5.]]
centers_no_weight=[[100.],[3.]]

### Actual Results

centers_with_weight=[[100.],[3.]]
centers_no_weight=[[100.],[3.]]

### Versions

```shell
System:
    python: 3.10.4 (tags/v3.10.4:9d38120, Mar 23 2022, 23:13:41) [MSC v.1929 64 bit (AMD64)]
executable: E:\WPy64-31040\python-3.10.4.amd64\python.exe
   machine: Windows-10-10.0.19045-SP0

Python dependencies:
      sklearn: 1.2.1
          pip: 22.3.1
   setuptools: 62.1.0
        numpy: 1.23.3
        scipy: 1.8.1
       Cython: 0.29.28
       pandas: 1.4.2
   matplotlib: 3.5.1
       joblib: 1.2.0
threadpoolctl: 3.1.0

Built with OpenMP: True

threadpoolctl info:
       user_api: blas
   internal_api: openblas
         prefix: libopenblas
       filepath: E:\WPy64-31040\python-3.10.4.amd64\Lib\site-packages\numpy\.libs\libopenblas.FB5AE2TYXYH2IJRDKGDGQ3XBKLKTF43H.gfortran-win_amd64.dll
        version: 0.3.20
threading_layer: pthreads
   architecture: Haswell
    num_threads: 12

       user_api: blas
   internal_api: openblas
         prefix: libopenblas
       filepath: E:\WPy64-31040\python-3.10.4.amd64\Lib\site-packages\scipy\.libs\libopenblas.XWYDX2IKJW2NMTWSFYNGFUWKQU3LYTCZ.gfortran-win_amd64.dll
        version: 0.3.17
threading_layer: pthreads
   architecture: Haswell
    num_threads: 12

       user_api: openmp
   internal_api: openmp
         prefix: vcomp
       filepath: E:\WPy64-31040\python-3.10.4.amd64\Lib\site-packages\sklearn\.libs\vcomp140.dll
        version: None
    num_threads: 12
```

Hints from issue discussion:
Thanks for the reproducible example.

`KMeans` **does** weight the data, but your example is an extreme case. Because `Kmeans` is a non-convex problem, the algorithm can get stuck in a local minimum, and not find the true minimum  of the optimization landscape. This is the reason why the code proposes to use multiple initializations, hoping that some initializations will not get stuck in poor local minima.

Importantly, **the initialization does not use sample weights**. So when using `init="k-means++"` (default), it gets a centroid on the outliers, and `Kmeans` cannot escape this local minimum, even with strong sample weights.

Instead, the optimization does not get stuck if we do one of the following changes:
- the outliers are less extremes (e.g. 10 instead of 100), because it allows the sample weights to remove the local minimum
- the low samples weights are exactly equal to zero, because it completely discard the outliers
- using `init="random"`, because some of the different init are able to avoid the local minimum
If we have less extreme weighting then ```init='random'``` does not work at all as shown in this example. The centers are always [[3.],[40.]] instead of circa [[1.],[5.]] for the weighted case.
```py
import numpy as np
from sklearn.cluster import KMeans
x = np.array([1, 1, 5, 5, 40, 40])
w = np.array([100,100,100,100,1,1]) # large weights for 1 and 5, small weights for 40
x=x.reshape(-1,1)# reshape to a 2-dimensional array requested for KMeans
centers_with_weight = KMeans(n_clusters=2, random_state=0,n_init=100,init='random').fit(x,sample_weight=w).cluster_centers_
centers_no_weight = KMeans(n_clusters=2, random_state=0,n_init=100).fit(x).cluster_centers_
```
Sure, you can find examples where `init="random"` also fails.
My point is that `KMeans` is non-convex, so it is never guaranteed to find the global minimum.

That being said, **it would be nice to have an initialization that uses the sample weights**. I don't know about `init="k-means++"`, but for `init="random"` we could use a non-uniform sampling relative to the sample weights, instead of the current uniform sampling.

Do you want to submit a pull-request?
yes
I want to work on this issue
Was this bug already resolved?
Not resolved yet, you are welcome to submit a pull-request if you'd like.
I already answered this question by "yes" 3 weeks ago. 
Is anyone working on this?

I'll create a draft and close it if someone will come forward

Commit message: Update _kmeans.py

Patch to review:
diff --git a/doc/whats_new/v1.3.rst b/doc/whats_new/v1.3.rst
index 7f4450975cd60..c902a0da15c54 100644
--- a/doc/whats_new/v1.3.rst
+++ b/doc/whats_new/v1.3.rst
@@ -107,9 +107,18 @@ Changelog
 
 - |API| The `sample_weight` parameter in `predict` for
   :meth:`cluster.KMeans.predict` and :meth:`cluster.MiniBatchKMeans.predict`
-  is now deprecated and will be removed in v1.5.
+  is now deprecated and will be removed in v1.5. 
   :pr:`25251` by :user:`Gleb Levitski <glevv>`.
 
+- |Feature| The `sample_weight` parameter now will be used in centroids
+  initialization for :class:`cluster.KMeans`, :class:`cluster.BisectingKMeans`
+  and :class:`cluster.MiniBatchKMeans`.
+  This change will break backward compatibility, since numbers generated
+  from same random seeds will be different.
+  :pr:`25752` by :user:`Gleb Levitski <glevv>`,
+  :user:`Jérémie du Boisberranger <jeremiedbb>`,
+  :user:`Guillaume Lemaitre <glemaitre>`.
+
 :mod:`sklearn.datasets`
 .......................
 
diff --git a/sklearn/cluster/_bisect_k_means.py b/sklearn/cluster/_bisect_k_means.py
index b860596c03540..486435cecec34 100644
--- a/sklearn/cluster/_bisect_k_means.py
+++ b/sklearn/cluster/_bisect_k_means.py
@@ -190,18 +190,18 @@ class BisectingKMeans(_BaseKMeans):
     --------
     >>> from sklearn.cluster import BisectingKMeans
     >>> import numpy as np
-    >>> X = np.array([[1, 2], [1, 4], [1, 0],
-    ...               [10, 2], [10, 4], [10, 0],
-    ...               [10, 6], [10, 8], [10, 10]])
+    >>> X = np.array([[1, 1], [10, 1], [3, 1],
+    ...               [10, 0], [2, 1], [10, 2],
+    ...               [10, 8], [10, 9], [10, 10]])
     >>> bisect_means = BisectingKMeans(n_clusters=3, random_state=0).fit(X)
     >>> bisect_means.labels_
-    array([2, 2, 2, 0, 0, 0, 1, 1, 1], dtype=int32)
+    array([0, 2, 0, 2, 0, 2, 1, 1, 1], dtype=int32)
     >>> bisect_means.predict([[0, 0], [12, 3]])
-    array([2, 0], dtype=int32)
+    array([0, 2], dtype=int32)
     >>> bisect_means.cluster_centers_
-    array([[10.,  2.],
-           [10.,  8.],
-           [ 1., 2.]])
+    array([[ 2., 1.],
+           [10., 9.],
+           [10., 1.]])
     """
 
     _parameter_constraints: dict = {
@@ -309,7 +309,12 @@ def _bisect(self, X, x_squared_norms, sample_weight, cluster_to_bisect):
         # Repeating `n_init` times to obtain best clusters
         for _ in range(self.n_init):
             centers_init = self._init_centroids(
-                X, x_squared_norms, self.init, self._random_state, n_centroids=2
+                X,
+                x_squared_norms=x_squared_norms,
+                init=self.init,
+                random_state=self._random_state,
+                n_centroids=2,
+                sample_weight=sample_weight,
             )
 
             labels, inertia, centers, _ = self._kmeans_single(
@@ -361,7 +366,8 @@ def fit(self, X, y=None, sample_weight=None):
 
         sample_weight : array-like of shape (n_samples,), default=None
             The weights for each observation in X. If None, all observations
-            are assigned equal weight.
+            are assigned equal weight. If not None, weights will be used in
+            initializations of centroids as well as fitting clusters.
 
         Returns
         -------
diff --git a/sklearn/cluster/_kmeans.py b/sklearn/cluster/_kmeans.py
index 11d2b81cd552f..8c02dd1ef55d9 100644
--- a/sklearn/cluster/_kmeans.py
+++ b/sklearn/cluster/_kmeans.py
@@ -26,7 +26,7 @@
 )
 from ..metrics.pairwise import euclidean_distances
 from ..metrics.pairwise import _euclidean_distances
-from ..utils.extmath import row_norms, stable_cumsum
+from ..utils.extmath import row_norms
 from ..utils.fixes import threadpool_limits
 from ..utils.fixes import threadpool_info
 from ..utils.sparsefuncs_fast import assign_rows_csr
@@ -160,7 +160,9 @@ def kmeans_plusplus(
     return centers, indices
 
 
-def _kmeans_plusplus(X, n_clusters, x_squared_norms, random_state, n_local_trials=None):
+def _kmeans_plusplus(
+    X, n_clusters, x_squared_norms, sample_weight, random_state, n_local_trials=None
+):
     """Computational component for initialization of n_clusters by
     k-means++. Prior validation of data is assumed.
 
@@ -206,7 +208,7 @@ def _kmeans_plusplus(X, n_clusters, x_squared_norms, random_state, n_local_trial
         n_local_trials = 2 + int(np.log(n_clusters))
 
     # Pick first center randomly and track index of point
-    center_id = random_state.randint(n_samples)
+    center_id = random_state.choice(n_samples, replace=False, p=sample_weight)
     indices = np.full(n_clusters, -1, dtype=int)
     if sp.issparse(X):
         centers[0] = X[center_id].toarray()
@@ -218,16 +220,18 @@ def _kmeans_plusplus(X, n_clusters, x_squared_norms, random_state, n_local_trial
     closest_dist_sq = _euclidean_distances(
         centers[0, np.newaxis], X, Y_norm_squared=x_squared_norms, squared=True
     )
-    current_pot = closest_dist_sq.sum()
+    current_pot = closest_dist_sq @ sample_weight
 
     # Pick the remaining n_clusters-1 points
     for c in range(1, n_clusters):
         # Choose center candidates by sampling with probability proportional
         # to the squared distance to the closest existing center
-        rand_vals = random_state.uniform(size=n_local_trials) * current_pot
-        candidate_ids = np.searchsorted(stable_cumsum(closest_dist_sq), rand_vals)
-        # XXX: numerical imprecision can result in a candidate_id out of range
-        np.clip(candidate_ids, None, closest_dist_sq.size - 1, out=candidate_ids)
+        candidate_ids = random_state.choice(
+            n_samples,
+            size=n_local_trials,
+            replace=False,
+            p=sample_weight * closest_dist_sq / current_pot,
+        )
 
         # Compute distances to center candidates
         distance_to_candidates = _euclidean_distances(
@@ -236,7 +240,7 @@ def _kmeans_plusplus(X, n_clusters, x_squared_norms, random_state, n_local_trial
 
         # update closest distances squared and potential for each candidate
         np.minimum(closest_dist_sq, distance_to_candidates, out=distance_to_candidates)
-        candidates_pot = distance_to_candidates.sum(axis=1)
+        candidates_pot = distance_to_candidates @ sample_weight.reshape(-1, 1)
 
         # Decide which candidate is the best
         best_candidate = np.argmin(candidates_pot)
@@ -323,7 +327,8 @@ def k_means(
 
     sample_weight : array-like of shape (n_samples,), default=None
         The weights for each observation in `X`. If `None`, all observations
-        are assigned equal weight.
+        are assigned equal weight. If not None, weights will be used in
+        initializations of centroids as well as fitting clusters.
 
     init : {'k-means++', 'random'}, callable or array-like of shape \
             (n_clusters, n_features), default='k-means++'
@@ -939,7 +944,14 @@ def _check_test_data(self, X):
         return X
 
     def _init_centroids(
-        self, X, x_squared_norms, init, random_state, init_size=None, n_centroids=None
+        self,
+        X,
+        x_squared_norms,
+        init,
+        random_state,
+        init_size=None,
+        n_centroids=None,
+        sample_weight=None,
     ):
         """Compute the initial centroids.
 
@@ -969,6 +981,10 @@ def _init_centroids(
             If left to 'None' the number of centroids will be equal to
             number of clusters to form (self.n_clusters)
 
+        sample_weight : ndarray of shape (n_samples,), default=None
+            The weights for each observation in X. If None, all observations
+            are assigned equal weight.
+
         Returns
         -------
         centers : ndarray of shape (n_clusters, n_features)
@@ -981,6 +997,7 @@ def _init_centroids(
             X = X[init_indices]
             x_squared_norms = x_squared_norms[init_indices]
             n_samples = X.shape[0]
+            sample_weight = sample_weight[init_indices]
 
         if isinstance(init, str) and init == "k-means++":
             centers, _ = _kmeans_plusplus(
@@ -988,9 +1005,15 @@ def _init_centroids(
                 n_clusters,
                 random_state=random_state,
                 x_squared_norms=x_squared_norms,
+                sample_weight=sample_weight,
             )
         elif isinstance(init, str) and init == "random":
-            seeds = random_state.permutation(n_samples)[:n_clusters]
+            seeds = random_state.choice(
+                n_samples,
+                size=n_clusters,
+                replace=False,
+                p=sample_weight / sample_weight.sum(),
+            )
             centers = X[seeds]
         elif _is_arraylike_not_scalar(self.init):
             centers = init
@@ -1412,7 +1435,8 @@ def fit(self, X, y=None, sample_weight=None):
 
         sample_weight : array-like of shape (n_samples,), default=None
             The weights for each observation in X. If None, all observations
-            are assigned equal weight.
+            are assigned equal weight. If not None, weights will be used in
+            initializations of centroids as well as fitting clusters.
 
             .. versionadded:: 0.20
 
@@ -1468,7 +1492,11 @@ def fit(self, X, y=None, sample_weight=None):
         for i in range(self._n_init):
             # Initialize centers
             centers_init = self._init_centroids(
-                X, x_squared_norms=x_squared_norms, init=init, random_state=random_state
+                X,
+                x_squared_norms=x_squared_norms,
+                init=init,
+                random_state=random_state,
+                sample_weight=sample_weight,
             )
             if self.verbose:
                 print("Initialization complete")
@@ -2015,7 +2043,8 @@ def fit(self, X, y=None, sample_weight=None):
 
         sample_weight : array-like of shape (n_samples,), default=None
             The weights for each observation in X. If None, all observations
-            are assigned equal weight.
+            are assigned equal weight. If not None, weights will be used in
+            initializations of centroids as well as fitting clusters.
 
             .. versionadded:: 0.20
 
@@ -2070,6 +2099,7 @@ def fit(self, X, y=None, sample_weight=None):
                 init=init,
                 random_state=random_state,
                 init_size=self._init_size,
+                sample_weight=sample_weight,
             )
 
             # Compute inertia on a validation set.
@@ -2170,7 +2200,8 @@ def partial_fit(self, X, y=None, sample_weight=None):
 
         sample_weight : array-like of shape (n_samples,), default=None
             The weights for each observation in X. If None, all observations
-            are assigned equal weight.
+            are assigned equal weight. If not None, weights will be used in
+            initializations of centroids as well as fitting clusters.
 
         Returns
         -------
@@ -2220,6 +2251,7 @@ def partial_fit(self, X, y=None, sample_weight=None):
                 init=init,
                 random_state=self._random_state,
                 init_size=self._init_size,
+                sample_weight=sample_weight,
             )
 
             # Initialize counts
diff --git a/sklearn/cluster/tests/test_bisect_k_means.py b/sklearn/cluster/tests/test_bisect_k_means.py
index 5f69ac1353a39..e1ab3e60f8d9f 100644
--- a/sklearn/cluster/tests/test_bisect_k_means.py
+++ b/sklearn/cluster/tests/test_bisect_k_means.py
@@ -4,34 +4,33 @@
 
 from sklearn.utils._testing import assert_array_equal, assert_allclose
 from sklearn.cluster import BisectingKMeans
+from sklearn.metrics import v_measure_score
 
 
 @pytest.mark.parametrize("bisecting_strategy", ["biggest_inertia", "largest_cluster"])
-def test_three_clusters(bisecting_strategy):
+@pytest.mark.parametrize("init", ["k-means++", "random"])
+def test_three_clusters(bisecting_strategy, init):
     """Tries to perform bisect k-means for three clusters to check
     if splitting data is performed correctly.
     """
-
-    # X = np.array([[1, 2], [1, 4], [1, 0],
-    #               [10, 2], [10, 4], [10, 0],
-    #               [10, 6], [10, 8], [10, 10]])
-
-    # X[0][1] swapped with X[1][1] intentionally for checking labeling
     X = np.array(
-        [[1, 2], [10, 4], [1, 0], [10, 2], [1, 4], [10, 0], [10, 6], [10, 8], [10, 10]]
+        [[1, 1], [10, 1], [3, 1], [10, 0], [2, 1], [10, 2], [10, 8], [10, 9], [10, 10]]
     )
     bisect_means = BisectingKMeans(
-        n_clusters=3, random_state=0, bisecting_strategy=bisecting_strategy
+        n_clusters=3,
+        random_state=0,
+        bisecting_strategy=bisecting_strategy,
+        init=init,
     )
     bisect_means.fit(X)
 
-    expected_centers = [[10, 2], [10, 8], [1, 2]]
-    expected_predict = [2, 0]
-    expected_labels = [2, 0, 2, 0, 2, 0, 1, 1, 1]
+    expected_centers = [[2, 1], [10, 1], [10, 9]]
+    expected_labels = [0, 1, 0, 1, 0, 1, 2, 2, 2]
 
-    assert_allclose(expected_centers, bisect_means.cluster_centers_)
-    assert_array_equal(expected_predict, bisect_means.predict([[0, 0], [12, 3]]))
-    assert_array_equal(expected_labels, bisect_means.labels_)
+    assert_allclose(
+        sorted(expected_centers), sorted(bisect_means.cluster_centers_.tolist())
+    )
+    assert_allclose(v_measure_score(expected_labels, bisect_means.labels_), 1.0)
 
 
 def test_sparse():
@@ -133,3 +132,16 @@ def test_float32_float64_equivalence(is_sparse):
 
     assert_allclose(km32.cluster_centers_, km64.cluster_centers_)
     assert_array_equal(km32.labels_, km64.labels_)
+
+
+def test_sample_weight_init_bisect(global_random_seed):
+    """Check that sample weight is used during init."""
+    rng = np.random.RandomState(global_random_seed)
+    X = rng.rand(20, 2)
+    sample_weight = rng.uniform(size=20)
+    kmeans = BisectingKMeans(
+        n_clusters=3, random_state=global_random_seed, init="random"
+    )
+    clusters_weighted = kmeans.fit(X, sample_weight=sample_weight).cluster_centers_
+    clusters = kmeans.fit(X).cluster_centers_
+    assert (clusters_weighted != clusters).any()
diff --git a/sklearn/cluster/tests/test_k_means.py b/sklearn/cluster/tests/test_k_means.py
index 9ae7510c45dba..5c83b39bbbb6f 100644
--- a/sklearn/cluster/tests/test_k_means.py
+++ b/sklearn/cluster/tests/test_k_means.py
@@ -1276,3 +1276,18 @@ def test_predict_does_not_change_cluster_centers(is_sparse):
 
     y_pred2 = kmeans.predict(X)
     assert_array_equal(y_pred1, y_pred2)
+
+
+@pytest.mark.parametrize("Estimator", [KMeans, MiniBatchKMeans])
+@pytest.mark.parametrize("init", ["k-means++", "random"])
+def test_sample_weight_init(Estimator, init, global_random_seed):
+    """Check that sample weight is used during init."""
+    rng = np.random.RandomState(global_random_seed)
+    X, _ = make_blobs(
+        n_samples=200, n_features=10, centers=10, random_state=global_random_seed
+    )
+    sample_weight = rng.uniform(size=200)
+    kmeans = Estimator(random_state=global_random_seed, init=init)
+    clusters_weighted = kmeans.fit(X, sample_weight=sample_weight).cluster_centers_
+    clusters = kmeans.fit(X).cluster_centers_
+    assert (clusters_weighted != clusters).any()


Code review comments:
diff_hunk: @@ -206,7 +208,7 @@ def _kmeans_plusplus(X, n_clusters, x_squared_norms, random_state, n_local_trial
         n_local_trials = 2 + int(np.log(n_clusters))
 
     # Pick first center randomly and track index of point
-    center_id = random_state.randint(n_samples)
+    center_id = random_state.choice(n_samples, replace=False, p=sample_weight)
path: sklearn/cluster/_kmeans.py
text: @user1:
probas need to sum to 1. replace=False is not necessary since we're only sampling 1 value.
```suggestion
    center_id = random_state.choice(n_samples, p=sample_weight / sample_weight.sum())
```

diff_hunk: @@ -206,7 +208,7 @@ def _kmeans_plusplus(X, n_clusters, x_squared_norms, random_state, n_local_trial
         n_local_trials = 2 + int(np.log(n_clusters))
 
     # Pick first center randomly and track index of point
-    center_id = random_state.randint(n_samples)
+    center_id = random_state.choice(n_samples, replace=False, p=sample_weight)
path: sklearn/cluster/_kmeans.py
text: @user2:
@user1 this is a breaking change even when sample_weight=None (transformed into ones at this point). Do you find it acceptable ?

@user1:
I think this is fine. Only the randomness is different but the behaviour is the same. We can document this change in the "Change models". Since we add support for `sample_weight` I think that we can go ahead here.