code/reproduce_section4.py

#!/usr/bin/env python3
"""Scaled reproduction of the synthetic GMM experiment in Section 4.

The script intentionally depends only on the Python standard library and NumPy.
PyTorch is imported only when ``--backend torch`` is requested.  Each trial uses
the same data and the same balanced random partition for Lloyd and Hartigan.

Examples
--------
Local smoke run::

    python reproduce_section4.py --profile smoke --backend numpy \
        --output-dir artifacts/section4-smoke

Single-GPU full run::

    python reproduce_section4.py --profile full --backend torch --device cuda \
        --output-dir artifacts/section4-full
"""

from __future__ import annotations

import argparse
import csv
import datetime as dt
import hashlib
import html
import json
import math
import os
import platform
import sys
from pathlib import Path
from typing import Any, Iterable, Sequence

import numpy as np


PROFILES = {
    # Deliberately small enough for a laptop while still crossing sigma^2 > n
    # and spanning two orders of magnitude in dimension.
    "smoke": {
        "dimensions": (20, 200, 2_000, 10_000),
        "noise_variances": (1.0, 100.0),
        "trials": 4,
    },
    # A scaled, substantive version of the official 20 x 20 x 100 grid.  This
    # retains the decisive regimes and the paper's largest dimension.
    "full": {
        "dimensions": (10, 30, 100, 300, 1_000, 3_000, 10_000),
        "noise_variances": (1.0, 10.0, 50.0, 100.0),
        "trials": 50,
    },
}

RAW_FIELDS = (
    "profile",
    "backend",
    "device",
    "dtype",
    "seed",
    "trial_seed",
    "trial",
    "k",
    "n",
    "dimension",
    "tau2",
    "sigma2",
    "high_noise_sigma2_gt_n",
    "init_nmi",
    "lloyd_nmi",
    "hartigan_nmi",
    "hartigan_minus_lloyd_nmi",
    "init_sse",
    "true_partition_sse",
    "lloyd_sse",
    "hartigan_sse",
    "lloyd_sse_over_true",
    "hartigan_sse_over_true",
    "init_lloyd_fixed",
    "init_hartigan_fixed",
    "lloyd_unchanged_from_init",
    "hartigan_unchanged_from_init",
    "lloyd_changed_fraction",
    "hartigan_changed_fraction",
    "lloyd_iterations",
    "hartigan_sweeps",
    "lloyd_converged",
    "hartigan_converged",
    "hartigan_fixed_is_lloyd_fixed",
)

SUMMARY_FIELDS = (
    "k",
    "n",
    "dimension",
    "tau2",
    "sigma2",
    "high_noise_sigma2_gt_n",
    "trials",
    "init_nmi_mean",
    "lloyd_nmi_mean",
    "lloyd_nmi_sd",
    "lloyd_nmi_se",
    "hartigan_nmi_mean",
    "hartigan_nmi_sd",
    "hartigan_nmi_se",
    "hartigan_minus_lloyd_nmi_mean",
    "init_lloyd_fixed_rate",
    "init_hartigan_fixed_rate",
    "lloyd_unchanged_rate",
    "hartigan_unchanged_rate",
    "lloyd_changed_fraction_mean",
    "hartigan_changed_fraction_mean",
    "lloyd_recovery_rate_nmi_ge_0_9",
    "hartigan_recovery_rate_nmi_ge_0_9",
    "lloyd_sse_over_true_mean",
    "hartigan_sse_over_true_mean",
    "lloyd_iterations_mean",
    "hartigan_sweeps_mean",
    "lloyd_converged_rate",
    "hartigan_converged_rate",
    "hartigan_fixed_is_lloyd_fixed_rate",
)


def _parse_int_list(value: str) -> tuple[int, ...]:
    try:
        parsed = tuple(int(piece.strip()) for piece in value.split(",") if piece.strip())
    except ValueError as exc:
        raise argparse.ArgumentTypeError("expected comma-separated integers") from exc
    if not parsed or any(item <= 0 for item in parsed):
        raise argparse.ArgumentTypeError("all dimensions must be positive")
    return parsed


def _parse_float_list(value: str) -> tuple[float, ...]:
    try:
        parsed = tuple(float(piece.strip()) for piece in value.split(",") if piece.strip())
    except ValueError as exc:
        raise argparse.ArgumentTypeError("expected comma-separated numbers") from exc
    if not parsed or any((not math.isfinite(item)) or item <= 0 for item in parsed):
        raise argparse.ArgumentTypeError("all noise variances must be finite and positive")
    return parsed


def _trial_seed(
    master_seed: int, k: int, n: int, dimension: int, sigma2: float, trial: int
) -> int:
    """Grid-order- and batch-size-independent seed for one paired trial."""
    key = f"section4|{master_seed}|{k}|{n}|{dimension}|{sigma2:.17g}|{trial}"
    digest = hashlib.blake2b(key.encode("ascii"), digest_size=8).digest()
    return int.from_bytes(digest, "little", signed=False)


