use dyn_stack::{PodStack, SizeOverflow, StackReq};
use faer_core::{
assert,
householder::apply_block_householder_sequence_transpose_on_the_right_in_place_with_conj,
inverse::invert_upper_triangular,
permutation::{permute_cols_in_place_req, permute_rows_in_place, Index, PermutationRef},
temp_mat_req, temp_mat_uninit, unzipped, zipped, ComplexField, Conj, Entity, MatMut, MatRef,
Parallelism,
};
use reborrow::*;
#[track_caller]
pub fn invert<I: Index, E: ComplexField>(
dst: MatMut<'_, E>,
qr_factors: MatRef<'_, E>,
householder_factor: MatRef<'_, E>,
col_perm: PermutationRef<'_, I, E>,
parallelism: Parallelism,
stack: PodStack<'_>,
) {
assert!(all(
qr_factors.nrows() == qr_factors.ncols(),
col_perm.len() == qr_factors.nrows(),
dst.ncols() == qr_factors.ncols(),
dst.nrows() == qr_factors.nrows(),
householder_factor.ncols() == Ord::min(qr_factors.nrows(), qr_factors.ncols()),
householder_factor.nrows() > 0
));
let mut dst = dst;
let mut stack = stack;
invert_upper_triangular(dst.rb_mut(), qr_factors, parallelism);
zipped!(dst.rb_mut())
.for_each_triangular_lower(faer_core::zip::Diag::Skip, |unzipped!(mut dst)| {
dst.write(E::faer_zero())
});
apply_block_householder_sequence_transpose_on_the_right_in_place_with_conj(
qr_factors,
householder_factor,
Conj::Yes,
dst.rb_mut(),
parallelism,
stack.rb_mut(),
);
permute_rows_in_place(dst.rb_mut(), col_perm.inverse(), stack)
}
#[track_caller]
pub fn invert_in_place<I: Index, E: ComplexField>(
qr_factors: MatMut<'_, E>,
householder_factor: MatRef<'_, E>,
col_perm: PermutationRef<'_, I, E>,
parallelism: Parallelism,
stack: PodStack<'_>,
) {
let (mut dst, stack) = temp_mat_uninit::<E>(qr_factors.nrows(), qr_factors.ncols(), stack);
let mut dst = dst.as_mut();
invert(
dst.rb_mut(),
qr_factors.rb(),
householder_factor,
col_perm,
parallelism,
stack,
);
zipped!(qr_factors, dst.rb()).for_each(|unzipped!(mut dst, src)| dst.write(src.read()));
}
pub fn invert_req<I: Index, E: Entity>(
qr_nrows: usize,
qr_ncols: usize,
blocksize: usize,
parallelism: Parallelism,
) -> Result<StackReq, SizeOverflow> {
let _ = qr_nrows;
let _ = parallelism;
StackReq::try_any_of([
temp_mat_req::<E>(blocksize, qr_ncols)?,
permute_cols_in_place_req::<I, E>(qr_nrows, qr_ncols)?,
])
}
pub fn invert_in_place_req<I: Index, E: Entity>(
qr_nrows: usize,
qr_ncols: usize,
blocksize: usize,
parallelism: Parallelism,
) -> Result<StackReq, SizeOverflow> {
StackReq::try_all_of([
temp_mat_req::<E>(qr_nrows, qr_ncols)?,
invert_req::<I, E>(qr_nrows, qr_ncols, blocksize, parallelism)?,
])
}
#[cfg(test)]
mod tests {
use super::*;
use crate::col_pivoting::compute::{qr_in_place, qr_in_place_req, recommended_blocksize};
use assert_approx_eq::assert_approx_eq;
use faer_core::{assert, c64, mul::matmul, Mat};
use rand::prelude::*;
use std::cell::RefCell;
macro_rules! make_stack {
($req: expr) => {
::dyn_stack::PodStack::new(&mut ::dyn_stack::GlobalPodBuffer::new($req.unwrap()))
};
}
type E = c64;
thread_local! {
static RNG: RefCell<StdRng> = RefCell::new(StdRng::seed_from_u64(0));
}
fn random_value() -> E {
RNG.with(|rng| {
let mut rng = rng.borrow_mut();
let rng = &mut *rng;
E::new(rng.gen(), rng.gen())
})
}
#[test]
fn test_invert() {
for n in [31, 32, 48, 65] {
let mat = Mat::from_fn(n, n, |_, _| random_value());
let blocksize = recommended_blocksize::<E>(n, n);
let mut qr = mat.clone();
let mut householder_factor = Mat::zeros(blocksize, n);
let parallelism = faer_core::Parallelism::Rayon(0);
let mut perm = vec![0usize; n];
let mut perm_inv = vec![0; n];
let (_, perm) = qr_in_place(
qr.as_mut(),
householder_factor.as_mut(),
&mut perm,
&mut perm_inv,
parallelism,
make_stack!(qr_in_place_req::<usize, E>(
n,
n,
blocksize,
parallelism,
Default::default()
)),
Default::default(),
);
let mut inv = Mat::zeros(n, n);
invert(
inv.as_mut(),
qr.as_ref(),
householder_factor.as_ref(),
perm.rb(),
parallelism,
make_stack!(invert_req::<usize, E>(n, n, blocksize, parallelism)),
);
let mut eye = Mat::zeros(n, n);
matmul(
eye.as_mut(),
inv.as_ref(),
mat.as_ref(),
None,
E::faer_one(),
Parallelism::None,
);
for i in 0..n {
for j in 0..n {
let target = if i == j {
E::faer_one()
} else {
E::faer_zero()
};
assert_approx_eq!(eye.read(i, j), target);
}
}
}
}
}