src/mapping.rs

/*!
Traits for mathematical functions.
*/

use std::marker::PhantomData;
use crate::types::{Float, HasRealField, HasScalarField};
use serde::Serialize;
use crate::loc::Loc;
use std::borrow::Cow;

/// Trait for application of `Self` as a mathematical function or operator on `X`.
pub trait Apply<X> {
    type Output;

    /// Compute the value of `self` at `x`.
    fn apply(&self, x : X) -> Self::Output;
}

/// This blanket implementation is a workaround helper to Rust trait system limitations.
///
/// It is introduced because the trait system does not allow blanket implementations of both
/// [`Apply<X>`] and [`Apply<&'a X>`]. With this, the latter is implemented automatically for
/// the reference, which can be sufficient to apply the operation in another blanket implementation.
impl<'a, T, X> Apply<X> for &'a T where T : Apply<X> {
    type Output = <T as Apply<X>>::Output;

    #[inline]
    fn apply(&self, x : X) -> Self::Output {
        (*self).apply(x)
    }
}

/// A mapping from `Domain` to `Codomain`.
///
/// This is automatically implemented when the relevant [`Apply`] are implemented.
pub trait Mapping<Domain> : Apply<Domain, Output=Self::Codomain>
                            + for<'a> Apply<&'a Domain, Output=Self::Codomain> {
    type Codomain;
}

impl<Domain, Codomain, T> Mapping<Domain> for T
where T : Apply<Domain, Output=Codomain> + for<'a> Apply<&'a Domain, Output=Codomain> {
    type Codomain = Codomain;
}

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

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

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

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


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

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


/// Trait for calculation the differential of `Self` as a mathematical function on `X`.
pub trait Differentiable<X> : Sized + HasRealField {
    type Derivative : HasRealField<RealField = Self::RealField>;

    /// Compute the differential of `self` at `x`.
    fn differential(&self, x : X) -> Self::Derivative;
}

// impl<'g, X, G : Differentiable<X>> Differentiable<X> for &'g G {
//     type Derivative = G::Derivative;
//     #[inline]
//     fn differential(&self, x : X) -> Self::Derivative {
//         (*self).differential(x)
//     }
// }

/// A differentiable mapping from `Domain` to [`Mapping::Codomain`], with differentials
/// `Differential`.
///
/// This is automatically implemented when the relevant [`Differentiate`] are implemented.
pub trait DifferentiableMapping<Domain>
: Mapping<Domain>
  + Differentiable<Domain, Derivative=Self::DerivativeDomain>
  + for<'a> Differentiable<&'a Domain, Derivative=Self::DerivativeDomain> {
    type DerivativeDomain : HasRealField<RealField=Self::RealField>;
    type Differential<'b> : Mapping<Domain, Codomain=Self::DerivativeDomain>
                            + HasRealField<RealField=Self::RealField>
                            + Clone where Self : 'b;

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


impl<Domain : Clone, Derivative, T> DifferentiableMapping<Domain> for T
where T : Mapping<Domain> + Clone
          + Differentiable<Domain, Derivative=Derivative>
          + for<'a> Differentiable<&'a Domain, Derivative=Derivative>,
      Derivative : HasRealField<RealField = T::RealField> {
    type DerivativeDomain = Derivative;
    type Differential<'b> = Differential<'b, Domain, Self> where Self : 'b + Clone;
    
    /// Form the differential mapping of `self`.
    fn diff(self) -> Self::Differential<'static> {
        Differential{ g : Cow::Owned(self), _space : PhantomData }
    }

    /// Form the differential mapping of `self`.
    fn diff_ref(&self) -> Self::Differential<'_> {
        Differential{ g : Cow::Borrowed(self), _space : PhantomData }
    }
}

/// A sum of [`Mapping`]s.
#[derive(Serialize, Debug, Clone)]
pub struct Sum<Domain, M : Mapping<Domain>> {
    components : Vec<M>,
    _domain : PhantomData<Domain>,
}