def _make_trial(
    *,
    master_seed: int,
    trial: int,
    k: int,
    n: int,
    dimension: int,
    tau2: float,
    sigma2: float,
    dtype: np.dtype[Any],
) -> tuple[np.ndarray, np.ndarray, np.ndarray, int]:
    seed = _trial_seed(master_seed, k, n, dimension, sigma2, trial)
    rng = np.random.default_rng(seed)
    true_labels = np.repeat(np.arange(k, dtype=np.int64), n // k)
    centers = rng.standard_normal((k, dimension), dtype=dtype)
    centers *= np.asarray(math.sqrt(tau2), dtype=dtype)
    noise = rng.standard_normal((n, dimension), dtype=dtype)
    noise *= np.asarray(math.sqrt(sigma2), dtype=dtype)
    data = centers[true_labels] + noise
    # This is the official random-partition construction: permute the balanced
    # true-label multiset.  Both algorithms receive these exact labels.
    initial_labels = rng.permutation(true_labels)
    return data, true_labels, initial_labels, seed


def _make_batch(
    trial_indices: Sequence[int],
    *,
    master_seed: int,
    k: int,
    n: int,
    dimension: int,
    tau2: float,
    sigma2: float,
    dtype: np.dtype[Any],
) -> tuple[np.ndarray, np.ndarray, np.ndarray, list[int]]:
    data = np.empty((len(trial_indices), n, dimension), dtype=dtype)
    true_labels = np.empty((len(trial_indices), n), dtype=np.int64)
    initial_labels = np.empty_like(true_labels)
    seeds: list[int] = []
    for row, trial in enumerate(trial_indices):
        x, y, initial, seed = _make_trial(
            master_seed=master_seed,
            trial=trial,
            k=k,
            n=n,
            dimension=dimension,
            tau2=tau2,
            sigma2=sigma2,
            dtype=dtype,
        )
        data[row] = x
        true_labels[row] = y
        initial_labels[row] = initial
        seeds.append(seed)
    expected = n // k
    counts = np.stack([(initial_labels == cluster).sum(axis=1) for cluster in range(k)], axis=1)
    assert np.all(counts == expected), "initial random partitions must remain balanced"
    return data, true_labels, initial_labels, seeds


def normalized_mutual_information(labels_a: np.ndarray, labels_b: np.ndarray) -> float:
    """Arithmetic-mean normalized mutual information, without scikit-learn."""
    a = np.asarray(labels_a).reshape(-1)
    b = np.asarray(labels_b).reshape(-1)
    if a.size != b.size or a.size == 0:
        raise ValueError("NMI inputs must have equal non-zero length")
    _, a_inverse = np.unique(a, return_inverse=True)
    _, b_inverse = np.unique(b, return_inverse=True)
    contingency = np.zeros(
        (int(a_inverse.max()) + 1, int(b_inverse.max()) + 1), dtype=np.float64
    )
    np.add.at(contingency, (a_inverse, b_inverse), 1.0)
    contingency /= float(a.size)
    p_a = contingency.sum(axis=1)
    p_b = contingency.sum(axis=0)
    nz_i, nz_j = np.nonzero(contingency)
    cells = contingency[nz_i, nz_j]
    mutual_information = float(
        np.sum(cells * np.log(cells / (p_a[nz_i] * p_b[nz_j])))
    )
    h_a = float(-np.sum(p_a[p_a > 0] * np.log(p_a[p_a > 0])))
    h_b = float(-np.sum(p_b[p_b > 0] * np.log(p_b[p_b > 0])))
    denominator = h_a + h_b
    if denominator == 0.0:
        return 1.0
    value = 2.0 * mutual_information / denominator
    # Roundoff can produce values a few ulps outside the mathematical range.
    return float(min(1.0, max(0.0, value)))


# ---------------------------------------------------------------------------
# NumPy implementation


def _centers_numpy(
    data: np.ndarray, labels: np.ndarray, k: int
) -> tuple[np.ndarray, np.ndarray]:
    batch, _, dimension = data.shape
    centers = np.zeros((batch, k, dimension), dtype=data.dtype)
    counts = np.zeros((batch, k), dtype=np.int64)
    for cluster in range(k):
        mask = labels == cluster
        counts[:, cluster] = mask.sum(axis=1)
        centers[:, cluster] = (data * mask[:, :, None]).sum(axis=1)
        centers[:, cluster] /= np.maximum(counts[:, cluster], 1)[:, None]
    return centers, counts


def _assign_numpy(data: np.ndarray, centers: np.ndarray) -> np.ndarray:
    squared_distances = np.sum(
        (data[:, :, None, :] - centers[:, None, :, :]) ** 2, axis=-1
    )
    return np.argmin(squared_distances, axis=2).astype(np.int64, copy=False)


def _loss_numpy(data: np.ndarray, labels: np.ndarray, k: int) -> np.ndarray:
    centers, _ = _centers_numpy(data, labels, k)
    batch_indices = np.arange(data.shape[0])[:, None]
    residual = data - centers[batch_indices, labels]
    return np.sum(residual * residual, axis=(1, 2), dtype=np.float64)


def _lloyd_numpy(
    data: np.ndarray, initial_labels: np.ndarray, k: int, max_sweeps: int
) -> dict[str, np.ndarray]:
    labels = initial_labels.copy()
    active = np.ones(data.shape[0], dtype=bool)
    iterations = np.zeros(data.shape[0], dtype=np.int64)
    initial_fixed = np.zeros(data.shape[0], dtype=bool)
    for sweep in range(max_sweeps):
        centers, _ = _centers_numpy(data, labels, k)
        proposed = _assign_numpy(data, centers)
        proposed[~active] = labels[~active]
        changed = np.any(proposed != labels, axis=1)
        if sweep == 0:
            initial_fixed = ~changed
        iterations[active] += 1
        labels[active] = proposed[active]
        active &= changed
        if not np.any(active):
            break
    centers, _ = _centers_numpy(data, labels, k)
    is_fixed = np.all(_assign_numpy(data, centers) == labels, axis=1)
    return {
        "labels": labels,
        "iterations": iterations,
        "converged": ~active,
        "initial_fixed": initial_fixed,
        "unchanged": np.all(labels == initial_labels, axis=1),
        "changed_fraction": np.mean(labels != initial_labels, axis=1),
        "sse": _loss_numpy(data, labels, k),
        "is_fixed": is_fixed,
    }


def _hartigan_numpy(
    data: np.ndarray, initial_labels: np.ndarray, k: int, max_sweeps: int
) -> dict[str, np.ndarray]:
    """Hartigan relocation with immediate center/count updates in sample order."""
    batch, n, _ = data.shape
    labels = initial_labels.copy()
    active = np.ones(batch, dtype=bool)
    iterations = np.zeros(batch, dtype=np.int64)
    initial_fixed = np.zeros(batch, dtype=bool)
    rows = np.arange(batch)

    for sweep in range(max_sweeps):
        before = labels.copy()
        active_at_start = active.copy()
        # Recomputing once per sweep prevents accumulated floating-point drift;
        # every accepted relocation below still updates means immediately.
        centers, counts = _centers_numpy(data, labels, k)
        for sample in range(n):
            old = labels[:, sample].copy()
            point = data[:, sample]
            squared_distances = np.sum(
                (point[:, None, :] - centers) ** 2, axis=2
            )
            count_float = counts.astype(data.dtype, copy=False)
            scores = squared_distances * count_float / (count_float + 1.0)
            old_count = counts[rows, old]
            source_distance = squared_distances[rows, old]
            source_score = np.where(
                old_count <= 1,
                -1.0,
                source_distance * old_count / np.maximum(old_count - 1, 1),
            )
            scores[rows, old] = source_score
            new = np.argmin(scores, axis=1).astype(np.int64, copy=False)
            new[~active_at_start] = old[~active_at_start]
            moved = active_at_start & (new != old)

            if np.any(moved):
                old_center = centers[rows, old].copy()
                new_center = centers[rows, new].copy()
                old_count_float = old_count.astype(data.dtype, copy=False)
                new_count = counts[rows, new]
                new_count_float = new_count.astype(data.dtype, copy=False)
                updated_old = (
                    old_center * old_count_float[:, None] - point
                ) / np.maximum(old_count_float - 1.0, 1.0)[:, None]
                updated_new = (
                    new_center * new_count_float[:, None] + point
                ) / (new_count_float + 1.0)[:, None]
                moved_rows = rows[moved]
                centers[moved_rows, old[moved]] = updated_old[moved]
                centers[moved_rows, new[moved]] = updated_new[moved]
                counts[moved_rows, old[moved]] -= 1
                counts[moved_rows, new[moved]] += 1
            labels[:, sample] = new

        changed = np.any(labels != before, axis=1)
        if sweep == 0:
            initial_fixed = ~changed
        iterations[active_at_start] += 1
        active &= changed
        if not np.any(active):
            break

    centers, _ = _centers_numpy(data, labels, k)
    lloyd_fixed = np.all(_assign_numpy(data, centers) == labels, axis=1)
    return {
        "labels": labels,
        "iterations": iterations,
        "converged": ~active,
        "initial_fixed": initial_fixed,
        "unchanged": np.all(labels == initial_labels, axis=1),
        "changed_fraction": np.mean(labels != initial_labels, axis=1),
        "sse": _loss_numpy(data, labels, k),
        "lloyd_fixed": lloyd_fixed,
    }


# ---------------------------------------------------------------------------
# Optional PyTorch implementation.  No torch symbol is imported at module load.


def _centers_torch(data: Any, labels: Any, k: int, torch: Any) -> tuple[Any, Any]:
    sums = []
    counts = []
    for cluster in range(k):
        mask = labels == cluster
        count = mask.sum(dim=1)
        counts.append(count)
        sums.append((data * mask.unsqueeze(2)).sum(dim=1))
    count_tensor = torch.stack(counts, dim=1)
    centers = torch.stack(sums, dim=1)
    centers = centers / count_tensor.clamp_min(1).unsqueeze(2).to(data.dtype)
    return centers, count_tensor


def _assign_torch(data: Any, centers: Any, torch: Any) -> Any:
    squared_distances = ((data.unsqueeze(2) - centers.unsqueeze(1)) ** 2).sum(dim=3)
    return torch.argmin(squared_distances, dim=2)


def _loss_torch(data: Any, labels: Any, k: int, torch: Any) -> Any:
    centers, _ = _centers_torch(data, labels, k, torch)
    rows = torch.arange(data.shape[0], device=data.device).unsqueeze(1)
    residual = data - centers[rows, labels]
    # Summation remains on-device; conversion to float64 happens only for the
    # small vector of per-trial results.
    return (residual * residual).sum(dim=(1, 2)).to(torch.float64)


def _lloyd_torch(
    data: Any, initial_labels: Any, k: int, max_sweeps: int, torch: Any
) -> dict[str, Any]:
    labels = initial_labels.clone()
    active = torch.ones(data.shape[0], dtype=torch.bool, device=data.device)
    iterations = torch.zeros(data.shape[0], dtype=torch.int64, device=data.device)
    initial_fixed = torch.zeros_like(active)
    for sweep in range(max_sweeps):
        centers, _ = _centers_torch(data, labels, k, torch)
        proposed = _assign_torch(data, centers, torch)
        proposed = torch.where(active.unsqueeze(1), proposed, labels)
        changed = torch.any(proposed != labels, dim=1)
        if sweep == 0:
            initial_fixed = ~changed
        iterations += active.to(torch.int64)
        labels = torch.where(active.unsqueeze(1), proposed, labels)
        active = active & changed
        if not bool(torch.any(active).item()):
            break
    centers, _ = _centers_torch(data, labels, k, torch)
    is_fixed = torch.all(_assign_torch(data, centers, torch) == labels, dim=1)
    return {
        "labels": labels,
        "iterations": iterations,
        "converged": ~active,
        "initial_fixed": initial_fixed,
        "unchanged": torch.all(labels == initial_labels, dim=1),
        "changed_fraction": (labels != initial_labels).to(torch.float64).mean(dim=1),
        "sse": _loss_torch(data, labels, k, torch),
        "is_fixed": is_fixed,
    }


def _hartigan_torch(
    data: Any, initial_labels: Any, k: int, max_sweeps: int, torch: Any
) -> dict[str, Any]:
    batch, n, _ = data.shape
    labels = initial_labels.clone()
    active = torch.ones(batch, dtype=torch.bool, device=data.device)
    iterations = torch.zeros(batch, dtype=torch.int64, device=data.device)
    initial_fixed = torch.zeros_like(active)
    rows = torch.arange(batch, device=data.device)

    for sweep in range(max_sweeps):
        before = labels.clone()
        active_at_start = active.clone()
        centers, counts = _centers_torch(data, labels, k, torch)
        for sample in range(n):
            old = labels[:, sample].clone()
            point = data[:, sample]
            squared_distances = ((point.unsqueeze(1) - centers) ** 2).sum(dim=2)
            count_float = counts.to(data.dtype)
            scores = squared_distances * count_float / (count_float + 1.0)
            old_count = counts[rows, old]
            source_distance = squared_distances[rows, old]
            source_score = torch.where(
                old_count <= 1,
                torch.full_like(source_distance, -1.0),
                source_distance
                * old_count.to(data.dtype)
                / (old_count - 1).clamp_min(1).to(data.dtype),
            )
            scores[rows, old] = source_score
            new = torch.argmin(scores, dim=1)
            new = torch.where(active_at_start, new, old)
            moved = active_at_start & (new != old)

            # Empty advanced-index assignments are valid, so this stays fully
            # asynchronous on CUDA instead of synchronizing once per sample.
            old_center = centers[rows, old].clone()
            new_center = centers[rows, new].clone()
            old_count_float = old_count.to(data.dtype)
            new_count = counts[rows, new]
            new_count_float = new_count.to(data.dtype)
            updated_old = (
                old_center * old_count_float.unsqueeze(1) - point
            ) / (old_count_float - 1.0).clamp_min(1.0).unsqueeze(1)
            updated_new = (
                new_center * new_count_float.unsqueeze(1) + point
            ) / (new_count_float + 1.0).unsqueeze(1)
            moved_rows = rows[moved]
            centers[moved_rows, old[moved]] = updated_old[moved]
            centers[moved_rows, new[moved]] = updated_new[moved]
            counts[moved_rows, old[moved]] -= 1
            counts[moved_rows, new[moved]] += 1
            labels[:, sample] = new

        changed = torch.any(labels != before, dim=1)
        if sweep == 0:
            initial_fixed = ~changed
        iterations += active_at_start.to(torch.int64)
        active = active & changed
        if not bool(torch.any(active).item()):
            break

    centers, _ = _centers_torch(data, labels, k, torch)
    lloyd_fixed = torch.all(_assign_torch(data, centers, torch) == labels, dim=1)
    return {
        "labels": labels,
        "iterations": iterations,
        "converged": ~active,
        "initial_fixed": initial_fixed,
        "unchanged": torch.all(labels == initial_labels, dim=1),
        "changed_fraction": (labels != initial_labels).to(torch.float64).mean(dim=1),
        "sse": _loss_torch(data, labels, k, torch),
        "lloyd_fixed": lloyd_fixed,
    }


def _torch_result_to_numpy(result: dict[str, Any]) -> dict[str, np.ndarray]:
    return {key: value.detach().cpu().numpy() for key, value in result.items()}


def _mean(rows: Sequence[dict[str, Any]], field: str) -> float:
    return float(np.mean([float(row[field]) for row in rows], dtype=np.float64))


def _sample_sd(rows: Sequence[dict[str, Any]], field: str) -> float:
    values = np.asarray([float(row[field]) for row in rows], dtype=np.float64)
    return float(np.std(values, ddof=1)) if len(values) > 1 else 0.0


def _aggregate(raw_rows: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
    groups: dict[tuple[int, int, int, float, float], list[dict[str, Any]]] = {}
    for row in raw_rows:
        key = (
            int(row["k"]),
            int(row["n"]),
            int(row["dimension"]),
            float(row["tau2"]),
            float(row["sigma2"]),
        )
        groups.setdefault(key, []).append(row)

    output: list[dict[str, Any]] = []
    for key in sorted(groups, key=lambda item: (item[4], item[2], item[0])):
        rows = groups[key]
        k, n, dimension, tau2, sigma2 = key
        lloyd_sd = _sample_sd(rows, "lloyd_nmi")
        hartigan_sd = _sample_sd(rows, "hartigan_nmi")
        count = len(rows)
        output.append(
            {
                "k": k,
                "n": n,
                "dimension": dimension,
                "tau2": tau2,
                "sigma2": sigma2,
                "high_noise_sigma2_gt_n": int(sigma2 > n),
                "trials": count,
                "init_nmi_mean": _mean(rows, "init_nmi"),
                "lloyd_nmi_mean": _mean(rows, "lloyd_nmi"),
                "lloyd_nmi_sd": lloyd_sd,
                "lloyd_nmi_se": lloyd_sd / math.sqrt(count),
                "hartigan_nmi_mean": _mean(rows, "hartigan_nmi"),
                "hartigan_nmi_sd": hartigan_sd,
                "hartigan_nmi_se": hartigan_sd / math.sqrt(count),
                "hartigan_minus_lloyd_nmi_mean": _mean(
                    rows, "hartigan_minus_lloyd_nmi"
                ),
                "init_lloyd_fixed_rate": _mean(rows, "init_lloyd_fixed"),
                "init_hartigan_fixed_rate": _mean(rows, "init_hartigan_fixed"),
                "lloyd_unchanged_rate": _mean(rows, "lloyd_unchanged_from_init"),
                "hartigan_unchanged_rate": _mean(rows, "hartigan_unchanged_from_init"),
                "lloyd_changed_fraction_mean": _mean(rows, "lloyd_changed_fraction"),
                "hartigan_changed_fraction_mean": _mean(rows, "hartigan_changed_fraction"),
                "lloyd_recovery_rate_nmi_ge_0_9": float(
                    np.mean([float(row["lloyd_nmi"]) >= 0.9 for row in rows])
                ),
                "hartigan_recovery_rate_nmi_ge_0_9": float(
                    np.mean([float(row["hartigan_nmi"]) >= 0.9 for row in rows])
                ),
                "lloyd_sse_over_true_mean": _mean(rows, "lloyd_sse_over_true"),
                "hartigan_sse_over_true_mean": _mean(rows, "hartigan_sse_over_true"),
                "lloyd_iterations_mean": _mean(rows, "lloyd_iterations"),
                "hartigan_sweeps_mean": _mean(rows, "hartigan_sweeps"),
                "lloyd_converged_rate": _mean(rows, "lloyd_converged"),
                "hartigan_converged_rate": _mean(rows, "hartigan_converged"),
                "hartigan_fixed_is_lloyd_fixed_rate": _mean(
                    rows, "hartigan_fixed_is_lloyd_fixed"
                ),
            }
        )
    return output


def _write_csv(path: Path, rows: Sequence[dict[str, Any]], fields: Sequence[str]) -> None:
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(rows)


def _polyline(points: Iterable[tuple[float, float]], css_class: str) -> str:
    coords = " ".join(f"{x:.2f},{y:.2f}" for x, y in points)
    return f'<polyline class="{css_class}" points="{coords}" />'


def _make_svg(summary: Sequence[dict[str, Any]]) -> str:
    noises = sorted({float(row["sigma2"]) for row in summary})
    panel_width = 270
    left = 72
    width = left + panel_width * len(noises) + 25
    height = 650
    top_y0, top_y1 = 85.0, 300.0
    bottom_y0, bottom_y1 = 375.0, 590.0
    pieces = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        'role="img" aria-labelledby="fig-title fig-desc">',
        '<title id="fig-title">Paired Lloyd and Hartigan synthetic GMM results</title>',
        '<desc id="fig-desc">Mean NMI and initial fixed-point rate versus dimension, split by noise variance.</desc>',
        """<style>
        text { font-family: ui-sans-serif, system-ui, sans-serif; fill: #172033; }
        .title { font-size: 18px; font-weight: 700; }
        .subtitle { font-size: 12px; fill: #4b5563; }
        .panel-title { font-size: 13px; font-weight: 700; }
        .tick { font-size: 10px; fill: #596579; }
        .axis { stroke: #64748b; stroke-width: 1; }
        .grid { stroke: #dbe2ea; stroke-width: 1; }
        .lloyd { fill: none; stroke: #d55e00; stroke-width: 2.4; }
        .hartigan { fill: none; stroke: #0072b2; stroke-width: 2.4; }
        .init { fill: none; stroke: #6b7280; stroke-width: 1.8; stroke-dasharray: 5 4; }
        .lloyd-dot { fill: #d55e00; }
        .hartigan-dot { fill: #0072b2; }
        .init-dot { fill: #6b7280; }
        .error { stroke-width: 1.1; opacity: .65; }
        </style>""",
        f'<text x="{width / 2:.1f}" y="25" text-anchor="middle" class="title">Section 4 scaled GMM reproduction</text>',
        f'<text x="{width / 2:.1f}" y="45" text-anchor="middle" class="subtitle">paired balanced partitions; curves show trial means (NMI error bars: ±1 SE)</text>',
        '<line x1="82" y1="61" x2="108" y2="61" class="lloyd"/><text x="113" y="65" class="tick">Lloyd</text>',
        '<line x1="165" y1="61" x2="191" y2="61" class="hartigan"/><text x="196" y="65" class="tick">Hartigan</text>',
        '<line x1="263" y1="61" x2="289" y2="61" class="init"/><text x="294" y="65" class="tick">initial NMI</text>',
        '<text x="16" y="195" transform="rotate(-90 16 195)" class="panel-title">Normalized mutual information</text>',
        '<text x="16" y="500" transform="rotate(-90 16 500)" class="panel-title">Initial partition fixed rate</text>',
    ]

    for column, sigma2 in enumerate(noises):
        rows = sorted(
            (row for row in summary if float(row["sigma2"]) == sigma2),
            key=lambda row: int(row["dimension"]),
        )
        dimensions = [int(row["dimension"]) for row in rows]
        log_min = math.log10(min(dimensions))
        log_max = math.log10(max(dimensions))
        if log_max == log_min:
            log_max += 1.0
        x0 = left + column * panel_width + 24
        x1 = left + (column + 1) * panel_width - 20

        def x_coord(dimension: int) -> float:
            return x0 + (math.log10(dimension) - log_min) / (log_max - log_min) * (x1 - x0)

        def y_coord(value: float, y0: float, y1: float) -> float:
            return y1 - min(1.0, max(0.0, value)) * (y1 - y0)

        regime = "high noise" if sigma2 > int(rows[0]["n"]) else "control"
        pieces.append(
            f'<text x="{(x0 + x1) / 2:.1f}" y="79" text-anchor="middle" class="panel-title">σ²={sigma2:g} · {regime}</text>'
        )
        for y0, y1 in ((top_y0, top_y1), (bottom_y0, bottom_y1)):
            pieces.append(f'<line x1="{x0}" y1="{y1}" x2="{x1}" y2="{y1}" class="axis"/>')
            pieces.append(f'<line x1="{x0}" y1="{y0}" x2="{x0}" y2="{y1}" class="axis"/>')
            for tick in (0.0, 0.5, 1.0):
                y = y_coord(tick, y0, y1)
                pieces.append(f'<line x1="{x0}" y1="{y:.2f}" x2="{x1}" y2="{y:.2f}" class="grid"/>')
                if column == 0:
                    pieces.append(f'<text x="{x0 - 7}" y="{y + 3:.2f}" text-anchor="end" class="tick">{tick:g}</text>')
            for dimension in dimensions:
                x = x_coord(dimension)
                pieces.append(f'<line x1="{x:.2f}" y1="{y1}" x2="{x:.2f}" y2="{y1 + 4}" class="axis"/>')
                pieces.append(f'<text x="{x:.2f}" y="{y1 + 17}" text-anchor="middle" class="tick">{dimension:g}</text>')
        pieces.append(f'<text x="{(x0 + x1) / 2:.1f}" y="625" text-anchor="middle" class="subtitle">dimension d (log scale)</text>')

        top_series = (
            ("init_nmi_mean", "init", "init-dot", None),
            ("lloyd_nmi_mean", "lloyd", "lloyd-dot", "lloyd_nmi_se"),
            ("hartigan_nmi_mean", "hartigan", "hartigan-dot", "hartigan_nmi_se"),
        )
        for field, line_class, dot_class, se_field in top_series:
            points = [(x_coord(int(row["dimension"])), y_coord(float(row[field]), top_y0, top_y1)) for row in rows]
            pieces.append(_polyline(points, line_class))
            for (x, y), row in zip(points, rows):
                if se_field is not None:
                    upper = y_coord(float(row[field]) + float(row[se_field]), top_y0, top_y1)
                    lower = y_coord(float(row[field]) - float(row[se_field]), top_y0, top_y1)
                    pieces.append(f'<line x1="{x:.2f}" y1="{upper:.2f}" x2="{x:.2f}" y2="{lower:.2f}" class="error {line_class}"/>')
                pieces.append(f'<circle cx="{x:.2f}" cy="{y:.2f}" r="3" class="{dot_class}"/>')

        bottom_series = (
            ("init_lloyd_fixed_rate", "lloyd", "lloyd-dot"),
            ("init_hartigan_fixed_rate", "hartigan", "hartigan-dot"),
        )
        for field, line_class, dot_class in bottom_series:
            points = [(x_coord(int(row["dimension"])), y_coord(float(row[field]), bottom_y0, bottom_y1)) for row in rows]
            pieces.append(_polyline(points, line_class))
            for x, y in points:
                pieces.append(f'<circle cx="{x:.2f}" cy="{y:.2f}" r="3" class="{dot_class}"/>')

    pieces.append("</svg>")
    return "\n".join(pieces)


def _make_html(
    *,
    svg: str,
    summary: Sequence[dict[str, Any]],
    config: dict[str, Any],
    metadata: dict[str, Any],
    checks: dict[str, Any],
) -> str:
    table_rows = []
    for row in summary:
        table_rows.append(
            "<tr>"
            f'<td>{int(row["dimension"])}</td>'
            f'<td>{float(row["sigma2"]):g}</td>'
            f'<td>{"yes" if row["high_noise_sigma2_gt_n"] else "no"}</td>'
            f'<td>{float(row["lloyd_nmi_mean"]):.3f}</td>'
            f'<td>{float(row["hartigan_nmi_mean"]):.3f}</td>'
            f'<td>{float(row["init_lloyd_fixed_rate"]):.3f}</td>'
            f'<td>{float(row["init_hartigan_fixed_rate"]):.3f}</td>'
            f'<td>{float(row["hartigan_minus_lloyd_nmi_mean"]):+.3f}</td>'
            "</tr>"
        )
    config_text = html.escape(json.dumps(config, indent=2, sort_keys=True))
    metadata_text = html.escape(json.dumps(metadata, indent=2, sort_keys=True))
    checked = int(checks["hartigan_fixed_points_checked"])
    return f"""<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Section 4 k-means reproduction</title>
<style>
body {{ margin: 0; color: #172033; background: #f5f7fb; font: 15px/1.5 ui-sans-serif, system-ui, sans-serif; }}
main {{ max-width: 1250px; margin: 0 auto; padding: 28px; }}
.card {{ background: white; border: 1px solid #dce3ec; border-radius: 12px; padding: 20px; margin: 16px 0; box-shadow: 0 2px 8px #1720330c; }}
h1, h2 {{ line-height: 1.2; }}
.lede {{ color: #4b5563; max-width: 82ch; }}
.figure {{ overflow-x: auto; }}
.figure svg {{ min-width: 760px; width: 100%; height: auto; }}
table {{ width: 100%; border-collapse: collapse; font-variant-numeric: tabular-nums; }}
th, td {{ padding: 7px 9px; border-bottom: 1px solid #e4e9f0; text-align: right; }}
th {{ background: #f8fafc; position: sticky; top: 0; }}
th:first-child, td:first-child {{ text-align: left; }}
pre {{ overflow: auto; padding: 14px; background: #101827; color: #e5edf8; border-radius: 8px; font-size: 12px; }}
.ok {{ color: #087f5b; font-weight: 650; }}
</style>
</head>
<body><main>
<h1>Section 4: synthetic GMM divergence</h1>
<p class="lede">A dependency-light, paired reproduction of Lloyd versus Hartigan. For every trial, centers are sampled from N(0, τ²I), balanced observations receive Gaussian noise N(0, σ²I), and the two algorithms start from the exact same balanced random label partition. Lloyd updates all labels simultaneously; Hartigan scans samples in fixed order and immediately applies each improving relocation.</p>
<section class="card figure">{svg}</section>
<section class="card">
<h2>Aggregate results</h2>
<div style="overflow:auto"><table>
<thead><tr><th>d</th><th>σ²</th><th>σ²&gt;n</th><th>Lloyd NMI</th><th>Hartigan NMI</th><th>Lloyd init fixed</th><th>Hartigan init fixed</th><th>NMI gap</th></tr></thead>
<tbody>{''.join(table_rows)}</tbody>
</table></div>
</section>
<section class="card">
<h2>Invariant check</h2>
<p class="ok">Passed: all {checked} converged Hartigan endpoints checked here are also fixed under a simultaneous Lloyd reassignment.</p>
<p>NMI uses arithmetic-mean entropy normalization and is implemented directly from the contingency table; no scikit-learn or plotting library is used.</p>
</section>
<section class="card"><h2>Configuration</h2><pre>{config_text}</pre></section>
<section class="card"><h2>Backend metadata</h2><pre>{metadata_text}</pre></section>
</main></body></html>"""


def _run(args: argparse.Namespace) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    profile = PROFILES[args.profile]
    dimensions = tuple(args.dimensions or profile["dimensions"])
    noise_variances = tuple(args.noise_variances or profile["noise_variances"])
    trials = int(args.trials if args.trials is not None else profile["trials"])
    if args.k < 2:
        raise ValueError("--k must be at least 2")
    if args.n <= args.k or args.n % args.k != 0:
        raise ValueError("--n must be greater than and divisible by --k")
    if trials <= 0 or args.max_sweeps <= 0:
        raise ValueError("--trials and --max-sweeps must be positive")
    if args.tau2 <= 0 or not math.isfinite(args.tau2):
        raise ValueError("--tau2 must be finite and positive")
    if args.batch_size < 0:
        raise ValueError("--batch-size cannot be negative")
    np_dtype = np.dtype(args.dtype)

    torch = None
    device_string = "cpu"
    backend_version = np.__version__
    device_detail = platform.processor() or platform.machine()
    if args.backend == "torch":
        try:
            import torch as imported_torch
        except ImportError as exc:
            raise RuntimeError(
                "--backend torch requires PyTorch; use --backend numpy locally"
            ) from exc
        torch = imported_torch
        torch.use_deterministic_algorithms(True)
        if hasattr(torch.backends, "cudnn"):
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
        if args.device == "auto":
            selected_device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            selected_device = args.device
        if selected_device.startswith("cuda") and not torch.cuda.is_available():
            raise RuntimeError("CUDA was requested but torch.cuda.is_available() is false")
        device = torch.device(selected_device)
        device_string = str(device)
        backend_version = str(torch.__version__)
        if device.type == "cuda":
            device_detail = torch.cuda.get_device_name(device)
        else:
            device_detail = platform.processor() or platform.machine()
    else:
        device = None
        if args.device not in ("auto", "cpu"):
            raise ValueError("NumPy backend supports only --device auto or cpu")

    batch_size = args.batch_size or (10 if args.backend == "torch" else 4)
    metadata = {
        "backend": args.backend,
        "backend_version": backend_version,
        "device": device_string,
        "device_detail": device_detail,
        "dtype": args.dtype,
        "numpy_version": np.__version__,
        "python_version": platform.python_version(),
        "platform": platform.platform(),
        "deterministic_algorithms": True,
        "trial_rng": "numpy.random.Generator(PCG64), one BLAKE2b-derived seed per trial",
        "batch_size": batch_size,
    }

    raw_rows: list[dict[str, Any]] = []
    for sigma2 in noise_variances:
        for dimension in dimensions:
            setting_rows: list[dict[str, Any]] = []
            for batch_start in range(0, trials, batch_size):
                trial_indices = list(range(batch_start, min(trials, batch_start + batch_size)))
                data_np, true_np, initial_np, trial_seeds = _make_batch(
                    trial_indices,
                    master_seed=args.seed,
                    k=args.k,
                    n=args.n,
                    dimension=dimension,
                    tau2=args.tau2,
                    sigma2=sigma2,
                    dtype=np_dtype,
                )
                if args.backend == "numpy":
                    lloyd = _lloyd_numpy(data_np, initial_np, args.k, args.max_sweeps)
                    hartigan = _hartigan_numpy(data_np, initial_np, args.k, args.max_sweeps)
                    true_sse = _loss_numpy(data_np, true_np, args.k)
                    init_sse = _loss_numpy(data_np, initial_np, args.k)
                else:
                    assert torch is not None and device is not None
                    data_tensor = torch.as_tensor(data_np, device=device)
                    true_tensor = torch.as_tensor(true_np, device=device)
                    initial_tensor = torch.as_tensor(initial_np, device=device)
                    with torch.no_grad():
                        lloyd_t = _lloyd_torch(
                            data_tensor, initial_tensor, args.k, args.max_sweeps, torch
                        )
                        hartigan_t = _hartigan_torch(
                            data_tensor, initial_tensor, args.k, args.max_sweeps, torch
                        )
                        true_sse_t = _loss_torch(data_tensor, true_tensor, args.k, torch)
                        init_sse_t = _loss_torch(data_tensor, initial_tensor, args.k, torch)
                    lloyd = _torch_result_to_numpy(lloyd_t)
                    hartigan = _torch_result_to_numpy(hartigan_t)
                    true_sse = true_sse_t.detach().cpu().numpy()
                    init_sse = init_sse_t.detach().cpu().numpy()

                converged_hartigan = np.asarray(hartigan["converged"], dtype=bool)
                hartigan_lloyd_fixed = np.asarray(hartigan["lloyd_fixed"], dtype=bool)
                if not np.all(hartigan_lloyd_fixed[converged_hartigan]):
                    bad = np.flatnonzero(
                        converged_hartigan & ~hartigan_lloyd_fixed
                    ).tolist()
                    raise AssertionError(
                        "Hartigan fixed point was not Lloyd-fixed for local batch rows "
                        f"{bad} at d={dimension}, sigma2={sigma2}"
                    )

                for local_index, trial in enumerate(trial_indices):
                    true_labels = true_np[local_index]
                    initial_labels = initial_np[local_index]
                    lloyd_labels = np.asarray(lloyd["labels"][local_index])
                    hartigan_labels = np.asarray(hartigan["labels"][local_index])
                    init_nmi = normalized_mutual_information(true_labels, initial_labels)
                    lloyd_nmi = normalized_mutual_information(true_labels, lloyd_labels)
                    hartigan_nmi = normalized_mutual_information(true_labels, hartigan_labels)
                    denominator = max(float(true_sse[local_index]), np.finfo(float).tiny)
                    row = {
                        "profile": args.profile,
                        "backend": args.backend,
                        "device": device_string,
                        "dtype": args.dtype,
                        "seed": args.seed,
                        "trial_seed": trial_seeds[local_index],
                        "trial": trial,
                        "k": args.k,
                        "n": args.n,
                        "dimension": dimension,
                        "tau2": args.tau2,
                        "sigma2": sigma2,
                        "high_noise_sigma2_gt_n": int(sigma2 > args.n),
                        "init_nmi": init_nmi,
                        "lloyd_nmi": lloyd_nmi,
                        "hartigan_nmi": hartigan_nmi,
                        "hartigan_minus_lloyd_nmi": hartigan_nmi - lloyd_nmi,
                        "init_sse": float(init_sse[local_index]),
                        "true_partition_sse": float(true_sse[local_index]),
                        "lloyd_sse": float(lloyd["sse"][local_index]),
                        "hartigan_sse": float(hartigan["sse"][local_index]),
                        "lloyd_sse_over_true": float(lloyd["sse"][local_index]) / denominator,
                        "hartigan_sse_over_true": float(hartigan["sse"][local_index]) / denominator,
                        "init_lloyd_fixed": int(bool(lloyd["initial_fixed"][local_index])),
                        "init_hartigan_fixed": int(bool(hartigan["initial_fixed"][local_index])),
                        "lloyd_unchanged_from_init": int(bool(lloyd["unchanged"][local_index])),
                        "hartigan_unchanged_from_init": int(bool(hartigan["unchanged"][local_index])),
                        "lloyd_changed_fraction": float(lloyd["changed_fraction"][local_index]),
                        "hartigan_changed_fraction": float(hartigan["changed_fraction"][local_index]),
                        "lloyd_iterations": int(lloyd["iterations"][local_index]),
                        "hartigan_sweeps": int(hartigan["iterations"][local_index]),
                        "lloyd_converged": int(bool(lloyd["converged"][local_index])),
                        "hartigan_converged": int(bool(hartigan["converged"][local_index])),
                        "hartigan_fixed_is_lloyd_fixed": int(
                            bool(hartigan_lloyd_fixed[local_index])
                            if converged_hartigan[local_index]
                            else False
                        ),
                    }
                    raw_rows.append(row)
                    setting_rows.append(row)

            print(
                f"d={dimension:>5} sigma2={sigma2:>6g} "
                f"Lloyd NMI={_mean(setting_rows, 'lloyd_nmi'):.3f} "
                f"Hartigan NMI={_mean(setting_rows, 'hartigan_nmi'):.3f} "
                f"Lloyd-init-fixed={_mean(setting_rows, 'init_lloyd_fixed'):.2f}",
                file=sys.stderr,
                flush=True,
            )
    return raw_rows, metadata


def _build_checks(raw_rows: Sequence[dict[str, Any]]) -> dict[str, Any]:
    converged_hartigan = [row for row in raw_rows if row["hartigan_converged"]]
    violations = [
        row for row in converged_hartigan if not row["hartigan_fixed_is_lloyd_fixed"]
    ]
    return {
        "paired_trials": len(raw_rows),
        "hartigan_fixed_points_checked": len(converged_hartigan),
        "hartigan_fixed_but_not_lloyd_fixed": len(violations),
        "hartigan_implies_lloyd_fixed_assertion_passed": len(violations) == 0,
        "all_lloyd_runs_converged": all(bool(row["lloyd_converged"]) for row in raw_rows),
        "all_hartigan_runs_converged": all(bool(row["hartigan_converged"]) for row in raw_rows),
    }


def _high_noise_indicators(summary: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
    indicators = []
    high_noises = sorted(
        {float(row["sigma2"]) for row in summary if row["high_noise_sigma2_gt_n"]}
    )
    for sigma2 in high_noises:
        rows = sorted(
            (row for row in summary if float(row["sigma2"]) == sigma2),
            key=lambda row: int(row["dimension"]),
        )
        if not rows:
            continue
        first, last = rows[0], rows[-1]
        indicators.append(
            {
                "sigma2": sigma2,
                "dimension_low": int(first["dimension"]),
                "dimension_high": int(last["dimension"]),
                "lloyd_initial_fixed_rate_low_d": float(first["init_lloyd_fixed_rate"]),
                "lloyd_initial_fixed_rate_high_d": float(last["init_lloyd_fixed_rate"]),
                "hartigan_initial_fixed_rate_high_d": float(last["init_hartigan_fixed_rate"]),
                "lloyd_nmi_high_d": float(last["lloyd_nmi_mean"]),
                "hartigan_nmi_high_d": float(last["hartigan_nmi_mean"]),
                "hartigan_minus_lloyd_nmi_high_d": float(
                    last["hartigan_minus_lloyd_nmi_mean"]
                ),
            }
        )
    return indicators


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Reproduce the paired synthetic-GMM divergence in Section 4."
    )
    parser.add_argument("--profile", choices=tuple(PROFILES), default="smoke")
    parser.add_argument("--backend", choices=("numpy", "torch"), default="numpy")
    parser.add_argument(
        "--device",
        default="auto",
        help="auto, cpu, cuda, or a concrete torch device such as cuda:0",
    )
    parser.add_argument("--dtype", choices=("float32", "float64"), default="float32")
    parser.add_argument("--output-dir", type=Path, default=Path("artifacts/section4"))
    parser.add_argument("--seed", type=int, default=29_589)
    parser.add_argument("--k", type=int, default=2)
    parser.add_argument("--n", type=int, default=40)
    parser.add_argument("--tau2", type=float, default=1.0)
    parser.add_argument(
        "--dimensions",
        type=_parse_int_list,
        help="optional comma-separated override of the profile dimensions",
    )
    parser.add_argument(
        "--noise-variances",
        type=_parse_float_list,
        help="optional comma-separated override of the profile noise variances",
    )
    parser.add_argument("--trials", type=int, help="optional trial-count override")
    parser.add_argument(
        "--batch-size",
        type=int,
        default=0,
        help="0 chooses 4 for NumPy or 10 for torch; seeds are batch-size invariant",
    )
    parser.add_argument("--max-sweeps", type=int, default=100)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    args.output_dir.mkdir(parents=True, exist_ok=True)
    raw_rows, metadata = _run(args)
    summary = _aggregate(raw_rows)
    checks = _build_checks(raw_rows)

    profile = PROFILES[args.profile]
    config = {
        "profile": args.profile,
        "backend": args.backend,
        "device": metadata["device"],
        "dtype": args.dtype,
        "seed": args.seed,
        "k": args.k,
        "n": args.n,
        "samples_per_cluster": args.n // args.k,
        "tau2": args.tau2,
        "dimensions": list(args.dimensions or profile["dimensions"]),
        "noise_variances": list(args.noise_variances or profile["noise_variances"]),
        "trials_per_setting": args.trials if args.trials is not None else profile["trials"],
        "batch_size": metadata["batch_size"],
        "max_sweeps": args.max_sweeps,
        "initialization": "balanced random permutation of the true-label multiset",
        "pairing": "identical data and initial labels for Lloyd and Hartigan",
        "lloyd_update": "simultaneous exact label-to-mean reassignment",
        "hartigan_update": "fixed sample order, sequential immediate relocations",
        "nmi_normalization": "2 * MI / (H(true) + H(predicted))",
        "recovery_threshold": "NMI >= 0.9",
    }

    paths = {
        "trials_csv": args.output_dir / "section4_trials.csv",
        "summary_csv": args.output_dir / "section4_summary.csv",
        "results_json": args.output_dir / "section4_results.json",
        "figure_svg": args.output_dir / "section4_divergence.svg",
        "report_html": args.output_dir / "section4_report.html",
    }
    _write_csv(paths["trials_csv"], raw_rows, RAW_FIELDS)
    _write_csv(paths["summary_csv"], summary, SUMMARY_FIELDS)
    svg = _make_svg(summary)
    paths["figure_svg"].write_text(svg, encoding="utf-8")
    report = _make_html(
        svg=svg, summary=summary, config=config, metadata=metadata, checks=checks
    )
    paths["report_html"].write_text(report, encoding="utf-8")
    payload = {
        "schema_version": 1,
        "generated_at_utc": dt.datetime.now(dt.timezone.utc).isoformat(),
        "experiment": "ICML 2026 paper 29589, scaled Section 4 synthetic GMM",
        "config": config,
        "backend_metadata": metadata,
        "checks": checks,
        "high_noise_indicators": _high_noise_indicators(summary),
        "summary": summary,
        "trials": raw_rows,
        "files": {key: path.name for key, path in paths.items()},
    }
    paths["results_json"].write_text(
        json.dumps(payload, indent=2, sort_keys=True, allow_nan=False) + "\n",
        encoding="utf-8",
    )
    print(
        json.dumps(
            {
                "status": "ok",
                "output_dir": str(args.output_dir.resolve()),
                "paired_trials": len(raw_rows),
                "files": {key: str(path.resolve()) for key, path in paths.items()},
            },
            indent=2,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())