import numpy as np
from sklearn.datasets import load_iris
X = load_iris().data # Load dataset (150 flowers, 4 features)
K = 3 # We want 3 clusters (groups)
n, d = X.shape # n=samples, d=features
p = np.ones((n, K))/K # Start: each point has equal chance in each cluster
m = X[np.random.choice(n, K, 0)] # Pick random points as initial cluster centers
s = np.array([np.eye(d)]*K) # Start: clusters are assumed circular (identity matrix)
for _ in range(10): # Repeat 10 times to improve clusters
for k in range(K): # ----- E STEP (Expectation) -----
diff = X - m[k] # Distance of all points from cluster k mean
p[:,k] = np.exp(-.5*np.sum(
[email protected](s[k])*diff,1))/np.sqrt(np.linalg.det(s[k]))
p /= p.sum(1, keepdims=True) # Convert to real probability (sum of each row = 1)
for k in range(K): # ----- M STEP (Maximization) -----
w = p[:,k] # Weight = how much each point belongs to cluster k
m[k] = (w@X)/w.sum() # Update cluster center (mean)
s[k] = ((X-m[k]).T*(w/w.sum()))@(X-m[k]) # Update cluster shape/spread
print("Cluster labels:", p.argmax(1)) # Final group = cluster with highest probability
