diff --git a/experiments/basics/kmeans.py b/experiments/basics/kmeans.py new file mode 100644 index 0000000..607bf6e --- /dev/null +++ b/experiments/basics/kmeans.py @@ -0,0 +1,56 @@ +"""Simple K-Means clustering implementation using NumPy.""" + +import numpy as np + + +def assign_centers(X: np.ndarray, centers: np.ndarray): + """Assign each data point to the nearest cluster center.""" + sq_dists = np.sum((X[:, None, :] - centers[None]) ** 2, axis=-1) + labels = np.argmin(sq_dists, axis=-1) + return labels + + +def update_centers(labels: np.ndarray, centers: np.ndarray): + """Update cluster centers based on current assignments.""" + return np.array([X[labels == k].mean(axis=0) for k in range(len(centers))]) + + +def kmeans( + X: np.ndarray, num_clusters: int, max_iters: int, tolerance: float = 1e-10 +) -> np.ndarray: + """Run K-Means clustering algorithm.""" + # Randomly initialize cluster centers + centers = X[np.random.choice(a=np.arange(len(X)), replace=False, size=num_clusters)] + labels = assign_centers(X, centers) + + # Iterate to refine centers and labels + for k in range(max_iters): + new_centers = update_centers(labels, centers) + labels = assign_centers(X, new_centers) + + rel_change = np.sum((new_centers - centers) ** 2) / np.sum(centers**2) + centers = new_centers + if rel_change < tolerance: + print(f"Converged after {k} iterations.") + break + + return labels + + +if __name__ == "__main__": + import matplotlib.pyplot as plt + + X = np.concatenate( + [ + np.random.randn(1000, 2), + np.random.randn(1000, 2) * 0.5 + 2.0, + np.random.randn(1000, 2) * 0.5 + 4.0, + ], + axis=0, + ) + + labels = kmeans(X, num_clusters=3, max_iters=100) + + plt.figure() + plt.scatter(X[:, 0], X[:, 1], c=labels, cmap="tab10") + plt.show()