lineal-rs/src/types/matrix/generic.rs

use std::ops::{Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Neg, Sub, SubAssign};
use std::ptr::swap;

use crate::{Float, Numeric, Primitive};
#[cfg(feature = "angular")]
use crate::Angular;

/// Struct representing a dense matrix.
#[repr(transparent)]
#[derive(Copy, Clone, Debug, PartialEq)]
pub struct GenericMatrix<T: Numeric, const ROWS: usize, const COLUMNS: usize> {
    pub data: [[T; COLUMNS]; ROWS],
}

/// Special kind of `GenericMatrix` where both dimensions are equal.
pub type SquareMatrix<T, const DIMENSION: usize> = GenericMatrix<T, DIMENSION, DIMENSION>;

/// Special kind of `GenericMatrix` with just one column.
pub type ColumnVector<T, const ROWS: usize> = GenericMatrix<T, ROWS, 1>;

/// Special kind of `GenericMatrix` with just one row.
pub type RowVector<T, const COLUMNS: usize> = GenericMatrix<T, 1, COLUMNS>;

/// A macro for easier creation of row vectors.
///
/// # Example
///
/// ```rust
/// # use lineal::{RowVector, rvec};
/// let a = rvec![1.0, 2.0, 3.0];
/// // is the same as
/// let b = RowVector {
///     data: [[1.0, 2.0, 3.0]],
/// };
/// ```
#[macro_export]
macro_rules! rvec {
    ($($value:expr),* $(,)?) => {
        {
            RowVector {
                data: [[$( $value, )*]],
            }
        }
    }
}

/// A macro for easier creation of column vectors.
///
/// # Example
///
/// ```rust
/// # use lineal::{ColumnVector, cvec};
/// let a = cvec![1.0, 2.0, 3.0];
/// // is the same as
/// let b = ColumnVector {
///     data: [[1.0], [2.0], [3.0]],
/// };
/// ```
#[macro_export]
macro_rules! cvec {
    ($($value:expr),* $(,)?) => {
        {
            ColumnVector {
                data: [$( [$value], )*],
            }
        }
    }
}

impl<T: Numeric, const ROWS: usize, const COLUMNS: usize> Default for GenericMatrix<T, ROWS, COLUMNS> {
    /// Create a matrix with all cells being zero.
    fn default() -> Self {
        let zero = T::whole(0);
        Self {
            data: [[zero; COLUMNS]; ROWS],
        }
    }
}

impl<T: Numeric, const ROWS: usize, const COLUMNS: usize> GenericMatrix<T, ROWS, COLUMNS> {
    /// Create a matrix with all cells being zero.
    pub fn new() -> Self {
        Self::default()
    }

    /// Create a transpose of the input matrix.
    ///
    /// # Example
    ///
    /// ```rust
    /// # use lineal::{ColumnVector, RowVector, cvec, rvec};
    /// let a = cvec![1, 2, 3];
    /// let b = rvec![1, 2, 3];
    /// assert_eq!(a.transposed(), b);
    /// ```
    pub fn transposed(&self) -> GenericMatrix<T, COLUMNS, ROWS> {
        let mut result = GenericMatrix::default();
        for r in 0..ROWS {
            for c in 0..COLUMNS {
                result.data[c][r] = self.data[r][c];
            }
        }
        result
    }
}

impl<T: Numeric, const DIMENSION: usize> SquareMatrix<T, DIMENSION> {
    /// Create a square identity matrix.
    ///
    /// In an identity matrix the main diagonal is filled with ones,
    /// while the rest if the cells are zero.
    ///
    /// # Example
    ///
    /// ```rust
    /// # use lineal::{SquareMatrix};
    /// let a: SquareMatrix<f32, 3> = SquareMatrix::identity();
    /// // is the same as
    /// let b: SquareMatrix<f32, 3> = SquareMatrix {
    ///     data: [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
    /// };
    /// ```
    pub fn identity() -> Self {
        let mut mat = Self::default();
        let one = T::whole(1);
        for i in 0..DIMENSION {
            mat.data[i][i] = one;
        }
        mat
    }

