automatica 1.0.0

//! # Matrices of polynomials
//!
//! `MatrixOfPoly` allows the definition of matrices of polynomials.

use std::{
    fmt::{self, Debug, Display, Formatter},
    ops::{Add, AddAssign, Index, IndexMut, Mul, Sub},
};

use ndarray::Array2;

use crate::{
    matrix::{MatrixMul, ScalarMul},
    wrappers::PP,
    Complex, Polynomial, Zero,
};

/// Polynomial matrix object
///
/// Contains the vector of coefficients form the lowest to the highest degree
///
/// P(x) = C0 + C1*x + C2*x^2 + ... =
///   ┌ ┌       ┐ ┌       ┐ ┌       ┐     ┐
/// = │ │a01 a02│,│a11 a12│,│a21 a22│, ...│
///   │ │a03 a04│ │a13 a14│ │a23 a24│     │
///   └ └       ┘ └       ┘ └       ┘     ┘
#[derive(Clone, Debug)]
pub(crate) struct PolyMatrix<T> {
    matrix_coeffs: Vec<Array2<T>>,
}

impl<T> PolyMatrix<T> {
    /// Degree of the polynomial matrix
    fn degree(&self) -> usize {
        assert!(
            !self.matrix_coeffs.is_empty(),
            "Degree is not defined on empty polynomial matrix"
        );
        self.matrix_coeffs.len() - 1
    }
}

impl<T> PolyMatrix<T>
where
    T: Clone + PartialEq + Zero,
{
    /// Create a new polynomial matrix given a slice of matrix coefficients.
    ///
    /// # Arguments
    ///
    /// * `matrix_coeffs` - slice of matrix coefficients
    pub(crate) fn new_from_coeffs(matrix_coeffs: &[Array2<T>]) -> Self {
        let shape = matrix_coeffs[0].shape();
        assert!(matrix_coeffs.iter().all(|c| c.shape() == shape));
        let mut pm = Self {
            matrix_coeffs: matrix_coeffs.into(),
        };
        pm.trim();
        debug_assert!(!pm.matrix_coeffs.is_empty());
        pm
    }

    /// Create a new polynomial matrix given an iterator of matrix coefficients.
    ///
    /// # Arguments
    ///
    /// * `matrix_iter` - iterator of matrix coefficients
    pub(crate) fn new_from_iter<II>(matrix_iter: II) -> Self
    where
        II: IntoIterator<Item = Array2<T>>,
    {
        let mut pm = Self {
            matrix_coeffs: matrix_iter.into_iter().collect(),
        };
        debug_assert!(!pm.matrix_coeffs.is_empty());
        let shape = pm.matrix_coeffs[0].shape();
        assert!(pm.matrix_coeffs.iter().all(|c| c.shape() == shape));
        pm.trim();
        pm
    }

    /// Trim the zeros coefficients of high degree terms
    fn trim(&mut self) {
        let rows = self.matrix_coeffs[0].nrows();
        let cols = self.matrix_coeffs[0].ncols();
        let zero = Array2::from_elem((rows, cols), T::zero());
        if let Some(p) = self.matrix_coeffs.iter().rposition(|c| c != zero) {
            self.matrix_coeffs.truncate(p + 1);
        } else {
            self.matrix_coeffs.resize(1, zero);
        }
    }
}

impl<T> PolyMatrix<T>
where
    T: Clone + Mul<Output = T> + PartialEq + Zero,
{
    /// Implementation of polynomial and matrix multiplication
    /// It's the polynomial matrix whose coefficients are the coefficients
    /// of the polynomial times the matrix
    ///
    /// # Arguments
    ///
    /// * `poly` - Polynomial
    /// * `matrix` - Matrix
    pub(crate) fn multiply(poly: &PP<T>, matrix: &Array2<T>) -> Self {
        let result = poly.0.as_slice().iter().map(|c| matrix.smul(c.0.clone()));
        Self::new_from_iter(result)
    }
}

impl<T> PolyMatrix<T>
where
    T: Add<Output = T> + Clone + Mul<Output = T> + PartialEq + Zero,
{
    /// Implementation of polynomial matrix and matrix multiplication
    ///
    /// `PolyMatrix` * `Matrix`
    pub(crate) fn right_mul(&self, rhs: &Array2<T>) -> Self {
        let result: Vec<_> = self.matrix_coeffs.iter().map(|l| l.mmul(rhs)).collect();
        Self::new_from_coeffs(&result)
    }

    /// Implementation of matrix and polynomial matrix multiplication
    ///
    /// `Matrix` * `PolyMatrix`
    pub(crate) fn left_mul(&self, lhs: &Array2<T>) -> Self {
        let res: Vec<_> = self.matrix_coeffs.iter().map(|r| lhs.mmul(r)).collect();
        Self::new_from_coeffs(&res)
    }
}

