numerical_analysis 0.2.0

#![cfg(feature = "std")]

use linear_isomorphic::*;
use rstar::primitives::GeomWithData;
use rstar::{Point, RTree};

/// Size of the neighbourhood the area estimate is built from, i.e. the `k` in
/// `a_i = pi * r_k(p_i)^2 / k`.
pub const KNN_K: usize = 8;

/// A `D`-dimensional position stored in the R-tree.
///
/// Coordinates are kept as `f64` regardless of the scalar of the vertex type:
/// the tree is only used to estimate a neighbourhood radius, so the conversion
/// costs nothing in accuracy that matters and keeps `rstar`'s numeric bounds
/// out of the public signatures.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct SamplePoint<const D: usize>
{
    coords: [f64; D],
}

impl<const D: usize> SamplePoint<D>
{
    /// Project a vertex onto its first `D` coordinates.
    pub fn from_vertex<Vert, S>(vertex: &Vert) -> Self
    where
        Vert: VectorSpace<Scalar = S>,
        S: linear_isomorphic::RealField,
    {
        let mut coords = [0.0; D];
        for (i, c) in coords.iter_mut().enumerate()
        {
            *c = vertex[i].to_f64().unwrap_or(0.0);
        }

        Self { coords }
    }
}

impl<const D: usize> Point for SamplePoint<D>
{
    type Scalar = f64;

    const DIMENSIONS: usize = D;

    fn generate(mut generator: impl FnMut(usize) -> f64) -> Self
    {
        let mut coords = [0.0; D];
        for (i, c) in coords.iter_mut().enumerate()
        {
            *c = generator(i);
        }

        Self { coords }
    }

    fn nth(&self, index: usize) -> f64 { self.coords[index] }

    fn nth_mut(&mut self, index: usize) -> &mut f64 { &mut self.coords[index] }
}

/// A sample position tagged with the index it was handed in with.
pub type IndexedSample<const D: usize> = GeomWithData<SamplePoint<D>, usize>;

/// Spatial index over a point cloud. Query it through `rstar`'s own methods,
/// e.g. `query.nearest_neighbor_iter(SamplePoint::from_vertex(&v))`.
pub type KnnQuery<const D: usize> = RTree<IndexedSample<D>>;

/// Build a [`KnnQuery`] over `samples`. Only the first `D` coordinates of each
/// vertex are used.
pub fn index_samples<const D: usize, Vert, S>(samples: &[(Vert, usize)]) -> KnnQuery<D>
where
    Vert: VectorSpace<Scalar = S>,
    S: linear_isomorphic::RealField,
{
    RTree::bulk_load(
        samples
            .iter()
            .map(|(v, vid)| IndexedSample::new(SamplePoint::from_vertex(v), *vid))
            .collect(),
    )
}

/// Radius of the `k`-point neighbourhood of `point`: the distance to the
/// furthest of its `k` nearest neighbours, `r_k(p_i)`.
///
/// The sample sitting at `point` is kept out of its own neighbourhood, which is
/// what `index` identifies. `None` when the query holds no other sample.
pub fn neighbourhood_radius<const D: usize>(
    query: &KnnQuery<D>,
    point: SamplePoint<D>,
    index: usize,
    k: usize,
) -> Option<f64>
{
    query
        .nearest_neighbor_iter_with_distance_2(point)
        .filter(|(sample, _)| sample.data != index)
        .take(k)
        .last()
        .map(|(_, distance_squared)| distance_squared.sqrt())
}

/// The area a sample carries under a k-nearest-neighbour density estimate,
/// `a_i = pi * r_k(p_i)^2 / k`. Zero for a sample with no neighbours.
pub fn knn_area_element<const D: usize, Vert, S>(
    query: &KnnQuery<D>,
    vertex: &Vert,
    index: usize,
    k: usize,
) -> S
where
    Vert: VectorSpace<Scalar = S>,
    S: linear_isomorphic::RealField,
{
    let point = SamplePoint::from_vertex(vertex);
    let radius = neighbourhood_radius(query, point, index, k).unwrap_or(0.0);

    S::from(core::f64::consts::PI * radius * radius / k as f64).unwrap_or_else(S::default)
}

pub fn integrate<Vert, Aarea, Function, S>(
    samples: impl Iterator<Item = (Vert, usize)>,
    area_element: Aarea,
    function: Function,
) -> S
where
    Aarea: Fn(&Vert, usize) -> S,
    S: linear_isomorphic::RealField,
    Function: Fn(&Vert, usize) -> S,
{
    let mut res = S::default();
    for (v, vid) in samples
    {
        let area = area_element(&v, vid);
        let val = function(&v, vid);

        res += val * area;
    }

    res
}

