Skip to content

Clustering

cluster() does not compute a 2-D layout first. That layout is an iterative SGD and takes minutes at a million points, and clustering does not require it.

cluster groups points directly, implemented here in numpy and scipy plus the native k-means kernel, with no external clustering library. It does not compute a 2-D layout first. That layout is an iterative SGD and takes minutes at a million points, and clustering does not need it.

from fastumap import cluster

labels = cluster(embeddings, method="kmeans", n_clusters=20)     # k-means++ / Lloyd
labels = cluster(embeddings, method="spectral", n_clusters=20)   # non-convex clusters
labels = cluster(embeddings, method="dbscan")                    # density: finds k, noise = -1
labels = cluster(embeddings, method="hdbscan")                   # density, no eps to pick

hdbscan is the hierarchical form of dbscan. It takes no eps, so clusters of different densities are found together. The root of the hierarchy is never a cluster, so a dataset containing a single group returns all -1.

kmeans and spectral take n_clusters and metric ("euclidean" or "cosine") and are deterministic given random_state. k-means uses k-means++ seeding and Lloyd iterations with a blocked assignment step, so the n×k distance matrix is never materialised. spectral reuses fastumap's normalised-eigenvector machinery and runs k-means on the result. For a visualisation, call umap_project separately, usually on a sample; clustering and visualisation are separate operations.

Further options:

  • return_centroids=True returns a ClusterResult(labels, centroids, inertia) instead of labels alone, which is what is needed to compare runs or choose k.
  • choose_k(x, k_min=2, k_max=10) selects n_clusters by the inertia elbow. It is a heuristic; it lands on the true k give or take one when the clusters are clear.
  • assign_clusters(x_new, centroids) labels new points against already-fitted centroids, without re-clustering.

Above 50,000 points the native k-means kernel (matrixmultiply SIMD GEMM with rayon over row blocks, compiled into the wheel) replaces the inner loop. It clusters 1,000,000 × 128 into 64 groups in about 10 s on a 16-core machine, measured on SIFT1M, against minutes for the numpy path.

Above 200,000 points it also switches to mini-batch k-means, which samples batch_size rows per iteration rather than sweeping all of them. Measured on SIFT1M inside a docker --cpus=2 container: 5.2 s against full Lloyd's 25.3 s, at 1.010× the inertia.

from fastumap import cluster, kmeans

labels = cluster(x, method="kmeans", n_clusters=64)   # mini-batch above 200k rows
labels, cent = kmeans(x, 64, batch_size=0)            # force full Lloyd at any size
labels, cent = kmeans(x, 64, batch_size=5000)         # force mini-batch, with an explicit batch

The switch is on n alone and never on the core count, so the same input and seed produce the same labels on any machine. One limitation: on perfectly separated clusters, mini-batch can leave two centroids inside one group, measured at 4.8× worse inertia on synthetic blobs 8σ apart. Its running-mean update cannot move a centroid back across an empty gap, where full Lloyd's recompute can. Real embeddings overlap and do not trigger this (MNIST measures 1.013× Lloyd), and batch_size=0 disables the switch.

Next