2.7. KMeans Clustering

KMeans partitions samples into a fixed number of clusters. It alternates between 2 steps: it assigns each point to its nearest centroid, then it recomputes each centroid as the mean of its members. RustyML’s implementation is a parallel Lloyd’s algorithm with k-means++ seeding and best-of-n restarts. If you know scikit-learn’s sklearn.cluster.KMeans, most of that mental model applies here too. The one default that differs on purpose is n_init. See 2.7.1 for the reason.

2.7.1. How the estimator works

One call to fit runs the whole procedure n_init times. It keeps the run with the lowest inertia. Within a run, centroids start from k-means++ seeding. The first center is a uniformly random sample.

Each later center is drawn from the remaining points. The probability is proportional to the squared distance to the nearest center already chosen (the classic D^2 roulette wheel). This spreads the initial centers out, and it gives Lloyd’s iteration a better starting point than uniform random seeding. If every candidate’s squared distance is zero, the implementation picks that center uniformly at random instead. This case happens only when all points duplicate the already-chosen centers.

After seeding, each iteration does 2 things. The assignment step assigns every point to the nearest centroid. The update step replaces each centroid with the mean of its assigned points.

Iteration stops when the centroids stop moving. The convergence test compares the summed squared shift of all centroids against a variance-scaled tolerance. The threshold equals the mean per-feature population variance of the data, multiplied by tol. So tol is a relative tolerance, not an absolute distance. This matches the convention scikit-learn uses. A converged solution is a fixed point: every centroid equals the mean of its assigned points.

A fit that runs out of budget, instead of converging, gets one extra assignment pass before it returns. The Lloyd loop labels points against the current centroids, then installs the updated centroids. So stopping at max_iterations would leave labels and inertia describing the old centroids, while get_centroids reports the new ones. predict(x) would then disagree with get_labels(). The final pass re-assigns points against the centroids the model actually stores. scikit-learn re-runs its final E-step for the same reason.

The one deliberate difference from scikit-learn is the restart count. n_init defaults to 10 here. scikit-learn’s n_init='auto' is 1 for k-means++. This is not an oversight. scikit-learn’s default rests on its greedy k-means++, which draws 2 + ln(k) candidates per center and keeps the best. So one seeding already has low variance there.

RustyML uses plain k-means++: a single D^2 draw per center. Restarts exist to compensate for that higher-variance seeding. Pass with_n_init(1) for exact scikit-learn parity. Expect the fitted result to change when you do. Restarts derive their seeds deterministically from random_state, so a seeded fit stays reproducible.

2.7.2. Constructing a model

The constructor takes 3 positional parameters and validates them eagerly. An invalid configuration fails at construction, not at fit.

pub fn new(n_clusters: usize, max_iterations: usize, tolerance: f64) -> Result<Self, Error>

Any violation returns Error::InvalidParameter: n_clusters == 0, max_iterations == 0, or a tolerance that is zero, negative, NaN, or infinite. Note the asymmetry with the data-validation errors. An out-of-range hyperparameter returns InvalidParameter. A non-finite value in the data returns NonFinite instead. This distinction is a crate-wide convention.

KMeans::default() gives you n_clusters = 8, max_iterations = 300, tolerance = 1e-4, n_init = 10, and no seed. This configuration is always valid.

The 2 remaining settings are builder steps. Each one consumes the instance and returns it, so you can chain them after new. with_random_state(seed) cannot fail. with_n_init(n) returns Result, because 0 restarts triggers Error::InvalidParameter:

let mut km = KMeans::new(3, 300, 1e-4)
    .unwrap()
    .with_n_init(1)          // scikit-learn parity, returns Result
    .unwrap()
    .with_random_state(42);  // returns Self

Without a seed, k-means++ draws from entropy, and each fit produces a different partition. With a seed, the fit is reproducible. See 2.7.6 for details. This reproducibility covers every restart too. Each restart derives its own sub-seed deterministically from the seed you set.

2.7.3. Fitting, predicting, and reading results

3 methods drive the model. They differ in what they return and what they mutate.

fit(&mut self, data) trains the model in place and returns Result<&mut Self, Error>. It computes and stores the centroids, the training-set labels, the inertia, and the iteration count. All 4 values describe the winning restart.

predict(&self, data) assigns each row of a new matrix to its nearest fitted centroid. It returns an owned Array1<isize>. It borrows self immutably and does not change any stored state, so you can call it as many times as you want after one fit.

