ear_algae/
matrix.rs

1use std::{array, ops::*};
2
3#[cfg(feature = "bytemuck")]
4use bytemuck::*;
5
6use crate::{
7    Vect,
8    op_wrapper::Sc,
9    ops::{Det, Dot},
10    traits::Field,
11};
12
13#[repr(transparent)]
14#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
15#[cfg_attr(feature = "bytemuck", derive(TransparentWrapper))]
16pub struct Mat<const N: usize, const M: usize, S: Field>(pub [[S; M]; N]);
17
18#[cfg(feature = "bytemuck")]
19unsafe impl<const N: usize, const M: usize, S: Field> Zeroable for Mat<N, M, S> where
20    [[S; M]; N]: Zeroable
21{
22}
23
24#[cfg(feature = "bytemuck")]
25unsafe impl<const N: usize, const M: usize, S: Field> Pod for Mat<N, M, S> where [[S; M]; N]: Pod {}
26
27impl<S: Field, const N: usize, const M: usize> Mat<N, M, S> {
28    pub fn from_scs(array: [[Sc<S>; M]; N]) -> Self {
29        Mat(array.map(|x| x.map(|x| x.0)))
30    }
31
32    pub fn map<T: Field>(self, f: impl Fn(S) -> T) -> Mat<N, M, T> {
33        Mat(self.0.map(|x| x.map(&f)))
34    }
35}
36
37impl<const N: usize, const M: usize> Mat<N, M, f32> {
38    pub fn to_f64(self) -> Mat<N, M, f64> {
39        self.map(|x| x as _)
40    }
41}
42
43impl<const N: usize, const M: usize> Mat<N, M, f64> {
44    pub fn to_f32(self) -> Mat<N, M, f32> {
45        self.map(|x| x as _)
46    }
47}
48
49impl<S: Field, const N: usize, const M: usize> Mat<N, M, S> {
50    pub fn try_index(&self, i: usize, j: usize) -> Option<S> {
51        if i > N || j > M {
52            None
53        } else {
54            Some(self[i][j])
55        }
56    }
57}
58
59impl<S: Field, const N: usize, const M: usize> Mat<N, M, S> {
60    fn row_swap_assign(&mut self, i1: usize, i2: usize) {
61        (self.0[i1], self.0[i2]) = (self.0[i2], self.0[i1])
62    }
63
64    fn row_sub_assign(&mut self, i1: usize, coefficient: S, i2: usize) {
65        for j in 0..M {
66            self.0[i1][j].sub_assign(self[i2][j].mul(coefficient));
67        }
68    }
69
70    fn row_div_assign(&mut self, i: usize, divisor: S) {
71        for j in 0..M {
72            self.0[i][j].div_assign(divisor);
73        }
74    }
75}
76
77struct AugMat<const N: usize, const M1: usize, const M2: usize, S: Field>(
78    pub Mat<N, M1, S>,
79    pub Mat<N, M2, S>,
80);
81
82impl<S: Field, const N: usize, const M1: usize, const M2: usize> AugMat<N, M1, M2, S> {
83    fn row_swap_assign(&mut self, i1: usize, i2: usize) {
84        self.0.row_swap_assign(i1, i2);
85        self.1.row_swap_assign(i1, i2);
86    }
87
88    fn row_sub_assign(&mut self, i1: usize, coefficient: S, i2: usize) {
89        self.0.row_sub_assign(i1, coefficient, i2);
90        self.1.row_sub_assign(i1, coefficient, i2);
91    }
92
93    fn row_div_assign(&mut self, i: usize, divisor: S) {
94        self.0.row_div_assign(i, divisor);
95        self.1.row_div_assign(i, divisor);
96    }
97}
98
99impl<S: Field, const N: usize> Mat<N, N, S> {
100    pub fn inverse(&self) -> Self {
101        let mut aug = AugMat(*self, Self::IDENT);
102        for j in 0..N {
103            {
104                let mut max = self[j][j].abs();
105                let mut i_max: usize = j;
106                for i in j + 1..N {
107                    let value = self[i][j].abs();
108                    if value > max {
109                        max = value;
110                        i_max = i;
111                    }
112                }
113                if j != i_max {
114                    aug.row_swap_assign(j, i_max);
115                }
116            }
117            aug.row_div_assign(j, aug.0[j][j]);
118            for i in j + 1..N {
119                let c = aug.0[i][j];
120                aug.row_sub_assign(i, c, j);
121            }
122        }
123        for i in (0..N - 1).rev() {
124            for j in i + 1..N {
125                aug.row_sub_assign(i, aug.0[i][j], j);
126            }
127        }
128        aug.1
129    }
130}
131
132impl<S: Field> Det for Mat<1, 1, S> {
133    type Output = S;
134
135    fn det(self) -> S {
136        self[0][0]
137    }
138}
139
140impl<S: Field> Det for Mat<2, 2, S> {
141    type Output = S;
142
143    fn det(self) -> S {
144        self[0][0].mul(self[1][1]).sub(self[0][1].mul(self[1][0]))
145    }
146}
147
148impl<S: Field> Det for Mat<3, 3, S> {
149    type Output = S;
150    fn det(self) -> S {
151        let a = Sc(self[0][0]);
152        let b = Sc(self[0][1]);
153        let c = Sc(self[0][2]);
154        let d = Sc(self[1][0]);
155        let e = Sc(self[1][1]);
156        let f = Sc(self[1][2]);
157        let g = Sc(self[2][0]);
158        let h = Sc(self[2][1]);
159        let i = Sc(self[2][2]);
160
161        //      a.mul(e.mul(i))  .add (b.mul(f.mul(g))) .add (c.mul(d.mul(h)))
162        //.sub (c.mul(e.mul(g))) .sub (b.mul(d.mul(i))) .sub (a.mul(f.mul(h)))
163
164        ((a * e * i + b * f * g + c * d * h) - (c * e * g + b * d * i + a * f * h)).0
165    }
166}
167
168impl<S: Field, const N: usize, const M: usize> Default for Mat<N, M, S> {
169    fn default() -> Self {
170        Self::ZERO
171    }
172}
173
174impl<S: Field, const N: usize, const M: usize> Mat<N, M, S> {
175    pub fn from_fn<F: Fn(usize, usize) -> S>(f: F) -> Self {
176        Mat(array::from_fn(|i| array::from_fn(|j| f(i, j))))
177    }
178
179    pub const ZERO: Self = Mat([[S::ZERO; M]; N]);
180}
181
182impl<S: Field, const N: usize> Mat<N, N, S> {
183    pub const IDENT: Self = {
184        let mut array = [[S::ZERO; N]; N];
185        let mut i = 0;
186        while i < N {
187            array[i][i] = S::ONE;
188            i += 1;
189        }
190        Self(array)
191    };
192}
193
194impl<S: Field, const N: usize, const M: usize> Index<usize> for Mat<N, M, S> {
195    type Output = [S; M];
196
197    fn index(&self, i: usize) -> &Self::Output {
198        &self.0[i]
199    }
200}
201
202impl<S: Field, const N: usize, const M: usize> Mat<N, M, S> {
203    pub fn row(&self, i: usize) -> Vect<M, S> {
204        Vect::from_fn(|j| self[i][j])
205    }
206
207    pub fn col(&self, j: usize) -> Vect<N, S> {
208        Vect::from_fn(|i| self[i][j])
209    }
210}
211
212impl<S: Field, const N: usize, const M: usize> Mul<Vect<M, S>> for Mat<N, M, S> {
213    type Output = Vect<N, S>;
214
215    fn mul(self, vector: Vect<M, S>) -> Self::Output {
216        Vect::from_fn(|i| self.row(i).dot(vector))
217    }
218}
219
220impl<S: Field, const N: usize, const M: usize> Mul<Mat<N, M, S>> for Vect<N, S> {
221    type Output = Vect<N, S>;
222
223    fn mul(self, matrix: Mat<N, M, S>) -> Self::Output {
224        Vect::from_fn(|j| self.dot(matrix.col(j)))
225    }
226}
227
228impl<S: Field, const N: usize, const M: usize, const P: usize> Mul<Mat<M, P, S>> for Mat<N, M, S> {
229    type Output = Mat<N, P, S>;
230
231    fn mul(self, other: Mat<M, P, S>) -> Self::Output {
232        Mat::from_fn(|i, j| self.row(i).dot(other.col(j)))
233    }
234}
235
236impl<S: Field> Mat<4, 4, S> {
237    pub fn flatten(self) -> [S; 16] {
238        array::from_fn(|i| self[i % 4][i / 4])
239    }
240}
241
242impl<S: Field, const N: usize, const M: usize> Mul<S> for Mat<N, M, S> {
243    type Output = Self;
244
245    fn mul(self, scalar: S) -> Self::Output {
246        Mat::from_fn(|i, j| self[i][j].mul(scalar))
247    }
248}
249
250impl<S: Field, const N: usize, const M: usize> Add for Mat<N, M, S> {
251    type Output = Self;
252
253    fn add(self, other: Self) -> Self::Output {
254        Mat::from_fn(|i, j| self[i][j].add(other[i][j]))
255    }
256}
257
258impl<S: Field, const N: usize, const M: usize> Sub for Mat<N, M, S> {
259    type Output = Self;
260    fn sub(self, other: Self) -> Self::Output {
261        Mat::from_fn(|i, j| self[i][j].sub(other[i][j]))
262    }
263}
264
265impl<S: Field, const N: usize, const M: usize> Neg for Mat<N, M, S> {
266    type Output = Self;
267    fn neg(self) -> Self::Output {
268        Mat::from_fn(|i, j| self[i][j].neg())
269    }
270}
ear_algae/matrix.rs

ear_algae/
matrix.rs