impl<T> PolyMatrix<T>
where
    T: Add<Output = T> + Clone + Mul<Output = T> + Sub<Output = T> + Zero,
{
    #[allow(dead_code)]
    /// Evaluate the polynomial matrix
    ///
    /// # Arguments
    /// * `s` - Matrix at which the polynomial matrix is evaluated.
    pub(crate) fn eval(&self, s: &Array2<Complex<T>>) -> Array2<Complex<T>> {
        // transform matrix_coeffs in complex numbers matrices
        //
        // ┌     ┐ ┌       ┐ ┌       ┐ ┌     ┐ ┌       ┐ ┌         ┐
        // │P1 P2│=│a01 a02│+│a11 a12│*│s1 s2│+│a21 a22│*│s1^2 s2^2│
        // │P3 P4│ │a03 a04│ │a13 a14│ │s1 s2│ │a23 a24│ │s1^2 s2^2│
        // └     ┘ └       ┘ └       ┘ └     ┘ └       ┘ └         ┘
        // `*` is the element by element multiplication
        // If i have a 2x2 matrix_coeff the result shall be a 2x2 matrix,
        // because otherwise i will sum P1 and P2 (P3 and P4)
        let rows = self.matrix_coeffs[0].nrows();
        let cols = self.matrix_coeffs[0].ncols();

        let mut result = Array2::from_elem((rows, cols), Complex::<T>::zero());

        for mc in self.matrix_coeffs.iter().rev() {
            let mcplx = mc.map(|x| Complex::new(x.clone(), T::zero()));
            result = result * s + mcplx;
        }
        result
    }
}

/// Implementation of polynomial matrices addition
impl<T> Add for PolyMatrix<T>
where
    T: Add<Output = T> + AddAssign + Clone + PartialEq + Zero,
{
    type Output = Self;

    fn add(mut self, mut rhs: Self) -> Self {
        // Check which polynomial matrix has the highest degree
        let mut result = if self.degree() < rhs.degree() {
            for (i, c) in self.matrix_coeffs.iter().enumerate() {
                rhs[i] += c;
            }
            rhs
        } else {
            for (i, c) in rhs.matrix_coeffs.iter().enumerate() {
                self[i] += c;
            }
            self
        };
        result.trim();
        result
    }
}

/// Implementation of read only indexing of polynomial matrix
/// returning its coefficients.
///
/// # Panics
///
/// Panics for out of bounds access.
impl<T> Index<usize> for PolyMatrix<T> {
    type Output = Array2<T>;

    fn index(&self, i: usize) -> &Self::Output {
        &self.matrix_coeffs[i]
    }
}

/// Implementation of mutable indexing of polynomial matrix
/// returning its coefficients.
///
/// # Panics
///
/// Panics for out of bounds access.
impl<T> IndexMut<usize> for PolyMatrix<T> {
    fn index_mut(&mut self, i: usize) -> &mut Self::Output {
        &mut self.matrix_coeffs[i]
    }
}

/// Implementation of polynomial matrix printing
impl<T: Display + PartialEq + Zero> Display for PolyMatrix<T> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        if self.degree() == 0 {
            return write!(f, "{}", self.matrix_coeffs[0]);
        }
        let mut s = String::new();
        let mut sep = "";
        for (i, c) in self.matrix_coeffs.iter().enumerate() {
            if c.iter().all(|x| x == &T::zero()) {
                continue;
            }
            s.push_str(sep);
            if i == 0 {
                s.push_str(&format!("{}", c));
            } else if i == 1 {
                s.push_str(&format!("+{}*s", c));
            } else {
                s.push_str(&format!("+{}*s^{}", c, i));
            }
            sep = " ";
        }

        write!(f, "{}", s)
    }
}

/// Polynomial matrix object
///
/// Contains the matrix of polynomials
///
/// `P(x) = [[P1, P2], [P3, P4]]`
#[derive(Debug)]
pub struct MatrixOfPoly<T> {
    matrix: Array2<PP<T>>,
}

/// Implementation methods for MP struct
impl<T> MatrixOfPoly<T> {
    /// Create a new polynomial matrix given a vector of polynomials.
    ///
    /// # Arguments
    ///
    /// * `rows` - number of rows of the matrix
    /// * `cols` - number of columns of the matrix
    /// * `data` - vector of polynomials in row major order
    ///
    /// # Panics
    ///
    /// Panics if the matrix cannot be build from given arguments.
    fn new(rows: usize, cols: usize, data: Vec<PP<T>>) -> Self {
        Self {
            matrix: Array2::from_shape_vec((rows, cols), data)
                .expect("Input data do not allow to create the matrix"),
        }
    }

