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: ParShape,
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	mut 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			// The pivot loop compares only later pivots with `tol`, which is
87			// zero when the largest diagonal entry is: that matrix has rank zero,
88			// as LAPACK `dpstrf` reports.
89			if ajj == zero::<T::Real>() {
90				rank = 0;
91				a.fill(zero());
92				break 'exit;
93			}
94			let tol = eps::<T::Real>() * from_f64::<T::Real>(n as f64) * &ajj;
95			let mut k = 0usize;
96			while k < n {
97				let bs = Ord::min(n - k, params.block_size);
98				for i in k..n {
99					dot_products[i] = zero::<T::Real>();
100				}
101				for j in k..k + bs {
102					if j == k {
103						for i in j..n {
104							diagonals[i] = a[(i, i)].real();
105						}
106					} else {
107						for i in j..n {
108							dot_products[i] =
109								&dot_products[i] + a[(i, j - 1)].abs2();
110							diagonals[i] = a[(i, i)].real() - &dot_products[i];
111						}
112					}
113					if j > 0 {
114						pvt = j;
115						ajj = zero::<T::Real>();
116						for i in j..n {
117							let aii = &diagonals[i];
118							if aii.is_nan() {
119								return Err(LltError::NonPositivePivot {
120									index: j,
121								});
122							}
123							if *aii > ajj {
124								pvt = i;
125								ajj = aii.copy();
126							}
127						}
128						if ajj < tol {
129							rank = j;
130							// the trailing columns of `a` still hold the input matrix;
131							// clear them so the stored factor is the rank-`rank` one
132							// that `reconstruct` reads back as `L L^H`
133							a.rb_mut().get_mut(.., j..).fill(zero());
134							break 'exit;
135						}
136					}
137					if pvt != j {
138						transposition_count += 1;
139						a[(pvt, pvt)] = a[(j, j)].copy();
140						crate::perm::swap_rows_idx(
141							a.rb_mut().get_mut(.., ..j),
142							j,
143							pvt,
144						);
145						crate::perm::swap_cols_idx(
146							a.rb_mut().get_mut(pvt + 1.., ..),
147							j,
148							pvt,
149						);
150						unsafe {
151							z!(
152								a.rb().get(j + 1..pvt, j).const_cast(),
153								a.rb()
154									.get(pvt, j + 1..pvt)
155									.const_cast()
156									.transpose_mut(),
157							)
158						}
159						.for_each(|uz!(a, b)| (*a, *b) = (b.conj(), a.conj()));
160						a[(pvt, j)] = a[(pvt, j)].conj();
161						let tmp = dot_products[j].copy();
162						dot_products[j] = dot_products[pvt].copy();
163						dot_products[pvt] = tmp;
164						perm.swap(j, pvt);
165					}
166					ajj = ajj.sqrt();
167					a[(j, j)] = ajj.to_cplx();
168					unsafe {
169						linalg::matmul::matmul(
170							a.rb().get(j + 1.., j).const_cast(),
171							Accum::Add,
172							a.rb().get(j + 1.., k..j),
173							a.rb().get(j, k..j).adjoint(),
174							-one::<T>(),
175							par.reborrow(),
176						);
177					}
178					let ajj = &ajj.recip();
179					z!(a.rb_mut().get_mut(j + 1.., j))
180						.for_each(|uz!(x)| *x = x.mul_real(ajj));
181				}
182				linalg::matmul::triangular::matmul(
183					unsafe { a.rb().get(k + bs.., k + bs..).const_cast() },
184					BlockStructure::TriangularLower,
185					Accum::Add,
186					a.rb().get(k + bs.., k..k + bs),
187					BlockStructure::Rectangular,
188					a.rb().get(k + bs.., k..k + bs).adjoint(),
189					BlockStructure::Rectangular,
190					-one::<T>(),
191					par.reborrow(),
192				);
193				k += bs;
194			}
195			rank = n;
196		}
197	}
198	for (i, p) in perm.iter().enumerate() {
199		perm_inv[p.zx()] = I::truncate(i);
200	}
201	unsafe {
202		Ok((
203			PivLltInfo {
204				rank,
205				transposition_count,
206			},
207			PermRef::new_unchecked(perm, perm_inv, n),
208		))
209	}
210}
faer/linalg/cholesky/llt_pivoting/factor.rs

faer/linalg/cholesky/llt_pivoting/
factor.rs