impl<Domain, M : Mapping<Domain> + HasScalarField> HasScalarField for Sum<Domain, M> {
    type Field = M::Field;
}

impl<Domain, M : Mapping<Domain>> Sum<Domain, M> {
    /// Construct from an iterator.
    pub fn new<I : Iterator<Item = M>>(iter : I) -> Self {
        Sum { components : iter.collect(), _domain : PhantomData }
    }

    /// Iterate over the component functions of the sum
    pub fn iter(&self) -> std::slice::Iter<'_, M> {
        self.components.iter()
    }
}


impl<Domain, M> Apply<Domain> for Sum<Domain, M>
where M : Mapping<Domain>,
      M::Codomain : std::iter::Sum {
    type Output = M::Codomain;

    fn apply(&self, x : Domain) -> Self::Output {
        self.components.iter().map(|c| c.apply(&x)).sum()
    }
}

impl<'a, Domain, M> Apply<&'a Domain> for Sum<Domain, M>
where M : Mapping<Domain>,
      M::Codomain : std::iter::Sum {
    type Output = M::Codomain;

    fn apply(&self, x : &'a Domain) -> Self::Output {
        self.components.iter().map(|c| c.apply(x)).sum()
    }
}

impl<Domain, M> Differentiable<Domain> for Sum<Domain, M>
where M : DifferentiableMapping<Domain>,
      M :: Codomain : std::iter::Sum,
      M :: DerivativeDomain : std::iter::Sum,
      Domain : Copy {

    type Derivative = M::DerivativeDomain;

    fn differential(&self, x : Domain) -> Self::Derivative {
        self.components.iter().map(|c| c.differential(x)).sum()
    }
}

/// Container for the differential [`Mapping`] of a [`Differentiable`] mapping.
#[derive(Clone, Debug)]
pub struct Differential<'a, X, G : 'a + DifferentiableMapping<X> + Clone> {
    g : Cow<'a, G>,
    _space : PhantomData<X>
}

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


impl<'a, X, G> HasScalarField for Differential<'a, X, G>
where G : 'a + DifferentiableMapping<X> + HasScalarField + Clone {
    type Field = G::Field;
}

impl<'b, X, G : DifferentiableMapping<X> + Clone> Apply<X> for Differential<'b, X, G> {
    type Output = G::DerivativeDomain;

    #[inline]
    fn apply(&self, x : X) -> Self::Output {
        self.g.differential(x)
    }
}

impl<'a, 'b, X, G : 'b + DifferentiableMapping<X> + Clone> Apply<&'a X>
for Differential<'b, X, G> {
    type Output = G::DerivativeDomain;

    #[inline]
    fn apply(&self, x : &'a X) -> Self::Output {
        self.g.differential(x)
    }
}


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

impl<X, F, G : Mapping<X, Codomain=Loc<F, 1>>> Apply<X> for FlattenedCodomain<X, F, G> {
    type Output = F;

    #[inline]
    fn apply(&self, x : X) -> Self::Output {
        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, F> : Mapping<X, Codomain=Loc<F, 1>> + 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, F, G : Sized + Mapping<X, Codomain=Loc<F, 1>>> FlattenCodomain<X, F> for G {}


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

impl<'a, X, F, G, const N : usize> HasScalarField for SlicedCodomain<'a, X, F, G, N>
where G : HasScalarField + Mapping<X, Codomain=Loc<F, N>> + Clone {
    type Field = G::Field;
}

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

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

impl<'a, 'b, X, F : Copy, G : Mapping<X, Codomain=Loc<F, N>> + Clone, const N : usize>
Apply<&'b X>
for SlicedCodomain<'a, X, F, G, N> {
    type Output = F;

    #[inline]
    fn apply(&self, x : &'b X) -> Self::Output {
        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, F : Copy, const N : usize> : Mapping<X, Codomain=Loc<F, N>> + Clone + 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 : Cow::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 : Cow::Borrowed(self), slice, _phantoms : PhantomData }
    }
}

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