qualia_core_db/solvers/learning/clustering/
kmeans.rs

1//! k-means clustering (ISL ch 12.4, PRML ch 9.1) — Lloyd's algorithm with
2//! k-means++ seeding, over a row-major feature matrix.
3//!
4//! Assign each point to its nearest centroid (squared Euclidean), recompute each
5//! centroid as the mean of its members, repeat to convergence. k-means++ seeding
6//! spreads the initial centroids to avoid poor local minima. Kernel-class
7//! `AllPairs` (the point↔centroid distances), dispatch-ready; deterministic given
8//! the seed.
9
10use crate::solvers::learning::LearningError;
11
12/// A fitted k-means model.
13#[derive(Debug, Clone)]
14pub struct KMeansModel {
15    /// `k × p` centroids, row-major.
16    pub centroids: Vec<f64>,
17    /// Cluster assignment per input row.
18    pub labels: Vec<usize>,
19    /// Within-cluster sum of squared distances (the objective; lower is tighter).
20    pub inertia: f64,
21    pub k: usize,
22    pub p: usize,
23    pub n_iter: usize,
24    pub converged: bool,
25}
26
27impl KMeansModel {
28    /// Index of the nearest centroid to a feature row.
29    pub fn predict_row(&self, x_row: &[f64]) -> usize {
30        nearest(&self.centroids, self.k, self.p, x_row).0
31    }
32}
33
34struct Lcg(u64);
35impl Lcg {
36    fn unit(&mut self) -> f64 {
37        self.0 = self
38            .0
39            .wrapping_mul(6364136223846793005)
40            .wrapping_add(1442695040888963407);
41        ((self.0 >> 11) as f64) / ((1u64 << 53) as f64)
42    }
43}
44
45#[inline]
46fn sq_dist(a: &[f64], b: &[f64]) -> f64 {
47    a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum()
48}
49
50/// (index, squared distance) of the nearest of `k` row-major centroids to `point`.
51fn nearest(centroids: &[f64], k: usize, p: usize, point: &[f64]) -> (usize, f64) {
52    let mut best = 0;
53    let mut best_d = f64::INFINITY;
54    for c in 0..k {
55        let d = sq_dist(&centroids[c * p..(c + 1) * p], point);
56        if d < best_d {
57            best_d = d;
58            best = c;
59        }
60    }
61    (best, best_d)
62}
63
64/// k-means++ seeding: choose `k` initial centroids spread by D²-weighting.
65fn kmeans_pp(x: &[f64], n: usize, p: usize, k: usize, rng: &mut Lcg) -> Vec<f64> {
66    let mut centroids = vec![0.0; k * p];
67    // First centroid uniformly at random.
68    let first = ((rng.unit() * n as f64) as usize).min(n - 1);
69    centroids[..p].copy_from_slice(&x[first * p..(first + 1) * p]);
70    let mut d2 = vec![0.0; n];
71    for c in 1..k {
72        // D²(point) = squared distance to the nearest chosen centroid.
73        let mut total = 0.0;
74        for i in 0..n {
75            let (_, d) = nearest(&centroids[..c * p], c, p, &x[i * p..(i + 1) * p]);
76            d2[i] = d;
77            total += d;
78        }
79        // Sample proportional to D².
80        let target = rng.unit() * total;
81        let mut acc = 0.0;
82        let mut chosen = n - 1;
83        for i in 0..n {
84            acc += d2[i];
85            if acc >= target {
86                chosen = i;
87                break;
88            }
89        }
90        centroids[c * p..(c + 1) * p].copy_from_slice(&x[chosen * p..(chosen + 1) * p]);
91    }
92    centroids
93}
94
95/// Fit k-means with `k` clusters. Fails closed: `InvalidDimension`,
96/// `InsufficientData` (`k == 0` or `k > n`).
97pub fn fit(
98    x: &[f64],
99    n: usize,
100    p: usize,
101    k: usize,
102    max_iter: usize,
103    seed: u64,
104) -> Result<KMeansModel, LearningError> {
105    if n == 0 || p == 0 || x.len() != n * p {
106        return Err(LearningError::InvalidDimension);
107    }
108    if k == 0 || k > n {
109        return Err(LearningError::InsufficientData);
110    }
111
112    let mut rng = Lcg(seed ^ 0x9E3779B97F4A7C15);
113    let mut centroids = kmeans_pp(x, n, p, k, &mut rng);
114    let mut labels = vec![0usize; n];
115    let mut converged = false;
116    let mut iters = 0;
117
118    for it in 1..=max_iter.max(1) {
119        iters = it;
120        // Assignment step — assign each point to its nearest centroid.
121        //
122        // Best-path-with-CPU-floor (mirrors `linear_algebra::gemm`): the point↔centroid
123        // squared-distance matrix (`n × k`, the `AllPairs` kernel) is the dominant cost
124        // when `n·k·p` is large, so above `GEMM_GPU_THRESHOLD` and with an accelerator
125        // present we compute it in one pass via `dispatch::pairwise_sq_dist_f64` (whose
126        // cross-term GEMM takes the best path on this machine) and argmin each row. Off
127        // accelerator, or sub-threshold, the EXACT per-point `nearest` loop runs —
128        // byte-identical to before, including its lowest-index tie-break.
129        let mut changed = false;
130        // GPU best-path (point↔centroid squared-distance matrix via the forge) only when
131        // it's compiled in (native + wgsl-forge). On wasm32 the exact per-point CPU loop runs.
132        #[cfg(all(not(target_arch = "wasm32"), feature = "wgsl-forge"))]
133        {
134            let work = n.saturating_mul(k).saturating_mul(p);
135            let caps = crate::wgsl_forge::dispatch::caps();
136            if (caps.cuda || caps.wgpu) && work >= crate::wgsl_forge::dispatch::GEMM_GPU_THRESHOLD {
137                let dist =
138                    crate::wgsl_forge::dispatch::pairwise_sq_dist_f64(x, &centroids, n, k, p);
139                for i in 0..n {
140                    let row = &dist[i * k..(i + 1) * k];
141                    let mut best = 0;
142                    let mut best_d = row[0];
143                    for (c, &d) in row.iter().enumerate().skip(1) {
144                        if d < best_d {
145                            best_d = d;
146                            best = c;
147                        }
148                    }
149                    if labels[i] != best {
150                        labels[i] = best;
151                        changed = true;
152                    }
153                }
154            } else {
155                for i in 0..n {
156                    let (c, _) = nearest(&centroids, k, p, &x[i * p..(i + 1) * p]);
157                    if labels[i] != c {
158                        labels[i] = c;
159                        changed = true;
160                    }
161                }
162            }
163        }
164        #[cfg(not(all(not(target_arch = "wasm32"), feature = "wgsl-forge")))]
165        {
166            for i in 0..n {
167                let (c, _) = nearest(&centroids, k, p, &x[i * p..(i + 1) * p]);
168                if labels[i] != c {
169                    labels[i] = c;
170                    changed = true;
171                }
172            }
173        }
174        // Update step: centroid = mean of its members; empty clusters keep place.
175        let mut sums = vec![0.0; k * p];
176        let mut counts = vec![0usize; k];
177        for i in 0..n {
178            let c = labels[i];
179            counts[c] += 1;
180            for j in 0..p {
181                sums[c * p + j] += x[i * p + j];
182            }
183        }
184        for c in 0..k {
185            if counts[c] > 0 {
186                for j in 0..p {
187                    centroids[c * p + j] = sums[c * p + j] / counts[c] as f64;
188                }
189            }
190        }
191        if !changed && it > 1 {
192            converged = true;
193            break;
194        }
195    }
196
197    // Final inertia.
198    let mut inertia = 0.0;
199    for i in 0..n {
200        inertia += sq_dist(
201            &centroids[labels[i] * p..(labels[i] + 1) * p],
202            &x[i * p..(i + 1) * p],
203        );
204    }
205
206    Ok(KMeansModel {
207        centroids,
208        labels,
209        inertia,
210        k,
211        p,
212        n_iter: iters,
213        converged,
214    })
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220
221    #[test]
222    fn recovers_three_separated_blobs() {
223        // Three tight clusters around (0,0), (10,10), (0,10).
224        let mut x = Vec::new();
225        let centers = [(0.0, 0.0), (10.0, 10.0), (0.0, 10.0)];
226        for &(cx, cy) in &centers {
227            for d in 0..5 {
228                x.push(cx + (d as f64 - 2.0) * 0.1);
229                x.push(cy + (d as f64 - 2.0) * 0.1);
230            }
231        }
232        let n = 15;
233        let m = fit(&x, n, 2, 3, 100, 1).unwrap();
234        assert!(m.converged);
235        // Each blob's 5 points share a label.
236        for blob in 0..3 {
237            let base = blob * 5;
238            let l = m.labels[base];
239            assert!(
240                (base..base + 5).all(|i| m.labels[i] == l),
241                "blob {blob} not pure"
242            );
243        }
244        // Three distinct labels used.
245        let mut used: Vec<usize> = m.labels.clone();
246        used.sort_unstable();
247        used.dedup();
248        assert_eq!(used.len(), 3);
249        // Tight clusters → small inertia.
250        assert!(m.inertia < 1.0, "inertia {}", m.inertia);
251    }
252
253    #[test]
254    fn single_cluster_centroid_is_the_mean() {
255        let x = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; // 3 points in 2-D
256        let m = fit(&x, 3, 2, 1, 50, 0).unwrap();
257        assert!((m.centroids[0] - 3.0).abs() < 1e-9); // mean of x-coords
258        assert!((m.centroids[1] - 4.0).abs() < 1e-9); // mean of y-coords
259        assert!(m.labels.iter().all(|&l| l == 0));
260    }
261
262    #[test]
263    fn guards() {
264        assert_eq!(
265            fit(&[1.0, 2.0], 1, 2, 3, 10, 0).unwrap_err(),
266            LearningError::InsufficientData
267        );
268        assert_eq!(
269            fit(&[1.0, 2.0, 3.0], 2, 2, 1, 10, 0).unwrap_err(),
270            LearningError::InvalidDimension
271        );
272    }
273}
qualia_core_db/solvers/learning/clustering/kmeans.rs

qualia_core_db/solvers/learning/clustering/
kmeans.rs
