numerical_analysis 0.1.2

#[cfg(feature = "std")]
use misc_iterators::{IterDimensions, integer_decomposition::ExactTupleSum};

use crate::EvalResult;
use crate::RealField;
#[cfg(not(target_arch = "spirv"))]
use crate::binomial;
use crate::legendre_roots::*;
use crate::to_index;

#[derive(Default, Clone, Debug)]
pub struct NormalizedMultiDimensionalLegendreBasis<
    S: RealField,
    const ARRAY_SIZE: usize = 20,
> {
    pub aabb_min: [S; 3],
    pub aabb_max: [S; 3],
}

impl<S: RealField, const ARRAY_SIZE: usize>
    NormalizedMultiDimensionalLegendreBasis<S, ARRAY_SIZE>
{
    pub fn new(aabb_min: [S; 3], aabb_max: [S; 3]) -> Self { Self { aabb_min, aabb_max } }

    pub fn aabb(&self) -> [[S; 3]; 2] { [self.aabb_min, self.aabb_max] }

    pub fn eval(
        &self,
        x: [S; 3],
        degree: usize,
        #[cfg(not(feature = "std"))] out: &mut [S; ARRAY_SIZE],
    ) -> EvalResult<S>
    {
        let construct = |i| {
            // Normalize input.
            let a = self.aabb_min[i];
            let b = self.aabb_max[i];
            let x_p = x[i] * (S::from(2).unwrap() / (b - a)) - ((a + b) / (b - a));

            #[cfg(feature = "std")]
            assert!(x_p >= S::from(-1.0001).unwrap() && x_p <= S::from(1.0001).unwrap());

            #[cfg(not(feature = "std"))]
            let mut res = [S::default(); 6];
            #[cfg(feature = "std")]
            let mut res = vec![S::default(); degree + 1];

            res[0] = S::from(1).unwrap();
            res[1] = x_p;

            let one = S::one();
            let two = S::from(2).unwrap();

            for j in 2..=degree
            {
                let n = S::from(j).unwrap();
                res[j] = (two * n - one) * x_p * res[j - 1] - (n - one) * res[j - 2];
                res[j] /= n;
            }

            for j in 0..=degree
            {
                let rho = S::from(j).unwrap();
                res[j] = res[j] * ((two * rho + one) / (b - a)).sqrt();
            }

            res
        };

        let evals = [construct(0), construct(1), construct(2)];

        #[cfg(not(feature = "std"))]
        let mut res = out;
        #[cfg(feature = "std")]
        let mut res = {
            let total_basis_count = binomial(degree + 3, 3);
            vec![S::from(0.).unwrap(); total_basis_count]
        };

        // Iterate over all integer tuples (i, j, k) such that
        // i + j + k <= degree.
        for i in 2..degree + 3
        {
            for j in 1..i
            {
                for k in 0..j
                {
                    let w = i - j - 1;
                    let v = j - k - 1;
                    let u = k;

                    let indices = [u, v, w];
                    let mut prod = S::from(1).unwrap();
                    for l in 0..3
                    {
                        prod *= evals[l][indices[l]];
                    }

                    res[to_index(&indices)] = prod;
                }
            }
        }

        #[cfg(feature = "std")]
        return res;

        #[cfg(not(feature = "std"))]
        []
    }

    #[cfg(feature = "std")]
    pub fn project<F>(
        &mut self,
        degree: usize,
        prior_degree: usize,
        mut function: F,
    ) -> Vec<S>
    where
        F: FnMut([S; 3]) -> S,
    {
        debug_assert!(prior_degree < degree);
        if degree == 0
        {
            return Vec::new();
        }
        // We use the nodes for the 4 * degree scheme in legendre interpolation.
        // This is a heuristic described in "Hierarchichal HP-Adaptive Sigend Distance
        // fields." We may want to use a different quadrature scheme in the
        // future.
        let n = 4 * degree;
        let nodes = legendre_nodes_lut(n);
        let weights = legendre_weights_lut(n);

        let total_basis_count = binomial(degree + 3, 3);

        let aabb_min = self.aabb_min;
        let aabb_max = self.aabb_max;
        let scale_quadrature = |x: [S; 3]| {
            // Normalize the quadrature point to the cell bounds.
            let one = S::from(1).unwrap();
            let two = S::from(2).unwrap();
            let point: [S; 3] = std::array::from_fn(|i| {
                let a = aabb_min[i];
                let b = aabb_max[i];
                ((b - a) * (x[i] + one) / two) + a
            });

            point
        };

        let mut coefficients = vec![S::default(); total_basis_count];
        let mut values = vec![S::default(); n * n * n];

        for multi_index in IterDimensions::new([n, n, n])
        {
            let mut weight = S::from(1.).unwrap();

            for (j, i) in multi_index.into_iter().enumerate()
            {
                weight *= S::from(weights[i]).unwrap() * (aabb_max[j] - aabb_min[j])
                    / S::from(2).unwrap();
            }

            let point = scale_quadrature(multi_index.map(|i| S::from(nodes[i]).unwrap()));

            let val = function(point);
            values[multi_index[2] * n * n + multi_index[1] * n + multi_index[0]] =
                val * weight;

            let vw = values[multi_index[2] * n * n + multi_index[1] * n + multi_index[0]];
            let quadrature_point =
                scale_quadrature(multi_index.map(|i| S::from(nodes[i]).unwrap()));
            let basis_eval = self.eval(quadrature_point, degree);

            for i in 2..degree + 3
            {
                for j in 1..i
                {
                    for k in 0..j
                    {
                        let w = i - j - 1;
                        let v = j - k - 1;
                        let u = k;

                        let power_indices = [u, v, w];
                        let index = to_index(&power_indices);
                        coefficients[index] += vw * basis_eval[index];
                    }
                }
            }
        }

        coefficients
    }

    #[cfg(feature = "std")]
    pub fn project_sparse<F>(
        &mut self,
        degree: usize,
        prior_degree: usize,
        mut function: F,
    ) -> (Vec<(S, [usize; 3])>, usize, [usize; 3])
    where
        F: FnMut([S; 3]) -> S,
    {
        debug_assert!(prior_degree < degree);
        if degree == 0
        {
            return (Vec::new(), 0, [0; 3]);
        }
        // We use the nodes for the 4 * degree scheme in legendre interpolation.
        // This is a heuristic described in "Hierarchichal HP-Adaptive Sigend Distance
        // fields." We may want to use a different quadrature scheme in the
        // future.
        let n = 4 * degree;
        let nodes = legendre_nodes_lut(n);
        let weights = legendre_weights_lut(n);

        let total_basis_count = binomial(degree + 3, 3);

        let aabb_min = self.aabb_min;
        let aabb_max = self.aabb_max;
        let scale_quadrature = |x: [S; 3]| {
            // Normalize the quadrature point to the cell bounds.
            let one = S::from(1).unwrap();
            let two = S::from(2).unwrap();
            let point: [S; 3] = std::array::from_fn(|i| {
                let a = aabb_min[i];
                let b = aabb_max[i];
                ((b - a) * (x[i] + one) / two) + a
            });

            point
        };

        let mut coefficients = vec![S::default(); total_basis_count];
        let mut values = vec![S::default(); n * n * n];

        for multi_index in IterDimensions::new([n, n, n])
        {
            let mut weight = S::from(1.).unwrap();

            for (j, i) in multi_index.into_iter().enumerate()
            {
                weight *= S::from(weights[i]).unwrap() * (aabb_max[j] - aabb_min[j])
                    / S::from(2).unwrap();
            }

            let point = scale_quadrature(multi_index.map(|i| S::from(nodes[i]).unwrap()));

            let val = function(point);
            values[multi_index[2] * n * n + multi_index[1] * n + multi_index[0]] =
                val * weight;

            let vw = values[multi_index[2] * n * n + multi_index[1] * n + multi_index[0]];
            let quadrature_point =
                scale_quadrature(multi_index.map(|i| S::from(nodes[i]).unwrap()));
            let basis_eval = self.eval(quadrature_point, degree);

            for i in 2..degree + 3
            {
                for j in 1..i
                {
                    for k in 0..j
                    {
                        let w = i - j - 1;
                        let v = j - k - 1;
                        let u = k;

                        let power_indices = [u, v, w];
                        let index = to_index(&power_indices);
                        coefficients[index] += vw * basis_eval[index];
                    }
                }
            }
        }

        let mut maxs = [0; 3];
        let mut coeff_counter = 0;
        let mut sparse_coefficients = Vec::new();
        for i in 2..degree + 3
        {
            for j in 1..i
            {
                for k in 0..j
                {
                    let w = i - j - 1;
                    let v = j - k - 1;
                    let u = k;

                    let power_indices = [u, v, w];
                    let index = to_index(&power_indices);
                    if coefficients[index] <= S::from(0.0001).unwrap()
                    {
                        continue;
                    }

                    maxs[0] = maxs[0].max(u);
                    maxs[1] = maxs[1].max(v);
                    maxs[2] = maxs[2].max(w);
                    sparse_coefficients.push((coefficients[index], power_indices));
                    coeff_counter += 1;
                }
            }
        }

        (sparse_coefficients, coeff_counter, maxs)
    }

    pub fn evaluate_function_vector(
        &self,
        x: [S; 3],
        degree: usize,
        #[cfg(feature = "std")] coefficients: &[S],
        #[cfg(not(feature = "std"))] coefficients: &[S; ARRAY_SIZE],
    ) -> S
    {
        #[cfg(feature = "std")]
        let basis_eval = self.eval(x, degree);

        #[cfg(not(feature = "std"))]
        let basis_eval = {
            let mut basis_eval = [S::default(); ARRAY_SIZE];
            self.eval(x, degree, &mut basis_eval);

            basis_eval
        };

        let mut res = S::default();
        for i in 0..basis_eval.len()
        {
            let v = basis_eval[i];
            let coeff = coefficients[i];

            res += v.clone() * coeff.clone();
        }

        res
    }

    pub fn evaluate_sparse_function_vector(
        &self,
        x: [S; 3],
        #[cfg(feature = "std")] coefficients: &[(S, [usize; 3])],
        #[cfg(not(feature = "std"))] coefficients: &[(S, [usize; 3]); ARRAY_SIZE],
        coeff_count: usize,
        max_degree: [usize; 3],
    ) -> S
    {
        let construct = |i| {
            // Normalize input.
            let a = self.aabb_min[i];
            let b = self.aabb_max[i];
            let x_p = x[i] * (S::from(2).unwrap() / (b - a)) - ((a + b) / (b - a));

            #[cfg(feature = "std")]
            assert!(x_p >= S::from(-1.0001).unwrap() && x_p <= S::from(1.0001).unwrap());

            #[cfg(not(feature = "std"))]
            let mut res = [S::default(); 6];
            #[cfg(feature = "std")]
            let mut res = vec![S::default(); max_degree[i] + 1];

            res[0] = S::from(1).unwrap();
            res[1] = x_p;

            let one = S::one();
            let two = S::from(2).unwrap();

            for j in 2..=max_degree[i]
            {
                let n = S::from(j).unwrap();
                res[j] = (two * n - one) * x_p * res[j - 1] - (n - one) * res[j - 2];
                res[j] /= n;
            }

            for j in 0..=max_degree[i]
            {
                let rho = S::from(j).unwrap();
                res[j] = res[j] * ((two * rho + one) / (b - a)).sqrt();
            }

            res
        };

        let evals = [construct(0), construct(1), construct(2)];

        let mut res = S::default();
        for i in 0..coeff_count
        {
            let (coeff, indices) = coefficients[i];
            let eval = coeff.clone()
                * evals[0][indices[0]]
                * evals[1][indices[1]]
                * evals[2][indices[2]];

            res += eval;
        }

        res
    }

    /// Estimate the current error of the approximation.
    // This is a heuristic comming from "Hierarchichal hp-Adaptive Signed Distance
    // Fields", equation 6.
    #[cfg(feature = "std")]
    pub fn error(coeffs: &[S], degree: usize) -> S
    {
        let mut error = S::default();
        for index in ExactTupleSum::<3>::new(degree)
        {
            let v = coeffs[to_index(&index)];
            error += v * v;
        }

        error
    }
}

