src/linsolve.rs

Sat, 22 Oct 2022 22:28:04 +0300

author
Tuomo Valkonen <tuomov@iki.fi>
date
Sat, 22 Oct 2022 22:28:04 +0300
changeset 2
ac84e995e119
parent 0
9f27689eb130
child 5
59dc4c5883f4
permissions
-rw-r--r--

Convert iteration utilities to GATs

0
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
1 /// Linear solvers for small problems.
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
2
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
3 use crate::types::Float;
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
4 use std::mem::MaybeUninit;
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
5
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
6 /// 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
7 /// $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
8 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
9 mut ab : [[F; N]; M]
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
10 ) -> [[F; K]; M] {
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
11 assert_eq!(M + K , N);
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
12
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
13 let mut k = 0;
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 // Convert to row-echelon form
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
16 for h in 0..(M-1) {
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
17 // Find pivotable column (has some non-zero entries in rows ≥ h)
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
18 'find_pivot: while k < N {
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
19 let (mut î, mut v) = (h, ab[h][k].abs());
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
20 // Find row ≥ h of maximum absolute value in this column
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
21 for i in (h+1)..M {
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
22 let ṽ = ab[i][k].abs();
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
23 if ṽ > v {
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
24 î = i;
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
25 v = ṽ;
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
26 }
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
27 }
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
28 if v > F::ZERO {
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
29 ab.swap(h, î);
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
30 for i in (h+1)..M {
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
31 let f = ab[i][k] / ab[h][k];
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
32 ab[i][k] = F::ZERO;
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
33 for j in (k+1)..N {
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
34 ab[i][j] -= ab[h][j]*f;
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
35 }
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
36 }
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
37 k += 1;
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
38 break 'find_pivot;
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
39 }
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
40 k += 1
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 }
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 // 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
45 // upper triangular matrix.
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
46 // 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
47 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
48 //unsafe { std::mem::MaybeUninit::uninit().assume_init() };
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
49 for i in (0..M).rev() {
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
50 for 𝓁 in 0..K {
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
51 let mut tmp = ab[i][M+𝓁];
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
52 for j in (i+1)..M {
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
53 tmp -= ab[i][j] * unsafe { *(x[j][𝓁].assume_init_ref()) };
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
54 }
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
55 tmp /= ab[i][i];
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
56 x[i][𝓁].write(tmp);
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
57 }
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
58 }
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
59 //unsafe { MaybeUninit::array_assume_init(x) };
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
60 let xinit = unsafe {
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
61 //core::intrinsics::assert_inhabited::<[[F; K]; M]>();
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
62 (&x as *const _ as *const [[F; K]; M]).read()
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
63 };
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
64
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
65 std::mem::forget(x);
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
66 xinit
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
67 }
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
68
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
69 /// 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
70 /// $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
71 #[inline]
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
72 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
73 let x : [[F; 1]; M] = linsolve0(ab);
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
74 unsafe { *((&x as *const [F; 1]) as *const [F; M] ) }
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
75 }
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
76
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 #[cfg(test)]
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
79 mod tests {
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
80 use super::*;
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
81
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
82 #[test]
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
83 fn linsolve_test() {
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
84 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
85 assert_eq!(linsolve(ab1), [3.0, 0.0]);
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
86 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
87 assert_eq!(linsolve(ab2), [0.0, 0.5, 0.0]);
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
88 }
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
89 }
9f27689eb130 Initialise new clean repository
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
90

mercurial