1#![allow(clippy::indexing_slicing)]
6
7use crate::f64::FloatExt;
8use crate::{One, ToSigned, Zero};
9use core::ops::{Add, AddAssign, Deref, DerefMut, Index, IndexMut, Mul, Neg, Sub};
10
11pub trait AsMatrix<const M: usize, const N: usize, T: Sized + Copy + Default> {
12 fn as_matrix(&self) -> Matrix<M, N, T>;
13}
14
15#[derive(Debug, Copy, Clone, Eq, PartialEq)]
16pub struct Matrix<const M: usize, const N: usize, T: Sized + Copy + Default> {
17 pub values: [[T; N]; M],
18}
19
20impl<const M: usize, const N: usize, T: Sized + Copy + Default> Matrix<M, N, T> {
21 #[must_use]
22 pub const fn new(values: [[T; N]; M]) -> Matrix<M, N, T> {
23 Matrix { values }
24 }
25
26 #[must_use]
28 pub fn submatrix(&self, rmin: usize, rmax: usize, cmin: usize, cmax: usize) -> Matrix<M, N, T> {
29 let mut values = [[T::default(); N]; M];
30 for i in rmin..rmax {
31 for j in cmin..cmax {
32 values[i - rmin][j - cmin] = self.values[i][j];
33 }
34 }
35 Matrix { values }
36 }
37 pub fn augment<const O: usize, const R: usize>(
38 &self,
39 other: &Matrix<M, O, T>,
40 ) -> Matrix<M, R, T> {
41 augment(self, other)
42 }
43
44 pub fn swap_rows(&mut self, r1: usize, r2: usize) {
47 assert!(r1 < M && r2 < M, "Values out of range");
48 self.values.swap(r1, r2);
49 }
50}
51fn augment<
52 const M: usize,
53 const N: usize,
54 const O: usize,
55 const R: usize,
56 T: Sized + Copy + Default,
57>(
58 first: &Matrix<M, N, T>,
59 second: &Matrix<M, O, T>,
60) -> Matrix<M, R, T> {
61 assert_eq!(R, N + O, "Values out of range");
62 let mut values = [[T::default(); R]; M];
63 for r in 0..M {
64 for c in 0..N {
65 values[r][c] = first[r][c];
66 }
67 for c in 0..O {
68 values[r][c + N] = second[r][c];
69 }
70 }
71 Matrix { values }
72}
73
74impl<const M: usize, const N: usize, T: Sized + Copy + Default> From<[[T; N]; M]>
75 for Matrix<M, N, T>
76{
77 fn from(value: [[T; N]; M]) -> Self {
78 Self { values: value }
79 }
80}
81impl<const M: usize, const N: usize, T: Sized + Copy + Default> From<&[[T; N]; M]>
82 for Matrix<M, N, T>
83{
84 fn from(value: &[[T; N]; M]) -> Self {
85 Self { values: *value }
86 }
87}
88impl<const M: usize, const N: usize, T: Sized + Copy + Default> AsMatrix<M, N, T> for [[T; N]; M] {
89 fn as_matrix(&self) -> Matrix<M, N, T> {
90 self.into()
91 }
92}
93impl<const M: usize, const N: usize, T: Default + Copy + Zero + AddAssign + Mul<Output = T>>
94 Matrix<M, N, T>
95{
96 #[must_use]
97 pub fn mul<const P: usize>(&self, other: Matrix<N, P, T>) -> Matrix<M, P, T> {
98 let mut out = [[T::ZERO; P]; M];
99 let mut m = 0;
100 while m < M {
101 let mut p = 0;
102 while p < P {
103 let mut n = 0;
104 let mut sum = T::ZERO;
105 while n < N {
106 sum += self.values[m][n] * other.values[n][p];
107 n += 1;
108 }
109 out[m][p] = sum;
110 p += 1;
111 }
112 m += 1;
113 }
114 Matrix { values: out }
115 }
116}
117#[derive(Debug, Copy, Clone, Eq, PartialEq)]
118pub struct LUPDecomposition<const M: usize, const N: usize, T: Sized + Copy + Default> {
119 pub lower: Matrix<M, N, T>,
120 pub upper: Matrix<M, N, T>,
121 pub permuted: Matrix<M, N, T>,
122}
123
124macro_rules! impl_square {
125 ($N:literal) => {
126 impl<T: Sized + Copy + Default + One> Matrix<$N, $N, T> {
127 #[must_use]
128 pub fn empty() -> Self {
129 Self {
130 values: [<[T; $N]>::default(); $N],
131 }
132 }
133 #[must_use]
134 pub fn identity() -> Self {
135 let mut out = Self::empty();
136 for i in 0..$N {
137 out[i][i] = One::ONE;
138 }
139 out
140 }
141 #[must_use]
142 pub fn transpose(&self) -> Self {
143 let mut out = Self::empty();
144 for i in 0..$N {
145 for j in 0..$N {
146 out[i][j] = self.values[j][i];
147 }
148 }
149 out
150 }
151 }
152 impl Matrix<$N, $N, f64> {
153 pub fn lup_decompose(&self) -> LUPDecomposition<$N, $N, f64> {
154 let mut l = Self::identity();
155 let mut u = self.clone();
156 let mut p = Self::identity();
157
158 for k in 0..$N {
159 let mut max = 0.0;
160 let mut pivot = k;
161 for i in k..$N {
162 if u[i][k].abs() > max {
163 max = u[i][k].abs();
164 pivot = i;
165 }
166 }
167
168 if pivot != k {
169 for j in 0..$N {
171 let temp = u[k][j];
172 u[k][j] = u[pivot][j];
173 u[pivot][j] = temp;
174
175 let temp = p[k][j];
176 p[k][j] = p[pivot][j];
177 p[pivot][j] = temp;
178
179 if j < k {
180 let temp = l[k][j];
181 l[k][j] = l[pivot][j];
182 l[pivot][j] = temp;
183 }
184 }
185 }
186
187 for j in (k + 1)..$N {
188 l[j][k] = u[j][k] / u[k][k];
189 for i in k..$N {
190 u[j][i] -= l[j][k] * u[k][i];
191 }
192 }
193 }
194
195 LUPDecomposition {
196 lower: l,
197 upper: u,
198 permuted: p,
199 }
200 }
201 }
202
203 impl LUPDecomposition<$N, $N, f64> {
204 pub fn solve_ax_eq_b(&self, b: &[f64; $N]) -> [f64; $N] {
205 let mut b2 = [0usize; $N];
206 for i in 0..$N {
207 for j in 0..$N {
208 if (self.permuted[i][j] - 1.).abs() < f64::EPSILON {
209 b2[i] = j;
210 break;
211 }
212 }
213 }
214 let mut y: [f64; $N] = Default::default();
216 for i in 1..$N {
217 let mut sum = 0.0;
218 for j in i..i {
219 sum += self.lower[i][j] * y[j];
220 }
221 y[i] = (b[b2[i]] - sum) / self.lower[i][i];
222 }
223 let mut x: [f64; $N] = Default::default();
224 for i in (1..$N).rev() {
225 let mut sum = 0.0;
226 for j in i..$N {
227 sum += self.upper[i][j] * x[j];
228 }
229 x[i] = (y[i] - sum) / self.upper[i][i];
230 }
231 x
239 }
240 }
241 };
242}
243impl_square!(2);
244impl_square!(3);
245impl_square!(4);
246impl_square!(5);
247impl_square!(6);
248impl_square!(7);
249impl_square!(8);
250impl_square!(9);
251impl_square!(10);
252
253impl Matrix<2, 2, f64> {
254 #[must_use]
255 pub fn rotation_counterclockwise(angle: f64) -> Self {
256 Self::new([[angle.cos(), -angle.sin()], [angle.sin(), angle.cos()]])
257 }
258 #[must_use]
259 pub fn rotate_counterclockwise(&self, angle: f64) -> Self {
260 self.mul(Self::rotation_counterclockwise(angle))
261 }
262 #[must_use]
263 pub fn rotation_clockwise(angle: f64) -> Self {
264 Self::new([[angle.cos(), angle.sin()], [-angle.sin(), angle.cos()]])
265 }
266 #[must_use]
267 pub fn rotate_clockwise(&self, angle: f64) -> Self {
268 self.mul(Self::rotation_clockwise(angle))
269 }
270
271 #[must_use]
272 pub fn sheered_x(factor: f64) -> Self {
273 Self::new([[1., factor], [0., 1.]])
274 }
275 #[must_use]
276 pub fn sheer_x(&self, factor: f64) -> Self {
277 self.mul(Self::sheered_x(factor))
278 }
279 #[must_use]
280 pub const fn sheered_y(factor: f64) -> Self {
281 Self::new([[1., 0.], [factor, 1.]])
282 }
283 #[must_use]
284 pub fn sheer_y(&self, factor: f64) -> Self {
285 self.mul(Self::sheered_y(factor))
286 }
287
288 #[must_use]
289 pub const fn scaled_x(factor: f64) -> Self {
290 Self::new([[factor, 0.], [0., 1.]])
291 }
292 #[must_use]
293 pub fn scale_x(&self, factor: f64) -> Self {
294 self.mul(Self::scaled_x(factor))
295 }
296
297 #[must_use]
298 pub const fn scaled_y(factor: f64) -> Self {
299 Self::new([[1., 0.], [0., factor]])
300 }
301 #[must_use]
302 pub fn scale_y(&self, factor: f64) -> Self {
303 self.mul(Self::scaled_y(factor))
304 }
305
306 #[must_use]
307 pub const fn scaled(factor: f64) -> Self {
308 Self::new([[factor, 0.], [0., factor]])
309 }
310 #[must_use]
311 pub fn scale(&self, factor: f64) -> Self {
312 self.mul(Self::scaled(factor))
313 }
314 #[must_use]
315 pub const fn reflected() -> Self {
316 Self::new([[-1., 0.], [0., -1.]])
317 }
318 #[must_use]
319 pub fn reflect(&self) -> Self {
320 self.mul(Self::reflected())
321 }
322
323 #[must_use]
324 pub const fn reflected_x() -> Self {
325 Self::new([[1., 0.], [0., -1.]])
326 }
327 #[must_use]
328 pub fn reflect_x(&self) -> Self {
329 self.mul(Self::reflected_x())
330 }
331 #[must_use]
332 pub const fn reflected_y() -> Self {
333 Self::new([[-1., 0.], [0., 1.]])
334 }
335 #[must_use]
336 pub fn reflect_y(&self) -> Self {
337 self.mul(Self::reflected_y())
338 }
339}
340impl<
341 T: Copy
342 + Default
343 + FloatExt<Type = T>
344 + One
345 + Zero
346 + Neg<Output = T>
347 + Add
348 + AddAssign<T>
349 + Mul<Output = T>,
350 > Matrix<3, 1, T>
351{
352 #[must_use]
353 pub fn translate(&self, x: T, y: T) -> Self {
354 Matrix::mul(
355 &Matrix::new([
356 [T::ONE, T::ZERO, x],
357 [T::ZERO, T::ONE, y],
358 [T::ZERO, T::ZERO, T::ONE],
359 ]),
360 *self,
361 )
362 }
363 cfg_feature_std! {
364 #[must_use]
365 pub fn rotate_x(&self, angle: T) -> Self {
366 Matrix::<3, 3, T>::rotated_x(angle).mul(*self)
367 }
368 #[must_use]
369 pub fn rotate_y(&self, angle: T) -> Self {
370 Matrix::<3, 3, T>::rotated_y(angle).mul(*self)
371 }
372 #[must_use]
373 pub fn rotate_z(&self, angle: T) -> Self {
374 Matrix::<3, 3, T>::rotated_z(angle).mul(*self)
375 }
376 #[must_use]
377 pub fn rotate_zyx(&self, x_angle: T, y_angle: T, z_angle: T) -> Self {
378 Matrix::<3, 3, T>::rotated_zyx(x_angle, y_angle, z_angle).mul(*self)
379 }
380 }
381}
382impl<
383 T: Copy
384 + Default
385 + FloatExt<Type = T>
386 + One
387 + Zero
388 + Neg<Output = T>
389 + Add
390 + AddAssign<T>
391 + Mul<Output = T>,
392 > Matrix<3, 3, T>
393{
394 cfg_feature_std! {
395 #[must_use]
396 pub fn rotated_x(angle: T) -> Self {
397 Self::new([
398 [T::ONE, T::ZERO, T::ZERO],
399 [T::ZERO, angle.cos(), -angle.sin()],
400 [T::ZERO, angle.sin(), angle.cos()],
401 ])
402 }
403 #[must_use]
404 pub fn rotate_x(&self, angle: T) -> Self {
405 self.mul(Self::rotated_x(angle))
406 }
407 #[must_use]
408 pub fn rotated_y(angle: T) -> Self {
409 Self::new([
410 [angle.cos(), T::ZERO, angle.sin()],
411 [T::ZERO, T::ONE, T::ZERO],
412 [-angle.sin(), T::ZERO, angle.cos()],
413 ])
414 }
415 #[must_use]
416 pub fn rotate_y(&self, angle: T) -> Self {
417 self.mul(Self::rotated_y(angle))
418 }
419 #[must_use]
420 pub fn rotated_z(angle: T) -> Self {
421 Self::new([
422 [angle.cos(), angle.sin(), T::ZERO],
423 [-angle.sin(), angle.cos(), T::ZERO],
424 [T::ZERO,T::ZERO, T::ONE],
425 ])
426 }
427 #[must_use]
428 pub fn rotate_z(&self, angle: T) -> Self {
429 self.mul(Self::rotated_z(angle))
430 }
431
432 #[must_use]
433 pub fn rotated_zyx(x_angle: T, y_angle: T, z_angle: T) -> Self {
434 Self::rotated_z(z_angle).rotate_y(y_angle).rotate_x(x_angle)
435 }
436 #[must_use]
437 pub fn rotate_zyx(&self, x_angle: T, y_angle: T, z_angle: T) -> Self {
438 self.rotate_z(z_angle).rotate_y(y_angle).rotate_x(x_angle)
439 }
440
441 #[must_use]
442 pub const fn scaled_x(factor: T) -> Self {
443 Self::new([[factor, T::ZERO, T::ZERO], [T::ZERO, T::ONE, T::ZERO], [T::ZERO, T::ZERO, T::ONE]])
444 }
445 #[must_use]
446 pub fn scale_x(&self, factor: T) -> Self {
447 self.mul(Self::scaled_x(factor))
448 }
449
450 #[must_use]
451 pub const fn scaled_y(factor: T) -> Self {
452 Self::new([[T::ONE, T::ZERO, T::ZERO], [T::ZERO, factor, T::ZERO], [T::ZERO, T::ZERO, T::ONE]])
453 }
454 #[must_use]
455 pub fn scale_y(&self, factor: T) -> Self {
456 self.mul(Self::scaled_y(factor))
457 }
458 #[must_use]
459 pub const fn scaled_z(factor: T) -> Self {
460 Self::new([[T::ONE, T::ZERO, T::ZERO], [T::ZERO, T::ONE, T::ZERO], [T::ZERO, T::ZERO, factor]])
461 }
462 #[must_use]
463 pub fn scale_z(&self, factor: T) -> Self {
464 self.mul(Self::scaled_z(factor))
465 }
466 }
467}
468
469impl<const M: usize, const N: usize, T: Sized + Copy + Default> Index<usize> for Matrix<M, N, T> {
470 type Output = [T; N];
471
472 fn index(&self, index: usize) -> &Self::Output {
473 self.values.index(index)
474 }
475}
476impl<const M: usize, const N: usize, T: Sized + Copy + Default> IndexMut<usize>
477 for Matrix<M, N, T>
478{
479 fn index_mut(&mut self, index: usize) -> &mut Self::Output {
480 self.values.index_mut(index)
481 }
482}
483
484impl<const M: usize, const N: usize, T: Sized + Copy + Default> Deref for Matrix<M, N, T> {
485 type Target = [[T; N]; M];
486
487 fn deref(&self) -> &Self::Target {
488 &self.values
489 }
490}
491impl<const M: usize, const N: usize, T: Sized + Copy + Default> DerefMut for Matrix<M, N, T> {
492 fn deref_mut(&mut self) -> &mut Self::Target {
493 &mut self.values
494 }
495}
496macro_rules! impl_mul {
498 ($($ty:ty)+) => {
499 impl<
500 const M: usize,
501 const N: usize,
502 const P: usize,
503 T: Sized + Copy + Default + Add + Mul + AddAssign<<T as Mul<T>>::Output>,
504 > Mul<Matrix<N, P, T>> for $($ty)+
505 {
506 type Output = Matrix<M, P, T>;
507 fn mul(self, other: Matrix<N, P, T>) -> Matrix<M, P, T> {
508 let mut out = [[T::default(); P]; M];
509 let mut m = 0;
510 while m < M {
511 let mut p = 0;
512 while p < P {
513 let mut n = 0;
514 let mut sum = T::default();
515 while n < N {
516 sum += self.values[m][n] * other.values[n][p];
517 n += 1;
518 }
519 out[m][p] = sum;
520 p += 1;
521 }
522 m += 1;
523 }
524 Matrix { values: out }
525 }
526 }
527 };
528}
529impl_mul!(Matrix<M, N, T>);
530impl_mul!(&Matrix<M, N, T>);
531impl_mul!(&mut Matrix<M, N, T>);
532
533impl<const M: usize, const N: usize, T: Sized + Copy + Default + Add<Output = T>> Add
535 for Matrix<M, N, T>
536{
537 type Output = Matrix<M, N, T>;
538 fn add(self, other: Matrix<M, N, T>) -> Matrix<M, N, T> {
539 let mut out = [[T::default(); N]; M];
540
541 for (i, ith) in out.iter_mut().enumerate().take(M) {
542 for (j, val) in ith.iter_mut().enumerate().take(N) {
543 *val = self.values[i][j] + other.values[i][j];
544 }
545 }
546 Matrix { values: out }
547 }
548}
549impl<const M: usize, const N: usize, T: Sized + Copy + Default + Mul<T, Output = T>> Mul<T>
551 for Matrix<M, N, T>
552{
553 type Output = Matrix<M, N, T>;
554 fn mul(self, other: T) -> Matrix<M, N, T> {
555 let mut out = [[T::default(); N]; M];
556
557 for (i, ith) in out.iter_mut().enumerate().take(M) {
558 for (j, val) in ith.iter_mut().enumerate().take(N) {
559 *val = self.values[i][j] * other;
560 }
561 }
562 Matrix { values: out }
563 }
564}
565
566impl<
567 const M: usize,
568 const N: usize,
569 T: Sized + Copy + Default + Add<Output = T> + Mul<T, Output = T> + ToSigned<Output = T>,
570 > Sub for Matrix<M, N, T>
571{
572 type Output = Matrix<M, N, T>;
573
574 fn sub(self, rhs: Self) -> Self::Output {
575 let v = rhs * <T as ToSigned>::negative_one();
576 self + v
577 }
578}
579
580#[cfg(test)]
581mod test {
582 use crate::math::{AsMatrix, LUPDecomposition, Matrix};
583 use core::ops::Deref;
584
585 #[test]
586 pub fn test_scalar() {
587 let mat = Matrix::new([[4, 0], [1, -9]]);
588 let res = mat * 2;
589 assert_eq!(res, Matrix::new([[8, 0], [2, -18]]));
590 }
591
592 #[test]
593 pub fn test_product() {
594 let m1 = Matrix::new([[1, 2, 3], [4, 5, 6]]);
595 let m2 = Matrix::new([[7, 8], [9, 10], [11, 12]]);
596 let res = m1 * m2;
597 assert_eq!(res, Matrix::new([[58, 64], [139, 154]]));
598 }
599
600 #[cfg(feature = "std")]
601 #[test]
602 pub fn test_rotate1() {
603 let m = [[3.], [7.], [4.]].as_matrix();
604 let [[x], [y], [z]] = *m.rotate_x(core::f64::consts::FRAC_PI_2).deref();
605 assert_eq_eps!(3., x, 2. * f64::EPSILON);
606 assert_eq_eps!(-4., y, 2. * f64::EPSILON);
607 assert_eq_eps!(7., z, 2. * f64::EPSILON);
608 }
609
610 #[cfg(feature = "std")]
611 #[test]
612 pub fn test_rotate2() {
613 let m = [[3.], [7.], [4.]].as_matrix();
614 let [[x], [y], [z]] = *m.rotate_y(core::f64::consts::FRAC_PI_2).deref();
615 assert_eq_eps!(4., x, 2. * f64::EPSILON);
616 assert_eq_eps!(7., y, 2. * f64::EPSILON);
617 assert_eq_eps!(-3., z, 2. * f64::EPSILON);
618 }
619
620 #[cfg(feature = "std")]
621 #[test]
622 pub fn test_rotate3() {
623 let m = [[3.], [7.], [4.]].as_matrix();
624 let [[x], [y], [z]] = *m.rotate_z(core::f64::consts::FRAC_PI_2).deref();
625 assert_eq_eps!(7., x, 2. * f64::EPSILON);
626 assert_eq_eps!(-3., y, 2. * f64::EPSILON);
627 assert_eq_eps!(4., z, 2. * f64::EPSILON);
628 }
629
630 #[cfg(feature = "std")]
631 #[test]
632 pub fn test_lup1() {
633 let a = [
634 [2., 0.0, 2., 0.6],
635 [3., 3., 4., -2.],
636 [5., 5., 4., 2.],
637 [-1., -2., 3.4, -1.],
638 ]
639 .as_matrix();
640 let LUPDecomposition {
641 lower,
642 upper,
643 permuted,
644 } = a.lup_decompose();
645 assert_eq!(
646 lower,
647 [
648 [1.0, 0.0, 0.0, 0.0],
649 [0.4, 1.0, 0.0, 0.0],
650 [-0.2, 0.5, 1.0, 0.0],
651 [0.6, -0.0, 0.4, 1.0]
652 ]
653 .as_matrix()
654 );
655 assert_eq!(
656 upper,
657 [
658 [5.0, 5.0, 4.0, 2.0],
659 [0.0, -2.0, 0.3999999999999999, -0.20000000000000007],
660 [0.0, 0.0, 4.0, -0.49999999999999994],
661 [0.0, 0.0, 0.0, -3.0]
662 ]
663 .as_matrix()
664 );
665 assert_eq!(
666 permuted,
667 [
668 [0.0, 0.0, 1.0, 0.0],
669 [1.0, 0.0, 0.0, 0.0],
670 [0.0, 0.0, 0.0, 1.0],
671 [0.0, 1.0, 0.0, 0.0]
672 ]
673 .as_matrix()
674 );
675
676 let c1 = permuted * a;
677 let c2 = lower * upper;
678 assert_eq!(c1, c2);
679 }
680
681 #[cfg(feature = "std")]
682 #[test]
683 pub fn test_solve1() {
684 let a = [[25., 5., 1.], [64., 8., 1.], [144., 12., 1.]].as_matrix();
685 let lup = a.lup_decompose();
686 println!("{:?}", lup);
687 let c1 = lup.permuted * a;
688 let c2 = lup.lower * lup.upper;
689 assert_eq!(c1, c2);
690
691 let b = [106.8, 177.2, 279.2];
692 let c = [b].as_matrix() * lup.permuted;
693 println!("{c:?}");
694 let res = lup.solve_ax_eq_b(&b);
695 println!("{:?}", res);
696 }
697
698 #[cfg(feature = "std")]
699 #[test]
700 pub fn test_lup2() {
701 let a = [[3., 17., 10.], [2., 4., -2.], [6., 18., -12.]].as_matrix();
702 let lup = a.lup_decompose();
703 assert_eq!(
704 lup.lower,
705 [[1., 0., 0.], [0.5, 1., 0.], [1. / 3., -0.25, 1.]].as_matrix()
706 );
707 assert_eq!(
708 lup.upper,
709 [[6., 18., -12.], [0., 8., 16.], [0., 0., 6.]].as_matrix()
710 );
711 }
712
713 #[cfg(feature = "std")]
714 #[test]
715 pub fn test_lup3() {
716 let a = [[2., 1., -5.], [4., 4., -4.], [1., 3., 1.]].as_matrix();
717 let lup = a.lup_decompose();
718 println!("{:?}", lup);
719 let c1 = lup.permuted * a;
720 let c2 = lup.lower * lup.upper;
721 assert_eq!(c1, c2);
722
723 let b = [5., 0., 6.];
724 let res = lup.solve_ax_eq_b(&b);
725 println!("{:?}", res);
726
727 let b2 = [res].as_matrix() * a;
728 println!("{:?}", b2);
729 }
730}