/// Integrate `function` over a bare point cloud, with no connectivity
/// information.
///
/// The area carried by each sample is estimated from the local point density,
/// `a_i = pi * r_k(p_i)^2 / k`, where `r_k(p_i)` is the radius of that sample's
/// own [`KNN_K`]-point neighbourhood, i.e. the distance to the furthest of its
/// [`KNN_K`] nearest neighbours.
///
/// `D` is the dimension of the ambient space to index the samples in, e.g.
/// `integrate_points::<3, _, _, _>(..)` for a surface in space.
pub fn integrate_points<const D: usize, Vert, Function, S>(
    samples: impl Iterator<Item = (Vert, usize)>,
    function: Function,
) -> S
where
    Vert: VectorSpace<Scalar = S>,
    S: linear_isomorphic::RealField,
    Function: Fn(&Vert, usize) -> S,
{
    let samples: Vec<(Vert, usize)> = samples.collect();
    if samples.is_empty()
    {
        return S::default();
    }

    let query = index_samples::<D, _, _>(&samples);

    // A sample is never its own neighbour, so a small cloud cannot fill a
    // `KNN_K`-point neighbourhood.
    let k = KNN_K.min(samples.len() - 1).max(1);

    integrate(
        samples.into_iter(),
        |v, vid| knn_area_element(&query, v, vid, k),
        function,
    )
}

#[cfg(test)]
mod tests
{
    use rand::rngs::StdRng;
    use rand::{Rng, SeedableRng};

    use super::*;

    // `nalgebra` rather than the internal `algebra` crate, which this published
    // crate cannot depend on.
    type Vec3 = nalgebra::Vector3<f64>;

    /// One draw from `N(0, 1)`, by the Box-Muller transform.
    fn standard_normal(rng: &mut StdRng) -> f64
    {
        // `random` yields `[0, 1)`, and the log needs `(0, 1]`.
        let u1 = 1.0 - rng.random::<f64>();
        let u2 = rng.random::<f64>();

        (-2.0 * u1.ln()).sqrt() * (core::f64::consts::TAU * u2).cos()
    }

    /// `count` points spread uniformly over the unit sphere. A vector of three
    /// independent standard normals is isotropic, so normalizing it leaves a
    /// direction with no preferred axis.
    fn unit_sphere_samples(count: usize) -> Vec<(Vec3, usize)>
    {
        let mut rng = StdRng::seed_from_u64(0x5eed_5b1e_0f77_3a91);

        (0..count)
            .map(|i| {
                let v = Vec3::new(
                    standard_normal(&mut rng),
                    standard_normal(&mut rng),
                    standard_normal(&mut rng),
                );

                (v.normalize(), i)
            })
            .collect()
    }

    #[test]
    fn integrates_the_area_of_the_sphere()
    {
        let samples = unit_sphere_samples(20_000);

        let area = integrate_points::<3, _, _, _>(samples.into_iter(), |_, _| 1.0);

        // `pi * r_k^2 / k` is unbiased for a Poisson process, and the sphere has
        // no boundary for the neighbourhoods to run off, so this converges on
        // the real area rather than to the lattice's `pi * 2 / 8` bias.
        let expected = 4.0 * core::f64::consts::PI;
        assert!(
            (area - expected).abs() < 0.001 * expected,
            "{area} vs {expected}"
        );
    }

    #[test]
    fn integrates_the_volume_of_the_sphere()
    {
        let samples = unit_sphere_samples(20_000);

        let volume =
            integrate_points::<3, _, _, _>(samples.into_iter(), |v, _| v.dot(v) / 3.);

        let expected = 4.0 / 3. * core::f64::consts::PI;
        assert!(
            (volume - expected).abs() < 0.001 * expected,
            "{volume} vs {expected}"
        );
    }

    #[test]
    fn integrates_a_varying_function_over_the_sphere()
    {
        let samples = unit_sphere_samples(20_000);

        // `\int z^2 dA = 4 pi / 3` over the unit sphere.
        let integral =
            integrate_points::<3, _, _, _>(samples.into_iter(), |v, _| v.z * v.z);

        let expected = 4.0 * core::f64::consts::PI / 3.0;
        assert!(
            (integral - expected).abs() < 0.01 * expected,
            "{integral} vs {expected}"
        );
    }
}