minibatch-kmeans 0.1.1__tar.gz

minibatch_kmeans-0.1.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Miroojin Bakshi
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

minibatch_kmeans-0.1.1/PKG-INFO

@@ -0,0 +1,100 @@
+Metadata-Version: 2.4
+Name: minibatch-kmeans
+Version: 0.1.1
+Summary: Mini-batch k-means clustering implemented in PyTorch
+Author: Miroojin Bakshi
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/mb-14/minibatch-kmeans
+Project-URL: Documentation, https://mb-14.github.io/minibatch-kmeans/
+Project-URL: Repository, https://github.com/mb-14/minibatch-kmeans
+Project-URL: Issues, https://github.com/mb-14/minibatch-kmeans/issues
+Keywords: k-means,clustering,pytorch,minibatch,machine-learning
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: scikit-learn>=1.3
+Requires-Dist: torch>=2.0
+Requires-Dist: tqdm>=4.66
+Provides-Extra: dev
+Requires-Dist: pdoc>=14; extra == "dev"
+Provides-Extra: eval
+Requires-Dist: clustering-benchmarks>=1.1; extra == "eval"
+Requires-Dist: fastkmeans>=0.5.0; extra == "eval"
+Requires-Dist: flash-kmeans>=0.2.0; extra == "eval"
+Requires-Dist: matplotlib>=3.8; extra == "eval"
+Requires-Dist: nvitop>=1.4; extra == "eval"
+Dynamic: license-file
+
+# minibatch-kmeans
+
+This package implements **mini-batch k-means** following [Sculley (WWW 2010)](https://doi.org/10.1145/1772690.1772862) (*Web-Scale K-Means Clustering*). [scikit-learn](https://scikit-learn.org/stable/modules/generated/sklearn.cluster.MiniBatchKMeans.html)'s `MiniBatchKMeans` is based on the same algorithm but is **CPU-only**; this library is for PyTorch pipelines that want native tensors and **GPU acceleration**.
+
+## Requirements
+
+- Python **3.10+** (see `pyproject.toml`)
+- **PyTorch** 2.x (`torch>=2.0`; use [pytorch.org](https://pytorch.org/get-started/locally/) for the wheel that matches your platform / CUDA).
+
+## Installation
+
+```bash
+pip install minibatch-kmeans
+```
+
+## Quick start
+
+```python
+import torch
+from minibatch_kmeans import MiniBatchKMeans
+
+X = torch.randn(5000, 4, dtype=torch.float32)
+km = MiniBatchKMeans(n_clusters=8, dtype=torch.float32)
+km.fit(X, batch_size=256, max_iter=20)
+labels = km.predict(X)
+```
+
+## Incremental training
+
+```python
+km = MiniBatchKMeans(n_clusters=8, dtype=torch.float32)
+for batch in X.split(256):
+    km.partial_fit(batch)
+labels = km.predict(X)
+```
+
+Example notebook: [notebooks/partial_fit_demo.ipynb](notebooks/partial_fit_demo.ipynb)
+
+## Documentation
+
+API reference: **[mb-14.github.io/minibatch-kmeans](https://mb-14.github.io/minibatch-kmeans/)**
+
+## Benchmarks & evaluation
+
+See **[evaluation/README.md](evaluation/README.md)** for benchmarks and evaluation.
+
+## Development
+
+```bash
+pip install -e ".[dev]"
+python -m unittest discover -s tests -v
+```
+
+Generate local API docs (HTML in `pdoc_html/`):
+
+```bash
+bash scripts/gen_api_docs.sh
+```
+
+## Related projects
+
+- [AnswerDotAI/fastkmeans](https://github.com/AnswerDotAI/fastkmeans)
+- [svg-project/flash-kmeans](https://github.com/svg-project/flash-kmeans)
+- [Faiss wiki](https://github.com/facebookresearch/faiss/wiki)
+- [rapidsai/cuml](https://github.com/rapidsai/cuml)
+

minibatch_kmeans-0.1.1/README.md

@@ -0,0 +1,66 @@
+# minibatch-kmeans
+
+This package implements **mini-batch k-means** following [Sculley (WWW 2010)](https://doi.org/10.1145/1772690.1772862) (*Web-Scale K-Means Clustering*). [scikit-learn](https://scikit-learn.org/stable/modules/generated/sklearn.cluster.MiniBatchKMeans.html)'s `MiniBatchKMeans` is based on the same algorithm but is **CPU-only**; this library is for PyTorch pipelines that want native tensors and **GPU acceleration**.
+
+## Requirements
+
+- Python **3.10+** (see `pyproject.toml`)
+- **PyTorch** 2.x (`torch>=2.0`; use [pytorch.org](https://pytorch.org/get-started/locally/) for the wheel that matches your platform / CUDA).
+
+## Installation
+
+```bash
+pip install minibatch-kmeans
+```
+
+## Quick start
+
+```python
+import torch
+from minibatch_kmeans import MiniBatchKMeans
+
+X = torch.randn(5000, 4, dtype=torch.float32)
+km = MiniBatchKMeans(n_clusters=8, dtype=torch.float32)
+km.fit(X, batch_size=256, max_iter=20)
+labels = km.predict(X)
+```
+
+## Incremental training
+
+```python
+km = MiniBatchKMeans(n_clusters=8, dtype=torch.float32)
+for batch in X.split(256):
+    km.partial_fit(batch)
+labels = km.predict(X)
+```
+
+Example notebook: [notebooks/partial_fit_demo.ipynb](notebooks/partial_fit_demo.ipynb)
+
+## Documentation
+
+API reference: **[mb-14.github.io/minibatch-kmeans](https://mb-14.github.io/minibatch-kmeans/)**
+
+## Benchmarks & evaluation
+
+See **[evaluation/README.md](evaluation/README.md)** for benchmarks and evaluation.
+
+## Development
+
+```bash
+pip install -e ".[dev]"
+python -m unittest discover -s tests -v
+```
+
+Generate local API docs (HTML in `pdoc_html/`):
+
+```bash
+bash scripts/gen_api_docs.sh
+```
+
+## Related projects
+
+- [AnswerDotAI/fastkmeans](https://github.com/AnswerDotAI/fastkmeans)
+- [svg-project/flash-kmeans](https://github.com/svg-project/flash-kmeans)
+- [Faiss wiki](https://github.com/facebookresearch/faiss/wiki)
+- [rapidsai/cuml](https://github.com/rapidsai/cuml)
+

minibatch_kmeans-0.1.1/evaluation/__init__.py

@@ -0,0 +1,3 @@
+"""Clustering evaluation (clustbench quality) and throughput harness (``perf_test``)."""
+
+from __future__ import annotations

minibatch_kmeans-0.1.1/evaluation/__main__.py

@@ -0,0 +1,4 @@
+from evaluation.run import main
+
+if __name__ == "__main__":
+    main()

minibatch_kmeans-0.1.1/evaluation/data.py

@@ -0,0 +1,121 @@
+"""Synthetic clustering datasets for benchmarking."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+
+def train_val_split(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    *,
+    train_fraction: float,
+    seed: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+    """
+    Random row permutation, then first ``floor(train_fraction * n)`` rows for train.
+
+    When ``train_fraction >= 1.0``, returns ``(x, y, None, None)`` — fit and optional
+    quality metrics use the full matrix (no holdout).
+
+    Raises
+    ------
+    ValueError
+        If ``train_fraction`` is not in ``(0, 1]``, or if the implied train/val sizes are invalid.
+    """
+    if train_fraction > 1.0 or train_fraction <= 0.0:
+        raise ValueError("train_fraction must be in (0, 1].")
+    n = x.shape[0]
+    if y.shape[0] != n:
+        raise ValueError("x and y must have the same number of rows.")
+    if train_fraction >= 1.0:
+        return x, y, None, None
+
+    g = torch.Generator()
+    g.manual_seed(seed)
+    perm = torch.randperm(n, generator=g)
+    n_train = int(math.floor(train_fraction * n))
+    n_val = n - n_train
+    if n_train < 1 or n_val < 1:
+        raise ValueError(
+            f"train_fraction={train_fraction} with n={n} yields n_train={n_train}, "
+            f"n_val={n_val}; need both >= 1."
+        )
+    p = perm.to(device=x.device)
+    x_perm = x[p]
+    y_perm = y[p]
+    x_train = x_perm[:n_train]
+    y_train = y_perm[:n_train]
+    x_val = x_perm[n_train:]
+    y_val = y_perm[n_train:]
+    return x_train, y_train, x_val, y_val
+
+
+def gaussian_mixture(
+    n_samples: int,
+    n_clusters: int,
+    n_features: int,
+    seed: int,
+    center_scale: float = 10.0,
+    noise_scale: float = 5.0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Structured clusters: points drawn as cluster_center + Gaussian noise.
+
+    Returns
+    -------
+    x : torch.Tensor, shape (n_samples, n_features), float32
+    y : torch.Tensor, shape (n_samples,), int64 — cluster index per point
+    """
+    g = torch.Generator()
+    g.manual_seed(seed)
+    centers = torch.randn(n_clusters, n_features, generator=g, dtype=torch.float32) * center_scale
+    cluster_indices = torch.randint(
+        0, n_clusters, (n_samples,), generator=g, dtype=torch.int64
+    )
+    x = torch.empty((n_samples, n_features), dtype=torch.float32)
+    batch = 100_000
+    for i in range(0, n_samples, batch):
+        end = min(i + batch, n_samples)
+        cidx = cluster_indices[i:end]
+        noise = torch.randn(end - i, n_features, generator=g, dtype=torch.float32) * noise_scale
+        x[i:end] = centers[cidx] + noise
+    return x, cluster_indices
+
+
+def gaussian_mixture_imbalanced(
+    n_samples: int,
+    n_clusters: int,
+    n_features: int,
+    seed: int,
+    *,
+    center_scale: float = 10.0,
+    noise_scale: float = 5.0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Same generative structure as ``gaussian_mixture``, but cluster assignment is
+    sampled from a **skewed** categorical (softmax of ``linspace(0, 2, k)``) so
+    high-index clusters have heavier mass than uniform sampling.
+    """
+    g = torch.Generator()
+    g.manual_seed(seed)
+    centers = torch.randn(n_clusters, n_features, generator=g, dtype=torch.float32) * center_scale
+    logits = torch.linspace(0.0, 2.0, n_clusters, dtype=torch.float32)
+    probs = torch.softmax(logits, dim=0)
+    cluster_indices = torch.multinomial(
+        probs, n_samples, replacement=True, generator=g
+    )
+    x = torch.empty((n_samples, n_features), dtype=torch.float32)
+    batch = 100_000
+    for i in range(0, n_samples, batch):
+        end = min(i + batch, n_samples)
+        cidx = cluster_indices[i:end]
+        noise = torch.randn(end - i, n_features, generator=g, dtype=torch.float32) * noise_scale
+        x[i:end] = centers[cidx] + noise
+    return x, cluster_indices
+
+
+def to_torch(x: torch.Tensor, device: torch.device) -> torch.Tensor:
+    return x.to(device=device, dtype=torch.float32)

minibatch_kmeans-0.1.1/evaluation/datasets.py

@@ -0,0 +1,49 @@
+"""Load clustering benchmark datasets via ``clustering-benchmarks`` (clustbench)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import torch
+
+
+@dataclass(frozen=True)
+class LoadedBenchmark:
+    """Torch-ready view of a clustbench dataset."""
+
+    battery: str
+    dataset: str
+    x: torch.Tensor  # float32 CPU dense ``(n, d)``
+    labels: list[np.ndarray]  # reference partitions (numpy, int)
+    n_clusters: np.ndarray  # k per reference
+
+
+def load_clustbench_dataset(
+    battery: str,
+    dataset: str,
+    *,
+    path: str,
+) -> LoadedBenchmark:
+    """
+    Load benchmark data the same way as upstream ``clustbench.load_dataset``.
+
+    ``path`` is the directory that contains battery folders (``sipu/``, ``fcps/``, …)—
+    equivalent to ``data_path`` in
+    https://clustering-benchmarks.gagolewski.com/weave/clustbench-usage.html
+    """
+    import clustbench
+
+    b: Any = clustbench.load_dataset(battery, dataset, path=path)
+    data = np.asarray(b.data, dtype=np.float64)
+    x = torch.from_numpy(data).to(dtype=torch.float32)
+    labels = [np.asarray(lab, dtype=np.int64) for lab in b.labels]
+    n_clusters = np.asarray(b.n_clusters, dtype=np.int64)
+    return LoadedBenchmark(
+        battery=battery,
+        dataset=dataset,
+        x=x,
+        labels=labels,
+        n_clusters=n_clusters,
+    )

minibatch_kmeans-0.1.1/evaluation/heuristics.py

@@ -0,0 +1,23 @@
+"""Dataset-scaled iteration counts for evaluation runs."""
+
+from __future__ import annotations
+
+# --- Iteration count (inverse to sample size) ----------------------------
+
+# At ``n == REFERENCE_N``, unclamped ``raw == BASE_NITER``; result is then clamped to [NITER_MIN, NITER_MAX].
+# ``REFERENCE_N`` is the median ``n`` over :data:`evaluation.presets.BENCHMARK_TASKS` (real suite).
+REFERENCE_N = 100_000
+NITER_MIN = 10
+NITER_MAX = 50
+BASE_NITER = (NITER_MIN + NITER_MAX) // 2
+
+
+def scaled_niter(n: int) -> int:
+    """
+    Lloyd / minibatch iterations for ``n`` points: ``BASE_NITER * REFERENCE_N / n``, rounded and
+    clamped to ``[NITER_MIN, NITER_MAX]``. Larger ``n`` → fewer iterations.
+    """
+    if n < 1:
+        raise ValueError("n must be >= 1")
+    raw = BASE_NITER * REFERENCE_N / n
+    return int(max(NITER_MIN, min(NITER_MAX, round(raw))))

minibatch_kmeans-0.1.1/evaluation/metrics.py

@@ -0,0 +1,238 @@
+"""Clustering metrics for benchmark evaluation (WCSS, ARI/NMI, clustbench reference scoring)."""
+
+from __future__ import annotations
+
+from typing import Callable
+
+import numpy as np
+import torch
+from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score
+
+from minibatch_kmeans._utils import _as_device
+
+# Rows of X processed per chunk in WCSS / assignment (avoids O(n*k) and full float64 X).
+_DEFAULT_CHUNK_ROWS = 8192
+
+
+def _row_chunk(x: torch.Tensor, start: int, end: int) -> torch.Tensor:
+    return x[start:end]
+
+
+def wcss_torch(
+    x: torch.Tensor,
+    centroids: torch.Tensor,
+    *,
+    chunk_rows: int = _DEFAULT_CHUNK_ROWS,
+    eval_device: torch.device | str | None = None,
+) -> float:
+    """
+    Within-cluster sum of squares (full data, current centroids).
+
+    Assigns each row to nearest centroid in Euclidean distance, then sums squared distances.
+    Accumulates in float64 for stability. Processes ``x`` in row chunks so peak memory stays
+    O(chunk_rows * max(d, k)) instead of O(n * k).
+
+    If ``eval_device`` is set and differs from ``x.device``, ``centroids`` are placed on
+    ``eval_device`` once and each chunk ``x[start:end]`` is moved there before distance work
+    (so ``x`` can remain on CPU while evaluation runs on GPU without materializing all of ``x``).
+    """
+    if chunk_rows < 1:
+        raise ValueError("chunk_rows must be >= 1")
+    if eval_device is None:
+        dev = x.device
+    else:
+        dev = _as_device(eval_device)
+
+    n = x.shape[0]
+
+    if eval_device is None:
+        staged = False
+        c = centroids
+    else:
+        staged = x.device != dev
+        c = centroids.to(device=dev, dtype=torch.float32)
+    c64 = c.double()
+    total = torch.zeros((), dtype=torch.float64, device=dev)
+    for start in range(0, n, chunk_rows):
+        end = min(start + chunk_rows, n)
+        if staged:
+            x64 = _row_chunk(x, start, end).to(dev).double()
+        else:
+            x64 = _row_chunk(x, start, end).double()
+        dists = torch.cdist(x64, c64, p=2.0) ** 2
+        total = total + dists.min(dim=1).values.sum()
+    return float(total.item())
+
+
+def assign_labels_torch(
+    x: torch.Tensor,
+    centroids: torch.Tensor,
+    *,
+    chunk_rows: int = _DEFAULT_CHUNK_ROWS,
+    eval_device: torch.device | str | None = None,
+) -> torch.Tensor:
+    """Nearest-centroid labels (same distance definition as ``wcss_torch``), chunked over rows."""
+    if chunk_rows < 1:
+        raise ValueError("chunk_rows must be >= 1")
+    if eval_device is None:
+        dev = x.device
+    else:
+        dev = _as_device(eval_device)
+
+    n = x.shape[0]
+
+    if eval_device is None:
+        staged = False
+        c = centroids
+    else:
+        staged = x.device != dev
+        c = centroids.to(device=dev, dtype=torch.float32)
+    c64 = c.double()
+    parts = []
+    for start in range(0, n, chunk_rows):
+        end = min(start + chunk_rows, n)
+        if staged:
+            x64 = _row_chunk(x, start, end).to(dev).double()
+        else:
+            x64 = _row_chunk(x, start, end).double()
+        dists = torch.cdist(x64, c64, p=2.0) ** 2
+        parts.append(dists.argmin(dim=1).to(torch.int64))
+    return torch.cat(parts, dim=0)
+
+
+def wcss_and_assign_labels_torch(
+    x: torch.Tensor,
+    centroids: torch.Tensor,
+    *,
+    chunk_rows: int = _DEFAULT_CHUNK_ROWS,
+    eval_device: torch.device | str | None = None,
+) -> tuple[float, torch.Tensor]:
+    """
+    Same as calling ``wcss_torch`` then ``assign_labels_torch``, but one pass over ``x``
+    (one ``cdist`` per chunk instead of two).
+    """
+    if chunk_rows < 1:
+        raise ValueError("chunk_rows must be >= 1")
+    if eval_device is None:
+        dev = x.device
+    else:
+        dev = _as_device(eval_device)
+
+    n = x.shape[0]
+
+    if eval_device is None:
+        staged = False
+        c = centroids
+    else:
+        staged = x.device != dev
+        c = centroids.to(device=dev, dtype=torch.float32)
+    c64 = c.double()
+    total = torch.zeros((), dtype=torch.float64, device=dev)
+    parts = []
+    for start in range(0, n, chunk_rows):
+        end = min(start + chunk_rows, n)
+        if staged:
+            x64 = _row_chunk(x, start, end).to(dev).double()
+        else:
+            x64 = _row_chunk(x, start, end).double()
+        dists = torch.cdist(x64, c64, p=2.0) ** 2
+        total = total + dists.min(dim=1).values.sum()
+        parts.append(dists.argmin(dim=1).to(torch.int64))
+    pred = torch.cat(parts, dim=0)
+    return float(total.item()), pred
+
+
+def _labels_to_numpy_int64(y: torch.Tensor) -> np.ndarray:
+    return y.detach().cpu().numpy().astype(np.int64, copy=False).ravel()
+
+
+def ari_score(y_true: torch.Tensor, y_pred: torch.Tensor) -> float:
+    """Adjusted Rand index (``sklearn.metrics.adjusted_rand_score``)."""
+    if y_true.numel() != y_pred.numel():
+        raise ValueError("y_true and y_pred must have the same length.")
+    return float(
+        adjusted_rand_score(
+            _labels_to_numpy_int64(y_true),
+            _labels_to_numpy_int64(y_pred),
+        )
+    )
+
+
+def nmi_score(y_true: torch.Tensor, y_pred: torch.Tensor) -> float:
+    """Normalized mutual information (``average_method='arithmetic'``)."""
+    return float(
+        normalized_mutual_info_score(
+            _labels_to_numpy_int64(y_true),
+            _labels_to_numpy_int64(y_pred),
+            average_method="arithmetic",
+        )
+    )
+
+
+def _labels_for_clustbench_get_score(y_pred: np.ndarray, k: int) -> np.ndarray:
+    """
+    ``clustbench.get_score`` requires ``results[k]`` with values in ``1..k`` (see clustbench/score.py).
+    Our k-means backends use nearest-centroid indices ``0..k-1``.
+    """
+    y = np.asarray(y_pred, dtype=np.int64)
+    lo, hi = int(np.min(y)), int(np.max(y))
+    if lo >= 1 and hi <= k:
+        return y
+    if lo >= 0 and hi <= k - 1:
+        return y + 1
+    raise ValueError(
+        f"Cannot map partition to clustbench 1..k labels: min={lo}, max={hi}, k={k}"
+    )
+
+
+def _metric_ari_numpy(y_true: np.ndarray, y_pred: np.ndarray) -> float:
+    return float(adjusted_rand_score(y_true, y_pred))
+
+
+def _metric_nmi_numpy(y_true: np.ndarray, y_pred: np.ndarray) -> float:
+    return float(
+        normalized_mutual_info_score(y_true, y_pred, average_method="arithmetic")
+    )
+
+
+def clustbench_best_score(
+    reference_labels: list[np.ndarray],
+    y_pred: np.ndarray,
+    k: int,
+    *,
+    metric: Callable[[np.ndarray, np.ndarray], float] | None = None,
+) -> float:
+    """
+    Best similarity to any reference partition, using clustbench's noise handling.
+
+    ``y_pred`` must be a length-``n`` integer partition with ``k`` clusters.
+    ``results`` is ``{k: y_pred}`` as required by ``clustbench.get_score``.
+    """
+    import clustbench
+
+    if metric is None:
+        metric = _metric_ari_numpy
+    y_pred_cb = _labels_for_clustbench_get_score(y_pred, k)
+    return float(
+        clustbench.get_score(
+            reference_labels,
+            {int(k): y_pred_cb},
+            metric=metric,
+            compute_max=True,
+        )
+    )
+
+
+def clustbench_best_ari_and_nmi(
+    reference_labels: list[np.ndarray],
+    y_pred: np.ndarray,
+    k: int,
+) -> tuple[float, float]:
+    """ARI and NMI each maximized over references (same ``y_pred``; two separate ``get_score`` calls)."""
+    ari = clustbench_best_score(
+        reference_labels, y_pred, k, metric=_metric_ari_numpy
+    )
+    nmi = clustbench_best_score(
+        reference_labels, y_pred, k, metric=_metric_nmi_numpy
+    )
+    return ari, nmi