Skip to main content

rusolver/
lu.rs

1// SPDX-License-Identifier: Apache-2.0
2use crate::{
3    Matrix, MatrixView, SolverError, Tolerance
4};
5use crate::numerics::{
6    checked, sum_iter
7};
8/// Reusable square LU with partial ROW pivoting. Convention: `P * A = L * U`.
9/// `pivots[k]` is the row interchanged with row k during elimination. The factors
10/// own a copy; the caller's input is never modified. Not a blocked BLAS kernel.
11#[derive(Clone, Debug)]
12pub struct Lu {
13    packed: Matrix, pivots: Vec<usize>, sign: i32
14}
15impl Lu {
16    pub fn factor(a: MatrixView<'_>, tolerance: Tolerance) -> Result<Self, SolverError> {
17        let n=a.square()?;
18        let cutoff=tolerance.threshold(a.max_abs())?;
19        let mut packed=a.to_owned()?;
20        let mut pivots=Vec::new();
21        pivots.try_reserve_exact(n).map_err(|_| SolverError::Allocation)?;
22        let mut sign=1;
23        for k in 0..n {
24            let mut pivot=k;
25            for i in k+1..n {
26                if packed.data[i*n+k].abs()>packed.data[pivot*n+k].abs() {
27                    pivot=i;
28                }
29            }
30            let value=packed.data[pivot*n+k];
31            if value.abs()<=cutoff {
32                return Err(SolverError::Singular {
33                    index: k, pivot: value, threshold: cutoff
34                });
35            }
36            pivots.push(pivot);
37            if pivot!=k {
38                for j in 0..n {
39                    packed.data.swap(k*n+j, pivot*n+j);
40                }
41                sign=-sign;
42            }
43            for i in k+1..n {
44                let ratio=checked(packed.data[i*n+k]/packed.data[k*n+k], "LU multiplier")?;
45                packed.data[i*n+k]=ratio;
46                for j in k+1..n {
47                    packed.data[i*n+j]=checked((-ratio).mul_add(packed.data[k*n+j], packed.data[i*n+j]), "LU update")?;
48                }
49            }
50        }
51        Ok(Self {
52            packed, pivots, sign
53        })
54    }
55    pub fn order(&self) -> usize {
56        self.packed.rows
57    }
58    pub fn pivots(&self) -> &[usize] {
59        &self.pivots
60    }
61    pub fn packed(&self) -> &Matrix {
62        &self.packed
63    }
64    pub fn lower(&self) -> Result<Matrix, SolverError> {
65        let n=self.order();
66        let mut l=Matrix::identity(n)?;
67        for i in 0..n {
68            for j in 0..i {
69                l.data[i*n+j]=self.packed.data[i*n+j];
70            }
71        }
72        Ok(l)
73    }
74    pub fn upper(&self) -> Result<Matrix, SolverError> {
75        let n=self.order();
76        let mut u=Matrix::zeros(n, n)?;
77        for i in 0..n {
78            for j in i..n {
79                u.data[i*n+j]=self.packed.data[i*n+j];
80            }
81        }
82        Ok(u)
83    }
84    /// Solve all RHS columns, reusing the factorization. RHS is not modified.
85    pub fn solve(&self, b: MatrixView<'_>) -> Result<Matrix, SolverError> {
86        let n=self.order();
87        let m=b.columns();
88        if b.rows()!=n {
89            return Err(SolverError::Shape("LU RHS row count"));
90        }
91        let mut x=b.to_owned()?;
92        for (k, &p) in self.pivots.iter().enumerate() {
93            if k!=p {
94                for j in 0..m {
95                    x.data.swap(k*m+j, p*m+j);
96                }
97            }
98        }
99        for i in 0..n {
100            for c in 0..m {
101                let sum=sum_iter((0..i).map(|j| self.packed.data[i*n+j]*x.data[j*m+c]))?;
102                x.data[i*m+c]=checked(x.data[i*m+c]-sum, "LU forward solve")?;
103            }
104        }
105        for i in (0..n).rev() {
106            for c in 0..m {
107                let sum=sum_iter((i+1..n).map(|j| self.packed.data[i*n+j]*x.data[j*m+c]))?;
108                x.data[i*m+c]=checked((x.data[i*m+c]-sum)/self.packed.data[i*n+i], "LU back solve")?;
109            }
110        }
111        Ok(x)
112    }
113    /// Solve `A^T X = B` without refactorizing A.
114    pub fn solve_transpose(&self, b: MatrixView<'_>) -> Result<Matrix, SolverError> {
115        let n=self.order();
116        let m=b.columns();
117        if b.rows()!=n {
118            return Err(SolverError::Shape("transposed LU RHS row count"));
119        }
120        let mut x=b.to_owned()?;
121        for i in 0..n {
122            for c in 0..m {
123                let sum=sum_iter((0..i).map(|j| self.packed.data[j*n+i]*x.data[j*m+c]))?;
124                x.data[i*m+c]=checked((x.data[i*m+c]-sum)/self.packed.data[i*n+i], "U transpose solve")?;
125            }
126        }
127        for i in (0..n).rev() {
128            for c in 0..m {
129                let sum=sum_iter((i+1..n).map(|j| self.packed.data[j*n+i]*x.data[j*m+c]))?;
130                x.data[i*m+c]=checked(x.data[i*m+c]-sum, "L transpose solve")?;
131            }
132        }
133        for k in (0..n).rev() {
134            let p=self.pivots[k];
135            if k!=p {
136                for j in 0..m {
137                    x.data.swap(k*m+j, p*m+j);
138                }
139            }
140        }
141        Ok(x)
142    }
143    /// Prefer solve(B) to explicitly forming an inverse for linear systems.
144    pub fn inverse(&self) -> Result<Matrix, SolverError> {
145        self.solve(Matrix::identity(self.order())?.view())
146    }
147    /// Determinant sign and log(abs(det)), avoiding a potentially overflowing product.
148    pub fn slogdet(&self) -> Result<(i32, f64), SolverError> {
149        let n=self.order();
150        let mut sign=self.sign;
151        for i in 0..n {
152            if self.packed.data[i*n+i]<0.0 {
153                sign=-sign;
154            }
155        }
156        let log=sum_iter((0..n).map(|i| self.packed.data[i*n+i].abs().ln()))?;
157        Ok((sign, log))
158    }
159}