    /// Create a square matrix with its main diagonal filled with the given values.
    ///
    /// # Example
    ///
    /// ```rust
    /// # use lineal::{SquareMatrix};
    /// let a = SquareMatrix::diagonal(&[1, 2, 3]);
    /// // is the same as
    /// let b = SquareMatrix {
    ///     data: [
    ///         [1, 0, 0],
    ///         [0, 2, 0],
    ///         [0, 0, 3],
    ///     ],
    /// };
    /// ```
    pub fn diagonal(values: &[T; DIMENSION]) -> Self {
        let mut mat = Self::default();
        for i in 0..DIMENSION {
            mat.data[i][i] = values[i];
        }
        mat
    }
}

impl<T: Numeric + Float + PartialOrd, const DIMENSION: usize> SquareMatrix<T, DIMENSION> {
    /// Apply an LUP decomposition to the matrix.
    ///
    /// # Safety
    ///
    /// This function is unsafe, because even when the result is false,
    /// the matrix may already have been modified.
    pub unsafe fn lup_decompose(&mut self, pivot: &mut Vec<usize>, tolerance: T) -> bool {
        pivot.clear();
        pivot.reserve(DIMENSION + 1);
        for i in 0..(DIMENSION + 1) {
            pivot.push(i);
        }

        let zero = T::whole(0);

        for i in 0..DIMENSION {
            let mut max_a = zero;
            let mut max_i = i;

            for k in i..DIMENSION {
                let abs_a = self[(k, i)].abs();
                if abs_a > max_a {
                    max_a = abs_a;
                    max_i = k;
                }
            }

            if max_a <= tolerance {
                return false;
            }

            if max_i != i {
                pivot.swap(i, max_i);

                let p_a: *mut [T; DIMENSION] = (&mut self.data[i]) as *mut [T; DIMENSION];
                let p_a_max: *mut [T; DIMENSION] = (&mut self.data[max_i]) as *mut [T; DIMENSION];
                swap(p_a, p_a_max);

                pivot[DIMENSION] += 1;
            }

            for j in (i + 1)..DIMENSION {
                let divisor = self[(i, i)];
                self[(j, i)] /= divisor;

                for k in (i + 1)..DIMENSION {
                    let a = self[(j, i)];
                    let b = self[(i, k)];
                    self[(j, k)] -= a * b;
                }
            }
        }

        true
    }

    /// Calculate the determinant of the input matrix.
    pub fn determinant(mut self, tolerance: T) -> Option<T> {
        let mut pivot = Vec::new();
        if !unsafe { self.lup_decompose(&mut pivot, tolerance) } {
            return None;
        }
        debug_assert_eq!(pivot.len(), DIMENSION + 1);

        let mut result = self[(0, 0)];
        for i in 1..DIMENSION {
            result *= self[(i, i)];
        }

        if (pivot[DIMENSION] - DIMENSION) % 2 == 0 {
            Some(result)
        } else {
            Some(-result)
        }
    }

    /// Solve the equation `self` * `x` = `b` for `x`.
    pub fn solve(mut self, b: &ColumnVector<T, DIMENSION>, tolerance: T) -> Option<RowVector<T, DIMENSION>> {
        let mut pivot = Vec::new();
        if !unsafe { self.lup_decompose(&mut pivot, tolerance) } {
            return None;
        }
        debug_assert_eq!(pivot.len(), DIMENSION + 1);

        let mut result = RowVector::new();

        for i in 0..DIMENSION {
            result[(0, i)] = b[(pivot[i], 0)];

            for k in 0..i {
                let factor = result[(0, k)];
                result[(0, i)] -= self[(i, k)] * factor;
            }
        }

        for i in (0..DIMENSION).rev() {
            for k in (i + 1)..DIMENSION {
                let factor = result[(0, k)];
                result[(0, i)] -= self[(i, k)] * factor;
            }

            result[(0, i)] /= self[(i, i)];
        }

        Some(result)
    }

