magnesia/linalg/
smatrix.rs

1use super::SVector;
2use crate::algebra::{Conj, MulRefs, One, Ring, Zero};
3use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign};
4
5/// A matrix type with static number of rows and columns.
6///
7/// # Example
8/// ```
9/// # use magnesia::linalg::SMatrix;
10/// let a = SMatrix::from([[0,1,2],[3,4,5]]);
11/// let b = SMatrix::from([[1,1,1],[1,1,1]]);
12/// let c = a + b;
13/// assert_eq!(c, SMatrix::from([[1,2,3],[4,5,6]]));
14/// ```
15#[derive(Copy, Clone, Debug, PartialEq, Eq)]
16pub struct SMatrix<T, const NUM_ROWS: usize, const NUM_COLS: usize>([[T; NUM_COLS]; NUM_ROWS]);
17
18impl<T: Zero, const NUM_ROWS: usize, const NUM_COLS: usize> Zero
19    for SMatrix<T, NUM_ROWS, NUM_COLS>
20{
21    /// Returns a matrix filled with zeros.
22    ///
23    /// # Example
24    ///
25    /// ```
26    /// # use magnesia::linalg::SMatrix;
27    /// # use magnesia::algebra::Zero;
28    /// let m = SMatrix::zero();
29    /// assert_eq!(m, SMatrix::from([[0,0],[0,0],[0,0]]));
30    /// ```
31    fn zero() -> Self {
32        Self([(); NUM_ROWS].map(|_| [(); NUM_COLS].map(|_| T::zero())))
33    }
34}
35
36impl<T: Zero + One, const NUM_ROWS: usize, const NUM_COLS: usize> One
37    for SMatrix<T, NUM_ROWS, NUM_COLS>
38{
39    /// Returns a matrix filled with ones on the diagonal and zeros everywhere else.
40    ///
41    /// # Example
42    ///
43    /// ```
44    /// # use magnesia::linalg::SMatrix;
45    /// # use magnesia::algebra::One;
46    /// let m = SMatrix::one();
47    /// assert_eq!(m, SMatrix::from([[1,0],[0,1],[0,0]]));
48    /// ```
49    fn one() -> Self {
50        let mut m = 0_usize;
51        Self([(); NUM_ROWS].map(|_| {
52            let mut n = 0_usize;
53            let row = [(); NUM_COLS].map(|_| {
54                let x = if m == n { T::one() } else { T::zero() };
55                n += 1;
56                x
57            });
58            m += 1;
59            row
60        }))
61    }
62}
63
64impl<T, const NUM_ROWS: usize, const NUM_COLS: usize> From<[[T; NUM_COLS]; NUM_ROWS]>
65    for SMatrix<T, NUM_ROWS, NUM_COLS>
66{
67    /// Creates a statically sized matrix from an array.
68    ///
69    /// # Example
70    /// ```
71    /// # use magnesia::linalg::SMatrix;
72    /// let a = SMatrix::from([[1,2],[3,4]]);
73    /// ```
74    fn from(coefficients: [[T; NUM_COLS]; NUM_ROWS]) -> Self {
75        Self(coefficients)
76    }
77}
78
79impl<T, const NUM_ROWS: usize, const NUM_COLS: usize> Add for SMatrix<T, NUM_ROWS, NUM_COLS>
80where
81    for<'a> T: AddAssign<&'a T>,
82{
83    type Output = Self;
84
85    /// Adds two matrices.
86    ///
87    /// # Example
88    /// ```
89    /// # use magnesia::linalg::SMatrix;
90    /// let a = SMatrix::from([[0,1,2],[3,4,5]]);
91    /// let b = SMatrix::from([[1,1,1],[1,1,1]]);
92    /// let c = a + b;
93    /// assert_eq!(c, SMatrix::from([[1,2,3],[4,5,6]]));
94    /// ```
95    fn add(mut self, other: Self) -> Self {
96        self += &other;
97        self
98    }
99}
100
101impl<T, const NUM_ROWS: usize, const NUM_COLS: usize> Add<&Self> for SMatrix<T, NUM_ROWS, NUM_COLS>
102where
103    for<'a> T: AddAssign<&'a T>,
104{
105    type Output = Self;
106
107    /// Adds two matrices.
108    ///
109    /// # Example
110    /// ```
111    /// # use magnesia::linalg::SMatrix;
112    /// let a = SMatrix::from([[0,1,2],[3,4,5]]);
113    /// let b = SMatrix::from([[1,1,1],[1,1,1]]);
114    /// let c = a + &b;
115    /// assert_eq!(c, SMatrix::from([[1,2,3],[4,5,6]]));
116    /// ```
117    fn add(mut self, other: &Self) -> Self {
118        self += other;
119        self
120    }
121}
122
123impl<T, const NUM_ROWS: usize, const NUM_COLS: usize> AddAssign<Self>
124    for SMatrix<T, NUM_ROWS, NUM_COLS>
125where
126    for<'a> T: AddAssign<&'a T>,
127{
128    /// Adds two matrices in-place.
129    ///
130    /// # Example
131    /// ```
132    /// # use magnesia::linalg::SMatrix;
133    /// let mut a = SMatrix::from([[0,1,2],[3,4,5]]);
134    /// let b = SMatrix::from([[1,1,1],[1,1,1]]);
135    /// a += b;
136    /// assert_eq!(a, SMatrix::from([[1,2,3],[4,5,6]]));
137    /// ```
138    fn add_assign(&mut self, other: Self) {
139        *self += &other;
140    }
141}
142
143impl<T, const NUM_ROWS: usize, const NUM_COLS: usize> AddAssign<&Self>
144    for SMatrix<T, NUM_ROWS, NUM_COLS>
145where
146    for<'a> T: AddAssign<&'a T>,
147{
148    /// Adds two matrices in-place.
149    ///
150    /// # Example
151    /// ```
152    /// # use magnesia::linalg::SMatrix;
153    /// let mut a = SMatrix::from([[0,1,2],[3,4,5]]);
154    /// let b = SMatrix::from([[1,1,1],[1,1,1]]);
155    /// a += &b;
156    /// assert_eq!(a, SMatrix::from([[1,2,3],[4,5,6]]));
157    /// ```
158    fn add_assign(&mut self, other: &Self) {
159        for (lo, ro) in self.0.iter_mut().zip(other.0.iter()) {
160            for (li, ri) in lo.iter_mut().zip(ro.iter()) {
161                *li += ri;
162            }
163        }
164    }
165}
166
167impl<T, const NUM_ROWS: usize, const NUM_COLS: usize> Sub for SMatrix<T, NUM_ROWS, NUM_COLS>
168where
169    for<'a> T: SubAssign<&'a T>,
170{
171    type Output = Self;
172
173    /// Subtracts two matrices.
174    ///
175    /// # Example
176    /// ```
177    /// # use magnesia::linalg::SMatrix;
178    /// let a = SMatrix::from([[1,2,3],[4,5,6]]);
179    /// let b = SMatrix::from([[1,1,1],[1,1,1]]);
180    /// let c = a - b;
181    /// assert_eq!(c, SMatrix::from([[0,1,2],[3,4,5]]));
182    /// ```
183    fn sub(mut self, other: Self) -> Self {
184        self -= &other;
185        self
186    }
187}
188
189impl<T, const NUM_ROWS: usize, const NUM_COLS: usize> Sub<&Self> for SMatrix<T, NUM_ROWS, NUM_COLS>
190where
191    for<'a> T: SubAssign<&'a T>,
192{
193    type Output = Self;
194
195    /// Subtracts two matrices.
196    ///
197    /// # Example
198    /// ```
199    /// # use magnesia::linalg::SMatrix;
200    /// let a = SMatrix::from([[1,2,3],[4,5,6]]);
201    /// let b = SMatrix::from([[1,1,1],[1,1,1]]);
202    /// let c = a - &b;
203    /// assert_eq!(c, SMatrix::from([[0,1,2],[3,4,5]]));
204    /// ```
205    fn sub(mut self, other: &Self) -> Self {
206        self -= other;
207        self
208    }
209}
210
211impl<T, const NUM_ROWS: usize, const NUM_COLS: usize> SubAssign<Self>
212    for SMatrix<T, NUM_ROWS, NUM_COLS>
213where
214    for<'a> T: SubAssign<&'a T>,
215{
216    /// Subtracts two matrices in-place.
217    ///
218    /// # Example
219    /// ```
220    /// # use magnesia::linalg::SMatrix;
221    /// let mut a = SMatrix::from([[1,2,3],[4,5,6]]);
222    /// let b = SMatrix::from([[1,1,1],[1,1,1]]);
223    /// a -= b;
224    /// assert_eq!(a, SMatrix::from([[0,1,2],[3,4,5]]));
225    /// ```
226    fn sub_assign(&mut self, other: Self) {
227        *self -= &other;
228    }
229}
230
231impl<T, const NUM_ROWS: usize, const NUM_COLS: usize> SubAssign<&Self>
232    for SMatrix<T, NUM_ROWS, NUM_COLS>
233where
234    for<'a> T: SubAssign<&'a T>,
235{
236    /// Subtracts two matrices in-place.
237    ///
238    /// # Example
239    /// ```
240    /// # use magnesia::linalg::SMatrix;
241    /// let mut a = SMatrix::from([[1,2,3],[4,5,6]]);
242    /// let b = SMatrix::from([[1,1,1],[1,1,1]]);
243    /// a -= &b;
244    /// assert_eq!(a, SMatrix::from([[0,1,2],[3,4,5]]));
245    /// ```
246    fn sub_assign(&mut self, other: &Self) {
247        for (lo, ro) in self.0.iter_mut().zip(other.0.iter()) {
248            for (li, ri) in lo.iter_mut().zip(ro.iter()) {
249                *li -= ri;
250            }
251        }
252    }
253}
254
255impl<'a, T, const L: usize, const M: usize, const N: usize> Mul<&'a SMatrix<T, M, N>>
256    for &'a SMatrix<T, L, M>
257where
258    for<'b> T: AddAssign<&'b T> + MulRefs + Zero,
259{
260    type Output = SMatrix<T, L, N>;
261
262    /// Multiplies two matrices
263    ///
264    /// # Example
265    /// ```
266    /// # use magnesia::linalg::SMatrix;
267    /// let a = SMatrix::from([[0,1,2],[3,4,5]]);
268    /// let b = SMatrix::from([[0,1],[2,3],[4,5]]);
269    /// let c = &a * &b;
270    /// assert_eq!(c, SMatrix::from([[10,13],[28,40]]));
271    /// ```
272    fn mul(self, other: &'a SMatrix<T, M, N>) -> Self::Output {
273        let mut l = 0_usize;
274        SMatrix([(); L].map(|_| {
275            let mut n = 0_usize;
276            let row = [(); N].map(|_| {
277                let mut sum = T::zero();
278                for m in 0..M {
279                    let prod = self.0[l][m].mul_refs(&other.0[m][n]);
280                    sum += &prod;
281                }
282                n += 1;
283                sum
284            });
285            l += 1;
286            row
287        }))
288    }
289}
290
291impl<T, const L: usize, const M: usize, const N: usize> Mul<SMatrix<T, M, N>> for SMatrix<T, L, M>
292where
293    for<'b> T: AddAssign<&'b T> + MulRefs + Zero,
294{
295    type Output = SMatrix<T, L, N>;
296
297    /// Multiplies two matrices
298    ///
299    /// # Example
300    /// ```
301    /// # use magnesia::linalg::SMatrix;
302    /// let a = SMatrix::from([[0,1,2],[3,4,5]]);
303    /// let b = SMatrix::from([[0,1],[2,3],[4,5]]);
304    /// let c = a * b;
305    /// assert_eq!(c, SMatrix::from([[10,13],[28,40]]));
306    /// ```
307    fn mul(self, other: SMatrix<T, M, N>) -> Self::Output {
308        &self * &other
309    }
310}
311
312impl<T, const M: usize, const N: usize> MulAssign<&SMatrix<T, N, N>> for SMatrix<T, M, N>
313where
314    for<'a> &'a Self: Mul<&'a SMatrix<T, N, N>, Output = Self>,
315{
316    /// Multiplies two matrices
317    ///
318    /// # Example
319    /// ```
320    /// # use magnesia::linalg::SMatrix;
321    /// let mut a = SMatrix::from([[0,1],[2,3],[4,5]]);
322    /// let b = SMatrix::from([[4,5],[6,7]]);
323    /// a *= &b;
324    /// assert_eq!(a, SMatrix::from([[6,7],[26,31],[46,55]]));
325    /// ```
326    fn mul_assign(&mut self, other: &SMatrix<T, N, N>) {
327        let r: &Self = self;
328        let result = r * other;
329        *self = result;
330    }
331}
332
333impl<T, const M: usize, const N: usize> MulAssign<SMatrix<T, N, N>> for SMatrix<T, M, N>
334where
335    for<'a> &'a Self: Mul<&'a SMatrix<T, N, N>, Output = Self>,
336{
337    /// Multiplies two matrices and assigns the result to the first operand.
338    ///
339    /// # Example
340    /// ```
341    /// # use magnesia::linalg::SMatrix;
342    /// let mut a = SMatrix::from([[0,1],[2,3],[4,5]]);
343    /// let b = SMatrix::from([[4,5],[6,7]]);
344    /// a *= b;
345    /// assert_eq!(a, SMatrix::from([[6,7],[26,31],[46,55]]));
346    /// ```
347    fn mul_assign(&mut self, other: SMatrix<T, N, N>) {
348        let r: &Self = self;
349        let result = r * &other;
350        *self = result;
351    }
352}
353
354impl<'a, T, const M: usize, const N: usize> Mul<&'a SVector<T, N>> for &'a SMatrix<T, M, N>
355where
356    for<'b> T: MulRefs + AddAssign<&'b T> + Zero,
357{
358    type Output = SVector<T, M>;
359
360    /// Implements multiplication of matrix by vector.
361    ///
362    /// # Example
363    /// ```
364    /// # use magnesia::linalg::SMatrix;
365    /// # use magnesia::linalg::SVector;
366    /// let m = SMatrix::from([[1,2], [3,4]]);
367    /// let u = SVector::from([1,2]);
368    /// let v = &m * &u;
369    /// assert_eq!(v, SVector::from([1*1 + 2*2, 1*3 + 2*4]));
370    /// ```
371    fn mul(self, vec: &'a SVector<T, N>) -> Self::Output {
372        let mut i = 0_usize;
373        SVector::from([(); M].map(|_| {
374            let mut sum = T::zero();
375            for (m, v) in self.0[i].iter().zip(vec) {
376                let prod = m.mul_refs(v);
377                sum += &prod;
378            }
379            i += 1;
380            sum
381        }))
382    }
383}
384
385impl<T: Neg<Output = T>, const M: usize, const N: usize> Neg for SMatrix<T, M, N> {
386    type Output = Self;
387    /// Negates a matrix in-place.
388    ///
389    /// # Example
390    /// ```
391    /// # use magnesia::linalg::SMatrix;
392    /// let mut m = SMatrix::from([[1,2,3],[4,5,6]]);
393    /// assert_eq!(-m, SMatrix::from([[-1,-2,-3],[-4,-5,-6]]));
394    fn neg(self) -> Self {
395        Self(self.0.map(|row| row.map(|x| -x)))
396    }
397}
398
399impl<T: Ring, const N: usize> Ring for SMatrix<T, N, N> {}
400
401impl<T: Conj, const M: usize, const N: usize> Conj for SMatrix<T, M, N> {
402    fn conj(self) -> Self {
403        Self(self.0.map(|row| row.map(|x| x.conj())))
404    }
405}
406
407#[test]
408fn test_conj_assign_on_i8_smatrix() {
409    let a = SMatrix::from([[1i8, 2i8], [3i8, 4i8]]);
410    let b = a.conj();
411    assert_eq!(a, b);
412}
413
414#[test]
415fn test_conj_assign_on_complex_i8_smatrix() {
416    use crate::algebra::Complex;
417    let a = SMatrix::from([
418        [Complex::new(1i8, 1i8), Complex::new(2i8, 2i8)],
419        [Complex::new(3i8, 3i8), Complex::new(4i8, 4i8)],
420    ]);
421    let b = SMatrix::from([
422        [Complex::new(1i8, -1i8), Complex::new(2i8, -2i8)],
423        [Complex::new(3i8, -3i8), Complex::new(4i8, -4i8)],
424    ]);
425    assert_eq!(a.conj(), b);
426}