pub struct LegendreIterator<S: RealField>
{
    x: S,
    prior: S,
    current: S,
    count: usize,
}

impl<S: RealField> LegendreIterator<S>
{
    pub fn new(x: S) -> Self
    {
        Self {
            prior: S::from(1.).unwrap(),
            current: x,
            count: 0,
            x,
        }
    }
}

impl<S: RealField> Iterator for LegendreIterator<S>
{
    type Item = S;

    fn next(&mut self) -> Option<Self::Item>
    {
        self.count += 1;

        if self.count == 1
        {
            return Some(self.prior);
        }
        else if self.count == 2
        {
            return Some(self.current);
        }

        let p = S::from(self.count - 1).unwrap();
        let one = S::from(1).unwrap();
        let p_inv = one / p;

        let next = p_inv
            * ((S::from(2).unwrap() * p - one) * self.x * self.current
                - (p - one) * self.prior);

        self.prior = self.current;
        self.current = next;

        Some(self.current)
    }
}

pub fn legendre_nodes_lut(order: usize) -> &'static [f64]
{
    let first = (order - 1) * order / 2;
    let last = first + order;

    &LEGENDRE_NODE_TABLES[first..last]
}

pub fn legendre_weights_lut(order: usize) -> &'static [f64]
{
    let first = (order - 1) * order / 2;
    let last = first + order;

    &LEGENDRE_WEIGHT_TABLES[first..last]
}

