Add benchmarks for spatial trees

spatial_bench/Cargo.toml

@@ -6,11 +6,13 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-bosque = "0.1.0"
+bosque = { git = "https://github.com/cavemanloverboy/bosque"}
 criterion = "0.5.1"
 csv = "1.3.0"
-fnntw = "0.4.1"
+fastrand = "2.0.1"
+kiddo = "4.0.0"
 nabo = "0.3.0"
+nalgebra = "0.32.3"
 rstar = "0.11.0"
 
 [[bench]]

spatial_bench/benches/spatial.rs

@@ -1,9 +1,78 @@
+use std::iter::repeat_with;
+
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
-use spatial_bench as sb;
+use fastrand::Rng;
+use spatial_bench::{
+    bosque::BosqueArena, kiddo::KiddoArena, nabo::NaboArena, read_augmented, rstar::RstarArena,
+    Point3, SpatialArena,
+};
+
+const N_NEURONS: usize = 1000;
+
+fn make_arena<S: SpatialArena>(pts: Vec<Vec<Point3>>) -> S {
+    let mut ar = S::default();
+    for p in pts.into_iter() {
+        ar.add_points(p);
+    }
+    ar
+}
+
+fn random_pairs(max: usize, n: usize, rng: &mut Rng) -> Vec<(usize, usize)> {
+    repeat_with(|| (rng.usize(0..max), rng.usize(0..max)))
+        .take(n)
+        .collect()
+}
+
+fn pair_queries<S: SpatialArena>(arena: &S, pairs: &[(usize, usize)]) {
+    for (q, t) in pairs {
+        arena.query_target(*q, *t);
+    }
+}
+
+fn read_augmented_fixed() -> Vec<Vec<Point3>> {
+    let mut rng = fastrand::Rng::with_seed(1991);
+    read_augmented(N_NEURONS, &mut rng, 20.0, 1.0)
+}
+
+pub fn bench_construction(c: &mut Criterion) {
+    let points = read_augmented_fixed();
+
+    let mut group = c.benchmark_group("construction");
+    group.bench_function("bosque", |b| {
+        b.iter(|| black_box(make_arena::<BosqueArena>(points.clone())))
+    });
+    group.bench_function("kiddo", |b| {
+        b.iter(|| black_box(make_arena::<KiddoArena>(points.clone())))
+    });
+    group.bench_function("nabo", |b| {
+        b.iter(|| black_box(make_arena::<NaboArena>(points.clone())))
+    });
+    group.bench_function("rstar", |b| {
+        b.iter(|| black_box(make_arena::<RstarArena>(points.clone())))
+    });
+}
+
+pub fn bench_queries(c: &mut Criterion) {
+    let points = read_augmented_fixed();
+    let n_pairs = 1_000;
+    let mut rng = fastrand::Rng::with_seed(1991);
+    let pairs = random_pairs(points.len(), n_pairs, &mut rng);
+    let mut group = c.benchmark_group("pairwise query");
+
+    let ar = make_arena::<BosqueArena>(points.clone());
+    group.bench_function("bosque", |b| {
+        b.iter(|| black_box(pair_queries(&ar, &pairs)))
+    });
+
+    let ar = make_arena::<KiddoArena>(points.clone());
+    group.bench_function("kiddo", |b| b.iter(|| black_box(pair_queries(&ar, &pairs))));
+
+    let ar = make_arena::<NaboArena>(points.clone());
+    group.bench_function("nabo", |b| b.iter(|| black_box(pair_queries(&ar, &pairs))));
 
-pub fn criterion_benchmark(c: &mut Criterion) {
-    c.bench_function("fib 20", |b| b.iter(|| fibonacci(black_box(20))));
+    let ar = make_arena::<RstarArena>(points.clone());
+    group.bench_function("rstar", |b| b.iter(|| black_box(pair_queries(&ar, &pairs))));
 }
 
-criterion_group!(benches, criterion_benchmark);
+criterion_group!(benches, bench_construction, bench_queries);
 criterion_main!(benches);

spatial_bench/src/bosque.rs

@@ -0,0 +1,50 @@
+use crate::{Point3, Precision, SpatialArena};
+
+#[derive(Default)]
+pub struct BosqueArena {
+    trees: Vec<Vec<Point3>>,
+    idxs: Vec<Vec<usize>>,
+}
+
+impl SpatialArena for BosqueArena {
+    fn add_points(&mut self, mut p: Vec<Point3>) -> usize {
+        let mut idxs: Vec<_> = (0..(p.len() as u32)).collect();
+        bosque::tree::build_tree_with_indices(p.as_mut(), idxs.as_mut());
+        let idx = self.len();
+        self.trees.push(p);
+        self.idxs
+            .push(idxs.into_iter().map(|i| i as usize).collect());
+        idx
+    }
+
+    fn query_target(&self, q: usize, t: usize) -> Vec<(usize, Precision)> {
+        let tgt = self.trees.get(t).unwrap();
+        let tgt_idxs = self.idxs.get(t).unwrap();
+        self.trees
+            .get(q)
+            .unwrap()
+            .iter()
+            .map(|p| {
+                let (d, idx) = bosque::tree::nearest_one(tgt.as_slice(), p);
+                (tgt_idxs[idx], d)
+            })
+            .collect()
+    }
+
+    fn local_query(&self, q: usize, neighborhood: usize) -> Vec<Vec<usize>> {
+        let t = self.trees.get(q).unwrap();
+        let idxs = self.idxs.get(q).unwrap();
+        t.iter()
+            .map(|p| {
+                bosque::tree::nearest_k(t.as_slice(), p, neighborhood)
+                    .into_iter()
+                    .map(|(_d, i)| idxs[i])
+                    .collect()
+            })
+            .collect()
+    }
+
+    fn len(&self) -> usize {
+        self.trees.len()
+    }
+}

spatial_bench/src/kiddo.rs

@@ -0,0 +1,50 @@
+use kiddo::{ImmutableKdTree, SquaredEuclidean};
+
+use crate::{Point3, Precision, SpatialArena};
+
+#[derive(Default)]
+pub struct KiddoArena {
+    trees: Vec<ImmutableKdTree<Precision, 3>>,
+    points: Vec<Vec<Point3>>,
+}
+
+impl SpatialArena for KiddoArena {
+    fn add_points(&mut self, p: Vec<Point3>) -> usize {
+        let idx = self.len();
+        self.trees.push(p.as_slice().into());
+        self.points.push(p);
+        idx
+    }
+
+    fn query_target(&self, q: usize, t: usize) -> Vec<(usize, Precision)> {
+        let tgt = self.trees.get(t).unwrap();
+        self.points
+            .get(q)
+            .unwrap()
+            .iter()
+            .map(|p| {
+                let nn = tgt.nearest_one::<SquaredEuclidean>(p);
+                (nn.item as usize, nn.distance.sqrt())
+            })
+            .collect()
+    }
+
+    fn local_query(&self, q: usize, neighborhood: usize) -> Vec<Vec<usize>> {
+        let tgt = self.trees.get(q).unwrap();
+        self.points
+            .get(q)
+            .unwrap()
+            .iter()
+            .map(|p| {
+                tgt.nearest_n::<SquaredEuclidean>(p, neighborhood)
+                    .into_iter()
+                    .map(|nn| nn.item as usize)
+                    .collect()
+            })
+            .collect()
+    }
+
+    fn len(&self) -> usize {
+        self.trees.len()
+    }
+}

spatial_bench/src/lib.rs

@@ -1,10 +1,18 @@
+use fastrand::Rng;
+use std::collections::HashMap;
 use std::fs::File;
 use std::path::PathBuf;
 
+pub mod bosque;
+pub mod kiddo;
+pub mod nabo;
+pub mod rstar;
+
 use csv::ReaderBuilder;
 
 pub type Precision = f64;
 pub type Point3 = [Precision; 3];
+pub const DIM: usize = 3;
 
 const NAMES: [&str; 20] = [
     "ChaMARCM-F000586_seg002",
@@ -75,3 +83,126 @@ fn read_points(name: &str) -> Vec<Point3> {
 fn read_all_points() -> Vec<Vec<Point3>> {
     NAMES.iter().map(|n| read_points(n)).collect()
 }
+
+pub trait SpatialArena: Default {
+    fn add_points(&mut self, p: Vec<Point3>) -> usize;
+
+    fn query_target(&self, q: usize, t: usize) -> Vec<(usize, Precision)>;
+
+    fn local_query(&self, q: usize, neighborhood: usize) -> Vec<Vec<usize>>;
+
+    fn len(&self) -> usize;
+
+    fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    fn all_locals(&self, neighborhood: usize) -> Vec<Vec<Vec<usize>>> {
+        (0..self.len())
+            .map(|idx| self.local_query(idx, neighborhood))
+            .collect()
+    }
+
+    fn all_v_all(&self) -> Vec<Vec<Vec<(usize, Precision)>>> {
+        let len = self.len();
+        (0..len)
+            .map(|q| {
+                (0..len)
+                    .map(move |t| self.query_target(q, t.clone()))
+                    .collect()
+            })
+            .collect()
+    }
+}
+
+#[derive(Default)]
+pub struct ArenaWrapper<S: SpatialArena>(S);
+
+impl<S: SpatialArena> SpatialArena for ArenaWrapper<S> {
+    fn add_points(&mut self, p: Vec<Point3>) -> usize {
+        self.0.add_points(p)
+    }
+
+    fn query_target(&self, q: usize, t: usize) -> Vec<(usize, Precision)> {
+        self.0.query_target(q, t)
+    }
+
+    fn local_query(&self, q: usize, neighborhood: usize) -> Vec<Vec<usize>> {
+        self.0.local_query(q, neighborhood)
+    }
+
+    fn len(&self) -> usize {
+        self.0.len()
+    }
+}
+
+struct PointAug<'a> {
+    rng: &'a mut Rng,
+    translation: Point3,
+    jitter_stdev: Precision,
+}
+
+fn box_muller_sin(u1: Precision, u2: Precision) -> Precision {
+    (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).sin()
+}
+
+fn random_translation(rng: &mut Rng, stdev: Precision) -> Point3 {
+    let u1 = rng.f64();
+    let u2 = rng.f64();
+    let u3 = rng.f64();
+    [
+        box_muller_sin(u1, u2) * stdev,
+        box_muller_sin(u2, u3) * stdev,
+        box_muller_sin(u1, u3) * stdev,
+    ]
+}
+
+impl<'a> PointAug<'a> {
+    pub fn new(translation: Point3, jitter_stdev: Precision, rng: &'a mut Rng) -> Self {
+        Self {
+            rng,
+            translation,
+            jitter_stdev,
+        }
+    }
+
+    pub fn new_random(
+        translation_stdev: Precision,
+        jitter_stdev: Precision,
+        rng: &'a mut Rng,
+    ) -> Self {
+        Self::new(
+            random_translation(rng, translation_stdev),
+            jitter_stdev,
+            rng,
+        )
+    }
+
+    pub fn augment(&mut self, orig: &Point3) -> Point3 {
+        let t2 = random_translation(self.rng, self.jitter_stdev);
+        [
+            orig[0] + self.translation[0] + t2[0],
+            orig[1] + self.translation[1] + t2[1],
+            orig[2] + self.translation[2] + t2[2],
+        ]
+    }
+
+    pub fn augment_all(&mut self, orig: &[Point3]) -> Vec<Point3> {
+        orig.iter().map(|p| self.augment(p)).collect()
+    }
+}
+
+pub fn read_augmented(
+    n: usize,
+    rng: &mut Rng,
+    translation_stdev: Precision,
+    jitter_stdev: Precision,
+) -> Vec<Vec<Point3>> {
+    let orig = read_all_points();
+    let mut aug = PointAug::new_random(translation_stdev, jitter_stdev, rng);
+    orig.iter()
+        .cycle()
+        .take(n)
+        .map(|p| aug.augment_all(p))
+        .collect()
+}