2.9. Mean Shift

Mean Shift is a clustering algorithm. Use it when you do not know how many clusters the data holds, and you do not want to guess a number. Unlike KMeans, which needs k up front, Mean Shift finds the number of clusters from the density of the data. It needs only 1 input from you: the bandwidth. The algorithm moves every seed point uphill on the density surface until it stops at a mode.

The modes that survive become the cluster centers. Their count comes from the data, not from you. This freedom comes at a price: the bandwidth. It is the only setting that controls everything. A wrong bandwidth is the only way to get a bad result.

RustyML exposes the algorithm through MeanShift and a free function, estimate_bandwidth. Both re-export from rustyml::machine_learning. The kernel, the merge rule, the noise label, and the bandwidth estimator all match scikit-learn 1.9.0. On a reference 10-point dataset, the cluster centers, their numbering, and every label match scikit-learn’s output exactly.

2.9.1. Mode seeking with a flat kernel

Picture a kernel density estimate laid over your points. Each sample adds a small bump. Where samples accumulate, the bumps stack into peaks.

Mean Shift places a candidate center at a seed location. It replaces the center with the mean of the data around it. That mean point lies in the direction where density increases. The center climbs to the nearest peak, a mode, over repeated iterations. Repeat the process from many seeds to find every mode.

The kernel decides which points weight that mean, and this detail decides the algorithm’s behavior. RustyML uses the flat kernel, the same kernel scikit-learn uses. For a center c, every point within bandwidth of it counts once. Every point outside it counts zero. The next center is the plain mean of the points inside the ball: mean{ x_i : ||c - x_i|| <= bandwidth }.

The window is a hard ball with radius bandwidth. The number of points it holds at convergence is the mode’s intensity. The merge phase below ranks modes by this intensity.

An earlier version of this crate used a Gaussian kernel weighted over the whole dataset. RustyML removed the Gaussian kernel instead of keeping it as an option. It has no scikit-learn counterpart, so nothing validates it. It also does not produce the window point count the merge rule needs. MeanShift has one kernel, and it has no kernel parameter.

At a very small bandwidth, a shifted center’s ball can end up empty. A naive fix divides by zero and collapses the center to the origin. That would add a false cluster at (0, 0, ...), unrelated to any real point. RustyML avoids this. It leaves the center where it is and stops iterating. This matches the early exit scikit-learn takes for an empty neighborhood.

A point with no neighbors becomes its own mode. Because of this rule, a bandwidth that shrinks toward zero degrades gracefully. The result becomes one cluster per isolated point, not meaningless output.

After every seed converges, MeanShift removes duplicate modes with scikit-learn’s method: intensity-ordered greedy suppression. First, it ranks the converged modes by how many points their window held, from most to least. Then it walks that order. It keeps each mode and discards every other mode within one bandwidth of it.

A kept center is a true density mode. An earlier version averaged the suppressed modes into the kept center instead. That approach dragged the center off the density peak. It also made the result depend on the order in which the pass processed the seeds.

The number of surviving centers is the cluster count. This count comes entirely from the data and the bandwidth. You never state the cluster count directly. Each input sample gets the label of its nearest surviving center.

This 2-phase structure, first converge many seeds and then merge them, explains why a bandwidth that is slightly too large still gives clean results. Even when seeds converge to slightly different spots, the merge step still combines them.

2.9.2. Constructing a MeanShift

MeanShift::new takes the bandwidth and returns a Result. A bandwidth that is not positive and finite is a usage error. RustyML does not clamp it silently:

let ms = MeanShift::new(2.0)?            // the only required argument
    .with_max_iter(300)?                 // returns Result, validates > 0
    .with_tolerance(1e-4)?               // returns Result, validates positive and finite
    .with_bin_seeding(true)              // returns Self, infallible toggle
    .with_cluster_all(true);             // returns Self, infallible toggle

The split in return types is deliberate, and it is easy to miss. The 2 convergence setters validate their argument and return Result<Self, Error>. They need the ? operator. The 2 boolean toggles cannot fail and return Self. You can chain them directly.

MeanShift::default() equals new(1.0) with every other setting at its default. This is convenient for a first look. It is rarely the bandwidth you actually want. See 2.9.4 for more.

Invalid arguments come back as Error::InvalidParameter from new, with_max_iter, and with_tolerance. This is the same error type described in 1.6. Error Handling.

