heapless_matrix/
matrix_ops.rs

1use heapless::Vec;
2
3use core::ops::{Add, AddAssign, Index, IndexMut, Mul, MulAssign, Sub, SubAssign};
4
5use core::cmp::PartialEq;
6
7use core::usize;
8
9use crate::{matrix_trait::MatrixTrait as _, Matrix};
10
11use core::panic;
12
13use core::clone::Clone;
14
15impl<const ROWS: usize, const COLS: usize> Index<usize> for Matrix<ROWS, COLS> {
16    type Output = Vec<f64, COLS>;
17
18    fn index(&self, index: usize) -> &Self::Output {
19        if index >= ROWS {
20            panic!("Row index is out of bounds");
21        }
22        &self.data[index]
23    }
24}
25
26impl<const ROWS: usize, const COLS: usize> IndexMut<usize> for Matrix<ROWS, COLS> {
27    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
28        if index >= ROWS {
29            panic!("Row index is out of bounds");
30        }
31        &mut self.data[index]
32    }
33}
34
35impl<const ROWS: usize, const COLS: usize> Add<Matrix<ROWS, COLS>> for Matrix<ROWS, COLS> {
36    type Output = Matrix<ROWS, COLS>;
37
38    fn add(self, rhs: Matrix<ROWS, COLS>) -> Self::Output {
39        let mut result: Matrix<ROWS, COLS> = Matrix::new().unwrap();
40
41        for i in 0..ROWS {
42            for j in 0..COLS {
43                result[i][j] = self[i][j] + rhs[i][j];
44            }
45        }
46
47        result
48    }
49}
50
51impl<const ROWS: usize, const COLS: usize> AddAssign<Matrix<ROWS, COLS>> for Matrix<ROWS, COLS> {
52    fn add_assign(&mut self, rhs: Matrix<ROWS, COLS>) {
53        for i in 0..ROWS {
54            for j in 0..COLS {
55                self[i][j] += rhs[i][j];
56            }
57        }
58    }
59}
60
61impl<const ROWS: usize, const COLS: usize> Sub<Matrix<ROWS, COLS>> for Matrix<ROWS, COLS> {
62    type Output = Matrix<ROWS, COLS>;
63
64    fn sub(mut self, rhs: Matrix<ROWS, COLS>) -> Self::Output {
65        for i in 0..ROWS {
66            for j in 0..COLS {
67                self[i][j] -= rhs[i][j];
68            }
69        }
70        self
71    }
72}
73
74impl<const ROWS: usize, const COLS: usize> SubAssign<Matrix<ROWS, COLS>> for Matrix<ROWS, COLS> {
75    fn sub_assign(&mut self, rhs: Matrix<ROWS, COLS>) {
76        for i in 0..ROWS {
77            for j in 0..COLS {
78                self[i][j] -= rhs[i][j]
79            }
80        }
81    }
82}
83
84impl<const LHS_ROWS: usize, const LHS_COLS: usize, const RHS_COLS: usize>
85    Mul<Matrix<LHS_COLS, RHS_COLS>> for Matrix<LHS_ROWS, LHS_COLS>
86{
87    type Output = Matrix<LHS_ROWS, RHS_COLS>;
88
89    fn mul(self, rhs: Matrix<LHS_COLS, RHS_COLS>) -> Self::Output {
90        let mut result: Matrix<LHS_ROWS, RHS_COLS> = Matrix::new().unwrap();
91
92        for i in 0..LHS_ROWS {
93            for j in 0..RHS_COLS {
94                let mut res = 0.;
95                for k in 0..LHS_COLS {
96                    res += self[i][k] * rhs[k][j];
97                }
98                result[i][j] = res;
99            }
100        }
101        result
102    }
103}
104
105///
106/// Function that returns the absolute value of a f64 number
107///
108fn abs(x: f64) -> f64 {
109    if x < 0.0 {
110        -x
111    } else {
112        x
113    }
114}
115
116///
117/// Function that checks the equality between two f64 numbers with the tolerance of epsilon
118///
119pub(crate) fn approx_equal(a: f64, b: f64, epsilon: f64) -> bool {
120    abs(a - b) < epsilon
121}
122
123impl<const ROWS: usize, const COLS: usize> PartialEq<Matrix<ROWS, COLS>> for Matrix<ROWS, COLS> {
124    fn eq(&self, other: &Matrix<ROWS, COLS>) -> bool {
125        for i in 0..ROWS {
126            for j in 0..COLS {
127                if !approx_equal(self[i][j], other[i][j], 1e-5) {
128                    return false;
129                }
130            }
131        }
132        true
133    }
134}
135
136impl<const N: usize> MulAssign<Matrix<N, N>> for Matrix<N, N> {
137    fn mul_assign(&mut self, rhs: Matrix<N, N>) {
138        let copy = self.clone();
139        for i in 0..N {
140            for j in 0..N {
141                let mut result = 0.;
142                for k in 0..N {
143                    result += copy[i][k] * rhs[k][j];
144                }
145                self[i][j] = result;
146            }
147        }
148    }
149}
150
151impl<const N: usize> MulAssign<&Matrix<N, N>> for Matrix<N, N> {
152    fn mul_assign(&mut self, rhs: &Matrix<N, N>) {
153        let copy = self.clone();
154        for i in 0..N {
155            for j in 0..N {
156                let mut result = 0.;
157                for k in 0..N {
158                    result += copy[i][k] * rhs[k][j];
159                }
160                self[i][j] = result;
161            }
162        }
163    }
164}
165
166impl<const ROWS: usize, const COLS: usize> MulAssign<f64> for Matrix<ROWS, COLS> {
167    fn mul_assign(&mut self, rhs: f64) {
168        for vec in self.iter_mut() {
169            for elem in vec.iter_mut() {
170                *elem *= rhs;
171            }
172        }
173    }
174}
175
176impl<const ROWS: usize, const COLS: usize> Mul<f64> for Matrix<ROWS, COLS> {
177    type Output = Matrix<ROWS, COLS>;
178
179    fn mul(self, rhs: f64) -> Self::Output {
180        let mut copy = self.clone();
181        copy *= rhs;
182        copy
183    }
184}
185
186impl<const ROWS: usize, const COLS: usize> Mul<Matrix<ROWS, COLS>> for f64 {
187    type Output = Matrix<ROWS, COLS>;
188
189    fn mul(self, rhs: Matrix<ROWS, COLS>) -> Self::Output {
190        rhs * (self)
191    }
192}