src/linsolve.rs

Wed, 07 Dec 2022 07:00:27 +0200

author
Tuomo Valkonen <tuomov@iki.fi>
date
Wed, 07 Dec 2022 07:00:27 +0200
changeset 18
2b75e98df693
parent 5
59dc4c5883f4
child 24
8efa7abce7c7
permissions
-rw-r--r--

Added tag v0.1.0 for changeset 51bfde513cfa

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

mercurial