Spatial bench

Cargo.toml

@@ -4,7 +4,10 @@ members = [
   "nblast-py",
   "nblast-rs",
   "nblast-js",
+  "spatial_bench",
 ]
 
+resolver = "2"
+
 # [profile.release]
 # debug = true

nblast-js/Cargo.toml

@@ -13,7 +13,7 @@ serde = {version = "1", features = ["derive"]}
 serde-wasm-bindgen = "0.5"
 wasm-bindgen = "0.2"
 js-sys = "0.3"
-nblast = { path = "../nblast-rs", version = "^0.5.0"}
+nblast = { path = "../nblast-rs", version = "^0.5.0", default-features = false, features = ["rstar"]}
 
 [lib]
 crate-type = ["cdylib"]

nblast-py/Cargo.toml

@@ -9,7 +9,7 @@ edition = "2018"
 [dependencies]
 pyo3 = { version = "0.18.2", features = ["extension-module"] }
 neurarbor = "0.2.0"
-nblast = { path = "../nblast-rs", version = "^0.5.0", features = ["parallel"] }
+nblast = { path = "../nblast-rs", version = "^0.5.0", features = ["parallel", "kiddo"] }
 numpy = "0.18"
 
 [lib]

nblast-py/src/lib.rs

@@ -12,8 +12,8 @@ use numpy::{IntoPyArray, PyArray1, PyArray2, PyReadonlyArray1, PyReadonlyArray2}
 
 use nblast::nalgebra::base::{Unit, Vector3};
 use nblast::{
-    BinLookup, NblastArena, Neuron, NeuronIdx, Precision, RangeTable, ScoreCalc,
-    ScoreMatrixBuilder, Symmetry, TangentAlpha,
+    neurons::kiddo::ExactKiddoTangentsAlphas as Neuron, BinLookup, NblastArena, NeuronIdx,
+    Precision, RangeTable, ScoreCalc, ScoreMatrixBuilder, Symmetry, TangentAlpha,
 };
 
 use nblast::rayon;
@@ -67,10 +67,11 @@ impl ArenaWrapper {
                 .as_array()
                 .rows()
                 .into_iter()
-                .map(|r| [r[0], r[1], r[2]]),
+                .map(|r| [r[0], r[1], r[2]])
+                .collect(),
             self.k,
         )
-        .map_err(PyErr::new::<exceptions::PyRuntimeError, _>)?;
+        .map_err(|e| PyErr::new::<PyValueError, _>(e))?;
         Ok(self.arena.add_neuron(neuron))
     }
 
@@ -109,10 +110,11 @@ impl ArenaWrapper {
                 .as_array()
                 .rows()
                 .into_iter()
-                .map(|r| [r[0], r[1], r[2]]),
+                .map(|r| [r[0], r[1], r[2]])
+                .collect(),
             tangents_alphas,
         )
-        .map_err(PyErr::new::<exceptions::PyRuntimeError, _>)?;
+        .map_err(|e| PyErr::new::<PyValueError, _>(e))?;
         Ok(self.arena.add_neuron(neuron))
     }
 
@@ -356,17 +358,17 @@ fn make_neurons_many(
             points_list
                 .into_par_iter()
                 .map(|ps| {
-                    Neuron::new(ps.into_iter().map(|p| [p[0], p[1], p[2]]), k)
-                        .expect("failed to construct neuron") // todo: error handling
+                    Neuron::new(ps.into_iter().map(|p| [p[0], p[1], p[2]]).collect(), k)
+                        .expect("Invalid neuron")
                 })
                 .collect()
         })
     } else {
         points_list
             .into_iter()
             .map(|ps| {
-                Neuron::new(ps.into_iter().map(|p| [p[0], p[1], p[2]]), k)
-                    .expect("failed to construct neuron") // todo: error handling
+                Neuron::new(ps.into_iter().map(|p| [p[0], p[1], p[2]]).collect(), k)
+                    .expect("invalid neuron")
             })
             .collect()
     }

nblast-rs/Cargo.toml

@@ -19,12 +19,13 @@ categories = ["algorithms", "science"]
 [dependencies]
 
 nalgebra = "0.31"
-rstar = "0.9"
+rstar = {version = "0.9", optional = true}
 fastrand = "1.9"
 rayon = { version = "1.5", optional = true }
 thiserror = "1.0"
 nabo = { version = "0.2.1", optional = true }
-# fnntw = "0.2"
+kiddo = { version = "4.0", optional = true }
+cfg-if = "1.0.0"
 
 [dev-dependencies]
 
@@ -33,6 +34,7 @@ csv = "1.1"
 serde = { version = "1", features = ["derive"] }
 
 [features]
+default = ["kiddo"]
 
 parallel = ["rayon"]
 

nblast-rs/benches/bench.rs

