use simba::scalar::ComplexField;
use simba::simd::SimdComplexField;
use crate::base::allocator::Allocator;
use crate::base::constraint::{SameNumberOfRows, ShapeConstraint};
use crate::base::dimension::{Dim, U1};
use crate::base::storage::{Storage, StorageMut};
use crate::base::{DVectorSlice, DefaultAllocator, Matrix, OMatrix, SquareMatrix, Vector};
impl<T: ComplexField, D: Dim, S: Storage<T, D, D>> SquareMatrix<T, D, S> {
#[must_use = "Did you mean to use solve_lower_triangular_mut()?"]
#[inline]
pub fn solve_lower_triangular<R2: Dim, C2: Dim, S2>(
&self,
b: &Matrix<T, R2, C2, S2>,
) -> Option<OMatrix<T, R2, C2>>
where
S2: Storage<T, R2, C2>,
DefaultAllocator: Allocator<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let mut res = b.clone_owned();
if self.solve_lower_triangular_mut(&mut res) {
Some(res)
} else {
None
}
}
#[must_use = "Did you mean to use solve_upper_triangular_mut()?"]
#[inline]
pub fn solve_upper_triangular<R2: Dim, C2: Dim, S2>(
&self,
b: &Matrix<T, R2, C2, S2>,
) -> Option<OMatrix<T, R2, C2>>
where
S2: Storage<T, R2, C2>,
DefaultAllocator: Allocator<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let mut res = b.clone_owned();
if self.solve_upper_triangular_mut(&mut res) {
Some(res)
} else {
None
}
}
pub fn solve_lower_triangular_mut<R2: Dim, C2: Dim, S2>(
&self,
b: &mut Matrix<T, R2, C2, S2>,
) -> bool
where
S2: StorageMut<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let cols = b.ncols();
for i in 0..cols {
if !self.solve_lower_triangular_vector_mut(&mut b.column_mut(i)) {
return false;
}
}
true
}
fn solve_lower_triangular_vector_mut<R2: Dim, S2>(&self, b: &mut Vector<T, R2, S2>) -> bool
where
S2: StorageMut<T, R2, U1>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let dim = self.nrows();
for i in 0..dim {
let coeff;
unsafe {
let diag = self.get_unchecked((i, i)).clone();
if diag.is_zero() {
return false;
}
coeff = b.vget_unchecked(i).clone() / diag;
*b.vget_unchecked_mut(i) = coeff.clone();
}
b.rows_range_mut(i + 1..)
.axpy(-coeff, &self.slice_range(i + 1.., i), T::one());
}
true
}
pub fn solve_lower_triangular_with_diag_mut<R2: Dim, C2: Dim, S2>(
&self,
b: &mut Matrix<T, R2, C2, S2>,
diag: T,
) -> bool
where
S2: StorageMut<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
if diag.is_zero() {
return false;
}
let dim = self.nrows();
let cols = b.ncols();
for k in 0..cols {
let mut bcol = b.column_mut(k);
for i in 0..dim - 1 {
let coeff = unsafe { bcol.vget_unchecked(i).clone() } / diag.clone();
bcol.rows_range_mut(i + 1..)
.axpy(-coeff, &self.slice_range(i + 1.., i), T::one());
}
}
true
}
pub fn solve_upper_triangular_mut<R2: Dim, C2: Dim, S2>(
&self,
b: &mut Matrix<T, R2, C2, S2>,
) -> bool
where
S2: StorageMut<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let cols = b.ncols();
for i in 0..cols {
if !self.solve_upper_triangular_vector_mut(&mut b.column_mut(i)) {
return false;
}
}
true
}
fn solve_upper_triangular_vector_mut<R2: Dim, S2>(&self, b: &mut Vector<T, R2, S2>) -> bool
where
S2: StorageMut<T, R2, U1>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let dim = self.nrows();
for i in (0..dim).rev() {
let coeff;
unsafe {
let diag = self.get_unchecked((i, i)).clone();
if diag.is_zero() {
return false;
}
coeff = b.vget_unchecked(i).clone() / diag;
*b.vget_unchecked_mut(i) = coeff.clone();
}
b.rows_range_mut(..i)
.axpy(-coeff, &self.slice_range(..i, i), T::one());
}
true
}
#[must_use = "Did you mean to use tr_solve_lower_triangular_mut()?"]
#[inline]
pub fn tr_solve_lower_triangular<R2: Dim, C2: Dim, S2>(
&self,
b: &Matrix<T, R2, C2, S2>,
) -> Option<OMatrix<T, R2, C2>>
where
S2: Storage<T, R2, C2>,
DefaultAllocator: Allocator<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let mut res = b.clone_owned();
if self.tr_solve_lower_triangular_mut(&mut res) {
Some(res)
} else {
None
}
}
#[must_use = "Did you mean to use tr_solve_upper_triangular_mut()?"]
#[inline]
pub fn tr_solve_upper_triangular<R2: Dim, C2: Dim, S2>(
&self,
b: &Matrix<T, R2, C2, S2>,
) -> Option<OMatrix<T, R2, C2>>
where
S2: Storage<T, R2, C2>,
DefaultAllocator: Allocator<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let mut res = b.clone_owned();
if self.tr_solve_upper_triangular_mut(&mut res) {
Some(res)
} else {
None
}
}
pub fn tr_solve_lower_triangular_mut<R2: Dim, C2: Dim, S2>(
&self,
b: &mut Matrix<T, R2, C2, S2>,
) -> bool
where
S2: StorageMut<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let cols = b.ncols();
for i in 0..cols {
if !self.xx_solve_lower_triangular_vector_mut(
&mut b.column_mut(i),
|e| e,
|a, b| a.dot(b),
) {
return false;
}
}
true
}
pub fn tr_solve_upper_triangular_mut<R2: Dim, C2: Dim, S2>(
&self,
b: &mut Matrix<T, R2, C2, S2>,
) -> bool
where
S2: StorageMut<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let cols = b.ncols();
for i in 0..cols {
if !self.xx_solve_upper_triangular_vector_mut(
&mut b.column_mut(i),
|e| e,
|a, b| a.dot(b),
) {
return false;
}
}
true
}
#[must_use = "Did you mean to use ad_solve_lower_triangular_mut()?"]
#[inline]
pub fn ad_solve_lower_triangular<R2: Dim, C2: Dim, S2>(
&self,
b: &Matrix<T, R2, C2, S2>,
) -> Option<OMatrix<T, R2, C2>>
where
S2: Storage<T, R2, C2>,
DefaultAllocator: Allocator<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let mut res = b.clone_owned();
if self.ad_solve_lower_triangular_mut(&mut res) {
Some(res)
} else {
None
}
}
#[must_use = "Did you mean to use ad_solve_upper_triangular_mut()?"]
#[inline]
pub fn ad_solve_upper_triangular<R2: Dim, C2: Dim, S2>(
&self,
b: &Matrix<T, R2, C2, S2>,
) -> Option<OMatrix<T, R2, C2>>
where
S2: Storage<T, R2, C2>,
DefaultAllocator: Allocator<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let mut res = b.clone_owned();
if self.ad_solve_upper_triangular_mut(&mut res) {
Some(res)
} else {
None
}
}
pub fn ad_solve_lower_triangular_mut<R2: Dim, C2: Dim, S2>(
&self,
b: &mut Matrix<T, R2, C2, S2>,
) -> bool
where
S2: StorageMut<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let cols = b.ncols();
for i in 0..cols {
if !self.xx_solve_lower_triangular_vector_mut(
&mut b.column_mut(i),
|e| e.conjugate(),
|a, b| a.dotc(b),
) {
return false;
}
}
true
}
pub fn ad_solve_upper_triangular_mut<R2: Dim, C2: Dim, S2>(
&self,
b: &mut Matrix<T, R2, C2, S2>,
) -> bool
where
S2: StorageMut<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let cols = b.ncols();
for i in 0..cols {
if !self.xx_solve_upper_triangular_vector_mut(
&mut b.column_mut(i),
|e| e.conjugate(),
|a, b| a.dotc(b),
) {
return false;
}
}
true
}
#[inline(always)]
fn xx_solve_lower_triangular_vector_mut<R2: Dim, S2>(
&self,
b: &mut Vector<T, R2, S2>,
conjugate: impl Fn(T) -> T,
dot: impl Fn(
&DVectorSlice<'_, T, S::RStride, S::CStride>,
&DVectorSlice<'_, T, S2::RStride, S2::CStride>,
) -> T,
) -> bool
where
S2: StorageMut<T, R2, U1>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let dim = self.nrows();
for i in (0..dim).rev() {
let dot = dot(&self.slice_range(i + 1.., i), &b.slice_range(i + 1.., 0));
unsafe {
let b_i = b.vget_unchecked_mut(i);
let diag = conjugate(self.get_unchecked((i, i)).clone());
if diag.is_zero() {
return false;
}
*b_i = (b_i.clone() - dot) / diag;
}
}
true
}
#[inline(always)]
fn xx_solve_upper_triangular_vector_mut<R2: Dim, S2>(
&self,
b: &mut Vector<T, R2, S2>,
conjugate: impl Fn(T) -> T,
dot: impl Fn(
&DVectorSlice<'_, T, S::RStride, S::CStride>,
&DVectorSlice<'_, T, S2::RStride, S2::CStride>,
) -> T,
) -> bool
where
S2: StorageMut<T, R2, U1>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let dim = self.nrows();
for i in 0..dim {
let dot = dot(&self.slice_range(..i, i), &b.slice_range(..i, 0));
unsafe {
let b_i = b.vget_unchecked_mut(i);
let diag = conjugate(self.get_unchecked((i, i)).clone());
if diag.is_zero() {
return false;
}
*b_i = (b_i.clone() - dot) / diag;
}
}
true
}
}
impl<T: SimdComplexField, D: Dim, S: Storage<T, D, D>> SquareMatrix<T, D, S> {
#[must_use = "Did you mean to use solve_lower_triangular_unchecked_mut()?"]
#[inline]
pub fn solve_lower_triangular_unchecked<R2: Dim, C2: Dim, S2>(
&self,
b: &Matrix<T, R2, C2, S2>,
) -> OMatrix<T, R2, C2>
where
S2: Storage<T, R2, C2>,
DefaultAllocator: Allocator<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let mut res = b.clone_owned();
self.solve_lower_triangular_unchecked_mut(&mut res);
res
}
#[must_use = "Did you mean to use solve_upper_triangular_unchecked_mut()?"]
#[inline]
pub fn solve_upper_triangular_unchecked<R2: Dim, C2: Dim, S2>(
&self,
b: &Matrix<T, R2, C2, S2>,
) -> OMatrix<T, R2, C2>
where
S2: Storage<T, R2, C2>,
DefaultAllocator: Allocator<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let mut res = b.clone_owned();
self.solve_upper_triangular_unchecked_mut(&mut res);
res
}
pub fn solve_lower_triangular_unchecked_mut<R2: Dim, C2: Dim, S2>(
&self,
b: &mut Matrix<T, R2, C2, S2>,
) where
S2: StorageMut<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
for i in 0..b.ncols() {
self.solve_lower_triangular_vector_unchecked_mut(&mut b.column_mut(i));
}
}
fn solve_lower_triangular_vector_unchecked_mut<R2: Dim, S2>(&self, b: &mut Vector<T, R2, S2>)
where
S2: StorageMut<T, R2, U1>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let dim = self.nrows();
for i in 0..dim {
let coeff;
unsafe {
let diag = self.get_unchecked((i, i)).clone();
coeff = b.vget_unchecked(i).clone() / diag;
*b.vget_unchecked_mut(i) = coeff.clone();
}
b.rows_range_mut(i + 1..)
.axpy(-coeff.clone(), &self.slice_range(i + 1.., i), T::one());
}
}
pub fn solve_lower_triangular_with_diag_unchecked_mut<R2: Dim, C2: Dim, S2>(
&self,
b: &mut Matrix<T, R2, C2, S2>,
diag: T,
) where
S2: StorageMut<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let dim = self.nrows();
let cols = b.ncols();
for k in 0..cols {
let mut bcol = b.column_mut(k);
for i in 0..dim - 1 {
let coeff = unsafe { bcol.vget_unchecked(i).clone() } / diag.clone();
bcol.rows_range_mut(i + 1..)
.axpy(-coeff, &self.slice_range(i + 1.., i), T::one());
}
}
}
pub fn solve_upper_triangular_unchecked_mut<R2: Dim, C2: Dim, S2>(
&self,
b: &mut Matrix<T, R2, C2, S2>,
) where
S2: StorageMut<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
for i in 0..b.ncols() {
self.solve_upper_triangular_vector_unchecked_mut(&mut b.column_mut(i))
}
}
fn solve_upper_triangular_vector_unchecked_mut<R2: Dim, S2>(&self, b: &mut Vector<T, R2, S2>)
where
S2: StorageMut<T, R2, U1>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let dim = self.nrows();
for i in (0..dim).rev() {
let coeff;
unsafe {
let diag = self.get_unchecked((i, i)).clone();
coeff = b.vget_unchecked(i).clone() / diag;
*b.vget_unchecked_mut(i) = coeff.clone();
}
b.rows_range_mut(..i)
.axpy(-coeff, &self.slice_range(..i, i), T::one());
}
}
#[must_use = "Did you mean to use tr_solve_lower_triangular_unchecked_mut()?"]
#[inline]
pub fn tr_solve_lower_triangular_unchecked<R2: Dim, C2: Dim, S2>(
&self,
b: &Matrix<T, R2, C2, S2>,
) -> OMatrix<T, R2, C2>
where
S2: Storage<T, R2, C2>,
DefaultAllocator: Allocator<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let mut res = b.clone_owned();
self.tr_solve_lower_triangular_unchecked_mut(&mut res);
res
}
#[must_use = "Did you mean to use tr_solve_upper_triangular_unchecked_mut()?"]
#[inline]
pub fn tr_solve_upper_triangular_unchecked<R2: Dim, C2: Dim, S2>(
&self,
b: &Matrix<T, R2, C2, S2>,
) -> OMatrix<T, R2, C2>
where
S2: Storage<T, R2, C2>,
DefaultAllocator: Allocator<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let mut res = b.clone_owned();
self.tr_solve_upper_triangular_unchecked_mut(&mut res);
res
}
pub fn tr_solve_lower_triangular_unchecked_mut<R2: Dim, C2: Dim, S2>(
&self,
b: &mut Matrix<T, R2, C2, S2>,
) where
S2: StorageMut<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
for i in 0..b.ncols() {
self.xx_solve_lower_triangular_vector_unchecked_mut(
&mut b.column_mut(i),
|e| e,
|a, b| a.dot(b),
)
}
}
pub fn tr_solve_upper_triangular_unchecked_mut<R2: Dim, C2: Dim, S2>(
&self,
b: &mut Matrix<T, R2, C2, S2>,
) where
S2: StorageMut<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
for i in 0..b.ncols() {
self.xx_solve_upper_triangular_vector_unchecked_mut(
&mut b.column_mut(i),
|e| e,
|a, b| a.dot(b),
)
}
}
#[must_use = "Did you mean to use ad_solve_lower_triangular_unchecked_mut()?"]
#[inline]
pub fn ad_solve_lower_triangular_unchecked<R2: Dim, C2: Dim, S2>(
&self,
b: &Matrix<T, R2, C2, S2>,
) -> OMatrix<T, R2, C2>
where
S2: Storage<T, R2, C2>,
DefaultAllocator: Allocator<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let mut res = b.clone_owned();
self.ad_solve_lower_triangular_unchecked_mut(&mut res);
res
}
#[must_use = "Did you mean to use ad_solve_upper_triangular_unchecked_mut()?"]
#[inline]
pub fn ad_solve_upper_triangular_unchecked<R2: Dim, C2: Dim, S2>(
&self,
b: &Matrix<T, R2, C2, S2>,
) -> OMatrix<T, R2, C2>
where
S2: Storage<T, R2, C2>,
DefaultAllocator: Allocator<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let mut res = b.clone_owned();
self.ad_solve_upper_triangular_unchecked_mut(&mut res);
res
}
pub fn ad_solve_lower_triangular_unchecked_mut<R2: Dim, C2: Dim, S2>(
&self,
b: &mut Matrix<T, R2, C2, S2>,
) where
S2: StorageMut<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
for i in 0..b.ncols() {
self.xx_solve_lower_triangular_vector_unchecked_mut(
&mut b.column_mut(i),
|e| e.simd_conjugate(),
|a, b| a.dotc(b),
)
}
}
pub fn ad_solve_upper_triangular_unchecked_mut<R2: Dim, C2: Dim, S2>(
&self,
b: &mut Matrix<T, R2, C2, S2>,
) where
S2: StorageMut<T, R2, C2>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
for i in 0..b.ncols() {
self.xx_solve_upper_triangular_vector_unchecked_mut(
&mut b.column_mut(i),
|e| e.simd_conjugate(),
|a, b| a.dotc(b),
)
}
}
#[inline(always)]
fn xx_solve_lower_triangular_vector_unchecked_mut<R2: Dim, S2>(
&self,
b: &mut Vector<T, R2, S2>,
conjugate: impl Fn(T) -> T,
dot: impl Fn(
&DVectorSlice<'_, T, S::RStride, S::CStride>,
&DVectorSlice<'_, T, S2::RStride, S2::CStride>,
) -> T,
) where
S2: StorageMut<T, R2, U1>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
let dim = self.nrows();
for i in (0..dim).rev() {
let dot = dot(&self.slice_range(i + 1.., i), &b.slice_range(i + 1.., 0));
unsafe {
let b_i = b.vget_unchecked_mut(i);
let diag = conjugate(self.get_unchecked((i, i)).clone());
*b_i = (b_i.clone() - dot) / diag;
}
}
}
#[inline(always)]
fn xx_solve_upper_triangular_vector_unchecked_mut<R2: Dim, S2>(
&self,
b: &mut Vector<T, R2, S2>,
conjugate: impl Fn(T) -> T,
dot: impl Fn(
&DVectorSlice<'_, T, S::RStride, S::CStride>,
&DVectorSlice<'_, T, S2::RStride, S2::CStride>,
) -> T,
) where
S2: StorageMut<T, R2, U1>,
ShapeConstraint: SameNumberOfRows<R2, D>,
{
for i in 0..self.nrows() {
let dot = dot(&self.slice_range(..i, i), &b.slice_range(..i, 0));
unsafe {
let b_i = b.vget_unchecked_mut(i);
let diag = conjugate(self.get_unchecked((i, i)).clone());
*b_i = (b_i.clone() - dot) / diag;
}
}
}
}