src/nalgebra_support.rs

/*!
Integration with nalgebra.

This module mainly implements [`Euclidean`], [`Norm`], [`Dot`], [`Linear`], etc. for [`nalgebra`]
matrices and vectors.
It also provides [`ToNalgebraRealField`] as a vomit-inducingly ugly workaround to nalgebra
force-feeding its own versions of the same basic mathematical methods on `f32` and `f64` as
[`num_traits`] does.
*/

use nalgebra::{
    Matrix, Storage, StorageMut, OMatrix, Dim, DefaultAllocator, Scalar,
    ClosedAddAssign, ClosedMulAssign, SimdComplexField, Vector, OVector, RealField,
    LpNorm, UniformNorm
};
use nalgebra::Norm as NalgebraNorm;
use nalgebra::base::constraint::{
    ShapeConstraint, SameNumberOfRows, SameNumberOfColumns
};
use nalgebra::base::dimension::*;
use nalgebra::base::allocator::Allocator;
use std::ops::Mul;
use num_traits::identities::{Zero, One};
use crate::linops::*;
use crate::euclidean::*;
use crate::mapping::{Space, BasicDecomposition};
use crate::types::Float;
use crate::norms::*;
use crate::instance::Instance;

impl<SM,N,M,E> Space for Matrix<E,N,M,SM>
where
    SM: Storage<E,N,M> + Clone,
    N : Dim, M : Dim, E : Scalar + Zero + One + ClosedAddAssign + ClosedMulAssign,
    DefaultAllocator : Allocator<N,M>,
{
    type Decomp = BasicDecomposition;
}

impl<SM,SV,N,M,K,E> Mapping<Matrix<E,M,K,SV>> for Matrix<E,N,M,SM>
where SM: Storage<E,N,M>, SV: Storage<E,M,K> + Clone,
        N : Dim, M : Dim, K : Dim, E : Scalar + Zero + One + ClosedAddAssign + ClosedMulAssign,
        DefaultAllocator : Allocator<N,K>,
        DefaultAllocator : Allocator<M,K>,
        DefaultAllocator : Allocator<N,M>,
        DefaultAllocator : Allocator<M,N> {
    type Codomain = OMatrix<E,N,K>;

    #[inline]
    fn apply<I : Instance<Matrix<E,M,K,SV>>>(
        &self, x : I
    ) -> Self::Codomain {
        x.either(|owned| self.mul(owned), |refr| self.mul(refr))
    }
}


impl<'a, SM,SV,N,M,K,E> Linear<Matrix<E,M,K,SV>> for Matrix<E,N,M,SM>
where SM: Storage<E,N,M>, SV: Storage<E,M,K> + Clone,
        N : Dim, M : Dim, K : Dim, E : Scalar + Zero + One + ClosedAddAssign + ClosedMulAssign,
        DefaultAllocator : Allocator<N,K>,
        DefaultAllocator : Allocator<M,K>,
        DefaultAllocator : Allocator<N,M>,
        DefaultAllocator : Allocator<M,N> {
}

impl<SM,SV1,SV2,N,M,K,E> GEMV<E, Matrix<E,M,K,SV1>, Matrix<E,N,K,SV2>> for Matrix<E,N,M,SM>
where SM: Storage<E,N,M>, SV1: Storage<E,M,K> + Clone, SV2: StorageMut<E,N,K>,
      N : Dim, M : Dim, K : Dim, E : Scalar + Zero + One + Float,
      DefaultAllocator : Allocator<N,K>,
      DefaultAllocator : Allocator<M,K>,
      DefaultAllocator : Allocator<N,M>,
      DefaultAllocator : Allocator<M,N> {

    #[inline]
    fn gemv<I : Instance<Matrix<E,M,K,SV1>>>(
        &self, y : &mut Matrix<E,N,K,SV2>, α : E, x : I, β : E
    ) {
        x.eval(|x̃| Matrix::gemm(y, α, self, x̃, β))
    }

    #[inline]
    fn apply_mut<'a, I : Instance<Matrix<E,M,K,SV1>>>(&self, y : &mut Matrix<E,N,K,SV2>, x : I) {
        x.eval(|x̃| self.mul_to(x̃, y))
    }
}

impl<SM,SV1,M,E> AXPY<E, Vector<E,M,SV1>> for Vector<E,M,SM>
where SM: StorageMut<E,M> + Clone, SV1: Storage<E,M> + Clone,
      M : Dim, E : Scalar + Zero + One + Float,
      DefaultAllocator : Allocator<M> {
    type Owned = OVector<E, M>;

    #[inline]
    fn axpy<I : Instance<Vector<E,M,SV1>>>(&mut self, α : E, x : I, β : E) {
        x.eval(|x̃| Matrix::axpy(self, α, x̃, β))
    }

    #[inline]
    fn copy_from<I : Instance<Vector<E,M,SV1>>>(&mut self, y : I) {
        y.eval(|ỹ| Matrix::copy_from(self, ỹ))
    }

    #[inline]
    fn set_zero(&mut self) {
        self.iter_mut().for_each(|e| *e = E::ZERO);
    }

    #[inline]
    fn similar_origin(&self) -> Self::Owned {
        OVector::zeros_generic(M::from_usize(self.len()), Const)
    }
}

