mod methods;
#[cfg(test)]
mod tests;
pub use methods::EigenIterator;
use crate::{Backend, Matrix, TruenoError};
const MAX_JACOBI_SWEEPS: usize = 50;
const CONVERGENCE_THRESHOLD: f32 = 1e-7;
#[allow(dead_code)]
const GPU_THRESHOLD: usize = 1000;
#[derive(Debug, Clone)]
pub struct SymmetricEigen {
pub(crate) eigenvalues: Vec<f32>,
pub(crate) eigenvectors: Matrix<f32>,
pub(crate) _sort_indices: Vec<usize>,
pub(crate) backend: Backend,
}
impl SymmetricEigen {
pub fn new(matrix: &Matrix<f32>) -> Result<Self, TruenoError> {
if matrix.rows() != matrix.cols() {
return Err(TruenoError::InvalidInput(format!(
"Matrix must be square for eigendecomposition, got {}x{}",
matrix.rows(),
matrix.cols()
)));
}
if matrix.rows() == 0 {
return Err(TruenoError::InvalidInput(
"Cannot compute eigendecomposition of empty matrix".to_string(),
));
}
let backend = Backend::select_best();
#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
{
let n = matrix.rows();
if n >= GPU_THRESHOLD && crate::backends::gpu::GpuBackend::is_available() {
return Self::compute_gpu(matrix);
}
}
Self::compute_jacobi(matrix, backend)
}
fn compute_jacobi(matrix: &Matrix<f32>, backend: Backend) -> Result<Self, TruenoError> {
let n = matrix.rows();
let mut a = matrix.as_slice().to_vec();
let frobenius_sq: f32 = a.iter().map(|x| x * x).sum();
let tolerance = CONVERGENCE_THRESHOLD * frobenius_sq.sqrt().max(1.0);
let mut v = vec![0.0f32; n * n];
for i in 0..n {
v[i * n + i] = 1.0;
}
for _sweep in 0..MAX_JACOBI_SWEEPS {
let mut converged = true;
for i in 0..n {
for j in (i + 1)..n {
let aij = a[i * n + j];
if aij.abs() < tolerance {
continue;
}
converged = false;
Self::jacobi_rotate(&mut a, &mut v, n, i, j, backend);
}
}
if converged {
let eigenvalues: Vec<f32> = (0..n).map(|i| a[i * n + i]).collect();
let mut indices: Vec<usize> = (0..n).collect();
indices.sort_by(|&i, &j| {
eigenvalues[j].partial_cmp(&eigenvalues[i]).unwrap_or(std::cmp::Ordering::Equal)
});
let sorted_eigenvalues: Vec<f32> =
indices.iter().map(|&i| eigenvalues[i]).collect();
let mut eigenvector_data = vec![0.0f32; n * n];
for (new_col, &old_col) in indices.iter().enumerate() {
for row in 0..n {
eigenvector_data[row * n + new_col] = v[row * n + old_col];
}
}
let eigenvectors = Matrix::from_vec(n, n, eigenvector_data)?;
return Ok(SymmetricEigen {
eigenvalues: sorted_eigenvalues,
eigenvectors,
_sort_indices: indices,
backend,
});
}
}
Err(TruenoError::InvalidInput(format!(
"Jacobi algorithm failed to converge after {} sweeps",
MAX_JACOBI_SWEEPS
)))
}
#[inline]
fn _find_max_off_diagonal(a: &[f32], n: usize) -> (usize, usize, f32) {
let mut max_val = 0.0f32;
let mut p = 0;
let mut q = 1;
for i in 0..n {
for j in (i + 1)..n {
let val = a[i * n + j].abs();
if val > max_val {
max_val = val;
p = i;
q = j;
}
}
}
(p, q, max_val)
}
#[inline]
fn jacobi_rotate(
a: &mut [f32],
v: &mut [f32],
n: usize,
p: usize,
q: usize,
_backend: Backend,
) {
let app = a[p * n + p];
let aqq = a[q * n + q];
let apq = a[p * n + q];
if apq.abs() < 1e-15 {
return;
}
let tau = (aqq - app) / (2.0 * apq);
let t = if tau >= 0.0 {
1.0 / (tau + (1.0 + tau * tau).sqrt())
} else {
-1.0 / (-tau + (1.0 + tau * tau).sqrt())
};
let c = 1.0 / (1.0 + t * t).sqrt();
let s = t * c;
a[p * n + p] = app - t * apq;
a[q * n + q] = aqq + t * apq;
a[p * n + q] = 0.0;
a[q * n + p] = 0.0;
for k in 0..n {
if k != p && k != q {
let akp = a[k * n + p];
let akq = a[k * n + q];
a[k * n + p] = c * akp - s * akq;
a[p * n + k] = a[k * n + p];
a[k * n + q] = s * akp + c * akq;
a[q * n + k] = a[k * n + q];
}
}
for k in 0..n {
let vkp = v[k * n + p];
let vkq = v[k * n + q];
v[k * n + p] = c * vkp - s * vkq;
v[k * n + q] = s * vkp + c * vkq;
}
}
#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
fn compute_gpu(matrix: &Matrix<f32>) -> Result<Self, TruenoError> {
use crate::backends::gpu::GpuBackend;
let n = matrix.rows();
let mut gpu = GpuBackend::new();
let (eigenvalues, eigenvector_data) =
gpu.symmetric_eigen(matrix.as_slice(), n).map_err(|e| {
TruenoError::InvalidInput(format!("GPU eigendecomposition failed: {}", e))
})?;
let mut indices: Vec<usize> = (0..n).collect();
indices.sort_by(|&i, &j| {
eigenvalues[j].partial_cmp(&eigenvalues[i]).unwrap_or(std::cmp::Ordering::Equal)
});
let sorted_eigenvalues: Vec<f32> = indices.iter().map(|&i| eigenvalues[i]).collect();
let mut sorted_eigenvector_data = vec![0.0f32; n * n];
for (new_col, &old_col) in indices.iter().enumerate() {
for row in 0..n {
sorted_eigenvector_data[row * n + new_col] = eigenvector_data[row * n + old_col];
}
}
let eigenvectors = Matrix::from_vec(n, n, sorted_eigenvector_data)?;
Ok(SymmetricEigen {
eigenvalues: sorted_eigenvalues,
eigenvectors,
_sort_indices: indices,
backend: Backend::GPU,
})
}
}