    /// Calculate the inverse of the input matrix.
    pub fn inverted(mut self, tolerance: T) -> Option<Self> {
        let mut pivot = Vec::new();
        if !unsafe { self.lup_decompose(&mut pivot, tolerance) } {
            return None;
        }
        debug_assert_eq!(pivot.len(), DIMENSION + 1);

        let mut result = Self::new();

        let zero = T::whole(0);
        let one = T::whole(1);

        for j in 0..DIMENSION {
            for i in 0..DIMENSION {
                result[(i, j)] = if pivot[i] == j { one } else { zero };

                for k in 0..i {
                    let factor = result[(k, j)];
                    result[(i, j)] -= self[(i, k)] * factor;
                }
            }

            for i in (0..DIMENSION).rev() {
                for k in (i + 1)..DIMENSION {
                    let factor = result[(k, j)];
                    result[(i, j)] -= self[(i, k)] * factor;
                }

                result[(i, j)] /= self[(i, i)];
            }
        }

        Some(result)
    }
}

impl<T: Numeric, const ROWS: usize, const COLUMNS: usize> From<[[T; COLUMNS]; ROWS]> for GenericMatrix<T, ROWS, COLUMNS> {
    fn from(data: [[T; COLUMNS]; ROWS]) -> Self {
        Self {
            data,
        }
    }
}

impl<T: Numeric, const ROWS: usize, const COLUMNS: usize> Index<(usize, usize)> for GenericMatrix<T, ROWS, COLUMNS> {
    type Output = T;

    fn index(&self, index: (usize, usize)) -> &Self::Output {
        let (row, column) = index;
        &self.data[row][column]
    }
}

impl<T: Numeric, const ROWS: usize, const COLUMNS: usize> IndexMut<(usize, usize)> for GenericMatrix<T, ROWS, COLUMNS> {
    fn index_mut(&mut self, index: (usize, usize)) -> &mut Self::Output {
        let (row, column) = index;
        &mut self.data[row][column]
    }
}

impl<T: Numeric + Neg<Output=T>, const ROWS: usize, const COLUMNS: usize> Neg for GenericMatrix<T, ROWS, COLUMNS> {
    type Output = Self;

    fn neg(self) -> Self::Output {
        let mut result = Self::default();
        for r in 0..ROWS {
            for c in 0..COLUMNS {
                result.data[r][c] = -self.data[r][c];
            }
        }
        result
    }
}

impl<T: Numeric, const ROWS: usize, const COLUMNS: usize> Add for GenericMatrix<T, ROWS, COLUMNS> {
    type Output = Self;

    fn add(mut self, rhs: Self) -> Self::Output {
        self += rhs;
        self
    }
}

impl<T: Numeric, const ROWS: usize, const COLUMNS: usize> AddAssign for GenericMatrix<T, ROWS, COLUMNS> {
    fn add_assign(&mut self, rhs: Self) {
        for r in 0..ROWS {
            for c in 0..COLUMNS {
                self.data[r][c] += rhs.data[r][c];
            }
        }
    }
}

impl<T: Numeric, const ROWS: usize, const COLUMNS: usize> Sub for GenericMatrix<T, ROWS, COLUMNS> {
    type Output = Self;

    fn sub(mut self, rhs: Self) -> Self::Output {
        self -= rhs;
        self
    }
}

impl<T: Numeric, const ROWS: usize, const COLUMNS: usize> SubAssign for GenericMatrix<T, ROWS, COLUMNS> {
    fn sub_assign(&mut self, rhs: Self) {
        for r in 0..ROWS {
            for c in 0..COLUMNS {
                self.data[r][c] -= rhs.data[r][c];
            }
        }
    }
}

impl<T: Numeric, const ROWS: usize, const COLUMNS: usize> Mul<T> for GenericMatrix<T, ROWS, COLUMNS> {
    type Output = Self;

    fn mul(mut self, rhs: T) -> Self::Output {
        self *= rhs;
        self
    }
}

impl<T: Numeric, const ROWS: usize, const COLUMNS: usize> MulAssign<T> for GenericMatrix<T, ROWS, COLUMNS> {
    fn mul_assign(&mut self, rhs: T) {
        for r in 0..ROWS {
            for c in 0..COLUMNS {
                self.data[r][c] *= rhs;
            }
        }
    }
}

