numerical_analysis 0.2.0

use itertools::iproduct;
use linear_isomorphic::RealField;

// Note: All products use logarithmic summation in iterated_product() for better
// numerical stability and to avoid overflow/underflow in intermediate calculations.

/// Computes $product^n_(i=1; i != k) (v_i - u)$ where $v_i$ are the `values`
/// provided by the iterator, `u` is the `fixed`` value (which should correspond
/// to the `index` in the value sequence), and $k$ is the `index` parameter.
pub fn lagrange_denominator<I, S>(values: I, fixed: S, index: usize) -> S
where
    I: Iterator<Item = S>,
    S: linear_isomorphic::RealField + std::iter::Product + std::iter::Sum,
{
    iterated_product(values.enumerate().filter_map(|(i, v)| {
        if i == index
        {
            assert!(v == fixed);
            None
        }
        else
        {
            Some(v - fixed)
        }
    }))
}

/// Computes the numerator of a lagrange polynomial of order 0, that is:
/// $product^n_(i=1; i!=k) v_i$ where $v_i$ are the iterator `values` and $k$ is
/// the selected `index`.
pub fn lagrange_numerator_order_0<I, S>(values: I, index: usize) -> S
where
    I: Iterator<Item = S>,
    S: linear_isomorphic::RealField + std::iter::Product + std::iter::Sum,
{
    iterated_product(
        values
            .enumerate()
            .filter_map(|(i, v)| if i == index { None } else { Some(v) }),
    )
}

/// Computes the numerator of a lagrange polynomial of order 1, i.e. the first
/// derivative of the lower order. That is: $sum^n_(k=1; k != j) product^n_(l=1;
/// l!=k, l!=j) v_l$ where $v_l$ are the iterator `values` and $j$ is the
/// selected `index`.
pub fn lagrange_numerator_order_1<I, S>(values: I, index: usize) -> S
where
    I: Iterator<Item = S> + ExactSizeIterator + Clone,
    S: linear_isomorphic::RealField + std::iter::Product + std::iter::Sum,
{
    -(0..values.len())
        .filter_map(|k| {
            if k == index
            {
                None
            }
            else
            {
                Some(iterated_product(values.clone().enumerate().filter_map(
                    |(l, v)| {
                        if l == k || l == index { None } else { Some(v) }
                    },
                )))
            }
        })
        .sum::<S>()
}

/// Computes the numerator of a lagrange polynomial of order 2, i.e. the seond
/// derivative of the lowest order. That is: $sum^n_(k=1; k != j)sum^n_(l=1; l
/// != j, l!= k) product^n_(m=1; m!=k, m!=j, m!=l) v_m$ where $v_m$ are the
/// iterator `values` and $j$ is the selected `index`.
pub fn lagrange_numerator_order_2<I, S>(values: I, index: usize) -> S
where
    I: Iterator<Item = S> + ExactSizeIterator + Clone,
    S: linear_isomorphic::RealField + std::iter::Product + std::iter::Sum,
{
    (0..values.len())
        .filter_map(|k| {
            if k == index
            {
                None
            }
            else
            {
                Some(
                    (0..values.len())
                        .filter_map(|l| {
                            if l == index || l == k
                            {
                                None
                            }
                            else
                            {
                                Some(iterated_product(
                                    values.clone().enumerate().filter_map(|(m, v)| {
                                        if m == k || m == l || m == index
                                        {
                                            None
                                        }
                                        else
                                        {
                                            Some(v)
                                        }
                                    }),
                                ))
                            }
                        })
                        .sum::<S>(),
                )
            }
        })
        .sum::<S>()
}

/// Use Lagrange Differentiation to compute the mixed partial derivatives up to
/// order two of a function. This method is very sensitive to machine epsilon,
/// if it is too small, it can give catastrophic errors such as 0 due to
/// truncation. Larger step sizes are safer.
///
/// Based on section 2.2 of [Computational assessment of curvatures and
/// principal directions of implicit surfaces from 3D salar data](
/// https://hal.science/hal-01486547/file/albin2016lncs.pdf)
pub fn partial_derivatives<F, S>(
    point: [S; 3],
    function: F,
    order: [u8; 3],
    step_size: S,
    sample_count: usize,
) -> S
where
    F: Fn(S, S, S) -> S + Clone,
    S: linear_isomorphic::RealField + std::iter::Product + std::iter::Sum,
{
    let n = sample_count;
    let centered_offset = step_size
        * (S::from(n / 2).unwrap()
            + S::from(0.5).unwrap() * S::from((n % 2 == 0) as u8).unwrap());

    let directional_pos = |j| -centered_offset + S::from(j).unwrap() * step_size;

    iproduct!(0..n, 0..n, 0..n)
        .map(|(i, j, k)| {
            let index = [i, j, k];
            let mut res = S::from(1.).unwrap();
            for l in 0..3
            {
                let skip_val = directional_pos(index[l]);
                let samples = (0..n).map(|m| directional_pos(m));
                let denom = lagrange_denominator(samples.clone(), skip_val, index[l]);

                let numerator = match order[l]
                {
                    0 => lagrange_numerator_order_0(samples.clone(), index[l]),
                    1 => lagrange_numerator_order_1(samples.clone(), index[l]),
                    2 => lagrange_numerator_order_2(samples.clone(), index[l]),
                    _ =>
                    {
                        panic!()
                    }
                };

                res *= numerator / denom;
            }

            let sample = function(
                point[0] + directional_pos(i),
                point[1] + directional_pos(j),
                point[2] + directional_pos(k),
            );

            res * sample
        })
        .sum()
}

fn iterated_product<I, S>(iter: I) -> S
where
    I: Iterator<Item = S>,
    S: RealField + std::iter::Sum + std::iter::Product,
{
    let mut sign = S::from(1.0).unwrap();
    let log_sum: S = iter
        .map(|v| {
            let coeff = if v.signum() < S::from(0.).unwrap()
            {
                S::from(-1.).unwrap()
            }
            else
            {
                S::from(1.).unwrap()
            };
            sign = sign * coeff;

            v.abs().ln()
        })
        .sum();

    log_sum.exp() * sign
}

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

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

    #[test]
    fn test_lagrange_differentiation()
    {
        let sdf_fun = |p: &DVec3| p.dot(&p) - 1.0;
        let test_fun = |x: f64, y: f64, z: f64| sdf_fun(&DVec3::new(x, y, z));
        let sdf_partial_x = |p: &DVec3| 2. * p.x;
        let _sdf_partial_y = |p: &DVec3| 2. * p.y;
        let _sdf_partial_z = |p: &DVec3| 2. * p.z;

        let test_point = DVec3::new(1., 0., 0.);
        let test_x_partial = partial_derivatives(
            [test_point.x, test_point.y, test_point.z],
            test_fun,
            [0, 2, 0],
            0.0000001,
            3,
        );

        let true_dx = sdf_partial_x(&test_point);
        assert!((true_dx - test_x_partial).abs() < 0.01);
    }
}