use crate::internal::{self, uninit};
use crate::linalg::vector::Vector;
use num_traits::{Float, One, Zero};
use std::fmt::{Debug, Display, Formatter};
use std::iter::Sum;
use std::ops::{Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Sub, SubAssign};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Matrix<T, const M: usize, const N: usize> {
cols: [Vector<T, M>; N],
}
impl<T, const M: usize, const N: usize> Matrix<T, M, N> {
pub const fn new(cols: [Vector<T, M>; N]) -> Self {
Self { cols }
}
pub fn gen<F>(f: F) -> Self
where
F: Fn(usize, usize) -> T,
{
let mut col_data = uninit::new_uninit_array::<Vector<T, M>, N>();
for (col_idx, col) in col_data.iter_mut().enumerate() {
col.write(Vector::gen(|row_idx| f(col_idx, row_idx)));
}
let col_data = unsafe { uninit::array_assume_init(col_data) };
Matrix::new(col_data)
}
pub fn submatrix<const P: usize, const Q: usize>(
self,
col_offset: usize,
row_offset: usize,
) -> Matrix<T, P, Q>
where
T: Clone,
{
debug_assert_eq!(P + row_offset, M);
debug_assert_eq!(Q + col_offset, N);
let mut submatrix_cols = uninit::new_uninit_array::<Vector<T, P>, Q>();
for (col_idx, col) in submatrix_cols.iter_mut().enumerate() {
let column = Vector::<T, P>::gen(|row_idx| {
self[col_idx + col_offset][row_idx + row_offset].clone()
});
col.write(column);
}
let submatrix_cols = unsafe { uninit::array_assume_init(submatrix_cols) };
Matrix::new(submatrix_cols)
}
pub fn transpose(self) -> Matrix<T, N, M>
where
T: Clone,
{
let mut new_cols = uninit::new_uninit_array::<Vector<T, N>, M>();
for (row_idx, new_col) in new_cols.iter_mut().enumerate() {
new_col.write(Vector::gen(|col_idx| self.cols[col_idx][row_idx].clone()));
}
let new_cols = unsafe { uninit::array_assume_init(new_cols) };
Matrix::new(new_cols)
}
pub fn map<F, U>(self, f: F) -> Matrix<U, M, N>
where
F: Fn(T) -> U,
{
let mut new_cols = uninit::new_uninit_array::<Vector<U, M>, N>();
for (new_col, old_col) in new_cols.iter_mut().zip(self.cols.into_iter()) {
new_col.write(old_col.map(&f));
}
let new_cols = unsafe { uninit::array_assume_init(new_cols) };
Matrix::new(new_cols)
}
pub fn map_mut<F>(&mut self, f: F)
where
F: Fn(&mut T),
{
for col in self.cols.iter_mut() {
col.map_mut(&f);
}
}
pub fn map_column<F>(self, f: F, col_idx: usize) -> Self
where
F: Fn(T) -> T,
T: Debug,
{
debug_assert!(col_idx < N);
let new_cols: Vec<_> = self
.cols
.into_iter()
.enumerate()
.map(|(i, col)| if i == col_idx { col.map(&f) } else { col })
.collect();
let new_cols: [Vector<T, M>; N] = new_cols.try_into().unwrap();
Matrix::new(new_cols)
}
pub fn map_column_mut<F>(&mut self, f: F, col_idx: usize)
where
F: Fn(&mut T),
T: Debug,
{
debug_assert!(col_idx < N);
self.cols
.iter_mut()
.enumerate()
.nth(col_idx)
.unwrap()
.1
.map_mut(f);
}
pub fn apply<F, U, V>(self, f: F, rhs: Matrix<U, M, N>) -> Matrix<V, M, N>
where
F: Fn(T, U) -> V,
{
let mut new_cols = uninit::new_uninit_array::<Vector<V, M>, N>();
for (new_col, (lhs, rhs)) in new_cols
.iter_mut()
.zip(self.cols.into_iter().zip(rhs.cols.into_iter()))
{
new_col.write(lhs.apply(&f, rhs));
}
let new_cols = unsafe { uninit::array_assume_init(new_cols) };
Matrix::new(new_cols)
}
pub fn apply_mut<F, U>(&mut self, f: F, rhs: Matrix<U, M, N>)
where
F: Fn(&mut T, U),
{
for (lhs_col, rhs_col) in self.cols.iter_mut().zip(rhs.cols.into_iter()) {
lhs_col.apply_mut(&f, rhs_col);
}
}
pub fn iter(&self) -> <&Self as IntoIterator>::IntoIter {
<&Self as IntoIterator>::into_iter(self)
}
pub fn into_iter(self) -> <Self as IntoIterator>::IntoIter {
<Self as IntoIterator>::into_iter(self)
}
pub fn iter_mut(&mut self) -> <&mut Self as IntoIterator>::IntoIter {
<&mut Self as IntoIterator>::into_iter(self)
}
pub fn row_echelon_form(self) -> Self
where
T: Float + Debug,
{
self.transpose().column_echelon_form().transpose()
}
pub fn reduced_row_echelon_form(self) -> Self
where
T: Float + Debug,
{
self.transpose().reduced_column_echelon_form().transpose()
}
pub fn column_echelon_form(self) -> Self
where
T: Float + Debug,
{
let mut arranged = self.sort_columns_by_leading_coefficient_index();
for i in 0..N {
let col_leading_coef_idx = arranged[i]
.iter()
.enumerate()
.find(|(_, ele)| ele.is_zero())
.map(|(i, _)| i);
if let Some(leading_idx) = col_leading_coef_idx {
let pivot_col = arranged[i];
for (j, col) in arranged.iter_mut().enumerate() {
if i == j {
continue;
}
if !col[leading_idx].is_zero() {
let pivot_ratio = col[leading_idx] / pivot_col[leading_idx];
*col = *col - pivot_col * pivot_ratio;
}
}
}
}
arranged
}
pub fn reduced_column_echelon_form(self) -> Self
where
T: Float + Debug,
{
let mut cef = self.column_echelon_form();
for (col, leading_idx) in cef.leading_coefficient_indices_mut() {
if let Some(leading_idx) = leading_idx {
let pivot_divisor = col[leading_idx];
*col = *col / pivot_divisor;
}
}
cef
}
pub fn sort_columns_by_leading_coefficient_index(self) -> Self
where
T: Zero + Clone + Debug,
{
let mut leading_coef_idxs: Vec<_> = self.leading_coefficient_indices().collect();
leading_coef_idxs
.sort_by(|(_, idx1), (_, idx2)| internal::option::option_ordering_max_none(idx1, idx2));
let sorted_cols: [Vector<T, M>; N] = leading_coef_idxs
.into_iter()
.map(|(col, _)| col)
.cloned()
.collect::<Vec<_>>()
.try_into()
.unwrap();
Matrix::new(sorted_cols)
}
fn leading_coefficient_indices(&self) -> impl Iterator<Item = (&Vector<T, M>, Option<usize>)>
where
T: Zero,
{
self.cols.iter().map(|col| {
let idx = col
.iter()
.enumerate()
.skip_while(|(_, ele)| ele.is_zero())
.map(|(idx, _)| idx)
.next();
(col, idx)
})
}
fn leading_coefficient_indices_mut(
&mut self,
) -> impl Iterator<Item = (&mut Vector<T, M>, Option<usize>)>
where
T: Zero,
{
self.cols.iter_mut().map(|col| {
let idx = col
.iter()
.enumerate()
.skip_while(|(_, ele)| ele.is_zero())
.map(|(idx, _)| idx)
.next();
(col, idx)
})
}
pub fn swap_columns(&mut self, col1_idx: usize, col2_idx: usize) {
debug_assert!(col1_idx < N);
debug_assert!(col2_idx < N);
self.cols.swap(col1_idx, col2_idx);
}
}
impl<T, const N: usize> Matrix<T, N, 1> {
pub fn new_column(col_data: [T; N]) -> Self {
Vector::new(col_data).into()
}
pub(super) fn into_vector(self) -> Vector<T, N> {
self.cols.into_iter().next().unwrap()
}
}
impl<T, const N: usize> Matrix<T, 1, N> {
pub fn new_row(row_data: [T; N]) -> Self {
let mut cols = uninit::new_uninit_array::<Vector<T, 1>, N>();
for (col, datum) in cols.iter_mut().zip(row_data.into_iter()) {
col.write(Vector::new1(datum));
}
let cols = unsafe { uninit::array_assume_init(cols) };
Matrix::new(cols)
}
}
impl<T, const M: usize, const N: usize> Add<Matrix<T, M, N>> for Matrix<T, M, N>
where
T: Add<T, Output = T>,
{
type Output = Matrix<T, M, N>;
fn add(self, rhs: Matrix<T, M, N>) -> Self::Output {
self.apply(T::add, rhs)
}
}
impl<T, const M: usize, const N: usize> AddAssign<Matrix<T, M, N>> for Matrix<T, M, N>
where
T: AddAssign<T>,
{
fn add_assign(&mut self, rhs: Matrix<T, M, N>) {
self.apply_mut(T::add_assign, rhs);
}
}
impl<T, const M: usize, const N: usize> Sub<Matrix<T, M, N>> for Matrix<T, M, N>
where
T: Sub<T, Output = T>,
{
type Output = Matrix<T, M, N>;
fn sub(self, rhs: Matrix<T, M, N>) -> Self::Output {
self.apply(T::sub, rhs)
}
}
impl<T, const M: usize, const N: usize> SubAssign<Matrix<T, M, N>> for Matrix<T, M, N>
where
T: SubAssign<T>,
{
fn sub_assign(&mut self, rhs: Matrix<T, M, N>) {
self.apply_mut(T::sub_assign, rhs)
}
}
impl<T, const M: usize, const N: usize> Mul<T> for Matrix<T, M, N>
where
T: Mul<T, Output = T> + Clone,
{
type Output = Matrix<T, M, N>;
fn mul(self, rhs: T) -> Self::Output {
self.map(|x| x * rhs.clone())
}
}
impl<T, const M: usize, const N: usize, const P: usize> Mul<Matrix<T, N, P>> for Matrix<T, M, N>
where
T: Clone + Mul<T, Output = T> + Sum,
{
type Output = Matrix<T, M, P>;
fn mul(self, rhs: Matrix<T, N, P>) -> Self::Output {
let tp = self.transpose();
let mut new_cols = uninit::new_uninit_array::<Vector<T, M>, P>();
for (new_col, rhs) in new_cols.iter_mut().zip(rhs.cols.into_iter()) {
let mut new_col_data = uninit::new_uninit_array::<T, M>();
for (new_col_datum, lhs) in new_col_data.iter_mut().zip(tp.cols.iter().cloned()) {
new_col_datum.write(lhs.clone().dot(rhs.clone()));
}
let new_col_data = unsafe { uninit::array_assume_init(new_col_data) };
new_col.write(Vector::new(new_col_data));
}
let new_cols = unsafe { uninit::array_assume_init(new_cols) };
Matrix::new(new_cols)
}
}
impl<T, const M: usize, const N: usize> Mul<Vector<T, N>> for Matrix<T, M, N>
where
T: Clone + Mul<T, Output = T> + Sum,
{
type Output = Vector<T, M>;
fn mul(self, rhs: Vector<T, N>) -> Self::Output {
(self * Matrix::from(rhs)).into()
}
}
impl<T, const M: usize, const N: usize> Div<T> for Matrix<T, M, N>
where
T: Div<T, Output = T> + Clone,
{
type Output = Matrix<T, M, N>;
fn div(self, rhs: T) -> Self::Output {
self.map(|x| x / rhs.clone())
}
}
impl<T, const M: usize, const N: usize> DivAssign<T> for Matrix<T, M, N>
where
T: DivAssign<T> + Clone,
{
fn div_assign(&mut self, rhs: T) {
self.map_mut(|x| *x /= rhs.clone());
}
}
impl<T, const M: usize, const N: usize> Zero for Matrix<T, M, N>
where
T: Zero,
{
fn zero() -> Self {
Matrix::gen(|_, _| T::zero())
}
fn is_zero(&self) -> bool {
self.cols.iter().all(|col| col.iter().all(|x| x.is_zero()))
}
}
impl<T, const M: usize> One for Matrix<T, M, M>
where
T: Clone + Mul<T, Output = T> + Sum + One + Zero,
{
fn one() -> Self {
Matrix::gen(|col_idx, row_idx| {
if col_idx == row_idx {
T::one()
} else {
T::zero()
}
})
}
}
impl<T, const M: usize, const N: usize> IntoIterator for Matrix<T, M, N> {
type Item = Vector<T, M>;
type IntoIter = std::array::IntoIter<Vector<T, M>, N>;
fn into_iter(self) -> Self::IntoIter {
self.cols.into_iter()
}
}
impl<'a, T, const M: usize, const N: usize> IntoIterator for &'a Matrix<T, M, N> {
type Item = &'a Vector<T, M>;
type IntoIter = std::slice::Iter<'a, Vector<T, M>>;
fn into_iter(self) -> Self::IntoIter {
self.cols.iter()
}
}
impl<'a, T, const M: usize, const N: usize> IntoIterator for &'a mut Matrix<T, M, N> {
type Item = &'a mut Vector<T, M>;
type IntoIter = std::slice::IterMut<'a, Vector<T, M>>;
fn into_iter(self) -> Self::IntoIter {
self.cols.iter_mut()
}
}
impl<T, const M: usize, const N: usize> Index<usize> for Matrix<T, M, N> {
type Output = Vector<T, M>;
fn index(&self, index: usize) -> &Self::Output {
&self.cols[index]
}
}
impl<T, const M: usize, const N: usize> IndexMut<usize> for Matrix<T, M, N> {
fn index_mut(&mut self, index: usize) -> &mut Self::Output {
&mut self.cols[index]
}
}
impl<T, const N: usize> From<Vector<T, N>> for Matrix<T, N, 1> {
fn from(value: Vector<T, N>) -> Self {
value.into_matrix()
}
}
impl<T, const M: usize, const N: usize> From<[[T; M]; N]> for Matrix<T, M, N> {
fn from(value: [[T; M]; N]) -> Self {
let cols = value.map(|row| Vector::new(row));
Matrix::new(cols)
}
}
impl<T, const M: usize, const N: usize> Display for Matrix<T, M, N>
where
T: Display,
{
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
for row_idx in 0..M {
for col_idx in 0..N {
write!(f, "{}\t", self.cols[col_idx][row_idx])?;
}
writeln!(f)?;
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn check_add() {
let lhs = Matrix::<isize, 3, 4>::from([[2, 3, 5], [7, 11, 13], [17, 19, 23], [29, 31, 37]]);
let rhs =
Matrix::<isize, 3, 4>::from([[41, 43, 47], [53, 59, 61], [67, 71, 73], [79, 83, 89]]);
let expected_sum = Matrix::<isize, 3, 4>::from([
[43, 46, 52],
[60, 70, 74],
[84, 90, 96],
[108, 114, 126],
]);
let actual_sum = lhs + rhs;
assert_eq!(expected_sum, actual_sum);
}
#[test]
fn check_sub() {
let lhs = Matrix::<isize, 2, 3>::from([[97, 101], [103, 107], [109, 113]]);
let rhs = Matrix::<isize, 2, 3>::from([[127, 131], [137, 139], [149, 151]]);
let expected_diff = Matrix::<isize, 2, 3>::from([[-30, -30], [-34, -32], [-40, -38]]);
let actual_diff = lhs - rhs;
assert_eq!(expected_diff, actual_diff);
}
#[test]
fn check_mul() {
let lhs = Matrix::<isize, 4, 2>::from([[2, 3, 5, 7], [11, 13, 17, 19]]);
let rhs = Matrix::<isize, 2, 3>::from([[23, 29], [31, 37], [41, 43]]);
let expected_prod = Matrix::<isize, 4, 3>::from([
[365, 446, 608, 712],
[469, 574, 784, 920],
[555, 682, 936, 1104],
]);
let actual_prod = lhs * rhs;
assert_eq!(expected_prod, actual_prod);
}
#[test]
fn check_transpose() {
let mat = Matrix::<isize, 3, 4>::from([[2, 3, 5], [7, 11, 13], [17, 19, 23], [29, 31, 37]]);
let expected_transpose =
Matrix::<isize, 4, 3>::from([[2, 7, 17, 29], [3, 11, 19, 31], [5, 13, 23, 37]]);
let actual_transpose = mat.transpose();
println!("{}", mat);
println!("{}", actual_transpose);
assert_eq!(expected_transpose, actual_transpose);
}
#[test]
fn check_rref() {
let mat = Matrix::<f64, 3, 4>::from([
[2.0, -3.0, -2.0],
[1.0, -1.0, 1.0],
[-1.0, 2.0, 2.0],
[8.0, -11.0, -3.0],
]);
let expected_rref = Matrix::<f64, 3, 4>::from([
[1.0, 0.0, 0.0],
[0.0, 1.0, 0.0],
[0.0, 0.0, 1.0],
[2.0, 3.0, -1.0],
]);
let actual_rref = mat.reduced_row_echelon_form();
assert_eq!(expected_rref, actual_rref);
}
}