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,
    ClosedMul, ClosedAdd, 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::types::Float;
use crate::norms::*;

impl<SM,SV,N,M,K,E> Apply<Matrix<E,M,K,SV>> for Matrix<E,N,M,SM>
where SM: Storage<E,N,M>, SV: Storage<E,M,K>,
        N : Dim, M : Dim, K : Dim, E : Scalar + ClosedMul + ClosedAdd + Zero + One,
        DefaultAllocator : Allocator<E,N,K>,
        DefaultAllocator : Allocator<E,M,K>,
        DefaultAllocator : Allocator<E,N,M>,
        DefaultAllocator : Allocator<E,M,N> {
    type Output = OMatrix<E,N,K>;

    #[inline]
    fn apply(&self, x : Matrix<E,M,K,SV>) -> Self::Output {
        self.mul(x)
    }
}

impl<'a, SM,SV,N,M,K,E> Apply<&'a Matrix<E,M,K,SV>> for Matrix<E,N,M,SM>
where SM: Storage<E,N,M>, SV: Storage<E,M,K>,
        N : Dim, M : Dim, K : Dim, E : Scalar + ClosedMul + ClosedAdd + Zero + One,
        DefaultAllocator : Allocator<E,N,K>,
        DefaultAllocator : Allocator<E,M,K>,
        DefaultAllocator : Allocator<E,N,M>,
        DefaultAllocator : Allocator<E,M,N> {
    type Output = OMatrix<E,N,K>;

    #[inline]
    fn apply(&self, x : &'a Matrix<E,M,K,SV>) -> Self::Output {
        self.mul(x)
    }
}

impl<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>,
        N : Dim, M : Dim, K : Dim, E : Float + Scalar + ClosedMul + ClosedAdd + Zero + One,
        DefaultAllocator : Allocator<E,N,K>,
        DefaultAllocator : Allocator<E,M,K>,
        DefaultAllocator : Allocator<E,N,M>,
        DefaultAllocator : Allocator<E,M,N> {
    type Codomain = OMatrix<E,N,K>;
}

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>, SV2: StorageMut<E,N,K>,
      N : Dim, M : Dim, K : Dim, E : Scalar + ClosedMul + ClosedAdd + Zero + One + Float,
      DefaultAllocator : Allocator<E,N,K>,
      DefaultAllocator : Allocator<E,M,K>,
      DefaultAllocator : Allocator<E,N,M>,
      DefaultAllocator : Allocator<E,M,N> {

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

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

impl<'a, SM,SV1,SV2,N,M,K,E> GEMV<E, &'a 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>, SV2: StorageMut<E,N,K>,
      N : Dim, M : Dim, K : Dim, E : Scalar + ClosedMul + ClosedAdd + Zero + One + Float,
      DefaultAllocator : Allocator<E,N,K>,
      DefaultAllocator : Allocator<E,M,K>,
      DefaultAllocator : Allocator<E,N,M>,
      DefaultAllocator : Allocator<E,M,N> {

    #[inline]
    fn gemv(&self, y : &mut Matrix<E,N,K,SV2>, α : E, x : &'a Matrix<E,M,K,SV1>, β : E) {
        Matrix::gemm(y, α, self, x, β)
    }

    #[inline]
    fn apply_mut(&self, y : &mut Matrix<E,N,K,SV2>, x : &'a Matrix<E,M,K,SV1>) {
        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>, SV1: Storage<E,M>,
      M : Dim, E : Scalar + ClosedMul + ClosedAdd + Zero + One + Float,
      DefaultAllocator : Allocator<E,M> {

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

    #[inline]
    fn copy_from(&mut self, y : Vector<E,M,SV1>) {
        Matrix::copy_from(self, &y)
    }
}

impl<'a, SM,SV1,M,E> AXPY<E, &'a Vector<E,M,SV1>> for Vector<E,M,SM>
where SM: StorageMut<E,M>, SV1: Storage<E,M>,
      M : Dim, E : Scalar + ClosedMul + ClosedAdd + Zero + One + Float,
      DefaultAllocator : Allocator<E,M> {

    #[inline]
    fn axpy(&mut self, α : E, x : &'a Vector<E,M,SV1>, β : E) {
        Matrix::axpy(self, α, x, β)
    }

    #[inline]
    fn copy_from(&mut self, y : &'a Vector<E,M,SV1>) {
        Matrix::copy_from(self, y)
    }
}

impl<SM,M,E> Projection<Linfinity> for Vector<E,M,SM>
where SM: StorageMut<E,M>,
      M : Dim, E : Scalar + ClosedMul + ClosedAdd + Zero + One + Float + RealField,
      DefaultAllocator : Allocator<E,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>, SV2: Storage<E,N,K>,
      N : Dim, M : Dim, K : Dim,
      E : Float + Scalar + ClosedMul + ClosedAdd + Zero + One + SimdComplexField,
      DefaultAllocator : Allocator<E,N,K>,
      DefaultAllocator : Allocator<E,M,K>,
      DefaultAllocator : Allocator<E,N,M>,
      DefaultAllocator : Allocator<E,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>>
for Vector<E,M,S>
where M : Dim,
      E : Float + Scalar + ClosedMul + ClosedAdd + Zero + One,
      S : Storage<E,M>,
      Si : Storage<E,M>,
      DefaultAllocator : Allocator<E,M> {

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

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

    #[inline]
    fn dot(&self, other : &'a 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,K,S> HasScalarField
for Matrix<E,M,K,S>
where M : Dim, K : Dim,
      S : Storage<E,M,K>,
      E : Float + Scalar,
      DefaultAllocator : Allocator<E,M,K> {
    type Field = E;
}

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

    type Output = OVector<E, M>;
    
    #[inline]
    fn similar_origin(&self) -> OVector<E, M> {
        OVector::zeros_generic(M::from_usize(self.len()), Const)
    }

    #[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
for Vector<E,M,S>
where M : DimName,
      S : StorageMut<E,M>,
      E : Float + Scalar + ClosedMul + ClosedAdd + Zero + One + RealField,
      DefaultAllocator : Allocator<E,M> {

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

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

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

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

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

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

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

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

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

impl<E,M,S> Dist<Linfinity>
for Vector<E,M,S>
where M : Dim,
      S : StorageMut<E,M>,
      E : Float + Scalar + ClosedMul + ClosedAdd + Zero + One + RealField,
      DefaultAllocator : Allocator<E,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 }
}