use super::blas::{Op, Side, Uplo, gemm, trsm, trsm_into, with_scratch};
use super::{Matrix, Scalar};
use crate::math::complex::Float;
use core::hint::select_unpredictable;
#[derive(Clone, Debug)]
pub struct Lu<T> {
lu: Matrix<T>,
perm: Vec<usize>,
odd: bool,
}
impl<T: Scalar> Matrix<T> {
pub fn lu(&self) -> Lu<T> {
let (m, n) = self.shape();
let mut lu = self.clone();
let mut perm: Vec<usize> = (0..m).collect();
let mut odd = false;
if m == n && n <= TINY {
match n {
0 => {}
1 => odd = tiny::<T, 1>(&mut lu.data, &mut perm),
2 => odd = tiny::<T, 2>(&mut lu.data, &mut perm),
3 => odd = tiny::<T, 3>(&mut lu.data, &mut perm),
4 => odd = tiny::<T, 4>(&mut lu.data, &mut perm),
5 => odd = tiny::<T, 5>(&mut lu.data, &mut perm),
6 => odd = tiny::<T, 6>(&mut lu.data, &mut perm),
7 => odd = tiny::<T, 7>(&mut lu.data, &mut perm),
8 => odd = tiny::<T, 8>(&mut lu.data, &mut perm),
9 => odd = tiny::<T, 9>(&mut lu.data, &mut perm),
10 => odd = tiny::<T, 10>(&mut lu.data, &mut perm),
11 => odd = tiny::<T, 11>(&mut lu.data, &mut perm),
12 => odd = tiny::<T, 12>(&mut lu.data, &mut perm),
13 => odd = tiny::<T, 13>(&mut lu.data, &mut perm),
14 => odd = tiny::<T, 14>(&mut lu.data, &mut perm),
15 => odd = tiny::<T, 15>(&mut lu.data, &mut perm),
_ => odd = tiny::<T, 16>(&mut lu.data, &mut perm),
}
return Lu { lu, perm, odd };
}
with_pivots(m.min(n), |pivots| {
factor_all(&mut lu.data, m, n, pivots);
for (i, &p) in pivots.iter().enumerate() {
if p != i {
perm.swap(i, p);
odd = !odd;
}
}
});
Lu { lu, perm, odd }
}
#[track_caller]
pub fn det(&self) -> T {
assert!(self.is_square(), "the matrix is not square");
self.lu().det()
}
#[track_caller]
pub fn inverse(&self) -> Option<Self> {
assert!(self.is_square(), "the matrix is not square");
match self.rows() {
0 => Some(self.clone()),
1 => tiny_inverse_inline::<T, 1>(self),
2 => tiny_inverse_inline::<T, 2>(self),
3 => tiny_inverse_inline::<T, 3>(self),
4 => tiny_inverse_inline::<T, 4>(self),
5 => tiny_inverse::<T, 5>(self),
6 => tiny_inverse::<T, 6>(self),
7 => tiny_inverse::<T, 7>(self),
8 => tiny_inverse::<T, 8>(self),
9 => tiny_inverse::<T, 9>(self),
10 => tiny_inverse::<T, 10>(self),
11 => tiny_inverse::<T, 11>(self),
12 => tiny_inverse::<T, 12>(self),
n => {
with_scratch(n * n, |lu| {
lu.copy_from_slice(&self.data);
with_pivots(n, |pivots| {
factor_all(lu, n, n, pivots);
if (0..n).any(|i| lu[i + i * n] == T::zero()) {
return None;
}
let mut x = Matrix::zeros(n, n);
inverse_from(lu, n, &mut x.data);
for (j, &p) in pivots.iter().enumerate().rev() {
if p != j {
let (a, b) = x.data.split_at_mut(p * n);
a[j * n..j * n + n].swap_with_slice(&mut b[..n]);
}
}
Some(x)
})
})
}
}
}
#[track_caller]
pub fn solve(&self, b: &Matrix<T>) -> Option<Self> {
assert!(self.is_square(), "the matrix is not square");
self.lu().solve(b)
}
#[track_caller]
pub fn powi(&self, n: i32) -> Option<Self> {
if n >= 0 {
Some(self.pow(n.unsigned_abs()))
} else {
Some(self.inverse()?.pow(n.unsigned_abs()))
}
}
}
impl<T: Scalar> Lu<T> {
pub fn l(&self) -> Matrix<T> {
let (m, n) = self.lu.shape();
Matrix::from_fn(m, m.min(n), |i, j| match i.cmp(&j) {
core::cmp::Ordering::Greater => self.lu[(i, j)],
core::cmp::Ordering::Equal => T::one(),
core::cmp::Ordering::Less => T::zero(),
})
}
pub fn u(&self) -> Matrix<T> {
let (m, n) = self.lu.shape();
Matrix::from_fn(m.min(n), n, |i, j| {
if i <= j { self.lu[(i, j)] } else { T::zero() }
})
}
pub fn p(&self) -> Matrix<T> {
let m = self.perm.len();
Matrix::from_fn(m, m, |i, j| {
if self.perm[i] == j {
T::one()
} else {
T::zero()
}
})
}
pub fn permutation(&self) -> &[usize] {
&self.perm
}
pub fn is_singular(&self) -> bool {
let (m, n) = self.lu.shape();
(0..m.min(n)).any(|i| self.lu[(i, i)] == T::zero())
}
#[track_caller]
pub fn det(&self) -> T {
assert!(self.lu.is_square(), "the matrix is not square");
let mut d = T::one();
for i in 0..self.lu.rows() {
d *= self.lu[(i, i)];
}
if self.odd { -d } else { d }
}
#[track_caller]
pub fn solve(&self, b: &Matrix<T>) -> Option<Matrix<T>> {
assert!(self.lu.is_square(), "the matrix is not square");
let n = self.lu.rows();
assert_eq!(b.rows(), n, "the right-hand side has not as many rows");
if self.is_singular() {
return None;
}
let mut x = Matrix::from_fn(n, b.cols(), |i, j| b[(self.perm[i], j)]);
let cols = x.cols();
trsm(
Side::Left,
Uplo::Lower,
Op::N,
true,
n,
cols,
&self.lu.data,
n,
&mut x.data,
n,
);
trsm(
Side::Left,
Uplo::Upper,
Op::N,
false,
n,
cols,
&self.lu.data,
n,
&mut x.data,
n,
);
Some(x)
}
#[track_caller]
pub fn inverse(&self) -> Option<Matrix<T>> {
assert!(self.lu.is_square(), "the matrix is not square");
if self.is_singular() {
return None;
}
let n = self.lu.rows();
let mut x = Matrix::zeros(n, n);
inverse_from(&self.lu.data, n, &mut x.data);
let mut to = self.perm.clone();
for j in 0..n {
while to[j] != j {
let k = to[j];
let (a, b) = x.data.split_at_mut(j.max(k) * n);
a[j.min(k) * n..j.min(k) * n + n].swap_with_slice(&mut b[..n]);
to.swap(j, k);
}
}
Some(x)
}
}
fn inverse_from<T: Scalar>(lu: &[T], n: usize, x: &mut [T]) {
upper_inverse(lu, n, x, n, n);
trsm(Side::Right, Uplo::Lower, Op::N, true, n, n, lu, n, x, n);
}
const INVERSE_LEAF: usize = 16;
fn upper_inverse<T: Scalar>(u: &[T], ldu: usize, w: &mut [T], ldw: usize, n: usize) {
if n <= INVERSE_LEAF {
for j in 0..n {
for i in 0..=j {
w[i + j * ldw] = if i == j { T::one() } else { T::zero() };
}
}
trsm(Side::Left, Uplo::Upper, Op::N, false, n, n, u, ldu, w, ldw);
return;
}
let n1 = split(n);
let n2 = n - n1;
let (w1, w2) = w.split_at_mut(n1 * ldw);
for j in 0..n2 {
let (src, dst) = (&u[(n1 + j) * ldu..], &mut w2[j * ldw..]);
for (d, s) in dst[..n1].iter_mut().zip(&src[..n1]) {
*d = -*s;
}
}
let u22 = &u[n1 + n1 * ldu..];
trsm(
Side::Left,
Uplo::Upper,
Op::N,
false,
n1,
n2,
u,
ldu,
w2,
ldw,
);
trsm(
Side::Right,
Uplo::Upper,
Op::N,
false,
n1,
n2,
u22,
ldu,
w2,
ldw,
);
upper_inverse(u, ldu, w1, ldw, n1);
upper_inverse(u22, ldu, &mut w2[n1..], ldw, n2);
}
#[inline(never)]
fn tiny_inverse<T: Scalar, const N: usize>(a: &Matrix<T>) -> Option<Matrix<T>> {
tiny_inverse_inline::<T, N>(a)
}
#[inline(always)]
#[allow(clippy::needless_range_loop)]
fn tiny_inverse_inline<T: Scalar, const N: usize>(a: &Matrix<T>) -> Option<Matrix<T>> {
let mut lu = [[T::zero(); N]; N];
lu.as_flattened_mut().copy_from_slice(&a.data);
let mut perm: [usize; N] = core::array::from_fn(|i| i);
let mut r = [T::zero(); N];
for j in 0..N {
let mut p = j;
let mut best = lu[j][j].abs1();
for i in j + 1..N {
let v = lu[j][i].abs1();
let greater = v > best;
best = select_unpredictable(greater, v, best);
p = select_unpredictable(greater, i, p);
}
for i in j + 1..N {
let swap = p == i;
for col in &mut lu {
let (x, y) = (col[j], col[i]);
col[j] = select_unpredictable(swap, y, x);
col[i] = select_unpredictable(swap, x, y);
}
let (x, y) = (perm[j], perm[i]);
perm[j] = select_unpredictable(swap, y, x);
perm[i] = select_unpredictable(swap, x, y);
}
let d = lu[j][j];
if d == T::zero() {
return None;
}
r[j] = T::one() / d;
if r[j].is_finite() {
for i in j + 1..N {
lu[j][i] *= r[j];
}
} else {
for i in j + 1..N {
lu[j][i] /= d;
}
}
for c in j + 1..N {
let u = lu[c][j];
for i in j + 1..N {
lu[c][i] = lu[c][i].mul_sub(lu[j][i], u);
}
}
}
let mut w = [[T::zero(); N]; N];
for j in 0..N {
w[j][j] = r[j];
for i in (0..j).rev() {
let mut s = T::zero();
for k in i + 1..=j {
s = s.mul_sub(lu[k][i], w[j][k]);
}
w[j][i] = s * r[i];
}
}
let mut x = [[T::zero(); N]; N];
for j in (0..N).rev() {
let mut col = w[j];
for k in j + 1..N {
let l = lu[j][k];
for (c, x) in col.iter_mut().zip(&x[k]) {
*c = c.mul_sub(*x, l);
}
}
x[j] = col;
}
let mut inverse = vec![T::zero(); N * N];
for (j, &p) in perm.iter().enumerate() {
inverse[p * N..p * N + N].copy_from_slice(&x[j]);
}
Some(Matrix::from_column_major(N, N, inverse))
}
const STACK_PIVOTS: usize = 32;
fn with_pivots<X>(k: usize, f: impl FnOnce(&mut [usize]) -> X) -> X {
if k <= STACK_PIVOTS {
f(&mut [0; STACK_PIVOTS][..k])
} else {
f(&mut vec![0; k])
}
}
fn factor_all<T: Scalar>(a: &mut [T], m: usize, n: usize, pivots: &mut [usize]) {
let k = m.min(n);
if k == 0 {
return;
}
if k <= UNBLOCKED {
unblocked(a, m, k, m, pivots);
} else {
with_scratch(work_size(k), |work| factor(a, m, k, m, pivots, work));
}
if n > k {
let (left, right) = a.split_at_mut(k * m);
swap_rows(right, m, n - k, pivots, 0);
trsm(
Side::Left,
Uplo::Lower,
Op::N,
true,
k,
n - k,
left,
m,
right,
m,
);
}
}
const NARROW: usize = 16;
const UNBLOCKED: usize = 24;
fn split(n: usize) -> usize {
(n / 2).next_multiple_of(NARROW).min(n - 1)
}
fn work_size(n: usize) -> usize {
if n <= NARROW {
return 0;
}
let n1 = split(n);
(n1 * (n - n1)).max(work_size(n1)).max(work_size(n - n1))
}
fn factor<T: Scalar>(
a: &mut [T],
m: usize,
n: usize,
ld: usize,
pivots: &mut [usize],
work: &mut [T],
) {
if n <= NARROW {
unblocked(a, m, n, ld, pivots);
return;
}
let n1 = split(n);
let n2 = n - n1;
factor(a, m, n1, ld, &mut pivots[..n1], work);
let (left, right) = a.split_at_mut(n1 * ld);
swap_rows(right, ld, n2, &pivots[..n1], 0);
let u12 = &mut work[..n1 * n2];
trsm_into(
Uplo::Lower,
Op::N,
true,
n1,
n2,
left,
ld,
right,
ld,
u12,
n1,
);
gemm(
m - n1,
n2,
n1,
-T::one(),
&left[n1..],
ld,
Op::N,
u12,
n1,
Op::N,
T::one(),
&mut right[n1..],
ld,
);
factor(&mut right[n1..], m - n1, n2, ld, &mut pivots[n1..], work);
for p in &mut pivots[n1..] {
*p += n1;
}
swap_rows(left, ld, n1, &pivots[n1..], n1);
}
fn unblocked<T: Scalar>(a: &mut [T], m: usize, n: usize, ld: usize, pivots: &mut [usize]) {
for j in 0..n {
let p = j + pivot(&a[j * ld + j..j * ld + m]);
pivots[j] = p;
if p != j {
for c in 0..n {
a.swap(j + c * ld, p + c * ld);
}
}
let (left, right) = a.split_at_mut(((j + 1) * ld).min(a.len()));
let col = &mut left[j * ld + j..j * ld + m];
let (d, l) = col.split_first_mut().unwrap();
if *d == T::zero() {
continue;
}
let r = T::one() / *d;
if r.is_finite() {
for x in l.iter_mut() {
*x *= r;
}
} else {
for x in l.iter_mut() {
*x /= *d;
}
}
for c in 0..n - j - 1 {
let col = &mut right[c * ld + j..c * ld + m];
let (u, x) = col.split_first_mut().unwrap();
let u = *u;
if u == T::zero() {
continue;
}
for (x, l) in x.iter_mut().zip(l.iter()) {
*x = x.mul_sub(*l, u);
}
}
}
}
const TINY: usize = 16;
#[inline(never)]
#[allow(clippy::needless_range_loop)]
fn tiny<T: Scalar, const N: usize>(data: &mut [T], perm: &mut [usize]) -> bool {
let mut a = [[T::zero(); N]; N];
a.as_flattened_mut().copy_from_slice(&data[..N * N]);
let perm: &mut [usize; N] = (&mut perm[..N]).try_into().unwrap();
let mut odd = false;
for j in 0..N {
let mut p = j;
let mut best = a[j][j].abs1();
for i in j + 1..N {
let v = a[j][i].abs1();
let greater = v > best;
best = select_unpredictable(greater, v, best);
p = select_unpredictable(greater, i, p);
}
odd ^= p != j;
for i in j + 1..N {
let swap = p == i;
for col in &mut a {
let (x, y) = (col[j], col[i]);
col[j] = select_unpredictable(swap, y, x);
col[i] = select_unpredictable(swap, x, y);
}
let (x, y) = (perm[j], perm[i]);
perm[j] = select_unpredictable(swap, y, x);
perm[i] = select_unpredictable(swap, x, y);
}
let d = a[j][j];
if d == T::zero() {
continue;
}
let r = T::one() / d;
if r.is_finite() {
for i in j + 1..N {
a[j][i] *= r;
}
} else {
for i in j + 1..N {
a[j][i] /= d;
}
}
for c in j + 1..N {
let u = a[c][j];
if u == T::zero() {
continue;
}
for i in j + 1..N {
a[c][i] = a[c][i].mul_sub(a[j][i], u);
}
}
}
data[..N * N].copy_from_slice(a.as_flattened());
odd
}
fn pivot<T: Scalar>(col: &[T]) -> usize {
let first = col[0].abs1();
if first.is_nan() {
return 0;
}
let best = col.iter().fold(first, |b, x| b.max(x.abs1()));
col.iter().position(|x| x.abs1() == best).unwrap_or(0)
}
fn swap_rows<T>(a: &mut [T], ld: usize, n: usize, pivots: &[usize], first: usize) {
for c in 0..n {
let col = &mut a[c * ld..];
for (i, &p) in pivots.iter().enumerate() {
let i = first + i;
if p != i {
col.swap(i, p);
}
}
}
}