The labels use a signed type, not an unsigned one. This lets every clustering estimator in the crate share one label type. DBSCAN and Mean Shift both use -1 for noise, and the shared type feeds the metrics in 5.3. Clustering Metrics without a conversion. k-means itself never returns a negative label.

fit_predict(&mut self, data) calls fit, then returns a clone of the training labels. Use it when you want only the labels for the data you just trained on. Use fit plus predict instead when you want to score fresh points.

The fitted state is exposed through getters that return None before fit:

fit reports 3 data errors. It returns Error::EmptyInput when the matrix has zero rows. It returns Error::NonFinite when any value is NaN or infinite. It returns Error::InvalidInput when there are fewer samples than clusters, because you cannot place k centroids with fewer than k points.

predict adds 2 more errors. It returns Error::NotFitted when called before fit. It returns Error::DimensionMismatch when the feature count does not match the training data. It also runs the same empty-input and non-finite checks as fit.

2.7.4. Clustering three blobs end to end

3 tight, well-separated blobs make a standard sanity check. Any correct k=3 run must give every blob its own cluster. This example uses a small, deterministic dataset, with no random numbers in the data itself, and a fixed seed. It runs instantly, and the structure of the result is predictable.

use rustyml::machine_learning::KMeans;
use ndarray::{array, Array2};

fn main() {
    // 3 blobs of 5 points, centered at (0,0), (10,0), (5,10).
    let data: Array2<f64> = array![
        [-0.05,  0.03], [ 0.04, -0.02], [ 0.01,  0.05], [-0.03, -0.04], [ 0.02,  0.01],
        [ 9.95,  0.03], [10.04, -0.02], [10.01,  0.05], [ 9.97, -0.04], [10.02,  0.01],
        [ 4.95, 10.03], [ 5.04,  9.98], [ 5.01, 10.05], [ 4.97,  9.96], [ 5.02, 10.01],
    ];

    let mut km = KMeans::new(3, 300, 1e-4).unwrap().with_random_state(42);
    km.fit(&data).unwrap();

    let labels = km.get_labels().unwrap();
    let centroids = km.get_centroids().unwrap();
    println!("labels:    {:?}", labels);
    println!("centroids: {:?}", centroids);
    println!("inertia:   {:.6}", km.get_inertia().unwrap());
    println!("iters:     {}", km.get_actual_iterations().unwrap());

    // Score fresh points sitting on each blob's true center.
    let new_points = array![[0.0, 0.0], [10.0, 0.0], [5.0, 10.0]];
    let predicted = km.predict(&new_points).unwrap();
    println!("new-point labels: {:?}", predicted);
}

The exact cluster indices are arbitrary. k-means numbers clusters by discovery order, so the seed decides which blob becomes cluster 0. The structure of the result stays fixed:

labels:    the 5 points of each blob share one index. The 3 blobs get 3 distinct indices
           (some permutation of 0, 1, 2).
centroids: shape (3, 2). Rows sit within about 0.1 of (0,0), (10,0), (5,10), in some order.
inertia:   a small positive f64 (sum of squared point-to-centroid distances).
iters:     a handful, well below max_iterations.
new-point labels: each test point maps to the same cluster index as the blob it sits on.

Cluster indices depend on the permutation, so never compare labels across 2 separately fitted models by equality. Instead, use a permutation-invariant metric, such as the adjusted Rand index in 5.3. Clustering Metrics.

2.7.5. Choosing k

k-means cannot tell you k. You must supply it. 2 techniques narrow down a good value, and RustyML gives you the raw material for both.

The elbow method plots inertia against k. Inertia falls monotonically as k grows, because more centroids can only reduce the sum of squared distances. At k = n it reaches zero. Look for the elbow, the point where the drop flattens out. Read inertia directly from get_inertia.

The silhouette score is more decisive, because it has an interior optimum instead of a monotone trend. For each point, it measures how much closer that point sits to its own cluster than to the nearest other cluster. It averages these values to a score in [-1, 1], where higher is better. The metrics module provides it as silhouette_score. It takes the feature matrix, the labels, and a DistanceCalculationMetric. It is defined only for 2..=n-1 distinct clusters, so it cannot score k = 1.

use rustyml::machine_learning::KMeans;
use rustyml::metrics::silhouette_score;
use rustyml::math::DistanceCalculationMetric;
use ndarray::{array, Array2};

