src/fe_model/p2_local_model.rs

/*!
Second order polynomical (P2) models on real intervals and planar 2D simplices.
*/

use super::base::{LocalModel, RealLocalModel};
use crate::euclidean::Euclidean;
use crate::instance::Instance;
use crate::linsolve::*;
use crate::loc::Loc;
use crate::sets::Cube;
use crate::sets::{NPolygon, Set, SpannedHalfspace};
use crate::types::*;
use numeric_literals::replace_float_literals;

/// Type for simplices of arbitrary dimension `N`.
///
/// The type parameter `D` indicates the number of nodes. (Rust's const generics do not currently
/// allow its automatic calculation from `N`.)
pub struct Simplex<F: Float, const N: usize, const D: usize>(pub [Loc<N, F>; D]);
/// A two-dimensional planar simplex
pub type PlanarSimplex<F> = Simplex<F, 2, 3>;
/// A real interval
pub type RealInterval<F> = Simplex<F, 1, 2>;

/// Calculates (a+b)/2
#[inline]
#[replace_float_literals(F::cast_from(literal))]
pub(crate) fn midpoint<F: Float, const N: usize>(a: &Loc<N, F>, b: &Loc<N, F>) -> Loc<N, F> {
    (a + b) / 2.0
}

impl<'a, F: Float> Set<Loc<1, F>> for RealInterval<F> {
    #[inline]
    fn contains<I: Instance<Loc<1, F>>>(&self, z: I) -> bool {
        z.eval(|&Loc([x])| {
            let &[Loc([x0]), Loc([x1])] = &self.0;
            (x0 < x && x < x1) || (x1 < x && x < x0)
        })
    }
}

impl<'a, F: Float> Set<Loc<2, F>> for PlanarSimplex<F> {
    #[inline]
    fn contains<I: Instance<Loc<2, F>>>(&self, z: I) -> bool {
        let &[x0, x1, x2] = &self.0;
        NPolygon([
            [x0, x1].spanned_halfspace(),
            [x1, x2].spanned_halfspace(),
            [x2, x0].spanned_halfspace(),
        ])
        .contains(z)
    }
}

trait P2Powers {
    type Output;
    type Diff;
    type Full;
    fn p2powers(&self) -> Self::Output;
    fn p2powers_full(&self) -> Self::Full;
    fn p2powers_diff(&self) -> Self::Diff;
}

#[replace_float_literals(F::cast_from(literal))]
impl<F: Float> P2Powers for Loc<1, F> {
    type Output = Loc<1, F>;
    type Full = Loc<3, F>;
    type Diff = Loc<1, Loc<1, F>>;

    #[inline]
    fn p2powers(&self) -> Self::Output {
        let &Loc([x0]) = self;
        [x0 * x0].into()
    }

    #[inline]
    fn p2powers_full(&self) -> Self::Full {
        let &Loc([x0]) = self;
        [1.0, x0, x0 * x0].into()
    }

    #[inline]
    fn p2powers_diff(&self) -> Self::Diff {
        let &Loc([x0]) = self;
        [[x0 + x0].into()].into()
    }
}

#[replace_float_literals(F::cast_from(literal))]
impl<F: Float> P2Powers for Loc<2, F> {
    type Output = Loc<3, F>;
    type Full = Loc<6, F>;
    type Diff = Loc<2, Loc<3, F>>;

    #[inline]
    fn p2powers(&self) -> Self::Output {
        let &Loc([x0, x1]) = self;
        [x0 * x0, x0 * x1, x1 * x1].into()
    }

    #[inline]
    fn p2powers_full(&self) -> Self::Full {
        let &Loc([x0, x1]) = self;
        [1.0, x0, x1, x0 * x0, x0 * x1, x1 * x1].into()
    }

    #[inline]
    fn p2powers_diff(&self) -> Self::Diff {
        let &Loc([x0, x1]) = self;
        [[x0 + x0, x1, 0.0].into(), [0.0, x0, x1 + x1].into()].into()
    }
}

/// A trait for generating second order polynomial model of dimension `N` on `Self´.
///
/// `Self` should present a subset aset of elements of the type [`Loc`]`<F, N>`.
pub trait P2Model<F: Num, const N: usize> {
    /// Implementation type of the second order polynomical model.
    /// Typically a [`P2LocalModel`].
    type Model: LocalModel<Loc<N, F>, F>;
    /// Generates a second order polynomial model of the function `g` on `Self`.
    fn p2_model<G: Fn(&Loc<N, F>) -> F>(&self, g: G) -> Self::Model;
}

