conspire/math/matrix/square/lu/
mod.rs

1#[cfg(test)]
2mod test;
3
4use super::{SquareMatrix, SquareMatrixError};
5use crate::{
6    ABS_TOL, REL_TOL,
7    math::{Scalar, Tensor, Vector, simd},
8};
9
10impl SquareMatrix {
11    /// Factorize the matrix using the LU decomposition.
12    pub fn factorize_lu(&self) -> Result<LuDecomposition, SquareMatrixError> {
13        let mut decomposition = LuDecomposition {
14            lu: self.clone(),
15            permutation: (0..self.len()).collect(),
16        };
17        decomposition.factorize()?;
18        Ok(decomposition)
19    }
20    /// Factorize the matrix into an existing LU decomposition of the same size.
21    pub fn factorize_lu_into(
22        &self,
23        decomposition: &mut LuDecomposition,
24    ) -> Result<(), SquareMatrixError> {
25        let LuDecomposition { lu, permutation } = decomposition;
26        lu.iter_mut().zip(self.iter()).for_each(|(lu_i, self_i)| {
27            lu_i.iter_mut()
28                .zip(self_i.iter())
29                .for_each(|(a, b)| *a = *b)
30        });
31        permutation
32            .iter_mut()
33            .enumerate()
34            .for_each(|(i, permutation_i)| *permutation_i = i);
35        decomposition.factorize()
36    }
37    /// Solve a system of linear equations using the LU decomposition.
38    pub fn solve_lu(&self, b: &Vector) -> Result<Vector, SquareMatrixError> {
39        Ok(self.factorize_lu()?.solve(b))
40    }
41}
42
43/// The LU decomposition of a square matrix.
44pub struct LuDecomposition {
45    lu: SquareMatrix,
46    permutation: Vec<usize>,
47}
48
49impl LuDecomposition {
50    fn factorize(&mut self) -> Result<(), SquareMatrixError> {
51        let Self { lu, permutation } = self;
52        let n = lu.len();
53        let mut largest = 0.0;
54        for i in 0..n {
55            let mut max_row = i;
56            let mut max_val = lu[i][i].abs();
57            for k in i + 1..n {
58                let candidate = lu[k][i].abs();
59                if candidate > max_val {
60                    max_row = k;
61                    max_val = candidate;
62                }
63            }
64            if max_row != i {
65                lu.0.swap(i, max_row);
66                permutation.swap(i, max_row);
67            }
68            largest = max_val.max(largest);
69            if max_val < ABS_TOL && max_val <= REL_TOL * largest {
70                return Err(SquareMatrixError::Singular);
71            }
72            let pivot = lu[i][i];
73            let (front, back) = lu.0.split_at_mut(i + 1);
74            let column = &front[i].as_slice()[i + 1..];
75            let mut count = 0;
76            for row in 0..back.len() {
77                let factor = back[row][i];
78                if factor != 0.0 {
79                    back[row][i] = factor / pivot;
80                    if row != count {
81                        back.swap(row, count);
82                        permutation.swap(i + 1 + row, i + 1 + count)
83                    }
84                    count += 1
85                }
86            }
87            if column.len() < 4 {
88                back[..count].iter_mut().for_each(|row| {
89                    let factor = row[i];
90                    row.as_mut_slice()[i + 1..]
91                        .iter_mut()
92                        .zip(column.iter())
93                        .for_each(|(row_k, column_k)| *row_k -= factor * column_k)
94                })
95            } else {
96                back[..count].chunks_mut(4).for_each(|chunk| {
97                    if let [a, b, c, d] = chunk {
98                        let u = [-a[i], -b[i], -c[i], -d[i]];
99                        simd::rank_one_quad(
100                            &mut a.as_mut_slice()[i + 1..],
101                            &mut b.as_mut_slice()[i + 1..],
102                            &mut c.as_mut_slice()[i + 1..],
103                            &mut d.as_mut_slice()[i + 1..],
104                            column,
105                            u,
106                        )
107                    } else {
108                        chunk.iter_mut().for_each(|row| {
109                            let factor = row[i];
110                            simd::axpy(&mut row.as_mut_slice()[i + 1..], column, factor)
111                        })
112                    }
113                })
114            }
115        }
116        Ok(())
117    }
118    /// An unfactorized decomposition sized to hold that of a matrix of the given length.
119    pub fn zero(len: usize) -> Self {
120        Self {
121            lu: SquareMatrix::zero(len),
122            permutation: (0..len).collect(),
123        }
124    }
125    /// Solve a system of linear equations for another right-hand side.
126    pub fn solve(&self, b: &Vector) -> Vector {
127        let mut x = Vector::zero(self.permutation.len());
128        self.solve_into(b, &mut x);
129        x
130    }
131    /// Solve a system of linear equations into an existing vector.
132    pub fn solve_into(&self, b: &Vector, x: &mut Vector) {
133        self.permutation
134            .iter()
135            .zip(x.iter_mut())
136            .for_each(|(&p_i, x_i)| *x_i = b[p_i]);
137        forward_substitution(x, &self.lu);
138        backward_substitution(x, &self.lu)
139    }
140}
141
142fn forward_substitution(x: &mut Vector, a: &SquareMatrix) {
143    a.iter().enumerate().for_each(|(i, a_i)| {
144        x[i] -= a_i
145            .iter()
146            .take(i)
147            .zip(x.iter().take(i))
148            .map(|(a_ij, x_j)| a_ij * x_j)
149            .sum::<Scalar>()
150    })
151}
152
153fn backward_substitution(x: &mut Vector, a: &SquareMatrix) {
154    a.0.iter().enumerate().rev().for_each(|(i, a_i)| {
155        x[i] -= a_i
156            .iter()
157            .skip(i + 1)
158            .zip(x.iter().skip(i + 1))
159            .map(|(a_ij, x_j)| a_ij * x_j)
160            .sum::<Scalar>();
161        x[i] /= a_i[i];
162    })
163}
conspire/math/matrix/square/lu/mod.rs

conspire/math/matrix/square/lu/
mod.rs
