k-means using numpy

Question

This is k-means implementation using Python (numpy). I believe there is room for improvement when it comes to computing distances (given I'm using a list comprehension, maybe I could also pack it in a numpy operation) and to compute the centroids using label-wise means (which I think also may be packed in a numpy operation).

import numpy as np
def k_means(data, k=2, max_iter=100):
 """Assigns data points into clusters using the k-means algorithm.
 Parameters
 ----------
 data : ndarray
 A 2D array containing data points to be clustered.
 k : int, optional
 Number of clusters (default = 2).
 max_iter : int, optional
 Number of maximum iterations
 Returns
 -------
 labels : ndarray
 A 1D array of labels for their respective input data points.
 """
 # data_max/data_min : array containing column-wise maximum/minimum values
 data_max = np.max(data, axis=0)
 data_min = np.min(data, axis=0)
 n_samples = data.shape[0]
 n_features = data.shape[1]
 # labels : array containing labels for data points, randomly initialized
 labels = np.random.randint(low=0, high=k, size=n_samples)
 # centroids : 2D containing centroids for the k-means algorithm
 # randomly initialized s.t. data_min <= centroid < data_max
 centroids = np.random.uniform(low=0., high=1., size=(k, n_features))
 centroids = centroids * (data_max - data_min) + data_min
 # k-means algorithm
 for i in range(max_iter):
 # distances : between datapoints and centroids
 distances = np.array(
 [np.linalg.norm(data - c, axis=1) for c in centroids])
 # new_labels : computed by finding centroid with minimal distance
 new_labels = np.argmin(distances, axis=0)
 if (labels == new_labels).all():
 # labels unchanged
 labels = new_labels
 print('Labels unchanged ! Terminating k-means.')
 break
 else:
 # labels changed
 # difference : percentage of changed labels
 difference = np.mean(labels != new_labels)
 print('%4f%% labels changed' % (difference * 100))
 labels = new_labels
 for c in range(k):
 # computing centroids by taking the mean over associated data points
 centroids[c] = np.mean(data[labels == c], axis=0)
 return labels

Answer 1

The code looks good, I wouldn't suggest anything as needing to be changed. In line with your question about possible ways to use Numpy better, I'll offer a few things you could try:

• You can pass low/high bounding arrays into numpy.randint, like .randint(low=data_min, high=data_max, size=(k, n_dimensions), rather than adjust the dimension-specific bounds after the fact. This is only done once, so it doesn't matter w.r.t. performance, but it cleans the code up a little.

• You can pack the distance computations into purely numpy code, but I don't know if I'd recommend it. Generally, K is very small compared to the size of your data, so your list comprehension is extremely cheap and readable. The alternative numpy code would abuse broadcasting to trick numpy into doing the looping for you, and would look something like this:

  distances = np.linalg.norm(np.expand_dims(d, 2) - np.expand_dims(centroids.T, 0), axis=1)
  new_labels = distances.argmin(axis=1)
  

  I wouldn't suggest using that code, though, since its performance gains are likely minimal (you can check for your own use case) and it's much less readable than what you currently have.

• You're right to be concerned about the centroid computation, since your current code is effectively O(K * N), whereas there's clearly a solution that's just O(N). The challenge, however, is to do so without losing the simplicity/readability of your code combined with numpy's inherent speed. Here's an interesting discussion about techniques to do a similar problem. Of those approaches, my quick tests say that the np.bincount method is fastest, but it doesn't work with multi-dimensional "weights", which is essentially what we need. My only thought that might be faster than your current approach is to write the code simply, then use Numba to JIT compile it to get C-level performance. That code would look like:

  @numba.jit
  def mean_vecs_by_group(mat, groups, num_groups=None):
   sum_vecs = np.zeros(shape=(num_groups, mat.shape[1]))
   counts = np.zeros(shape=(num_groups,))
   for g, vec in zip(groups, mat):
   sum_vecs[g] += vec
   counts[g] += 1
   return sum_vecs / np.expand_dims(counts, -1)
  

  However, in my test runs, this performs 1.5-10x slower than your current numpy code. I don't know of any other way in Numpy to do this all in one iteration through the array (if someone else knows, please comment). I suspect that in most cases, your current code is both the most maintainable solution and has very competitive performance.