@@ -6,6 +6,8 @@ use bencher::{benchmark_group, benchmark_main, Bencher};
 use csv::ReaderBuilder;
 use fastrand::Rng;
 
+#[cfg(feature = "kiddo")]
+use nblast::neurons::kiddo::{ExactKiddoTangentsAlphas, KiddoTangentsAlphas};
 #[cfg(feature = "nabo")]
 use nblast::neurons::nabo::NaboTangentsAlphas;
 use nblast::neurons::rstar::RStarTangentsAlphas;
@@ -264,6 +266,16 @@ fn bench_construction_nabo(b: &mut Bencher) {
     b.iter(|| NaboTangentsAlphas::new(points.clone(), N_NEIGHBORS))
 }
 
+fn bench_construction_kiddo(b: &mut Bencher) {
+    let points = read_points(NAMES[0]);
+    b.iter(|| KiddoTangentsAlphas::new(points.clone(), N_NEIGHBORS))
+}
+
+fn bench_construction_exact_kiddo(b: &mut Bencher) {
+    let points = read_points(NAMES[0]);
+    b.iter(|| ExactKiddoTangentsAlphas::new(points.clone(), N_NEIGHBORS))
+}
+
 fn bench_query_nabo(b: &mut Bencher) {
     let score_fn = get_score_fn();
     let query = NaboTangentsAlphas::new(read_points(NAMES[0]), N_NEIGHBORS);
@@ -272,6 +284,22 @@ fn bench_query_nabo(b: &mut Bencher) {
     b.iter(|| query.query(&target, false, &score_fn))
 }
 
+fn bench_query_kiddo(b: &mut Bencher) {
+    let score_fn = get_score_fn();
+    let query = KiddoTangentsAlphas::new(read_points(NAMES[0]), N_NEIGHBORS).unwrap();
+    let target = KiddoTangentsAlphas::new(read_points(NAMES[1]), N_NEIGHBORS).unwrap();
+
+    b.iter(|| query.query(&target, false, &score_fn))
+}
+
+fn bench_query_exact_kiddo(b: &mut Bencher) {
+    let score_fn = get_score_fn();
+    let query = ExactKiddoTangentsAlphas::new(read_points(NAMES[0]), N_NEIGHBORS).unwrap();
+    let target = ExactKiddoTangentsAlphas::new(read_points(NAMES[1]), N_NEIGHBORS).unwrap();
+
+    b.iter(|| query.query(&target, false, &score_fn))
+}
+
 fn bench_construction_with_tangents_rstar(b: &mut Bencher) {
     let points = read_points(NAMES[0]);
     let neuron = RStarTangentsAlphas::new(&points, N_NEIGHBORS).expect("couldn't parse");
@@ -375,6 +403,28 @@ fn bench_all_to_all_serial_nabo(b: &mut Bencher) {
     b.iter(|| arena.queries_targets(&idxs, &idxs, false, &None, None));
 }
 
+fn bench_all_to_all_serial_kiddo(b: &mut Bencher) {
+    let mut arena = NblastArena::new(get_score_fn(), false);
+    let mut idxs = Vec::new();
+    for name in NAMES.iter() {
+        let points = read_points(name);
+        idxs.push(arena.add_neuron(KiddoTangentsAlphas::new(points, N_NEIGHBORS).unwrap()));
+    }
+
+    b.iter(|| arena.queries_targets(&idxs, &idxs, false, &None, None));
+}
+
+fn bench_all_to_all_serial_exact_kiddo(b: &mut Bencher) {
+    let mut arena = NblastArena::new(get_score_fn(), false);
+    let mut idxs = Vec::new();
+    for name in NAMES.iter() {
+        let points = read_points(name);
+        idxs.push(arena.add_neuron(ExactKiddoTangentsAlphas::new(points, N_NEIGHBORS).unwrap()));
+    }
+
+    b.iter(|| arena.queries_targets(&idxs, &idxs, false, &None, None));
+}
+
 #[cfg(feature = "parallel")]
 fn bench_all_to_all_parallel(b: &mut Bencher) {
     let mut arena = NblastArena::new(get_score_fn(), false).with_threads(0);
@@ -455,6 +505,22 @@ benchmark_group!(
     bench_all_to_all_serial_nabo,
 );
 
+#[cfg(feature = "kiddo")]
+benchmark_group!(
+    impl_kiddo,
+    bench_construction_kiddo,
+    bench_query_kiddo,
+    bench_all_to_all_serial_kiddo,
+);
+
+#[cfg(feature = "kiddo")]
+benchmark_group!(
+    impl_exact_kiddo,
+    bench_construction_exact_kiddo,
+    bench_query_exact_kiddo,
+    bench_all_to_all_serial_exact_kiddo,
+);
+
 benchmark_group!(
     arena,
     bench_arena_query,
@@ -472,4 +538,4 @@ benchmark_group!(
     bench_smatbuild_rstar_quantiles,
 );
 
-benchmark_main!(impl_rstar, impl_nabo, arena, smat);
+benchmark_main!(impl_rstar, impl_nabo, impl_kiddo, arena, smat);

nblast-rs/src/lib.rs

@@ -61,7 +61,7 @@
 //! ```
 //! use nblast::{NblastArena, ScoreCalc, Neuron, Symmetry};
 //!
-//! // create a lookup table for the point match scores
+//! // Create a lookup table for the point match scores
 //! let smat = ScoreCalc::table_from_bins(
 //!   vec![0.0, 0.1, 0.25, 0.5, 1.0, 5.0, f64::INFINITY], // distance thresholds
 //!   vec![0.0, 0.2, 0.4, 0.6, 0.8, 1.0], // dot product thresholds
@@ -75,7 +75,9 @@
 //!   ],
 //! ).expect("could not build score matrix");
 //!
-//! // create an arena to hold your neurons with this score function,
+//! // See the ScoreMatrixBuilder for constructing a score matrix from test data.
+//!
+//! // Create an arena to hold your neurons with this score function,
 //! // whether it should scale the dot products by the colinearity value,
 //! // and how many threads to use (default serial)
 //! let mut arena = NblastArena::new(smat, false).with_threads(2);
@@ -91,7 +93,7 @@
 //! }
 //!
 
-//! // add some neurons built from points and a neighborhood size,
+//! // Add some neurons built from points and a neighborhood size,
 //! // returning their indices in the arena
 //! let idx1 = arena.add_neuron(
 //!     Neuron::new(random_points(6, &mut rng), 5).expect("cannot construct neuron")
@@ -125,9 +127,11 @@ mod table_lookup;
 pub use table_lookup::{BinLookup, NdBinLookup, RangeTable};
 
 pub mod neurons;
-use neurons::rstar::points_to_rtree_tangents_alphas;
 pub use neurons::{NblastNeuron, Neuron, QueryNeuron, TargetNeuron};
 
+#[cfg(not(any(feature = "nabo", feature = "rstar", feature = "kiddo")))]
+compile_error!("one of 'nabo', 'rstar', or 'kiddo' features must be enabled");
+
 /// Floating point precision type used internally
 pub type Precision = f64;
 /// 3D point type used internally
@@ -327,17 +331,11 @@ pub struct PointsTangentsAlphas {
 }
 
 impl PointsTangentsAlphas {
-    /// Calculates tangents from the given points.
-    /// Note that this constructs a spatial index in order to calculate the tangents,
-    /// and then throws it away: you may as well use a [TargetNeuron](trait.TargetNeuron.html)
-    /// type, with regards to performance.
-    /// `k` is the number of points tangents will be calculated with,
-    /// and includes the point itself.
-    pub fn new(points: Vec<Point3>, k: usize) -> Result<Self, &'static str> {
-        points_to_rtree_tangents_alphas(points.iter(), k).map(|(_, tangents_alphas)| Self {
+    pub fn new(points: Vec<Point3>, tangents_alphas: Vec<TangentAlpha>) -> Self {
+        Self {
             points,
             tangents_alphas,
-        })
+        }
     }
 }
 
@@ -884,8 +882,7 @@ mod test {
     #[test]
     fn construct() {
         let points = make_points(&[0., 0., 0.], &[1., 0., 0.], 10);
-        PointsTangentsAlphas::new(points.clone(), N_NEIGHBORS).expect("Query construction failed");
-        Neuron::new(&points, N_NEIGHBORS).expect("Target construction failed");
+        Neuron::new(points, N_NEIGHBORS);
     }
 
     fn is_close(val1: Precision, val2: Precision) -> bool {
@@ -1092,13 +1089,10 @@ mod test {
             RangeTable::new_from_bins(vec![dist_thresholds, dot_thresholds], cells).unwrap(),
         );
 
-        let query = Neuron::new(&make_points(&[0., 0., 0.], &[1., 0., 0.], 10), N_NEIGHBORS)
-            .expect("Construction failed");
-        let target = Neuron::new(
-            &make_points(&[0.5, 0., 0.], &[1.1, 0., 0.], 10),
-            N_NEIGHBORS,
-        )
-        .expect("Construction failed");
+        let query =
+            Neuron::new(make_points(&[0., 0., 0.], &[1., 0., 0.], 10), N_NEIGHBORS).unwrap();
+        let target =
+            Neuron::new(make_points(&[0.5, 0., 0.], &[1.1, 0., 0.], 10), N_NEIGHBORS).unwrap();
 
         let mut arena = NblastArena::new(score_calc, false);
         let q_idx = arena.add_neuron(query);