Simple k-means clustering for PQ codebook training. Uses proper k-means++ initialization (weighted d² sampling) with a deterministic seed so training is reproducible across runs.
(data: &[&[f32]], dim: usize, k: usize, max_iter: usize)
| 272 | /// Uses proper k-means++ initialization (weighted d² sampling) with a |
| 273 | /// deterministic seed so training is reproducible across runs. |
| 274 | fn kmeans(data: &[&[f32]], dim: usize, k: usize, max_iter: usize) -> Vec<Vec<f32>> { |
| 275 | let n = data.len(); |
| 276 | if n == 0 || k == 0 { |
| 277 | return Vec::new(); |
| 278 | } |
| 279 | let k = k.min(n); // Can't have more centroids than data points. |
| 280 | |
| 281 | // K-means++ initialization with deterministic xorshift. |
| 282 | let mut rng = crate::hnsw::Xorshift64::new(0xC0FF_EEDE_ADBE_EF42); |
| 283 | |
| 284 | let mut centroids: Vec<Vec<f32>> = Vec::with_capacity(k); |
| 285 | centroids.push(data[0].to_vec()); |
| 286 | |
| 287 | let mut min_dists = vec![f32::MAX; n]; |
| 288 | // Update against the first centroid. |
| 289 | for (i, point) in data.iter().enumerate() { |
| 290 | let d = l2_sub(point, ¢roids[0]); |
| 291 | if d < min_dists[i] { |
| 292 | min_dists[i] = d; |
| 293 | } |
| 294 | } |
| 295 | |
| 296 | for _ in 1..k { |
| 297 | let total: f64 = min_dists.iter().map(|&d| d as f64).sum(); |
| 298 | let next_idx = if total < f64::EPSILON { |
| 299 | // All points coincide with existing centroids. |
| 300 | 0 |
| 301 | } else { |
| 302 | let target = rng.next_f64() * total; |
| 303 | let mut acc = 0.0f64; |
| 304 | let mut chosen = n - 1; |
| 305 | for (i, &d) in min_dists.iter().enumerate() { |
| 306 | acc += d as f64; |
| 307 | if acc >= target { |
| 308 | chosen = i; |
| 309 | break; |
| 310 | } |
| 311 | } |
| 312 | chosen |
| 313 | }; |
| 314 | centroids.push(data[next_idx].to_vec()); |
| 315 | // Incrementally update min_dists against the new centroid. |
| 316 | let last = centroids.last().expect("just pushed"); |
| 317 | for (i, point) in data.iter().enumerate() { |
| 318 | let d = l2_sub(point, last); |
| 319 | if d < min_dists[i] { |
| 320 | min_dists[i] = d; |
| 321 | } |
| 322 | } |
| 323 | } |
| 324 | |
| 325 | // K-means iterations. |
| 326 | let mut assignments = vec![0usize; n]; |
| 327 | for _ in 0..max_iter { |
| 328 | // Assignment step. |
| 329 | let mut changed = false; |
| 330 | for (i, point) in data.iter().enumerate() { |
| 331 | let mut best = 0; |