automatica 1.0.0

//! Matrix and vector methods

use std::{
    any::Any,
    fmt::Debug,
    ops::{Add, Div, Mul, Sub},
};

use nalgebra::DMatrix;
use ndarray::{s, Array1, Array2, ArrayView1};

use crate::{Abs, Max, One, RelativeEq, Zero};

/// Trait for the calculation of the matrix trace
pub(crate) trait Trace<T> {
    /// Trace of a matrix, i.e. sum of the elements on the principal diagonal
    fn trace(&self) -> T;
}

impl<T> Trace<T> for Array2<T>
where
    T: Add<Output = T> + Clone + Zero,
{
    fn trace(&self) -> T {
        self.diag().fold(T::zero(), |acc, d| acc + d.clone())
    }
}

/// Multiplication trait between matrix and scalar
pub(crate) trait ScalarMul<T> {
    type Output;
    /// Multiplication between matrix and scalar
    fn smul(self, x: T) -> Self::Output;
}

/// Scalar matrix multiplication `c_ij = a_ij * x`
impl<T> ScalarMul<T> for &Array2<T>
where
    T: Clone + Mul<Output = T>,
{
    type Output = Array2<T>;
    fn smul(self, x: T) -> Self::Output {
        self.mapv(|a_ij| a_ij * x.clone())
    }
}

/// Scalar matrix inplace multiplication `a_ij = a_ij * x`
impl<T> ScalarMul<T> for Array2<T>
where
    T: Clone + Mul<Output = T>,
{
    type Output = Self;
    fn smul(mut self, x: T) -> Self::Output {
        self.mapv_inplace(|a_ij| a_ij * x.clone());
        self
    }
}

/// Scalar vector multiplication `c_i = a_i * x`
impl<T> ScalarMul<T> for &Array1<T>
where
    T: Clone + Mul<Output = T>,
{
    type Output = Array1<T>;
    fn smul(self, x: T) -> Self::Output {
        self.mapv(|v_i| v_i * x.clone())
    }
}

/// Scalar vector inplace multiplication `a_i = a_i * x`
impl<T> ScalarMul<T> for Array1<T>
where
    T: Clone + Mul<Output = T>,
{
    type Output = Self;
    fn smul(mut self, x: T) -> Self::Output {
        self.mapv_inplace(|v_i| v_i * x.clone());
        self
    }
}

/// Scalar vector multiplication `c_i = a_i * x`
impl<T> ScalarMul<T> for ArrayView1<'_, T>
where
    T: Clone + Mul<Output = T>,
{
    type Output = Array1<T>;
    fn smul(self, x: T) -> Self::Output {
        self.mapv(|v_i| v_i * x.clone())
    }
}

/// Matrix multiplication trait.
pub(crate) trait MatrixMul<Rhs = Self> {
    type Output;
    /// Matrix multiplication: `C = A * B`
    fn mmul(self, b: Rhs) -> Self::Output;
}

/// Matrix multiplication trait.
pub(crate) trait MatrixMulSca<T, Rhs = Self> {
    type Output;
    /// Matrix multiplication with scalar scale: `D = c * A * B`
    fn mmul_scale(self, b: Rhs, c: T) -> Self::Output;
}

/// Matrix multiplication for `Array2`
impl<T> MatrixMul for &Array2<T>
where
    T: Add<Output = T> + Clone + Mul<Output = T> + Zero,
{
    type Output = Array2<T>;

    #[allow(clippy::many_single_char_names)]
    fn mmul(self, b: Self) -> Self::Output {
        let a = self;
        let m = a.nrows();
        let n = a.ncols();
        let p = b.ncols();
        debug_assert_eq!(n, b.nrows());
        Array2::from_shape_fn((m, p), |(i, j)| {
            a.row(i)
                .iter()
                .zip(b.column(j))
                .fold(T::zero(), |acc, (ai, bj)| acc + ai.clone() * bj.clone())
        })
    }
}

/// Matrix multiplication for `Array2`
impl<T> MatrixMulSca<T> for &Array2<T>
where
    T: Add<Output = T> + Clone + Mul<Output = T> + Zero,
{
    type Output = Array2<T>;

    #[allow(clippy::many_single_char_names)]
    fn mmul_scale(self, b: &Array2<T>, c: T) -> Array2<T>
    where
        T: Add<Output = T> + Clone + Mul<Output = T> + Zero,
    {
        let a = self;
        let m = a.nrows();
        let n = a.ncols();
        let p = b.ncols();
        debug_assert_eq!(n, b.nrows());
        Array2::from_shape_fn((m, p), |(i, j)| {
            a.row(i)
                .iter()
                .zip(b.column(j))
                .fold(T::zero(), |acc, (ai, bj)| acc + ai.clone() * bj.clone())
                * c.clone()
        })
    }
}