/* Implemented automatically as Euclidean.
impl<SM,M,E> Projection<E, L2> for Vector<E,M,SM>
where SM: StorageMut<E,M> + Clone,
      M : Dim, E : Scalar + Zero + One + Float + RealField,
      DefaultAllocator : Allocator<M> {
    #[inline]
    fn proj_ball_mut(&mut self, ρ : E, _ : L2) {
        let n = self.norm(L2);
        if n > ρ {
            self.iter_mut().for_each(|v| *v *= ρ/n)
        }
    }
}*/

impl<SM,M,E> Projection<E, Linfinity> for Vector<E,M,SM>
where SM: StorageMut<E,M> + Clone,
      M : Dim, E : Scalar + Zero + One + Float + RealField,
      DefaultAllocator : Allocator<M> {
    #[inline]
    fn proj_ball_mut(&mut self, ρ : E, _ : Linfinity) {
        self.iter_mut().for_each(|v| *v = num_traits::clamp(*v, -ρ, ρ))
    }
}

impl<'own,SV1,SV2,SM,N,M,K,E> Adjointable<Matrix<E,M,K,SV1>, Matrix<E,N,K,SV2>>
for Matrix<E,N,M,SM>
where SM: Storage<E,N,M>, SV1: Storage<E,M,K> + Clone, SV2: Storage<E,N,K> + Clone,
      N : Dim, M : Dim, K : Dim, E : Scalar + Zero + One + SimdComplexField,
      DefaultAllocator : Allocator<N,K>,
      DefaultAllocator : Allocator<M,K>,
      DefaultAllocator : Allocator<N,M>,
      DefaultAllocator : Allocator<M,N> {
    type AdjointCodomain = OMatrix<E,M,K>;
    type Adjoint<'a> = OMatrix<E,M,N> where SM : 'a;

    #[inline]
    fn adjoint(&self) -> Self::Adjoint<'_> {
        Matrix::adjoint(self)
    }
}

impl<E,M,S,Si> Dot<Vector<E,M,Si>,E>
for Vector<E,M,S>
where M : Dim,
      E : Float + Scalar + Zero + One,
      S : Storage<E,M>,
      Si : Storage<E,M>,
      DefaultAllocator : Allocator<M> {

    #[inline]
    fn dot(&self, other : &Vector<E,M,Si>) -> E {
        Vector::<E,M,S>::dot(self, other)
    }
}

/// This function is [`nalgebra::EuclideanNorm::metric_distance`] without the `sqrt`.
#[inline]
fn metric_distance_squared<T, R1, C1, S1, R2, C2, S2>(
    /*ed: &EuclideanNorm,*/
    m1: &Matrix<T, R1, C1, S1>,
    m2: &Matrix<T, R2, C2, S2>,
) -> T::SimdRealField
where
    T:  SimdComplexField,
    R1: Dim,
    C1: Dim,
    S1: Storage<T, R1, C1>,
    R2: Dim,
    C2: Dim,
    S2: Storage<T, R2, C2>,
    ShapeConstraint: SameNumberOfRows<R1, R2> + SameNumberOfColumns<C1, C2>,
{
    m1.zip_fold(m2, T::SimdRealField::zero(), |acc, a, b| {
        let diff = a - b;
        acc + diff.simd_modulus_squared()
    })
}

// TODO: should allow different input storages in `Euclidean`.

impl<E,M,S> Euclidean<E>
for Vector<E,M,S>
where M : Dim,
      S : StorageMut<E,M> + Clone,
      E : Float + Scalar + Zero + One + RealField,
      DefaultAllocator : Allocator<M> {

    type Output = OVector<E, M>;
    
    #[inline]
    fn norm2_squared(&self) -> E {
        Vector::<E,M,S>::norm_squared(self)
    }

    #[inline]
    fn dist2_squared(&self, other : &Self) -> E {
        metric_distance_squared(self, other)
    }
}

impl<E,M,S> StaticEuclidean<E>
for Vector<E,M,S>
where M : DimName,
      S : StorageMut<E,M> + Clone,
      E : Float + Scalar + Zero + One + RealField,
      DefaultAllocator : Allocator<M> {

    #[inline]
    fn origin() -> OVector<E, M> {
        OVector::zeros()
    }
}

/// The default norm for `Vector` is [`L2`].
impl<E,M,S> Normed<E>
for Vector<E,M,S>
where M : Dim,
      S : Storage<E,M> + Clone,
      E : Float + Scalar + Zero + One + RealField,
      DefaultAllocator : Allocator<M> {

    type NormExp = L2;

    #[inline]
    fn norm_exponent(&self) -> Self::NormExp {
        L2
    }

    #[inline]
    fn is_zero(&self) -> bool {
        Vector::<E,M,S>::norm_squared(self) == E::ZERO
    }
}

