Repository: scipy/scipy
File: scipy/cluster/vq.py (Python)

Code before:
    ...                       [ 0.4,1.8],
    ...                       [ 1.0,1.0]])
    >>> book = np.array((features[0],features[2]))
    >>> _kmeans(features,book)
    (array([[ 1.7       ,  2.4       ],
           [ 0.73333333,  1.13333333]]), 0.40563916697728591)

    """
    xp = np if xp is None else xp
    code_book = guess
    diff = xp.inf
    prev_avg_dists = deque([diff], maxlen=2)

    np_obs = np.asarray(obs)
    while diff > thresh:
        # compute membership and distances between obs and code_book
        obs_code, distort = vq(obs, code_book, check_finite=False)
        prev_avg_dists.append(xp.mean(distort, axis=-1))
        # recalc code_book as centroids of associated obs
        obs_code = np.asarray(obs_code)
        code_book, has_members = _vq.update_cluster_means(np_obs, obs_code,
                                                          code_book.shape[0])
        code_book = code_book[has_members]
        code_book = xp.asarray(code_book)
        diff = xp.abs(prev_avg_dists[0] - prev_avg_dists[1])
    _,final_distortions = vq(obs, code_book, check_finite=False)
    final_distortions_avg = xp.mean(final_distortions, axis=-1)
    return code_book, final_distortions_avg


@xp_capabilities(cpu_only=True, jax_jit=False, allow_dask_compute=True)
@_transition_to_rng("seed")
def kmeans(obs, k_or_guess, iter=20, thresh=1e-5, check_finite=True,
           *, rng=None):
    """
    Performs k-means on a set of observation vectors forming k clusters.

    The k-means algorithm adjusts the classification of the observations
    into clusters and updates the cluster centroids until the position of
    the centroids is stable over successive iterations. In this
    implementation of the algorithm, the stability of the centroids is
    determined by comparing the absolute value of the change in the average
    Euclidean distance between the observations and their corresponding
    centroids against a threshold. This yields
    a code book mapping centroids to codes and vice versa.

    Parameters
    ----------
    obs : ndarray
       Each row of the M by N array is an observation vector. The
       columns are the features seen during each observation.

Review comment (suggestion):
```suggestion

    _, final_distortions = vq(obs, code_book, check_finite=False)
```

Code after:
    ...                       [ 0.4,1.8],
    ...                       [ 1.0,1.0]])
    >>> book = np.array((features[0],features[2]))
    >>> _kmeans(features,book)
    (array([[ 1.7       ,  2.4       ],
           [ 0.73333333,  1.13333333]]), 0.40563916697728591)

    """
    xp = np if xp is None else xp
    code_book = guess
    diff = xp.inf
    prev_avg_dists = deque([diff], maxlen=2)

    np_obs = np.asarray(obs)
    while diff > thresh:
        # compute membership and distances between obs and code_book
        obs_code, distort = vq(obs, code_book, check_finite=False)
        prev_avg_dists.append(xp.mean(distort, axis=-1))
        # recalc code_book as centroids of associated obs
        obs_code = np.asarray(obs_code)
        code_book, has_members = _vq.update_cluster_means(np_obs, obs_code,
                                                          code_book.shape[0])
        code_book = code_book[has_members]
        code_book = xp.asarray(code_book)
        diff = xp.abs(prev_avg_dists[0] - prev_avg_dists[1])

    _, final_distortions = vq(obs, code_book, check_finite=False)
    final_distortions_avg = xp.mean(final_distortions, axis=-1)
    return code_book, final_distortions_avg


@xp_capabilities(cpu_only=True, jax_jit=False, allow_dask_compute=True)
@_transition_to_rng("seed")
def kmeans(obs, k_or_guess, iter=20, thresh=1e-5, check_finite=True,
           *, rng=None):
    """
    Performs k-means on a set of observation vectors forming k clusters.

    The k-means algorithm adjusts the classification of the observations
    into clusters and updates the cluster centroids until the position of
    the centroids is stable over successive iterations. In this
    implementation of the algorithm, the stability of the centroids is
    determined by comparing the absolute value of the change in the average
    Euclidean distance between the observations and their corresponding
    centroids against a threshold. This yields
    a code book mapping centroids to codes and vice versa.

    Parameters
    ----------
    obs : ndarray
       Each row of the M by N array is an observation vector. The