src/nalgebra_support.rs

/*!
Integration with nalgebra.

This module mainly implements [`Euclidean`], [`Norm`], [`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 crate::euclidean::*;
use crate::instance::{Decomposition, Instance, Ownable, Space};
use crate::linops::*;
use crate::norms::*;
use crate::types::Float;
use nalgebra::base::allocator::Allocator;
use nalgebra::base::constraint::{DimEq, SameNumberOfColumns, SameNumberOfRows, ShapeConstraint};
use nalgebra::base::dimension::*;
use nalgebra::{
    ClosedAddAssign, ClosedMulAssign, DefaultAllocator, Dim, LpNorm, Matrix, MatrixView, OMatrix,
    OVector, RawStorage, RealField, Scalar, SimdComplexField, Storage, StorageMut, UniformNorm,
    Vector, U1,
};
use num_traits::identities::{One, Zero};
use std::ops::Mul;

impl<S, M, N, E> Ownable for Matrix<E, M, N, S>
where
    S: Storage<E, M, N>,
    M: Dim,
    N: Dim,
    E: Scalar + Zero + One,
    DefaultAllocator: Allocator<M, N>,
{
    type OwnedVariant = OMatrix<E, M, N>;

    #[inline]
    fn into_owned(self) -> Self::OwnedVariant {
        Matrix::into_owned(self)
    }

    /// Returns an owned instance of a reference.
    fn clone_owned(&self) -> Self::OwnedVariant {
        Matrix::clone_owned(self)
    }
}

trait StridesOk<E, N, M = U1, S = <DefaultAllocator as Allocator<N, M>>::Buffer<E>>:
    DimEq<Dyn, S::RStride>
    + DimEq<Dyn, S::CStride>
    + DimEq<Dyn, <<DefaultAllocator as Allocator<N, M>>::Buffer<E> as RawStorage<E, N, M>>::RStride>
    + DimEq<Dyn, <<DefaultAllocator as Allocator<N, M>>::Buffer<E> as RawStorage<E, N, M>>::CStride>
where
    S: RawStorage<E, N, M>,
    E: Scalar,
    N: Dim,
    M: Dim,
    DefaultAllocator: Allocator<N, M>,
{
}

impl<S, E, N, M> StridesOk<E, N, M, S> for ShapeConstraint
where
    ShapeConstraint: DimEq<Dyn, S::RStride>
        + DimEq<Dyn, S::CStride>
        + DimEq<
            Dyn,
            <<DefaultAllocator as Allocator<N, M>>::Buffer<E> as RawStorage<E, N, M>>::RStride,
        > + DimEq<
            Dyn,
            <<DefaultAllocator as Allocator<N, M>>::Buffer<E> as RawStorage<E, N, M>>::CStride,
        >,
    S: Storage<E, N, M>,
    E: Scalar,
    N: Dim,
    M: Dim,
    DefaultAllocator: Allocator<N, M>,
{
}

impl<SM, N, M, E> Space for Matrix<E, N, M, SM>
where
    SM: Storage<E, N, M>,
    N: Dim,
    M: Dim,
    E: Scalar + Zero + One,
    DefaultAllocator: Allocator<N, M>,
    ShapeConstraint: StridesOk<E, N, M, SM> + StridesOk<E, N, M>,
{
    type OwnedSpace = OMatrix<E, N, M>;
    type Decomp = MatrixDecomposition;
}

#[derive(Copy, Clone, Debug)]
pub struct MatrixDecomposition;

/// nalgebra [`Storage`] that's convertible to dynamic strides (practically all, but we need this
/// to work around the type system).
trait ViewableStorage<E, M, K = U1>: Storage<E, M, K>
where
    M: Dim,
    K: Dim,
    E: Scalar,
    ShapeConstraint: StridesOk<E, M, K, Self>,
    DefaultAllocator: Allocator<M, K>,
{
}

impl<M, K, E, S> ViewableStorage<E, M, K> for S
where
    S: Storage<E, M, K>,
    M: Dim,
    K: Dim,
    E: Scalar,
    ShapeConstraint: StridesOk<E, M, K, Self>,
    DefaultAllocator: Allocator<M, K>,
{
}

/// nalgebra [`StorageMut`] that's convertible to dynamic strides (practically all, but we need this
/// to work around the type system).
trait ViewableStorageMut<E, M, K = U1>: StorageMut<E, M, K>
where
    M: Dim,
    K: Dim,
    E: Scalar,
    ShapeConstraint: StridesOk<E, M, K, Self>,
    DefaultAllocator: Allocator<M, K>,
{
}

impl<M, K, E, S> ViewableStorageMut<E, M, K> for S
where
    S: StorageMut<E, M, K>,
    M: Dim,
    K: Dim,
    E: Scalar,
    ShapeConstraint: StridesOk<E, M, K, Self>,
    DefaultAllocator: Allocator<M, K>,
{
}

impl<E, M, K, S> Decomposition<Matrix<E, M, K, S>> for MatrixDecomposition
where
    S: ViewableStorage<E, M, K>,
    M: Dim,
    K: Dim,
    E: Scalar + Zero + One,
    DefaultAllocator: Allocator<M, K>,
    ShapeConstraint: StridesOk<E, M, K, S> + StridesOk<E, M, K>,
{
    type OwnedInstance = OMatrix<E, M, K>;

    type Decomposition<'b>
        = OMatrix<E, M, K>
    where
        Matrix<E, M, K, S>: 'b;
    type Reference<'b>
        = MatrixView<'b, E, M, K, Dyn, Dyn>
    where
        Matrix<E, M, K, S>: 'b;

    #[inline]
    fn lift<'b>(r: Self::Reference<'b>) -> Self::Decomposition<'b>
    where
        S: 'b,
    {
        r.into_owned()
    }
}

impl<S1, S2, M, K, E> Instance<Matrix<E, M, K, S1>, MatrixDecomposition> for Matrix<E, M, K, S2>
where
    S1: ViewableStorage<E, M, K>,
    S2: ViewableStorage<E, M, K>,
    M: Dim,
    K: Dim,
    E: Scalar + Zero + One,
    DefaultAllocator: Allocator<M, K>,
    ShapeConstraint: StridesOk<E, M, K, S1> + StridesOk<E, M, K, S2>,
{
    fn eval_decompose<'b, R>(self, f: impl FnOnce(OMatrix<E, M, K>) -> R) -> R
    where
        Self: 'b,
    {
        f(self.into_owned())
    }

    fn eval_ref_decompose<'b, R>(
        &'b self,
        f: impl FnOnce(<MatrixDecomposition as Decomposition<Matrix<E, M, K, S1>>>::Reference<'b>) -> R,
    ) -> R
    where
        Self: 'b,
        Matrix<E, M, K, S1>: 'b,
    {
        f(self.as_view::<M, K, Dyn, Dyn>())
    }

    #[inline]
    fn own(self) -> OMatrix<E, M, K> {
        self.into_owned()
    }
}

impl<'a, S1, S2, M, K, E> Instance<Matrix<E, M, K, S1>, MatrixDecomposition>
    for &'a Matrix<E, M, K, S2>
where
    S1: ViewableStorage<E, M, K>,
    S2: ViewableStorage<E, M, K>,
    M: Dim,
    K: Dim,
    E: Scalar + Zero + One,
    DefaultAllocator: Allocator<M, K>,
    ShapeConstraint: StridesOk<E, M, K, S1> + StridesOk<E, M, K, S2>,
{
    fn eval_decompose<'b, R>(self, f: impl FnOnce(OMatrix<E, M, K>) -> R) -> R
    where
        Self: 'b,
    {
        f(self.into_owned())
    }

    fn eval_ref_decompose<'b, R>(
        &'b self,
        f: impl FnOnce(<MatrixDecomposition as Decomposition<Matrix<E, M, K, S1>>>::Reference<'b>) -> R,
    ) -> R
    where
        Self: 'b,
        Matrix<E, M, K, S1>: 'b,
    {
        f((*self).as_view::<M, K, Dyn, Dyn>())
    }

    #[inline]
    fn own(self) -> OMatrix<E, M, K> {
        self.into_owned()
    }
}

impl<SM, SV, N, M, K, E> Mapping<Matrix<E, M, K, SV>> for Matrix<E, N, M, SM>
where
    SM: ViewableStorage<E, N, M>,
    SV: ViewableStorage<E, M, K> + Clone,
    N: Dim,
    M: Dim,
    K: Dim,
    E: Scalar + Zero + One + ClosedAddAssign + ClosedMulAssign,
    DefaultAllocator: Allocator<N, K> + Allocator<M, K> + Allocator<N, M> + Allocator<M, N>,
    ShapeConstraint: StridesOk<E, N, M, SM> + StridesOk<E, M, K, SV> + StridesOk<E, N, K>,
{
    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: ViewableStorage<E, N, M>,
    SV: ViewableStorage<E, M, K> + Clone,
    N: Dim,
    M: Dim,
    K: Dim,
    E: Scalar + Zero + One + ClosedAddAssign + ClosedMulAssign,
    DefaultAllocator: Allocator<N, K> + Allocator<M, K> + Allocator<N, M> + Allocator<M, N>,
    ShapeConstraint: StridesOk<E, N, M, SM> + StridesOk<E, M, K, SV> + StridesOk<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: ViewableStorage<E, N, M>,
    SV1: ViewableStorage<E, M, K> + Clone,
    SV2: ViewableStorageMut<E, N, K>,
    N: Dim,
    M: Dim,
    K: Dim,
    E: Scalar + Zero + One + Float,
    DefaultAllocator: Allocator<N, K> + Allocator<M, K> + Allocator<N, M> + Allocator<M, N>,
    ShapeConstraint: StridesOk<E, N, M, SM> + StridesOk<E, M, K, SV1> + StridesOk<E, N, K, SV2>,
{
    #[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<S, M, N, E> VectorSpace for Matrix<E, M, N, S>
where
    S: ViewableStorage<E, M, N>,
    M: Dim,
    N: Dim,
    E: Scalar + Zero + One + Float,
    DefaultAllocator: Allocator<M, N>,
    ShapeConstraint: StridesOk<E, M, N, S>,
{
    type Field = E;
    type Owned = OMatrix<E, M, N>;

    #[inline]
    fn similar_origin(&self) -> Self::Owned {
        let (n, m) = self.shape_generic();
        OMatrix::zeros_generic(n, m)
    }
}

impl<SM, SV1, M, N, E> AXPY<Matrix<E, M, N, SV1>> for Matrix<E, M, N, SM>
where
    SM: ViewableStorageMut<E, M, N>,
    SV1: ViewableStorage<E, M, N>,
    M: Dim,
    N: Dim,
    E: Scalar + Zero + One + Float,
    DefaultAllocator: Allocator<M, N>,
    ShapeConstraint: StridesOk<E, M, N, SM> + StridesOk<E, M, N, SV1>,
{
    #[inline]
    fn axpy<I: Instance<Matrix<E, M, N, SV1>>>(&mut self, α: E, x: I, β: E) {
        x.eval(|x̃| {
            assert_eq!(self.ncols(), x̃.ncols());
            // nalgebra does not implement axpy for matrices, and flattenining
            // also seems difficult, so loop over columns.
            for (mut y, ỹ) in self.column_iter_mut().zip(x̃.column_iter()) {
                Vector::axpy(&mut y, α, &ỹ, β)
            }
        })
    }

    #[inline]
    fn copy_from<I: Instance<Matrix<E, M, N, 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);
    }
}

/* 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: ViewableStorageMut<E, M> + Clone,
    M: Dim,
    E: Scalar + Zero + One + Float + RealField,
    DefaultAllocator: Allocator<M>,
    ShapeConstraint: StridesOk<E, M, U1, SM>,
{
    #[inline]
    fn proj_ball(self, ρ: E, exp: Linfinity) -> <Self as Space>::OwnedSpace {
        let mut owned = self.into_owned();
        owned.proj_ball_mut(ρ, exp);
        owned
    }
}

impl<SM, M, E> ProjectionMut<E, Linfinity> for Vector<E, M, SM>
where
    SM: ViewableStorageMut<E, M> + Clone,
    M: Dim,
    E: Scalar + Zero + One + Float + RealField,
    DefaultAllocator: Allocator<M>,
    ShapeConstraint: StridesOk<E, M, U1, SM>,
{
    #[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: ViewableStorage<E, N, M>,
    SV1: ViewableStorage<E, M, K> + Clone,
    SV2: ViewableStorage<E, N, K> + Clone,
    N: Dim,
    M: Dim,
    K: Dim,
    E: Scalar + Zero + One + SimdComplexField,
    DefaultAllocator: Allocator<N, K> + Allocator<M, K> + Allocator<N, M> + Allocator<M, N>,
    ShapeConstraint: StridesOk<E, N, M, SM>
        + StridesOk<E, M, K, SV1>
        + StridesOk<E, N, K, SV2>
        + StridesOk<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)
    }
}

/// 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: ViewableStorage<E, M>,
    E: Float + Scalar + Zero + One + RealField,
    DefaultAllocator: Allocator<M>,
    ShapeConstraint: StridesOk<E, M, U1, S>,
{
    type OwnedEuclidean = OVector<E, M>;

    #[inline]
    fn dot<I: Instance<Self>>(&self, other: I) -> E {
        other.eval_ref_decompose(|ref r| Vector::<E, M, S>::dot(self, r))
    }

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

    #[inline]
    fn dist2_squared<I: Instance<Self>>(&self, other: I) -> E {
        other.eval_ref_decompose(|ref r| metric_distance_squared(self, r))
    }
}

impl<E, M, S> StaticEuclidean<E> for Vector<E, M, S>
where
    M: DimName,
    S: ViewableStorage<E, M>,
    E: Float + Scalar + Zero + One + RealField,
    DefaultAllocator: Allocator<M>,
    ShapeConstraint: StridesOk<E, M, U1, S>,
{
    #[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: ViewableStorage<E, M>,
    E: Float + Scalar + Zero + One + RealField,
    DefaultAllocator: Allocator<M>,
    ShapeConstraint: StridesOk<E, M, U1, S>,
{
    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: ViewableStorage<E, M>,
    E: Float + Scalar + Zero + One + RealField,
    DefaultAllocator: Allocator<M>,
    ShapeConstraint: StridesOk<E, M, U1, S>,
{
    type DualSpace = OVector<E, M>;

    fn dual_origin(&self) -> OVector<E, M> {
        OVector::zeros_generic(M::from_usize(self.len()), Const)
    }
}

impl<E, M, S> Norm<L1, E> for Vector<E, M, S>
where
    M: Dim,
    S: ViewableStorage<E, M>,
    E: Float + Scalar + Zero + One + RealField,
    DefaultAllocator: Allocator<M>,
    ShapeConstraint: StridesOk<E, M, U1, S>,
{
    #[inline]
    fn norm(&self, _: L1) -> E {
        nalgebra::Norm::norm(&LpNorm(1), self)
    }
}

impl<E, M, S> Dist<L1, E> for Vector<E, M, S>
where
    M: Dim,
    S: ViewableStorage<E, M> + Clone,
    E: Float + Scalar + Zero + One + RealField,
    DefaultAllocator: Allocator<M>,
    ShapeConstraint: StridesOk<E, M, U1, S>,
{
    #[inline]
    fn dist<I: Instance<Self>>(&self, other: I, _: L1) -> E {
        other.eval_ref_decompose(|ref r| nalgebra::Norm::metric_distance(&LpNorm(1), self, r))
    }
}

impl<E, M, S> Norm<L2, E> for Vector<E, M, S>
where
    M: Dim,
    S: ViewableStorage<E, M>,
    E: Float + Scalar + Zero + One + RealField,
    DefaultAllocator: Allocator<M>,
    ShapeConstraint: StridesOk<E, M, U1, S>,
{
    #[inline]
    fn norm(&self, _: L2) -> E {
        nalgebra::Norm::norm(&LpNorm(2), self)
    }
}

impl<E, M, S> Dist<L2, E> for Vector<E, M, S>
where
    M: Dim,
    S: ViewableStorage<E, M> + Clone,
    E: Float + Scalar + Zero + One + RealField,
    DefaultAllocator: Allocator<M>,
    ShapeConstraint: StridesOk<E, M, U1, S>,
{
    #[inline]
    fn dist<I: Instance<Self>>(&self, other: I, _: L2) -> E {
        other.eval_ref_decompose(|ref r| nalgebra::Norm::metric_distance(&LpNorm(2), self, r))
    }
}

impl<E, M, S> Norm<Linfinity, E> for Vector<E, M, S>
where
    M: Dim,
    S: ViewableStorage<E, M>,
    E: Float + Scalar + Zero + One + RealField,
    DefaultAllocator: Allocator<M>,
    ShapeConstraint: StridesOk<E, M, U1, S>,
{
    #[inline]
    fn norm(&self, _: Linfinity) -> E {
        nalgebra::Norm::norm(&UniformNorm, self)
    }
}

impl<E, M, S> Dist<Linfinity, E> for Vector<E, M, S>
where
    M: Dim,
    S: ViewableStorage<E, M> + Clone,
    E: Float + Scalar + Zero + One + RealField,
    DefaultAllocator: Allocator<M>,
    ShapeConstraint: StridesOk<E, M, U1, S>,
{
    #[inline]
    fn dist<I: Instance<Self>>(&self, other: I, _: Linfinity) -> E {
        other.eval_ref_decompose(|ref r| nalgebra::Norm::metric_distance(&UniformNorm, self, r))
    }
}

/// 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
    }
}