te1d/

linalg.rs

//! Solves linear systems and handles matrices.
//! * The utilities in this module is not at all fast. For massive computations, use other specialized linear algebra packages.
//! * By matrices, we mean [`ndarray::Array2`].

use std::fmt;

use ndarray::prelude::*;

use crate::calc::ColVec;


/// Error type for linear algebraic computations
pub struct LinalgError {
    repr: LinalgErrorKind,
}

impl fmt::Display for LinalgErrorKind {
    /// Shows a human-readable description of the [`LinalgErrorKind`].
    ///
    /// This is similar to `impl Display for Error`, but doesn't require first converting to Error.
    ///
    /// # Examples
    /// ```
    /// use te1d::linalg::LinalgErrorKind;
    /// assert_eq!("singular matrix", LinalgErrorKind::SingularMatrixError.to_string());
    /// ```
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.write_str(self.as_str())
    }
}

impl fmt::Debug for LinalgError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(&self.repr, f)
    }
}

impl LinalgError {
    pub fn new(error_kind: LinalgErrorKind) -> Self {
        LinalgError { repr: error_kind }
    }

    #[inline]
    pub fn kind(&self) -> LinalgErrorKind {
        self.repr
    }
}

/// A list specifying general categories of errors in linear algebraic computation
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum LinalgErrorKind {
    /// the matrix is singular, that is, not invertible
    SingularMatrixError,
}

impl LinalgErrorKind {
    pub(crate) fn as_str(&self) -> &'static str {
        use LinalgErrorKind::*;
        // Strictly alphabetical, please.  (Sadly rustfmt cannot do this yet.)
        match *self {
            SingularMatrixError => "singular matrix",
        }
    }
}


/// Find the solution `x` of the linear system `mat * x = b` using Gaussian elimination.
/// 
/// # Arguments
/// * `square_mat`: a square matrix.
/// 
/// # Examples
/// ```
/// use ndarray::{Array2, array};
/// use te1d::prelude::*;
/// use te1d::linalg::{solve_gauss, LinalgErrorKind};
/// 
/// let mat: Array2<f64> = array![[1.0, 0.0, 1.0], [0.0, -3.0, 1.0], [2.0, 1.0, 3.0]];
/// let b: ColVec = array![6.0, 7.0, 15.0];
/// let x = solve_gauss(&mat, &b).unwrap();
/// let x_exact: ColVec = array![2.0, -1.0, 4.0];
/// 
/// assert!(all_close(&x, &x_exact, 1e-05, 1e-08));
/// 
/// // Error for singular matrix
/// let mat: Array2<f64> = array![[0.0, 0.0], [1.0, 2.0]];
/// let b: ColVec = array![1.0, 2.0];
/// let msg = match solve_gauss(&mat, &b) {
///     Ok(b) => String::from("no error"),
///     Err(err) => err.kind().to_string(),
/// };
/// assert_eq!(msg, LinalgErrorKind::SingularMatrixError.to_string());
/// ```
/// 
/// # Panics
/// * when `mat` is not square.
/// * when the number of rows of `mat` and the length of `b` are not the same.
pub fn solve_gauss(square_mat: &Array2<f64>, b: &ColVec) -> Result<ColVec, LinalgError> {
    if !square_mat.is_square() {
        panic!("`mat` must be a square matrix!");
    }
    let n = square_mat.shape()[0];
    if b.len() != n {
        panic!("the number of rows of `mat` and the length of `b` must be the same!");
    }

    let mut mat = square_mat.to_owned();
    let mut b = b.to_owned();
    let mut pivot_row: usize;
    let mut pivot_value: f64;
    let mut target_row_pivot_value: f64;

    // perform upper triangulation
    for row in 0..n {
        // set the pivot row
        pivot_row = row_of_abs_max(&mat, row, row, n-1);
        if row != pivot_row {
            swap_rows(&mut mat, row, pivot_row);
            b.swap(row, pivot_row);
        }
        // normalize the pivot row and `b`
        pivot_value = mat[[row, row]];
        if pivot_value.abs() <= f64::EPSILON {
            return Err(LinalgError::new(LinalgErrorKind::SingularMatrixError));
        }
        mat[[row, row]] = 1.0;
        mat.row_mut(row).slice_mut(s![row+1..]).mapv_inplace(|e| { e/pivot_value });
        b[row] /= pivot_value;
        // make upper triangular
        for target_row in row+1..n {
            target_row_pivot_value = mat[[target_row, row]];
            mat[[target_row, row]] = 0.0;
            for col in row+1..n {
                mat[[target_row, col]] -= target_row_pivot_value * mat[[row, col]];
            }
            b[target_row] -= target_row_pivot_value * b[row];
        }
    }

    // perform diagonalization
    for row in (0..n).rev() {
        for target_row in 0..row {
            target_row_pivot_value = mat[[target_row, row]];
            mat[[target_row, row]] = 0.0;
            for col in row+1..n {
                mat[[target_row, col]] -= target_row_pivot_value * mat[[row, col]];
            }
            b[target_row] -= target_row_pivot_value * b[row];
        }
    }

    Ok(b)
}