2.9.3. Fitting, predicting, and reading the results

fit takes a 2-D array, one sample per row. It runs the algorithm and returns &mut Self. predict maps new points to the learned centers and returns Array1<isize>. fit_predict does both, and returns the training labels directly.

The labels are signed because -1 is reserved for noise, the same convention DBSCAN and scikit-learn use. This convention lets the output of any clustering estimator feed any metric in 5.3. Clustering Metrics without a conversion. After fit, getters expose everything the model found.

use ndarray::Array2;
use rustyml::machine_learning::MeanShift;

fn main() {
    // 2 tight blobs: 5 points near (0, 0) and 5 near (20, 20).
    let data = Array2::from_shape_vec(
        (10, 2),
        vec![
            -0.1, 0.0, 0.1, 0.0, 0.0, -0.1, 0.0, 0.1, 0.0, 0.0, // blob A
            19.9, 20.0, 20.1, 20.0, 20.0, 19.9, 20.0, 20.1, 20.0, 20.0, // blob B
        ],
    )
    .unwrap();

    let mut ms = MeanShift::new(2.0).unwrap();
    let labels = ms.fit_predict(&data).unwrap();

    let centers = ms.get_cluster_centers().unwrap();
    println!("clusters found: {}", centers.nrows()); // emerges from the data: 2
    println!("labels: {:?}", labels);
    println!("samples per center: {:?}", ms.get_n_samples_per_center().unwrap());
    println!("iterations run: {}", ms.get_actual_iterations().unwrap());
}

The getters fall into 2 groups: results and echoed settings. get_cluster_centers returns Option<&Array2<f64>>, one row per cluster. get_labels returns Option<&Array1<isize>>. get_n_samples_per_center returns Option<&Array1<usize>>, the number of input samples assigned to each center. With cluster_all = true, these counts sum to the sample count. With cluster_all = false, the -1 outliers are excluded from the sum.

get_actual_iterations returns Option<usize>, the largest iteration count over all seeds. Use it to tell whether the run converged or hit max_iter. All 4 of these getters return None before fit and Some after.

The remaining getters, get_bandwidth, get_max_iterations, get_tolerance, get_bin_seeding, and get_cluster_all, simply read back the configuration.

predict has 4 failure modes. Calling it before fit returns Error::NotFitted. Passing an empty array returns Error::EmptyInput. Passing points whose feature count differs from the training data returns Error::DimensionMismatch. Passing data with NaN or infinite values returns Error::NonFinite. None of these errors are recoverable by retrying, so treat them as programming errors that surface at runtime.

2.9.4. Bandwidth: the hyperparameter that decides everything

Everything about a Mean Shift run flows from the bandwidth. It sets the radius of the flat kernel’s window. It also sets the merge radius, and the cutoff for what counts as an outlier when cluster_all is off. The bandwidth controls how far each seed can see, how strongly nearby modes collapse into one, and where noise begins.

A bandwidth too small causes over-segmentation. Seeds cannot reach across a cluster’s own spread, and modes multiply. In the extreme case, every well-separated point becomes its own cluster. A bandwidth too large causes under-segmentation instead. Distant blobs pull on each other until their modes drift together, and the merge step fuses them into one. Eventually the whole dataset becomes a single cluster.

No default bandwidth works for arbitrary data. The correct value is a length scale in the units of your features.

When you have no prior estimate, estimate_bandwidth gives a data-driven starting point. It takes the data and an optional quantile (default 0.3). It also takes an optional n_samples subsample size (default: all rows, clamped to the dataset size) and an optional random_state.

It computes k = max(1, floor(n * quantile)). Then it measures each point’s distance to its (k - 1)-th nearest neighbor, and returns the mean of those distances. This is a local-density statistic. It answers how far a typical point is from the edge of its own neighborhood, which is exactly what a bandwidth needs to be.

The (k - 1) term reproduces scikit-learn’s off-by-one behavior, because its neighbor query counts the query point itself. The result agrees with scikit-learn 1.9.0 to within 1e-14. A neighborhood of one point yields 0.0, the same result scikit-learn gives.

An earlier version returned a quantile of the whole pairwise-distance distribution instead. That is a global spread measure. It runs far larger than a bandwidth should on clustered data. It collapsed everything into a single cluster. If you tuned a bandwidth against the old estimator, re-estimate it now.

