faer/linalg/cholesky/llt_pivoting/
factor.rs

1use crate::assert;
2use crate::internal_prelude::*;
3pub use linalg::cholesky::llt::factor::LltError;
4use linalg::matmul::triangular::BlockStructure;
5#[derive(Copy, Clone, Debug)]
6pub struct PivLltParams {
7	pub block_size: usize,
8	#[doc(hidden)]
9	pub non_exhaustive: NonExhaustive,
10}
11impl Default for PivLltParams {
12	#[inline]
13	fn default() -> Self {
14		Self {
15			block_size: 128,
16			non_exhaustive: NonExhaustive(()),
17		}
18	}
19}
20impl<T> Auto<T> for PivLltParams {
21	#[inline]
22	fn auto() -> Self {
23		Self {
24			block_size: 128,
25			non_exhaustive: NonExhaustive(()),
26		}
27	}
28}
29#[derive(Copy, Clone, Debug)]
30pub struct PivLltInfo {
31	/// numerical rank of the matrix
32	pub rank: usize,
33	/// number of transpositions that make up the permutation
34	pub transposition_count: usize,
35}
36#[inline]
37pub fn cholesky_in_place_scratch<I: Index, T: ComplexField>(
38	dim: usize,
39	par: Par,
40	params: Spec<PivLltParams, T>,
41) -> StackReq {
42	_ = par;
43	_ = params;
44	temp_mat_scratch::<T::Real>(dim, 2)
45}
46#[track_caller]
47pub fn cholesky_in_place<'out, I: Index, T: ComplexField>(
48	a: MatMut<'_, T>,
49	perm: &'out mut [I],
50	perm_inv: &'out mut [I],
51	par: Par,
52	stack: &mut MemStack,
53	params: Spec<PivLltParams, T>,
54) -> Result<(PivLltInfo, PermRef<'out, I>), LltError> {
55	assert!(a.nrows() == a.ncols());
56	let n = a.nrows();
57	assert!(n <= I::Signed::MAX.zx());
58	let mut rank = n;
59	let mut transposition_count = 0;
60	'exit: {
61		if n > 0 {
62			let mut a = a;
63			for (i, p) in perm.iter_mut().enumerate() {
64				*p = I::truncate(i);
65			}
66			let (mut work1, stack) =
67				unsafe { temp_mat_uninit::<T::Real, _, _>(n, 1, stack) };
68			let (mut work2, _) =
69				unsafe { temp_mat_uninit::<T::Real, _, _>(n, 1, stack) };
70			let work1 = work1.as_mat_mut();
71			let work2 = work2.as_mat_mut();
72			let mut dot_products = work1.col_mut(0);
73			let mut diagonals = work2.col_mut(0);
74			let mut ajj = zero::<T::Real>();
75			let mut pvt = 0usize;
76			for i in 0..n {
77				let aii = a[(i, i)].real();
78				if aii < zero::<T::Real>() || aii.is_nan() {
79					return Err(LltError::NonPositivePivot { index: 0 });
80				}
81				if aii > ajj {
82					ajj = aii;
83					pvt = i;
84				}
85			}
86			let tol = eps::<T::Real>() * from_f64::<T::Real>(n as f64) * &ajj;
87			let mut k = 0usize;
88			while k < n {
89				let bs = Ord::min(n - k, params.block_size);
90				for i in k..n {
91					dot_products[i] = zero::<T::Real>();
92				}
93				for j in k..k + bs {
94					if j == k {
95						for i in j..n {
96							diagonals[i] = a[(i, i)].real();
97						}
98					} else {
99						for i in j..n {
100							dot_products[i] =
101								&dot_products[i] + a[(i, j - 1)].abs2();
102							diagonals[i] = a[(i, i)].real() - &dot_products[i];
103						}
104					}
105					if j > 0 {
106						pvt = j;
107						ajj = zero::<T::Real>();
108						for i in j..n {
109							let aii = &diagonals[i];
110							if aii.is_nan() {
111								return Err(LltError::NonPositivePivot {
112									index: j,
113								});
114							}
115							if *aii > ajj {
116								pvt = i;
117								ajj = aii.copy();
118							}
119						}
120						if ajj < tol {
121							rank = j;
122							// the trailing columns of `a` still hold the input matrix;
123							// clear them so the stored factor is the rank-`rank` one
124							// that `reconstruct` reads back as `L L^H`
125							a.rb_mut().get_mut(.., j..).fill(zero());
126							break 'exit;
127						}
128					}
129					if pvt != j {
130						transposition_count += 1;
131						a[(pvt, pvt)] = a[(j, j)].copy();
132						crate::perm::swap_rows_idx(
133							a.rb_mut().get_mut(.., ..j),
134							j,
135							pvt,
136						);
137						crate::perm::swap_cols_idx(
138							a.rb_mut().get_mut(pvt + 1.., ..),
139							j,
140							pvt,
141						);
142						unsafe {
143							z!(
144								a.rb().get(j + 1..pvt, j).const_cast(),
145								a.rb()
146									.get(pvt, j + 1..pvt)
147									.const_cast()
148									.transpose_mut(),
149							)
150						}
151						.for_each(|uz!(a, b)| (*a, *b) = (b.conj(), a.conj()));
152						a[(pvt, j)] = a[(pvt, j)].conj();
153						let tmp = dot_products[j].copy();
154						dot_products[j] = dot_products[pvt].copy();
155						dot_products[pvt] = tmp;
156						perm.swap(j, pvt);
157					}
158					ajj = ajj.sqrt();
159					a[(j, j)] = ajj.to_cplx();
160					unsafe {
161						linalg::matmul::matmul(
162							a.rb().get(j + 1.., j).const_cast(),
163							Accum::Add,
164							a.rb().get(j + 1.., k..j),
165							a.rb().get(j, k..j).adjoint(),
166							-one::<T>(),
167							par,
168						);
169					}
170					let ajj = &ajj.recip();
171					z!(a.rb_mut().get_mut(j + 1.., j))
172						.for_each(|uz!(x)| *x = x.mul_real(ajj));
173				}
174				linalg::matmul::triangular::matmul(
175					unsafe { a.rb().get(k + bs.., k + bs..).const_cast() },
176					BlockStructure::TriangularLower,
177					Accum::Add,
178					a.rb().get(k + bs.., k..k + bs),
179					BlockStructure::Rectangular,
180					a.rb().get(k + bs.., k..k + bs).adjoint(),
181					BlockStructure::Rectangular,
182					-one::<T>(),
183					par,
184				);
185				k += bs;
186			}
187			rank = n;
188		}
189	}
190	for (i, p) in perm.iter().enumerate() {
191		perm_inv[p.zx()] = I::truncate(i);
192	}
193	unsafe {
194		Ok((
195			PivLltInfo {
196				rank,
197				transposition_count,
198			},
199			PermRef::new_unchecked(perm, perm_inv, n),
200		))
201	}
202}
faer/linalg/cholesky/llt_pivoting/factor.rs

faer/linalg/cholesky/llt_pivoting/
factor.rs
