src/mapping.rs

/*!
Traits for mathematical functions.
*/

use crate::error::DynResult;
use crate::instance::MyCow;
pub use crate::instance::{BasicDecomposition, ClosedSpace, Decomposition, Instance, Space};
use crate::loc::Loc;
use crate::norms::{Norm, NormExponent};
use crate::operator_arithmetic::{Constant, Weighted};
use crate::types::{ClosedMul, Float, Num};
use std::marker::PhantomData;
use std::ops::Mul;

/// A mapping from `Domain` to `Self::Codomain`.
pub trait Mapping<Domain: Space> {
    type Codomain: ClosedSpace;

    /// Compute the value of `self` at `x`.
    fn apply<I: Instance<Domain>>(&self, x: I) -> Self::Codomain;

    #[inline]
    /// Form the composition `self ∘ other`
    fn compose<X: Space, T: Mapping<X, Codomain = Domain>>(self, other: T) -> Composition<Self, T>
    where
        Self: Sized,
    {
        Composition { outer: self, inner: other, intermediate_norm_exponent: () }
    }

    #[inline]
    /// Form the composition `self ∘ other`, assigning a norm to the inermediate space
    fn compose_with_norm<F, X, T, E>(self, other: T, norm: E) -> Composition<Self, T, E>
    where
        Self: Sized,
        X: Space,
        T: Mapping<X, Codomain = Domain>,
        E: NormExponent,
        Domain: Norm<E, F>,
        F: Num,
    {
        Composition { outer: self, inner: other, intermediate_norm_exponent: norm }
    }

    /// Multiply `self` by the scalar `a`.
    #[inline]
    fn weigh<C>(self, a: C) -> Weighted<Self, C>
    where
        Self: Sized,
        C: Constant,
        Self::Codomain: ClosedMul<C::Type>,
    {
        Weighted { weight: a, base_fn: self }
    }
}

/// Automatically implemented shorthand for referring to [`Mapping`]s from [`Loc<N, F>`] to `F`.
pub trait RealMapping<const N: usize, F: Float = f64>: Mapping<Loc<N, F>, Codomain = F> {}

impl<F: Float, T, const N: usize> RealMapping<N, F> for T where T: Mapping<Loc<N, F>, Codomain = F> {}

/// A helper trait alias for referring to [`Mapping`]s from [`Loc<N, F>`] to [`Loc<M, F>`].
pub trait RealVectorField<const N: usize, const M: usize, F: Float = f64>:
    Mapping<Loc<N, F>, Codomain = Loc<M, F>>
{
}

impl<F: Float, T, const N: usize, const M: usize> RealVectorField<N, M, F> for T where
    T: Mapping<Loc<N, F>, Codomain = Loc<M, F>>
{
}

/// A differentiable mapping from `Domain` to [`Mapping::Codomain`], with differentials
/// `Differential`.
///
/// This is automatically implemented when [`DifferentiableImpl`] is.
pub trait DifferentiableMapping<Domain: Space>: Mapping<Domain> {
    type DerivativeDomain: ClosedSpace;
    type Differential<'b>: Mapping<Domain, Codomain = Self::DerivativeDomain>
    where
        Self: 'b;

    /// Calculate differential at `x`
    fn differential<I: Instance<Domain>>(&self, x: I) -> Self::DerivativeDomain;

    /// Calculate differential and value at `x`
    fn apply_and_differential<I: Instance<Domain>>(
        &self,
        x: I,
    ) -> (Self::Codomain, Self::DerivativeDomain) {
        x.eval_ref(|xo| (self.apply(xo), self.differential(xo)))
    }

    /// Form the differential mapping of `self`.
    fn diff(self) -> Self::Differential<'static>;

    /// Form the differential mapping of `self`.
    fn diff_ref(&self) -> Self::Differential<'_>;
}

/// Automatically implemented shorthand for referring to differentiable [`Mapping`]s from
/// [`Loc<N, F>`] to `F`.
pub trait DifferentiableRealMapping<const N: usize, F: Float>:
    DifferentiableMapping<Loc<N, F>, Codomain = F, DerivativeDomain = Loc<N, F>>
{
}

impl<F: Float, T, const N: usize> DifferentiableRealMapping<N, F> for T where
    T: DifferentiableMapping<Loc<N, F>, Codomain = F, DerivativeDomain = Loc<N, F>>
{
}