A quantile of 0.3 gives a typical short-to-medium neighborhood radius. This value tends to land near the intra-cluster scale. Larger quantiles bias toward larger bandwidths and fewer clusters.

use ndarray::Array2;
use rustyml::machine_learning::{MeanShift, estimate_bandwidth};

fn main() {
    // 3 well-separated blobs, 12 tight points each.
    let mut v: Vec<f64> = Vec::new();
    for (cx, cy) in [(0.0, 0.0), (10.0, 0.0), (5.0, 9.0)] {
        for k in 0..12u32 {
            v.push(cx + ((k * 7) % 5) as f64 * 0.05 - 0.1);
            v.push(cy + ((k * 3) % 5) as f64 * 0.05 - 0.1);
        }
    }
    let data = Array2::from_shape_vec((36, 2), v).unwrap();

    // A reasonable starting point straight from the data.
    let bw = estimate_bandwidth(&data, Some(0.3), None, Some(0)).unwrap();
    println!("estimated bandwidth: {:.3}", bw);

    // Sweep: small over-segments, large collapses everything into one.
    for bandwidth in [0.05_f64, 0.5, 3.0, 30.0] {
        let mut ms = MeanShift::new(bandwidth).unwrap();
        ms.fit(&data).unwrap();
        let k = ms.get_cluster_centers().unwrap().nrows();
        println!("bandwidth {bandwidth:>5} -> {k} clusters");
    }
}

The cluster count moves with the bandwidth in the direction you would expect. Because the exact counts depend on the data, treat the following as the shape of the output, not literal numbers:

estimated bandwidth: <small positive value>
bandwidth  0.05 -> many clusters      (blobs fragment; over-segmentation)
bandwidth   0.5 -> one cluster per blob
bandwidth     3 -> one cluster per blob
bandwidth    30 -> a single cluster   (all blobs merged; under-segmentation)

Follow a simple workflow. Call estimate_bandwidth once. Fit at that value. Look at the cluster count. Then increase the bandwidth if you got too many clusters, or decrease it if you got too few. Validate the choice with a metric that does not need ground-truth labels, such as the silhouette score in 5.3. Clustering Metrics.

estimate_bandwidth measures distances directly. Standardize your features first (see 4.2. Standardization and Normalization) to make the estimate meaningful when feature scales differ.

2.9.5. Bin seeding, cluster_all, and the outlier label

By default, every input point is a seed. This is the most thorough option. As 2.9.7 explains, it is also why fitting is deterministic. On dense datasets, using every point as a seed is wasteful, because thousands of seeds inside one blob all climb to the same mode.

with_bin_seeding(true) fixes this. It quantizes the feature space onto a grid whose cells have side length bandwidth, then keeps 1 representative seed per occupied cell. Fewer seeds mean fewer uphill walks and a faster fit.

The cost is a coarser, approximate seeding. A mode whose basin never contains a grid representative can be missed. Bin seeding is worth the trade when the seed loop dominates the run time and the data is dense enough to fill whole cells. On small or sparse data, it saves little, and it can only cost resolution.

cluster_all decides what happens to points that do not really belong to any mode. With the default true, MeanShift forces every point to its nearest center. Labels always fall in 0..n_clusters, and there is no concept of noise. Set it to false, and any point farther than one bandwidth from every center gets the label -1. This is scikit-learn’s noise value, the same value DBSCAN uses. The rule applies to the training labels from fit and to fresh points from predict.

The -1 label replaced an older sentinel value equal to n_clusters. Downstream code that counted distinct labels read that old sentinel as a real extra cluster. If your code compares a label against the cluster count, change it to label < 0.

use ndarray::Array2;
use rustyml::machine_learning::MeanShift;

fn main() {
    let data = Array2::from_shape_vec(
        (10, 2),
        vec![
            -0.1, 0.0, 0.1, 0.0, 0.0, -0.1, 0.0, 0.1, 0.0, 0.0,
            19.9, 20.0, 20.1, 20.0, 20.0, 19.9, 20.0, 20.1, 20.0, 20.0,
        ],
    )
    .unwrap();

    let mut ms = MeanShift::new(2.0).unwrap().with_cluster_all(false);
    ms.fit(&data).unwrap();

    // (10, 10) sits ~14 units from both blobs, far beyond the bandwidth of 2.0.
    let probe = Array2::from_shape_vec((1, 2), vec![10.0, 10.0]).unwrap();
    let pred = ms.predict(&probe).unwrap();

    if pred[0] < 0 {
        println!("outlier: label {}", pred[0]); // -1
    } else {
        println!("assigned to cluster {}", pred[0]);
    }
}