/// A local second order polynomical model of dimension `N` with `E` edges
pub struct P2LocalModel<
    F: Num,
    const N: usize,
    const E: usize, /*, const V : usize, const Q : usize*/
> {
    a0: F,
    a1: Loc<N, F>,
    a2: Loc<E, F>,
    //node_values : Loc<V, F>,
    //edge_values : Loc<Q, F>,
}

//
// 1D planar model construction
//

impl<F: Float> RealInterval<F> {
    #[inline]
    pub fn midpoints(&self) -> [Loc<1, F>; 1] {
        let [ref n0, ref n1] = &self.0;
        let n01 = midpoint(n0, n1);
        [n01]
    }
}

impl<F: Float> P2LocalModel<F, 1, 1 /*, 2, 0*/> {
    /// Creates a new 1D second order polynomical model based on three nodal coordinates and
    /// corresponding function values.
    #[inline]
    pub fn new(&[n0, n1, n01]: &[Loc<1, F>; 3], &[v0, v1, v01]: &[F; 3]) -> Self {
        let p = move |x: &Loc<1, F>, v: F| {
            let Loc([c, d, e]) = x.p2powers_full();
            [c, d, e, v]
        };
        let [a0, a1, a11] = linsolve([p(&n0, v0), p(&n1, v1), p(&n01, v01)]);
        P2LocalModel {
            a0: a0,
            a1: [a1].into(),
            a2: [a11].into(),
            //node_values : [v0, v1].into(),
            //edge_values: [].into(),
        }
    }
}

impl<F: Float> P2Model<F, 1> for RealInterval<F> {
    type Model = P2LocalModel<F, 1, 1 /*, 2, 0*/>;

    #[inline]
    fn p2_model<G: Fn(&Loc<1, F>) -> F>(&self, g: G) -> Self::Model {
        let [n01] = self.midpoints();
        let [n0, n1] = self.0;
        let vals = [g(&n0), g(&n1), g(&n01)];
        let nodes = [n0, n1, n01];
        Self::Model::new(&nodes, &vals)
    }
}

//
// 2D planar model construction
//

impl<F: Float> PlanarSimplex<F> {
    #[inline]
    /// Returns the midpoints of all the edges of the simplex
    pub fn midpoints(&self) -> [Loc<2, F>; 3] {
        let [ref n0, ref n1, ref n2] = &self.0;
        let n01 = midpoint(n0, n1);
        let n12 = midpoint(n1, n2);
        let n20 = midpoint(n2, n0);
        [n01, n12, n20]
    }
}

impl<F: Float> P2LocalModel<F, 2, 3 /*, 3, 3*/> {
    /// Creates a new 2D second order polynomical model based on six nodal coordinates and
    /// corresponding function values.
    #[inline]
    pub fn new(
        &[n0, n1, n2, n01, n12, n20]: &[Loc<2, F>; 6],
        &[v0, v1, v2, v01, v12, v20]: &[F; 6],
    ) -> Self {
        let p = move |x: &Loc<2, F>, v: F| {
            let Loc([c, d, e, f, g, h]) = x.p2powers_full();
            [c, d, e, f, g, h, v]
        };
        let [a0, a1, a2, a11, a12, a22] = linsolve([
            p(&n0, v0),
            p(&n1, v1),
            p(&n2, v2),
            p(&n01, v01),
            p(&n12, v12),
            p(&n20, v20),
        ]);
        P2LocalModel {
            a0: a0,
            a1: [a1, a2].into(),
            a2: [a11, a12, a22].into(),
            //node_values : [v0, v1, v2].into(),
            //edge_values: [v01, v12, v20].into(),
        }
    }
}

impl<F: Float> P2Model<F, 2> for PlanarSimplex<F> {
    type Model = P2LocalModel<F, 2, 3 /*, 3, 3*/>;

    #[inline]
    fn p2_model<G: Fn(&Loc<2, F>) -> F>(&self, g: G) -> Self::Model {
        let midpoints = self.midpoints();
        let [ref n0, ref n1, ref n2] = self.0;
        let [ref n01, ref n12, ref n20] = midpoints;
        let vals = [g(n0), g(n1), g(n2), g(n01), g(n12), g(n20)];
        let nodes = [*n0, *n1, *n2, *n01, *n12, *n20];
        Self::Model::new(&nodes, &vals)
    }
}

