MCPcopy Create free account
hub / github.com/NodeDB-Lab/nodedb / kmeans_centroids

Function kmeans_centroids

nodedb-vector/src/ivf.rs:221–306  ·  view source on GitHub ↗
(data: &[&[f32]], dim: usize, k: usize, max_iter: usize)

Source from the content-addressed store, hash-verified

219}
220
221fn kmeans_centroids(data: &[&[f32]], dim: usize, k: usize, max_iter: usize) -> Vec<Vec<f32>> {
222 let n = data.len();
223 let k = k.min(n);
224 if k == 0 {
225 return Vec::new();
226 }
227
228 let mut centroids: Vec<Vec<f32>> = vec![data[0].to_vec()];
229 let mut min_dists = vec![f32::MAX; n];
230
231 // Initialize min_dists against the first centroid.
232 for (i, point) in data.iter().enumerate() {
233 let d = distance(point, &centroids[0], DistanceMetric::L2);
234 if d < min_dists[i] {
235 min_dists[i] = d;
236 }
237 }
238
239 let mut rng = crate::hnsw::Xorshift64::new(0xC0FF_EEDE_ADBE_EF42);
240 for _ in 1..k {
241 let total: f64 = min_dists.iter().map(|&d| d as f64).sum();
242 let next_idx = if total < f64::EPSILON {
243 0
244 } else {
245 let target = rng.next_f64() * total;
246 let mut acc = 0.0f64;
247 let mut chosen = n - 1;
248 for (i, &d) in min_dists.iter().enumerate() {
249 acc += d as f64;
250 if acc >= target {
251 chosen = i;
252 break;
253 }
254 }
255 chosen
256 };
257 centroids.push(data[next_idx].to_vec());
258 let last = centroids.last().expect("just pushed");
259 for (i, point) in data.iter().enumerate() {
260 let d = distance(point, last, DistanceMetric::L2);
261 if d < min_dists[i] {
262 min_dists[i] = d;
263 }
264 }
265 }
266
267 let mut assignments = vec![0usize; n];
268 for _ in 0..max_iter {
269 let mut changed = false;
270 for (i, point) in data.iter().enumerate() {
271 let mut best = 0;
272 let mut best_d = f32::MAX;
273 for (c, centroid) in centroids.iter().enumerate() {
274 let d = distance(point, centroid, DistanceMetric::L2);
275 if d < best_d {
276 best_d = d;
277 best = c;
278 }

Callers 1

trainMethod · 0.85

Calls 9

sumMethod · 0.80
distanceFunction · 0.50
lenMethod · 0.45
iterMethod · 0.45
next_f64Method · 0.45
pushMethod · 0.45
to_vecMethod · 0.45
expectMethod · 0.45
lastMethod · 0.45

Tested by

no test coverage detected