/// Return the index of the element having the largest absolute value among a given column.
fn row_of_abs_max(mat: &Array2<f64>, col: usize, start_row: usize, end_row: usize) -> usize {
    let mut max_value: f64 = mat[[start_row, col]].abs();
    let mut result: usize = start_row;
    for i in start_row+1..end_row+1 {
        let cur_value = mat[[i, col]].abs();
        if cur_value > max_value {
            max_value = cur_value;
            result = i;
        }
    }

    result
}


/// Swap two rows of a matrix.
/// 
/// # Examples
/// ```
/// use ndarray::prelude::*;
/// use te1d::prelude::*;
/// use te1d::linalg::swap_rows;
/// 
/// let mut mat: Array2<f64> = array![[0.0, 1.0], [1.0, 2.0], [2.0, 3.0]];
/// let result: Array2<f64> = array![[2.0, 3.0], [1.0, 2.0], [0.0, 1.0]];
/// 
/// swap_rows(&mut mat, 0, 2);
/// assert_eq!(mat, result);
/// ```
#[inline]
pub fn swap_rows(mat: &mut Array2<f64>, row1: usize, row2: usize) {
    for i in 0..mat.shape()[1] {
        mat.swap([row1, i], [row2, i]);
    }
}


/// Swap two columns of a matrix.
/// 
/// # Examples
/// ```
/// use ndarray::prelude::*;
/// use te1d::prelude::*;
/// use te1d::linalg::swap_cols;
/// 
/// let mut mat: Array2<f64> = array![[0.0, 1.0, 2.0], [1.0, 2.0, 3.0]];
/// let result: Array2<f64> = array![[2.0, 1.0, 0.0], [3.0, 2.0, 1.0]];
/// 
/// swap_cols(&mut mat, 0, 2);
/// assert_eq!(mat, result);
/// ```
#[inline]
pub fn swap_cols(mat: &mut Array2<f64>, col1: usize, col2: usize) {
    for i in 0..mat.shape()[0] {
        mat.swap([i, col1], [i, col2]);
    }
}


/// Compute the outer product of two [`ColVec`]s.
/// For two column vectors `u` and `v`, the result is `u*v^T`.
/// 
/// # Examples
/// ```
/// use ndarray::prelude::*;
/// use te1d::prelude::*;
/// use te1d::linalg::outer_product;
/// 
/// let u: ColVec = array![1.0, 2.0, 3.0];
/// let v: ColVec = array![4.0, 5.0];
/// let result: Array2<f64> = outer_product(&u, &v);
/// let sol: Array2<f64> = array![[4.0, 5.0], [8.0, 10.0], [12.0, 15.0]];
/// 
/// assert_eq!(&result, &sol);
/// ```
pub fn outer_product(u: &ColVec, v: &ColVec) -> Array2<f64> {
    let m: usize = u.len();
    let n: usize = v.len();
    let mut result: Array2<f64> = Array::default((m, n));
    for i in 0..m {
        for j in 0..n {
            result[[i, j]] = u[i] * v[j];
        }
    }

    result
}


#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_row_of_abs_max() {
        let mat: Array2<f64> = array![[1.0, 2.0], [-5.0, -2.0], [2.0, -7.0], [3.0, 9.0]];
        let result = row_of_abs_max(&mat, 1, 1, 2);
        assert_eq!(result, 2 as usize);
    }

    #[test]
    #[should_panic]
    fn test_panic_solve_gauss_when_not_square() {
        let mat: Array2<f64> = array![[1.0, 2.0], [-5.0, -2.0], [2.0, -7.0], [3.0, 9.0]];
        let b = array![0.0, 1.0, 2.0, 3.0];
        solve_gauss(&mat, &b).unwrap();   // must panic: not a square matrix

    }

    #[test]
    #[should_panic]
    fn test_panic_solve_gauss_when_size_mistmatch() {
        let mat: Array2<f64> = array![[1.0, 2.0], [3.0, 4.0]];
        let b = array![0.0, 1.0, 2.0, 3.0];
        solve_gauss(&mat, &b).unwrap();   // must panic: size does not match
    }
}