impl<T> MatrixMul<Array2<T>> for &Array2<T>
where
    T: Add<Output = T> + Clone + Mul<Output = T> + Zero,
{
    type Output = Array2<T>;

    #[allow(clippy::many_single_char_names)]
    fn mmul(self, b: Array2<T>) -> Self::Output {
        self.mmul(&b)
    }
}

/// Matrix-vector multiplication for `Array2` and `Array1`
impl<T> MatrixMul<&Array1<T>> for &Array2<T>
where
    T: Add<Output = T> + Clone + Mul<Output = T> + Zero,
{
    type Output = Array1<T>;

    #[allow(clippy::many_single_char_names)]
    fn mmul(self, b: &Array1<T>) -> Self::Output {
        let a = self;
        let m = a.nrows();
        let n = a.ncols();
        debug_assert_eq!(n, b.len());
        Array1::from_shape_fn(m, |i| {
            a.row(i)
                .iter()
                .zip(b)
                .fold(T::zero(), |acc, (ai, bj)| acc + ai.clone() * bj.clone())
        })
    }
}

/// Matrix-vector multiplication for `ndarray::Array2` and `ndarray::Array1`
impl<T> MatrixMul<Array1<T>> for &Array2<T>
where
    T: Add<Output = T> + Clone + Mul<Output = T> + Zero,
{
    type Output = Array1<T>;

    #[allow(clippy::many_single_char_names)]
    fn mmul(self, b: Array1<T>) -> Self::Output {
        self.mmul(&b)
    }
}

/// Matrix-vector multiplication for `Array2` and `slice`
impl<T> MatrixMul<&[T]> for &Array2<T>
where
    T: Add<Output = T> + Clone + Mul<Output = T> + Zero,
{
    type Output = Array1<T>;

    #[allow(clippy::many_single_char_names)]
    fn mmul(self, b: &[T]) -> Self::Output {
        let a = self;
        let m = a.nrows();
        let n = a.ncols();
        debug_assert_eq!(n, b.len());
        Array1::from_shape_fn(m, |i| {
            a.row(i)
                .iter()
                .zip(b)
                .fold(T::zero(), |acc, (ai, bj)| acc + ai.clone() * bj.clone())
        })
    }
}

/// Maximum value of the absolute values of the elements of the array.
pub(crate) fn max_vabs<T>(mut v: Array1<T>) -> T
where
    T: Abs + Max + Clone,
{
    v.mapv_inplace(|x| x.abs());
    let mut it = v.iter();
    // Safe unwrap since there is always an element in the array.
    let first = it.next().unwrap().clone();
    it.fold(first, |acc, x| acc.max(x))
}

/// Relative equality between arrays element-wise.
pub(crate) fn relative_eq<T>(a: &Array1<T>, b: &Array1<T>, epsilon: &T) -> bool
where
    T: Clone + RelativeEq,
{
    a.iter().zip(b).all(|(x, y)| x.relative_eq(y, epsilon))
}

/// Identity square matrix of size n, all 1 in the principal diagonal, 0 elsewhere
pub(crate) fn identity<T>(n: usize) -> Array2<T>
where
    T: Clone + One + Zero,
{
    let mut temp = Array2::from_elem((n, n), T::zero());
    for d in temp.diag_mut() {
        *d = T::one();
    }
    temp
}

/// LUP decomposition
#[derive(Clone, Debug)]
pub(crate) struct Lup<T> {
    /// LU matrix
    lu: Array2<T>,
    /// Permutation vector
    p: Array1<usize>,
}

impl<T> Lup<T>
where
    T: Add<Output = T> + Clone + Div<Output = T> + Mul<Output = T> + Sub<Output = T> + Zero,
{
    /// Solve the linear system: `Ax=b => LUx=Pb`
    pub(crate) fn solve(&self, b: &Array1<T>) -> Array1<T> {
        let n = self.lu.nrows();
        let mut y = vec![T::zero(); n];
        let mut x = Array1::from_elem(n, T::zero());
        self.forward_substitution(n, b, &mut y);
        self.backward_substitution(n, &y, &mut x);
        x
    }

    /// Forward substitution method
    fn forward_substitution(&self, n: usize, b: &Array1<T>, y: &mut [T]) {
        y[0] = b[self.p[0]].clone();
        for i in 1..n {
            y[i] = b[self.p[i]].clone()
                - (0..i).fold(T::zero(), |acc, j| {
                    acc + self.lu[(i, j)].clone() * y[j].clone()
                });
        }
    }

    /// Backward substitution method
    fn backward_substitution(&self, n: usize, y: &[T], x: &mut Array1<T>) {
        x[n - 1] = y[n - 1].clone() / self.lu[(n - 1, n - 1)].clone();
        for i in (0..(n - 1)).rev() {
            x[i] = (y[i].clone()
                - ((i + 1)..n).fold(T::zero(), |acc, j| {
                    acc + self.lu[(i, j)].clone() * x[j].clone()
                }))
                / self.lu[(i, i)].clone();
        }
    }

    /// Solve the linear system: `Ax=b => LUx=Pb`,
    /// the result `x` is written inside `b`.
    pub(crate) fn solve_mut(&self, b: &mut Array1<T>) {
        let n = self.lu.nrows();
        let mut y = vec![T::zero(); n];
        self.forward_substitution(n, b, &mut y);
        let x = b; // Store `x` vector in `b`.
        self.backward_substitution(n, &y, x);
    }
}

