use crate::dist::sqdist;
pub(crate) const FLOOR: usize = 256;
const KEEP: usize = 256;
const WALK: usize = 16;
const ROUNDS: usize = 3;
#[derive(Debug, Default)]
pub(crate) struct Coarse {
points: Vec<f32>,
under: Vec<Vec<u32>>,
owner: Vec<u32>,
built: usize,
}
impl Coarse {
pub(crate) fn ready(&self) -> bool {
!self.under.is_empty()
}
#[cfg(test)]
pub(crate) fn anchors(&self) -> usize {
self.under.len()
}
pub(crate) fn stale(&self, n: usize) -> bool {
if n < FLOOR {
return self.ready();
}
!self.ready() || n * 4 > self.built * 5 || n * 5 < self.built * 4
}
pub(crate) fn rebuild(&mut self, centroids: &[f32], dim: usize, n: usize) {
self.points.clear();
self.under.clear();
self.owner.clear();
self.built = n;
if n < FLOOR {
return;
}
let a = (n as f64).sqrt().ceil() as usize;
let stride = n / a;
self.points.reserve(a * dim);
for i in 0..a {
let at = (i * stride).min(n - 1) * dim;
self.points.extend_from_slice(¢roids[at..at + dim]);
}
self.under = vec![Vec::new(); a];
self.owner = vec![0; n];
self.assign(centroids, dim, n);
for _ in 0..ROUNDS {
self.recentre(centroids, dim);
self.assign(centroids, dim, n);
}
}
fn assign(&mut self, centroids: &[f32], dim: usize, n: usize) {
for list in &mut self.under {
list.clear();
}
for p in 0..n {
let x = ¢roids[p * dim..(p + 1) * dim];
let owner = self.nearest_anchor(x, dim);
self.owner[p] = owner as u32;
self.under[owner].push(p as u32);
}
}
fn recentre(&mut self, centroids: &[f32], dim: usize) {
for (a, list) in self.under.iter().enumerate() {
if list.is_empty() {
continue;
}
let at = &mut self.points[a * dim..(a + 1) * dim];
at.fill(0.0);
for &p in list {
let x = ¢roids[p as usize * dim..(p as usize + 1) * dim];
for (s, v) in at.iter_mut().zip(x) {
*s += *v;
}
}
let by = list.len() as f32;
for s in at {
*s /= by;
}
}
}
pub(crate) fn added(&mut self, p: usize, x: &[f32], dim: usize) {
if !self.ready() {
return;
}
debug_assert_eq!(p, self.owner.len(), "a partition is added at the end");
let owner = self.nearest_anchor(x, dim);
self.owner.push(owner as u32);
self.under[owner].push(p as u32);
}
pub(crate) fn dropped(&mut self, p: usize) {
if !self.ready() {
return;
}
let last = self.owner.len() - 1;
let owner = self.owner[p] as usize;
self.under[owner].retain(|&q| q as usize != p);
if p != last {
let moved = self.owner[last] as usize;
for q in &mut self.under[moved] {
if *q as usize == last {
*q = p as u32;
}
}
self.owner[p] = moved as u32;
}
self.owner.pop();
}
pub(crate) fn moved(&mut self, p: usize, x: &[f32], dim: usize) {
if !self.ready() {
return;
}
let was = self.owner[p] as usize;
let now = self.nearest_anchor(x, dim);
if now == was {
return;
}
self.under[was].retain(|&q| q as usize != p);
self.under[now].push(p as u32);
self.owner[p] = now as u32;
}
pub(crate) fn shortlist(&self, x: &[f32], dim: usize, out: &mut Vec<u32>) {
out.clear();
let mut near = [(0u32, f32::INFINITY); WALK];
let mut held = 0usize;
for i in 0..self.under.len() {
let d = sqdist(x, &self.points[i * dim..(i + 1) * dim]);
if held == WALK && d >= near[WALK - 1].1 {
continue;
}
let mut at = held.min(WALK - 1);
while at > 0 && near[at - 1].1 > d {
near[at] = near[at - 1];
at -= 1;
}
near[at] = (i as u32, d);
held = (held + 1).min(WALK);
}
for &(i, _) in &near[..held] {
let list = &self.under[i as usize];
out.extend_from_slice(list);
if out.len() >= KEEP && out.len() > list.len() {
break;
}
}
}
fn nearest_anchor(&self, x: &[f32], dim: usize) -> usize {
let mut best = 0;
let mut at = f32::INFINITY;
for i in 0..self.under.len() {
let d = sqdist(x, &self.points[i * dim..(i + 1) * dim]);
if d < at {
at = d;
best = i;
}
}
best
}
}
#[cfg(test)]
mod tests {
use super::*;
use yo_common::Rng;
fn spread(n: usize, dim: usize, seed: u64) -> Vec<f32> {
let mut rng = Rng::new(seed);
(0..n * dim)
.map(|_| (rng.next_u64() >> 40) as f32 / (1u32 << 24) as f32)
.collect()
}
fn sized(n: usize, dim: usize) -> (usize, usize) {
if cfg!(miri) { (FLOOR, 4) } else { (n, dim) }
}
fn built(n: usize, dim: usize, seed: u64) -> (Coarse, Vec<f32>) {
let centroids = spread(n, dim, seed);
let mut c = Coarse::default();
c.rebuild(¢roids, dim, n);
(c, centroids)
}
#[test]
fn a_small_collection_gets_no_layer_at_all() {
let (c, _) = built(FLOOR - 1, 8, 1);
assert!(!c.ready(), "under the floor there is nothing to look at");
assert!(!c.stale(FLOOR - 1), "and nothing to rebuild");
}
fn intact(c: &Coarse, n: usize) {
assert_eq!(c.owner.len(), n, "one owner per partition");
let mut seen = vec![0usize; n];
for (a, list) in c.under.iter().enumerate() {
for &p in list {
assert!((p as usize) < n, "anchor {a} holds partition {p} of {n}");
assert_eq!(c.owner[p as usize] as usize, a, "partition {p} disagrees");
seen[p as usize] += 1;
}
}
for (p, times) in seen.iter().enumerate() {
assert_eq!(*times, 1, "partition {p} is under {times} anchors");
}
}
#[test]
fn every_centroid_ends_up_under_exactly_one_anchor() {
let (n, dim) = sized(1000, 16);
let (c, _) = built(n, dim, 2);
assert!(c.ready());
assert_eq!(c.anchors(), if cfg!(miri) { 16 } else { 32 });
intact(&c, n);
}
#[test]
fn adding_and_dropping_partitions_keeps_the_lists_straight() {
let (n, dim) = sized(400, 16);
let (mut c, mut centroids) = built(n, dim, 3);
let mut n = n;
let extra = spread(50, dim, 4);
for i in 0..50 {
let x = &extra[i * dim..(i + 1) * dim];
centroids.extend_from_slice(x);
c.added(n, x, dim);
n += 1;
}
intact(&c, n);
for p in [7usize, 0, 100] {
let last = n - 1;
centroids.copy_within(last * dim..(last + 1) * dim, p * dim);
centroids.truncate(last * dim);
c.dropped(p);
n -= 1;
intact(&c, n);
}
c.dropped(n - 1);
n -= 1;
intact(&c, n);
}
#[test]
fn a_centroid_that_moves_moves_between_anchors() {
let (n, dim) = sized(500, 16);
let (mut c, mut centroids) = built(n, dim, 5);
let target: Vec<f32> = c.points[9 * dim..10 * dim].to_vec();
centroids[3 * dim..4 * dim].copy_from_slice(&target);
c.moved(3, &target, dim);
assert_eq!(c.owner[3], 9, "it belongs to the anchor it is sitting on");
intact(&c, n);
}
#[test]
#[cfg_attr(
miri,
ignore = "the count is the claim: a recall rate over five hundred queries against a thousand centroids, which is nothing at all at a size Miri can afford"
)]
fn the_nearest_centroid_is_almost_always_on_the_shortlist() {
let dim = 32;
let n = 1000;
let (c, centroids) = built(n, dim, 6);
let queries = spread(500, dim, 7);
let mut out = Vec::new();
let mut found = 0;
for q in 0..500 {
let x = &queries[q * dim..(q + 1) * dim];
let truth = (0..n)
.min_by(|&i, &j| {
sqdist(x, ¢roids[i * dim..(i + 1) * dim])
.total_cmp(&sqdist(x, ¢roids[j * dim..(j + 1) * dim]))
})
.unwrap();
c.shortlist(x, dim, &mut out);
assert!(
out.len() >= KEEP,
"a shortlist of {} is too short",
out.len()
);
assert!(
out.len() < n,
"a shortlist of everything is not a shortlist"
);
if out.contains(&(truth as u32)) {
found += 1;
}
}
assert!(
found >= 390,
"the nearest centroid was on the shortlist {found} times in 500"
);
}
#[test]
#[cfg_attr(
miri,
ignore = "the count is the claim: ten times as many centroids and the same amount of work, and the bigger of the two sizes is the point"
)]
fn a_bigger_collection_does_not_mean_a_bigger_shortlist() {
let dim = 16;
let x = spread(1, dim, 9);
let mut out = Vec::new();
let mut sizes = Vec::new();
for n in [900usize, 9000] {
let (c, _) = built(n, dim, 8);
c.shortlist(&x, dim, &mut out);
assert!(
out.len() >= KEEP,
"a shortlist of {} is too short",
out.len()
);
sizes.push(out.len());
}
assert!(
sizes[1] < sizes[0] * 2,
"{} candidates at 900 centroids and {} at 9000",
sizes[0],
sizes[1]
);
}
#[test]
fn the_layer_is_rebuilt_when_the_collection_has_moved_a_quarter() {
let (n, dim) = sized(1000, 8);
let (c, _) = built(n, dim, 10);
assert!(!c.stale(n));
assert!(!c.stale(n * 11 / 10), "a tenth is not worth a rebuild");
assert!(c.stale(n * 7 / 5), "nearly a half is");
assert!(c.stale(n * 7 / 10), "and so is shrinking by a third");
}
}