Skip to main content

faer/linalg/cholesky/bunch_kaufman/
factor.rs

1use crate::internal_prelude::*;
2use crate::{assert, perm};
3use linalg::matmul::triangular::BlockStructure;
4const TOP_BIT: usize = 1 << (usize::BITS - 1);
5/// pivoting strategy for choosing the pivots
6#[derive(Copy, Clone, Debug, PartialEq, Eq)]
7#[non_exhaustive]
8pub enum PivotingStrategy {
9	/// deprecated, corresponds to partial pivoting
10	#[deprecated]
11	Diagonal,
12	/// searches for the k-th pivot in the k-th column
13	Partial,
14	/// searches for the k-th pivot in the k-th column, as well as the tail of
15	/// the diagonal of the matrix
16	PartialDiag,
17	/// searches for pivots that are locally optimal
18	Rook,
19	/// searches for pivots that are locally optimal, as well as the tail of the
20	/// diagonal of the matrix
21	RookDiag,
22	/// searches for pivots that are globally optimal
23	Full,
24}
25/// tuning parameters for the decomposition
26#[derive(Copy, Clone, Debug)]
27pub struct LbltParams {
28	/// pivoting strategy
29	pub pivoting: PivotingStrategy,
30	/// block size of the algorithm
31	pub block_size: usize,
32	/// threshold at which size parallelism should be disabled
33	pub par_threshold: usize,
34	#[doc(hidden)]
35	pub non_exhaustive: NonExhaustive,
36}
37fn swap_self_adjoint<T: ComplexField>(A: MatMut<'_, T>, i: usize, j: usize) {
38	assert_ne!(i, j);
39	let mut A = A;
40	let (i, j) = (Ord::min(i, j), Ord::max(i, j));
41	perm::swap_cols_idx(A.rb_mut().get_mut(j + 1.., ..), i, j);
42	perm::swap_rows_idx(A.rb_mut().get_mut(.., ..i), i, j);
43	let tmp = A[(i, i)].real();
44	A[(i, i)] = A[(j, j)].as_real();
45	A[(j, j)] = tmp.to_cplx();
46	A[(j, i)] = A[(j, i)].conj();
47	let (Ai, Aj) = A.split_at_row_mut(j);
48	let Ai = Ai.get_mut(i + 1..j, i);
49	let Aj = Aj.get_mut(0, i + 1..j).transpose_mut();
50	zip!(Ai, Aj).for_each(|unzip!(x, y): Zip!(&mut _, &mut _)| {
51		(*x, *y) = (y.conj(), x.conj());
52	});
53}
54fn rank_1_update_and_argmax_fallback<'M, 'N, T: ComplexField>(
55	A: MatMut<'_, T, Dim<'N>, Dim<'N>>,
56	L: ColRef<'_, T, Dim<'N>>,
57	d: T::Real,
58	start: IdxInc<'N>,
59	end: IdxInc<'N>,
60) -> (usize, usize, T::Real) {
61	let mut A = A;
62	let n = A.nrows();
63	let mut max_j = n.idx(0);
64	let mut max_i = n.idx(0);
65	let mut max_offdiag = zero();
66	for j in start.to(end) {
67		for i in j.next().to(n.end()) {
68			A[(i, j)] = &A[(i, j)] - (&L[i] * L[j].conj()).mul_real(&d);
69			let val = A[(i, j)].abs2();
70			if val > max_offdiag {
71				max_offdiag = val;
72				max_i = i;
73				max_j = j;
74			}
75		}
76	}
77	(*max_i, *max_j, max_offdiag)
78}
79fn rank_2_update_and_argmax_fallback<'N, T: ComplexField>(
80	A: MatMut<'_, T, Dim<'N>, Dim<'N>>,
81	L0: ColRef<'_, T, Dim<'N>>,
82	L1: ColRef<'_, T, Dim<'N>>,
83	d: T::Real,
84	d00: T::Real,
85	d11: T::Real,
86	d10: T,
87	start: IdxInc<'N>,
88	end: IdxInc<'N>,
89) -> (usize, usize, T::Real) {
90	let mut A = A;
91	let n = A.nrows();
92	let mut max_j = n.idx(0);
93	let mut max_i = n.idx(0);
94	let mut max_offdiag = zero();
95	for j in start.to(end) {
96		let x0 = L0[j].copy();
97		let x1 = L1[j].copy();
98		let w0 = (x0.mul_real(&d11) - &x1 * &d10).mul_real(&d);
99		let w1 = (x1.mul_real(&d00) - &x0 * d10.conj()).mul_real(&d);
100		for i in j.next().to(n.end()) {
101			A[(i, j)] = &A[(i, j)] - &L0[i] * w0.conj() - &L1[i] * w1.conj();
102			let val = A[(i, j)].abs2();
103			if val > max_offdiag {
104				max_offdiag = val;
105				max_i = i;
106				max_j = j;
107			}
108		}
109	}
110	(*max_i, *max_j, max_offdiag)
111}
112fn rank_1_update_and_argmax_seq<'M, 'N, T: ComplexField>(
113	A: MatMut<'_, T, Dim<'N>, Dim<'N>>,
114	L: ColRef<'_, T, Dim<'N>>,
115	d: T::Real,
116	start: IdxInc<'N>,
117	end: IdxInc<'N>,
118) -> (usize, usize, T::Real) {
119	rank_1_update_and_argmax_fallback(A, L, d, start, end)
120}
121fn rank_2_update_and_argmax_seq<'N, T: ComplexField>(
122	A: MatMut<'_, T, Dim<'N>, Dim<'N>>,
123	L0: ColRef<'_, T, Dim<'N>>,
124	L1: ColRef<'_, T, Dim<'N>>,
125	d: T::Real,
126	d00: T::Real,
127	d11: T::Real,
128	d10: T,
129	start: IdxInc<'N>,
130	end: IdxInc<'N>,
131) -> (usize, usize, T::Real) {
132	rank_2_update_and_argmax_fallback(A, L0, L1, d, d00, d11, d10, start, end)
133}
134fn rank_1_update_and_argmax<T: ComplexField>(
135	A: MatMut<'_, T>,
136	L: ColRef<'_, T>,
137	d: T::Real,
138	par: Par,
139) -> (usize, usize, T::Real) {
140	with_dim!(N, A.nrows());
141	match par {
142		Par::Seq => rank_1_update_and_argmax_seq(
143			A.as_shape_mut(N, N),
144			L.as_row_shape(N),
145			d,
146			IdxInc::ZERO,
147			N.end(),
148		),
149		#[cfg(feature = "rayon")]
150		Par::Rayon(nthreads) => {
151			use rayon::prelude::*;
152			let nthreads = nthreads.get();
153			let n = *N;
154			assert!((n as u64) < (1u64 << 50));
155			let idx_to_col_start = |idx: usize| {
156				let idx_as_percent = idx as f64 / nthreads as f64;
157				let col_start_percent =
158					1.0f64 - libm::sqrt(1.0f64 - idx_as_percent);
159				(col_start_percent * n as f64) as usize
160			};
161			let mut r =
162				alloc::vec![(0usize, 0usize, zero::< T::Real > ()); nthreads];
163			spindle::for_each(
164				nthreads,
165				r.par_iter_mut().enumerate(),
166				|(idx, out)| {
167					let A = unsafe { A.rb().const_cast() };
168					let start = N.idx_inc(idx_to_col_start(idx));
169					let end = N.idx_inc(idx_to_col_start(idx + 1));
170					*out = rank_1_update_and_argmax_seq(
171						A.as_shape_mut(N, N),
172						L.as_row_shape(N),
173						d.copy(),
174						start,
175						end,
176					);
177				},
178			);
179			r.into_iter()
180				.max_by(|(_, _, a), (_, _, b)| {
181					if a == b {
182						core::cmp::Ordering::Equal
183					} else if a > b {
184						core::cmp::Ordering::Greater
185					} else {
186						core::cmp::Ordering::Less
187					}
188				})
189				.unwrap()
190		},
191	}
192}
193fn rank_2_update_and_argmax<'N, T: ComplexField>(
194	A: MatMut<'_, T>,
195	L0: ColRef<'_, T>,
196	L1: ColRef<'_, T>,
197	d: T::Real,
198	d00: T::Real,
199	d11: T::Real,
200	d10: T,
201	par: Par,
202) -> (usize, usize, T::Real) {
203	with_dim!(N, A.nrows());
204	match par {
205		Par::Seq => rank_2_update_and_argmax_seq(
206			A.as_shape_mut(N, N),
207			L0.as_row_shape(N),
208			L1.as_row_shape(N),
209			d,
210			d00,
211			d11,
212			d10,
213			IdxInc::ZERO,
214			N.end(),
215		),
216		#[cfg(feature = "rayon")]
217		Par::Rayon(nthreads) => {
218			use rayon::prelude::*;
219			let nthreads = nthreads.get();
220			let n = *N;
221			assert!((n as u64) < (1u64 << 50));
222			let idx_to_col_start = |idx: usize| {
223				let idx_as_percent = idx as f64 / nthreads as f64;
224				let col_start_percent =
225					1.0f64 - libm::sqrt(1.0f64 - idx_as_percent);
226				(col_start_percent * n as f64) as usize
227			};
228			let mut r =
229				alloc::vec![(0usize, 0usize, zero::< T::Real > ()); nthreads];
230			spindle::for_each(
231				nthreads,
232				r.par_iter_mut().enumerate(),
233				|(idx, out)| {
234					let A = unsafe { A.rb().const_cast() };
235					let start = N.idx_inc(idx_to_col_start(idx));
236					let end = N.idx_inc(idx_to_col_start(idx + 1));
237					*out = rank_2_update_and_argmax_seq(
238						A.as_shape_mut(N, N),
239						L0.as_row_shape(N),
240						L1.as_row_shape(N),
241						d.copy(),
242						d00.copy(),
243						d11.copy(),
244						d10.copy(),
245						start,
246						end,
247					);
248				},
249			);
250			r.into_iter()
251				.max_by(|(_, _, a), (_, _, b)| {
252					if a == b {
253						core::cmp::Ordering::Equal
254					} else if a < b {
255						core::cmp::Ordering::Less
256					} else {
257						core::cmp::Ordering::Greater
258					}
259				})
260				.unwrap()
261		},
262	}
263}
264fn lblt_full_piv<T: ComplexField>(
265	A: MatMut<'_, T>,
266	subdiag: DiagMut<'_, T>,
267	pivots: &mut [usize],
268	par: Par,
269	params: LbltParams,
270) {
271	let alpha = ((one::<T::Real>() + from_f64::<T::Real>(17.0).sqrt())
272		* from_f64::<T::Real>(0.125))
273	.abs2();
274	let mut A = A;
275	let mut subdiag = subdiag.column_vector_mut();
276	let mut par = par;
277	let n = A.nrows();
278	let scale_fwd = &A.norm_max();
279	let scale_bwd = &scale_fwd.recip();
280	zip!(A.rb_mut()).for_each(|unzip!(x)| *x = x.mul_real(scale_bwd));
281	let mut max_i = 0;
282	let mut max_j = 0;
283	let mut max_offdiag = zero();
284	for j in 0..n {
285		for i in j + 1..n {
286			let val = A[(i, j)].abs2();
287			if val > max_offdiag {
288				max_offdiag = val;
289				max_i = i;
290				max_j = j;
291			}
292		}
293	}
294	let mut k = 0;
295	while k < n {
296		if max_offdiag == zero() {
297			break;
298		}
299		let (mut Aprev, mut A) =
300			A.rb_mut().get_mut(k.., ..).split_at_col_mut(k);
301		let mut subdiag = subdiag.rb_mut().get_mut(k..);
302		let pivots = &mut pivots[k..];
303		let n = A.nrows();
304		let mut max_s = 0;
305		let mut max_diag = zero();
306		for s in 0..n {
307			let val = A[(s, s)].abs2();
308			if val > max_diag {
309				max_diag = val;
310				max_s = s;
311			}
312		}
313		let npiv;
314		let i0;
315		let i1;
316		if max_diag >= &alpha * &max_offdiag {
317			npiv = 1;
318			i0 = max_s;
319			i1 = usize::MAX;
320		} else {
321			npiv = 2;
322			i0 = max_j;
323			i1 = max_i;
324		}
325		let rem = n - npiv;
326		if rem * rem < params.par_threshold {
327			par = Par::Seq;
328		}
329		if i0 != 0 {
330			swap_self_adjoint(A.rb_mut(), 0, i0);
331			perm::swap_rows_idx(Aprev.rb_mut(), 0, i0);
332		}
333		if npiv == 2 && i1 != 1 {
334			swap_self_adjoint(A.rb_mut(), 1, i1);
335			perm::swap_rows_idx(Aprev.rb_mut(), 1, i1);
336		}
337		if npiv == 1 {
338			let diag = A[(0, 0)].real();
339			let diag_inv = &diag.recip();
340			subdiag[0] = zero();
341			let (_, _, L, mut A) = A.rb_mut().split_at_mut(1, 1);
342			let n = A.nrows();
343			let mut L = L.col_mut(0);
344			zip!(L.rb_mut()).for_each(|unzip!(x)| *x = x.mul_real(diag_inv));
345			for i in 0..n {
346				A[(i, i)] = (A[(i, i)].real() - &diag * L[i].abs2()).to_cplx();
347			}
348			if n < params.par_threshold {}
349			if n != 0 {
350				(max_i, max_j, max_offdiag) =
351					rank_1_update_and_argmax(A.rb_mut(), L.rb(), diag, par);
352			}
353		} else {
354			let a00 = A[(0, 0)].real();
355			let a11 = A[(1, 1)].real();
356			let a10 = A[(1, 0)].copy();
357			subdiag[0] = a10.copy();
358			subdiag[1] = zero();
359			A[(1, 0)] = zero();
360			let d10 = a10.abs();
361			let d10_inv = &d10.recip();
362			let d00 = a00 * d10_inv;
363			let d11 = a11 * d10_inv;
364			let t = (&d00 * &d11 - one::<T::Real>()).recip();
365			let d10 = a10.mul_real(d10_inv);
366			let d = t * d10_inv;
367			let (_, _, L, mut A) = A.rb_mut().split_at_mut(2, 2);
368			let (mut L0, mut L1) = L.two_cols_mut(0, 1);
369			let n = A.nrows();
370			if n != 0 {
371				(max_i, max_j, max_offdiag) = rank_2_update_and_argmax(
372					A.rb_mut(),
373					L0.rb(),
374					L1.rb(),
375					d.copy(),
376					d00.copy(),
377					d11.copy(),
378					d10.copy(),
379					par,
380				);
381			}
382			for j in 0..n {
383				let x0 = &L0[j].copy();
384				let x1 = &L1[j].copy();
385				let w0 = (x0.mul_real(&d11) - x1 * &d10).mul_real(&d);
386				let w1 = (x1.mul_real(&d00) - x0 * d10.conj()).mul_real(&d);
387				A[(j, j)] =
388					(&A[(j, j)] - &L0[j] * w0.conj() - &L1[j] * w1.conj())
389						.as_real();
390				L0[j] = w0;
391				L1[j] = w1;
392			}
393		}
394		if npiv == 2 {
395			pivots[0] = (i0 + k) | TOP_BIT;
396			pivots[1] = (i1 + k) | TOP_BIT;
397		} else {
398			pivots[0] = i0 + k;
399		}
400		k += npiv;
401	}
402	while k < n {
403		let (mut Aprev, mut A) =
404			A.rb_mut().get_mut(k.., ..).split_at_col_mut(k);
405		let mut subdiag = subdiag.rb_mut().get_mut(k..);
406		let pivots = &mut pivots[k..];
407		let n = A.nrows();
408		let mut max_s = 0;
409		let mut max_diag = zero();
410		for s in 0..n {
411			let val = A[(s, s)].abs2();
412			if val > max_diag {
413				max_diag = val;
414				max_s = s;
415			}
416		}
417		if max_s != 0 {
418			let (mut A0, mut As) = A.rb_mut().two_cols_mut(0, max_s);
419			core::mem::swap(&mut A0[0], &mut As[max_s]);
420			perm::swap_rows_idx(Aprev.rb_mut(), 0, max_s);
421		}
422		subdiag[0] = zero();
423		pivots[0] = max_s + k;
424		k += 1;
425	}
426	zip!(A.rb_mut().diagonal_mut().column_vector_mut())
427		.for_each(|unzip!(x)| *x = x.mul_real(scale_fwd));
428	zip!(subdiag.rb_mut()).for_each(|unzip!(x)| *x = x.mul_real(scale_fwd));
429}
430#[track_caller]
431fn l1_argmax<T: ComplexField>(col: ColRef<'_, T>) -> (Option<usize>, T::Real) {
432	let n = col.nrows();
433	if n == 0 {
434		return (None, zero());
435	}
436	let mut i = 0;
437	let mut best = zero();
438	for j in 0..n {
439		let val = col[j].abs1();
440		if val > best {
441			best = val;
442			i = j;
443		}
444	}
445	(Some(i), best)
446}
447#[track_caller]
448fn offdiag_argmax<T: ComplexField>(
449	A: MatRef<'_, T>,
450	idx: usize,
451) -> (Option<usize>, T::Real) {
452	let (mut col_argmax, col_max) = l1_argmax(A.rb().get(idx + 1.., idx));
453	col_argmax.as_mut().map(|col_argmax| *col_argmax += idx + 1);
454	let (row_argmax, row_max) = l1_argmax(A.rb().get(idx, ..idx).transpose());
455	if col_max > row_max {
456		(col_argmax, col_max)
457	} else {
458		(row_argmax, row_max)
459	}
460}
461fn update_and_offdiag_argmax<T: ComplexField>(
462	mut dst: ColMut<'_, T>,
463	Wl: MatRef<'_, T>,
464	Al: MatRef<'_, T>,
465	Ar: MatRef<'_, T>,
466	i0: usize,
467	par: Par,
468) -> (Option<usize>, T::Real) {
469	let n = Al.nrows();
470	for j in 0..i0 {
471		dst[j] = Ar[(i0, j)].conj();
472	}
473	dst[i0] = zero();
474	for j in i0 + 1..n {
475		dst[j] = Ar[(j, i0)].copy();
476	}
477	linalg::matmul::matmul(
478		dst.rb_mut(),
479		Accum::Add,
480		Al.rb(),
481		Wl.row(i0).adjoint(),
482		-one::<T>(),
483		par,
484	);
485	dst[i0] = zero();
486	let ret = l1_argmax(dst.rb());
487	dst[i0] = Ar[(i0, i0)].as_real();
488	if n == 1 { (None, zero()) } else { ret }
489}
490#[inline(never)]
491fn lblt_blocked_step<T: ComplexField>(
492	alpha: T::Real,
493	W: MatMut<'_, T>,
494	A_left: MatMut<'_, T>,
495	A: MatMut<'_, T>,
496	subdiag: DiagMut<'_, T>,
497	pivots: &mut [usize],
498	rook: bool,
499	diagonal: bool,
500	par: Par,
501) -> usize {
502	let mut A = A;
503	let mut A_left = A_left;
504	let mut subdiag = subdiag;
505	let mut W = W;
506	let n = A.nrows();
507	let block_size = W.ncols();
508	assert!(all(
509		A.nrows() == n,
510		A.ncols() == n,
511		W.nrows() == n,
512		subdiag.dim() == n,
513		block_size >= 2,
514	));
515	let kmax = Ord::min(block_size - 1, n);
516	let mut k = 0usize;
517	while k < kmax {
518		let mut A = A.rb_mut();
519		let mut W = W.rb_mut();
520		let mut subdiag = subdiag.rb_mut().column_vector_mut().get_mut(k..);
521		let A_left = A_left.rb_mut().get_mut(k.., ..);
522		let (mut Wl, mut Wr) = W.rb_mut().get_mut(k.., ..).split_at_col_mut(k);
523		let (mut Al, mut Ar) = A.rb_mut().get_mut(k.., ..).split_at_col_mut(k);
524		let mut Al = Al.rb_mut();
525		let mut Wr = Wr.rb_mut().get_mut(.., ..2);
526		let npiv;
527		let mut i0 = if diagonal {
528			l1_argmax(Ar.rb().diagonal().column_vector()).0.unwrap()
529		} else {
530			0
531		};
532		let mut i1 = usize::MAX;
533		let mut nothing_to_do = false;
534		let (mut Wr0, mut Wr1) = Wr.rb_mut().two_cols_mut(0, 1);
535		let (r, mut gamma_i) = update_and_offdiag_argmax(
536			Wr0.rb_mut(),
537			Wl.rb(),
538			Al.rb(),
539			Ar.rb(),
540			i0,
541			par,
542		);
543		if k + 1 == n || gamma_i == zero() {
544			nothing_to_do = true;
545			npiv = 1;
546		} else if Ar[(i0, i0)].real().abs() >= &alpha * &gamma_i {
547			npiv = 1;
548		} else {
549			i1 = r.unwrap();
550			if rook {
551				loop {
552					let (s, gamma_r) = update_and_offdiag_argmax(
553						Wr1.rb_mut(),
554						Wl.rb(),
555						Al.rb(),
556						Ar.rb(),
557						i1,
558						par,
559					);
560					if Ar[(i1, i1)].abs1() >= &alpha * &gamma_r {
561						npiv = 1;
562						i0 = i1;
563						i1 = usize::MAX;
564						Wr0.copy_from(&Wr1);
565						break;
566					} else if s == Some(i0) || gamma_i == gamma_r {
567						npiv = 2;
568						break;
569					} else {
570						i0 = i1;
571						i1 = s.unwrap();
572						gamma_i = gamma_r;
573						Wr0.copy_from(&Wr1);
574					}
575				}
576			} else {
577				let (_, gamma_r) = update_and_offdiag_argmax(
578					Wr1.rb_mut(),
579					Wl.rb(),
580					Al.rb(),
581					Ar.rb(),
582					i1,
583					par,
584				);
585				if Ar[(i0, i0)].real().abs()
586					>= (&alpha * &gamma_r) * (&gamma_r / &gamma_i)
587				{
588					npiv = 1;
589				} else if Ar[(i1, i1)].real().abs() >= &alpha * &gamma_r {
590					npiv = 1;
591					i0 = i1;
592					i1 = usize::MAX;
593					Wr0.copy_from(&Wr1);
594				} else {
595					npiv = 2;
596				}
597			}
598		}
599		if npiv == 2 && i0 > i1 {
600			perm::swap_cols_idx(Wr.rb_mut(), 0, 1);
601			(i0, i1) = (i1, i0);
602		}
603		let mut Wr = Wr.rb_mut().get_mut(.., ..npiv);
604		'next_iter: {
605			if i0 != 0 {
606				swap_self_adjoint(Ar.rb_mut(), 0, i0);
607				perm::swap_rows_idx(Al.rb_mut(), 0, i0);
608				perm::swap_rows_idx(Wl.rb_mut(), 0, i0);
609				perm::swap_rows_idx(Wr.rb_mut(), 0, i0);
610			}
611			if npiv == 2 && i1 != 1 {
612				swap_self_adjoint(Ar.rb_mut(), 1, i1);
613				perm::swap_rows_idx(Al.rb_mut(), 1, i1);
614				perm::swap_rows_idx(Wl.rb_mut(), 1, i1);
615				perm::swap_rows_idx(Wr.rb_mut(), 1, i1);
616			}
617			if nothing_to_do {
618				break 'next_iter;
619			}
620			if npiv == 1 {
621				let W0 = Wr.rb_mut().col_mut(0);
622				let diag = &W0[0].real();
623				let diag_inv = &diag.recip();
624				subdiag[0] = zero();
625				let (_, _, L, mut A) = Ar.rb_mut().split_at_mut(1, 1);
626				let W0 = W0.rb().get(1..);
627				let n = A.nrows();
628				let mut L = L.col_mut(0);
629				zip!(W0, L.rb_mut()).for_each(
630					|unzip!(w, a): Zip!(&T, &mut T)| {
631						*a = w.mul_real(diag_inv);
632					},
633				);
634				for j in 0..n {
635					A[(j, j)] =
636						(A[(j, j)].real() - diag * &L[j].abs2()).to_cplx();
637				}
638			} else {
639				let a00 = Wr[(0, 0)].real();
640				let a11 = Wr[(1, 1)].real();
641				let a10 = Wr[(1, 0)].copy();
642				subdiag[0] = a10.copy();
643				subdiag[1] = zero();
644				Wr[(1, 0)] = zero();
645				Ar[(1, 0)] = zero();
646				let d10 = a10.abs();
647				let d10_inv = &d10.recip();
648				let d00 = a00 * d10_inv;
649				let d11 = a11 * d10_inv;
650				let t = (&d00 * &d11 - one::<T::Real>()).recip();
651				let d10 = a10.mul_real(d10_inv);
652				let d = t * d10_inv;
653				let (_, _, L, mut A) = Ar.rb_mut().split_at_mut(2, 2);
654				let (mut L0, mut L1) = L.two_cols_mut(0, 1);
655				let Wr = Wr.rb().get(2.., ..);
656				let W0 = Wr.col(0);
657				let W1 = Wr.col(1);
658				let n = A.nrows();
659				for j in 0..n {
660					let x0 = &W0[j];
661					let x1 = &W1[j];
662					let w0 = (x0.mul_real(&d11) - x1 * &d10).mul_real(&d);
663					let w1 = (x1.mul_real(&d00) - x0 * d10.conj()).mul_real(&d);
664					A[(j, j)] =
665						(&A[(j, j)] - &W0[j] * w0.conj() - &W1[j] * w1.conj())
666							.as_real();
667					L0[j] = w0;
668					L1[j] = w1;
669				}
670			}
671		}
672		let offset = A_left.ncols();
673		if npiv == 2 {
674			pivots[k] = (offset + i0 + k) | TOP_BIT;
675			pivots[k + 1] = (offset + i1 + k) | TOP_BIT;
676		} else {
677			pivots[k] = offset + i0 + k;
678		}
679		k += npiv;
680	}
681	let W = W.rb().get(k.., ..k);
682	let (_, _, Al, mut Ar) = A.rb_mut().split_at_mut(k, k);
683	let Al = Al.rb();
684	linalg::matmul::triangular::matmul(
685		Ar.rb_mut(),
686		BlockStructure::StrictTriangularLower,
687		Accum::Add,
688		W,
689		BlockStructure::Rectangular,
690		Al.adjoint(),
691		BlockStructure::Rectangular,
692		-one::<T>(),
693		par,
694	);
695	for j in 0..n - k {
696		Ar[(j, j)] = Ar[(j, j)].as_real();
697	}
698	k
699}
700fn lblt_blocked<T: ComplexField>(
701	A: MatMut<'_, T>,
702	subdiag: DiagMut<'_, T>,
703	pivots: &mut [usize],
704	block_size: usize,
705	rook: bool,
706	diagonal: bool,
707	par: Par,
708	stack: &mut MemStack,
709) {
710	let ref alpha = (one::<T::Real>() + from_f64::<T::Real>(17.0).sqrt())
711		* from_f64::<T::Real>(0.125);
712	let mut A = A;
713	let mut subdiag = subdiag.column_vector_mut();
714	let n = A.nrows();
715	let mut k = 0;
716	while k < n {
717		let (_, _, mut A_left, A_right) = A.rb_mut().split_at_mut(k, k);
718		let (mut W, _) =
719			unsafe { temp_mat_uninit::<T, _, _>(n - k, block_size, stack) };
720		let W = W.as_mat_mut();
721		let next;
722		if block_size < 2 || n - k <= block_size {
723			lblt_unblocked(
724				alpha.copy(),
725				A_left.rb_mut(),
726				A_right,
727				subdiag.rb_mut().get_mut(k..).as_diagonal_mut(),
728				&mut pivots[k..],
729				rook,
730				diagonal,
731				par,
732			);
733			next = n;
734		} else {
735			let block_size = lblt_blocked_step(
736				alpha.copy(),
737				W,
738				A_left.rb_mut(),
739				A_right,
740				subdiag.rb_mut().get_mut(k..).as_diagonal_mut(),
741				&mut pivots[k..],
742				rook,
743				diagonal,
744				par,
745			);
746			next = k + block_size;
747		}
748		let pivots = &pivots[k..next];
749		let A_left = A.rb_mut().get_mut(.., ..k);
750		if A_left.ncols() > 0 {
751			match par {
752				Par::Seq => {
753					for mut col in A_left.col_iter_mut() {
754						for (i, &j) in core::iter::zip(k..next, pivots) {
755							let j = j & !TOP_BIT;
756							linalg::lu::partial_pivoting::factor::swap_elems(
757								col.rb_mut(),
758								i,
759								j,
760							);
761						}
762					}
763				},
764				#[cfg(feature = "rayon")]
765				Par::Rayon(nthreads) => {
766					let nthreads = nthreads.get();
767					spindle::for_each(
768						nthreads,
769						A_left.par_col_iter_mut(),
770						|mut col| {
771							for (i, &j) in core::iter::zip(k..next, pivots) {
772								let j = j & !TOP_BIT;
773								linalg::lu::partial_pivoting::factor::swap_elems(col.rb_mut(), i, j);
774							}
775						},
776					);
777				},
778			}
779		}
780		k = next;
781	}
782}
783#[inline(never)]
784fn lblt_unblocked<T: ComplexField>(
785	alpha: T::Real,
786	A_left: MatMut<'_, T>,
787	A: MatMut<'_, T>,
788	subdiag: DiagMut<'_, T>,
789	pivots: &mut [usize],
790	rook: bool,
791	diagonal: bool,
792	par: Par,
793) {
794	let _ = par;
795	let mut A = A;
796	let mut A_left = A_left;
797	let mut subdiag = subdiag;
798	let n = A.nrows();
799	assert!(all(A.nrows() == n, A.ncols() == n, subdiag.dim() == n));
800	let mut k = 0usize;
801	while k < n {
802		let (_, _, mut L_prev, mut A) = A.rb_mut().split_at_mut(k, k);
803		let mut subdiag = subdiag.rb_mut().column_vector_mut().get_mut(k..);
804		let A_left = A_left.rb_mut().get_mut(k.., ..);
805		let npiv;
806		let mut i0 = if diagonal {
807			l1_argmax(A.rb().diagonal().column_vector()).0.unwrap()
808		} else {
809			0
810		};
811		let mut i1 = usize::MAX;
812		let (r, mut gamma_i) = offdiag_argmax(A.rb(), i0);
813		let mut nothing_to_do = false;
814		if k + 1 == n || gamma_i == zero() {
815			nothing_to_do = true;
816			npiv = 1;
817		} else if A[(i0, i0)].real().abs() >= &alpha * &gamma_i {
818			npiv = 1;
819		} else {
820			i1 = r.unwrap();
821			if rook {
822				loop {
823					let (s, gamma_r) = offdiag_argmax(A.rb(), i1);
824					if A[(i1, i1)].abs1() >= &alpha * &gamma_r {
825						npiv = 1;
826						i0 = i1;
827						i1 = usize::MAX;
828						break;
829					} else if gamma_i == gamma_r {
830						npiv = 2;
831						break;
832					} else {
833						i0 = i1;
834						i1 = s.unwrap();
835						gamma_i = gamma_r;
836					}
837				}
838			} else {
839				let (_, gamma_r) = offdiag_argmax(A.rb(), i1);
840				if A[(i0, i0)].real().abs()
841					>= (&alpha * &gamma_r) * (&gamma_r / &gamma_i)
842				{
843					npiv = 1;
844				} else if A[(i1, i1)].real().abs() >= &alpha * &gamma_r {
845					npiv = 1;
846					i0 = i1;
847				} else {
848					npiv = 2;
849				}
850			}
851		}
852		if npiv == 2 && i0 > i1 {
853			(i0, i1) = (i1, i0);
854		}
855		'next_iter: {
856			if i0 != 0 {
857				swap_self_adjoint(A.rb_mut(), 0, i0);
858				perm::swap_rows_idx(L_prev.rb_mut(), 0, i0);
859			}
860			if npiv == 2 && i1 != 1 {
861				swap_self_adjoint(A.rb_mut(), 1, i1);
862				perm::swap_rows_idx(L_prev.rb_mut(), 1, i1);
863			}
864			if nothing_to_do {
865				break 'next_iter;
866			}
867			if npiv == 1 {
868				let diag = A[(0, 0)].real();
869				let diag_inv = diag.recip();
870				subdiag[0] = zero();
871				let (_, _, L, A) = A.rb_mut().split_at_mut(1, 1);
872				let L = L.col_mut(0);
873				rank1_update(A, L, diag_inv);
874			} else {
875				let a00 = A[(0, 0)].real();
876				let a11 = A[(1, 1)].real();
877				let a10 = A[(1, 0)].copy();
878				subdiag[0] = a10.copy();
879				subdiag[1] = zero();
880				A[(1, 0)] = zero();
881				let d10 = a10.abs();
882				let d10_inv = &d10.recip();
883				let d00 = a00 * d10_inv;
884				let d11 = a11 * d10_inv;
885				let t = (&d00 * &d11 - one::<T::Real>()).recip();
886				let d10 = a10.mul_real(d10_inv);
887				let d = t * d10_inv;
888				let (_, _, L, A) = A.rb_mut().split_at_mut(2, 2);
889				let (L0, L1) = L.two_cols_mut(0, 1);
890				rank2_update(A, L0, L1, d, d00, d10, d11);
891			}
892		}
893		let offset = A_left.ncols();
894		if npiv == 2 {
895			pivots[k] = (offset + i0 + k) | TOP_BIT;
896			pivots[k + 1] = (offset + i1 + k) | TOP_BIT;
897		} else {
898			pivots[k] = offset + i0 + k;
899		}
900		k += npiv;
901	}
902}
903impl<T: ComplexField> Auto<T> for LbltParams {
904	fn auto() -> Self {
905		Self {
906			pivoting: PivotingStrategy::PartialDiag,
907			block_size: 64,
908			par_threshold: 128 * 128,
909			non_exhaustive: NonExhaustive(()),
910		}
911	}
912}
913pub fn rank2_update<'a, T: ComplexField>(
914	mut A: MatMut<'a, T>,
915	mut L0: ColMut<'a, T>,
916	mut L1: ColMut<'a, T>,
917	d: T::Real,
918	d00: T::Real,
919	d10: T,
920	d11: T::Real,
921) {
922	if const { T::SIMD_CAPABILITIES.is_simd() } {
923		if let (Some(A), Some(L0), Some(L1)) = (
924			A.rb_mut().try_as_col_major_mut(),
925			L0.rb_mut().try_as_col_major_mut(),
926			L1.rb_mut().try_as_col_major_mut(),
927		) {
928			rank2_update_simd(A, L0, L1, d, d00, d10, d11);
929		} else {
930			rank2_update_fallback(A, L0, L1, d, d00, d10, d11);
931		}
932	} else {
933		rank2_update_fallback(A, L0, L1, d, d00, d10, d11);
934	}
935}
936pub fn rank2_update_simd<'a, T: ComplexField>(
937	A: MatMut<'a, T, usize, usize, ContiguousFwd>,
938	L0: ColMut<'a, T, usize, ContiguousFwd>,
939	L1: ColMut<'a, T, usize, ContiguousFwd>,
940	d: T::Real,
941	d00: T::Real,
942	d10: T,
943	d11: T::Real,
944) {
945	struct Impl<'a, T: ComplexField> {
946		A: MatMut<'a, T, usize, usize, ContiguousFwd>,
947		L0: ColMut<'a, T, usize, ContiguousFwd>,
948		L1: ColMut<'a, T, usize, ContiguousFwd>,
949		d: T::Real,
950		d00: T::Real,
951		d10: T,
952		d11: T::Real,
953	}
954	impl<T: ComplexField> pulp::WithSimd for Impl<'_, T> {
955		type Output = ();
956
957		#[inline(always)]
958		fn with_simd<S: pulp::Simd>(self, simd: S) {
959			let Self {
960				mut A,
961				mut L0,
962				mut L1,
963				d,
964				d00,
965				d10,
966				d11,
967			} = self;
968			let n = A.nrows();
969			for j in 0..n {
970				let x0 = &L0[j];
971				let x1 = &L1[j];
972				let w0 = (x0.mul_real(&d11) - x1 * &d10).mul_real(&d);
973				let w1 = (x1.mul_real(&d00) - x0 * d10.conj()).mul_real(&d);
974				with_dim!({
975					let subrange_len = n - j;
976				});
977				{
978					let mut A = A
979						.rb_mut()
980						.get_mut(j.., j)
981						.as_row_shape_mut(subrange_len);
982					let L0 = L0.rb().get(j..).as_row_shape(subrange_len);
983					let L1 = L1.rb().get(j..).as_row_shape(subrange_len);
984					let simd =
985						SimdCtx::<T, S>::new(T::simd_ctx(simd), subrange_len);
986					let w0_conj = w0.conj();
987					let w1_conj = w1.conj();
988					let w0_conj_neg = -w0_conj;
989					let w1_conj_neg = -w1_conj;
990					let w0_splat = simd.splat(&w0_conj_neg);
991					let w1_splat = simd.splat(&w1_conj_neg);
992					simd_iter!(for i in [simd.indices()] {
993						let mut acc = simd.read(A.rb(), i);
994						let l0_val = simd.read(L0, i);
995						let l1_val = simd.read(L1, i);
996						acc = simd.mul_add(l0_val, w0_splat, acc);
997						acc = simd.mul_add(l1_val, w1_splat, acc);
998						simd.write(A.rb_mut(), i, acc);
999					});
1000				}
1001				A[(j, j)] = A[(j, j)].real().to_cplx();
1002				L0[j] = w0;
1003				L1[j] = w1;
1004			}
1005		}
1006	}
1007	dispatch!(
1008		Impl {
1009			A,
1010			L0,
1011			L1,
1012			d,
1013			d00,
1014			d10,
1015			d11
1016		},
1017		Impl,
1018		T
1019	)
1020}
1021pub fn rank2_update_fallback<'a, T: ComplexField>(
1022	mut A: MatMut<'a, T>,
1023	mut L0: ColMut<'a, T>,
1024	mut L1: ColMut<'a, T>,
1025	d: T::Real,
1026	d00: T::Real,
1027	d10: T,
1028	d11: T::Real,
1029) {
1030	let n = A.nrows();
1031	for j in 0..n {
1032		let x0 = &L0[j].copy();
1033		let x1 = &L1[j].copy();
1034		let w0 = (x0.mul_real(&d11) - x1 * &d10).mul_real(&d);
1035		let w1 = (x1.mul_real(&d00) - x0 * d10.conj()).mul_real(&d);
1036		for i in j..n {
1037			A[(i, j)] = &A[(i, j)] - &L0[i] * w0.conj() - &L1[i] * w1.conj();
1038		}
1039		A[(j, j)] = A[(j, j)].as_real();
1040		L0[j] = w0;
1041		L1[j] = w1;
1042	}
1043}
1044pub fn rank1_update<'a, T: ComplexField>(
1045	mut A: MatMut<'a, T>,
1046	mut L0: ColMut<'a, T>,
1047	d: T::Real,
1048) {
1049	if const { T::SIMD_CAPABILITIES.is_simd() } {
1050		if let (Some(A), Some(L0)) = (
1051			A.rb_mut().try_as_col_major_mut(),
1052			L0.rb_mut().try_as_col_major_mut(),
1053		) {
1054			rank1_update_simd(A, L0, d);
1055		} else {
1056			rank1_update_fallback(A, L0, d);
1057		}
1058	} else {
1059		rank1_update_fallback(A, L0, d);
1060	}
1061}
1062pub fn rank1_update_simd<'a, T: ComplexField>(
1063	A: MatMut<'a, T, usize, usize, ContiguousFwd>,
1064	L0: ColMut<'a, T, usize, ContiguousFwd>,
1065	d: T::Real,
1066) {
1067	struct Impl<'a, T: ComplexField> {
1068		A: MatMut<'a, T, usize, usize, ContiguousFwd>,
1069		L0: ColMut<'a, T, usize, ContiguousFwd>,
1070		d: T::Real,
1071	}
1072	impl<T: ComplexField> pulp::WithSimd for Impl<'_, T> {
1073		type Output = ();
1074
1075		#[inline(always)]
1076		fn with_simd<S: pulp::Simd>(self, simd: S) {
1077			let Self { mut A, mut L0, d } = self;
1078			let n = A.nrows();
1079			for j in 0..n {
1080				let x0 = &L0[j];
1081				let w0 = x0.mul_real(&d);
1082				with_dim!({
1083					let subrange_len = n - j;
1084				});
1085				{
1086					let mut A = A
1087						.rb_mut()
1088						.get_mut(j.., j)
1089						.as_row_shape_mut(subrange_len);
1090					let L0 = L0.rb().get(j..).as_row_shape(subrange_len);
1091					let simd =
1092						SimdCtx::<T, S>::new(T::simd_ctx(simd), subrange_len);
1093					let w0_conj = w0.conj();
1094					let w0_conj_neg = -w0_conj;
1095					let w0_splat = simd.splat(&w0_conj_neg);
1096					simd_iter!(for i in [simd.indices()] {
1097						let mut acc = simd.read(A.rb(), i);
1098						let l0_val = simd.read(L0, i);
1099						acc = simd.mul_add(l0_val, w0_splat, acc);
1100						simd.write(A.rb_mut(), i, acc);
1101					});
1102				}
1103				A[(j, j)] = A[(j, j)].real().to_cplx();
1104				L0[j] = w0;
1105			}
1106		}
1107	}
1108	dispatch!(Impl { A, L0, d }, Impl, T)
1109}
1110pub fn rank1_update_fallback<'a, T: ComplexField>(
1111	mut A: MatMut<'a, T>,
1112	mut L0: ColMut<'a, T>,
1113	d: T::Real,
1114) {
1115	let n = A.nrows();
1116	for j in 0..n {
1117		let x0 = &L0[j];
1118		let w0 = x0.mul_real(&d);
1119		for i in j..n {
1120			A[(i, j)] = &A[(i, j)] - &L0[i] * w0.conj();
1121		}
1122		A[(j, j)] = A[(j, j)].real().to_cplx();
1123		L0[j] = w0;
1124	}
1125}
1126/// computes the layout of required workspace for performing an $LBL^\top$
1127/// decomposition
1128pub fn cholesky_in_place_scratch<I: Index, T: ComplexField>(
1129	dim: usize,
1130	par: Par,
1131	params: Spec<LbltParams, T>,
1132) -> StackReq {
1133	let params = params.config;
1134	let _ = par;
1135	let mut bs = params.block_size;
1136	if bs < 2 || dim <= bs {
1137		bs = 0;
1138	}
1139	StackReq::new::<usize>(dim).and(temp_mat_scratch::<T>(dim, bs))
1140}
1141/// info about the result of the $LBL^\top$ factorization
1142#[derive(Copy, Clone, Debug)]
1143pub struct LbltInfo {
1144	/// number of pivoting transpositions
1145	pub transposition_count: usize,
1146}
1147/// computes the $LBL^\top$ factorization of $A$ and stores the factorization in
1148/// `matrix` and `subdiag`
1149///
1150/// the diagonal of the block diagonal matrix is stored on the diagonal
1151/// of `matrix`, while the subdiagonal elements of the blocks are stored in
1152/// `subdiag`
1153///
1154/// # panics
1155///
1156/// panics if the input matrix is not square
1157///
1158/// this can also panic if the provided memory in `stack` is insufficient (see
1159/// [`cholesky_in_place_scratch`]).
1160#[track_caller]
1161pub fn cholesky_in_place<'out, I: Index, T: ComplexField>(
1162	A: MatMut<'_, T>,
1163	subdiag: DiagMut<'_, T>,
1164	perm: &'out mut [I],
1165	perm_inv: &'out mut [I],
1166	par: Par,
1167	stack: &mut MemStack,
1168	params: Spec<LbltParams, T>,
1169) -> (LbltInfo, PermRef<'out, I>) {
1170	let params = params.config;
1171	let truncate = <I::Signed as SignedIndex>::truncate;
1172	let n = A.nrows();
1173	assert!(all(
1174		A.nrows() == A.ncols(),
1175		subdiag.dim() == n,
1176		perm.len() == n,
1177		perm_inv.len() == n
1178	));
1179	#[cfg(feature = "perf-warn")]
1180	if A.row_stride().unsigned_abs() != 1 && crate::__perf_warn!(CHOLESKY_WARN)
1181	{
1182		if A.col_stride().unsigned_abs() == 1 {
1183			log::warn!(
1184				target : "faer_perf",
1185				"$LBL^\top$ decomposition prefers column-major
1186			matrix. Found row-major matrix."
1187			);
1188		} else {
1189			log::warn!(
1190				target : "faer_perf",
1191				"$LBL^\top$ decomposition prefers column-major
1192			matrix. Found matrix with generic strides."
1193			);
1194		}
1195	}
1196	let (mut pivots, stack) = stack.make_with::<usize>(n, |_| 0);
1197	let pivots = &mut *pivots;
1198	let mut bs = params.block_size;
1199	if bs < 2 || n <= bs {
1200		bs = 0;
1201	}
1202	let (rook, diagonal) = match params.pivoting {
1203		PivotingStrategy::Partial => (false, false),
1204		PivotingStrategy::PartialDiag => (false, true),
1205		PivotingStrategy::Rook => (true, false),
1206		PivotingStrategy::RookDiag => (true, true),
1207		_ => (false, false),
1208	};
1209	if params.pivoting == PivotingStrategy::Full {
1210		lblt_full_piv(A, subdiag, pivots, par, params);
1211	} else {
1212		lblt_blocked(A, subdiag, pivots, bs, rook, diagonal, par, stack);
1213	}
1214	for (i, p) in perm.iter_mut().enumerate() {
1215		*p = I::from_signed(truncate(i));
1216	}
1217	let mut transposition_count = 0usize;
1218	for i in 0..n {
1219		let p = pivots[i] & !TOP_BIT;
1220		if i != p {
1221			transposition_count += 1;
1222		}
1223		perm.swap(i, p);
1224	}
1225	for (i, &p) in perm.iter().enumerate() {
1226		perm_inv[p.to_signed().zx()] = I::from_signed(truncate(i));
1227	}
1228	(
1229		LbltInfo {
1230			transposition_count,
1231		},
1232		unsafe { PermRef::new_unchecked(perm, perm_inv, n) },
1233	)
1234}