Demonstrating and visualizing cluster centroid updates by using the first four iterations in Python?

Question

I managed to create a dataset with 6 clusters and visualize it with the code below, now I would like to visualize demonstration of update of the cluster centroids in KMeans algorithm. This demonstration should include first four iterations by generating 2×2-axis figure

Here is my code:

# import statements
from sklearn.datasets import make_blobs
import numpy as np
import matplotlib.pyplot as plt
# create blobs
data = make_blobs(n_samples=200, n_features=6, centers=6, cluster_std=1.6, random_state=50)
# create np array for data points
points = data[0]
# create scatter plot
plt.scatter(data[0][:,0], data[0][:,1], c=data[1], cmap='jet',marker="+",label="Original Data")
plt.xlim(-15,15)
plt.ylim(-15,15)
plt.show()

How can I implement this algorithm? It has been asked before using R but, I would like to do it in python. Can you help me visualize the first 4 iterations?

Output should be like:

Answer 1

You can use the plt.scatter() and plt.subplots() to achieve this as follows:

import matplotlib.pyplot as plt
from sklearn.datasets import make_blobs
data = make_blobs(n_samples=200, n_features=8, 
                           centers=6, cluster_std=1.8,random_state=101)

fig, ax = plt.subplots(nrows=2, ncols=2,figsize=(10,10))

from sklearn.cluster import KMeans
c=d=0
for i in range(4):
    ax[c,d].title.set_text(f"{i+1} iteration points:")
    kmeans = KMeans(n_clusters=6,random_state=0,max_iter=i+1)
    kmeans.fit(data[0])
    centroids=kmeans.cluster_centers_
    ax[c,d].scatter(data[0][:,0],data[0][:,1],c=data[1],cmap='brg')
    ax[c,d].scatter(kmeans.cluster_centers_[:, 0], kmeans.cluster_centers_[:, 1], s=200, c='black')
    d+=1
    if d==2:
        c+=1
        d=0

This will produce:

Answer 2

You could either write a KMeans algorithm yourself such that you can perform each update step-by-step and easily incorporate the plotting of the points into your code, but could probably also use the KMeans implementation from scikit-learn. To stop the algorithm to fully converge you could limit max_iter to [1, 2, 3, 4] since you want to plot just the first four iterations, and you can extract the cluster centroids and labels from the cluster_centers_ and labels_ attributes respectively.