When you enable cluster_all = false, test for label < 0. Do not assume labels are dense. A naive labels.iter().max() no longer tells you the cluster count if every point happens to be noise. get_cluster_centers().unwrap().nrows() gives the reliable count.

2.9.6. Cost, convergence, and parallelism

Mean Shift has quadratic cost. Consider this before you run it on a large dataset. Each seed’s iteration touches all n points in d dimensions, to check which points fall inside the window and to average them. One seed costs O(iterations * n * d). With the default seeding, every point is a seed, so there are n seeds. A full fit costs roughly O(iterations * n^2 * d).

This is the same asymptotic class as DBSCAN’s pairwise scan. It is considerably heavier than KMeans’ O(iterations * n * k * d), whose k is usually far below n. Bin seeding cuts the seed count from n to the number of occupied cells. This reduces the constant, but it does not change the quadratic per-iteration term.

The convergence bound is max_iter (default 300). A seed stops early once its shift falls below tol. get_actual_iterations reports the largest count any seed needed. A value pinned at max_iter signals that some seed never settled.

The implementation runs in parallel only when doing so helps. In fit, the per-seed uphill walks are independent. They run across a Rayon pool once the total work (seeds times samples times features) clears RustyML’s calibrated scan-class gate. This gate defaults to 262,144 element-operations, and you can adjust it through crate::tuning. Below that gate, the loop stays serial, because the overhead of forking is not worth it.

Inside each seed, the implementation casts the weighted-mean computations as matrix-vector products. It deliberately keeps those serial when the seed axis alone already fills the pool. This avoids nested Rayon forks that would compete with each other. predict parallelizes its nearest-center scan under the same gate, keyed on samples times clusters times features.

For most datasets, you get parallelism automatically, without changing any setting. The tuning knobs, and the reasoning behind these thresholds, live in 7.3. Performance Tuning and Parallelism.

2.9.7. Reproducibility and persistence

Fitting a MeanShift is deterministic. It takes no random seed. With the default seeding, it starts from every point. With bin seeding, it starts from a fixed grid representative per cell. Neither approach draws from a random number generator, so two fits on identical data produce byte-identical centers and labels.

This is a real convenience over KMeans, whose centroid initialization is randomized and needs a seed to reproduce. See 7.1. Reproducibility and Random Seeds for the broader picture.

The only place randomness enters this module is estimate_bandwidth, and only when you ask it to subsample. If n_samples is smaller than the dataset, estimate_bandwidth shuffles indices to pick the subset. Pass a fixed random_state when you need the estimate itself to be reproducible. Requesting all rows, the default for n_samples, removes the randomness entirely, because there is nothing to sample.

A fitted model serializes to a compact postcard binary through save_to_path, and restores with load_from_path. The saved file carries the centers, labels, hyperparameters, and training metadata. A reloaded model predicts identically without a re-fit.

use ndarray::Array2;
use rustyml::machine_learning::MeanShift;

fn main() {
    let data = Array2::from_shape_vec(
        (6, 2),
        vec![0.0, 0.0, 0.1, 0.1, -0.1, 0.0, 10.0, 10.0, 10.1, 9.9, 9.9, 10.0],
    )
    .unwrap();

    let mut ms = MeanShift::new(2.0).unwrap();
    ms.fit(&data).unwrap();

    let path = "mean_shift_model.bin";
    ms.save_to_path(path).unwrap();
    let restored = MeanShift::load_from_path(path).unwrap();

    let before = ms.predict(&data).unwrap();
    let after = restored.predict(&data).unwrap();
    assert_eq!(before, after); // identical after the round trip

    std::fs::remove_file(path).unwrap();
    println!("round-trip predictions match");
}

One caveat applies to model files saved before the kernel and merge rule changed. The centers inside those files came from the old Gaussian-kernel run. A reloaded model clusters differently from a fresh fit on the same data. This happens silently, because nothing about the file announces which algorithm produced it. Re-fit the model and save it again.

Persist the model when the fit is expensive and the data is stable. Downstream services can then load and predict cheaply. 7.2. Model Persistence in Depth covers the format details and versioning caveats.