K-means clustering post artifacts, in Python (alkzar.cl)

kmeans.py

"""K-means clustering post artifacts, in Python.

Regenerates the plots and animation embedded in
content/posts/2018-05-17-a-brief-post-about-k-means.md. The original 2018
post was written as an .Rmd (R/reticulate) that ran a pure-python k-means
implementation but rendered every plot through ggplot2 and stitched the
per-iteration PNGs into a GIF with an external tool. This script reproduces
the same idea end to end in Python: same simulated data (two bivariate
normal clusters, mu=(0,0)/(17,17)), same k-means algorithm, but a single
continuous animation instead of stills glued together.

Styled with the site's Flexoki palette (styles/main.css) so it fits the
blog's look, following the same approach as scripts/berkson_paradox.py.

Usage:
    python3 scripts/kmeans.py

Requires: numpy, matplotlib, pillow
Writes:   content/static/img/kmeans-post/initial_plot.png
          content/static/img/kmeans-post/cost_plot.png
          content/static/img/k-means_process.gif
"""

from pathlib import Path

import numpy as np
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter

# ── Flexoki palette (light theme, styles/main.css) ──────────────────────────
BG = "#FFFCF0"       # --bg
UI = "#E6E4D9"       # --ui (borders / grid)
TX = "#100F0F"       # --tx (primary text)
TX_2 = "#6F6E69"      # --tx-2 (secondary text / axis labels)
BLUE = "#205EA6"      # --syn-blue (cluster 0 / centroid 0)
RED = "#AF3029"       # --syn-red (cluster 1 / centroid 1)
ORANGE = "#BC5215"    # --syn-orange (true group B, initial_plot only)

ROOT = Path(__file__).resolve().parent.parent
IMG_DIR = ROOT / "content/static/img/kmeans-post"
GIF_PATH = ROOT / "content/static/img/k-means_process.gif"

CLUSTER_COLORS = [BLUE, RED]


def euclidean_norm(x, y):
    """Return the euclidean norm between x and y."""
    return np.sqrt(np.sum((x - y) ** 2))


def k_means_history(df, K, max_iter=50, seed=0):
    """Run k-means, recording (centroids, assignment, cost) at every step.

    Same three-step algorithm described in the post: initialize
    representatives, assign each point to its closest one, then move each
    representative to the mean of its assigned points. Returns the full
    history so the caller can animate the process, plus the converged
    (c, z, J) triple.
    """
    np.random.seed(seed)
    nrow = df.shape[0]
    z = df[np.random.choice(nrow, K), :]
    history = []
    J = []
    for iteration in range(max_iter):
        distances = np.array([[euclidean_norm(x, zi) for zi in z] for x in df])
        c = distances.argmin(axis=1)
        cost = (distances[np.arange(nrow), c] ** 2).mean()
        J.append(cost)
        z_new = np.array([df[c == k].mean(axis=0) for k in range(K)])
        history.append({"z_before": z.copy(), "c": c.copy(), "z_after": z_new.copy()})
        if np.allclose(z_new, z):
            break
        z = z_new
    return history, c, z, J


def style_axes(ax, title):
    ax.set_facecolor(BG)
    ax.set_xlabel("x1", color=TX_2)
    ax.set_ylabel("x2", color=TX_2)
    ax.tick_params(colors=TX_2, length=0)
    for spine in ax.spines.values():
        spine.set_color(UI)
    ax.set_title(title, color=TX, fontsize=13, pad=12)


def plot_initial_data(df, labels):
    fig, ax = plt.subplots(figsize=(6, 5.4), dpi=200)
    fig.patch.set_facecolor(BG)
    colors = np.where(labels == 0, BLUE, ORANGE)
    ax.scatter(df[:, 0], df[:, 1], s=8, alpha=0.55, linewidths=0, c=colors)
    style_axes(ax, "Simulated data from two bivariate\nnormal distributions")
    fig.subplots_adjust(left=0.12, right=0.96, top=0.84, bottom=0.11)
    IMG_DIR.mkdir(parents=True, exist_ok=True)
    out = IMG_DIR / "initial_plot.png"
    fig.savefig(out)
    plt.close(fig)
    print(f"wrote {out}")


def plot_cost(J):
    fig, ax = plt.subplots(figsize=(6, 4.2), dpi=200)
    fig.patch.set_facecolor(BG)
    iters = np.arange(1, len(J) + 1)
    ax.plot(iters, J, color=BLUE, linewidth=1.6, zorder=2)
    ax.scatter(iters, J, color=BLUE, s=35, zorder=3)
    ax.set_xticks(iters)
    style_axes(ax, "Cost function J per iteration")
    ax.set_xlabel("Iteration (i)", color=TX_2)
    ax.set_ylabel("J(i)", color=TX_2)
    fig.subplots_adjust(left=0.14, right=0.96, top=0.88, bottom=0.14)
    out = IMG_DIR / "cost_plot.png"
    fig.savefig(out)
    plt.close(fig)
    print(f"wrote {out}")


