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							a[(j, j)] = ajj.to_cplx();
123							break 'exit;
124						}
125					}
126					if pvt != j {
127						transposition_count += 1;
128						a[(pvt, pvt)] = a[(j, j)].copy();
129						crate::perm::swap_rows_idx(
130							a.rb_mut().get_mut(.., ..j),
131							j,
132							pvt,
133						);
134						crate::perm::swap_cols_idx(
135							a.rb_mut().get_mut(pvt + 1.., ..),
136							j,
137							pvt,
138						);
139						unsafe {
140							z!(
141								a.rb().get(j + 1..pvt, j).const_cast(),
142								a.rb()
143									.get(pvt, j + 1..pvt)
144									.const_cast()
145									.transpose_mut(),
146							)
147						}
148						.for_each(|uz!(a, b)| (*a, *b) = (b.conj(), a.conj()));
149						a[(pvt, j)] = a[(pvt, j)].conj();
150						let tmp = dot_products[j].copy();
151						dot_products[j] = dot_products[pvt].copy();
152						dot_products[pvt] = tmp;
153						perm.swap(j, pvt);
154					}
155					ajj = ajj.sqrt();
156					a[(j, j)] = ajj.to_cplx();
157					unsafe {
158						linalg::matmul::matmul(
159							a.rb().get(j + 1.., j).const_cast(),
160							Accum::Add,
161							a.rb().get(j + 1.., k..j),
162							a.rb().get(j, k..j).adjoint(),
163							-one::<T>(),
164							par,
165						);
166					}
167					let ajj = &ajj.recip();
168					z!(a.rb_mut().get_mut(j + 1.., j))
169						.for_each(|uz!(x)| *x = x.mul_real(ajj));
170				}
171				linalg::matmul::triangular::matmul(
172					unsafe { a.rb().get(k + bs.., k + bs..).const_cast() },
173					BlockStructure::TriangularLower,
174					Accum::Add,
175					a.rb().get(k + bs.., k..k + bs),
176					BlockStructure::Rectangular,
177					a.rb().get(k + bs.., k..k + bs).adjoint(),
178					BlockStructure::Rectangular,
179					-one::<T>(),
180					par,
181				);
182				k += bs;
183			}
184			rank = n;
185		}
186	}
187	for (i, p) in perm.iter().enumerate() {
188		perm_inv[p.zx()] = I::truncate(i);
189	}
190	unsafe {
191		Ok((
192			PivLltInfo {
193				rank,
194				transposition_count,
195			},
196			PermRef::new_unchecked(perm, perm_inv, n),
197		))
198	}
199}
faer/linalg/cholesky/llt_pivoting/factor.rs

faer/linalg/cholesky/llt_pivoting/
factor.rs
