src/nalgebra_support.rs

///! Integration with nalgebra

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

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 : 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>;

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

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<'a>(&self, y : &mut Matrix<E,N,K,SV2>, x : &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<SM,M,E> Projection<E, 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 : 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>,E>
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)
    }
}

/// 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>,
      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<E>
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<E, 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<E, 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<E, 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<E, 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<E, 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<E, 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`
/// while allowing nalgebra to be used in subroutines.
pub trait ToNalgebraRealField : Float {
    type NalgebraType : RealField;
    type MixedType : RealField + Float;

    fn to_nalgebra(self) -> Self::NalgebraType;
    fn to_nalgebra_mixed(self) -> Self::MixedType;

    fn from_nalgebra(t : Self::NalgebraType) -> 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 }
}