def build_schedule(history, move_frames=10, hold_frames=16, converge_hold=26):
    """Frame-by-frame (kind, step_idx, label, t) schedule for the animation.

    step_idx indexes into `history`. "assign" recolours points to the
    incoming cluster assignment while centroids stay put (a hard decision,
    so no fade needed there); "update" glides centroids from their old
    position to the new mean over `move_frames`, `t` going 0 -> 1.
    """
    frames = [("init", None, None, 1.0)] * hold_frames

    for i, step in enumerate(history):
        assign_label = f"Iteration {i + 1} — assign cluster"
        frames += [("assign", i, assign_label, 1.0)] * hold_frames

        update_label = f"Iteration {i + 1} — update centroids"
        for f in range(move_frames):
            frames.append(("update", i, update_label, (f + 1) / move_frames))
        frames += [("update", i, update_label, 1.0)] * hold_frames

    frames += [("converged", len(history) - 1, "Convergence!", 1.0)] * converge_hold
    return frames


def render_gif(df, history):
    plt.rcParams.update({"font.size": 12})
    fig, ax = plt.subplots(figsize=(6, 5), dpi=140)
    fig.patch.set_facecolor(BG)
    ax.set_facecolor(BG)

    scatter = ax.scatter(df[:, 0], df[:, 1], s=8, alpha=0.6, linewidths=0, color=TX_2)
    centroids = ax.scatter([], [], s=220, marker="o", edgecolors=TX, linewidths=1.6,
                            zorder=4)
    title = ax.set_title("", fontsize=14, pad=12, color=TX)

    ax.tick_params(colors=TX_2, length=0)
    for spine in ax.spines.values():
        spine.set_color(UI)
    ax.set_xlabel("x1", color=TX_2)
    ax.set_ylabel("x2", color=TX_2)
    fig.subplots_adjust(left=0.12, right=0.96, top=0.9, bottom=0.12)

    schedule = build_schedule(history)

    def lerp(a, b, t):
        return a + (b - a) * t

    def frame(i):
        entry = schedule[i]
        kind = entry[0]

        if kind == "init":
            z0 = history[0]["z_before"]
            scatter.set_color(TX_2)
            centroids.set_offsets(z0)
            centroids.set_color(CLUSTER_COLORS)
            title.set_text("Initialize centroids")
            return scatter, centroids, title

        step_idx, label, t = entry[1], entry[2], entry[3]
        step = history[step_idx]

        if kind == "assign":
            scatter.set_color(np.array(CLUSTER_COLORS)[step["c"]])
            centroids.set_offsets(step["z_before"])
        elif kind == "update":
            scatter.set_color(np.array(CLUSTER_COLORS)[step["c"]])
            pos = lerp(step["z_before"], step["z_after"], t)
            centroids.set_offsets(pos)
        else:  # converged
            scatter.set_color(np.array(CLUSTER_COLORS)[step["c"]])
            centroids.set_offsets(step["z_after"])

        centroids.set_color(CLUSTER_COLORS)
        title.set_text(label)
        return scatter, centroids, title

    anim = FuncAnimation(fig, frame, frames=len(schedule), blit=False)
    GIF_PATH.parent.mkdir(parents=True, exist_ok=True)
    anim.save(GIF_PATH, writer=PillowWriter(fps=20))
    plt.close(fig)
    print(f"wrote {GIF_PATH} ({GIF_PATH.stat().st_size / 1024:.0f} KB)")


def main():
    mu_a, cov_a = [0, 0], [[1, 0], [0, 50]]
    mu_b, cov_b = [17, 17], [[15, 0], [0, 12]]

    np.random.seed(0)
    x1_a, x2_a = np.random.multivariate_normal(mu_a, cov_a, 1500).T
    np.random.seed(0)
    x1_b, x2_b = np.random.multivariate_normal(mu_b, cov_b, 1500).T

    labels = np.concatenate([np.zeros(1500), np.ones(1500)])
    df = np.vstack([np.column_stack([x1_a, x2_a]), np.column_stack([x1_b, x2_b])])

    plot_initial_data(df, labels)

    history, c, z, J = k_means_history(df, K=2, seed=0)
    print("cost per iteration:", [round(j, 2) for j in J])
    plot_cost(J)
    render_gif(df, history)


if __name__ == "__main__":
    main()