1use crate::internal_prelude::*;
2use crate::utils;
3use equator::assert;
4use linalg::matmul::{matmul, triangular as tr};
5pub fn generalized_hessenberg_scratch<T: ComplexField>(
8 n: usize,
9 params: GeneralizedHessenbergParams,
10) -> StackReq {
11 if params.block_size <= 1 || n <= params.block_size {
12 StackReq::EMPTY
13 } else {
14 linalg::temp_mat_scratch::<T>(
15 n,
16 match params.block_size.checked_mul(6) {
17 Some(n) => n,
18 None => return StackReq::OVERFLOW,
19 },
20 )
21 }
22}
23pub(crate) fn make_givens<T: ComplexField>(f: T, g: T) -> (T::Real, T, T) {
24 if g == zero() {
25 (one(), zero(), f)
26 } else if f == zero() {
27 let c = zero::<T::Real>();
28 let d = g.abs();
29 let d_inv = d.recip();
30 let s = g.conj().mul_real(&d_inv);
31 let r = d.to_cplx();
32 (c, s, r)
33 } else {
34 let f1 = f.abs();
35 let g1 = g.abs();
36 let h1 = f1.hypot(g1);
37 let ref h1_inv = h1.recip();
38 let c = f1 * h1_inv;
39 let r = f.mul_real(c.recip());
40 let s = g.conj() * r.mul_real(h1_inv).mul_real(h1_inv);
41 (c, s, r)
42 }
43}
44pub(crate) fn rot<T: ComplexField>(
45 c: T::Real,
46 s: T,
47 x: RowMut<'_, T>,
48 y: RowMut<'_, T>,
49) {
50 let (c, s) = &(c, s);
51 zip!(x, y).for_each(|unzip!(x, y): Zip!(&mut _, &mut _)| {
52 (*x, *y) = (x.mul_real(c) + &*y * s, y.mul_real(c) - &*x * s.conj());
53 });
54}
55#[derive(Copy, Clone, Debug)]
57pub struct GeneralizedHessenbergParams {
58 pub block_size: usize,
60 pub blocking_threshold: usize,
62 #[doc(hidden)]
63 pub non_exhaustive: NonExhaustive,
64}
65impl<T: ComplexField> Auto<T> for GeneralizedHessenbergParams {
66 fn auto() -> Self {
67 Self {
68 block_size: 32,
69 blocking_threshold: 256 * 256,
70 non_exhaustive: NonExhaustive(()),
71 }
72 }
73}
74pub(crate) fn trot<T: ComplexField>(
75 c: T::Real,
76 s: T,
77 x: ColMut<'_, T>,
78 y: ColMut<'_, T>,
79) {
80 rot(c, -s.conj(), x.transpose_mut(), y.transpose_mut());
81}
82fn generalized_hessenberg_unblocked<T: ComplexField>(
83 A: MatMut<'_, T>,
84 B: MatMut<'_, T>,
85 Q: Option<MatMut<'_, T>>,
86 Z: Option<MatMut<'_, T>>,
87 Q_is_I: bool,
88 Z_is_I: bool,
89 par: Par,
90 stack: &mut MemStack,
91 params: GeneralizedHessenbergParams,
92) {
93 _ = params;
94 _ = stack;
95 _ = Q_is_I;
96 _ = Z_is_I;
97 _ = par;
98 let mut A = A;
99 let mut B = B;
100 let mut Q = Q;
101 let mut Z = Z;
102 let n = A.nrows();
103 let (Q_nrows, Q_ncols) =
104 Q.rb().map(|Q| (Q.nrows(), Q.ncols())).unwrap_or((n, n));
105 let (Z_nrows, Z_ncols) =
106 Z.rb().map(|Z| (Z.nrows(), Z.ncols())).unwrap_or((n, n));
107 assert!(all(
108 A.nrows() == n,
109 A.ncols() == n,
110 B.nrows() == n,
111 B.ncols() == n,
112 Q_nrows == n,
113 Q_ncols == n,
114 Z_nrows == n,
115 Z_ncols == n,
116 ));
117 if n <= 2 {
118 return;
119 }
120 for jcol in 0..n - 2 {
121 for jrow in (jcol + 2..n).rev() {
122 let (c, s, r) =
123 make_givens(A[(jrow - 1, jcol)].copy(), A[(jrow, jcol)].copy());
124 A[(jrow - 1, jcol)] = r;
125 A[(jrow, jcol)] = zero();
126 let (x, y) = A
127 .rb_mut()
128 .get_mut(.., jcol + 1..)
129 .two_rows_mut(jrow - 1, jrow);
130 rot(c.copy(), s.copy(), x, y);
131 let (x, y) = B
132 .rb_mut()
133 .get_mut(.., jrow - 1..)
134 .two_rows_mut(jrow - 1, jrow);
135 rot(c.copy(), s.copy(), x, y);
136 if let Some(mut Q) = Q.rb_mut() {
137 let (x, y) = Q.rb_mut().two_cols_mut(jrow - 1, jrow);
138 rot(c.copy(), s.conj(), x.transpose_mut(), y.transpose_mut());
139 }
140 let (c, s, r) =
141 make_givens(B[(jrow, jrow)].copy(), B[(jrow, jrow - 1)].copy());
142 B[(jrow, jrow)] = r;
143 B[(jrow, jrow - 1)] = zero();
144 let (x, y) = A.rb_mut().two_cols_mut(jrow, jrow - 1);
145 rot(c.copy(), s.copy(), x.transpose_mut(), y.transpose_mut());
146 let (x, y) =
147 B.rb_mut().get_mut(..jrow, ..).two_cols_mut(jrow, jrow - 1);
148 rot(c.copy(), s.copy(), x.transpose_mut(), y.transpose_mut());
149 if let Some(mut Z) = Z.rb_mut() {
150 let (x, y) = Z.rb_mut().two_cols_mut(jrow, jrow - 1);
151 rot(c, s, x.transpose_mut(), y.transpose_mut());
152 }
153 }
154 }
155}
156fn apply_U_from_the_right<T: ComplexField>(
157 M: MatMut<'_, T>,
158 U: MatRef<'_, T>,
159 par: Par,
160 stack: &mut MemStack,
161) {
162 let n = U.nrows() / 2;
163 let (U11, U12, U21, U22) = U.split_at(n, n);
164 let mut M = M;
165 let (mut tmp, _) = unsafe {
166 linalg::temp_mat_uninit::<T, _, _>(M.nrows(), M.ncols(), stack)
167 };
168 let mut tmp = tmp.as_mat_mut();
169 let (mut T1, mut T2) = tmp.rb_mut().split_at_col_mut(n);
170 let (M1, M2) = M.rb().split_at_col(n);
171 utils::thread::join_raw(
172 |mut par| {
173 matmul(T1.rb_mut(), Accum::Replace, M1.rb(), U11, one(), par.reborrow());
174 tr::matmul(
175 T1.rb_mut(),
176 tr::BlockStructure::Rectangular,
177 Accum::Add,
178 M2.rb(),
179 tr::BlockStructure::Rectangular,
180 U21,
181 tr::BlockStructure::TriangularUpper,
182 one(),
183 par,
184 );
185 },
186 |mut par| {
187 tr::matmul(
188 T2.rb_mut(),
189 tr::BlockStructure::Rectangular,
190 Accum::Replace,
191 M1.rb(),
192 tr::BlockStructure::Rectangular,
193 U12,
194 tr::BlockStructure::TriangularLower,
195 one(),
196 par.reborrow(),
197 );
198 matmul(T2.rb_mut(), Accum::Add, M2.rb(), U22, one(), par);
199 },
200 par,
201 );
202 M.copy_from(tmp);
203}
204pub fn generalized_hessenberg<T: ComplexField>(
224 A: MatMut<'_, T>,
225 B: MatMut<'_, T>,
226 Q_inout: Option<MatMut<'_, T>>,
227 Z_inout: Option<MatMut<'_, T>>,
228 mut par: Par,
229 stack: &mut MemStack,
230 params: GeneralizedHessenbergParams,
231) {
232 let n = A.nrows();
233 let (Q_nrows, Q_ncols) = Q_inout
234 .rb()
235 .map(|Q| (Q.nrows(), Q.ncols()))
236 .unwrap_or((n, n));
237 let (Z_nrows, Z_ncols) = Z_inout
238 .rb()
239 .map(|Z| (Z.nrows(), Z.ncols()))
240 .unwrap_or((n, n));
241 let Q_is_I = Q_inout.rb().is_some_and(|Q| Q.is_identity());
242 let Z_is_I = Z_inout.rb().is_some_and(|Z| Z.is_identity());
243 assert!(all(
244 A.nrows() == n,
245 A.ncols() == n,
246 B.nrows() == n,
247 B.ncols() == n,
248 Q_nrows == n,
249 Q_ncols == n,
250 Z_nrows == n,
251 Z_ncols == n,
252 ));
253 if n <= 2 {
254 return;
255 }
256 if params.block_size <= 1 || A.nrows() <= params.block_size {
257 return generalized_hessenberg_unblocked(
258 A, B, Q_inout, Z_inout, Q_is_I, Z_is_I, par, stack, params,
259 );
260 }
261 let mut A = A;
262 let mut B = B;
263 let mut Q = Q_inout;
264 let mut Z = Z_inout;
265 let mut jcol = 0;
266 while jcol < n - 2 {
267 let nnb = params.block_size.clamp(1, n - 2 - jcol);
268 let top = if jcol < 2 { 0 } else { jcol + 1 };
269 let n2nb = ((n - 2 - jcol) / nnb).saturating_sub(1);
270 let nblst = (n - 1 - jcol) - n2nb * nnb;
271 alloca!('stack: {
272 let mut work0 = unsafe { mat![uninit::<T>, nblst, nblst] };
273 let mut work1 =
274 unsafe { mat![uninit::<T>, 2 * nnb, 2 * nnb * n2nb] };
275 });
276 work0.fill(zero());
277 work0.rb_mut().diagonal_mut().fill(one());
278 work1.fill(zero());
279 for i in 0..n2nb {
280 work1
281 .rb_mut()
282 .subcols_mut(2 * nnb * i, 2 * nnb)
283 .diagonal_mut()
284 .fill(one());
285 }
286 for j in jcol..jcol + nnb {
287 for i in (j + 2..n).rev() {
288 let (c, s, r) =
289 make_givens(A[(i - 1, j)].copy(), A[(i, j)].copy());
290 A[(i - 1, j)] = r;
291 A[(i, j)] = c.to_cplx();
292 B[(i, j)] = s;
293 }
294 let mut len = 2 + j - jcol;
295 let jrow = j + n2nb * nnb + 2;
296 for i in (jrow..n).rev() {
297 let ref c = A[(i, j)].real();
298 let ref s = B[(i, j)].copy();
299 let col = nblst - (n - i) - 1;
300 for jj in nblst - len..nblst {
301 let ref temp0 = work0[(jj, col)].copy();
302 let ref temp1 = work0[(jj, col + 1)].copy();
303 work0[(jj, col + 1)] = temp1.mul_real(c) - s * temp0;
304 work0[(jj, col)] = temp0.mul_real(c) + s.conj() * temp1;
305 }
306 len += 1;
307 }
308 let mut j0 = jrow;
309 let mut idx = 0usize;
310 while j0 > j + 2 {
311 j0 -= nnb;
312 let mut start = nnb;
313 let mut len = 2 + j - jcol;
314 let mut col = nnb + j - jcol;
315 let mut U = work1.rb_mut().submatrix_mut(
316 0,
317 idx * 2 * nnb,
318 2 * nnb,
319 2 * nnb,
320 );
321 for i in (j0..j0 + nnb).rev() {
322 let ref c = A[(i, j)].real();
323 let ref s = B[(i, j)].copy();
324 col -= 1;
325 start -= 1;
326 for jj in start..start + len {
327 let ref temp0 = U[(jj, col)].copy();
328 let ref temp1 = U[(jj, col + 1)].copy();
329 U[(jj, col + 1)] = temp1.mul_real(c) - s * temp0;
330 U[(jj, col)] = s.conj() * temp1 + temp0.mul_real(c);
331 }
332 len += 1;
333 }
334 idx += 1;
335 }
336 for jj in (j + 1..n).rev() {
337 for i in (j + 2..Ord::min(n, jj + 2)).rev() {
338 let ref c = A[(i, j)].real();
339 let ref s = B[(i, j)].copy();
340 let ref temp0 = B[(i - 1, jj)].copy();
341 let ref temp1 = B[(i, jj)].copy();
342 B[(i, jj)] = temp1.mul_real(c) - s.conj() * temp0;
343 B[(i - 1, jj)] = temp0.mul_real(c) + s * temp1;
344 }
345 if jj + 1 < n {
346 let (c, s, r) = make_givens(
347 B[(jj + 1, jj + 1)].copy(),
348 B[(jj + 1, jj)].copy(),
349 );
350 B[(jj + 1, jj + 1)] = r;
351 B[(jj + 1, jj)] = zero();
352 let (bjj, bjj1) = B
353 .rb_mut()
354 .get_mut(top..jj + 1, ..)
355 .two_cols_mut(jj, jj + 1);
356 rot(
357 c.copy(),
358 s.copy(),
359 bjj1.transpose_mut(),
360 bjj.transpose_mut(),
361 );
362 A[(jj + 1, j)] = c.to_cplx();
363 B[(jj + 1, j)] = -s.conj();
364 }
365 }
366 let jj = (n - 2 - j) % 3;
367 let mut i = n - 1 - j;
368 while i > jj + 1 {
369 i -= 3;
370 let ref c0 = A[(j + i + 1, j)].real();
371 let ref s0 = -&B[(j + i + 1, j)];
372 let ref c1 = A[(j + i + 2, j)].real();
373 let ref s1 = -&B[(j + i + 2, j)];
374 let ref c2 = A[(j + i + 3, j)].real();
375 let ref s2 = -&B[(j + i + 3, j)];
376 for k in top..n {
377 let ref temp0 = A[(k, j + i + 0)].copy();
378 let ref temp1 = A[(k, j + i + 1)].copy();
379 let ref temp2 = A[(k, j + i + 2)].copy();
380 let ref temp3 = A[(k, j + i + 3)].copy();
381 A[(k, j + i + 3)] = temp3.mul_real(c2) + temp2 * s2.conj();
382 let temp2 = temp2.mul_real(c2) - temp3 * s2;
383 A[(k, j + i + 2)] = temp2.mul_real(c1) + temp1 * s1.conj();
384 let temp1 = temp1.mul_real(c1) - temp2 * s1;
385 A[(k, j + i + 1)] = temp1.mul_real(c0) + temp0 * s0.conj();
386 let temp0 = temp0.mul_real(c0) - temp1 * s0;
387 A[(k, j + i + 0)] = temp0;
388 }
389 }
390 for i in (1..jj + 1).rev() {
391 let c = A[(j + i + 1, j)].real();
392 let s = B[(j + i + 1, j)].copy();
393 let (aj1, aj) = A
394 .rb_mut()
395 .get_mut(top.., ..)
396 .two_cols_mut(j + i + 1, j + i);
397 trot(c, s, aj1, aj);
398 }
399 if j < jcol + nnb - 1 {
400 let len = j - jcol + 1;
401 let jrow = n - nblst;
402 {
403 let (mut work2, _) = unsafe {
404 linalg::temp_mat_uninit::<T, _, _>(nblst, 1, stack)
405 };
406 let mut work2 = work2.as_mat_mut().col_mut(0);
407 matmul(
408 work2.rb_mut().get_mut(..len),
409 Accum::Replace,
410 work0.rb().get(.., ..len).adjoint(),
411 A.rb().get(jrow.., j + 1),
412 one(),
413 par.reborrow(),
414 );
415 tr::matmul(
416 work2.rb_mut().get_mut(len..),
417 tr::BlockStructure::Rectangular,
418 Accum::Replace,
419 work0.rb().get(..nblst - len, len..).adjoint(),
420 tr::BlockStructure::TriangularUpper,
421 A.rb().get(jrow..jrow + nblst - len, j + 1),
422 tr::BlockStructure::Rectangular,
423 one(),
424 par.reborrow(),
425 );
426 matmul(
427 work2.rb_mut().get_mut(len..),
428 Accum::Add,
429 work0.rb().get(nblst - len.., len..).adjoint(),
430 A.rb().get(jrow + nblst - len.., j + 1),
431 one(),
432 par.reborrow(),
433 );
434 for i in jrow..jrow + nblst {
435 A[(i, j + 1)] = work2[i - jrow].copy();
436 }
437 }
438 let mut j0 = jrow;
439 let mut idx = 0usize;
440 while j0 > jcol + nnb {
441 j0 -= nnb;
442 let U = work1.rb().submatrix(
443 0,
444 idx * 2 * nnb,
445 nnb + len,
446 nnb + len,
447 );
448 let (U11, U12, U21, U22) = U.split_at(nnb, len);
449 let (mut work2, _) = unsafe {
450 linalg::temp_mat_uninit::<T, _, _>(nnb + len, 1, stack)
451 };
452 let mut work2 = work2.as_mat_mut().col_mut(0);
453 matmul(
454 work2.rb_mut().get_mut(..len),
455 Accum::Replace,
456 U11.adjoint(),
457 A.rb().get(j0..j0 + nnb, j + 1),
458 one(),
459 par.reborrow(),
460 );
461 tr::matmul(
462 work2.rb_mut().get_mut(..len),
463 tr::BlockStructure::Rectangular,
464 Accum::Add,
465 U21.adjoint(),
466 tr::BlockStructure::TriangularLower,
467 A.rb().get(j0 + nnb..j0 + nnb + len, j + 1),
468 tr::BlockStructure::Rectangular,
469 one(),
470 par.reborrow(),
471 );
472 tr::matmul(
473 work2.rb_mut().get_mut(len..),
474 tr::BlockStructure::Rectangular,
475 Accum::Replace,
476 U12.adjoint(),
477 tr::BlockStructure::TriangularUpper,
478 A.rb().get(j0..j0 + nnb, j + 1),
479 tr::BlockStructure::Rectangular,
480 one(),
481 par.reborrow(),
482 );
483 matmul(
484 work2.rb_mut().get_mut(len..),
485 Accum::Add,
486 U22.adjoint(),
487 A.rb().get(j0 + nnb..j0 + nnb + len, j + 1),
488 one(),
489 par.reborrow(),
490 );
491 for i in j0..j0 + len + nnb {
492 A[(i, j + 1)] = work2[i - j0].copy();
493 }
494 idx += 1;
495 }
496 }
497 }
498 let cola = n - jcol - nnb;
499 let j = n - nblst;
500 {
501 let (mut work2, _) = unsafe {
502 linalg::temp_mat_uninit::<T, _, _>(nblst, cola, stack)
503 };
504 let mut work2 = work2.as_mat_mut();
505 matmul(
506 work2.rb_mut(),
507 Accum::Replace,
508 work0.rb().adjoint(),
509 A.rb().get(j.., jcol + nnb..),
510 one(),
511 par.reborrow(),
512 );
513 A.rb_mut().get_mut(j.., jcol + nnb..).copy_from(&work2);
514 }
515 let mut j0 = j;
516 let mut idx = 0usize;
517 while j0 > jcol + nnb {
518 j0 -= nnb;
519 let U = work1.rb().submatrix(0, idx * 2 * nnb, 2 * nnb, 2 * nnb);
520 let (U11, U12, U21, U22) = U.split_at(nnb, nnb);
521 let (mut work2, _) = unsafe {
522 linalg::temp_mat_uninit::<T, _, _>(2 * nnb, cola, stack)
523 };
524 let mut work2 = work2.as_mat_mut();
525 let (mut work2_top, mut work2_bot) =
526 work2.rb_mut().split_at_row_mut(nnb);
527 utils::thread::join_raw(
528 |mut par| {
529 matmul(
530 work2_top.rb_mut(),
531 Accum::Replace,
532 U11.adjoint(),
533 A.rb().get(j0..j0 + nnb, jcol + nnb..),
534 one(),
535 par.reborrow(),
536 );
537 tr::matmul(
538 work2_top.rb_mut(),
539 tr::BlockStructure::Rectangular,
540 Accum::Add,
541 U21.adjoint(),
542 tr::BlockStructure::TriangularLower,
543 A.rb().get(j0 + nnb..j0 + 2 * nnb, jcol + nnb..),
544 tr::BlockStructure::Rectangular,
545 one(),
546 par,
547 );
548 },
549 |mut par| {
550 tr::matmul(
551 work2_bot.rb_mut(),
552 tr::BlockStructure::Rectangular,
553 Accum::Replace,
554 U12.adjoint(),
555 tr::BlockStructure::TriangularUpper,
556 A.rb().get(j0..j0 + nnb, jcol + nnb..),
557 tr::BlockStructure::Rectangular,
558 one(),
559 par.reborrow(),
560 );
561 matmul(
562 work2_bot.rb_mut(),
563 Accum::Add,
564 U22.adjoint(),
565 A.rb().get(j0 + nnb..j0 + 2 * nnb, jcol + nnb..),
566 one(),
567 par,
568 );
569 },
570 par.reborrow(),
571 );
572 A.rb_mut()
573 .get_mut(j0..j0 + 2 * nnb, jcol + nnb..)
574 .copy_from(&work2);
575 idx += 1;
576 }
577 if let Some(mut Q) = Q.rb_mut() {
578 let topq;
579 let nh;
580 if Q_is_I {
581 topq = Ord::max(1, j - jcol);
582 nh = n - topq;
583 } else {
584 topq = 0;
585 nh = n;
586 }
587 {
588 let (mut work2, _) = unsafe {
589 linalg::temp_mat_uninit::<T, _, _>(nh, nblst, stack)
590 };
591 let mut work2 = work2.as_mat_mut();
592 matmul(
593 work2.rb_mut(),
594 Accum::Replace,
595 Q.rb().get(topq.., j..),
596 work0.rb(),
597 one(),
598 par.reborrow(),
599 );
600 Q.rb_mut().get_mut(topq.., j..).copy_from(&work2);
601 }
602 let mut j0 = j;
603 let mut idx = 0usize;
604 while j0 > jcol + nnb {
605 j0 -= nnb;
606 let topq;
607 if Q_is_I {
608 topq = Ord::max(1, j0 - jcol);
609 } else {
610 topq = 0;
611 }
612 let U =
613 work1.rb().submatrix(0, idx * 2 * nnb, 2 * nnb, 2 * nnb);
614 apply_U_from_the_right(
615 Q.rb_mut().get_mut(topq.., j0..j0 + 2 * nnb),
616 U,
617 par.reborrow(),
618 stack,
619 );
620 idx += 1;
621 }
622 }
623 if Z.is_some() || top > 0 {
624 work0.fill(zero());
625 work0.rb_mut().diagonal_mut().fill(one());
626 work1.fill(zero());
627 for i in 0..n2nb {
628 work1
629 .rb_mut()
630 .subcols_mut(2 * nnb * i, 2 * nnb)
631 .diagonal_mut()
632 .column_vector_mut()
633 .fill(one());
634 }
635 for j in jcol..jcol + nnb {
636 let mut len = 2 + j - jcol;
637 let jrow = j + n2nb * nnb + 2;
638 for i in (jrow..n).rev() {
639 let ref c = A[(i, j)].real();
640 let ref s = B[(i, j)].copy();
641 A[(i, j)] = zero();
642 B[(i, j)] = zero();
643 let col = nblst - (n - i) - 1;
644 for jj in nblst - len..nblst {
645 let ref temp0 = work0[(jj, col)].copy();
646 let ref temp1 = work0[(jj, col + 1)].copy();
647 work0[(jj, col + 1)] =
648 temp1.mul_real(c) - temp0 * s.conj();
649 work0[(jj, col)] = temp0.mul_real(c) + temp1 * s;
650 }
651 len += 1;
652 }
653 let mut j0 = jrow;
654 let mut idx = 0usize;
655 while j0 > j + 2 {
656 j0 -= nnb;
657 let mut len = 2 + j - jcol;
658 let mut col = nnb + j - jcol;
659 let mut start = nnb;
660 let mut U = work1.rb_mut().submatrix_mut(
661 0,
662 idx * 2 * nnb,
663 2 * nnb,
664 2 * nnb,
665 );
666 for i in (j0..j0 + nnb).rev() {
667 let ref c = A[(i, j)].real();
668 let ref s = B[(i, j)].copy();
669 A[(i, j)] = zero();
670 B[(i, j)] = zero();
671 col -= 1;
672 start -= 1;
673 for jj in start..start + len {
674 let ref temp0 = U[(jj, col)].copy();
675 let ref temp1 = U[(jj, col + 1)].copy();
676 U[(jj, col + 1)] =
677 temp1.mul_real(c) - temp0 * s.conj();
678 U[(jj, col)] = temp0.mul_real(c) + temp1 * s;
679 }
680 len += 1;
681 }
682 idx += 1;
683 }
684 }
685 } else {
686 zip!(A.rb_mut().get_mut(jcol + 2.., jcol..n - jcol - 2))
687 .for_each(|unzip!(x)| *x = zero());
688 zip!(B.rb_mut().get_mut(jcol + 2.., jcol..n - jcol - 2))
689 .for_each(|unzip!(x)| *x = zero());
690 }
691 if top > 0 {
692 for mut M in [A.rb_mut(), B.rb_mut()] {
693 let j = n - nblst;
694 {
695 let (mut work2, _) = unsafe {
696 linalg::temp_mat_uninit::<T, _, _>(top, nblst, stack)
697 };
698 let mut work2 = work2.as_mat_mut();
699 matmul(
700 work2.rb_mut(),
701 Accum::Replace,
702 M.rb().get(..top, j..),
703 work0.rb(),
704 one(),
705 par.reborrow(),
706 );
707 M.rb_mut().get_mut(..top, j..).copy_from(&work2);
708 }
709 let mut j0 = j;
710 let mut idx = 0usize;
711 while j0 > jcol + nnb {
712 j0 -= nnb;
713 let U = work1.rb().submatrix(
714 0,
715 idx * 2 * nnb,
716 2 * nnb,
717 2 * nnb,
718 );
719 apply_U_from_the_right(
720 M.rb_mut().get_mut(..top, j0..j0 + 2 * nnb),
721 U,
722 par.reborrow(),
723 stack,
724 );
725 idx += 1;
726 }
727 }
728 }
729 if let Some(mut Z) = Z.rb_mut() {
730 let topq;
731 let nh;
732 if Z_is_I {
733 topq = Ord::max(1, j - jcol);
734 nh = n - topq;
735 } else {
736 topq = 0;
737 nh = n;
738 }
739 {
740 let (mut work2, _) = unsafe {
741 linalg::temp_mat_uninit::<T, _, _>(nh, nblst, stack)
742 };
743 let mut work2 = work2.as_mat_mut();
744 matmul(
745 work2.rb_mut(),
746 Accum::Replace,
747 Z.rb().get(topq.., j..),
748 work0.rb(),
749 one(),
750 par.reborrow(),
751 );
752 Z.rb_mut().get_mut(topq.., j..).copy_from(&work2);
753 }
754 let mut j0 = j;
755 let mut idx = 0usize;
756 while j0 > jcol + nnb {
757 j0 -= nnb;
758 let topq;
759 if Z_is_I {
760 topq = Ord::max(1, j0 - jcol);
761 } else {
762 topq = 0;
763 }
764 let U =
765 work1.rb().submatrix(0, idx * 2 * nnb, 2 * nnb, 2 * nnb);
766 apply_U_from_the_right(
767 Z.rb_mut().get_mut(topq.., j0..j0 + 2 * nnb),
768 U,
769 par.reborrow(),
770 stack,
771 );
772 idx += 1;
773 }
774 }
775 jcol += nnb;
776 }
777}
778#[cfg(test)]
779mod tests {
780 use super::*;
781 use crate::{linalg, stats};
782 use dyn_stack::MemBuffer;
783 use equator::assert;
784 use stats::prelude::*;
785 #[test]
786 fn test_givens() {
787 let rng = &mut StdRng::seed_from_u64(0);
788 let rand = ComplexDistribution::new(StandardUniform, StandardUniform);
789 let mut sample = || -> c32 { rand.sample(rng) };
790 let shift = c32::new(0.5, 0.5);
791 let x = sample() - shift;
792 let y = sample() - shift;
793 let (c, s, r) = make_givens(x, y);
794 assert!((x * c + y * s - r).norm() < 1e-6);
795 assert!((-x * s.conj() + y * c).norm() < 1e-6);
796 }
797 #[test]
798 fn test_hessenberg() {
799 let rng = &mut StdRng::seed_from_u64(0);
800 for n in [12, 35, 128, 255] {
801 let rand = stats::CwiseMatDistribution {
802 nrows: n,
803 ncols: n,
804 dist: ComplexDistribution::new(StandardNormal, StandardNormal),
805 };
806 let mut sample = || -> Mat<c64> { rand.rand(rng) };
807 let A = sample();
808 let mut B = sample();
809 zip!(&mut B).for_each_triangular_lower(
810 linalg::zip::Diag::Skip,
811 |unzip!(x)| {
812 *x = 0.0.into();
813 },
814 );
815 let B = B;
816 let mut mem =
817 MemBuffer::new(generalized_hessenberg_scratch::<c64>(
818 n,
819 GeneralizedHessenbergParams {
820 block_size: 32,
821 ..auto!(c64)
822 },
823 ));
824 for gen_hessenberg in [
825 |A: MatMut<'_, _>,
826 B: MatMut<'_, _>,
827 Q: Option<MatMut<'_, _>>,
828 Z: Option<MatMut<'_, _>>,
829 _,
830 _,
831 par,
832 stack: &mut MemStack,
833 params| {
834 generalized_hessenberg(A, B, Q, Z, par, stack, params)
835 },
836 generalized_hessenberg_unblocked,
837 ] {
838 for (Q_is_I, Z_is_I) in
839 [(false, false), (true, true), (true, false), (false, true)]
840 {
841 let Q0 = &if Q_is_I {
842 Mat::identity(n, n)
843 } else {
844 UnitaryMat {
845 dim: n,
846 standard_normal: ComplexDistribution::new(
847 StandardNormal,
848 StandardNormal,
849 ),
850 }
851 .rand::<Mat<c64>>(rng)
852 };
853 let Z0 = &if Z_is_I {
854 Mat::identity(n, n)
855 } else {
856 UnitaryMat {
857 dim: n,
858 standard_normal: ComplexDistribution::new(
859 StandardNormal,
860 StandardNormal,
861 ),
862 }
863 .rand::<Mat<c64>>(rng)
864 };
865 let mut Q = Q0.to_owned();
866 let mut Z = Z0.to_owned();
867 let mut H = A.clone();
868 let mut T = B.clone();
869 gen_hessenberg(
870 H.as_mut(),
871 T.as_mut(),
872 Some(Q.as_mut()),
873 Some(Z.as_mut()),
874 Q_is_I,
875 Z_is_I,
876 Par::Seq,
877 MemStack::new(&mut mem),
878 GeneralizedHessenbergParams {
879 block_size: 32,
880 ..auto!(c64)
881 },
882 );
883 Q = Q0.adjoint() * &Q;
884 Z = Z0.adjoint() * &Z;
885 assert!((&Q * &H * Z.adjoint() - &A).norm_max() < 1e-13);
886 assert!((&Q * &T * Z.adjoint() - &B).norm_max() < 1e-13);
887 for j in 0..n {
888 for i in j + 1..n {
889 assert!(T[(i, j)] == c64::ZERO);
890 }
891 for i in j + 2..n {
892 assert!(H[(i, j)] == c64::ZERO);
893 }
894 }
895 }
896 }
897 }
898 }
899}