1use crate::{
3 Matrix, MatrixView, SolverError, Tolerance
4};
5use crate::numerics::{
6 checked, sum_iter
7};
8#[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 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 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 pub fn inverse(&self) -> Result<Matrix, SolverError> {
145 self.solve(Matrix::identity(self.order())?.view())
146 }
147 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}