macro_rules! impl_local_model {
    ($n:literal, $e:literal, $v:literal, $q:literal) => {
        impl<F: Float> LocalModel<Loc<$n, F>, F> for P2LocalModel<F, $n, /*, $v, $q*/ $e> {
            #[inline]
            fn value(&self, x: &Loc<$n, F>) -> F {
                self.a0 + x.dot(&self.a1) + x.p2powers().dot(&self.a2)
            }

            #[inline]
            fn differential(&self, x: &Loc<$n, F>) -> Loc<$n, F> {
                self.a1 + x.p2powers_diff().map(|di| di.dot(&self.a2))
            }
        }
    };
}

impl_local_model!(1, 1, 2, 0);
impl_local_model!(2, 3, 3, 3);

//
// Minimisation
//

#[replace_float_literals(F::cast_from(literal))]
impl<F: Float> P2LocalModel<F, 1, 1 /*, 2, 0*/> {
    /// Minimises the model along the edge `[x0, x1]`.
    #[inline]
    fn minimise_edge(&self, x0: Loc<1, F>, x1: Loc<1, F>) -> (Loc<1, F>, F) {
        let &P2LocalModel {
            a1: Loc([a1]),
            a2: Loc([a11]),
            //node_values : Loc([v0, v1]),
            ..
        } = self;
        // We do this in cases, first trying for an interior solution, then edges.
        // For interior solution, first check determinant; no point trying if non-positive
        if a11 > 0.0 {
            // An interior solution x[1] has to satisfy
            //   2a₁₁*x[1] + a₁ =0
            // This gives
            let t = -a1 / (2.0 * a11);
            let (Loc([t0]), Loc([t1])) = (x0, x1);
            if (t0 <= t && t <= t1) || (t1 <= t && t <= t0) {
                let x = [t].into();
                let v = self.value(&x);
                return (x, v);
            }
        }

        let v0 = self.value(&x0);
        let v1 = self.value(&x1);
        if v0 < v1 {
            (x0, v0)
        } else {
            (x1, v1)
        }
    }
}

impl<'a, F: Float> RealLocalModel<RealInterval<F>, Loc<1, F>, F>
    for P2LocalModel<F, 1, 1 /*, 2, 0*/>
{
    #[inline]
    fn minimise(&self, &Simplex([x0, x1]): &RealInterval<F>) -> (Loc<1, F>, F) {
        self.minimise_edge(x0, x1)
    }
}

#[replace_float_literals(F::cast_from(literal))]
impl<F: Float> P2LocalModel<F, 2, 3 /*, 3, 3*/> {
    /// Minimise the 2D model along the edge `[x0, x1] = {x0 + t(x1 - x0) | t ∈ [0, 1] }`.
    #[inline]
    fn minimise_edge(
        &self,
        x0: &Loc<2, F>,
        x1: &Loc<2, F>, /*, v0 : F, v1 : F*/
    ) -> (Loc<2, F>, F) {
        let &P2LocalModel {
            a0,
            a1: Loc([a1, a2]),
            a2: Loc([a11, a12, a22]),
            ..
        } = self;
        let &Loc([x00, x01]) = x0;
        let d @ Loc([d0, d1]) = x1 - x0;
        let b0 = a0 + a1 * x00 + a2 * x01 + a11 * x00 * x00 + a12 * x00 * x01 + a22 * x01 * x01;
        let b1 = a1 * d0
            + a2 * d1
            + 2.0 * a11 * d0 * x00
            + a12 * (d0 * x01 + d1 * x00)
            + 2.0 * a22 * d1 * x01;
        let b11 = a11 * d0 * d0 + a12 * d0 * d1 + a22 * d1 * d1;
        let edge_1d_model = P2LocalModel {
            a0: b0,
            a1: Loc([b1]),
            a2: Loc([b11]),
            //node_values : Loc([v0, v1]),
        };
        let (Loc([t]), v) = edge_1d_model.minimise_edge(0.0.into(), 1.0.into());
        (x0 + d * t, v)
    }
}

