use crate::{
assert,
linalg::{
matmul::triangular::{self, BlockStructure},
temp_mat_req, temp_mat_uninit,
triangular_inverse::invert_lower_triangular,
},
ComplexField, Entity, MatMut, MatRef, Parallelism,
};
use dyn_stack::{PodStack, SizeOverflow, StackReq};
use reborrow::*;
fn invert_lower_impl<E: ComplexField>(
dst: MatMut<'_, E>,
cholesky_factor: Option<MatRef<'_, E>>,
parallelism: Parallelism,
stack: &mut PodStack,
) {
let cholesky_factor = match cholesky_factor {
Some(cholesky_factor) => cholesky_factor,
None => dst.rb(),
};
let n = cholesky_factor.nrows();
let (mut tmp, _) = temp_mat_uninit::<E>(n, n, stack);
let mut tmp = tmp.as_mut();
invert_lower_triangular(tmp.rb_mut(), cholesky_factor, parallelism);
triangular::matmul(
dst,
BlockStructure::TriangularLower,
tmp.rb().adjoint(),
BlockStructure::TriangularUpper,
tmp.rb(),
BlockStructure::TriangularLower,
None,
E::faer_one(),
parallelism,
);
}
pub fn invert_lower_req<E: Entity>(
dimension: usize,
parallelism: Parallelism,
) -> Result<StackReq, SizeOverflow> {
let _ = parallelism;
temp_mat_req::<E>(dimension, dimension)
}
pub fn invert_lower_in_place_req<E: Entity>(
dimension: usize,
parallelism: Parallelism,
) -> Result<StackReq, SizeOverflow> {
invert_lower_req::<E>(dimension, parallelism)
}
#[track_caller]
pub fn invert_lower_in_place<E: ComplexField>(
cholesky_factor: MatMut<'_, E>,
parallelism: Parallelism,
stack: &mut PodStack,
) {
assert!(cholesky_factor.nrows() == cholesky_factor.ncols());
invert_lower_impl(cholesky_factor, None, parallelism, stack);
}
#[track_caller]
pub fn invert_lower<E: ComplexField>(
dst: MatMut<'_, E>,
cholesky_factor: MatRef<'_, E>,
parallelism: Parallelism,
stack: &mut PodStack,
) {
assert!(all(
cholesky_factor.nrows() == cholesky_factor.ncols(),
dst.nrows() == cholesky_factor.nrows(),
dst.ncols() == cholesky_factor.ncols(),
));
invert_lower_impl(dst, Some(cholesky_factor), parallelism, stack);
}