feat (knn/kde) density estimation

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"

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<T>(n_samples: usize) -> Vec<T>
+where
+    StandardUniform: rand::distr::Distribution<T>,
+{
+    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);

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<S, X>(x: &X, samples: &S, bandwidth: Option<f64>) -> Result<f64, DensityError>
+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::<f64>())
+    }
+}
+
+#[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::<Vec<_>>();
+        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::<Vec<_>>();
+
+        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);
+    }
+}

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, S>(x: &X, samples: &S, bandwidth: Option<f64>) -> Result<f64, DensityError>
+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::<Vec<_>>();
+        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::<Vec<_>>();
+
+        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)));
+    }
+}