/// Helper trait for implementing [`DifferentiableMapping`]
pub trait DifferentiableImpl<X: Space>: Sized {
    type Derivative: ClosedSpace;

    /// Compute the differential of `self` at `x`, consuming the input.
    fn differential_impl<I: Instance<X>>(&self, x: I) -> Self::Derivative;

    fn apply_and_differential_impl<I: Instance<X>>(
        &self,
        x: I,
    ) -> (Self::Codomain, Self::Derivative)
    where
        Self: Mapping<X>,
    {
        x.eval_ref(|xo| (self.apply(xo), self.differential_impl(xo)))
    }
}

impl<T, Domain> DifferentiableMapping<Domain> for T
where
    Domain: Space,
    T: Mapping<Domain> + DifferentiableImpl<Domain>,
{
    type DerivativeDomain = T::Derivative;
    type Differential<'b>
        = Differential<'b, Domain, Self>
    where
        Self: 'b;

    #[inline]
    fn differential<I: Instance<Domain>>(&self, x: I) -> Self::DerivativeDomain {
        self.differential_impl(x)
    }

    #[inline]
    fn apply_and_differential<I: Instance<Domain>>(
        &self,
        x: I,
    ) -> (T::Codomain, Self::DerivativeDomain) {
        self.apply_and_differential_impl(x)
    }

    fn diff(self) -> Differential<'static, Domain, Self> {
        Differential { g: MyCow::Owned(self), _space: PhantomData }
    }

    fn diff_ref(&self) -> Differential<'_, Domain, Self> {
        Differential { g: MyCow::Borrowed(self), _space: PhantomData }
    }
}

/// Container for the differential [`Mapping`] of a [`DifferentiableMapping`].
pub struct Differential<'a, X, G> {
    g: MyCow<'a, G>,
    _space: PhantomData<X>,
}

impl<'a, X, G> Differential<'a, X, G> {
    pub fn base_fn(&self) -> &G {
        &self.g
    }
}

impl<'a, X, G> Mapping<X> for Differential<'a, X, G>
where
    X: Space,
    G: DifferentiableMapping<X>,
{
    type Codomain = G::DerivativeDomain;

    #[inline]
    fn apply<I: Instance<X>>(&self, x: I) -> Self::Codomain {
        (*self.g).differential(x)
    }
}

/// Container for flattening [`Loc`]`<F, 1>` codomain of a [`Mapping`] to `F`.
pub struct FlattenedCodomain<X, F, G> {
    g: G,
    _phantoms: PhantomData<(X, F)>,
}

impl<F, X, G> Mapping<X> for FlattenedCodomain<X, F, G>
where
    F: ClosedSpace,
    X: Space,
    G: Mapping<X, Codomain = Loc<1, F>>,
{
    type Codomain = F;

    #[inline]
    fn apply<I: Instance<X>>(&self, x: I) -> Self::Codomain {
        self.g.apply(x).flatten1d()
    }
}

/// An auto-trait for constructing a [`FlattenCodomain`] structure for
/// flattening the codomain of a [`Mapping`] from [`Loc`]`<F, 1>` to `F`.
pub trait FlattenCodomain<X: Space, F>: Mapping<X, Codomain = Loc<1, F>> + Sized {
    /// Flatten the codomain from [`Loc`]`<F, 1>` to `F`.
    fn flatten_codomain(self) -> FlattenedCodomain<X, F, Self> {
        FlattenedCodomain { g: self, _phantoms: PhantomData }
    }
}

impl<X: Space, F, G: Sized + Mapping<X, Codomain = Loc<1, F>>> FlattenCodomain<X, F> for G {}

/// Container for dimensional slicing [`Loc`]`<N, F>` codomain of a [`Mapping`] to `F`.
pub struct SlicedCodomain<'a, X, F, G, const N: usize> {
    g: MyCow<'a, G>,
    slice: usize,
    _phantoms: PhantomData<(X, F)>,
}

impl<'a, X, F, G, const N: usize> Mapping<X> for SlicedCodomain<'a, X, F, G, N>
where
    X: Space,
    F: Copy + ClosedSpace,
    G: Mapping<X, Codomain = Loc<N, F>>,
{
    type Codomain = F;

    #[inline]
    fn apply<I: Instance<X>>(&self, x: I) -> Self::Codomain {
        let tmp: [F; N] = (*self.g).apply(x).into();
        // Safety: `slice_codomain` below checks the range.
        unsafe { *tmp.get_unchecked(self.slice) }
    }
}

