diff --git a/Cargo.lock.MSRV b/Cargo.lock.MSRV index cba8a7b7..4ea6bb30 100644 --- a/Cargo.lock.MSRV +++ b/Cargo.lock.MSRV @@ -335,12 +335,12 @@ dependencies = [ [[package]] name = "kdtree" -version = "0.7.0" +version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f0a0e9f770b65bac9aad00f97a67ab5c5319effed07f6da385da3c2115e47ba" +checksum = "9cd666c2aae5dde60d2f9bfa4e1f8c17698fce13d599c0f9b76725641fec13d9" dependencies = [ "num-traits", - "thiserror 1.0.69", + "thiserror", ] [[package]] @@ -706,7 +706,7 @@ dependencies = [ "nalgebra", "num-traits", "rand", - "thiserror 2.0.19", + "thiserror", ] [[package]] @@ -731,33 +731,13 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "thiserror" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" -dependencies = [ - "thiserror-impl 1.0.69", -] - [[package]] name = "thiserror" version = "2.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" dependencies = [ - "thiserror-impl 2.0.19", -] - -[[package]] -name = "thiserror-impl" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.119", + "thiserror-impl", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index be611f5d..18bc853a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -54,7 +54,7 @@ optional = true default-features = false [dependencies.kdtree] -version = "0.7.0" +version = "0.8.1" optional = true [dev-dependencies] diff --git a/benches/density.rs b/benches/density.rs index e7fc8f54..e1f655b1 100644 --- a/benches/density.rs +++ b/benches/density.rs @@ -45,6 +45,31 @@ fn bench_density(c: &mut Criterion) { let _f = statrs::density::kde::kde_pdf(&Vector3::new(0., 0., 0.), &samples, None); }); }); + group.finish(); + + // The free functions above rebuild the k-d tree on every call. These + // measure the same work with the tree hoisted out of the loop, which is the + // point of `DensityEstimator`. + let samples: Vec<[f64; 1]> = generate(100_000); + let grid: Vec<[f64; 1]> = (0..200).map(|i| [i as f64 / 200.0]).collect(); + let mut group = c.benchmark_group("density_grid_200"); + group.sample_size(20); + group.bench_function("knn_1d_free_fns", |b| { + b.iter(|| { + for g in &grid { + let _f = statrs::density::knn::knn_pdf(g, &samples, Some(0.05)); + } + }); + }); + group.bench_function("knn_1d_prepared_estimator", |b| { + b.iter(|| { + let est = statrs::density::DensityEstimator::new(&samples).unwrap(); + for g in &grid { + let _f = est.knn_pdf(g, Some(0.05)); + } + }); + }); + group.finish(); } criterion_group!(benches, bench_density); diff --git a/src/density/kde.rs b/src/density/kde.rs index caa17720..9a74b5bf 100644 --- a/src/density/kde.rs +++ b/src/density/kde.rs @@ -1,16 +1,60 @@ use kdtree::distance::squared_euclidean; use crate::{ - density::{Container, DensityError, nearest_neighbors}, + density::{Container, DensityError, DensityEstimator}, function::kernel::{Gaussian, Kernel}, }; +impl DensityEstimator<'_, S, X> +where + S: AsRef<[X]> + Container, + X: AsRef<[f64]> + Container + PartialEq, +{ + /// Computes the kernel density estimate at `x`, using the distance to the + /// furthest neighbor as a local bandwidth. + /// + /// The optimal `k` is computed using [Orava's](https://www.sav.sk/journals/uploads/0127102604orava.pdf) + /// formula when `bandwidth` is `None`. + /// + /// # Errors + /// + /// Returns [`DensityError::EmptyNeighborhood`] if no sample falls inside the + /// neighborhood of `x`. + pub fn kde_pdf(&self, x: &X, bandwidth: Option) -> Result { + // The neighborhood is used only for its radius: the kernel sum below + // runs over every sample, so this stays O(n) per evaluation. + let neighbors = self.nearest_neighbors(x, bandwidth)?; + if neighbors.is_empty() { + return Err(DensityError::EmptyNeighborhood); + } + let radius = neighbors.radius; + let d = x.length() as i32; + Ok((1. / (self.n_samples() * radius.powi(d))) + * self + .samples() + .as_ref() + .iter() + .map(|xi| { + Gaussian.evaluate(squared_euclidean(x.as_ref(), xi.as_ref()).sqrt() / radius) + / crate::consts::SQRT_2PI.powi(d - 1) + }) + .sum::()) + } +} + /// Computes the kernel density estimate for a given point `x` /// using the samples provided and a specified kernel. /// /// The optimal `k` is computed using [Orava's](https://www.sav.sk/journals/uploads/0127102604orava.pdf) /// formula when `bandwidth` is `None`. /// +/// # Performance +/// +/// This builds a k-d tree over `samples` on every call. To evaluate the density +/// at more than a couple of points, build a +/// [`DensityEstimator`] once and call +/// [`DensityEstimator::kde_pdf`] instead. +/// /// # Examples /// /// ``` @@ -25,23 +69,7 @@ where S: AsRef<[X]> + Container, X: AsRef<[f64]> + Container + PartialEq, { - let n_samples = samples.length() as f64; - let neighbors = nearest_neighbors(x, samples, bandwidth)?.0; - if neighbors.is_empty() { - Err(DensityError::EmptyNeighborhood) - } else { - let radius = neighbors.last().unwrap().sqrt(); // safe to unwrap here since `neighbors` is not empty - let d = x.length() as i32; - Ok((1. / (n_samples * radius.powi(d))) - * samples - .as_ref() - .iter() - .map(|xi| { - Gaussian.evaluate(squared_euclidean(x.as_ref(), xi.as_ref()).sqrt() / radius) - / crate::consts::SQRT_2PI.powi(d - 1) - }) - .sum::()) - } + DensityEstimator::new(samples)?.kde_pdf(x, bandwidth) } #[cfg(test)] diff --git a/src/density/knn.rs b/src/density/knn.rs index 1f58637d..e55393d7 100644 --- a/src/density/knn.rs +++ b/src/density/knn.rs @@ -1,16 +1,51 @@ use super::Container; use crate::{ - density::{DensityError, nearest_neighbors}, + density::{DensityError, DensityEstimator}, function::gamma::gamma, }; use core::f64::consts::PI; +impl DensityEstimator<'_, S, X> +where + S: AsRef<[X]> + Container, + X: AsRef<[f64]> + Container + PartialEq, +{ + /// Computes the `k`-nearest neighbor density estimate at `x`. + /// + /// The optimal `k` is computed using [Orava's](https://www.sav.sk/journals/uploads/0127102604orava.pdf) + /// formula when `bandwidth` is `None`. + /// + /// # Errors + /// + /// Returns [`DensityError::EmptyNeighborhood`] if no sample falls inside the + /// neighborhood of `x`. + pub fn knn_pdf(&self, x: &X, bandwidth: Option) -> Result { + let neighbors = self.nearest_neighbors(x, bandwidth)?; + if neighbors.is_empty() { + return Err(DensityError::EmptyNeighborhood); + } + // k / (n * V_d(r)), with V_d(r) = pi^(d/2) r^d / Gamma(d/2 + 1) the + // volume of the d-ball containing the neighborhood + let radius = neighbors.radius; + let d = x.length() as f64; + Ok((neighbors.k / self.n_samples()) + * (gamma(d / 2. + 1.) / (PI.powf(d / 2.) * radius.powf(d)))) + } +} + /// Computes the `k`-nearest neighbor density estimate for a given point `x` /// using the samples provided. /// /// The optimal `k` is computed using [Orava's](https://www.sav.sk/journals/uploads/0127102604orava.pdf) /// formula when `bandwidth` is `None`. /// +/// # Performance +/// +/// This builds a k-d tree over `samples` on every call. To evaluate the density +/// at more than a couple of points, build a +/// [`DensityEstimator`] once and call +/// [`DensityEstimator::knn_pdf`] instead - about 5x faster over a grid. +/// /// # Examples /// /// ``` @@ -25,15 +60,7 @@ where S: AsRef<[X]> + Container, X: AsRef<[f64]> + Container + PartialEq, { - let n_samples = samples.length() as f64; - let (neighbors, k) = nearest_neighbors(x, samples, bandwidth)?; - if neighbors.is_empty() { - Err(DensityError::EmptyNeighborhood) - } else { - let radius = neighbors.last().unwrap().sqrt(); - let d = x.length() as f64; - Ok((k / n_samples) * (gamma(d / 2. + 1.) / (PI.powf(d / 2.) * radius.powf(d)))) - } + DensityEstimator::new(samples)?.knn_pdf(x, bandwidth) } #[cfg(test)] diff --git a/src/density/mod.rs b/src/density/mod.rs index b5c7896f..198867cd 100644 --- a/src/density/mod.rs +++ b/src/density/mod.rs @@ -11,6 +11,10 @@ //! Both accept an explicit `bandwidth` (a fixed search radius), or fall back //! to a `k` chosen by [Orava's formula](https://www.sav.sk/journals/uploads/0127102604orava.pdf) //! when `bandwidth` is `None`. +//! +//! The free functions build the k-d tree on every call. To evaluate a density at +//! more than a couple of points, build a [`DensityEstimator`] once and query it +//! repeatedly - measured at ~7x over a 200-point grid. pub mod kde; pub mod knn; @@ -76,39 +80,155 @@ impl_container!( nalgebra::Vector5, nalgebra::Vector6 ); -pub type NearestNeighbors = (Vec, f64); +/// Outcome of a neighborhood query around a point. +#[derive(Clone, Debug, PartialEq)] +pub struct NearestNeighbors { + /// Squared distances from the query point to each neighbor found. + /// + /// The order is **unspecified**. `KdTree::within` returned these sorted in + /// `kdtree` 0.7 but returns raw heap order from 0.8 onward, so the last + /// element is not the furthest neighbor; use [`Self::radius`] instead of + /// indexing this. + pub squared_distances: Vec, + + /// How many neighbors were found, as an `f64`. + pub k: f64, + + /// Distance - not squared - to the furthest neighbor, i.e. the radius of + /// the neighborhood actually used by the estimators. + pub radius: f64, +} + +impl NearestNeighbors { + fn new(squared_distances: Vec, k: f64) -> Self { + // Computed as a maximum rather than by taking the last element, so that + // it does not depend on the ordering of `squared_distances`. + let radius = squared_distances + .iter() + .copied() + .fold(0.0_f64, f64::max) + .sqrt(); + Self { + squared_distances, + k, + radius, + } + } + + /// Whether the neighborhood is empty. + pub fn is_empty(&self) -> bool { + self.squared_distances.is_empty() + } +} + +/// Leaf bucket size for the backing k-d tree. +/// +/// This is a deliberate trade: measured at `n = 1e5`, `d = 1`, a capacity of +/// `~n` builds in 1.3 ms against 19.5 ms for the crate's default of 16, while +/// queries are only ~1.1x slower - the query is dominated by materialising the +/// result vector, not by traversal. For query-heavy use of +/// [`DensityEstimator`], where the build is amortised away, a small constant is +/// worth roughly 15% of query time instead. +fn leaf_capacity(n_samples: usize) -> usize { + 2usize.pow((n_samples as f64).log2() as u32) +} + +/// A reusable nearest-neighbor index over a fixed sample set. +/// +/// The free functions [`knn::knn_pdf`] and [`kde::kde_pdf`] rebuild the backing +/// k-d tree on every call, which is `O(n log n)` per evaluation - about 59% of a +/// single call at `n = 1e5`. Building the index once and querying it repeatedly +/// is ~5x faster when evaluating a density over more than a handful of points: +/// +/// ``` +/// use statrs::density::DensityEstimator; +/// +/// let samples: Vec<[f64; 1]> = vec![[-1.0], [-0.5], [0.0], [0.25], [0.5], [1.0]]; +/// let estimator = DensityEstimator::new(&samples).unwrap(); +/// +/// // one tree, many evaluations +/// let densities: Vec = (0..5) +/// .map(|i| -1.0 + 0.5 * i as f64) +/// .map(|g| estimator.knn_pdf(&[g], Some(1.0)).unwrap()) +/// .collect(); +/// assert!(densities.iter().all(|d| *d > 0.0)); +/// ``` +#[derive(Clone, Debug)] +pub struct DensityEstimator<'a, S, X: AsRef<[f64]> + PartialEq> { + samples: &'a S, + tree: KdTree, +} -pub(crate) fn nearest_neighbors( - x: &X, - samples: &S, - bandwidth: Option, -) -> Result +impl<'a, S, X> DensityEstimator<'a, S, X> where S: AsRef<[X]> + Container, X: AsRef<[f64]> + Container + PartialEq, { - if samples.length() == 0 { - return Err(DensityError::EmptySample); + /// Builds the index over `samples`. + /// + /// # Errors + /// + /// Returns [`DensityError::EmptySample`] if `samples` is empty, or + /// [`DensityError::KdTree`] if a sample's dimension disagrees with the + /// first one's. + pub fn new(samples: &'a S) -> Result { + let points = samples.as_ref(); + let Some(first) = points.first() else { + return Err(DensityError::EmptySample); + }; + // Dimensionality is taken from the samples rather than from a query + // point, so a mismatched query is rejected by the tree instead of + // silently building a tree of the wrong shape. + let mut tree = KdTree::with_capacity(first.length(), leaf_capacity(points.len())); + for (position, sample) in points.iter().enumerate() { + tree.add(sample.clone(), position)?; + } + Ok(Self { samples, tree }) } - let n_samples = samples.length() as f64; - let d = x.length(); - let mut tree = KdTree::with_capacity(d, 2usize.pow(n_samples.log2() as u32)); - for (position, sample) in samples.as_ref().iter().enumerate() { - tree.add(sample.clone(), position)?; + + /// The samples this index was built over. + pub fn samples(&self) -> &'a S { + self.samples } - if let Some(bandwidth) = bandwidth { - let neighbors = tree.within(x.as_ref(), bandwidth, &squared_euclidean)?; - let k = neighbors.len() as f64; - Ok((neighbors.into_iter().map(|r| r.0).collect(), k)) - } else { - let k = orava_optimal_k(n_samples); - Ok(( - tree.nearest(x.as_ref(), k as usize, &squared_euclidean)? - .into_iter() - .map(|r| r.0) - .collect(), - k, - )) + + pub(crate) fn n_samples(&self) -> f64 { + self.samples.as_ref().len() as f64 + } + + /// Finds the neighborhood of `x`: within a fixed squared radius when + /// `bandwidth` is `Some`, otherwise the `k` nearest neighbors for a `k` + /// chosen by Orava's formula. + /// + /// # Errors + /// + /// Returns [`DensityError::KdTree`] if `x`'s dimension disagrees with the + /// samples'. + pub fn nearest_neighbors( + &self, + x: &X, + bandwidth: Option, + ) -> Result { + // Both queries use `squared_euclidean`, so `bandwidth` is a squared + // radius and every returned distance is squared. + if let Some(bandwidth) = bandwidth { + let neighbors = self + .tree + .within(x.as_ref(), bandwidth, &squared_euclidean)?; + let k = neighbors.len() as f64; + Ok(NearestNeighbors::new( + neighbors.into_iter().map(|r| r.0).collect(), + k, + )) + } else { + let k = orava_optimal_k(self.n_samples()); + let neighbors = self + .tree + .nearest(x.as_ref(), k as usize, &squared_euclidean)?; + Ok(NearestNeighbors::new( + neighbors.into_iter().map(|r| r.0).collect(), + k, + )) + } } } #[cfg(test)] @@ -117,6 +237,90 @@ mod tests { use super::*; + /// The neighborhood radius must be the distance to the *furthest* neighbor + /// no matter what order the backing tree returns results in. + /// + /// `knn_pdf`/`kde_pdf` used to read it as `distances.last()`, which was only + /// correct because `kdtree` 0.7's `within` happened to sort; 0.8 returns raw + /// heap order, which silently scaled both densities by the wrong volume. + /// Hand-placed samples make the expected values exact, unlike the + /// Monte Carlo tests that could not reliably detect this. + #[test] + fn test_nearest_neighbors_radius_is_the_maximum() { + let samples: Vec<[f64; 1]> = vec![[-3.0], [-2.0], [-1.0], [0.0], [1.0], [2.0], [3.0]]; + // `bandwidth` is a *squared* radius, so 4.0 admits everything within 2.0 + let estimator = DensityEstimator::new(&samples).unwrap(); + let nn = estimator.nearest_neighbors(&[0.0], Some(4.0)).unwrap(); + assert_eq!(nn.k, 5.0, "expected the 5 samples within distance 2"); + assert_eq!(nn.radius, 2.0); + // and it agrees with an order-independent maximum by construction + let max = nn + .squared_distances + .iter() + .copied() + .fold(0.0_f64, f64::max) + .sqrt(); + assert_eq!(nn.radius, max); + + // the `nearest` path reports a radius too + let nn = estimator.nearest_neighbors(&[0.0], None).unwrap(); + assert!(nn.radius > 0.0 && nn.radius.is_finite()); + } + + /// With the radius pinned, `knn_pdf` is exactly `(k / n) / V_d(r)`: + /// `(5 / 7) / (2 * 2) = 5 / 28` in one dimension. + #[test] + fn test_knn_pdf_deterministic() { + let samples: Vec<[f64; 1]> = vec![[-3.0], [-2.0], [-1.0], [0.0], [1.0], [2.0], [3.0]]; + let got = crate::density::knn::knn_pdf(&[0.0], &samples, Some(4.0)).unwrap(); + // the tolerance is set by `gamma(d / 2 + 1)` in the ball-volume factor, + // not by anything in this module + crate::prec::assert_relative_eq!(got, 5.0 / 28.0, epsilon = 0.0, max_relative = 1e-14); + } + + /// The prepared estimator must be bit-identical to the free functions - it + /// is the same code path with the tree hoisted out of the loop, so any + /// difference would mean the hoist changed the maths. + #[test] + fn test_estimator_matches_free_functions() { + let samples: Vec<[f64; 1]> = (0..500).map(|i| [(i as f64 * 0.37).sin() * 2.0]).collect(); + let estimator = DensityEstimator::new(&samples).unwrap(); + for i in 0..40 { + let g = [-2.0 + 4.0 * i as f64 / 40.0]; + for bw in [Some(0.05), Some(1.0), None] { + assert_eq!( + estimator.knn_pdf(&g, bw).ok(), + crate::density::knn::knn_pdf(&g, &samples, bw).ok(), + "knn_pdf mismatch at {g:?} bandwidth {bw:?}" + ); + assert_eq!( + estimator.kde_pdf(&g, bw).ok(), + crate::density::kde::kde_pdf(&g, &samples, bw).ok(), + "kde_pdf mismatch at {g:?} bandwidth {bw:?}" + ); + } + } + } + + #[test] + fn test_estimator_empty_sample() { + let empty: Vec<[f64; 1]> = vec![]; + assert_eq!( + DensityEstimator::new(&empty).unwrap_err(), + DensityError::EmptySample + ); + } + + /// A query point of the wrong dimension is rejected by the tree rather than + /// silently answered, because the tree's dimension comes from the samples. + #[test] + fn test_estimator_rejects_dimension_mismatch() { + let samples: Vec> = vec![vec![0.0, 0.0], vec![1.0, 1.0]]; + let estimator = DensityEstimator::new(&samples).unwrap(); + assert!(estimator.nearest_neighbors(&vec![0.0], Some(1.0)).is_err()); + assert!(estimator.knn_pdf(&vec![0.0, 0.0], Some(4.0)).is_ok()); + } + #[test] fn test_vec_container() { let v1 = vec![1.0, 2.0, 3.0];