1use super::{Inner, Matrix, PlusTimes, Scalar, Side, Tiled, Triangle, filled, identity, transpose_into};
11use crate::{Native, float::{FUSED, sealed::Core}, parallel, slice::appended};
12use num_complex::Complex;
13use num_traits::{Float as _, One as _, Zero as _};
14use std::{borrow::Cow, cmp::Ordering, marker::PhantomData, mem::MaybeUninit};
15
16pub(super) mod eigen;
17pub(super) mod general;
18mod secular;
19pub(super) mod svd;
20mod truncated;
21#[doc(hidden)]
23pub const EIGH_CONQUER: usize = 64;
24#[doc(hidden)]
26pub const SVD_SPLIT: usize = 64;
27
28fn cholesky_block(n: usize) -> usize { (n / 8).clamp(32, 128) }
30
31fn qr_block(p: usize) -> usize { (p / 16).clamp(32, 64) }
33
34const SMALL: usize = 96;
36
37fn reflections<T: Tiled>(w: &[T], taus: &[T], n: usize) -> Vec<T> {
39 let mut qt = identity(n);
40 if n < 2 { return qt; }
41 if n <= SMALL {
42 let mut v = filled(n, T::zero());
44 for (k, &tau) in taus.iter().enumerate().rev() {
45 if tau == T::zero() { continue; }
46 let v = &mut v[..n - k - 1];
47 v[0] = T::one();
48 v[1..].copy_from_slice(&w[k * n + k + 2..(k + 1) * n]);
49 for row in qt[(k + 1) * n..].chunks_exact_mut(n) {
50 let row = &mut row[k + 1..];
51 let s = tau * dotc(v, row);
52 axpy(row, -s, v);
53 }
54 }
55 return qt;
56 }
57 let p = n - 1;
58 transpose_into(&blocked(w, taus, n).q(), p, &mut qt[n + 1..], n, [p, p], |x| x);
59 qt
60}
61
62fn blocked<T: Tiled>(w: &[T], taus: &[T], n: usize) -> Reflections<T> {
64 let p = n - 1;
65 let blocks = blocks_of(&taus[..p], qr_block(p), |k0, k1| (k0..k1).flat_map(|k| w[k * n + k0 + 1..(k + 1) * n].iter().copied()).collect());
66 Reflections { m: p, blocks }
67}
68
69fn qr_width(p: usize) -> usize { (p / 8).clamp(32, 128) }
71pub(super) fn stride(n: usize) -> usize { if n > 0 && n % 128 == 0 { n + 8 } else { n } }
73const ONE_AT_A_TIME: usize = 64;
75#[doc(hidden)]
77pub const FEW_COLUMNS: usize = 12;
78const LEAF: usize = 16;
80#[cfg(native)]
82const NS: usize = 32;
83
84pub trait Factor: Scalar {
86 fn matrix_lu_into(a: &[Self], n: usize, work: &mut LuWork<Self>);
88 fn matrix_lu_solve(lu: LuRef<'_, Self>, b: &[Self], k: usize) -> Vec<Self>;
90 fn matrix_lu_inverse_into(lu: LuRef<'_, Self>, x: &mut Vec<Self>, out: &mut [Self]);
92 fn matrix_cholesky(a: &[Self], n: usize) -> Option<Vec<Self>>;
94 fn matrix_qr(a: &[Self], size: [usize; 2], complete: bool) -> (Vec<Self>, Vec<Self>);
96 fn matrix_qr_into(a: &[Self], size: [usize; 2], complete: bool, work: &mut QrWork<Self>, q: &mut [Self], r: &mut [Self]);
98 fn matrix_svd(a: &[Self], size: [usize; 2], vectors: bool) -> Svd<Self>;
100 fn matrix_truncated_svd(a: &[Self], size: [usize; 2], k: usize) -> Svd<Self>;
102 fn matrix_svd_solve(svd: SvdRef<'_, Self>, b: Option<&[Self]>, k: usize) -> Vec<Self>;
104 fn matrix_lstsq(a: &[Self], b: &[Self], size: [usize; 2], k: usize) -> (Vec<Self>, Vec<Self::Real>);
106 fn matrix_eigh(a: &[Self], n: usize, vectors: bool) -> Option<Eigh<Self>>;
108 fn matrix_eig(a: &[Self], n: usize, vectors: bool) -> Option<Eig<Self::Real>>;
110}
111
112macro_rules! factor {
113 ($($t:ty),+) => {$(
114 impl Factor for $t {
115 fn matrix_lu_into(a: &[$t], n: usize, work: &mut LuWork<$t>) { lu_factor(a, n, work) }
116 fn matrix_lu_solve(lu: LuRef<'_, $t>, b: &[$t], k: usize) -> Vec<$t> { lu_solve(lu, b, k) }
117 fn matrix_lu_inverse_into(lu: LuRef<'_, $t>, x: &mut Vec<$t>, out: &mut [$t]) { lu_inverse(lu, x, out) }
118 fn matrix_cholesky(a: &[$t], n: usize) -> Option<Vec<$t>> { cholesky_factor(a, n) }
119 fn matrix_qr(a: &[$t], size: [usize; 2], complete: bool) -> (Vec<$t>, Vec<$t>) { qr_factor(a, size, complete) }
120 fn matrix_qr_into(a: &[$t], size: [usize; 2], complete: bool, work: &mut QrWork<$t>, q: &mut [$t], r: &mut [$t]) { qr_into(a, size, complete, work, q, r) }
121 fn matrix_svd(a: &[$t], size: [usize; 2], vectors: bool) -> Svd<$t> { svd_factor(a, size, vectors) }
122 fn matrix_truncated_svd(a: &[$t], size: [usize; 2], k: usize) -> Svd<$t> { truncated::truncated_factor(a, size, k) }
123 fn matrix_svd_solve(svd: SvdRef<'_, $t>, b: Option<&[$t]>, k: usize) -> Vec<$t> { svd_solve(svd, b, k) }
124 fn matrix_lstsq(a: &[$t], b: &[$t], size: [usize; 2], k: usize) -> (Vec<$t>, Vec<<$t as Scalar>::Real>) { lstsq_solve(a, b, size, k) }
125 fn matrix_eigh(a: &[$t], n: usize, vectors: bool) -> Option<Eigh<$t>> { eigh_factor(a, n, vectors) }
126 fn matrix_eig(a: &[$t], n: usize, vectors: bool) -> Option<Eig<<$t as Scalar>::Real>> { eig_factor(a, n, vectors) }
127 }
128 )+};
129}
130factor!(f32, f64, Complex<f32>, Complex<f64>);
131
132#[derive(Clone, Copy, Debug, PartialEq)]
134pub struct Lu<T, F = Vec<T>, P = Vec<usize>> { pub n: usize, pub factors: F, pub perm: P, items: PhantomData<T> }
135
136pub type LuRef<'a, T> = Lu<T, &'a [T], &'a [usize]>;
138
139impl<T, F: AsRef<[T]>, P: AsRef<[usize]>> Lu<T, F, P> {
140 pub fn new(n: usize, factors: F, perm: P) -> Self { Lu { n, factors, perm, items: PhantomData } }
142 pub fn as_ref(&self) -> LuRef<'_, T> { Lu::new(self.n, self.factors.as_ref(), self.perm.as_ref()) }
144}
145
146impl<T: Factor, F: AsRef<[T]>, P: AsRef<[usize]>> Lu<T, F, P> {
147 pub fn solve(&self, b: &[T], k: usize) -> Vec<T> { T::matrix_lu_solve(self.as_ref(), b, k) }
149 pub fn inverse(&self) -> Vec<T> {
151 let mut out = filled(self.n * self.n, T::zero());
152 T::matrix_lu_inverse_into(self.as_ref(), &mut Vec::new(), &mut out);
153 out
154 }
155 pub fn diagonal(&self) -> impl Iterator<Item = T> + '_ { let (n, factors) = (self.n, self.factors.as_ref()); (0..n).map(move |i| factors[i * n + i]) }
157 pub fn lower(&self) -> Vec<T> {
159 let (n, factors) = (self.n, self.factors.as_ref());
160 (0..n * n).map(|k| {
161 let (i, j) = (k / n, k % n);
162 match i.cmp(&j) { Ordering::Greater => factors[j * n + i], Ordering::Equal => T::one(), Ordering::Less => T::zero() }
163 }).collect()
164 }
165 pub fn upper(&self) -> Vec<T> {
167 let (n, factors) = (self.n, self.factors.as_ref());
168 (0..n * n).map(|k| {
169 let (i, j) = (k / n, k % n);
170 if j >= i { factors[j * n + i] } else { T::zero() }
171 }).collect()
172 }
173 pub fn determinant(&self) -> T { let product = self.diagonal().fold(T::one(), |p, d| p * d); if odd(self.perm.as_ref()) { -product } else { product } }
175 pub fn log_determinant(&self) -> (T, T::Real) {
177 let start = (if odd(self.perm.as_ref()) { -T::one() } else { T::one() }, T::Real::zero());
178 self.diagonal().fold(start, |(sign, log), d| {
179 let m = d.modulus();
180 if m == T::Real::zero() { (T::zero(), T::Real::neg_infinity()) } else { (sign * d.mul_real(m.recip()), log + m.ln()) }
181 })
182 }
183}
184
185fn odd(perm: &[usize]) -> bool {
187 let cycles = (0..perm.len()).filter(|&start| {
189 let mut i = perm[start];
190 while i > start { i = perm[i] }
191 i == start
192 }).count();
193 (perm.len() - cycles) % 2 == 1
194}
195
196pub struct LuWork<T> { lu: Lu<T>, scratch: Vec<T>, pivots: Vec<usize>, x: Vec<T> }
198
199impl<T> Default for LuWork<T> {
200 fn default() -> Self { LuWork { lu: Lu::new(0, Vec::new(), Vec::new()), scratch: Vec::new(), pivots: Vec::new(), x: Vec::new() } }
201}
202
203impl<T: Factor> LuWork<T> {
204 pub fn factor(&mut self, a: &[T], n: usize) -> LuRef<'_, T> { T::matrix_lu_into(a, n, self); self.lu.as_ref() }
206 pub fn solve(&mut self, a: &[T], b: &[T], size: [usize; 2], k: usize) -> Vec<T> {
208 if self.regular(a, size) { self.lu.as_ref().solve(b, k) } else { lstsq(a, b, size, k).0 }
209 }
210 pub fn pinv_into(&mut self, a: &[T], size: [usize; 2], out: &mut [T]) {
212 if self.inverse_into(a, size, out).is_none() { out.copy_from_slice(&svd(a, size).pseudoinverse()) }
213 }
214 pub fn inverse_into(&mut self, a: &[T], size: [usize; 2], out: &mut [T]) -> Option<T> {
216 if !self.regular(a, size) { return None; }
217 T::matrix_lu_inverse_into(self.lu.as_ref(), &mut self.x, out);
218 Some(self.lu.as_ref().determinant())
219 }
220 fn regular(&mut self, a: &[T], [m, n]: [usize; 2]) -> bool {
222 if m != n || n == 0 { return false; }
223 let largest = greatest(a, T::modulus);
224 let negligible = T::Real::epsilon() * <T::Real as crate::Float>::narrow(n as f64) * largest;
225 self.factor(a, n).diagonal().all(|d| !(d.modulus() <= negligible))
226 }
227}
228
229pub struct QrWork<T> { at: Vec<T>, taus: Vec<T>, qt: Vec<T> }
231
232impl<T> Default for QrWork<T> { fn default() -> Self { QrWork { at: Vec::new(), taus: Vec::new(), qt: Vec::new() } } }
233
234impl<T: Factor> QrWork<T> {
235 pub fn qr_into(&mut self, a: &[T], size: [usize; 2], complete: bool, q: &mut [T], r: &mut [T]) { T::matrix_qr_into(a, size, complete, self, q, r) }
237}
238
239#[derive(Clone, Copy, Debug, PartialEq)]
241pub struct Svd<T: Scalar, U = Vec<T>, S = Vec<<T as Scalar>::Real>> { pub size: [usize; 2], pub u: U, pub s: S, pub v: U, items: PhantomData<T> }
242
243pub type SvdRef<'a, T> = Svd<T, &'a [T], &'a [<T as Scalar>::Real]>;
245
246impl<T: Scalar, U: AsRef<[T]>, S: AsRef<[T::Real]>> Svd<T, U, S> {
247 pub fn new(size: [usize; 2], u: U, s: S, v: U) -> Self { Svd { size, u, s, v, items: PhantomData } }
249 pub fn as_ref(&self) -> SvdRef<'_, T> { Svd::new(self.size, self.u.as_ref(), self.s.as_ref(), self.v.as_ref()) }
251}
252
253impl<T: Factor, U: AsRef<[T]>, S: AsRef<[T::Real]>> Svd<T, U, S> {
254 pub fn solve(&self, b: &[T], k: usize) -> Vec<T> { T::matrix_svd_solve(self.as_ref(), Some(b), k) }
256 pub fn pseudoinverse(&self) -> Vec<T> { T::matrix_svd_solve(self.as_ref(), None, self.size[0]) }
258}
259
260pub fn lu<T: Factor>(a: &[T], n: usize) -> Lu<T> {
262 let mut work = LuWork::default();
263 work.factor(a, n);
264 work.lu
265}
266pub fn cholesky<T: Factor>(a: &[T], n: usize) -> Option<Vec<T>> { T::matrix_cholesky(a, n) }
268pub fn qr<T: Factor>(a: &[T], size: [usize; 2]) -> (Vec<T>, Vec<T>) { T::matrix_qr(a, size, false) }
270pub fn svd<T: Factor>(a: &[T], size: [usize; 2]) -> Svd<T> { T::matrix_svd(a, size, true) }
272pub fn singular_values<T: Factor>(a: &[T], size: [usize; 2]) -> Vec<T::Real> { T::matrix_svd(a, size, false).s }
274pub fn truncated_svd<T: Factor>(a: &[T], size: [usize; 2], k: usize) -> Svd<T> { T::matrix_truncated_svd(a, size, k) }
276pub fn lstsq<T: Factor>(a: &[T], b: &[T], size: [usize; 2], k: usize) -> (Vec<T>, Vec<T::Real>) { T::matrix_lstsq(a, b, size, k) }
278pub fn solve<T: Factor>(a: &[T], b: &[T], size: [usize; 2], k: usize) -> Vec<T> { LuWork::default().solve(a, b, size, k) }
280pub fn pinv<T: Factor>(a: &[T], [m, n]: [usize; 2]) -> Vec<T> {
282 let mut out = filled(m * n, T::zero());
283 LuWork::default().pinv_into(a, [m, n], &mut out);
284 out
285}
286fn safe_low<R: Native>() -> R { R::min_positive_value().sqrt() / R::epsilon() }
288pub fn rank<T: Factor>(a: &[T], size: [usize; 2]) -> usize {
290 let s = singular_values(a, size);
291 let floor = cutoff(&s, size[0].max(size[1]));
292 s.iter().filter(|&&x| x > floor).count()
293}
294
295#[derive(Clone, Copy, Debug, PartialEq, Eq)]
297pub enum Norm {
298 One,
300 Two,
302 Inf,
304 Frobenius,
306}
307
308pub fn norm<T: Factor>(a: &[T], [m, n]: [usize; 2], p: Norm) -> T::Real {
310 let zero = T::Real::zero();
311 if m == 0 || n == 0 { return zero; }
312 let largest = |x: &mut dyn Iterator<Item = T::Real>| x.fold(zero, |l, x| if x > l || x.is_nan() { x } else { l });
313 match p {
314 Norm::One => {
315 let mut sums = vec![zero; n];
316 for row in a.chunks_exact(n) { for (s, x) in sums.iter_mut().zip(row) { *s = *s + x.modulus() } }
317 largest(&mut sums.into_iter())
318 }
319 Norm::Two => singular_values(a, [m, n])[0],
320 Norm::Inf => largest(&mut a.chunks_exact(n).map(|row| row.iter().fold(zero, |s, x| s + x.modulus()))),
321 Norm::Frobenius => crate::stats::norm(a, [1, a.len(), 1], 2.0)[0],
322 }
323}
324
325pub fn cond<T: Factor>(a: &[T], [m, n]: [usize; 2], p: Norm) -> T::Real {
327 if p == Norm::Two {
328 let s = singular_values(a, [m, n]);
329 return match (s.first(), s.last()) {
330 (Some(&large), Some(&small)) => { if large == T::Real::zero() { T::Real::infinity() } else { large / small } } _ => T::Real::nan()
331 };
332 }
333 if m != n { return T::Real::infinity(); }
334 let mut inverse = filled(n * n, T::zero());
335 match LuWork::default().inverse_into(a, [m, n], &mut inverse) { Some(_) => norm(a, [m, n], p) * norm(&inverse, [n, n], p), None => T::Real::infinity() }
336}
337
338#[derive(Clone, Debug, PartialEq)]
340pub struct Eigh<T: Scalar> { pub values: Vec<T::Real>, pub vectors: Vec<T> }
341
342pub fn eigh<T: Factor>(a: &[T], n: usize) -> Option<Eigh<T>> { T::matrix_eigh(a, n, true) }
344pub fn hermitian_eigenvalues<T: Factor>(a: &[T], n: usize) -> Option<Vec<T::Real>> { T::matrix_eigh(a, n, false).map(|f| f.values) }
346
347#[derive(Clone, Debug, PartialEq)]
349pub struct Eig<T> { pub values: Vec<Complex<T>>, pub vectors: Vec<Complex<T>> }
350
351pub fn eig<T: Factor>(a: &[T], n: usize) -> Option<Eig<T::Real>> { T::matrix_eig(a, n, true) }
353
354pub fn eigenvalues<T: Factor>(a: &[T], n: usize) -> Option<Vec<Complex<T::Real>>> { T::matrix_eig(a, n, false).map(|f| f.values) }
356
357pub fn from_svd<T: Factor>(u: &[T], s: &[T], v: &[T], [m, p, n]: [usize; 3]) -> Vec<T> where PlusTimes: Inner<T, Out = T> {
359 PlusTimes.inner(&scaled_columns(u, s, p), &adjoint(v, n, p), [m, p, n])
360}
361
362pub fn from_cholesky<T: Factor>(l: &[T], n: usize) -> Vec<T> where PlusTimes: Inner<T, Out = T> { PlusTimes.inner(l, &adjoint(l, n, n), [n; 3]) }
364
365pub fn from_eigen<T: Factor>(values: &[T], vectors: &[T], n: usize) -> Option<Vec<T>> where PlusTimes: Inner<T, Out = T> {
367 let mut inverse = filled(n * n, T::zero());
368 LuWork::default().inverse_into(vectors, [n, n], &mut inverse)?;
369 let lengths = crate::stats::norm(vectors, [1, n, n], 2.0);
371 let scaled: Vec<T> = vectors.chunks_exact(n).flat_map(|row| row.iter().zip(&lengths).map(|(&x, &l)| x.div_real(l))).collect();
372 let scaled_inverse: Vec<T> = inverse.chunks_exact(n).zip(&lengths).flat_map(|(row, &l)| row.iter().map(move |&x| x.mul_real(l))).collect();
373 let condition = norm(&scaled, [n, n], Norm::One) * norm(&scaled_inverse, [n, n], Norm::One);
374 (condition <= <T::Real as num_traits::Float>::epsilon().sqrt().recip()).then(|| PlusTimes.inner(&scaled_columns(vectors, values, n), &inverse, [n; 3]))
375}
376
377pub fn from_lu<T: Factor>(l: &[T], u: &[T], p: &[T], n: usize) -> Option<Vec<T>> where PlusTimes: Inner<T, Out = T> {
379 let mut seen = vec![false; n];
380 let order = p.iter().map(|&x| {
381 let i = x.re();
382 if x != T::from_real(i) || i.fract() != T::Real::zero() { return None; }
383 let i = num_traits::ToPrimitive::to_usize(&i).filter(|&i| i < n && !seen[i])?;
384 seen[i] = true;
385 Some(i)
386 }).collect::<Option<Vec<_>>>()?;
387 let product = PlusTimes.inner(l, u, [n; 3]);
388 let mut a = filled(n * n, T::zero());
389 for (i, &row) in order.iter().enumerate() { a[row * n..(row + 1) * n].copy_from_slice(&product[i * n..(i + 1) * n]) }
390 Some(a)
391}
392
393pub fn is_hermitian<T: Scalar>(a: &[T], n: usize) -> bool { (0..n).all(|i| (0..n).all(|j| a[i * n + j] == a[j * n + i].conj())) }
395
396fn scaled_columns<T: Scalar>(a: &[T], d: &[T], columns: usize) -> Vec<T> {
398 a.chunks_exact(columns.max(1)).flat_map(|row| row.iter().zip(d).map(|(&x, &s)| x * s)).collect()
399}
400
401fn finite<T: Scalar>(x: &[T]) -> bool { x.iter().all(|x| x.re().is_finite() && x.im().is_finite()) }
403
404fn svd_factor<T: Tiled>(a: &[T], [m, n]: [usize; 2], vectors: bool) -> Svd<T> {
406 assert_eq!(a.len(), m * n, "`a` must hold `m`×`n` items");
407 let p = m.min(n);
408 if p == 0 { return Svd::new([m, n], Vec::new(), Vec::new(), Vec::new()); }
409 let (u, s, v) = finite(a).then(|| tall_svd(a, [m, n], vectors)).flatten().unwrap_or_else(|| {
410 let nan = T::Real::nan();
411 let filled = |len| if vectors { vec![T::from_real(nan); len] } else { Vec::new() };
412 (filled(m * p), vec![nan; p], filled(n * p))
413 });
414 Svd::new([m, n], u, s, v)
415}
416
417fn tall_svd<T: Tiled>(a: &[T], [m, n]: [usize; 2], vectors: bool) -> Option<(Vec<T>, Vec<T::Real>, Vec<T>)> {
419 if 6 * m < 11 * n { return T::svd(a, [m, n], vectors); }
420 if tall_for_qr([m, n]) && let Some((r, q)) = tall_qr(a, [m, n], if vectors { Tall::Q } else { Tall::R }) {
421 let (ur, s, v) = T::svd(&r, [n, n], vectors)?;
422 let u = if vectors { T::multiply(Matrix::rows(&q, n), Matrix::rows(&ur, n), [m, n, n]) } else { Vec::new() };
423 return Some((u, s, v));
424 }
425 let (reflections, r) = reflect(a, [m, n]);
426 let (ur, s, v) = T::svd(&r, [n, n], vectors)?;
427 let mut u = Vec::new();
428 if vectors {
429 u = filled(m * n, T::zero());
430 u[..n * n].copy_from_slice(&ur);
431 reflections.apply(&mut u, n);
432 }
433 Some((u, s, v))
434}
435
436fn eigh_factor<T: Tiled>(a: &[T], n: usize, vectors: bool) -> Option<Eigh<T>> {
438 assert_eq!(a.len(), n * n, "`a` must hold `n`×`n` items");
439 if n == 0 { return Some(Eigh { values: Vec::new(), vectors: Vec::new() }); }
440 (0..n).all(|i| finite(&a[i * n..=i * n + i])).then(|| T::eigh(a, n, vectors)).flatten()
441}
442
443fn eig_factor<T: Tiled>(a: &[T], n: usize, vectors: bool) -> Option<Eig<T::Real>> {
445 assert_eq!(a.len(), n * n, "`a` must hold `n`×`n` items");
446 if n == 0 { return Some(Eig { values: Vec::new(), vectors: Vec::new() }); }
447 finite(a).then(|| T::eig(a, n, vectors)).flatten()
448}
449
450#[inline(always)]
452pub(crate) fn axpy<T: Scalar>(y: &mut [T], a: T, x: &[T]) { for (y, &x) in y.iter_mut().zip(x) { *y = x.mul_acc(a, *y) } }
453
454#[inline(always)]
456fn axpyc<T: Scalar>(y: &mut [T], a: T, x: &[T]) { for (y, &x) in y.iter_mut().zip(x) { *y = x.conj().mul_acc(a, *y) } }
457
458#[inline(always)]
460pub(crate) fn dot<T: Scalar>(x: &[T], y: &[T]) -> T { sum_of_products(x, y, |a| a) }
461
462#[inline(always)]
464fn dotc<T: Scalar>(x: &[T], y: &[T]) -> T { sum_of_products(x, y, T::conj) }
465
466#[inline(always)]
468fn sum_of_products<T: Scalar>(x: &[T], y: &[T], f: impl Fn(T) -> T) -> T {
469 let mut s = [T::zero(); 8];
470 let ((xc, xr), (yc, yr)) = (x.as_chunks::<8>(), y.as_chunks::<8>());
471 for (a, b) in xc.iter().zip(yc) { for k in 0..8 { s[k] = f(a[k]).mul_acc(b[k], s[k]) } }
472 xr.iter().zip(yr).fold(s.iter().fold(T::zero(), |t, &v| t + v), |t, (&a, &b)| f(a).mul_acc(b, t))
473}
474
475fn largest<T: Scalar>(x: &[T]) -> usize { let top = greatest(x, T::abs1); x.iter().position(|v| v.abs1() == top).unwrap_or(0) }
477
478fn plain_hypot<R: Native>(x: R, y: R) -> R { x.mla::<FUSED>(x, y * y).sqrt() }
480
481pub(crate) fn givens<R: Native>(y: R, z: R) -> (R, R, R) {
483 let r = plain_hypot(y, z);
484 if r == R::zero() { (R::one(), R::zero(), R::zero()) } else { (y / r, z / r, r) }
485}
486
487fn greatest<T: Copy, R: Native>(x: &[T], f: impl Fn(T) -> R) -> R { crate::reduce::folded(x, R::zero(), |t, v| f(v).maximum(t), R::maximum) }
489
490fn safe_scale<T: Scalar>(a: &[T]) -> Option<T::Real> {
492 let top = greatest(a, T::abs1);
493 let top = if top.is_finite() { top } else { greatest(a, |x| x.re().abs().max(x.im().abs())) };
494 let low = safe_low::<T::Real>();
495 if !top.is_finite() || top == T::Real::zero() || (low..=low.recip()).contains(&top) { return None; }
496 unit_scale(top)
497}
498
499fn unit_scale<R: Native>(top: R) -> Option<R> {
501 let limit = R::max_value().log2().floor();
502 Some((R::one() + R::one()).powi((-top.log2().floor()).max(-limit).min(limit).to_i32()?))
503}
504
505fn divide<T: Scalar>(x: &mut [T], d: T) {
507 if d.abs1() >= T::Real::min_positive_value() { let r = d.reciprocal(); for v in x { *v = *v * r } } else { for v in x { *v = v.quotient(d) } }
508}
509
510pub(super) fn transpose_by<T: Copy + num_traits::Zero + Send + Sync>(a: &[T], ls: usize, size: [usize; 2], f: impl Fn(T) -> T + Sync) -> Vec<T> {
512 let mut t = Vec::new();
513 transpose_by_into(a, [ls, size[0]], size, f, &mut t);
514 t
515}
516pub(super) fn transpose_by_into<T: Copy + num_traits::Zero + Send + Sync>(
518 a: &[T], [ls, lt]: [usize; 2], [rows, cols]: [usize; 2], f: impl Fn(T) -> T + Sync, t: &mut Vec<T>) {
519 t.clear();
520 if rows * cols == 0 { return; }
521 unsafe {
522 appended(t, lt * cols, |slots| {
523 parallel::split(slots, lt, &|start, part| {
524 transpose_into(&a[start / lt..], ls, part, lt, [rows, part.len() / lt], |x| MaybeUninit::new(f(x)));
525 for row in part.chunks_exact_mut(lt) { row[rows..].fill(MaybeUninit::new(T::zero())) }
526 true
527 })
528 })
529 };
530}
531
532fn transpose<T: Copy + num_traits::Zero + Send + Sync>(a: &[T], m: usize, n: usize) -> Vec<T> { transpose_by(a, n, [m, n], |x| x) }
534
535pub(super) fn spaced<T: Copy + num_traits::Zero>(a: &[T], [rows, cols]: [usize; 2], ld: usize) -> Vec<T> {
537 if ld == cols { return a[..rows * cols].to_vec(); }
538 let mut x = Vec::with_capacity(rows * ld);
539 for row in a[..rows * cols].chunks_exact(cols) {
540 x.extend_from_slice(row);
541 x.resize(x.len() + ld - cols, T::zero());
542 }
543 x
544}
545
546pub(super) fn compact<T: Copy>(x: &mut Vec<T>, [rows, cols]: [usize; 2], ld: usize) {
548 if ld != cols { for i in 1..rows { x.copy_within(i * ld..i * ld + cols, i * cols) } }
549 x.truncate(rows * cols);
550}
551
552fn adjoint<T: Scalar>(a: &[T], m: usize, n: usize) -> Vec<T> { transpose_by(a, n, [m, n], T::conj) }
554
555#[inline(always)]
557fn two_rows<T>(a: &mut [T], width: usize, p: usize, q: usize) -> (&mut [T], &mut [T]) {
558 let (head, tail) = a.split_at_mut(q * width);
559 (&mut head[p * width..(p + 1) * width], &mut tail[..width])
560}
561
562#[inline(always)]
564fn rotate<T: Scalar>(x: &mut [T], y: &mut [T], c: T::Real, s: T::Real) {
565 for (x, y) in x.iter_mut().zip(y) { (*x, *y) = (x.mul_real(c) - y.mul_real(s), x.mul_real(s) + y.mul_real(c)) }
566}
567
568fn transposed<T: Scalar>(f: &[T], n: usize, k0: usize, k1: usize) -> Vec<T> { transpose_by(&f[k0 * n + k0..], n, [f.len() / n - k0, k1 - k0], |x| x) }
570
571fn swap_rows<T: Send>(columns: &mut [T], ld: usize, pivots: &[usize], first: usize) {
573 parallel::split(columns, ld, &|_, columns| {
574 for column in columns.chunks_exact_mut(ld) { for i in 0..pivots.len() { column.swap(first + i, pivots[i]) } }
575 true
576 });
577}
578
579fn lu_panel<T: Tiled>(f: &mut [T], [n, ld]: [usize; 2], [c0, c1]: [usize; 2], pivots: &mut [usize], scratch: &mut Vec<T>) {
581 if c1 - c0 > LEAF {
582 let mid = (c0 + c1) / 2;
583 lu_panel(f, [n, ld], [c0, mid], pivots, scratch);
584 let (left, right) = f.split_at_mut(mid * ld);
585 let (left, k, w) = (&*left, mid - c0, c1 - mid);
586 let lw = stride(w);
587 let half = &mut right[..w * ld];
588 swap_rows(half, ld, &pivots[c0..mid], c0);
589 transpose_by_into(&half[c0..], [ld, lw], [w, k], |x| x, scratch);
590 T::triangular(scratch, lw, Matrix::rows(&left[c0 * ld + c0..], ld).transpose(), [k, w], Side::Left, Triangle::UnitUpper);
591 let u12 = &scratch[..];
592 parallel::split(half, ld, &|start, columns| {
593 let count = columns.len() / ld;
594 transpose_into(&u12[start / ld..], lw, &mut columns[c0..], ld, [k, count], |x| x);
595 true
596 });
597 T::gemm(&mut half[mid..], ld, Matrix::rows(u12, lw).transpose(), Matrix::rows(&left[c0 * ld + mid..], ld), [w, k, n - mid], -T::one(), T::one());
598 lu_panel(f, [n, ld], [mid, c1], pivots, scratch);
599 return swap_rows(&mut f[c0 * ld..mid * ld], ld, &pivots[mid..c1], mid);
600 }
601 for c in c0..c1 {
602 let p = c + largest(&f[c * ld + c..c * ld + n]);
603 pivots[c] = p;
604 if p != c { for j in c0..c1 { f.swap(j * ld + c, j * ld + p) } }
605 let pivot = f[c * ld + c];
606 if pivot == T::zero() { continue; }
607 let (left, right) = f.split_at_mut((c + 1) * ld);
608 let l = &mut left[c * ld + c + 1..c * ld + n];
609 divide(l, pivot);
610 for column in right[..(c1 - c - 1) * ld].chunks_exact_mut(ld) {
611 let u = column[c];
612 axpy(&mut column[c + 1..n], -u, l);
613 }
614 }
615}
616fn lu_factor<T: Tiled>(a: &[T], n: usize, work: &mut LuWork<T>) {
618 assert_eq!(a.len(), n * n, "`a` must hold `n`×`n` items");
619 let LuWork { lu, scratch, pivots, .. } = work;
620 let ld = stride(n);
621 lu.n = n;
622 transpose_by_into(a, [n, ld], [n, n], |x| x, &mut lu.factors);
623 pivots.clear();
624 pivots.resize(n, 0);
625 lu_panel(&mut lu.factors, [n, ld], [0, n], pivots, scratch);
626 compact(&mut lu.factors, [n, n], ld);
627 lu.perm.clear();
628 lu.perm.extend(0..n);
629 for c in 0..n { lu.perm.swap(c, pivots[c]) }
630}
631
632fn lu_solve<T: Tiled>(lu: LuRef<'_, T>, b: &[T], k: usize) -> Vec<T> {
634 let (n, f) = (lu.n, lu.factors);
635 assert_eq!(b.len(), n * k, "`b` must hold `n`×`k` items");
636 let mut x: Vec<T> = lu.perm.iter().flat_map(|&p| b[p * k..(p + 1) * k].iter().copied()).collect();
637 let p = Matrix::rows(f, n).transpose();
638 T::triangular(&mut x, k, p, [n, k], Side::Left, Triangle::UnitUpper);
639 T::triangular(&mut x, k, p, [n, k], Side::Left, Triangle::Lower);
640 x
641}
642
643#[cfg_attr(not(native), allow(unused_variables))]
645fn lu_inverse<T: Tiled>(lu: LuRef<'_, T>, x: &mut Vec<T>, out: &mut [T]) {
646 let (n, p) = (lu.n, Matrix::rows(lu.factors, lu.n).transpose());
647 #[cfg(accelerate)]
648 if super::accelerated() {
649 out.fill(T::zero());
650 for (c, &r) in lu.perm.iter().enumerate() { out[c * n + r] = T::one() }
651 T::triangular(out, n, p, [n, n], Side::Left, Triangle::UnitUpper);
652 return T::triangular(out, n, p, [n, n], Side::Left, Triangle::Lower);
653 }
654 #[cfg(native)]
655 {
656 let ld = stride(n);
657 x.clear();
658 x.resize(n * ld, T::zero());
659 for i in 0..n { x[i * ld + i] = T::one() }
660 forward(p, [0, n], x, [ld, n], NS, true, true);
661 T::triangular(x, ld, p, [n, n], Side::Left, Triangle::Lower);
662 for r in 0..n { for (c, &q) in lu.perm.iter().enumerate() { out[r * n + q] = x[r * ld + c] } }
663 }
664}
665
666#[cfg(native)]
668fn forward<T: Tiled>(l: Matrix<'_, T>, [r0, r1]: [usize; 2], x: &mut [T], [ldx, k]: [usize; 2], block: usize, lower: bool, unit: bool) {
669 if !lower && r1 - r0 > block {
670 let mid = (r0 + r1) / 2;
671 let (top, bottom) = x.split_at_mut((mid - r0) * ldx);
672 forward(l, [r0, mid], top, [ldx, k], block, false, unit);
673 T::gemm(bottom, ldx, l.offset(mid, r0), Matrix::rows(top, ldx), [r1 - mid, mid - r0, k], -T::one(), T::one());
674 return forward(l, [mid, r1], bottom, [ldx, k], block, false, unit);
675 }
676 for i0 in (r0..r1).step_by(block) {
677 let i1 = (i0 + block).min(r1);
678 let width = if lower { i1 } else { k };
679 let (done, rest) = x.split_at_mut(((i1 - r0) * ldx).min(x.len()));
680 let rows = &mut done[(i0 - r0) * ldx..];
681 for i in i0..i1 {
682 let (head, tail) = rows.split_at_mut((i - i0) * ldx);
683 for j in i0..i { axpy(&mut tail[..width], -l.at(i, j), &head[(j - i0) * ldx..(j - i0) * ldx + width]) }
684 if !unit { divide(&mut tail[..width], l.at(i, i)) }
685 }
686 let rows = &*rows;
687 if i1 < r1 { T::gemm(rest, ldx, l.offset(i1, i0), Matrix::rows(rows, ldx), [r1 - i1, i1 - i0, width], -T::one(), T::one()) }
688 }
689}
690
691#[cfg(native)]
693fn backward<T: Tiled>(u: Matrix<'_, T>, [r0, r1]: [usize; 2], x: &mut [T], [ldx, k]: [usize; 2], unit: bool) {
694 if r1 - r0 > NS {
695 let mid = (r0 + r1) / 2;
696 let (top, bottom) = x.split_at_mut((mid - r0) * ldx);
697 backward(u, [mid, r1], bottom, [ldx, k], unit);
698 T::gemm(top, ldx, u.offset(r0, mid), Matrix::rows(bottom, ldx), [mid - r0, r1 - mid, k], -T::one(), T::one());
699 return backward(u, [r0, mid], top, [ldx, k], unit);
700 }
701 for i in (r0..r1).rev() {
702 let (done, below) = x.split_at_mut(((i - r0 + 1) * ldx).min(x.len()));
703 let xi = &mut done[(i - r0) * ldx..][..k];
704 for j in i + 1..r1 { axpy(xi, -u.at(i, j), &below[(j - i - 1) * ldx..][..k]) }
705 if !unit { divide(xi, u.at(i, i)) }
706 }
707}
708
709#[cfg(native)]
711pub(super) fn triangular_native<T: Tiled>(x: &mut [T], ldx: usize, a: Matrix<'_, T>, [m, n]: [usize; 2], side: Side, triangle: Triangle) {
712 if m == 0 || n == 0 { return; }
713 if matches!(side, Side::Right) {
714 let mut xt = transpose_by(x, ldx, [m, n], T::conj);
715 triangular_native(&mut xt, m, a.adjoint(), [n, m], Side::Left, triangle);
716 return transpose_into(&xt, m, x, ldx, [n, m], T::conj);
717 }
718 let (lower, unit) = (matches!(triangle, Triangle::Lower) != a.op.transposed(), matches!(triangle, Triangle::UnitUpper));
719 if n == 1 && ldx == 1 && single(x, a, m, lower, unit) { return; }
720 if lower { forward(a, [0, m], x, [ldx, n], if n >= m { LEAF } else { NS }, false, unit) } else { backward(a, [0, m], x, [ldx, n], unit) }
721}
722
723#[cfg(native)]
725fn single<T: Scalar>(x: &mut [T], a: Matrix<'_, T>, m: usize, lower: bool, unit: bool) -> bool {
726 let (d, ld) = (a.data, a.ld);
727 let solved = |v: T, i: usize| if unit { v } else { v.quotient(d[i * ld + i]) };
728 match a.op {
729 super::Op::Transpose if lower => {
730 for j in 0..m {
731 x[j] = solved(x[j], j);
732 let xj = x[j];
733 axpy(&mut x[j + 1..], -xj, &d[j * ld + j + 1..j * ld + m]);
734 }
735 }
736 super::Op::Transpose => {
737 for j in (0..m).rev() {
738 x[j] = solved(x[j], j);
739 let xj = x[j];
740 axpy(&mut x[..j], -xj, &d[j * ld..j * ld + j]);
741 }
742 }
743 super::Op::Normal if lower => { for i in 0..m { x[i] = solved(x[i] - dot(&d[i * ld..i * ld + i], &x[..i]), i) } }
744 super::Op::Normal => { for i in (0..m).rev() { x[i] = solved(x[i] - dot(&d[i * ld + i + 1..i * ld + m], &x[i + 1..]), i) } }
745 _ => return false,
746 }
747 true
748}
749
750fn cholesky_panel<T: Tiled>(panel: &mut [T], [rows, ld]: [usize; 2], [c0, c1]: [usize; 2]) -> bool {
752 if c1 - c0 > LEAF {
753 let mid = (c0 + c1) / 2;
754 if !cholesky_panel(panel, [rows, ld], [c0, mid]) { return false; }
755 let (left, right) = panel.split_at_mut(mid * ld);
756 let l = Matrix::rows(&left[c0 * ld + mid..], ld);
757 T::gemm(&mut right[mid..], ld, l.adjoint(), l, [c1 - mid, mid - c0, rows - mid], -T::one(), T::one());
758 return cholesky_panel(panel, [rows, ld], [mid, c1]);
759 }
760 for c in c0..c1 {
761 let d = panel[c * ld + c];
762 if !(d.re() > T::Real::zero()) { return false; }
763 let (left, right) = panel.split_at_mut((c + 1) * ld);
764 let d = T::from_real(d.re().sqrt());
765 left[c * ld + c] = d;
766 let lc = &mut left[c * ld + c + 1..c * ld + rows];
767 divide(lc, d);
768 for (j, column) in right[..(c1 - c - 1) * ld].chunks_exact_mut(ld).enumerate() {
770 let below = &lc[j..];
771 axpy(&mut column[c + 1 + j..rows], -below[0].conj(), below);
772 }
773 }
774 true
775}
776fn cholesky_factor<T: Tiled>(a: &[T], n: usize) -> Option<Vec<T>> {
778 assert_eq!(a.len(), n * n, "`a` must hold `n`×`n` items");
779 let ld = stride(n);
780 let mut l = spaced(a, [n, n], ld);
781 let (width, mut panel) = (cholesky_block(n), Vec::new());
782 for k0 in (0..n).step_by(width) {
783 let k1 = (k0 + width).min(n);
784 let (nb, rows, w) = (k1 - k0, n - k0, n - k1);
785 let pl = stride(rows);
786 transpose_by_into(&l[k0 * ld + k0..], [ld, pl], [rows, nb], |x| x, &mut panel);
787 if !cholesky_panel(&mut panel, [rows, pl], [0, nb]) { return None; }
788 transpose_into(&panel, pl, &mut l[k0 * ld + k0..], ld, [nb, rows], |x| x);
789 if w == 0 { break; }
790 T::gram(&mut l[k1 * ld + k1..], ld, Matrix::rows(&panel[nb..], pl).transpose(), [w, nb], -T::Real::one(), T::Real::one());
791 }
792 compact(&mut l, [n, n], ld);
793 for i in 0..n { l[i * n + i + 1..(i + 1) * n].fill(T::zero()) }
794 Some(l)
795}
796
797fn reflector<T: Scalar>(x: &mut [T]) -> T {
799 let (head, rest) = x.split_at_mut(1);
800 let (x0, below) = (head[0], dotc(rest, rest).re());
801 let zero = T::Real::zero();
802 if below == zero && x0.im() == zero { return T::zero(); }
803 let (re, im) = (x0.re(), x0.im());
804 let norm = re.mla::<FUSED>(re, im.mla::<FUSED>(im, below)).sqrt();
805 let beta = T::from_real(if re >= zero { -norm } else { norm });
806 let scale = (x0 - beta).reciprocal();
807 for v in rest.iter_mut() { *v = *v * scale }
808 head[0] = beta;
809 (beta - x0).div_real(beta.re())
810}
811
812#[inline(always)]
814fn apply_reflector<T: Scalar>(x: &mut [T], tau: T, tail: &[T]) {
815 let tw = tau.conj() * (x[0] + dotc(tail, &x[1..]));
816 x[0] = x[0] - tw;
817 axpy(&mut x[1..], -tw, tail);
818}
819
820fn triangular<T: Tiled>(panel: &[T], rows: usize, [c0, c1]: [usize; 2], t: &mut [T]) {
822 if c1 - c0 > LEAF {
823 let mid = (c0 + c1) / 2;
824 triangular(panel, rows, [c0, mid], t);
825 triangular(panel, rows, [mid, c1], t);
826 return joined(panel, rows, [c0, mid, c1], t);
827 }
828 let nb = panel.len() / rows;
829 for c in c0..c1 {
830 let tau = t[c * nb + c];
831 if tau == T::zero() { continue; }
832 let vc = &panel[c * rows + c + 1..(c + 1) * rows];
833 let s: Vec<T> = (c0..c).map(|j| panel[j * rows + c].conj() + dotc(&panel[j * rows + c + 1..(j + 1) * rows], vc)).collect();
835 for i in c0..c { t[i * nb + c] = -tau * dot(&t[i * nb + i..i * nb + c], &s[i - c0..]) }
836 }
837}
838
839fn joined<T: Tiled>(panel: &[T], rows: usize, [c0, mid, c1]: [usize; 3], t: &mut [T]) {
841 let ([k1, k2], nb, below) = ([mid - c0, c1 - mid], panel.len() / rows, rows - mid);
842 let v2 = vectors(panel, rows, [mid, c1]);
843 let s = T::multiply(Matrix::rows(&panel[c0 * rows + mid..], rows).conjugate(), Matrix::rows(&v2, below).transpose(), [k1, below, k2]);
844 let s2 = T::multiply(Matrix::rows(&s, k2), Matrix::rows(&t[mid * nb + mid..], nb), [k1, k2, k2]);
845 let t12 = T::multiply(Matrix::rows(&t[c0 * nb + c0..], nb), Matrix::rows(&s2, k2), [k1, k1, k2]);
846 for i in 0..k1 { for j in 0..k2 { t[(c0 + i) * nb + mid + j] = -t12[i * k2 + j] } }
847}
848
849fn explicit_vectors<T: Scalar>(panel: &mut [T], rows: usize) {
851 for (c, v) in panel.chunks_exact_mut(rows).enumerate() {
852 v[..c].fill(T::zero());
853 v[c] = T::one();
854 }
855}
856
857fn vectors<T: Scalar>(panel: &[T], rows: usize, [c0, c1]: [usize; 2]) -> Vec<T> {
859 let mut v = Vec::with_capacity((c1 - c0) * (rows - c0));
860 for column in panel[c0 * rows..c1 * rows].chunks_exact(rows) { v.extend_from_slice(&column[c0..]) }
861 explicit_vectors(&mut v, rows - c0);
862 v
863}
864
865fn reflect_columns<T: Scalar>(panel: &mut [T], rows: usize, [c0, c1, end]: [usize; 3], mut keep: impl FnMut(usize, T)) {
867 for c in c0..c1 {
868 let (left, right) = panel.split_at_mut((c + 1) * rows);
869 let v = &mut left[c * rows + c..];
870 let tau = reflector(v);
871 keep(c, tau);
872 if tau == T::zero() { continue; }
873 for column in right[..(end - c - 1) * rows].chunks_exact_mut(rows) { apply_reflector(&mut column[c..], tau, &v[1..]) }
875 }
876}
877
878fn panel_qr<T: Tiled>(panel: &mut [T], rows: usize, [c0, c1]: [usize; 2], t: &mut [T]) {
880 let nb = panel.len() / rows;
881 if c1 - c0 <= LEAF {
882 reflect_columns(panel, rows, [c0, c1, c1], |c, tau| t[c * nb + c] = tau);
883 return triangular(panel, rows, [c0, c1], t);
884 }
885 let mid = (c0 + c1) / 2;
886 let ([k1, k2], r) = ([mid - c0, c1 - mid], rows - c0);
887 panel_qr(panel, rows, [c0, mid], t);
888 let vt = vectors(panel, rows, [c0, mid]);
889 let x = &mut panel[mid * rows..c1 * rows];
891 let w = T::multiply(Matrix::rows(&x[c0..], rows), Matrix::rows(&vt, r).adjoint(), [k2, r, k1]);
892 let z = T::multiply(Matrix::rows(&w, k1), Matrix::rows(&t[c0 * nb + c0..], nb).conjugate(), [k2, k1, k1]);
893 T::gemm(&mut x[c0..], rows, Matrix::rows(&z, k1), Matrix::rows(&vt, r), [k2, k1, r], -T::one(), T::one());
894 panel_qr(panel, rows, [mid, c1], t);
895 joined(panel, rows, [c0, mid, c1], t)
896}
897
898struct Block<T> { k0: usize, nb: usize, vt: Vec<T>, t: Vec<T> }
900
901impl<T: Tiled> Block<T> {
902 fn new(k0: usize, panel: Vec<T>, taus: &[T]) -> Self {
904 let nb = taus.len();
905 let mut t = filled(nb * nb, T::zero());
906 for (c, &tau) in taus.iter().enumerate() { t[c * nb + c] = tau }
907 triangular(&panel, panel.len() / nb, [0, nb], &mut t);
908 Self::from_panel(k0, nb, panel, t)
909 }
910
911 fn from_panel(k0: usize, nb: usize, mut panel: Vec<T>, t: Vec<T>) -> Self {
913 let rows = panel.len() / nb;
914 explicit_vectors(&mut panel, rows);
915 Self { k0, nb, vt: panel, t }
916 }
917
918 fn apply(&self, x: &mut [T], ld: usize, [c0, c1]: [usize; 2], adjoint: bool) {
920 let (k0, nb, rows, cols) = (self.k0, self.nb, self.vt.len() / self.nb, c1 - c0);
921 let t = Matrix::rows(&self.t, nb);
923 let t = if adjoint { t } else { t.adjoint() };
924 let v = Matrix::rows(&self.vt, rows).transpose();
925 let wt = T::multiply(Matrix::rows(&x[k0 * ld + c0..], ld).adjoint(), v, [cols, rows, nb]);
926 let twt = T::multiply(Matrix::rows(&wt, nb), t, [cols, nb, nb]);
927 T::gemm(&mut x[k0 * ld + c0..], ld, v, Matrix::rows(&twt, nb).adjoint(), [rows, nb, cols], -T::one(), T::one());
928 }
929}
930
931fn blocks_of<T: Tiled>(taus: &[T], width: usize, panel: impl Fn(usize, usize) -> Vec<T>) -> Vec<Block<T>> {
933 let p = taus.len();
934 (0..p).step_by(width).map(|k0| {
935 let k1 = (k0 + width).min(p);
936 Block::new(k0, panel(k0, k1), &taus[k0..k1])
937 }).collect()
938}
939
940struct Reflections<T> { m: usize, blocks: Vec<Block<T>> }
942
943fn reflect<T: Tiled>(a: &[T], [m, n]: [usize; 2]) -> (Reflections<T>, Vec<T>) {
945 assert_eq!(a.len(), m * n, "`a` must hold `m`×`n` items");
946 let p = m.min(n);
947 let (ld, scale) = (stride(n), safe_scale(a));
948 let mut f = spaced(a, [m, n], ld);
949 if let Some(s) = scale { for x in &mut f { *x = x.mul_real(s) } }
950 let mut blocks = Vec::new();
951 let width = qr_width(p);
952 for k0 in (0..p).step_by(width) {
953 let k1 = (k0 + width).min(p);
954 let (nb, rows) = (k1 - k0, m - k0);
955 let mut panel = transposed(&f, ld, k0, k1);
956 let mut t = filled(nb * nb, T::zero());
957 panel_qr(&mut panel, rows, [0, nb], &mut t);
958 for r in 0..nb { for c in r..nb { f[(k0 + r) * ld + k0 + c] = panel[c * rows + r] } }
959 let block = Block::from_panel(k0, nb, panel, t);
960 if k1 < n { block.apply(&mut f, ld, [k1, n], true) }
962 blocks.push(block);
963 }
964 let r = (0..p * n).map(|k| {
966 let (i, j) = (k / n, k % n);
967 if j < i { T::zero() } else { scale.map_or(f[i * ld + j], |s| f[i * ld + j].div_real(s)) }
968 }).collect();
969 (Reflections { m, blocks }, r)
970}
971
972impl<T: Tiled> Reflections<T> {
973 fn each(&self) -> impl DoubleEndedIterator<Item = (usize, T, &[T])> {
975 self.blocks.iter().flat_map(move |b| {
976 let rows = self.m - b.k0;
977 (0..b.nb).map(move |c| (b.k0 + c, b.t[c * b.nb + c], &b.vt[c * rows + c + 1..(c + 1) * rows]))
978 })
979 }
980
981 fn q(&self) -> Vec<T> {
983 let (m, p) = (self.m, self.blocks.last().map_or(0, |b| b.k0 + b.nb));
984 if p <= ONE_AT_A_TIME {
985 let mut q = filled(m * p, T::zero());
986 one_at_a_time(m, p, self.each(), &mut Vec::new(), &mut q);
987 return q;
988 }
989 let ld = stride(p);
990 let mut q = filled(m * ld, T::zero());
991 for i in 0..p { q[i * ld + i] = T::one() }
992 for block in self.blocks.iter().rev() { block.apply(&mut q, ld, [block.k0, p], false) }
993 compact(&mut q, [m, p], ld);
994 q
995 }
996
997 fn apply(&self, x: &mut [T], k: usize) { for block in self.blocks.iter().rev() { block.apply(x, k, [0, k], false) } }
999
1000 fn apply_adjoint(&self, x: &mut [T], k: usize) {
1002 if k > FEW_COLUMNS { for block in &self.blocks { block.apply(x, k, [0, k], true) } return; }
1003 let m = self.m;
1004 let mut xt = transpose(x, m, k);
1005 for (row, tau, v) in self.each() {
1006 if tau == T::zero() { continue; }
1007 for column in xt.chunks_exact_mut(m) { apply_reflector(&mut column[row..], tau, v) }
1008 }
1009 transpose_into(&xt, m, x, k, [k, m], |x| x);
1010 }
1011}
1012
1013fn leaf([m, n]: [usize; 2], complete: bool) -> bool { !complete && m.min(n) <= LEAF }
1015
1016fn qr_factor<T: Tiled>(a: &[T], [m, n]: [usize; 2], complete: bool) -> (Vec<T>, Vec<T>) {
1018 if !complete && tall_for_qr([m, n]) && let Some((r, q)) = tall_qr(a, [m, n], Tall::Q) { return (q, r); }
1019 if leaf([m, n], complete) {
1020 let p = m.min(n);
1021 let (mut q, mut r) = (filled(m * p, T::zero()), filled(p * n, T::zero()));
1022 leaf_qr(a, [m, n], &mut QrWork::default(), &mut q, &mut r);
1023 return (q, r);
1024 }
1025 let (reflections, mut r) = reflect(a, [m, n]);
1026 if !complete { return (reflections.q(), r); }
1027 let mut q = identity(m);
1028 reflections.apply(&mut q, m);
1029 r.resize(m * n, T::zero());
1030 (q, r)
1031}
1032
1033fn qr_into<T: Tiled>(a: &[T], size: [usize; 2], complete: bool, work: &mut QrWork<T>, q: &mut [T], r: &mut [T]) {
1035 if leaf(size, complete) && !tall_for_qr(size) { return leaf_qr(a, size, work, q, r); }
1036 let (q0, r0) = qr_factor(a, size, complete);
1037 q.copy_from_slice(&q0);
1038 r.copy_from_slice(&r0);
1039}
1040
1041fn leaf_qr<T: Tiled>(a: &[T], [m, n]: [usize; 2], work: &mut QrWork<T>, q: &mut [T], r: &mut [T]) {
1043 let (p, scale) = (m.min(n), safe_scale(a));
1044 let QrWork { at, taus, qt } = work;
1045 at.resize(m * n, T::zero());
1046 transpose_into(a, n, at, m, [m, n], |x| x);
1047 if let Some(s) = scale { for x in at.iter_mut() { *x = x.mul_real(s) } }
1048 taus.clear();
1049 taus.resize(p, T::zero());
1050 reflect_columns(at, m, [0, p, n], |c, tau| taus[c] = tau);
1051 for i in 0..p {
1053 for j in 0..n { r[i * n + j] = if j < i { T::zero() } else { scale.map_or(at[j * m + i], |s| at[j * m + i].div_real(s)) } }
1054 }
1055 one_at_a_time(m, p, (0..p).map(|c| (c, taus[c], &at[c * m + c + 1..(c + 1) * m])), qt, q)
1056}
1057
1058fn one_at_a_time<'a, T: Tiled + 'a>(m: usize, p: usize, reflections: impl DoubleEndedIterator<Item = (usize, T, &'a [T])>, qt: &mut Vec<T>, q: &mut [T]) {
1060 qt.clear();
1061 qt.extend((0..p * m).map(|k| if k / m == k % m { T::one() } else { T::zero() }));
1062 for (row, tau, v) in reflections.rev() {
1063 if tau == T::zero() { continue; }
1064 for column in qt[row * m..].chunks_exact_mut(m) { apply_reflector(&mut column[row..], tau.conj(), v) }
1066 }
1067 transpose_into(qt, m, q, p, [p, m], |x| x);
1068}
1069
1070fn svd_solve<T: Tiled>(svd: SvdRef<'_, T>, b: Option<&[T]>, k: usize) -> Vec<T> {
1072 let ([m, n], p) = (svd.size, svd.s.len());
1073 if let Some(b) = b { assert_eq!(b.len(), m * k, "`b` must hold `m`×`k` items") }
1074 assert!(svd.u.len() == m * p && svd.v.len() == n * p, "the SVD must hold its singular vectors");
1075 pseudo_solve(svd, b, k, m.max(n))
1076}
1077
1078fn pseudo_solve<T: Tiled>(svd: SvdRef<'_, T>, b: Option<&[T]>, k: usize, scale: usize) -> Vec<T> {
1080 let ([m, n], p) = (svd.size, svd.s.len());
1081 if k == 0 { return Vec::new(); }
1082 let floor = cutoff(&svd.s, scale);
1083 let mut y = match b { Some(b) => T::multiply(Matrix::rows(svd.u, p).adjoint(), Matrix::rows(b, k), [p, m, k]), None => adjoint(&svd.u, m, p) };
1084 for (row, &s) in y.chunks_exact_mut(k).zip(svd.s) { let d = if s > floor { s.recip() } else { T::Real::zero() }; for x in row { *x = x.mul_real(d) } }
1085 T::multiply(Matrix::rows(svd.v, p), Matrix::rows(&y, k), [n, p, k])
1086}
1087
1088fn cutoff<R: crate::Float>(s: &[R], scale: usize) -> R { R::epsilon() * R::narrow(scale as f64) * s.first().copied().unwrap_or(R::zero()) }
1090
1091fn lstsq_solve<T: Tiled>(a: &[T], b: &[T], [m, n]: [usize; 2], k: usize) -> (Vec<T>, Vec<T::Real>) {
1093 assert_eq!(b.len(), m * k, "`b` must hold `m`×`k` items");
1094 if m < n { let svd = svd_factor(a, [m, n], true); return (pseudo_solve(svd.as_ref(), Some(b), k, n), svd.s); }
1095 let reduced = tall_for_qr([m, n]).then(|| tall_qr(a, [m, n], Tall::Project(b, k))).flatten();
1096 let (r, y) = reduced.unwrap_or_else(|| {
1097 let (reflections, r) = reflect(a, [m, n]);
1098 let mut y = b.to_vec();
1099 reflections.apply_adjoint(&mut y, k);
1100 y.truncate(n * k);
1101 (r, y)
1102 });
1103 let svd = svd_factor(&r, [n, n], true);
1104 (pseudo_solve(svd.as_ref(), Some(&y), k, m), svd.s)
1105}
1106
1107const TALL_BLOCK: usize = 1 << 18;
1109const NEARLY_ORTHONORMAL: f64 = 0.25;
1111
1112fn tall_for_qr([m, n]: [usize; 2]) -> bool { (n >= if super::accelerated() { 16 } else { 20 } && m >= (4 * n).max(n * n / 4)) || threaded_work(m, n, 1 << 21) }
1114
1115fn threaded_work(m: usize, n: usize, work: usize) -> bool {
1117 (super::accelerated() || !parallel::serial()) && m >= 16 * n && m.saturating_mul(n.saturating_mul(n)) >= work
1118}
1119
1120enum Tall<'a, T> { Q, R, Project(&'a [T], usize) }
1122
1123fn tall_qr<T: Tiled>(a: &[T], [m, n]: [usize; 2], want: Tall<'_, T>) -> Option<(Vec<T>, Vec<T>)> {
1127 let scale = safe_scale(a);
1128 let scaled = match scale { Some(s) => Cow::Owned(a.iter().map(|x| x.mul_real(s)).collect()), None => Cow::Borrowed(a) };
1129 let (a, rows, k) = (&scaled[..], (TALL_BLOCK / n).max(1), if let Tall::Project(_, k) = want { k } else { 0 });
1130 let count = parallel::threads().min(m.div_ceil(rows));
1131 let range = |p: usize| p * m / count..(p + 1) * m / count;
1132 let mut grams: Vec<_> = (0..count).map(|_| filled(n * n, T::zero())).collect();
1133 parallel::pieces_of(&mut grams, 1, usize::MAX, &|p, g| {
1134 let (start, end) = (range(p).start, range(p).end);
1135 for x in (start..end).step_by(rows).map(|s| &a[s * n..end.min(s + rows) * n]) {
1136 T::gram(&mut g[0], n, Matrix::rows(x, n).adjoint(), [n, x.len() / n], T::Real::one(), T::Real::one());
1137 }
1138 true
1139 });
1140 let l1 = cholesky_factor(&summed(grams.iter()), n)?;
1141 let mut q = if matches!(want, Tall::Q) { a.to_vec() } else { Vec::new() };
1142 let mut rest = &mut q[..];
1143 let outs: Vec<&mut [T]> = (0..count).map(|p| {
1144 let (out, tail) = std::mem::take(&mut rest).split_at_mut(if matches!(want, Tall::Q) { range(p).len() * n } else { 0 });
1145 rest = tail;
1146 out
1147 }).collect();
1148 let mut parts: Vec<_> = outs.into_iter().map(|out| (out, filled(n * n, T::zero()), filled(n * k, T::zero()), Vec::new())).collect();
1149 parallel::pieces_of(&mut parts, 1, usize::MAX, &|p, part| {
1150 let (out, gram, y, scratch) = &mut part[0];
1151 let (start, end) = (range(p).start, range(p).end);
1152 for r in (start..end).step_by(rows).map(|s| s..end.min(s + rows)) {
1153 let h = r.len();
1154 let src = &a[r.start * n..r.end * n];
1155 let q1 = if out.is_empty() {
1156 scratch.clear();
1157 scratch.extend_from_slice(src);
1158 &mut scratch[..]
1159 } else { &mut out[(r.start - start) * n..(r.end - start) * n] };
1160 T::triangular(q1, n, Matrix::rows(&l1, n).adjoint(), [h, n], Side::Right, Triangle::Lower);
1161 let qt = Matrix::rows(q1, n).adjoint();
1162 T::gram(gram, n, qt, [n, h], T::Real::one(), T::Real::one());
1163 if let Tall::Project(b, k) = want { T::gemm(y, k, qt, Matrix::rows(&b[r.start * k..r.end * k], k), [n, h, k], T::one(), T::one()); }
1164 }
1165 true
1166 });
1167 let (g2, mut y) = (summed(parts.iter().map(|part| &part.1)), summed(parts.iter().map(|part| &part.2)));
1168 let far = (0..n * n).filter(|&c| c % n <= c / n).fold(T::Real::zero(), |sum, c| {
1169 let (diagonal, v) = (c % n == c / n, g2[c].modulus());
1170 let v = if diagonal { (g2[c] - T::one()).modulus() } else { v };
1171 sum + if diagonal { v * v } else { v * v + v * v }
1172 });
1173 if !(far.sqrt() <= <T::Real as crate::Float>::narrow(NEARLY_ORTHONORMAL)) { return None; }
1174 let l2 = cholesky_factor(&g2, n)?;
1175 let l = T::multiply(Matrix::rows(&l1, n), Matrix::rows(&l2, n), [n, n, n]);
1176 let mut r = adjoint(&l, n, n);
1177 if let Some(s) = scale { for v in &mut r { *v = v.div_real(s) } }
1178 if !matches!(want, Tall::Q) { T::triangular(&mut y, k, Matrix::rows(&l2, n), [n, k], Side::Left, Triangle::Lower); return Some((r, y)); }
1179 parallel::pieces_of(&mut q, rows * n, usize::MAX, &|_, band| {
1180 let h = band.len() / n;
1181 T::triangular(band, n, Matrix::rows(&l2, n).adjoint(), [h, n], Side::Right, Triangle::Lower);
1182 true
1183 });
1184 Some((r, q))
1185}
1186
1187fn summed<'a, T: Scalar + 'a>(mut parts: impl Iterator<Item = &'a Vec<T>>) -> Vec<T> {
1189 let mut sum = parts.next().cloned().unwrap_or_default();
1190 for part in parts { for (s, &v) in sum.iter_mut().zip(part) { *s = *s + v } }
1191 sum
1192}