diff --git a/Cargo.lock.MSRV b/Cargo.lock.MSRV index 9d5ddc56..576dcd57 100644 --- a/Cargo.lock.MSRV +++ b/Cargo.lock.MSRV @@ -333,6 +333,16 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "kdtree" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f0a0e9f770b65bac9aad00f97a67ab5c5319effed07f6da385da3c2115e47ba" +dependencies = [ + "num-traits", + "thiserror 1.0.69", +] + [[package]] name = "libc" version = "0.2.186" @@ -638,7 +648,7 @@ checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn 3.0.0", + "syn 3.0.1", ] [[package]] @@ -685,9 +695,11 @@ dependencies = [ "anyhow", "approx", "criterion", + "kdtree", "nalgebra", "num-traits", "rand", + "thiserror 2.0.19", ] [[package]] @@ -703,15 +715,55 @@ dependencies = [ [[package]] name = "syn" -version = "3.0.0" +version = "3.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2fac314a64dc9a36e61a9eb4261a5e9bbfbc922b27e518af97bc32b926cf967" +checksum = "5edbec4ed188954a10c12c038215f8ce7606b2d5c973cd8dc43e8795065c5f2f" dependencies = [ "proc-macro2", "quote", "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", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.1", +] + [[package]] name = "tinytemplate" version = "1.2.1" diff --git a/Cargo.toml b/Cargo.toml index 950afb4c..ed8d0133 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,16 +24,24 @@ name = "order_statistics" harness = false required-features = ["rand", "std"] +[[bench]] +name = "density" +harness = false +required-features = ["rand", "std", "kde"] + [features] default = ["std", "nalgebra", "rand"] std = ["nalgebra?/std", "rand?/std"] # at the moment, all nalgebra features needs std nalgebra = ["dep:nalgebra", "std"] rand = ["dep:rand", "nalgebra?/rand-no-std", "rand?/std_rng"] +# kd-tree backed density estimation (src/density), implemented in terms of nalgebra vectors +kde = ["dep:kdtree", "nalgebra"] [dependencies] approx = "0.5.0" num-traits = "0.2.14" +thiserror = { version = "2.0", default-features = false } [dependencies.rand] version = "0.10.0" @@ -45,6 +53,10 @@ version = "0.35" optional = true default-features = false +[dependencies.kdtree] +version = "0.7.0" +optional = true + [dev-dependencies] criterion = "0.8" anyhow = "1.0" diff --git a/benches/density.rs b/benches/density.rs new file mode 100644 index 00000000..e7fc8f54 --- /dev/null +++ b/benches/density.rs @@ -0,0 +1,52 @@ +extern crate criterion; +extern crate rand; +extern crate statrs; +use criterion::{Criterion, criterion_group, criterion_main}; +use nalgebra::{Vector1, Vector3}; +use rand::RngExt; +use rand::SeedableRng; +use rand::distr::StandardUniform; +use rand::rngs::StdRng; + +fn generate(n_samples: usize) -> Vec +where + StandardUniform: rand::distr::Distribution, +{ + let mut rng = StdRng::seed_from_u64(42); + (0..n_samples).map(|_| rng.random()).collect() +} + +fn bench_density(c: &mut Criterion) { + let samples = generate(100_000); + let mut group = c.benchmark_group("density"); + group.bench_function("knn_density_1d", |b| { + b.iter(|| { + let _f = statrs::density::knn::knn_pdf(&[0.], &samples, None); + }); + }); + + let samples = generate(100_000); + group.bench_function("knn_density_3d", |b| { + b.iter(|| { + let _f = statrs::density::knn::knn_pdf(&[0., 0., 0.], &samples, None); + }); + }); + + let samples = generate(100_000); + group.bench_function("kde_density_1d", |b| { + b.iter(|| { + let _f = statrs::density::kde::kde_pdf(&Vector1::new(0.), &samples, None); + }); + }); + + let samples = generate(100_000); + group.bench_function("kde_density_3d", |b| { + b.iter(|| { + let _f = statrs::density::kde::kde_pdf(&Vector3::new(0., 0., 0.), &samples, None); + }); + }); +} + +criterion_group!(benches, bench_density); + +criterion_main!(benches); diff --git a/src/density/kde.rs b/src/density/kde.rs new file mode 100644 index 00000000..caa17720 --- /dev/null +++ b/src/density/kde.rs @@ -0,0 +1,89 @@ +use kdtree::distance::squared_euclidean; + +use crate::{ + density::{Container, DensityError, nearest_neighbors}, + function::kernel::{Gaussian, Kernel}, +}; + +/// 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`. +/// +/// # Examples +/// +/// ``` +/// use statrs::density::kde::kde_pdf; +/// +/// let samples: Vec<[f64; 1]> = vec![[-1.0], [0.0], [1.0]]; +/// let density = kde_pdf(&[0.0], &samples, Some(1.0)).unwrap(); +/// assert!(density > 0.0); +/// ``` +pub fn kde_pdf(x: &X, samples: &S, bandwidth: Option) -> Result +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::()) + } +} + +#[cfg(test)] +mod tests { + use core::f32::consts::PI; + + use super::*; + use crate::distribution::Normal; + use crate::function::kernel::Kernel; + use nalgebra::{Vector1, Vector2}; + use rand::SeedableRng; + use rand::distr::Distribution; + use rand::rngs::StdRng; + + #[test] + fn test_kde_pdf() { + let law = Normal::new(0., 1.).unwrap(); + let mut rng = StdRng::seed_from_u64(42); + let gaussian = crate::function::kernel::Gaussian; + let samples_1d = (0..100000) + .map(|_| Vector1::new(law.sample(&mut rng))) + .collect::>(); + let x = Vector1::new(0.); + let kde_density_with_bandwidth = kde_pdf(&x, &samples_1d, Some(0.05)); + let kde_density = kde_pdf(&x, &samples_1d, None); + let reference_value = gaussian.evaluate(0.); + assert!(kde_density.is_ok()); + assert!(kde_density_with_bandwidth.is_ok()); + assert!((kde_density.unwrap() - reference_value).abs() < 2e-2); + assert!((kde_density_with_bandwidth.unwrap() - reference_value).abs() < 3e-2); + + let samples_2d = (0..100000) + .map(|_| Vector2::new(law.sample(&mut rng), law.sample(&mut rng))) + .collect::>(); + + let x = Vector2::new(0., 0.); + let kde_density_with_bandwidth = kde_pdf(&x, &samples_2d, Some(0.05)); + let kde_density = kde_pdf(&x, &samples_2d, None); + let reference_value = 1. / (2. * PI) as f64; + assert!(kde_density.is_ok()); + assert!(kde_density_with_bandwidth.is_ok()); + assert!((kde_density.unwrap() - reference_value).abs() < 2e-2); + assert!((kde_density_with_bandwidth.unwrap() - reference_value).abs() < 3e-2); + } +} diff --git a/src/density/knn.rs b/src/density/knn.rs new file mode 100644 index 00000000..1f58637d --- /dev/null +++ b/src/density/knn.rs @@ -0,0 +1,89 @@ +use super::Container; +use crate::{ + density::{DensityError, nearest_neighbors}, + function::gamma::gamma, +}; +use core::f64::consts::PI; + +/// 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`. +/// +/// # Examples +/// +/// ``` +/// use statrs::density::knn::knn_pdf; +/// +/// let samples: Vec<[f64; 1]> = vec![[-1.0], [0.0], [1.0]]; +/// let density = knn_pdf(&[0.0], &samples, Some(1.0)).unwrap(); +/// assert!(density > 0.0); +/// ``` +pub fn knn_pdf(x: &X, samples: &S, bandwidth: Option) -> Result +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)))) + } +} + +#[cfg(test)] +mod tests { + use core::f32::consts::PI; + + use super::*; + use crate::distribution::Normal; + use crate::function::kernel::Kernel; + use nalgebra::{Vector1, Vector2}; + use rand::SeedableRng; + use rand::distr::Distribution; + use rand::rngs::StdRng; + + #[test] + fn test_knn_pdf() { + let law = Normal::new(0., 1.).unwrap(); + let mut rng = StdRng::seed_from_u64(42); + let gaussian = crate::function::kernel::Gaussian; + let samples_1d = (0..100000) + .map(|_| Vector1::new(law.sample(&mut rng))) + .collect::>(); + let x = Vector1::new(0.); + let knn_density_with_bandwidth = knn_pdf(&x, &samples_1d, Some(0.05)); + let knn_density = knn_pdf(&x, &samples_1d, None); + let reference_value = gaussian.evaluate(0.); + assert!(knn_density.is_ok()); + assert!(knn_density_with_bandwidth.is_ok()); + assert!((knn_density.unwrap() - reference_value).abs() < 2e-2); + assert!((knn_density_with_bandwidth.unwrap() - reference_value).abs() < 3e-2); + + let samples_2d = (0..100000) + .map(|_| Vector2::new(law.sample(&mut rng), law.sample(&mut rng))) + .collect::>(); + + let x = Vector2::new(0., 0.); + let knn_density_with_bandwidth = knn_pdf(&x, &samples_2d, Some(0.05)); + let knn_density = knn_pdf(&x, &samples_2d, None); + let reference_value = 1. / (2. * PI) as f64; + assert!(knn_density.is_ok()); + assert!(knn_density_with_bandwidth.is_ok()); + assert!((knn_density.unwrap() - reference_value).abs() < 2e-2); + assert!((knn_density_with_bandwidth.unwrap() - reference_value).abs() < 3e-2); + } + + #[test] + fn test_knn_pdf_empty_samples() { + let samples: Vec<[f64; 1]> = vec![]; + let x = 3.0; + let result = knn_pdf(&[x], &samples, None); + assert!(matches!(result, Err(DensityError::EmptySample))); + } +} diff --git a/src/density/mod.rs b/src/density/mod.rs new file mode 100644 index 00000000..b5c7896f --- /dev/null +++ b/src/density/mod.rs @@ -0,0 +1,127 @@ +//! Nearest-neighbor [density estimation](https://en.wikipedia.org/wiki/Multivariate_kernel_density_estimation) +//! for samples in R^d, backed by a k-d tree for neighbor search. +//! +//! Two estimators are provided, differing in how they turn a neighborhood +//! into a density: +//! - [`knn::knn_pdf`] uses the distance to the `k`-th nearest neighbor directly. +//! - [`kde::kde_pdf`] additionally weights every sample in that neighborhood +//! by a Gaussian kernel, using the `k`-th neighbor's distance as a local +//! bandwidth. +//! +//! 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`. + +pub mod kde; +pub mod knn; +use kdtree::{ErrorKind, KdTree, distance::squared_euclidean}; +use thiserror::Error; + +/// Errors that can occur when estimating a density from a sample. +#[derive(Error, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub enum DensityError { + /// The k-d tree backing the nearest-neighbor search could not be built or queried. + #[error("K-d tree error: {0}")] + KdTree(ErrorKind), + + /// The sample provided was empty, so no density can be estimated. + #[error("No samples provided")] + EmptySample, + + /// No sample points fell within the queried neighborhood. + #[error("No neighbors found")] + EmptyNeighborhood, +} + +impl From for DensityError { + fn from(err: ErrorKind) -> Self { + DensityError::KdTree(err) + } +} + +fn orava_optimal_k(n_samples: f64) -> f64 { + // Adapted from K-nearest neighbour kernel density estimation, the choice of optimal k; Jan Orava 2012 + (0.587 * n_samples.powf(4.0 / 5.0)).round().max(1.) +} + +/// Handles variable/point types for which nearest neighbors can be computed. +pub trait Container: Clone { + type Elem; + fn length(&self) -> usize; +} + +macro_rules! impl_container { + ($($t:ty),*) => { + $( + impl Container for $t { + type Elem = T; + fn length(&self) -> usize { + self.len() + } + + } + )* + }; +} +impl_container!( + [T; 1], + [T; 2], + [T; 3], + Vec, + nalgebra::Vector1, + nalgebra::Vector2, + nalgebra::Vector3, + nalgebra::Vector4, + nalgebra::Vector5, + nalgebra::Vector6 +); +pub type NearestNeighbors = (Vec, f64); + +pub(crate) fn nearest_neighbors( + x: &X, + samples: &S, + bandwidth: Option, +) -> Result +where + S: AsRef<[X]> + Container, + X: AsRef<[f64]> + Container + PartialEq, +{ + if samples.length() == 0 { + return Err(DensityError::EmptySample); + } + 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)?; + } + 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, + )) + } +} +#[cfg(test)] +mod tests { + use nalgebra::Vector3; + + use super::*; + + #[test] + fn test_vec_container() { + let v1 = vec![1.0, 2.0, 3.0]; + assert_eq!(v1.length(), 3); + let v2 = Vector3::new(1.0, 2.0, 3.0); + assert_eq!(v2.length(), 3); + } +} diff --git a/src/lib.rs b/src/lib.rs index ddd5a30b..8307ed18 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -55,6 +55,9 @@ #![cfg_attr(not(feature = "std"), no_std)] pub mod consts; +#[cfg(feature = "kde")] +#[cfg_attr(docsrs, doc(cfg(feature = "kde")))] +pub mod density; pub mod distribution; pub mod euclid; pub mod function;