use faer::linalg::solvers::DenseSolveCore;
use faer::{Mat, Side};
use ndarray::{Array1, Array2, Axis};
use robust_rs_core::error::RobustError;
pub(crate) fn to_faer(a: &Array2<f64>) -> Mat<f64> {
Mat::from_fn(a.nrows(), a.ncols(), |i, j| a[[i, j]])
}
pub(crate) fn from_faer(m: &Mat<f64>) -> Array2<f64> {
Array2::from_shape_fn((m.nrows(), m.ncols()), |(i, j)| m[(i, j)])
}
pub(crate) fn mean(x: &Array2<f64>) -> Array1<f64> {
x.mean_axis(Axis(0)).expect("non-empty data")
}
pub(crate) fn center(x: &Array2<f64>, loc: &Array1<f64>) -> Array2<f64> {
let (n, p) = x.dim();
Array2::from_shape_fn((n, p), |(i, j)| x[[i, j]] - loc[j])
}
pub(crate) fn mean_covariance(x: &Array2<f64>) -> (Array1<f64>, Array2<f64>) {
let (m, _p) = x.dim();
let mu = mean(x);
let centered = center(x, &mu);
let denom = if m > 1 { (m - 1) as f64 } else { 1.0 };
let cov = centered.t().dot(¢ered) / denom;
(mu, cov)
}
pub(crate) fn spd_inverse_logdet(a: &Array2<f64>) -> Result<(Array2<f64>, f64), RobustError> {
let p = a.nrows();
let llt = to_faer(a)
.llt(Side::Lower)
.map_err(|_| RobustError::SingularDesign)?;
let l = llt.L();
let mut logdet = 0.0;
for i in 0..p {
logdet += l[(i, i)].ln();
}
logdet *= 2.0;
if !logdet.is_finite() {
return Err(RobustError::SingularDesign);
}
Ok((from_faer(&llt.inverse()), logdet))
}
pub(crate) fn spd_logdet(a: &Array2<f64>) -> Result<f64, RobustError> {
let p = a.nrows();
let llt = to_faer(a)
.llt(Side::Lower)
.map_err(|_| RobustError::SingularDesign)?;
let l = llt.L();
let mut logdet = 0.0;
for i in 0..p {
logdet += l[(i, i)].ln();
}
logdet *= 2.0;
if logdet.is_finite() {
Ok(logdet)
} else {
Err(RobustError::SingularDesign)
}
}
pub(crate) fn mahalanobis_sq(x: &Array2<f64>, loc: &Array1<f64>, inv: &Array2<f64>) -> Array1<f64> {
let centered = center(x, loc); let tmp = centered.dot(inv); (¢ered * &tmp).sum_axis(Axis(1))
}
pub(crate) fn symmetric_eigen(a: &Array2<f64>) -> Result<(Array1<f64>, Array2<f64>), RobustError> {
let p = a.nrows();
let eig = to_faer(a)
.self_adjoint_eigen(Side::Lower)
.map_err(|_| RobustError::SingularDesign)?;
let s = eig.S().column_vector();
let u = eig.U();
let values = Array1::from_shape_fn(p, |i| s[i]);
let vectors = Array2::from_shape_fn((p, p), |(i, j)| u[(i, j)]);
Ok((values, vectors))
}