Wed, 07 Dec 2022 07:00:27 +0200
Added tag v0.1.0 for changeset 51bfde513cfa
5 | 1 | /*! |
2 | Linear equation solvers for small problems stored in Rust arrays. | |
3 | */ | |
0 | 4 | |
5 | use crate::types::Float; | |
6 | use std::mem::MaybeUninit; | |
7 | ||
8 | /// Gaussian elimination for $AX=B$, where $A$ and $B$ are both stored in `ab`, | |
9 | /// $A \in \mathbb{R}^{M \times M}$ and $X, B \in \mathbb{R}^{M \times K}$. | |
10 | pub fn linsolve0<F : Float, const M : usize, const N : usize, const K : usize>( | |
11 | mut ab : [[F; N]; M] | |
12 | ) -> [[F; K]; M] { | |
13 | assert_eq!(M + K , N); | |
14 | ||
15 | let mut k = 0; | |
16 | ||
17 | // Convert to row-echelon form | |
18 | for h in 0..(M-1) { | |
19 | // Find pivotable column (has some non-zero entries in rows ≥ h) | |
20 | 'find_pivot: while k < N { | |
21 | let (mut î, mut v) = (h, ab[h][k].abs()); | |
22 | // Find row ≥ h of maximum absolute value in this column | |
23 | for i in (h+1)..M { | |
24 | let ṽ = ab[i][k].abs(); | |
25 | if ṽ > v { | |
26 | î = i; | |
27 | v = ṽ; | |
28 | } | |
29 | } | |
30 | if v > F::ZERO { | |
31 | ab.swap(h, î); | |
32 | for i in (h+1)..M { | |
33 | let f = ab[i][k] / ab[h][k]; | |
34 | ab[i][k] = F::ZERO; | |
35 | for j in (k+1)..N { | |
36 | ab[i][j] -= ab[h][j]*f; | |
37 | } | |
38 | } | |
39 | k += 1; | |
40 | break 'find_pivot; | |
41 | } | |
42 | k += 1 | |
43 | } | |
44 | } | |
45 | ||
46 | // Solve UAX=UB for X where UA with U presenting the transformations above an | |
47 | // upper triangular matrix. | |
48 | // This use of MaybeUninit assumes F : Copy. Otherwise undefined behaviour may occur. | |
49 | let mut x : [[MaybeUninit<F>; K]; M] = core::array::from_fn(|_| MaybeUninit::uninit_array::<K>() ); | |
50 | //unsafe { std::mem::MaybeUninit::uninit().assume_init() }; | |
51 | for i in (0..M).rev() { | |
52 | for 𝓁 in 0..K { | |
53 | let mut tmp = ab[i][M+𝓁]; | |
54 | for j in (i+1)..M { | |
55 | tmp -= ab[i][j] * unsafe { *(x[j][𝓁].assume_init_ref()) }; | |
56 | } | |
57 | tmp /= ab[i][i]; | |
58 | x[i][𝓁].write(tmp); | |
59 | } | |
60 | } | |
61 | //unsafe { MaybeUninit::array_assume_init(x) }; | |
62 | let xinit = unsafe { | |
63 | //core::intrinsics::assert_inhabited::<[[F; K]; M]>(); | |
64 | (&x as *const _ as *const [[F; K]; M]).read() | |
65 | }; | |
66 | ||
67 | std::mem::forget(x); | |
68 | xinit | |
69 | } | |
70 | ||
71 | /// Gaussian elimination for $Ax=b$, where $A$ and $b$ are both stored in `ab`, | |
72 | /// $A \in \mathbb{R}^{M \times M}$ and $x, b \in \mathbb{R}^M$. | |
73 | #[inline] | |
74 | pub fn linsolve<F : Float, const M : usize, const N : usize>(ab : [[F; N]; M]) -> [F; M] { | |
75 | let x : [[F; 1]; M] = linsolve0(ab); | |
76 | unsafe { *((&x as *const [F; 1]) as *const [F; M] ) } | |
77 | } | |
78 | ||
79 | ||
80 | #[cfg(test)] | |
81 | mod tests { | |
82 | use super::*; | |
83 | ||
84 | #[test] | |
85 | fn linsolve_test() { | |
86 | let ab1 = [[1.0, 2.0, 3.0], [2.0, 1.0, 6.0]]; | |
87 | assert_eq!(linsolve(ab1), [3.0, 0.0]); | |
88 | let ab2 = [[1.0, 2.0, 0.0, 1.0], [4.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0]]; | |
89 | assert_eq!(linsolve(ab2), [0.0, 0.5, 0.0]); | |
90 | } | |
91 | } | |
92 |