fn main() {
    let data: Array2<f64> = array![
        [-0.05,  0.03], [ 0.04, -0.02], [ 0.01,  0.05], [-0.03, -0.04], [ 0.02,  0.01],
        [ 9.95,  0.03], [10.04, -0.02], [10.01,  0.05], [ 9.97, -0.04], [10.02,  0.01],
        [ 4.95, 10.03], [ 5.04,  9.98], [ 5.01, 10.05], [ 4.97,  9.96], [ 5.02, 10.01],
    ];

    println!(" k   inertia   silhouette");
    for k in 1..=5usize {
        let mut km = KMeans::new(k, 300, 1e-4).unwrap().with_random_state(42);
        let labels = km.fit_predict(&data).unwrap();
        let inertia = km.get_inertia().unwrap();
        if k >= 2 {
            let s = silhouette_score(&data, &labels, DistanceCalculationMetric::Euclidean);
            println!("{k:2}  {inertia:8.4}   {s:7.4}");
        } else {
            println!("{k:2}  {inertia:8.4}      (n/a)");
        }
    }
}

On these 3 clean blobs, inertia drops sharply from k = 1 to k = 3, then flattens. The silhouette score peaks at k = 3. Both point at the true structure. Real data is messier, so the silhouette’s interior maximum tends to give a clearer signal. See 5.3. Clustering Metrics for the full list of metrics, including Davies-Bouldin and Calinski-Harabasz, which give independent second opinions.

2.7.6. Reproducible clusters: local seeds and the global seed

k-means++ is randomized, so an unseeded model produces a different partition on every fit. There are 2 ways to pin it down, and they combine in a predictable way.

The local seed is set with with_random_state(seed). It is self-contained: it seeds that estimator’s initialization RNG directly. It ignores any global state, and it never affects other components.

Each of the n_init restarts derives its own sub-seed from it, through a fixed mixing step, instead of by advancing one shared RNG. So a restart’s seeding does not depend on how much randomness the earlier restarts consumed. 2 models built with the same local seed, and fitted on the same data, produce identical centroids, labels, and inertia. The iteration count matches too, down to the last bit.

The global seed, set with rustyml::random::set_global_seed(seed), fixes every unseeded randomized component built and fitted on the same thread afterward. This includes k-means, the neural-network initializers, train_test_split, and more, all from one call. This mirrors Keras’ global-seed behavior. An unseeded k-means model draws an independent sub-seed from the global stream at fit time. To reproduce a run, set the global seed again before you re-fit.

2 caveats matter here. The global seed is thread-local, so set it on the thread that fits your models. Unseeded components also consume the stream in fit order, so their reproducibility depends on that order. A with_random_state seed avoids both problems. This is why it is the right tool for pinning a single estimator.

use rustyml::machine_learning::KMeans;
use rustyml::random::{set_global_seed, clear_global_seed};
use ndarray::{array, Array2};

fn main() {
    let data: Array2<f64> = array![
        [0.0, 0.0], [0.1, 0.0], [0.0, 0.1],
        [10.0, 0.0], [10.1, 0.0], [10.0, 0.1],
        [5.0, 10.0], [5.1, 10.0], [5.0, 10.1],
    ];

    // Local seed: identical results, independent of any global state.
    let mut a = KMeans::new(3, 300, 1e-4).unwrap().with_random_state(42);
    let mut b = KMeans::new(3, 300, 1e-4).unwrap().with_random_state(42);
    a.fit(&data).unwrap();
    b.fit(&data).unwrap();
    assert_eq!(a.get_labels().unwrap(), b.get_labels().unwrap());

    // Global seed: reset it before each fit to reproduce an unseeded run.
    set_global_seed(7);
    let mut c = KMeans::new(3, 300, 1e-4).unwrap();
    c.fit(&data).unwrap();
    let labels_c = c.get_labels().unwrap().clone();

    set_global_seed(7);
    let mut d = KMeans::new(3, 300, 1e-4).unwrap();
    d.fit(&data).unwrap();
    assert_eq!(&labels_c, d.get_labels().unwrap());

    clear_global_seed();
}

This determinism holds because the parallel arithmetic is designed for reproducibility, not just for speed. See 2.7.8 for how. For the full seeding model across the crate, see 7.1. Reproducibility and Random Seeds.

2.7.7. Failure modes and gotchas

Empty clusters. When Lloyd’s assignment leaves a centroid with no points, RustyML does not drop the cluster. It also does not leave a stale centroid in place. Each iteration reseeds every empty cluster with the single point farthest from its own assigned centroid, the point that contributes most to inertia. This keeps the model at exactly n_clusters centroids, and moves it out of the degenerate state on the next pass. So get_centroids always returns n_clusters rows. On pathological data, such as heavy duplication or k close to n, the labels can still resolve to fewer than k distinct values.