    /// Get access to a reference of the internal matrix.
    pub(crate) fn matrix(&self) -> &Array2<PP<T>> {
        &self.matrix
    }

    /// Extract the polynomial from the matrix if is the only one.
    pub(crate) fn single(&self) -> Option<&PP<T>> {
        if self.matrix.shape() == [1, 1] {
            self.matrix.first()
        } else {
            None
        }
    }
}

/// Implement read only indexing of the matrix of polynomials.
///
/// # Panics
///
/// Panics for out of bounds access.
impl<T> Index<[usize; 2]> for MatrixOfPoly<T> {
    type Output = PP<T>;

    fn index(&self, i: [usize; 2]) -> &Self::Output {
        &self.matrix[i]
    }
}

impl<T> MatrixOfPoly<T>
where
    T: Clone,
{
    /// Get a copy of the polynomial of the matrix at index `i`
    ///
    /// # Panics
    ///
    /// Panics for out of bounds access.
    #[must_use]
    pub fn get_unchecked(&self, i: [usize; 2]) -> impl Polynomial<T> {
        self.matrix[i].to_unwrapped_vec()
    }
}

impl<T> MatrixOfPoly<T>
where
    T: Clone + PartialEq + Zero,
{
    /// Set a polynomial of the matrix at index `i`
    ///
    /// # Panics
    ///
    /// Panics for out of bounds access.
    pub fn set_unchecked(&mut self, i: [usize; 2], value: &impl Polynomial<T>) {
        self.matrix[i] = PP::new_from_coeffs(value.as_slice());
    }
}

/// Implement conversion between different representations.
impl<T: Clone + PartialEq + Zero> From<PolyMatrix<T>> for MatrixOfPoly<T> {
    fn from(pm: PolyMatrix<T>) -> Self {
        // coeffs :
        // ┌ ┌       ┐ ┌       ┐    ┌       ┐        ┐
        // │ │a01 a02│+│a11 a12│*x +│a21 a22│*x^2 ...│
        // │ │a03 a04│ │a13 a14│    │a23 a24│        │
        // └ └       ┘ └       ┘    └       ┘        ┘
        let coeffs = pm.matrix_coeffs; // vector of matrices
        let rows = coeffs[0].nrows();
        let cols = coeffs[0].ncols();

        // Each vector contains the corresponding matrix in "coeffs" variable,
        // so each vector contains the coefficients of the polynomial
        // with the same order (increasing).
        // vectorized_coeffs:
        // [[a01 a02 a03 a04] + [a11 a12 a13 a14]*x + [a21 a22 a23 a24]*x^2 ...]
        let vectorized_coeffs: Vec<Vec<_>> = coeffs
            .iter()
            .map(|c| c.iter().cloned().collect::<Vec<_>>())
            .collect();

        // Create a vector containing the vectors of coefficients of a single
        // polynomial in row major mode with respect to the initial
        // vector of matrices.
        // tmp:
        // [[a01 a11 a21], [a02 a12 a22], [a03 a13 a23], [a04 a14 a24]]
        let mut tmp: Vec<Vec<T>> = vec![vec![]; rows * cols];
        for order in vectorized_coeffs {
            for (i, value) in order.into_iter().enumerate() {
                tmp[i].push(value);
            }
        }

        let polys: Vec<_> = tmp.iter().map(|p| PP::new_from_coeffs(p)).collect();
        Self::new(rows, cols, polys)
    }
}