#[replace_float_literals(F::cast_from(literal))]
impl<'a, F: Float> RealLocalModel<PlanarSimplex<F>, Loc<2, F>, F>
    for P2LocalModel<F, 2, 3 /*, 3, 3*/>
{
    #[inline]
    fn minimise(&self, el: &PlanarSimplex<F>) -> (Loc<2, F>, F) {
        let &P2LocalModel {
            a1: Loc([a1, a2]),
            a2: Loc([a11, a12, a22]),
            //node_values : Loc([v0, v1, v2]),
            ..
        } = self;

        // We do this in cases, first trying for an interior solution, then edges.
        // For interior solution, first check determinant; no point trying if non-positive
        let r = 2.0 * (a11 * a22 - a12 * a12);
        if r > 0.0 {
            // An interior solution (x[1], x[2]) has to satisfy
            //   2a₁₁*x[1] + 2a₁₂*x[2]+a₁ =0 and 2a₂₂*x[1] + 2a₁₂*x[1]+a₂=0
            // This gives
            let x = [(a22 * a1 - a12 * a2) / r, (a12 * a1 - a11 * a2) / r].into();
            if el.contains(&x) {
                return (x, self.value(&x));
            }
        }

        let &[ref x0, ref x1, ref x2] = &el.0;
        let mut min_edge = self.minimise_edge(x0, x1);
        let more_edge = [self.minimise_edge(x1, x2), self.minimise_edge(x2, x0)];

        for edge in more_edge {
            if edge.1 < min_edge.1 {
                min_edge = edge;
            }
        }

        min_edge
    }
}

#[replace_float_literals(F::cast_from(literal))]
impl<'a, F: Float> RealLocalModel<Cube<2, F>, Loc<2, F>, F>
    for P2LocalModel<F, 2, 3 /*, 3, 3*/>
{
    #[inline]
    fn minimise(&self, el: &Cube<2, F>) -> (Loc<2, F>, F) {
        let &P2LocalModel {
            a1: Loc([a1, a2]),
            a2: Loc([a11, a12, a22]),
            //node_values : Loc([v0, v1, v2]),
            ..
        } = self;

        // We do this in cases, first trying for an interior solution, then edges.
        // For interior solution, first check determinant; no point trying if non-positive
        let r = 2.0 * (a11 * a22 - a12 * a12);
        if r > 0.0 {
            // An interior solution (x[1], x[2]) has to satisfy
            //   2a₁₁*x[1] + 2a₁₂*x[2]+a₁ =0 and 2a₂₂*x[1] + 2a₁₂*x[1]+a₂=0
            // This gives
            let x = [(a22 * a1 - a12 * a2) / r, (a12 * a1 - a11 * a2) / r].into();
            if el.contains(&x) {
                return (x, self.value(&x));
            }
        }

        let [x0, x1, x2, x3] = el.corners();
        let mut min_edge = self.minimise_edge(&x0, &x1);
        let more_edge = [
            self.minimise_edge(&x1, &x2),
            self.minimise_edge(&x2, &x3),
            self.minimise_edge(&x3, &x0),
        ];

        for edge in more_edge {
            if edge.1 < min_edge.1 {
                min_edge = edge;
            }
        }

        min_edge
    }
}

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

    #[test]
    fn p2_model_1d_test() {
        let vertices = [Loc([0.0]), Loc([1.0])];
        let domain = Simplex(vertices);
        // A simple quadratic function for which the approximation is exact on reals,
        // and appears exact on f64 as well.
        let f = |&Loc([x]): &Loc<1, f64>| x * x + x + 1.0;
        let model = domain.p2_model(f);
        let xs = [Loc([0.5]), Loc([0.25])];

        for x in vertices.iter().chain(xs.iter()) {
            assert_eq!(model.value(&x), f(&x));
        }

        assert_eq!(model.minimise(&domain), (Loc([0.0]), 1.0));
    }

    #[test]
    fn p2_model_2d_test() {
        let vertices = [Loc([0.0, 0.0]), Loc([1.0, 0.0]), Loc([1.0, 1.0])];
        let domain = Simplex(vertices);
        // A simple quadratic function for which the approximation is exact on reals,
        // and appears exact on f64 as well.
        let f = |&Loc([x, y]): &Loc<2, f64>| -(x * x + x * y + x - 2.0 * y + 1.0);
        let model = domain.p2_model(f);
        let xs = [Loc([0.5, 0.5]), Loc([0.25, 0.25])];

        for x in vertices.iter().chain(xs.iter()) {
            assert_eq!(model.value(&x), f(&x));
        }

        assert_eq!(model.minimise(&domain), (Loc([1.0, 0.0]), -3.0));
    }
}