Sensitivity to feature scale. k-means minimizes Euclidean distance. So a feature measured in thousands dominates one measured in fractions, and the clustering effectively ignores the small-scale feature. There is no built-in standardization. Standardize your data first, with the tools in 4.2. Standardization and Normalization, whenever your features sit on different scales. This is the most common reason a k-means result looks wrong.

Sensitivity to outliers. Centroids are plain means, so a few extreme points can drag them off the dense region of their cluster and inflate inertia. If your data has outliers, either trim them, or use a density-based method that treats them as noise. 2.8. DBSCAN labels outliers explicitly. 2.9. Mean Shift finds modes without a preset k.

Non-globular clusters. k-means partitions space into Voronoi cells around the centroids. So it can recover only roughly convex, similarly sized blobs. Concentric rings or elongated manifolds defeat it, regardless of k. That geometry is what DBSCAN and the manifold methods in 2.12. t-SNE are for.

Too few samples. Requesting more clusters than you have points returns Error::InvalidInput. It does not silently clamp k. If k depends on your data, check n_clusters <= data.nrows() before you fit.

2.7.8. Parallelism and performance

The assignment step is the cost center. It compares every point against every centroid, each iteration, so this is where the parallelism lives. Each iteration computes all point-to-centroid projections as one parallel matrix product (see 6.2. Matrix Multiplication). It then finds each point’s nearest centroid with a short arg-min scan.

The centroid update accumulates each point’s row into its cluster’s running sum. It does this as a deterministic blocked reduction (see 6.3. Parallel Reductions). The k-means++ seeding parallelizes its distance passes the same way.

The word “deterministic” matters here. Each of these stages parallelizes only above a calibrated work threshold. The reduction always sums in a fixed block order, not in thread-arrival order. This is why the parallel path gives the same answer as the serial path. It is also why a seeded fit is reproducible on one machine, regardless of how many threads rayon uses.

Floating-point addition is not associative. Summing in whatever order the threads happen to finish would leak the thread count into the result. Across different machines or build targets, tiny last-bit differences in the arithmetic backend are still possible. Within one machine, a fixed seed gives a fixed answer.

For small inputs, everything runs single-threaded. This avoids rayon’s coordination overhead on data that does not need it. 3 thresholds control the crossover. The f64 scan gate covers the arg-min pass. The f64 sum gate covers the centroid-sum accumulation. The f64 cheap-map gate covers the division of each centroid by its cluster size.

The scan and sum gates live in rustyml::tuning::reduction. The cheap-map gate lives in rustyml::tuning::elementwise. Move them if you are clustering unusual shapes and want to shift a threshold. Most users never touch them. 7.3. Performance Tuning and Parallelism covers these knobs.

Building with the show_progress feature draws a live inertia and iteration bar during fit. This helps on large datasets, and it costs nothing when the feature is off.

2.7.9. Saving and loading a fitted model

KMeans derives serde’s Serialize and Deserialize. The generated save_to_path and load_from_path methods persist the whole model, including the centroids, labels, hyperparameters, and metadata, to a compact postcard binary blob. The file extension does not matter, because the format is always binary postcard. A loaded model predicts identically to the original, down to the last bit of the centroids.

The serialized layout gained the n_init field when restarts were added to the crate. So a blob written by an older version will not load. Re-fit the model and save it again instead.

use rustyml::machine_learning::KMeans;
use ndarray::{array, Array2};
use std::fs::remove_file;

fn main() {
    let data: Array2<f64> = array![
        [0.0, 0.0], [0.1, 0.0], [0.0, 0.1],
        [10.0, 0.0], [10.1, 0.0], [10.0, 0.1],
        [5.0, 10.0], [5.1, 10.0], [5.0, 10.1],
    ];

    let mut km = KMeans::new(3, 300, 1e-4).unwrap().with_random_state(42);
    km.fit(&data).unwrap();
    km.save_to_path("kmeans_model.bin").unwrap();

    let loaded = KMeans::load_from_path("kmeans_model.bin").unwrap();
    let original = km.predict(&data).unwrap();
    let restored = loaded.predict(&data).unwrap();
    assert_eq!(original, restored);

    remove_file("kmeans_model.bin").unwrap();
}

I/O and deserialization problems surface as Error::Io. A missing file on load is the common case. See 7.2. Model Persistence in Depth for the persistence format, versioning concerns, and how it works with the rest of the crate.