{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3",
   "language": "python"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "name": "clustering.ipynb"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# clustering \u2014 Python demo\n\nNumerical companion to the entry [clustering](https://dictionaryofml.org/terms/clustering.html) of the [Dictionary of Applied Machine Learning](https://dictionaryofml.org/): it recomputes what the entry states and prints one line per check.\n\nRequires NumPy and Matplotlib only, and uses fixed seeds, so the printed numbers reproduce exactly. Generated from [`pythondemos/clustering.py`](https://dictionaryofml.org/terms/clustering.py); CC BY 4.0."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# Notebook shim: the script resolves output paths relative to __file__,\n# which a notebook kernel does not define; everything lands in the\n# working directory instead.\nimport os\n__file__ = os.path.join(os.getcwd(), \"clustering.py\")\nos.makedirs(\"pythondemos\", exist_ok=True)"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "#!/usr/bin/env python3\n\"\"\"\nclustering.py \u2014 Foreground/background segmentation of the cow image via\nhard clustering (k-means) and soft clustering (Gaussian mixture model).\n\nPixels of ``assets/CowsAustria.jpg`` are used as a 3-dimensional dataset\nin RGB space. Two clusters are fit:\n\n    1. k-means  (hard assignment: each pixel \u2192 one cluster)\n    2. GMM/EM   (soft assignment: each pixel \u2192 posterior over clusters)\n\nThe script writes two PNGs into ``pythondemos/`` for quick preview and a\nsingle preview PDF showing original / hard / soft side by side.\n\nRun from the repo root:\n\n    python3 pythondemos/clustering.py\n\"\"\"\nfrom __future__ import annotations\n\nfrom pathlib import Path\n\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom matplotlib.image import imread\n\nREPO = Path(__file__).resolve().parent.parent\nIMG_PATH = REPO / \"assets\" / \"BergSee.jpg\"\nOUT_DIR = REPO / \"pythondemos\"\n\n# Export the PNG panels at 300 ppi for their printed size. Each panel is\n# placed in a 0.31\\textwidth minipage (\\textwidth = 390pt), so its printed\n# width is 0.31 * 390 / 72.27 in; at 300 ppi that is ~502 px. Saving the\n# full-resolution source (2816 px wide) embedded ~30x more pixels than the\n# page can show and dominated the book PDF size.\nTEXTWIDTH_PT = 390.0\nPT_PER_INCH = 72.27\nPANEL_FRACTION = 0.31\nTARGET_DPI = 300\nMAX_WIDTH_PX = round(PANEL_FRACTION * TEXTWIDTH_PT / PT_PER_INCH * TARGET_DPI)\n\nRNG = np.random.default_rng(0)\n\n\ndef save_png_300dpi(img01: np.ndarray, path: Path, max_width_px: int = MAX_WIDTH_PX) -> None:\n    \"\"\"Save an (H, W, 3) array in [0,1] as PNG, downscaled so its printed\n    width (MAX_WIDTH_PX) is TARGET_DPI ppi. Pillow is matplotlib's PNG\n    backend (imread/imsave already rely on it), so this adds no dependency.\"\"\"\n    from PIL import Image\n    rgb8 = (np.clip(img01, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8)\n    im = Image.fromarray(rgb8)\n    if im.width > max_width_px:\n        new_h = round(im.height * max_width_px / im.width)\n        im = im.resize((max_width_px, new_h), Image.LANCZOS)\n    im.save(path)\n\n\ndef _kmeanspp_init(X: np.ndarray, k: int) -> np.ndarray:\n    \"\"\"k-means++ seeding: spreads initial centroids across the dataset.\"\"\"\n    n = len(X)\n    centroids = [X[RNG.integers(n)]]\n    for _ in range(k - 1):\n        d2 = np.min(\n            ((X[:, None, :] - np.stack(centroids)[None, :, :]) ** 2).sum(-1),\n            axis=1,\n        )\n        probs = d2 / d2.sum()\n        centroids.append(X[RNG.choice(n, p=probs)])\n    return np.stack(centroids)\n\n\ndef kmeans(X: np.ndarray, k: int, n_iter: int = 20) -> tuple[np.ndarray, np.ndarray]:\n    \"\"\"Lloyd's algorithm with k-means++ seeding.\"\"\"\n    centroids = _kmeanspp_init(X, k).copy()\n    for _ in range(n_iter):\n        d2 = ((X[:, None, :] - centroids[None, :, :]) ** 2).sum(-1)\n        labels = d2.argmin(axis=1)\n        for c in range(k):\n            mask = labels == c\n            if mask.any():\n                centroids[c] = X[mask].mean(axis=0)\n    return centroids, labels\n\n\ndef gmm_em(X: np.ndarray, k: int, n_iter: int = 40) -> np.ndarray:\n    \"\"\"Fit a diagonal-covariance GMM via EM. Returns soft assignments (n,k).\"\"\"\n    n, d = X.shape\n    idx = RNG.choice(n, size=k, replace=False)\n    mu = X[idx].copy()\n    var = np.tile(X.var(axis=0) + 1e-3, (k, 1))\n    pi = np.full(k, 1.0 / k)\n    for _ in range(n_iter):\n        # E-step: log N(x | mu_c, diag(var_c))\n        log_p = np.empty((n, k))\n        for c in range(k):\n            diff = X - mu[c]\n            log_p[:, c] = (\n                np.log(pi[c] + 1e-12)\n                - 0.5 * np.sum(np.log(2 * np.pi * var[c]))\n                - 0.5 * np.sum(diff ** 2 / var[c], axis=1)\n            )\n        log_p -= log_p.max(axis=1, keepdims=True)\n        resp = np.exp(log_p)\n        resp /= resp.sum(axis=1, keepdims=True)\n        # M-step\n        nk = resp.sum(axis=0) + 1e-12\n        pi = nk / n\n        mu = (resp.T @ X) / nk[:, None]\n        for c in range(k):\n            diff = X - mu[c]\n            var[c] = (resp[:, c, None] * diff ** 2).sum(axis=0) / nk[c] + 1e-4\n    return resp\n\n\ndef main() -> None:\n    img = imread(IMG_PATH).astype(np.float32) / 255.0\n    H, W, _ = img.shape\n    # Downsample for speed; label every pixel at the end.\n    step = 4\n    small = img[::step, ::step].reshape(-1, 3)\n\n    K = 2\n\n    # ---- hard clustering ----\n    centroids, _ = kmeans(small, k=K, n_iter=25)\n    X = img.reshape(-1, 3)\n    d2 = ((X[:, None, :] - centroids[None, :, :]) ** 2).sum(-1)\n    hard_labels = d2.argmin(axis=1)\n    # Pick one cluster to highlight: the brightest centroid, which\n    # corresponds to the sky/mountain region.\n    target = int(np.argmax(centroids.sum(axis=1)))\n    # Hard: keep the original RGB for pixels assigned to the target\n    # cluster; black out the rest.\n    hard_mask = (hard_labels == target).reshape(H, W, 1).astype(np.float32)\n    hard_img = img * hard_mask\n\n    # ---- soft clustering ----\n    # Fit a GMM on the full image, then soften the posteriors with a\n    # temperature so that the visualization shows graded membership\n    # rather than a near-binary assignment.\n    # Fit the GMM, then inflate the per-component variances to widen\n    # each Gaussian and re-score with the posterior formula. Larger\n    # variances give softer (more uniform) posteriors, producing\n    # smoother transitions between clusters.\n    resp = gmm_em(X, k=K, n_iter=25)\n    # Recover component means and recompute posteriors with inflated\n    # diagonal covariance (same variance in each component for the\n    # visualization \u2014 we only need softer posteriors, not a better\n    # fit).\n    nk = resp.sum(axis=0) + 1e-12\n    mu = (resp.T @ X) / nk[:, None]\n    pi = nk / nk.sum()\n    var = np.full((K, 3), 0.25)  # wide, uniform covariance\n    log_p = np.empty((len(X), K))\n    for c in range(K):\n        diff = X - mu[c]\n        log_p[:, c] = (\n            np.log(pi[c] + 1e-12)\n            - 0.5 * np.sum(np.log(2 * np.pi * var[c]))\n            - 0.5 * np.sum(diff ** 2 / var[c], axis=1)\n        )\n    log_p -= log_p.max(axis=1, keepdims=True)\n    resp = np.exp(log_p)\n    resp /= resp.sum(axis=1, keepdims=True)\n    # Align the GMM component order with k-means centroids by matching\n    # GMM means to the k-means centroids (greedy match).\n    mu_full = np.stack([\n        (resp[:, c, None] * X).sum(axis=0)\n        / (resp[:, c].sum() + 1e-12)\n        for c in range(K)\n    ])\n    perm = []\n    used = set()\n    for c in range(K):\n        order = np.argsort(np.linalg.norm(mu_full - centroids[c], axis=1))\n        for j in order:\n            if j not in used:\n                perm.append(j)\n                used.add(j)\n                break\n    resp = resp[:, perm]\n    # Soft: modulate each pixel's original RGB by the posterior\n    # probability that it belongs to the target cluster.\n    soft_weight = resp[:, target].reshape(H, W, 1)\n    soft_img = img * soft_weight\n\n    save_png_300dpi(img, OUT_DIR / \"clustering_original.png\")\n    save_png_300dpi(hard_img, OUT_DIR / \"clustering_hard.png\")\n    save_png_300dpi(soft_img, OUT_DIR / \"clustering_soft.png\")\n\n    # ---- preview PDF ----\n    fig, axes = plt.subplots(1, 3, figsize=(10, 3.2))\n    axes[0].imshow(img)\n    axes[0].set_title(\"original\")\n    axes[1].imshow(np.clip(hard_img, 0, 1))\n    axes[1].set_title(f\"hard: k-means (k={K})\")\n    axes[2].imshow(np.clip(soft_img, 0, 1))\n    axes[2].set_title(f\"soft: GMM posterior (k={K})\")\n    for ax in axes:\n        ax.set_xticks([])\n        ax.set_yticks([])\n    fig.tight_layout()\n    fig.savefig(OUT_DIR / \"clustering.png\", dpi=110)\n\n    print(f\"wrote {OUT_DIR/'clustering.png'}\")\n    print(f\"wrote {OUT_DIR/'clustering_original.png'}\")\n    print(f\"wrote {OUT_DIR/'clustering_hard.png'}\")\n    print(f\"wrote {OUT_DIR/'clustering_soft.png'}\")\n\n\nif __name__ == \"__main__\":\n    main()\n"
  }
 ]
}