Repository: scikit-learn/scikit-learn
File: sklearn/cluster/_kmeans.py (Python)

Code before:
        # Attributes to monitor the convergence
        self._ewa_inertia = None
        self._ewa_inertia_min = None
        self._no_improvement = 0

        # Initialize number of samples seen since last reassignment
        self._n_since_last_reassign = 0

        n_steps = (self.max_iter * n_samples) // self._batch_size

        with _get_threadpool_controller().limit(limits=1, user_api="blas"):
            # Perform the iterative optimization until convergence
            for i in range(n_steps):
                # Sample a minibatch from the full dataset
                minibatch_indices = random_state.choice(
                    n_samples,
                    self._batch_size,
                    p=sample_weight / np.sum(sample_weight),
                    replace=True,
                )

                # Perform the actual update step on the minibatch data
                # Note, I am not sure how sample weights are used here
                # So left it in, it seems like the weight sums are updated using
                # sample weights so need some help here to understand the
                # _minibatch_update_dense/sparse code
                batch_inertia = _mini_batch_step(
                    X=X[minibatch_indices],
                    sample_weight=sample_weight[minibatch_indices],
                    centers=centers,
                    centers_new=centers_new,
                    weight_sums=self._counts,
                    random_state=random_state,
                    random_reassign=self._random_reassign(),
                    reassignment_ratio=self.reassignment_ratio,
                    verbose=self.verbose,
                    n_threads=self._n_threads,
                )

                if self._tol > 0.0:
                    centers_squared_diff = np.sum((centers_new - centers) ** 2)
                else:
                    centers_squared_diff = 0

                centers, centers_new = centers_new, centers

                # Monitor convergence and do early stopping if necessary
                if self._mini_batch_convergence(
                    i, n_steps, n_samples, centers_squared_diff, batch_inertia
                ):
                    break

Review comment (performance):
Like in KBinsDiscretizer, we should not use sample weights after using them for weighted sampling.
So let's pass an array of ones as sample weights here.

> I am not sure how sample weights are used here
   So left it in, it seems like the weight sums are updated using
   sample weights so need some help here to understand

weight_sum is the sum of weights of all points belonging to each clusters. It's used to track clusters where there are very few points (more precisely points that add up to a small weight) and reassign them to a different cluster.

I don't think we have to modify `_mini_batch_step`. It's still useful that it handles sample weights because it's also used in partial_fit where there is no sampling so sample weights must be passed.

Code after:

            # Compute inertia on a validation set.
            _, inertia = _labels_inertia_threadpool_limit(
                X_valid,
                sample_weight_valid,
                cluster_centers,
                n_threads=self._n_threads,
            )

            if self.verbose:
                print(f"Inertia for init {init_idx + 1}/{self._n_init}: {inertia}")
            if best_inertia is None or inertia < best_inertia:
                init_centers = cluster_centers
                best_inertia = inertia

        centers = init_centers
        centers_new = np.empty_like(centers)

        # Initialize counts
        self._counts = np.zeros(self.n_clusters, dtype=X.dtype)

        # Attributes to monitor the convergence
        self._ewa_inertia = None
        self._ewa_inertia_min = None
        self._no_improvement = 0

        # Initialize number of samples seen since last reassignment
        self._n_since_last_reassign = 0

        sum_of_weights = np.sum(sample_weight)

        n_steps = (self.max_iter * n_samples) // self._batch_size
        normalized_sample_weight = sample_weight / sum_of_weights
        unit_sample_weight = np.ones_like(sample_weight, shape=(self._batch_size,))

        with _get_threadpool_controller().limit(limits=1, user_api="blas"):
            # Perform the iterative optimization until convergence
            for i in range(n_steps):
                # Sample a minibatch from the full dataset
                minibatch_indices = random_state.choice(
                    n_samples,
                    self._batch_size,
                    p=normalized_sample_weight,
                    replace=True,
                )
                # Perform the actual update step on the minibatch data
                # Note: since the sampling of the minibatch is sample_weight aware,
                # we pass fixed unit weights to the `_mini_batch_step` call to avoid
                # accounting for the weights twice. Also note that `_mini_batch_step`
                # can be called with non-unit weights when the caller constructs
                # the batches explicitly by calling the public `partial_fit` method