KMeans vectorized implementation updating cluster centroids. Numpy pro

Viewed 488

I'mn implementing kmeans in python. In a single iteration i have computed the center labels for each 150 points:

label = 
array([0, 1, 2, 3, 4, 5, 6, 7, 3, 1, 5, 7, 1, 2, 5, 5, 5, 0, 5, 4, 0, 4,
       6, 7, 7, 1, 7, 0, 0, 3, 3, 0, 5, 5, 1, 1, 0, 4, 3, 7, 0, 1, 3, 7,
       5, 1, 4, 3, 0, 7, 5, 5, 5, 5, 5, 5, 5, 3, 5, 5, 3, 5, 5, 5, 5, 5,
       5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
       5, 5, 5, 5, 5, 3, 5, 5, 5, 5, 1, 5, 5, 5, 5, 5, 5, 5, 3, 5, 5, 5,
       5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
       5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5], dtype=int64)

and initial 8 centers:

centers =
array([[5.1, 3.5, 1.4, 0.2],
       [4.9, 3. , 1.4, 0.2],
       [4.7, 3.2, 1.3, 0.2],
       [4.6, 3.1, 1.5, 0.2],
       [5. , 3.6, 1.4, 0.2],
       [5.4, 3.9, 1.7, 0.4],
       [4.6, 3.4, 1.4, 0.3],
       [5. , 3.4, 1.5, 0.2]])

X is the iris data X.shape=(150, 4):

X =
array(  [5.1, 3.8, 1.5, 0.3],
   [5.4, 3.4, 1.7, 0.2],
   [5.1, 3.7, 1.5, 0.4],
   [4.6, 3.6, 1. , 0.2],
   [5.1, 3.3, 1.7, 0.5],
   [4.8, 3.4, 1.9, 0.2],
   [5. , 3. , 1.6, 0.2],
   [5. , 3.4, 1.6, 0.4],
   [5.2, 3.5, 1.5, 0.2],
   [5.2, 3.4, 1.4, 0.2],
   [4.7, 3.2, 1.6, 0.2],
   [4.8, 3.1, 1.6, 0.2],
   [5.4, 3.4, 1.5, 0.4],
            ...

Now I want to based on current center labels to update centers. That means iterate unique values in label. Then extract all corresponding points in X, calculate center based on all extracted points. Finally update centers. For example, in first iteration, extract all elements in X with label 0. Then calculate center(mean of each dimension). Then update centers[0] to new center. And so on for label 1, 2...

That's one iteration in original kmeans algorithm. My problem is how can i write this step in a numpy vectorized way instead of looping over.

1 Answers

Update centers

You can use boolean array indexing and computation along an axis to only explicitly iterate over the clusters instead of over every single data point.

K = 8
for k in range(K):
    centers[k] = X[label==k].mean(axis=0)

Update labels

This can also be done by iterating over all clusters:

distances = np.empty(shape=(X.shape[0], K))
for k in range(K):
    distances[:, k] = np.sqrt(np.sum((X - centers[k])**2, axis=1))
labels = distances.argmin(axis=1)

But it can also be done without an explicit loop by exploiting that matrix multiplication is the pairwise dot product.

squared_distances = np.sum(centers**2, axis=1) + (np.sum(X**2, axis=1) - 2*centers @ X.T).T
squared_distances[np.isclose(squared_distances, 0)] = 0  # self-distance can become slightly negative with this method (floating point precision problem)
distances = np.sqrt(squared_distances)
labels = distances.argmin(axis=1)
Related