/// LUP decomposition `PA=LU`
pub(crate) fn lup<T>(m: Array2<T>) -> Option<Lup<T>>
where
    T: Abs
        + Add<Output = T>
        + Clone
        + Div<Output = T>
        + Mul<Output = T>
        + PartialOrd
        + Sub<Output = T>
        + Zero,
{
    let mut a = m;
    let n = a.nrows();
    let mut pi: Array1<_> = (0_usize..n).collect();
    for k in 0_usize..n {
        let mut p = T::zero();
        let mut kp = 0;
        for i in k..n {
            let aik = a[(i, k)].abs();
            if aik > p {
                p = aik;
                kp = i;
            }
        }
        if p.is_zero() {
            return None;
        }
        pi.swap(k, kp);
        for i in 0..n {
            a.swap((k, i), (kp, i));
        }
        for i in (k + 1)..n {
            a[(i, k)] = a[(i, k)].clone() / a[(k, k)].clone();
            for j in (k + 1)..n {
                a[(i, j)] = a[(i, j)].clone() - a[(i, k)].clone() * a[(k, j)].clone();
            }
        }
    }
    Some(Lup { lu: a, p: pi })
}

/// Matrix inversion through LUP decomposition.
pub(crate) fn inverse<T>(a: Array2<T>) -> Option<Array2<T>>
where
    T: Abs
        + Add<Output = T>
        + Clone
        + Div<Output = T>
        + Mul<Output = T>
        + One
        + PartialOrd
        + Sub<Output = T>
        + Zero,
{
    let n = a.nrows();
    let lup = lup(a)?;
    let mut inv = Array2::from_elem((n, n), T::zero());
    for i in 0..n {
        let x = lup.solve(&Array1::from_shape_fn(n, |j| {
            if j == i {
                T::one()
            } else {
                T::zero()
            }
        }));
        inv.slice_mut(s![.., i]).assign(&x);
    }
    Some(inv)
}

/// From `Array2` to `DMatrix`
pub(crate) fn ar_to_dm<T>(ar: &Array2<T>) -> DMatrix<T>
where
    T: Any + Copy + Debug + PartialEq,
{
    DMatrix::from_iterator(ar.nrows(), ar.ncols(), ar.t().iter().copied())
}

#[cfg(test)]
mod tests {
    use ndarray::{arr1, arr2};

    use super::*;

    #[test]
    #[allow(clippy::float_cmp)]
    fn matrix_trace() {
        let m11 = -13.;
        let m22 = 7.;
        let m = arr2(&[[m11, 3.], [1., m22]]);
        assert_eq!(m11 + m22, m.trace());
    }

    #[test]
    fn scalar_matrix_multiplication_ref() {
        let a = arr2(&[[2., 4.], [6., 8.]]);
        let x = 4.;
        let c = ScalarMul::smul(&a, x);
        assert_eq!(c, x * a);
    }

    #[test]
    fn scalar_matrix_multiplication_value() {
        let a = arr2(&[[2., 4.], [6., 8.]]);
        let x = 4.;
        let c = ScalarMul::smul(a.clone(), x);
        assert_eq!(c, x * a);
    }

    #[test]
    fn scalar_vector_multiplication() {
        let a = arr1(&[2., 4., 6., 8.]);
        let x = 4.;
        let c = (&a).smul(x);
        assert_eq!(c, x * &a);
    }

    #[test]
    #[allow(clippy::many_single_char_names)]
    fn matrix_matrix_multiplication() {
        let a = arr2(&[[2., 4.], [6., 8.]]);
        let b = arr2(&[[2., 0.], [0., 2.]]);
        let c = a.mmul(&b);
        assert_eq!(c, a.dot(&b));
        let d = arr2(&[[-3.], [2.]]);
        let e = c.mmul(&d);
        assert_eq!(e, c.dot(&d));
    }

    #[test]
    fn matrix_matrix_multiplication_value() {
        let a = arr2(&[[2., 4.], [6., 8.]]);
        let b = arr2(&[[2., 0.], [0., 2.]]);
        let c = a.mmul(b.clone());
        assert_eq!(c, a.dot(&b));
    }

