From f395a7ae6e86b44a4ca0eb8959a9f2e481ee294c Mon Sep 17 00:00:00 2001
From: Chris Barnes <chrislloydbarnes@gmail.com>
Date: Sat, 20 Jan 2024 12:15:18 +0000
Subject: [PATCH] Add benchmarks for spatial trees

---
 spatial_bench/Cargo.toml         |   6 +-
 spatial_bench/benches/spatial.rs |  77 +++++++++++++++++-
 spatial_bench/src/bosque.rs      |  50 ++++++++++++
 spatial_bench/src/kiddo.rs       |  50 ++++++++++++
 spatial_bench/src/lib.rs         | 131 +++++++++++++++++++++++++++++++
 spatial_bench/src/nabo.rs        |  92 ++++++++++++++++++++++
 spatial_bench/src/rstar.rs       |  55 +++++++++++++
 7 files changed, 455 insertions(+), 6 deletions(-)
 create mode 100644 spatial_bench/src/bosque.rs
 create mode 100644 spatial_bench/src/kiddo.rs
 create mode 100644 spatial_bench/src/nabo.rs
 create mode 100644 spatial_bench/src/rstar.rs

diff --git a/spatial_bench/Cargo.toml b/spatial_bench/Cargo.toml
index 048af85..53de860 100644
--- a/spatial_bench/Cargo.toml
+++ b/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]]
diff --git a/spatial_bench/benches/spatial.rs b/spatial_bench/benches/spatial.rs
index 33ecd2c..16cc191 100644
--- a/spatial_bench/benches/spatial.rs
+++ b/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);
diff --git a/spatial_bench/src/bosque.rs b/spatial_bench/src/bosque.rs
new file mode 100644
index 0000000..d22c1d4
--- /dev/null
+++ b/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()
+    }
+}
diff --git a/spatial_bench/src/kiddo.rs b/spatial_bench/src/kiddo.rs
new file mode 100644
index 0000000..67c8ee4
--- /dev/null
+++ b/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()
+    }
+}
diff --git a/spatial_bench/src/lib.rs b/spatial_bench/src/lib.rs
index 09b4a94..9f11df2 100644
--- a/spatial_bench/src/lib.rs
+++ b/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()
+}
diff --git a/spatial_bench/src/nabo.rs b/spatial_bench/src/nabo.rs
new file mode 100644
index 0000000..87df039
--- /dev/null
+++ b/spatial_bench/src/nabo.rs
@@ -0,0 +1,92 @@
+use crate::{Point3, Precision, SpatialArena};
+use nabo::{KDTree, NotNan, Point};
+
+#[derive(Clone, Copy, Debug, Default)]
+struct NaboPointWithIndex {
+    pub point: [NotNan<Precision>; 3],
+    pub index: usize,
+}
+
+impl NaboPointWithIndex {
+    pub fn new(point: &Point3, index: usize) -> Self {
+        Self {
+            point: [
+                NotNan::new(point[0]).unwrap(),
+                NotNan::new(point[1]).unwrap(),
+                NotNan::new(point[2]).unwrap(),
+            ],
+            index,
+        }
+    }
+}
+
+impl From<NaboPointWithIndex> for Point3 {
+    fn from(p: NaboPointWithIndex) -> Self {
+        [*p.point[0], *p.point[1], *p.point[2]]
+    }
+}
+
+impl Point<Precision> for NaboPointWithIndex {
+    fn get(&self, i: u32) -> nabo::NotNan<Precision> {
+        self.point[i as usize]
+    }
+
+    fn set(&mut self, i: u32, value: nabo::NotNan<Precision>) {
+        self.point[i as usize] = value;
+    }
+
+    const DIM: u32 = 3;
+}
+
+#[derive(Default)]
+pub struct NaboArena {
+    trees: Vec<KDTree<Precision, NaboPointWithIndex>>,
+    points: Vec<Vec<NaboPointWithIndex>>,
+}
+
+impl SpatialArena for NaboArena {
+    fn add_points(&mut self, p: Vec<Point3>) -> usize {
+        let all_p: Vec<_> = p
+            .iter()
+            .enumerate()
+            .map(|(idx, p)| NaboPointWithIndex::new(p, idx))
+            .collect();
+        let t = KDTree::new(&all_p);
+        let idx = self.len();
+        self.trees.push(t);
+        self.points.push(all_p);
+        idx
+    }
+
+    fn query_target(&self, q: usize, t: usize) -> Vec<(usize, Precision)> {
+        let t_tree = self.trees.get(t).unwrap();
+        self.points
+            .get(q)
+            .unwrap()
+            .iter()
+            .map(|p| {
+                let n = t_tree.knn(1, p).pop().unwrap();
+                (n.point.index, n.dist2.sqrt())
+            })
+            .collect()
+    }
+
+    fn local_query(&self, q: usize, neighborhood: usize) -> Vec<Vec<usize>> {
+        let t = self.trees.get(q).unwrap();
+        self.points
+            .get(q)
+            .unwrap()
+            .iter()
+            .map(|p| {
+                t.knn(neighborhood as u32, p)
+                    .into_iter()
+                    .map(|n| n.point.index)
+                    .collect()
+            })
+            .collect()
+    }
+
+    fn len(&self) -> usize {
+        self.trees.len()
+    }
+}
diff --git a/spatial_bench/src/rstar.rs b/spatial_bench/src/rstar.rs
new file mode 100644
index 0000000..31757c3
--- /dev/null
+++ b/spatial_bench/src/rstar.rs
@@ -0,0 +1,55 @@
+use rstar::{primitives::GeomWithData, PointDistance, RTree};
+
+use crate::{Point3, Precision, SpatialArena};
+
+pub type RsPoint = GeomWithData<Point3, usize>;
+
+#[derive(Debug, Default)]
+pub struct RstarArena {
+    trees: Vec<rstar::RTree<RsPoint>>,
+}
+
+impl SpatialArena for RstarArena {
+    fn add_points(&mut self, p: Vec<Point3>) -> usize {
+        let t = RTree::bulk_load(
+            p.into_iter()
+                .enumerate()
+                .map(|(idx, p)| RsPoint::new(p, idx))
+                .collect(),
+        );
+        let idx = self.len();
+        self.trees.push(t);
+        idx
+    }
+
+    fn query_target(&self, q: usize, t: usize) -> Vec<(usize, Precision)> {
+        let t_tree = self.trees.get(t).unwrap();
+
+        self.trees
+            .get(q)
+            .unwrap()
+            .iter()
+            .map(|p| {
+                let neighb = t_tree.nearest_neighbor(p.geom()).unwrap();
+                let dist = p.distance_2(neighb.geom()).sqrt();
+                (p.data, dist)
+            })
+            .collect()
+    }
+
+    fn local_query(&self, q: usize, neighborhood: usize) -> Vec<Vec<usize>> {
+        let tree = self.trees.get(q).unwrap();
+        tree.iter()
+            .map(|p| {
+                tree.nearest_neighbor_iter(p.geom())
+                    .take(neighborhood)
+                    .map(|n| n.data)
+                    .collect()
+            })
+            .collect()
+    }
+
+    fn len(&self) -> usize {
+        self.trees.len()
+    }
+}