/// Legendre quadrature is only defined for the interval [-1, 1], beware.
pub fn quadrature<F>(count: usize, mut function: F) -> f64
where
    F: FnMut(f64) -> f64,
{
    let nodes = legendre_nodes_lut(count);
    let weights = legendre_weights_lut(count);

    let mut integral = 0.;
    for (n, w) in nodes.iter().zip(weights.iter())
    {
        integral += w * function(*n);
    }

    integral
}

#[cfg(test)]
mod tests
{
    use super::*;

    #[test]
    fn test_multidimensional_legendre_quadrature()
    {
        let mut poly = NormalizedMultiDimensionalLegendreBasis::<_, 20>::new(
            [-1., -1., -1.],
            [1., 1., 1.],
        );

        use nalgebra::Vector3;
        let mut poly = NormalizedMultiDimensionalLegendreBasis::<_, 20>::new(
            [-1., -1., -1.],
            [1., 1., 1.],
        );

        let sdf = |x: [f64; 3]| {
            let v = Vector3::from(x);
            v.norm_squared() - 1.
        };
        let coeffs = poly.project(2, 0, &sdf);
    }

    #[test]
    fn test_multidimensional_legendre_quadrature_sparse()
    {
        let mut poly = NormalizedMultiDimensionalLegendreBasis::<_, 20>::new(
            [-1., -1., -1.],
            [1., 1., 1.],
        );

        use nalgebra::Vector3;
        let mut poly = NormalizedMultiDimensionalLegendreBasis::<_, 20>::new(
            [-1., -1., -1.],
            [1., 1., 1.],
        );

        let sdf = |x: [f64; 3]| {
            let v = Vector3::from(x);
            v.norm_squared() - 1.
        };

        let (coeffs, coeff_count, max_degrees) = poly.project_sparse(2, 0, &sdf);
        let res = poly.evaluate_sparse_function_vector(
            [0., 0., 0.],
            &coeffs,
            coeff_count,
            max_degrees,
        );
        println!("{}", res);

        let (coeffs, coeff_count, max_degrees) = poly.project_sparse(2, 0, &sdf);
        let res = poly.evaluate_sparse_function_vector(
            [0., 1., 0.],
            &coeffs,
            coeff_count,
            max_degrees,
        );
        println!("{}", res);

        let (coeffs, coeff_count, max_degrees) = poly.project_sparse(2, 0, &sdf);
        let res = poly.evaluate_sparse_function_vector(
            [0., 2., 0.],
            &coeffs,
            coeff_count,
            max_degrees,
        );
        println!("{}", res);
    }