/// Implementation of matrix of polynomials printing
impl<T> Display for MatrixOfPoly<T>
where
    T: Display + PartialOrd + Zero,
{
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "{}", self.matrix)
    }
}

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

    use crate::wrappers::PWrapper;

    use super::*;

    #[test]
    fn trim_empty() {
        let v = vec![Array2::<f32>::zeros((1, 2)), Array2::zeros((1, 2))];
        let pm = PolyMatrix::new_from_coeffs(&v);
        assert_eq!(Array2::<f32>::zeros((1, 2)), pm.matrix_coeffs[0]);
    }

    #[test]
    fn right_mul() {
        let v = vec![
            arr2(&[[1.0_f32, 2.], [3., 4.]]),
            arr2(&[[1., 0.], [0., 1.]]),
        ];
        let pm = PolyMatrix::new_from_coeffs(&v);
        let res = pm.right_mul(&arr2(&[[5., 0.], [0., 5.]]));
        assert_eq!(arr2(&[[5., 10.], [15., 20.]]), res.matrix_coeffs[0]);
        dbg!(&res);
    }

    #[test]
    fn left_mul() {
        let v = vec![
            arr2(&[[1.0_f32, 2.], [3., 4.]]),
            arr2(&[[1., 0.], [0., 1.]]),
        ];
        let pm = PolyMatrix::new_from_coeffs(&v);
        let res = pm.left_mul(&arr2(&[[5., 0.], [0., 5.]]));
        assert_eq!(arr2(&[[5., 10.], [15., 20.]]), res[0]);
    }

    #[test]
    fn eval() {
        let v = vec![
            arr2(&[[1.0_f32, 2.], [3., 4.]]),
            arr2(&[[1., 0.], [0., 1.]]),
        ];
        let pm = PolyMatrix::new_from_coeffs(&v);
        let res = pm.eval(&arr2(&[
            [Complex::new(5., 0.), Complex::zero()],
            [Complex::zero(), Complex::new(0., 5.)],
        ]));
        assert_eq!(
            arr2(&[
                [Complex::new(6., 0.), Complex::new(2., 0.)],
                [Complex::new(3., 0.), Complex::new(4., 5.)]
            ]),
            res
        );
    }

    #[test]
    fn add() {
        let v = vec![
            arr2(&[[1.0_f64, 2.], [3., 4.]]),
            arr2(&[[1., 0.], [0., 1.]]),
        ];
        let pm = PolyMatrix::new_from_coeffs(&v);
        let res = pm.clone() + pm;
        assert_eq!(arr2(&[[2., 4.], [6., 8.]]), res[0]);
        assert_eq!(arr2(&[[2., 0.], [0., 2.]]), res[1]);
    }

    #[test]
    fn format_polymatrix_zero_degree() {
        let v = vec![arr2(&[[1., 0.]])];
        let pm = PolyMatrix::new_from_coeffs(&v);
        let expected = pm.matrix_coeffs[0].to_string();
        assert_eq!(expected, format!("{}", &pm));
    }

    #[test]
    fn format_polymatrix_zero_coefficient() {
        let v = vec![Array2::<f32>::zeros((1, 2)), arr2(&[[1., 0.]])];
        let pm = PolyMatrix::new_from_coeffs(&v);
        assert!(format!("{}", &pm).starts_with('+'));
    }

    #[test]
    fn mp_creation() {
        let c = [4.3, 5.32];
        let p = PP::new_from_coeffs(&c);
        let v = vec![p.clone(), p.clone(), p.clone(), p];
        let mp = MatrixOfPoly::new(2, 2, v);
        let expected = "[[4.3 +5.32s, 4.3 +5.32s],\n [4.3 +5.32s, 4.3 +5.32s]]";
        assert_eq!(expected, format!("{}", &mp));
    }

    #[test]
    #[allow(clippy::float_cmp)]
    fn indexing() {
        let c = [4.3, 5.32];
        let p = PP::new_from_coeffs(&c);
        let v = vec![p.clone(), p.clone(), p.clone(), p];
        let mut mp = MatrixOfPoly::new(2, 2, v);
        assert_eq!(c, mp.get_unchecked([1, 1]).as_slice());
        let p2 = vec![1., 2.];
        mp.set_unchecked([1, 1], &p2);
        assert_eq!(p2, mp.get_unchecked([1, 1]).as_slice());
    }

    #[test]
    fn single() {
        let v = vec![PP::new_from_coeffs(&[4.3, 5.32])];
        let mp = MatrixOfPoly::new(1, 1, v);
        assert_eq!(1, mp.matrix().len());
        let res = mp.single();
        assert!(res.is_some());
        assert_relative_eq!(
            14.94,
            res.unwrap().0.eval(&PWrapper(2.)).0,
            max_relative = 1e-10
        );
    }

    #[test]
    fn single_fail() {
        let c = [4.3, 5.32];
        let p = PP::new_from_coeffs(&c);
        let v = vec![p.clone(), p.clone(), p.clone(), p];
        let mp = MatrixOfPoly::new(2, 2, v);
        let res = mp.single();
        assert!(res.is_none());
    }

    #[test]
    fn matrixofpoly_index() {
        let c = [4.3, 5.32];
        let p = PP::new_from_coeffs(&c);
        let v = vec![p.clone(), p.clone(), p.clone(), p.clone()];
        let mp = MatrixOfPoly::new(2, 2, v);
        assert_eq!(p, mp[[1, 0]]);
    }
}