1use crate::internal_prelude::*;
2use crate::{assert, perm};
3use linalg::matmul::triangular::BlockStructure;
4const TOP_BIT: usize = 1 << (usize::BITS - 1);
5#[derive(Copy, Clone, Debug, PartialEq, Eq)]
7#[non_exhaustive]
8pub enum PivotingStrategy {
9 #[deprecated]
11 Diagonal,
12 Partial,
14 PartialDiag,
17 Rook,
19 RookDiag,
22 Full,
24}
25#[derive(Copy, Clone, Debug)]
27pub struct LbltParams {
28 pub pivoting: PivotingStrategy,
30 pub block_size: usize,
32 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}
1126pub 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#[derive(Copy, Clone, Debug)]
1143pub struct LbltInfo {
1144 pub transposition_count: usize,
1146}
1147#[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}