numerical_linear_algebra/
matrix.rs

1use std::{
2    fmt::Display,
3    iter::Sum,
4    ops::{Add, AddAssign, Index, IndexMut, Mul, MulAssign, Sub, SubAssign},
5};
6
7use num_complex::{Complex64, ComplexFloat};
8use rayon::iter::{IntoParallelRefMutIterator, ParallelBridge, ParallelIterator};
9use tabled::{
10    builder::Builder,
11    settings::{Margin, Style},
12};
13
14use crate::Vector;
15
16#[derive(Debug, Clone)]
17pub struct Matrix<const M: usize, const N: usize> {
18    data: [[Complex64; N]; M],
19}
20
21impl<const M: usize, const N: usize> Index<[usize; 2]> for Matrix<M, N> {
22    type Output = Complex64;
23
24    fn index(&self, [i, j]: [usize; 2]) -> &Self::Output {
25        &self.data[i][j]
26    }
27}
28
29impl<const M: usize, const N: usize> IndexMut<[usize; 2]> for Matrix<M, N> {
30    fn index_mut(&mut self, index: [usize; 2]) -> &mut Self::Output {
31        &mut self.data[index[0]][index[1]]
32    }
33}
34
35impl<const M: usize, const N: usize> AddAssign for Matrix<M, N> {
36    fn add_assign(&mut self, rhs: Self) {
37        self.iter_mut()
38            .zip(rhs.iter())
39            .par_bridge()
40            .for_each(|(left, right)| *left += right);
41    }
42}
43
44impl<const M: usize, const N: usize> Add for Matrix<M, N> {
45    type Output = Matrix<M, N>;
46
47    fn add(self, rhs: Self) -> Self::Output {
48        let mut result = self.clone();
49        result += rhs;
50        result
51    }
52}
53
54impl<const M: usize, const N: usize> Sum for Matrix<M, N> {
55    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
56        iter.reduce(|acc, x| acc + x).unwrap_or(Matrix::ZEROS)
57    }
58}
59
60impl<const M: usize, const N: usize> SubAssign for Matrix<M, N> {
61    fn sub_assign(&mut self, rhs: Self) {
62        self.iter_mut()
63            .zip(rhs.iter())
64            .par_bridge()
65            .for_each(|(left, right)| *left -= right);
66    }
67}
68
69impl<const M: usize, const N: usize> Sub for Matrix<M, N> {
70    type Output = Matrix<M, N>;
71
72    fn sub(self, rhs: Self) -> Self::Output {
73        let mut result = self.clone();
74        result -= rhs;
75        result
76    }
77}
78
79impl<const M: usize, const N: usize, const P: usize> Mul<Matrix<N, P>> for Matrix<M, N> {
80    type Output = Matrix<M, P>;
81
82    fn mul(self, rhs: Matrix<N, P>) -> Self::Output {
83        let mut result = Matrix::ZEROS;
84
85        for i in 0..M {
86            for j in 0..P {
87                for k in 0..N {
88                    result[[i, j]] += self[[i, k]] * rhs[[k, j]];
89                }
90            }
91        }
92
93        result
94    }
95}
96
97impl<const M: usize, const N: usize> MulAssign<Complex64> for Matrix<M, N> {
98    fn mul_assign(&mut self, rhs: Complex64) {
99        self.iter_mut().par_bridge().for_each(|x| *x *= rhs);
100    }
101}
102
103impl<const M: usize, const N: usize> Mul<Complex64> for Matrix<M, N> {
104    type Output = Matrix<M, N>;
105
106    fn mul(self, rhs: Complex64) -> Self::Output {
107        let mut result = self.clone();
108        result *= rhs;
109        result
110    }
111}
112
113impl<const M: usize, const N: usize> Mul<Matrix<M, N>> for Complex64 {
114    type Output = Matrix<M, N>;
115
116    fn mul(self, rhs: Matrix<M, N>) -> Self::Output {
117        rhs * self
118    }
119}
120
121impl<const M: usize, const N: usize> PartialEq for Matrix<M, N> {
122    fn eq(&self, other: &Self) -> bool {
123        let epsilon = self.epsilon();
124        self.iter()
125            .zip(other.iter())
126            .par_bridge()
127            .all(|(a, b)| (a - b).abs() < epsilon)
128    }
129}
130
131impl<const M: usize, const N: usize> Matrix<M, N> {
132    #[inline]
133    pub const fn new(data: [[Complex64; N]; M]) -> Self {
134        Self { data }
135    }
136
137    pub fn from_column_vectors(columns: [Vector<M>; N]) -> Self {
138        Matrix::<N, M>::new(columns.map(|column| column.transpose().data[0])).transpose()
139    }
140
141    pub fn to_column_vectors(self) -> [Vector<M>; N] {
142        self.transpose()
143            .data
144            .map(|row| Matrix::new([row]).transpose())
145    }
146
147    pub const ZEROS: Self = Matrix::new([[Complex64::ZERO; N]; M]);
148
149    pub fn iter(&self) -> impl Iterator<Item = &Complex64> {
150        self.data.iter().flatten()
151    }
152
153    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut Complex64> {
154        self.data.iter_mut().flatten()
155    }
156
157    pub fn epsilon(&self) -> f64 {
158        self.iter()
159            .par_bridge()
160            .map(|x| x.abs())
161            .max_by(|a, b| a.total_cmp(b))
162            .unwrap()
163            * self.data.len().max(self.data[0].len()) as f64
164            * f64::EPSILON
165    }
166
167    pub fn swap_row(&mut self, i: usize, j: usize) {
168        self.data.swap(i, j);
169    }
170
171    pub fn multiply_row_by_scalar(&mut self, i: usize, scalar: Complex64) {
172        self.data[i]
173            .par_iter_mut()
174            .for_each(|entry| *entry *= scalar);
175    }
176
177    pub fn add_row_with_scalar(&mut self, i: usize, j: usize, scalar: Complex64) {
178        for k in 0..N {
179            let value = scalar * self[[i, k]];
180            self[[j, k]] += value;
181        }
182    }
183
184    pub fn transpose(&self) -> Matrix<N, M> {
185        let mut result = Matrix::ZEROS;
186
187        for i in 0..M {
188            for j in 0..N {
189                result[[j, i]] = self[[i, j]];
190            }
191        }
192
193        result
194    }
195
196    pub fn gaussian_elimination(self) -> (Matrix<M, N>, Matrix<M, M>, usize) {
197        let mut row_echelon_matrix = self;
198        let mut operation_matrix = Matrix::<M, M>::identity();
199        let mut row_swap_count = 0;
200        let size = M.min(N);
201
202        for i in 0..size {
203            let max_entry_index = (i..size)
204                .max_by(|&j, &k| {
205                    row_echelon_matrix[[j, i]]
206                        .abs()
207                        .total_cmp(&row_echelon_matrix[[k, i]].abs())
208                })
209                .unwrap();
210            if row_echelon_matrix[[max_entry_index, i]] == Complex64::ZERO {
211                continue;
212            }
213            if max_entry_index != i {
214                row_echelon_matrix.swap_row(i, max_entry_index);
215                operation_matrix.swap_row(i, max_entry_index);
216                row_swap_count += 1;
217            }
218
219            for j in i + 1..size {
220                let scalar = -row_echelon_matrix[[j, i]] / row_echelon_matrix[[i, i]];
221                row_echelon_matrix.add_row_with_scalar(i, j, scalar);
222                operation_matrix.add_row_with_scalar(i, j, scalar);
223            }
224        }
225
226        (row_echelon_matrix, operation_matrix, row_swap_count)
227    }
228
229    pub fn rank(&self) -> usize {
230        let epsilon = self.epsilon();
231        let (row_echelon_matrix, _, _) = self.clone().gaussian_elimination();
232        row_echelon_matrix
233            .data
234            .iter()
235            .map(|row| row.iter().any(|&entry| entry.abs() >= epsilon))
236            .filter(|&x| x)
237            .count()
238    }
239}
240
241impl<const M: usize> Matrix<M, M> {
242    pub fn diagonal(data: [Complex64; M]) -> Self {
243        let mut result = Matrix::ZEROS;
244        for i in 0..M {
245            result[[i, i]] = data[i];
246        }
247
248        result
249    }
250
251    pub fn identity() -> Self {
252        Self::diagonal([Complex64::ONE; M])
253    }
254
255    pub fn determinant(&self) -> Complex64 {
256        let (row_echelon_matrix, _, row_swap_count) = self.clone().gaussian_elimination();
257        let sign = if row_swap_count % 2 == 0 { 1. } else { -1. };
258        sign * (0..M)
259            .map(|i| row_echelon_matrix[[i, i]])
260            .product::<Complex64>()
261    }
262
263    pub fn gauss_jordan_elimination(self) -> (Matrix<M, M>, Matrix<M, M>) {
264        let epsilon = self.epsilon();
265        let (mut reduced_row_echelon_matrix, mut operation_matrix, _) = self.gaussian_elimination();
266
267        for i in 1..M {
268            if reduced_row_echelon_matrix[[i, i]].abs() <= epsilon {
269                continue;
270            }
271            for j in 0..=i - 1 {
272                let scalar =
273                    -reduced_row_echelon_matrix[[j, i]] / reduced_row_echelon_matrix[[i, i]];
274
275                reduced_row_echelon_matrix.add_row_with_scalar(i, j, scalar);
276                operation_matrix.add_row_with_scalar(i, j, scalar);
277            }
278        }
279
280        for i in 0..M {
281            if reduced_row_echelon_matrix[[i, i]].abs() <= epsilon {
282                continue;
283            }
284
285            let scalar = 1. / reduced_row_echelon_matrix[[i, i]];
286            reduced_row_echelon_matrix.multiply_row_by_scalar(i, scalar);
287            operation_matrix.multiply_row_by_scalar(i, scalar);
288        }
289
290        (reduced_row_echelon_matrix, operation_matrix)
291    }
292
293    pub fn inverse(self) -> Option<Matrix<M, M>> {
294        let (reduced_row_echelon_matrix, operation_matrix) = self.gauss_jordan_elimination();
295        if reduced_row_echelon_matrix == Matrix::<M, M>::identity() {
296            Some(operation_matrix)
297        } else {
298            None
299        }
300    }
301
302    pub fn qr_decomposition(self) -> (Matrix<M, M>, Matrix<M, M>) {
303        let q = Matrix::from_column_vectors(Vector::<M>::gram_schimidt_process(
304            self.clone().to_column_vectors(),
305        ));
306        let r = q.transpose() * self;
307
308        (q, r)
309    }
310
311    pub fn schur_form(self, iterations: usize) -> Matrix<M, M> {
312        let mut result = self;
313
314        for _ in 0..iterations {
315            let (q, r) = result.qr_decomposition();
316            result = r * q;
317        }
318
319        result
320    }
321
322    pub fn eigenvalues(&self, iterations: usize) -> [Complex64; M] {
323        let schur_form = self.clone().schur_form(iterations);
324        let mut eigenvalues: [Complex64; M] = (0..M)
325            .map(|i| schur_form[[i, i]])
326            .collect::<Vec<_>>()
327            .try_into()
328            .unwrap();
329
330        eigenvalues.sort_by(|a, b| b.re().total_cmp(&a.abs()));
331
332        eigenvalues
333    }
334
335    pub fn eigenvectors(&self, iterations: usize) -> [Vector<M>; M] {
336        (0..iterations)
337            .fold(self.clone().qr_decomposition().0, |q, _| {
338                (self.clone() * q).qr_decomposition().0
339            })
340            .to_column_vectors()
341    }
342}
343
344impl<const M: usize, const N: usize> Display for Matrix<M, N> {
345    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
346        let mut table_builder = Builder::default();
347        for i in 0..M {
348            table_builder.push_record(self.data[i].iter().map(|x| x.to_string()));
349        }
350        let mut table = table_builder.build();
351        table.with(Style::blank());
352        table.with(Margin::new(4, 0, 0, 0));
353
354        write!(f, "[\n{}\n]", table)
355    }
356}
numerical_linear_algebra/matrix.rs

numerical_linear_algebra/
matrix.rs