/// An auto-trait for constructing a [`FlattenCodomain`] structure for
/// flattening the codomain of a [`Mapping`] from [`Loc`]`<F, 1>` to `F`.
pub trait SliceCodomain<X: Space, const N: usize, F: Copy = f64>:
    Mapping<X, Codomain = Loc<N, F>> + Sized
{
    /// Flatten the codomain from [`Loc`]`<F, 1>` to `F`.
    fn slice_codomain(self, slice: usize) -> SlicedCodomain<'static, X, F, Self, N> {
        assert!(slice < N);
        SlicedCodomain { g: MyCow::Owned(self), slice, _phantoms: PhantomData }
    }

    /// Flatten the codomain from [`Loc`]`<F, 1>` to `F`.
    fn slice_codomain_ref(&self, slice: usize) -> SlicedCodomain<'_, X, F, Self, N> {
        assert!(slice < N);
        SlicedCodomain { g: MyCow::Borrowed(self), slice, _phantoms: PhantomData }
    }
}

impl<X: Space, F: Copy, G: Sized + Mapping<X, Codomain = Loc<N, F>>, const N: usize>
    SliceCodomain<X, N, F> for G
{
}

/// The composition S ∘ T. `E` is for storing a `NormExponent` for the intermediate space.
pub struct Composition<S, T, E = ()> {
    pub outer: S,
    pub inner: T,
    pub intermediate_norm_exponent: E,
}

impl<S, T, X, E> Mapping<X> for Composition<S, T, E>
where
    X: Space,
    T: Mapping<X>,
    S: Mapping<T::Codomain>,
{
    type Codomain = S::Codomain;

    #[inline]
    fn apply<I: Instance<X>>(&self, x: I) -> Self::Codomain {
        self.outer.apply(self.inner.apply(x))
    }
}

/// Helper trait for implementing [`DifferentiableMapping`]
impl<S, T, X, E, Y> DifferentiableImpl<X> for Composition<S, T, E>
where
    X: Space,
    T: DifferentiableImpl<X> + Mapping<X>,
    S: DifferentiableImpl<T::Codomain>,
    E: Copy,
    //Composition<S::Derivative, T::Derivative, E>: Space,
    S::Derivative: Mul<T::Derivative, Output = Y>,
    Y: ClosedSpace,
{
    //type Derivative = Composition<S::Derivative, T::Derivative, E>;
    type Derivative = Y;

    /// Compute the differential of `self` at `x`, consuming the input.
    fn differential_impl<I: Instance<X>>(&self, x: I) -> Self::Derivative {
        // Composition {
        //     outer: self
        //         .outer
        //         .differential_impl(self.inner.apply(x.ref_instance())),
        //     inner: self.inner.differential_impl(x),
        //     intermediate_norm_exponent: self.intermediate_norm_exponent,
        // }

        self.outer
            .differential_impl(x.eval_ref(|r| self.inner.apply(r)))
            * self.inner.differential_impl(x)
    }
}

mod dataterm;
pub use dataterm::DataTerm;

/// Trait for indicating that `Self` is Lipschitz with respect to the (semi)norm `D`.
pub trait Lipschitz<M> {
    /// The type of floats
    type FloatType: Float;

    /// Returns the Lipschitz factor of `self` with respect to the (semi)norm `D`.
    fn lipschitz_factor(&self, seminorm: M) -> DynResult<Self::FloatType>;
}

/// Helper trait for implementing [`Lipschitz`] for mappings that implement [`DifferentiableImpl`].
pub trait LipschitzDifferentiableImpl<X: Space, M>: DifferentiableImpl<X> {
    type FloatType: Float;

    /// Compute the lipschitz factor of the derivative of `f`.
    fn diff_lipschitz_factor(&self, seminorm: M) -> DynResult<Self::FloatType>;
}

impl<'b, M, X, A> Lipschitz<M> for Differential<'b, X, A>
where
    X: Space,
    A: LipschitzDifferentiableImpl<X, M>,
{
    type FloatType = A::FloatType;

    fn lipschitz_factor(&self, seminorm: M) -> DynResult<Self::FloatType> {
        (*self.g).diff_lipschitz_factor(seminorm)
    }
}