1use std::{
2 f32::consts::PI,
3 fmt::{Debug, Display},
4 ops::{Add, AddAssign, Index, IndexMut, Mul, MulAssign, Sub, SubAssign},
5};
6
7use crate::{
8 scalar::{MulAdd, Scalar},
9 vector::{Dot, Vector},
10 V,
11};
12
13#[derive(Clone)]
14pub struct Matrix<K> {
15 pub _d: Vec<K>,
16 pub rows: usize,
17 pub cols: usize,
18}
19
20pub trait Transpose<K> {
21 fn transpose(&self) -> Matrix<K>;
22}
23
24#[macro_export]
25macro_rules! M {
26 ($values:expr) => {
27 Matrix::from($values)
28 };
29}
30
31impl<K: Scalar> Debug for Matrix<K> {
32 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
33 f.write_str("[")?;
34 for i in 0..self.rows {
35 write!(f, "{:?}", &self[i])?;
36 if i != self.rows - 1 {
37 writeln!(f, ",")?;
38 }
39 }
40 f.write_str("]")?;
41 Ok(())
42 }
43}
44
45impl<K: Scalar> Display for Matrix<K> {
46 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
47 f.write_str("[")?;
48 for i in 0..self.rows {
49 write!(f, "{:?}", &self[i])?;
50 if i != self.rows - 1 {
51 writeln!(f, ",")?;
52 }
53 }
54 f.write_str("]")?;
55 Ok(())
56 }
57}
58
59impl<K: Clone, const N: usize, const D: usize> From<[[K; D]; N]> for Matrix<K> {
60 fn from(value: [[K; D]; N]) -> Self {
61 let mut vec = Vec::with_capacity(N * D);
62 (0..N).for_each(|i| {
63 for j in 0..D {
64 vec.push(value[i][j].clone());
65 }
66 });
67
68 Matrix {
69 rows: N,
70 cols: D,
71 _d: vec,
72 }
73 }
74}
75
76impl<K> Index<usize> for Matrix<K> {
77 type Output = [K];
78
79 fn index(&self, index: usize) -> &[K] {
80 let start = index * self.cols;
81 &self._d[start..start + self.cols]
82 }
83}
84
85impl<K> IndexMut<usize> for Matrix<K> {
86 fn index_mut(&mut self, index: usize) -> &mut [K] {
87 let start = index * self.cols;
88 &mut self._d[start..start + self.cols]
89 }
90}
91
92impl<K: Scalar> Add for Matrix<K> {
93 type Output = Self;
94
95 fn add(self, other: Self) -> Self::Output {
96 assert_eq!(
97 self.shape(),
98 other.shape(),
99 "matrices must be the same size"
100 );
101
102 let mut vec = Vec::with_capacity(self.rows * self.cols);
103 for i in 0..self._d.len() {
104 vec.push(self._d[i] + other._d[i]);
105 }
106
107 Matrix {
108 _d: vec,
109 cols: self.cols,
110 rows: self.rows,
111 }
112 }
113}
114
115impl<K: Scalar> AddAssign<&Matrix<K>> for Matrix<K> {
116 fn add_assign(&mut self, rhs: &Matrix<K>) {
117 assert_eq!(self.shape(), rhs.shape(), "matrices must be the same size");
118
119 for i in 0..self._d.len() {
120 self._d[i] += rhs._d[i]
121 }
122 }
123}
124
125impl<K: Scalar> Sub for Matrix<K> {
126 type Output = Self;
127
128 fn sub(self, other: Self) -> Self::Output {
129 assert_eq!(
130 self.shape(),
131 other.shape(),
132 "matrices must be the same size"
133 );
134
135 let mut vec = Vec::with_capacity(self.rows * self.cols);
136 for i in 0..self._d.len() {
137 vec.push(self._d[i] - other._d[i]);
138 }
139
140 Matrix {
141 _d: vec,
142 cols: self.cols,
143 rows: self.rows,
144 }
145 }
146}
147
148impl<K: Scalar> SubAssign<&Matrix<K>> for Matrix<K> {
149 fn sub_assign(&mut self, rhs: &Matrix<K>) {
150 assert_eq!(self.shape(), rhs.shape(), "matrices must be the same size");
151
152 for i in 0..self._d.len() {
153 self._d[i] -= rhs._d[i];
154 }
155 }
156}
157
158impl<K: Scalar + Mul<U, Output = K>, U: Scalar> Mul<U> for Matrix<K> {
159 type Output = Matrix<K>;
160
161 fn mul(self, a: U) -> Self::Output {
162 let mut vec = Vec::with_capacity(self._d.len());
163 for i in 0..self._d.len() {
164 vec.push(self._d[i] * a);
165 }
166
167 Matrix {
168 _d: vec,
169 cols: self.cols,
170 rows: self.rows,
171 }
172 }
173}
174
175impl<K: Scalar + MulAssign<U>, U: Scalar> MulAssign<&U> for Matrix<K> {
176 fn mul_assign(&mut self, rhs: &U) {
177 for i in 0..self._d.len() {
178 self._d[i] *= *rhs;
179 }
180 }
181}
182
183impl<K: Scalar + Mul<U, Output = K> + MulAdd<U, K>, U: Scalar>
184 MulAdd<U, Matrix<K>> for Matrix<K>
185{
186 fn mul_add(self, a: &U, b: &Matrix<K>) -> Self {
187 assert!(self.shape() == b.shape(), "matrices must be the same size");
188
189 let mut vec = Vec::with_capacity(self.rows * self.cols);
190
191 for i in 0..self._d.len() {
192 vec.push(self._d[i].mul_add(a, &b._d[i]));
193 }
194
195 Matrix {
196 _d: vec,
197 rows: self.rows,
198 cols: self.cols,
199 }
200 }
201}
202
203impl<'a, K: Scalar> Mul<&Vector<K>> for &Matrix<K>
204where
205 [K]: Dot<K>,
206{
207 type Output = Vector<K>;
208
209 fn mul(self, rhs: &Vector<K>) -> Self::Output {
210 assert_eq!(
211 self.shape().0,
212 rhs.size(),
213 "bad input for matrix and vector column multiplication"
214 );
215
216 let mut vec = Vec::with_capacity(rhs.size());
217 for i in 0..rhs.size() {
218 vec.push(self[i].dot(rhs));
219 }
220 V!(vec)
221 }
222}
223
224impl<K: Scalar> Mul<&Matrix<K>> for &Matrix<K> {
225 type Output = Matrix<K>;
226
227 fn mul(self, rhs: &Matrix<K>) -> Self::Output {
228 assert_eq!(
229 self.cols, rhs.rows,
230 "bad input for matrix and matrix multiplication"
231 );
232
233 let cols = rhs.cols;
234 let rows = self.rows;
235 let mut vec = Vec::with_capacity(rows * cols);
236
237 for i in 0..rows {
238 for j in 0..cols {
239 let mut val = self[i][0] * rhs[0][j];
240 for r in 1..self.cols {
241 val += self[i][r] * rhs[r][j];
242 }
243 vec.push(val);
244 }
245 }
246
247 Matrix {
248 _d: vec,
249 rows,
250 cols,
251 }
252 }
253}
254
255impl<K: Scalar> Matrix<K> {
256 pub fn shape(&self) -> (usize, usize) {
257 (self.rows, self.cols)
258 }
259
260 pub fn is_square(&self) -> bool {
261 self.rows == self.cols
262 }
263
264 pub fn add(&mut self, v: &Matrix<K>) {
265 *self += v;
266 }
267
268 pub fn sub(&mut self, v: &Matrix<K>) {
269 *self -= v;
270 }
271
272 pub fn scl(&mut self, a: K) {
273 *self *= &a;
274 }
275
276 pub fn mul_vec(&self, vec: &Vector<K>) -> Vector<K>
277 where
278 [K]: Dot<K>,
279 {
280 self * vec
281 }
282
283 pub fn mul_mat(&self, mat: &Matrix<K>) -> Matrix<K> {
284 self * mat
285 }
286
287 pub fn trace(&self) -> Option<K> {
288 assert!(self.is_square(), "matrix must be squared");
289
290 match self.rows {
291 0 => None,
292 _ => Some((0..self.rows).map(|i| self[i][i]).sum()),
293 }
294 }
295
296 pub fn row_echelon(&self) -> Matrix<K> {
297 let mut ret = self.clone();
298 ret.row_echelon_mut();
299 ret
300 }
301
302 pub fn row_echelon_mut(&mut self) -> &mut Self {
303 let mut col: usize = 0;
304
305 for r in 0..self.rows {
306 let mut p = None;
307
308 for j in col..self.cols {
309 for i in r..self.rows {
310 if self[i][j].is_non_zero() {
311 p = Some((i, j));
312 break;
313 }
314 }
315
316 if p.is_some() {
317 break;
318 }
319 }
320
321 if p.is_none() {
322 break;
323 }
324 let cur = p.unwrap();
325 col = cur.1;
326 if r != cur.0 {
327 for c in col..self.cols {
328 let tmp = self[r][c];
329 self[r][c] = self[cur.0][c];
330 self[cur.0][c] = tmp;
331 }
332 }
333
334 let v = self[r][col];
335
336 self[r][col] = K::one();
337 for i in col + 1..self.cols {
338 self[r][i] /= v;
339 }
340
341 for i in 0..self.rows {
342 if i == r {
343 continue;
344 }
345 let cur = self[i][col];
346 self[i][col] = K::default();
347 for j in col + 1..self.cols {
348 let x = self[r][j];
349 self[i][j] -= x * cur;
350 }
351 }
352
353 col += 1;
354 }
355
356 self
357 }
358
359 pub fn determinant(&self) -> K {
360 assert!(self.is_square(), "matrix must be squared");
361
362 fn det<K: Scalar>(
363 mat: &[K],
364 row_size: usize,
365 cur: usize,
366 skip_cols: &[bool],
367 ) -> K {
368 let mut cols = Vec::with_capacity(skip_cols.len());
369
370 for (col, skip) in skip_cols.iter().enumerate() {
371 if !*skip {
372 cols.push(col);
373 }
374 }
375
376 match cur {
377 0 => K::default(),
378 1 => mat[0],
379 2 => {
380 mat[cols[0]] * mat[row_size + cols[1]]
381 - mat[cols[1]] * mat[row_size + cols[0]]
382 }
383 _ => {
384 let mut result = K::default();
385 for (i, col) in cols.iter().enumerate() {
386 let mut m_col = skip_cols.to_vec();
387 m_col[*col] = true;
388 let ret = mat[*col]
389 * det(&mat[row_size..], row_size, cur - 1, &m_col);
390 result += if i % 2 == 1 { -ret } else { ret };
391 m_col[*col] = false;
392 }
393 result
394 }
395 }
396 }
397
398 det(
399 self._d.as_slice(),
400 self.rows,
401 self.rows,
402 &vec![false; self.cols],
403 )
404 }
405
406 pub fn inverse(&self) -> Result<Matrix<K>, &'static str> {
407 assert!(self.is_square(), "matrix must be squared");
408
409 assert_ne!(self.determinant(), K::default(), "matrix must be singular");
410
411 let mut aug_m = Matrix {
412 _d: vec![K::default(); self.rows * self.cols * 2],
413 cols: self.cols * 2,
414 rows: self.rows,
415 };
416
417 for row in 0..self.rows {
418 for col in 0..self.cols {
419 aug_m._d[row * (self.cols * 2) + col] = self[row][col];
420 }
421 aug_m._d[row * (self.cols * 2) + self.cols + row] = K::one();
422 }
423
424 aug_m.row_echelon_mut();
425
426 let mut vec = Vec::with_capacity(self.rows * self.cols);
427
428 for row in aug_m._d.chunks(self.cols * 2) {
429 for &v in row.iter().skip(self.cols) {
430 vec.push(v);
431 }
432 }
433
434 Ok(Matrix {
435 _d: vec,
436 cols: self.cols,
437 rows: self.rows,
438 })
439 }
440
441 pub fn rank(&self) -> usize {
442 let mat = self.row_echelon();
443 let mut rank = 0;
444
445 let mut cur = 0;
446 for row in mat._d.chunks(self.cols) {
447 for (j, v) in row.iter().enumerate().skip(cur) {
448 if v.is_non_zero() {
449 cur = j + 1;
450 rank += 1;
451 break;
452 }
453 }
454 }
455
456 rank
457 }
458
459 pub fn is_independent(&self) -> bool {
460 self.rank() == self.rows
461 }
462
463 pub fn row_space(&self) -> Matrix<K> {
464 let mat = self.row_echelon();
465 let mut vec = vec![];
466
467 let mut cur = 0;
468 let mut rank = 0;
469 for row in mat._d.chunks(self.cols) {
470 for (j, v) in row.iter().enumerate().skip(cur) {
471 if *v != K::default() {
472 vec.copy_from_slice(row);
473 cur = j + 1;
474 rank += 1;
475 break;
476 }
477 }
478 }
479
480 Matrix {
481 cols: self.cols,
482 rows: rank,
483 _d: vec,
484 }
485 }
486
487 pub fn col_space(&self) -> Matrix<K> {
488 let mat = self.row_echelon();
489 let mut vec = vec![];
490
491 let mut cur = 0;
492 let mut rank = 0;
493 for row in mat._d.chunks(self.cols) {
494 for (j, v) in row.iter().enumerate().skip(cur) {
495 if *v != K::default() {
496 for r in self._d.chunks(self.cols) {
497 vec.push(r[cur])
498 }
499 cur = j + 1;
500 rank += 1;
501 break;
502 }
503 }
504 }
505
506 Matrix {
507 cols: self.rows,
508 rows: rank,
509 _d: vec,
510 }
511 }
512}
513
514pub fn projection(fov: f32, ratio: f32, n: f32, f: f32) -> Matrix<f32> {
515 let tan = ((fov / 2.) * (PI / 180.)).tan();
553 let t = n * tan;
554 let r = t * ratio;
555 let (l, b) = (-r, -t);
556
557 M!([
558 [2. * n / (r - l), 0., (r + l) / (r - l), 0.], [0., 2. * n / (t - b), (t + b) / (t - b), 0.], [0., 0., f / (n - f), (n * f) / (n - f)], [0., 0., -1., 0.], ])
563}