K-Means initialization using k-means||

    This is algorithm 2 in Scalable K-Means++ (2012).

def init_scalable(
    X, n_clusters, random_state=None, max_iter=None, oversampling_factor=2
):
    """K-Means initialization using k-means||

    This is algorithm 2 in Scalable K-Means++ (2012).
    """

    logger.info("Initializing with k-means||")
    # Step 1: Initialize Centers
    idx = 0
    centers = da.compute(X[idx, np.newaxis])[0]
    c_idx = {idx}

    # Step 2: Initialize cost
    cost, = compute(evaluate_cost(X, centers))

    if cost == 0:
        n_iter = 0
    else:
        n_iter = int(np.round(np.log(cost)))

    if max_iter is not None:
        n_iter = min(max_iter, n_iter)

    # Steps 3 - 6: update candidate Centers
    for i in range(n_iter):
        with _timer(
            "init iteration %2d/%2d , %2d centers" % (i + 1, n_iter, len(c_idx)),
            _logger=logger,
        ):
            new_idxs = _sample_points(X, centers, oversampling_factor, random_state)
            new_idxs = set(*compute(new_idxs))
            c_idx |= new_idxs

        # Sort before slicing, for better performance / memory
        # usage with the scheduler.
        # See https://github.com/dask/dask-ml/issues/39
        centers = X[sorted(c_idx)].compute()

    # XXX: scikit-learn doesn't have weighted k-means.
    # The paper weights each center by the number of points closest to it.
    # https://stackoverflow.com/a/37198799/1889400 claims you can scale the
    # features before clustering, but that doesn't seem right.
    # I think that replicating the *points*, proportional to the number of
    # original points closest to the candidate centers, would be a better way
    # to do that.

    if len(centers) < n_clusters:
        logger.warning("Found fewer than %d clusters in init.", n_clusters)
        # supplement with random
        need = n_clusters - len(centers)
        locs = sorted(
            random_state.choice(
                np.arange(0, len(X)), size=need, replace=False, chunks=len(X)
            )
        )
        extra = X[locs].compute()
        return np.vstack([centers, extra])
    else:
        # Step 7, 8 without weights
        # dask RandomState objects aren't valid for scikit-learn
        rng2 = (
            random_state.randint(0, np.iinfo("i4").max - 1, chunks=())
            .compute(scheduler="single-threaded")
            .item()
        )
        km = sk_k_means.KMeans(n_clusters, random_state=rng2)
        km.fit(centers)

    return km.cluster_centers_