1use crate::assert;
2use crate::internal_prelude::*;
3use crate::linalg::matmul::internal::*;
4use linalg::matmul::triangular::BlockStructure;
5use pulp::Simd;
6#[inline(always)]
7fn simd_cholesky_row_batch<'N, T: ComplexField, S: Simd>(
8 simd: T::SimdCtx<S>,
9 A: MatMut<'_, T, Dim<'N>, Dim<'N>, ContiguousFwd>,
10 D: RowMut<'_, T, Dim<'N>>,
11 start: IdxInc<'N>,
12 is_llt: bool,
13 regularize: bool,
14 eps: T::Real,
15 delta: T::Real,
16 signs: Option<&Array<'N, i8>>,
17) -> Result<usize, usize> {
18 let mut A = A;
19 let mut D = D;
20 let n = A.ncols();
21 with_dim!(TAIL, *n - *start);
22 let simd = SimdCtx::<T, S>::new_force_mask(simd, TAIL);
23 let (idx_head, indices, idx_tail) = simd.indices();
24 assert!(idx_head.is_none());
25 let Some(idx_tail) = idx_tail else { panic!() };
26 let mut count = 0usize;
27 for j in n.indices() {
28 with_dim!(LEFT, *j);
29 let (A_0, Aj) = A.rb_mut().split_at_col_mut(j.into());
30 let A_0 = A_0.as_col_shape(LEFT);
31 let A10 = A_0.subrows(start, TAIL);
32 let mut Aj = Aj.col_mut(0).subrows_mut(start, TAIL);
33 {
34 let D = D.rb().subcols(IdxInc::ZERO, LEFT);
35 let mut Aj = Aj.rb_mut();
36 let mut iter = indices.clone();
37 let i0 = iter.next();
38 let i1 = iter.next();
39 let i2 = iter.next();
40 match (i0, i1, i2) {
41 (None, None, None) => {
42 let mut Aij = simd.read(Aj.rb(), idx_tail);
43 for k in LEFT.indices() {
44 let Ak = A10.col(k);
45 let D = D[k].real();
46 let D = if is_llt { one() } else { D };
47 let Ajk = simd.splat(&A_0[(j, k)].conj().mul_real(-D));
48 let Aik = simd.read(Ak, idx_tail);
49 Aij = simd.mul_add(Ajk, Aik, Aij);
50 }
51 simd.write(Aj.rb_mut(), idx_tail, Aij);
52 },
53 (Some(i0), None, None) => {
54 let mut A0j = simd.read(Aj.rb(), i0);
55 let mut Aij = simd.read(Aj.rb(), idx_tail);
56 for k in LEFT.indices() {
57 let Ak = A10.col(k);
58 let D = D[k].real();
59 let D = if is_llt { one() } else { D };
60 let Ajk = simd.splat(&A_0[(j, k)].conj().mul_real(-D));
61 let A0k = simd.read(Ak, i0);
62 let Aik = simd.read(Ak, idx_tail);
63 A0j = simd.mul_add(Ajk, A0k, A0j);
64 Aij = simd.mul_add(Ajk, Aik, Aij);
65 }
66 simd.write(Aj.rb_mut(), i0, A0j);
67 simd.write(Aj.rb_mut(), idx_tail, Aij);
68 },
69 (Some(i0), Some(i1), None) => {
70 let mut A0j = simd.read(Aj.rb(), i0);
71 let mut A1j = simd.read(Aj.rb(), i1);
72 let mut Aij = simd.read(Aj.rb(), idx_tail);
73 for k in LEFT.indices() {
74 let Ak = A10.col(k);
75 let D = D[k].real();
76 let D = if is_llt { one() } else { D };
77 let Ajk = simd.splat(&A_0[(j, k)].conj().mul_real(-D));
78 let A0k = simd.read(Ak, i0);
79 let A1k = simd.read(Ak, i1);
80 let Aik = simd.read(Ak, idx_tail);
81 A0j = simd.mul_add(Ajk, A0k, A0j);
82 A1j = simd.mul_add(Ajk, A1k, A1j);
83 Aij = simd.mul_add(Ajk, Aik, Aij);
84 }
85 simd.write(Aj.rb_mut(), i0, A0j);
86 simd.write(Aj.rb_mut(), i1, A1j);
87 simd.write(Aj.rb_mut(), idx_tail, Aij);
88 },
89 (Some(i0), Some(i1), Some(i2)) => {
90 let mut A0j = simd.read(Aj.rb(), i0);
91 let mut A1j = simd.read(Aj.rb(), i1);
92 let mut A2j = simd.read(Aj.rb(), i2);
93 let mut Aij = simd.read(Aj.rb(), idx_tail);
94 for k in LEFT.indices() {
95 let Ak = A10.col(k);
96 let D = D[k].real();
97 let D = if is_llt { one() } else { D };
98 let Ajk = simd.splat(&A_0[(j, k)].conj().mul_real(-D));
99 let A0k = simd.read(Ak, i0);
100 let A1k = simd.read(Ak, i1);
101 let A2k = simd.read(Ak, i2);
102 let Aik = simd.read(Ak, idx_tail);
103 A0j = simd.mul_add(Ajk, A0k, A0j);
104 A1j = simd.mul_add(Ajk, A1k, A1j);
105 A2j = simd.mul_add(Ajk, A2k, A2j);
106 Aij = simd.mul_add(Ajk, Aik, Aij);
107 }
108 simd.write(Aj.rb_mut(), i0, A0j);
109 simd.write(Aj.rb_mut(), i1, A1j);
110 simd.write(Aj.rb_mut(), i2, A2j);
111 simd.write(Aj.rb_mut(), idx_tail, Aij);
112 },
113 _ => {
114 unreachable!();
115 },
116 }
117 }
118 let D = D.rb_mut().at_mut(j);
119 if *j >= *start {
120 let j_row = TAIL.idx(*j - *start);
121 let mut diag = Aj[j_row].real();
122 if regularize {
123 let sign = if is_llt {
124 1
125 } else {
126 if let Some(signs) = signs { signs[j] } else { 0 }
127 };
128 let small_or_negative = diag <= eps;
129 let minus_small_or_positive = diag >= -&eps;
130 if sign == 1 && small_or_negative {
131 diag = delta.copy();
132 count += 1;
133 } else if sign == -1i8 && minus_small_or_positive {
134 diag = -δ
135 } else {
136 if small_or_negative && minus_small_or_positive {
137 if diag < zero() {
138 diag = -δ
139 } else {
140 diag = delta.copy();
141 }
142 }
143 }
144 }
145 let j = j;
146 let diag = if is_llt {
147 if !(diag > zero()) {
148 *D = diag.to_cplx();
149 return Err(*j);
150 }
151 diag.sqrt()
152 } else {
153 diag.copy()
154 };
155 *D = diag.to_cplx();
156 if diag == zero() || !diag.is_finite() {
157 return Err(*j);
158 }
159 }
160 let diag = D.real();
161 {
162 let mut Aj = Aj.rb_mut();
163 let inv = simd.splat_real(&diag.recip());
164 for i in indices.clone() {
165 let mut Aij = simd.read(Aj.rb(), i);
166 Aij = simd.mul_real(Aij, inv);
167 simd.write(Aj.rb_mut(), i, Aij);
168 }
169 {
170 let mut Aij = simd.read(Aj.rb(), idx_tail);
171 Aij = simd.mul_real(Aij, inv);
172 simd.write(Aj.rb_mut(), idx_tail, Aij);
173 }
174 }
175 }
176 Ok(count)
177}
178#[inline(always)]
179fn simd_cholesky_matrix<T: ComplexField, S: Simd>(
180 simd: T::SimdCtx<S>,
181 A: MatMut<'_, T, usize, usize, ContiguousFwd>,
182 D: RowMut<'_, T, usize>,
183 is_llt: bool,
184 regularize: bool,
185 eps: T::Real,
186 delta: T::Real,
187 signs: Option<&[i8]>,
188) -> Result<usize, usize> {
189 let N = A.ncols();
190 let block_size =
191 4 * (core::mem::size_of::<T::SimdVec<S>>() / core::mem::size_of::<T>());
192 let mut A = A;
193 let mut D = D;
194 let mut count = 0;
195 let mut j = 0;
196 while j < N {
197 let block_size = Ord::min(block_size, N - j);
198 let j_next = j + block_size;
199 with_dim!(HEAD, j_next);
200 let A = A.rb_mut().submatrix_mut(0, 0, HEAD, HEAD);
201 let D = D.rb_mut().subcols_mut(0, HEAD);
202 let signs = signs.map(|signs| Array::from_ref(&signs[..*HEAD], HEAD));
203 count += simd_cholesky_row_batch(
204 simd,
205 A,
206 D,
207 HEAD.idx_inc(j),
208 is_llt,
209 regularize,
210 eps.clone(),
211 delta.clone(),
212 signs,
213 )?;
214 j += block_size;
215 }
216 Ok(count)
217}
218fn simd_cholesky<T: ComplexField>(
219 A: MatMut<'_, T>,
220 D: RowMut<'_, T>,
221 is_llt: bool,
222 regularize: bool,
223 eps: T::Real,
224 delta: T::Real,
225 signs: Option<&[i8]>,
226) -> Result<usize, usize> {
227 struct Impl<'a, T: ComplexField> {
228 A: MatMut<'a, T, usize, usize, ContiguousFwd>,
229 D: RowMut<'a, T>,
230 is_llt: bool,
231 regularize: bool,
232 eps: T::Real,
233 delta: T::Real,
234 signs: Option<&'a [i8]>,
235 }
236 impl<'a, T: ComplexField> pulp::WithSimd for Impl<'a, T> {
237 type Output = Result<usize, usize>;
238
239 #[inline(always)]
240 fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
241 let Self {
242 A,
243 D,
244 is_llt,
245 regularize,
246 eps,
247 delta,
248 signs,
249 } = self;
250 let simd = T::simd_ctx(simd);
251 if A.nrows() > 0 {
252 simd_cholesky_matrix(
253 simd, A, D, is_llt, regularize, eps, delta, signs,
254 )
255 } else {
256 Ok(0)
257 }
258 }
259 }
260 let mut A = A;
261 if const { T::SIMD_CAPABILITIES.is_simd() } {
262 if let Some(A) = A.rb_mut().try_as_col_major_mut() {
263 dispatch!(
264 Impl {
265 A,
266 D,
267 is_llt,
268 regularize,
269 eps,
270 delta,
271 signs,
272 },
273 Impl,
274 T
275 )
276 } else {
277 cholesky_fallback(
278 A,
279 D,
280 is_llt,
281 regularize,
282 eps.clone(),
283 delta.clone(),
284 signs,
285 )
286 }
287 } else {
288 cholesky_fallback(
289 A,
290 D,
291 is_llt,
292 regularize,
293 eps.clone(),
294 delta.clone(),
295 signs,
296 )
297 }
298}
299fn cholesky_fallback<T: ComplexField>(
300 A: MatMut<'_, T>,
301 D: RowMut<'_, T>,
302 is_llt: bool,
303 regularize: bool,
304 eps: T::Real,
305 delta: T::Real,
306 signs: Option<&[i8]>,
307) -> Result<usize, usize> {
308 let n = A.nrows();
309 let mut count = 0;
310 let mut A = A;
311 let mut D = D;
312 for j in 0..n {
313 for i in j..n {
314 let mut sum = zero::<T>();
315 for k in 0..j {
316 let D = D[k].real();
317 let D = if is_llt { one() } else { D };
318 sum = sum + A[(j, k)].conj() * A[(i, k)].mul_real(D);
319 }
320 A[(i, j)] = &A[(i, j)] - sum;
321 }
322 let D = D.rb_mut().at_mut(j);
323 let mut diag = A[(j, j)].real();
324 if regularize {
325 let sign = if is_llt {
326 1
327 } else {
328 if let Some(signs) = signs { signs[j] } else { 0 }
329 };
330 let small_or_negative = diag <= eps;
331 let minus_small_or_positive = diag >= -&eps;
332 if sign == 1 && small_or_negative {
333 diag = delta.copy();
334 count += 1;
335 } else if sign == -1i8 && minus_small_or_positive {
336 diag = -δ
337 } else {
338 if small_or_negative && minus_small_or_positive {
339 if diag < zero() {
340 diag = -δ
341 } else {
342 diag = delta.copy();
343 }
344 }
345 }
346 }
347 let diag = if is_llt {
348 if !(diag > zero()) {
349 *D = diag.to_cplx();
350 return Err(j);
351 }
352 diag.sqrt()
353 } else {
354 diag.copy()
355 };
356 *D = diag.to_cplx();
357 if diag == zero() || !diag.is_finite() {
358 return Err(j);
359 }
360 let inv = &diag.recip();
361 for i in j..n {
362 A[(i, j)] = A[(i, j)].mul_real(inv);
363 }
364 }
365 Ok(count)
366}
367pub(crate) fn cholesky_recursion_right_looking<T: ComplexField>(
368 A: MatMut<'_, T>,
369 D: RowMut<'_, T>,
370 recursion_threshold: usize,
371 block_size: usize,
372 is_llt: bool,
373 regularize: bool,
374 eps: &T::Real,
375 delta: &T::Real,
376 signs: Option<&[i8]>,
377 par: Par,
378) -> Result<usize, usize> {
379 let n = A.ncols();
380 if n <= recursion_threshold {
381 simd_cholesky(
382 A,
383 D,
384 is_llt,
385 regularize,
386 eps.clone(),
387 delta.clone(),
388 signs,
389 )
390 } else {
391 let mut count = 0;
392 let block_size = Ord::min(n.next_power_of_two() / 2, block_size);
393 let mut A = A;
394 let mut D = D;
395 let mut j = 0;
396 while j < n {
397 let block_size = Ord::min(block_size, n - j);
398 let (mut A00, A01, mut A10, mut A11) = A
399 .rb_mut()
400 .get_mut(j.., j..)
401 .split_at_mut(block_size, block_size);
402 let mut D0 = D.rb_mut().subcols_mut(j, block_size);
403 let mut L10xD0 = A01.transpose_mut();
404 let signs = signs.map(|signs| &signs[j..][..block_size]);
405 match cholesky_recursion_right_looking(
406 A00.rb_mut(),
407 D0.rb_mut(),
408 recursion_threshold,
409 block_size,
410 is_llt,
411 regularize,
412 eps,
413 delta,
414 signs,
415 par,
416 ) {
417 Ok(local_count) => count += local_count,
418 Err(fail_idx) => return Err(j + fail_idx),
419 }
420 let A00 = A00.rb();
421 if is_llt {
422 linalg::triangular_solve::solve_lower_triangular_in_place(
423 A00.conjugate(),
424 A10.rb_mut().transpose_mut(),
425 par,
426 )
427 } else {
428 linalg::triangular_solve::solve_unit_lower_triangular_in_place(
429 A00.conjugate(),
430 A10.rb_mut().transpose_mut(),
431 par,
432 )
433 }
434 let mut A10 = A10.rb_mut();
435 if is_llt {
436 linalg::matmul::triangular::matmul(
437 A11.rb_mut(),
438 BlockStructure::TriangularLower,
439 Accum::Add,
440 A10.rb(),
441 BlockStructure::Rectangular,
442 A10.rb().adjoint(),
443 BlockStructure::Rectangular,
444 -one::<T>(),
445 par,
446 );
447 } else {
448 if has_spicy_matmul::<T>() {
449 for k in 0..block_size {
450 let d = &D0[k].real().recip();
451 for i in j + block_size..n {
452 let i = i - (j + block_size);
453 A10[(i, k)] = A10[(i, k)].mul_real(d);
454 }
455 }
456 spicy_matmul::<usize, T>(
457 A11.rb_mut(),
458 BlockStructure::TriangularLower,
459 None,
460 None,
461 Accum::Add,
462 A10.rb(),
463 Conj::No,
464 A10.rb().transpose(),
465 Conj::Yes,
466 Some(D0.rb().transpose().as_diagonal()),
467 -one::<T>(),
468 par,
469 MemStack::new(&mut []),
470 );
471 } else {
472 for k in 0..block_size {
473 let d = &D0[k].real().recip();
474 for i in j + block_size..n {
475 let i = i - (j + block_size);
476 let a = A10[(i, k)].copy();
477 A10[(i, k)] = A10[(i, k)].mul_real(d);
478 L10xD0[(i, k)] = a;
479 }
480 }
481 linalg::matmul::triangular::matmul(
482 A11.rb_mut(),
483 BlockStructure::TriangularLower,
484 Accum::Add,
485 A10,
486 BlockStructure::Rectangular,
487 L10xD0.adjoint(),
488 BlockStructure::Rectangular,
489 -one::<T>(),
490 par,
491 );
492 }
493 };
494 j += block_size;
495 }
496 Ok(count)
497 }
498}
499pub(crate) fn cholesky_block_left_looking<T: ComplexField>(
500 A: MatMut<'_, T>,
501 D: RowMut<'_, T>,
502 right_looking_threshold: usize,
503 recursion_threshold: usize,
504 block_size: usize,
505 is_llt: bool,
506 regularize: bool,
507 eps: &T::Real,
508 delta: &T::Real,
509 signs: Option<&[i8]>,
510 par: Par,
511) -> Result<usize, usize> {
512 let n = A.nrows();
513 let cholesky =
514 |A: MatMut<'_, T>, D: RowMut<'_, T>, signs: Option<&[i8]>| {
515 cholesky_recursion_right_looking(
516 A,
517 D,
518 recursion_threshold,
519 block_size,
520 is_llt,
521 regularize,
522 eps,
523 delta,
524 signs,
525 par,
526 )
527 };
528 if true || n < right_looking_threshold {
529 cholesky(A, D, signs)
530 } else {
531 let mut A = A;
532 let mut D = D;
533 let mut count = 0;
534 let mut j = 0;
535 while j < n {
536 let bj = Ord::min(n - j, right_looking_threshold);
537 let (_, A01, AL, mut AR) = A.rb_mut().split_at_mut(j, j);
538 let AL = AL.rb();
539 let mut AL0xD0 = A01.get_mut(.., ..bj).transpose_mut();
540 let (AL0, AL1) = AL.split_at_row(bj);
541 let (mut AR0, mut AR1) =
542 AR.rb_mut().get_mut(.., ..bj).split_at_row_mut(bj);
543 let (D0, D1) = D.rb_mut().split_at_col_mut(j);
544 let D0 = D0.rb();
545 let mut D1 = D1.get_mut(..bj);
546 if is_llt {
547 linalg::matmul::triangular::matmul(
548 AR0.rb_mut(),
549 BlockStructure::TriangularLower,
550 Accum::Add,
551 AL0.rb(),
552 BlockStructure::Rectangular,
553 AL0.rb().adjoint(),
554 BlockStructure::Rectangular,
555 -one::<T>(),
556 par,
557 );
558 linalg::matmul::matmul(
559 AR1.rb_mut(),
560 Accum::Add,
561 AL1.rb(),
562 AL0.rb().adjoint(),
563 -one::<T>(),
564 par,
565 );
566 } else {
567 if has_spicy_matmul::<T>() {
568 spicy_matmul::<usize, T>(
569 AR0.rb_mut(),
570 BlockStructure::TriangularLower,
571 None,
572 None,
573 Accum::Add,
574 AL0.rb(),
575 Conj::No,
576 AL0.rb().transpose(),
577 Conj::Yes,
578 Some(D0.rb().transpose().as_diagonal()),
579 -one::<T>(),
580 par,
581 MemStack::new(&mut []),
582 );
583 spicy_matmul::<usize, T>(
584 AR1.rb_mut(),
585 BlockStructure::Rectangular,
586 None,
587 None,
588 Accum::Add,
589 AL1.rb(),
590 Conj::No,
591 AL0.rb().transpose(),
592 Conj::Yes,
593 Some(D0.rb().transpose().as_diagonal()),
594 -one::<T>(),
595 par,
596 MemStack::new(&mut []),
597 );
598 } else {
599 for i in 0..j {
600 let d = &D0[i].real();
601 for k in 0..bj {
602 AL0xD0[(k, i)] = AL0[(k, i)].mul_real(d);
603 }
604 }
605 let AL0xD0 = AL0xD0.rb();
606 linalg::matmul::triangular::matmul(
607 AR0.rb_mut(),
608 BlockStructure::TriangularLower,
609 Accum::Add,
610 AL0,
611 BlockStructure::Rectangular,
612 AL0xD0.adjoint(),
613 BlockStructure::Rectangular,
614 -one::<T>(),
615 par,
616 );
617 linalg::matmul::matmul(
618 AR1.rb_mut(),
619 Accum::Add,
620 AL1,
621 AL0xD0.adjoint(),
622 -one::<T>(),
623 par,
624 );
625 }
626 };
627 match cholesky(
628 AR0.rb_mut(),
629 D1.rb_mut(),
630 signs.map(|signs| &signs[j..][..bj]),
631 ) {
632 Ok(local_count) => count += local_count,
633 Err(fail_idx) => return Err(j + fail_idx),
634 }
635 if is_llt {
636 linalg::triangular_solve::solve_lower_triangular_in_place(
637 AR0.conjugate(),
638 AR1.rb_mut().transpose_mut(),
639 par,
640 );
641 } else {
642 linalg::triangular_solve::solve_unit_lower_triangular_in_place(
643 AR0.conjugate(),
644 AR1.rb_mut().transpose_mut(),
645 par,
646 );
647 for k in 0..bj {
648 let d = &D1[k].real().recip();
649 for i in j + bj..n {
650 let i = i - (j + bj);
651 AR1[(i, k)] = AR1[(i, k)].mul_real(d);
652 }
653 }
654 }
655 j += bj;
656 }
657 Ok(count)
658 }
659}
660#[derive(Copy, Clone, Debug)]
664pub struct LdltRegularization<'a, T> {
665 pub dynamic_regularization_signs: Option<&'a [i8]>,
667 pub dynamic_regularization_delta: T,
669 pub dynamic_regularization_epsilon: T,
671}
672#[derive(Copy, Clone, Debug)]
674pub struct LdltInfo {
675 pub dynamic_regularization_count: usize,
677}
678#[derive(Copy, Clone, Debug)]
680pub enum LdltError {
681 ZeroPivot { index: usize },
682}
683impl core::fmt::Display for LdltError {
684 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
685 core::fmt::Debug::fmt(self, f)
686 }
687}
688impl core::error::Error for LdltError {}
689impl<T: RealField> Default for LdltRegularization<'_, T> {
690 fn default() -> Self {
691 Self {
692 dynamic_regularization_signs: None,
693 dynamic_regularization_delta: zero(),
694 dynamic_regularization_epsilon: zero(),
695 }
696 }
697}
698#[derive(Copy, Clone, Debug)]
699pub struct LdltParams {
700 pub recursion_threshold: usize,
701 pub block_size: usize,
702 #[doc(hidden)]
703 pub non_exhaustive: NonExhaustive,
704}
705impl<T: ComplexField> Auto<T> for LdltParams {
706 #[inline]
707 fn auto() -> Self {
708 Self {
709 recursion_threshold: 64,
710 block_size: 128,
711 non_exhaustive: NonExhaustive(()),
712 }
713 }
714}
715#[inline]
716pub fn cholesky_in_place_scratch<T: ComplexField>(
717 dim: usize,
718 par: Par,
719 params: Spec<LdltParams, T>,
720) -> StackReq {
721 _ = par;
722 _ = params;
723 temp_mat_scratch::<T>(dim, 1)
724}
725pub fn cholesky_in_place<T: ComplexField>(
726 A: MatMut<'_, T>,
727 regularization: LdltRegularization<'_, T::Real>,
728 par: Par,
729 stack: &mut MemStack,
730 params: Spec<LdltParams, T>,
731) -> Result<LdltInfo, LdltError> {
732 let params = params.config;
733 let n = A.nrows();
734 let mut D = unsafe { temp_mat_uninit(n, 1, stack).0 };
735 let D = D.as_mat_mut();
736 let mut D = D.col_mut(0).transpose_mut();
737 let mut A = A;
738 let ret = match cholesky_block_left_looking(
739 A.rb_mut(),
740 D.rb_mut(),
741 params.block_size,
742 params.recursion_threshold,
743 params.block_size,
744 false,
745 regularization.dynamic_regularization_delta > zero()
746 && regularization.dynamic_regularization_epsilon > zero(),
747 ®ularization.dynamic_regularization_epsilon,
748 ®ularization.dynamic_regularization_delta,
749 regularization
750 .dynamic_regularization_signs
751 .map(|signs| signs),
752 par,
753 ) {
754 Ok(count) => Ok(LdltInfo {
755 dynamic_regularization_count: count,
756 }),
757 Err(index) => Err(LdltError::ZeroPivot { index }),
758 };
759 let init = if let Err(LdltError::ZeroPivot { index }) = ret {
760 index + 1
761 } else {
762 n
763 };
764 for i in 0..init {
765 A[(i, i)] = D[i].copy();
766 }
767 ret
768}
769#[cfg(test)]
770mod tests {
771 use super::*;
772 use crate::stats::prelude::*;
773 use crate::utils::approx::*;
774 use crate::{Mat, Row, assert, c64};
775 #[test]
776 fn test_simd_cholesky() {
777 let rng = &mut StdRng::seed_from_u64(0);
778 type T = c64;
779 for n in 0..=64 {
780 for f in [cholesky_fallback::<T>, simd_cholesky::<T>] {
781 for llt in [true, false] {
782 let approx_eq = CwiseMat(ApproxEq {
783 abs_tol: 1e-12,
784 rel_tol: 1e-12,
785 });
786 let A = CwiseMatDistribution {
787 nrows: n,
788 ncols: n,
789 dist: ComplexDistribution::new(
790 StandardNormal,
791 StandardNormal,
792 ),
793 }
794 .rand::<Mat<c64>>(rng);
795 let A = &A * &A.adjoint();
796 let A = A.as_ref().as_shape(n, n);
797 let mut L = A.cloned();
798 let mut L = L.as_mut();
799 let mut D = Row::zeros(n);
800 let mut D = D.as_mut();
801 f(L.rb_mut(), D.rb_mut(), llt, false, 0.0, 0.0, None)
802 .unwrap();
803 for j in 0..n {
804 for i in 0..j {
805 L[(i, j)] = c64::ZERO;
806 }
807 }
808 let L = L.rb().as_dyn_stride();
809 if llt {
810 assert!(L * L.adjoint() ~ A);
811 } else {
812 assert!(L * D.as_diagonal() * L.adjoint() ~ A);
813 };
814 }
815 }
816 }
817 }
818 #[test]
819 fn test_cholesky() {
820 let rng = &mut StdRng::seed_from_u64(0);
821 for n in [2, 4, 8, 31, 127, 240] {
822 for llt in [false, true] {
823 let approx_eq = CwiseMat(ApproxEq {
824 abs_tol: 1e-12,
825 rel_tol: 1e-12,
826 });
827 let A = CwiseMatDistribution {
828 nrows: n,
829 ncols: n,
830 dist: ComplexDistribution::new(
831 StandardNormal,
832 StandardNormal,
833 ),
834 }
835 .rand::<Mat<c64>>(rng);
836 let A = &A * &A.adjoint();
837 let A = A.as_ref();
838 let mut L = A.cloned();
839 let mut L = L.as_mut();
840 let mut D = Row::zeros(n);
841 let mut D = D.as_mut();
842 cholesky_recursion_right_looking(
843 L.rb_mut(),
844 D.rb_mut(),
845 32,
846 32,
847 llt,
848 false,
849 &0.0,
850 &0.0,
851 None,
852 Par::Seq,
853 )
854 .unwrap();
855 for j in 0..n {
856 for i in 0..j {
857 L[(i, j)] = c64::ZERO;
858 }
859 }
860 let L = L.rb().as_dyn_stride();
861 if llt {
862 assert!(L * L.adjoint() ~ A);
863 } else {
864 assert!(L * D.as_diagonal() * L.adjoint() ~ A);
865 };
866 }
867 }
868 }
869}