    #[test]
    fn matrix_vector_multiplication() {
        let a = arr2(&[[2., 4.], [6., 8.]]);
        let b = arr1(&[2., -3.]);
        let c = a.mmul(&b);
        assert_eq!(c, a.dot(&b));
    }

    #[test]
    #[allow(clippy::float_cmp)]
    fn maximum_absolute_value() {
        let expected = 28.;
        let v = arr1(&[2., -expected, 0.32, 8.123]);
        assert_eq!(expected, max_vabs(v));
    }

    #[test]
    #[allow(clippy::float_cmp)]
    fn identity_matrix() {
        let id = identity::<f32>(10);
        let res = id.mmul(&id);
        assert_eq!(res[(0, 0)], 1.);
        assert_eq!(id, res);
    }

    #[test]
    fn lup_decomposition() {
        // Cormen Figure 28.2
        let a = arr2(&[
            [2., 0., 2., 0.6],
            [3., 3., 4., -2.],
            [5., 5., 4., 2.],
            [-1., -2., 3.4, -1.],
        ]);
        let lup = lup(a).unwrap();

        assert_relative_eq!(5., lup.lu[(0, 0)], max_relative = 1e-10);
        assert_relative_eq!(5., lup.lu[(0, 1)], max_relative = 1e-10);
        assert_relative_eq!(4., lup.lu[(0, 2)], max_relative = 1e-10);
        assert_relative_eq!(2., lup.lu[(0, 3)], max_relative = 1e-10);
        assert_relative_eq!(0.4, lup.lu[(1, 0)], max_relative = 1e-10);
        assert_relative_eq!(-2., lup.lu[(1, 1)], max_relative = 1e-10);
        assert_relative_eq!(0.4, lup.lu[(1, 2)], max_relative = 1e-10);
        assert_relative_eq!(-0.2, lup.lu[(1, 3)], max_relative = 1e-10);
        assert_relative_eq!(-0.2, lup.lu[(2, 0)], max_relative = 1e-10);
        assert_relative_eq!(0.5, lup.lu[(2, 1)], max_relative = 1e-10);
        assert_relative_eq!(4., lup.lu[(2, 2)], max_relative = 1e-10);
        assert_relative_eq!(-0.5, lup.lu[(2, 3)], max_relative = 1e-10);
        assert_relative_eq!(0.6, lup.lu[(3, 0)], max_relative = 1e-10);
        assert_relative_eq!(0., lup.lu[(3, 1)], max_relative = 1e-10);
        assert_relative_eq!(0.4, lup.lu[(3, 2)], max_relative = 1e-10);
        assert_relative_eq!(-3., lup.lu[(3, 3)], max_relative = 1e-10);

        assert_eq!(2, lup.p[0]);
        assert_eq!(0, lup.p[1]);
        assert_eq!(3, lup.p[2]);
        assert_eq!(1, lup.p[3]);
    }

    #[test]
    fn system_solve() {
        // Cormen paragraph 28.1
        let a = arr2(&[[1., 2., 0.], [3., 4., 4.], [5., 6., 3.]]);
        let lup = lup(a).unwrap();

        assert_relative_eq!(5., lup.lu[(0, 0)], max_relative = 1e-10);
        assert_relative_eq!(6., lup.lu[(0, 1)], max_relative = 1e-10);
        assert_relative_eq!(3., lup.lu[(0, 2)], max_relative = 1e-10);
        assert_relative_eq!(0.2, lup.lu[(1, 0)], max_relative = 1e-10);
        assert_relative_eq!(0.8, lup.lu[(1, 1)], max_relative = 1e-10);
        assert_relative_eq!(-0.6, lup.lu[(1, 2)], max_relative = 1e-10);
        assert_relative_eq!(0.6, lup.lu[(2, 0)], max_relative = 1e-10);
        assert_relative_eq!(0.5, lup.lu[(2, 1)], max_relative = 1e-10);
        assert_relative_eq!(2.5, lup.lu[(2, 2)], max_relative = 1e-10);

        assert_eq!(2, lup.p[0]);
        assert_eq!(0, lup.p[1]);
        assert_eq!(1, lup.p[2]);

        let x = lup.solve(&arr1(&[3., 7., 8.]));

        assert_relative_eq!(-1.4, x[0], max_relative = 1e-10);
        assert_relative_eq!(2.2, x[1], max_relative = 1e-10);
        assert_relative_eq!(0.6, x[2], max_relative = 1e-10);
    }

    #[test]
    fn matrix_inverse() {
        // Wikipedia.
        let a = arr2(&[[1., 2.], [2., 3.]]);
        let inv = arr2(&[[-3., 2.], [2., -1.]]);
        assert_eq!(inv, inverse(a).unwrap());
    }

    #[test]
    fn smul() {}
}