use crate::linear_algebra::{Matrix, Vector};
use crate::scalar::Numeric;
use crate::utils::error_codes::CalcError;
#[derive(Debug, Clone, Copy)]
#[must_use]
pub struct Cholesky<const N: usize, T = f64> {
pub(crate) l: Matrix<N, N, T>,
}
impl<const N: usize, T: Numeric> Matrix<N, N, T> {
pub fn cholesky(self) -> Result<Cholesky<N, T>, CalcError> {
let mut l = Matrix::zeros();
for j in 0..N {
let mut d = self[(j, j)];
for k in 0..j {
d -= l[(j, k)] * l[(j, k)];
}
if d <= T::ZERO {
return Err(CalcError::NotPositiveDefinite);
}
let ljj = d.sqrt();
l[(j, j)] = ljj;
for i in (j + 1)..N {
let mut s = self[(i, j)];
for k in 0..j {
s -= l[(i, k)] * l[(j, k)];
}
l[(i, j)] = s / ljj;
}
}
Ok(Cholesky { l })
}
}
impl<const N: usize, T: Numeric> Cholesky<N, T> {
pub fn l(&self) -> Matrix<N, N, T> {
self.l
}
#[inline]
#[must_use]
pub fn determinant(&self) -> T {
let mut det = T::ONE;
for i in 0..N {
det *= self.l[(i, i)] * self.l[(i, i)];
}
det
}
pub fn solve(&self, b: Vector<N, T>) -> Vector<N, T> {
let mut x: [T; N] = core::array::from_fn(|i| b[i]);
for i in 0..N {
let mut sum = x[i];
for (j, &xj) in x.iter().enumerate().take(i) {
sum -= self.l[(i, j)] * xj;
}
x[i] = sum / self.l[(i, i)];
}
for i in (0..N).rev() {
let mut sum = x[i];
for (j, &xj) in x.iter().enumerate().skip(i + 1) {
sum -= self.l[(j, i)] * xj;
}
x[i] = sum / self.l[(i, i)];
}
Vector::new(x)
}
pub fn solve_matrix<const K: usize>(&self, b: Matrix<N, K, T>) -> Matrix<N, K, T> {
let mut result = Matrix::zeros();
for c in 0..K {
let x = self.solve(b.column(c));
for r in 0..N {
result[(r, c)] = x[r];
}
}
result
}
pub fn inverse(&self) -> Matrix<N, N, T> {
self.solve_matrix(Matrix::identity())
}
}