impl<E,M,S> HasDual<E>
for Vector<E,M,S>
where M : Dim,
      S : Storage<E,M> + Clone,
      E : Float + Scalar + Zero + One + RealField,
      DefaultAllocator : Allocator<M> {
    // TODO: Doesn't work with different storage formats.
    type DualSpace = Self;
}

impl<E,M,S> Norm<E, L1>
for Vector<E,M,S>
where M : Dim,
      S : Storage<E,M>,
      E : Float + Scalar + Zero + One + RealField,
      DefaultAllocator : Allocator<M> {

    #[inline]
    fn norm(&self, _ : L1) -> E {
        LpNorm(1).norm(self)
    }
}

impl<E,M,S> Dist<E, L1>
for Vector<E,M,S>
where M : Dim,
      S : Storage<E,M>,
      E : Float + Scalar + Zero + One + RealField,
      DefaultAllocator : Allocator<M> {
    #[inline]
    fn dist(&self, other : &Self, _ : L1) -> E {
        LpNorm(1).metric_distance(self, other)
    }
}

impl<E,M,S> Norm<E, L2>
for Vector<E,M,S>
where M : Dim,
      S : Storage<E,M>,
      E : Float + Scalar + Zero + One + RealField,
      DefaultAllocator : Allocator<M> {

    #[inline]
    fn norm(&self, _ : L2) -> E {
        LpNorm(2).norm(self)
    }
}

impl<E,M,S> Dist<E, L2>
for Vector<E,M,S>
where M : Dim,
      S : Storage<E,M>,
      E : Float + Scalar + Zero + One + RealField,
      DefaultAllocator : Allocator<M> {
    #[inline]
    fn dist(&self, other : &Self, _ : L2) -> E {
        LpNorm(2).metric_distance(self, other)
    }
}

impl<E,M,S> Norm<E, Linfinity>
for Vector<E,M,S>
where M : Dim,
      S : Storage<E,M>,
      E : Float + Scalar + Zero + One + RealField,
      DefaultAllocator : Allocator<M> {

    #[inline]
    fn norm(&self, _ : Linfinity) -> E {
        UniformNorm.norm(self)
    }
}

impl<E,M,S> Dist<E, Linfinity>
for Vector<E,M,S>
where M : Dim,
      S : Storage<E,M>,
      E : Float + Scalar + Zero + One + RealField,
      DefaultAllocator : Allocator<M> {
    #[inline]
    fn dist(&self, other : &Self, _ : Linfinity) -> E {
        UniformNorm.metric_distance(self, other)
    }
}

/// Helper trait to hide the symbols of [`nalgebra::RealField`].
///
/// By assuming `ToNalgebraRealField` intead of `nalgebra::RealField` as a trait bound,
/// functions can piggyback `nalgebra::RealField` without exponsing themselves to it.
/// Thus methods from [`num_traits`] can be used directly without similarly named methods
/// from [`nalgebra`] conflicting with them. Only when absolutely necessary to work with
/// nalgebra, one can convert to the nalgebra view of the same type using the methods of
/// this trait.
pub trait ToNalgebraRealField : Float {
    /// The nalgebra type corresponding to this type. Usually same as `Self`.
    ///
    /// This type only carries `nalgebra` traits.
    type NalgebraType : RealField;
    /// The “mixed” type corresponding to this type. Usually same as `Self`.
    ///
    /// This type carries both `num_traits` and `nalgebra` traits.
    type MixedType : RealField + Float;

    /// Convert to the nalgebra view of `self`.
    fn to_nalgebra(self) -> Self::NalgebraType;

    /// Convert to the mixed (nalgebra and num_traits) view of `self`.
    fn to_nalgebra_mixed(self) -> Self::MixedType;

    /// Convert from the nalgebra view of `self`.
    fn from_nalgebra(t : Self::NalgebraType) -> Self;

    /// Convert from the mixed (nalgebra and num_traits) view to `self`.
    fn from_nalgebra_mixed(t : Self::MixedType) -> Self;
}

impl ToNalgebraRealField for f32 {
    type NalgebraType = f32;
    type MixedType = f32;

    #[inline]
    fn to_nalgebra(self) -> Self::NalgebraType { self }

    #[inline]
    fn to_nalgebra_mixed(self) -> Self::MixedType { self }

    #[inline]
    fn from_nalgebra(t : Self::NalgebraType) -> Self { t }

    #[inline]
    fn from_nalgebra_mixed(t : Self::MixedType) -> Self { t }

}

impl ToNalgebraRealField for f64 {
    type NalgebraType = f64;
    type MixedType = f64;

    #[inline]
    fn to_nalgebra(self) -> Self::NalgebraType { self }

    #[inline]
    fn to_nalgebra_mixed(self) -> Self::MixedType { self }

    #[inline]
    fn from_nalgebra(t : Self::NalgebraType) -> Self { t }

    #[inline]
    fn from_nalgebra_mixed(t : Self::MixedType) -> Self { t }
}