impl<T: Numeric, const ROWS: usize, const COLUMNS: usize> Div<T> for GenericMatrix<T, ROWS, COLUMNS> {
    type Output = Self;

    fn div(mut self, rhs: T) -> Self::Output {
        self /= rhs;
        self
    }
}

impl<T: Numeric, const ROWS: usize, const COLUMNS: usize> DivAssign<T> for GenericMatrix<T, ROWS, COLUMNS> {
    fn div_assign(&mut self, rhs: T) {
        for r in 0..ROWS {
            for c in 0..COLUMNS {
                self.data[r][c] /= rhs;
            }
        }
    }
}

impl<T: Numeric + Float + PartialOrd, const DIMENSION: usize> Div<SquareMatrix<T, DIMENSION>> for ColumnVector<T, DIMENSION> {
    type Output = RowVector<T, DIMENSION>;

    fn div(self, rhs: SquareMatrix<T, DIMENSION>) -> Self::Output {
        rhs.solve(&self, T::value(1e-9)).unwrap()
    }
}

impl<T: Numeric, const ROWS: usize, const COMMON: usize, const COLUMNS: usize> Mul<GenericMatrix<T, COMMON, COLUMNS>> for GenericMatrix<T, ROWS, COMMON> {
    type Output = GenericMatrix<T, ROWS, COLUMNS>;

    fn mul(self, rhs: GenericMatrix<T, COMMON, COLUMNS>) -> Self::Output {
        let mut result = Self::Output::default();
        for i in 0..ROWS {
            for j in 0..COLUMNS {
                for k in 0..COMMON {
                    result.data[i][j] += self.data[i][k] * rhs.data[k][j];
                }
            }
        }
        result
    }
}

/// Apply a transformation matrix to a column vector.
impl<T: Numeric + Float + Primitive> Mul<ColumnVector<T, 3>> for SquareMatrix<T, 4> {
    type Output = ColumnVector<T, 3>;

    fn mul(self, rhs: ColumnVector<T, 3>) -> Self::Output {
        let x_old = rhs[(0, 0)];
        let y_old = rhs[(1, 0)];
        let z_old = rhs[(2, 0)];

        cvec![
            self[(0,0)] * x_old + self[(0, 1)] * y_old + self[(0, 2)] * z_old + self[(0, 3)],
            self[(1,0)] * x_old + self[(1, 1)] * y_old + self[(1, 2)] * z_old + self[(1, 3)],
            self[(2,0)] * x_old + self[(2, 1)] * y_old + self[(2, 2)] * z_old + self[(2, 3)],
        ]
    }
}

impl<T: Numeric + Float + Primitive> SquareMatrix<T, 4> {
    /// Create a 3D translation matrix.
    pub fn translation(x: T, y: T, z: T) -> Self {
        let mut result = SquareMatrix::identity();
        result[(0, 3)] = x;
        result[(1, 3)] = y;
        result[(2, 3)] = z;
        result
    }

    /// Create a 3D scale matrix.
    pub fn scale(x: T, y: T, z: T) -> Self {
        let one = T::whole(1);
        SquareMatrix::diagonal(&[x, y, z, one])
    }

    fn rotation_impl(angle: T, x: T, y: T, z: T) -> Self {
        let zero = T::whole(0);
        let one = T::whole(1);

        let c = angle.cos();
        let s = angle.sin();
        let omc = one - c;

        Self::from([
            [x * x * omc + c, x * y * omc - z * s, x * z * omc + y * s, zero],
            [y * x * omc + z * s, y * y * omc + c, y * z * omc - x * s, zero],
            [x * z * omc - y * s, y * z * omc + x * s, z * z * omc + c, zero],
            [zero, zero, zero, one],
        ])
    }

    /// Create a 3D rotation matrix.
    ///
    /// `angle` is expected to be in radians.
    #[cfg(not(feature = "angular"))]
    pub fn rotation(angle: T, x: T, y: T, z: T) -> Self {
        Self::rotation_impl(angle, x, y, z)
    }

