use ndarray::{Array1, Array2, Array3, ArrayView1, ArrayView2, Axis, Zip, s};
use num_complex::Complex64;
use sprs::{CsMat, FillInReduction, TriMat};
use sprs_ldl::Ldl;
use sprs_suitesparse_ldl::LdlNumeric;
use crate::{Laplacian, error::GspError};
pub(super) type SparseFactor = LdlNumeric;
#[derive(Debug, Clone)]
pub(super) struct BranchUpdate {
pub(super) i: usize,
pub(super) j: usize,
pub(super) sqrt_w: f64,
}
#[derive(Debug, Clone)]
pub(super) struct PoleState {
pub(super) factor: SparseFactor,
pub(super) solved_cols: Vec<Array1<f64>>,
pub(super) gram_inv: Array2<f64>,
}
pub(super) fn check_scale(scale: f64) -> Result<(), GspError> {
if !scale.is_finite() || scale <= 0.0 {
return Err(GspError::InvalidScales(format!(
"scale must be finite and > 0, got {scale}"
)));
}
Ok(())
}
pub(super) fn check_scales(scales: &[f64]) -> Result<(), GspError> {
if scales.is_empty() {
return Err(GspError::InvalidScales(
"at least one scale is required".to_string(),
));
}
for &s in scales {
check_scale(s)?;
}
Ok(())
}
pub(super) fn check_signal_rows(n_signal_rows: usize, n_vertices: usize) -> Result<(), GspError> {
if n_signal_rows != n_vertices {
return Err(GspError::Dimensions(format!(
"signal rows {n_signal_rows} does not match graph size {n_vertices}"
)));
}
Ok(())
}
pub(super) fn shifted_matrix(l: &Laplacian, q: f64) -> Laplacian {
let n = l.rows();
let mut tri = TriMat::<f64>::with_capacity((n, n), l.nnz() + n);
for (col, vec) in l.outer_iterator().enumerate() {
for (row, val) in vec.iter() {
tri.add_triplet(row, col, *val);
}
}
for i in 0..n {
tri.add_triplet(i, i, q);
}
tri.to_csc()
}
pub(super) fn affine_sparse(l: &Laplacian, alpha: f64, beta: f64) -> Laplacian {
let n = l.rows();
let mut tri = TriMat::<f64>::with_capacity((n, n), l.nnz() + n);
for (col, vec) in l.outer_iterator().enumerate() {
for (row, val) in vec.iter() {
tri.add_triplet(row, col, alpha * *val);
}
}
for i in 0..n {
tri.add_triplet(i, i, beta);
}
tri.to_csc()
}
pub(super) fn factorize_shifted_matrix(shifted: &Laplacian) -> Result<SparseFactor, GspError> {
Ok(Ldl::new()
.fill_in_reduction(FillInReduction::CAMDSuiteSparse)
.numeric_c(shifted.view())?)
}
pub(super) fn solve_multi_rhs(
factor: &SparseFactor,
b: ArrayView2<f64>,
) -> Result<Array2<f64>, GspError> {
let n = b.nrows();
let m = b.ncols();
let mut out = Array2::<f64>::zeros((n, m));
let mut rhs = vec![0.0; n];
if let Some(out_data) = out.as_slice_memory_order_mut() {
for col in 0..m {
for row in 0..n {
rhs[row] = b[[row, col]];
}
let x = factor.solve(&rhs);
if x.len() != n {
return Err(GspError::SingularSolve(
"solver returned incorrect output length".to_string(),
));
}
for (row, &v) in x.iter().enumerate() {
out_data[row * m + col] = v;
}
}
return Ok(out);
}
for col in 0..m {
for row in 0..n {
rhs[row] = b[[row, col]];
}
let x = factor.solve(&rhs);
if x.len() != n {
return Err(GspError::SingularSolve(
"solver returned incorrect output length".to_string(),
));
}
for (row, &v) in x.iter().enumerate() {
out[[row, col]] = v;
}
}
Ok(out)
}
pub(super) fn solve_vector(
factor: &SparseFactor,
rhs: ArrayView1<f64>,
) -> Result<Array1<f64>, GspError> {
let x = factor.solve(&rhs.to_vec());
if x.len() != rhs.len() {
return Err(GspError::SingularSolve(
"solver returned incorrect output length".to_string(),
));
}
Ok(Array1::from(x))
}
pub(super) fn update_vector(n: usize, upd: &BranchUpdate) -> Array1<f64> {
let mut c = Array1::<f64>::zeros(n);
c[upd.i] = upd.sqrt_w;
c[upd.j] = -upd.sqrt_w;
c
}
pub(super) fn dot_update(upd: &BranchUpdate, x: ArrayView1<'_, f64>) -> f64 {
upd.sqrt_w * (x[upd.i] - x[upd.j])
}
pub(super) fn invert_dense(a: Array2<f64>) -> Result<Array2<f64>, GspError> {
let n = a.nrows();
if n != a.ncols() {
return Err(GspError::Dimensions(format!(
"dense inverse expects square matrix, got {}x{}",
a.nrows(),
a.ncols()
)));
}
if n == 0 {
return Ok(Array2::zeros((0, 0)));
}
let data = a
.as_slice()
.ok_or_else(|| GspError::Factorization("matrix storage is not contiguous".to_string()))?;
let dm = nalgebra::DMatrix::<f64>::from_row_slice(n, n, data);
let inv = dm.try_inverse().ok_or_else(|| {
GspError::SingularSolve("failed to invert Woodbury gram matrix".to_string())
})?;
let mut out = Array2::<f64>::zeros((n, n));
for r in 0..n {
for c in 0..n {
out[[r, c]] = inv[(r, c)];
}
}
Ok(out)
}
pub(super) fn sparse_dense_mul_into(a: &CsMat<f64>, x: &Array2<f64>, y: &mut Array2<f64>) {
debug_assert_eq!(a.cols(), x.nrows());
debug_assert_eq!(y.nrows(), a.rows());
debug_assert_eq!(y.ncols(), x.ncols());
y.fill(0.0);
let n_rhs = x.ncols();
if n_rhs == 0 {
return;
}
if let (Some(x_data), Some(y_data)) = (x.as_slice_memory_order(), y.as_slice_memory_order_mut())
{
for (col, vec) in a.outer_iterator().enumerate() {
let x_row = &x_data[col * n_rhs..(col + 1) * n_rhs];
for (row, val) in vec.iter() {
let y_row = &mut y_data[row * n_rhs..(row + 1) * n_rhs];
for (dst, &src) in y_row.iter_mut().zip(x_row.iter()) {
*dst += *val * src;
}
}
}
return;
}
for (col, vec) in a.outer_iterator().enumerate() {
for (row, val) in vec.iter() {
for k in 0..n_rhs {
y[[row, k]] += *val * x[[col, k]];
}
}
}
}
pub(super) fn accumulate(w: &mut Array3<f64>, z: &Array2<f64>, c: ndarray::ArrayView1<'_, f64>) {
for d in 0..c.len() {
let coeff = c[d];
if coeff == 0.0 {
continue;
}
w.slice_mut(s![.., .., d]).scaled_add(coeff, &z.view());
}
}
pub(super) fn apply_direct_term(
output: &mut Array3<f64>,
signal: ArrayView2<f64>,
direct: ArrayView1<'_, f64>,
) -> Result<(), GspError> {
if direct.is_empty() {
return Ok(());
}
let n_channels = output.len_of(Axis(2));
if direct.len() != n_channels {
return Err(GspError::Dimensions(format!(
"direct term length {} does not match residue dimension {}",
direct.len(),
n_channels
)));
}
for channel in 0..n_channels {
let coeff = direct[channel];
if coeff == 0.0 {
continue;
}
output
.slice_mut(s![.., .., channel])
.scaled_add(coeff, &signal);
}
Ok(())
}
pub(super) fn combine_complex_matrix(real: Array2<f64>, imag: Array2<f64>) -> Array2<Complex64> {
let mut out = Array2::<Complex64>::zeros(real.raw_dim());
Zip::from(&mut out)
.and(&real)
.and(&imag)
.for_each(|o, &r, &i| *o = Complex64::new(r, i));
out
}
pub(super) fn combine_complex_vec(
real: Vec<Array2<f64>>,
imag: Vec<Array2<f64>>,
) -> Result<Vec<Array2<Complex64>>, GspError> {
if real.len() != imag.len() {
return Err(GspError::Dimensions(format!(
"mismatched complex decomposition lengths: real {}, imag {}",
real.len(),
imag.len()
)));
}
Ok(real
.into_iter()
.zip(imag)
.map(|(r, i)| combine_complex_matrix(r, i))
.collect())
}