1use anyhow::{anyhow, Result};
  2use std::ops::{Add, Mul};
  3
  4#[derive(Debug, Clone, PartialEq)]
  5pub struct Matrix<T>(pub Vec<Vec<T>>);
  6
  7impl<T> Matrix<T>
  8where
  9    T: Copy + Add<Output = T> + Mul<Output = T> + Default + std::fmt::Debug,
 10{
 11    // TODO maybe rm this method, move it to tests only
 12    pub fn new(rows: usize, cols: usize, value: T) -> Self {
 13        Matrix(vec![vec![value; cols]; rows])
 14    }
 15
 16    pub fn add(&self, other: &Matrix<T>) -> Result<Matrix<T>> {
 17        if self.0.len() != other.0.len() || self.0[0].len() != other.0[0].len() {
 18            return Err(anyhow!("dimensions don't match"));
 19        }
 20
 21        let r = self
 22            .0
 23            .iter()
 24            .zip(&other.0)
 25            .map(|(row1, row2)| {
 26                row1.iter()
 27                    .zip(row2)
 28                    .map(|(a, b)| *a + *b)
 29                    .collect::<Vec<T>>()
 30            })
 31            .collect::<Vec<Vec<T>>>();
 32
 33        Ok(Matrix(r))
 34    }
 35
 36    pub fn mul(&self, other: &Matrix<T>) -> Result<Matrix<T>> {
 37        let rows_a = self.0.len();
 38        let cols_a = self.0[0].len();
 39        let rows_b = other.0.len();
 40        let cols_b = other.0[0].len();
 41
 42        if cols_a != rows_b {
 43            return Err(anyhow!("self.n_cols != other.n_rows"));
 44        }
 45
 46        let mut r = vec![vec![T::default(); cols_b]; rows_a];
 47
 48        for i in 0..rows_a {
 49            for j in 0..cols_b {
 50                for k in 0..cols_a {
 51                    r[i][j] = r[i][j] + self.0[i][k] * other.0[k][j];
 52                }
 53            }
 54        }
 55
 56        Ok(Matrix(r))
 57    }
 58    pub fn mul_vec(&self, v: &Vec<T>) -> Result<Vec<T>> {
 59        let rows = self.0.len();
 60        let cols = self.0[0].len();
 61
 62        if cols != v.len() {
 63            return Err(anyhow!(
 64                "Number of columns in matrix does not match the length of the vector"
 65            ));
 66        }
 67
 68        let mut r = vec![T::default(); rows];
 69
 70        for i in 0..rows {
 71            for j in 0..cols {
 72                r[i] = r[i] + self.0[i][j] * v[j];
 73            }
 74        }
 75        Ok(r)
 76    }
 77
 78    pub fn transpose(&self) -> Matrix<T> {
 79        let rows = self.0.len();
 80        let cols = self.0[0].len();
 81        let mut r = vec![vec![T::default(); rows]; cols];
 82
 83        for i in 0..rows {
 84            for j in 0..cols {
 85                r[j][i] = self.0[i][j];
 86            }
 87        }
 88
 89        Matrix(r)
 90    }
 91
 92    pub fn scalar_mul(&self, scalar: T) -> Matrix<T> {
 93        let r = self
 94            .0
 95            .iter()
 96            .map(|row| row.iter().map(|&val| val * scalar).collect::<Vec<T>>())
 97            .collect::<Vec<Vec<T>>>();
 98
 99        Matrix(r)
100    }
101}
102
103// WIP. Currently uses ndarray, ndarray_linalg, num_complex to solve a system of
104// linear equations A*x=b for x.
105use crate::C;
106impl Matrix<C<f64>> {
107    pub fn solve(&self, b: &Vec<C<f64>>) -> Result<Vec<C<f64>>> {
108        use ndarray::{Array1, Array2};
109        use ndarray_linalg::Solve;
110        use num_complex::Complex64;
111
112        let m: Array2<Complex64> = Array2::from_shape_vec(
113            (self.0.len(), self.0[0].len()),
114            self.0
115                .clone()
116                .into_iter()
117                .flatten()
118                .map(|e| Complex64::new(e.re, e.im))
119                .collect(),
120        )
121        .unwrap();
122        let v: Array1<Complex64> = Array1::from_shape_vec(
123            b.len(),
124            b.iter().map(|e| Complex64::new(e.re, e.im)).collect(),
125        )
126        .unwrap();
127        let r = m.solve(&v)?;
128        let r: Vec<C<f64>> = r.into_iter().map(|e| C::<f64>::new(e.re, e.im)).collect();
129        Ok(r)
130    }
131}
132impl Matrix<f64> {
133    // tmp (rm)
134    pub fn solve(&self, b: Vec<f64>) -> Result<Vec<f64>> {
135        use ndarray::{Array1, Array2};
136        use ndarray_linalg::Solve;
137
138        let m: Array2<f64> = Array2::from_shape_vec(
139            (self.0.len(), self.0[0].len()),
140            self.0.clone().into_iter().flatten().collect(),
141        )
142        .unwrap();
143        let v: Array1<f64> = Array1::from_shape_vec(b.len(), b).unwrap();
144        let r = m.solve(&v)?;
145        let r: Vec<f64> = r.into_iter().map(|e| e).collect();
146        Ok(r)
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153
154    #[test]
155    fn test_add() -> Result<()> {
156        let a = Matrix::new(2, 3, 1);
157        let b = Matrix::new(2, 3, 2);
158        let expected = Matrix::new(2, 3, 3);
159        assert_eq!(a.add(&b).unwrap(), expected);
160        Ok(())
161    }
162
163    #[test]
164    fn test_mul() -> Result<()> {
165        let a = Matrix::new(2, 3, 1);
166        let b = Matrix::new(3, 2, 1);
167        let expected = Matrix::new(2, 2, 3); // 2x3 * 3x2 = 2x2 matrix (with all values 3)
168        assert_eq!(a.mul(&b).unwrap(), expected);
169        Ok(())
170    }
171
172    #[test]
173    fn test_transpose() -> Result<()> {
174        let a = Matrix::new(2, 3, 1);
175        let expected = Matrix::new(3, 2, 1);
176        assert_eq!(a.transpose(), expected);
177        Ok(())
178    }
179
180    #[test]
181    fn test_scalar_mul() -> Result<()> {
182        let a = Matrix::new(2, 3, 1);
183        let expected = Matrix::new(2, 3, 3);
184        assert_eq!(a.scalar_mul(3), expected);
185        Ok(())
186    }
187}