    #[test]
    fn test_legendre_quadrature()
    {
        let res = quadrature(4, |x| x * x);
        assert!((res - (2. / 3.)).abs() < 0.00001);

        let res = quadrature(5, |x| x.sin());
        assert!(res.abs() < 0.00001);

        let res = quadrature(5, |x| x.cos());
        assert!((res - 2. * 1_f64.sin()).abs() < 0.00001);

        let res = quadrature(5, |x| x.exp());
        let e = std::f64::consts::E;
        assert!((res - (e - 1. / e)).abs() < 0.00001);
    }

    #[test]
    fn test_legendre_nodes()
    {
        let n = 5;
        let weights = legendre_weights_lut(n);
        for (index, xi) in legendre_nodes_lut(n).iter().enumerate()
        {
            let mut p_n2 = 1.;
            let mut p_n1 = *xi;

            for i in 2..=n
            {
                let i = i as f64;
                let next = ((2. * i - 1.) * xi * p_n1 - (i - 1.) * p_n2) / i;

                p_n2 = p_n1;
                p_n1 = next;
            }

            let n = n as f64;
            let n1 = n * p_n2;
            let d = 1. - xi * xi;
            let w = if n > 0. { (2. * d) / (n1 * n1) } else { 2. };

            assert_eq!(w, weights[index]);
        }
    }

    #[test]
    fn test_legendre_iterator()
    {
        let iter = LegendreIterator::new(0.5);
        let expected = vec![
            1.,
            0.5,
            -0.125,
            -0.4375,
            -0.2890625,
            0.08984375,
            0.3232421875,
            0.22314453125,
            -0.073638916015625,
            -0.2678985595703125,
        ];

        for (v, e) in iter.take(10).zip(expected.iter())
        {
            assert!(v == *e);
        }
    }
}