    /// Create a 3D rotation matrix.
    #[cfg(feature = "angular")]
    pub fn rotation<A: Angular<T>>(angle: A, x: T, y: T, z: T) -> Self {
        Self::rotation_impl(angle.radiant(), x, y, z)
    }
}

#[cfg(test)]
mod tests {
    use crate::{ColumnVector, Complex, cplx, Float, GenericMatrix, Numeric, RowVector, SquareMatrix};

    fn matrices_are_equal<T: Numeric + Float + PartialOrd, const ROWS: usize, const COLUMNS: usize>(
        a: &GenericMatrix<T, ROWS, COLUMNS>,
        b: &GenericMatrix<T, ROWS, COLUMNS>,
        tolerance: T,
    ) -> bool {
        for r in 0..ROWS {
            for c in 0..COLUMNS {
                if (a[(r, c)] - b[(r, c)]).abs() > tolerance {
                    return false;
                }
            }
        }
        true
    }

    #[test]
    fn identity_matrix() {
        let mat: SquareMatrix<f64, 3> = SquareMatrix::identity();
        let expected = SquareMatrix::from([
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
            [0.0, 0.0, 1.0],
        ]);

        assert_eq!(mat, expected);
    }

    #[test]
    fn diagonal_matrix() {
        let mat: SquareMatrix<f32, 3> = SquareMatrix::diagonal(&[1.0, 2.0, 3.0]);
        let expected = SquareMatrix {
            data: [[1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 3.0]],
        };

        assert_eq!(mat, expected);
    }

    #[test]
    #[should_panic = "index out of bounds"]
    fn out_of_bounds_access() {
        let mat: SquareMatrix<f32, 1> = SquareMatrix::identity();
        mat[(1, 0)];
    }

    #[test]
    fn transposing_matrix() {
        let mat = GenericMatrix {
            data: [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]],
        }.transposed();
        let expected = GenericMatrix {
            data: [[1.0, 3.0, 5.0], [2.0, 4.0, 6.0]],
        };

        assert_eq!(mat, expected);
    }

    #[test]
    fn negating_matrix() {
        let mat = -GenericMatrix {
            data: [[1, -1], [2, -2]],
        };
        let expected = GenericMatrix {
            data: [[-1, 1], [-2, 2]],
        };

        assert_eq!(mat, expected);
    }

    #[test]
    fn dot_product_of_vectors() {
        let a = rvec![1, 3, -5];
        let b = cvec![4, -2, -1];
        let product = a * b;

        assert_eq!(product[(0, 0)], 3);
    }

    #[test]
    fn matrix_product_of_vectors() {
        let a = cvec![1, 3, -5];
        let b = rvec![4, -2, -1];
        let product = a * b;
        let expected = SquareMatrix {
            data: [
                [4, -2, -1],
                [12, -6, -3],
                [-20, 10, 5],
            ],
        };

        assert_eq!(product, expected);
    }

    #[test]
    fn complex_matrix_multiplication() {
        let a = SquareMatrix {
            data: [
                [cplx!(2.0, 1.0), cplx!(i = 5.0,)],
                [cplx!(3.0), cplx!(3.0, -4.0)],
            ],
        };
        let b = SquareMatrix {
            data: [
                [cplx!(1.0, -1.0), cplx!(4.0, 2.0)],
                [cplx!(1.0, -6.0), cplx!(3.0)],
            ],
        };
        let mat = a * b;
        let expected = SquareMatrix::from([
            [cplx!(33.0, 4.0), cplx!(6.0, 23.0)],
            [cplx!(-18.0, -25.0), cplx!(21.0, -6.0)],
        ]);

        assert_eq!(mat, expected);
    }

    #[test]
    fn invert_matrix() {
        let tolerance = 1e-9;
        let mat = SquareMatrix::from([
            [2.0, 1.0],
            [6.0, 4.0],
        ]);
        let inverted = mat.inverted(tolerance).unwrap();
        let expected = SquareMatrix::from([
            [2.0, -0.5],
            [-3.0, 1.0],
        ]);

        assert!(matrices_are_equal(&inverted, &expected, tolerance));
        assert!(matrices_are_equal(&(inverted * mat), &SquareMatrix::identity(), tolerance));
    }
}