diff --git a/Cargo.toml b/Cargo.toml index 08120c51..a2e642b0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -52,6 +52,11 @@ optional = true paste = "0.1" criterion = "0.3" +[[bench]] +name = "truncated_eig" +harness = false + [[bench]] name = "eigh" harness = false + diff --git a/benches/truncated_eig.rs b/benches/truncated_eig.rs new file mode 100644 index 00000000..b24aac7c --- /dev/null +++ b/benches/truncated_eig.rs @@ -0,0 +1,41 @@ +#[macro_use] +extern crate criterion; + +use criterion::Criterion; +use ndarray::*; +use ndarray_linalg::*; + +macro_rules! impl_teig { + ($n:expr) => { + paste::item! { + fn [](c: &mut Criterion) { + c.bench_function(&format!("truncated_eig{}", $n), |b| { + let a: Array2 = random(($n, $n)); + let a = a.t().dot(&a); + + b.iter(move || { + let _result = TruncatedEig::new(a.clone(), TruncatedOrder::Largest).decompose(1); + }) + }); + c.bench_function(&format!("truncated_eig{}_t", $n), |b| { + let a: Array2 = random(($n, $n).f()); + let a = a.t().dot(&a); + + b.iter(|| { + let _result = TruncatedEig::new(a.clone(), TruncatedOrder::Largest).decompose(1); + }) + }); + } + } + }; +} + +impl_teig!(4); +impl_teig!(8); +impl_teig!(16); +impl_teig!(32); +impl_teig!(64); +impl_teig!(128); + +criterion_group!(teig, teig4, teig8, teig16, teig32, teig64, teig128); +criterion_main!(teig); diff --git a/examples/solveh.rs b/examples/solveh.rs index 8ddc73a0..30b2897b 100644 --- a/examples/solveh.rs +++ b/examples/solveh.rs @@ -10,7 +10,7 @@ fn solve() -> Result<(), error::LinalgError> { let b: Array1 = random(3); println!("b = {:?}", &b); let x = a.solveh(&b)?; - println!("Ax = {:?}", a.dot(&x));; + println!("Ax = {:?}", a.dot(&x)); Ok(()) } diff --git a/examples/truncated_eig.rs b/examples/truncated_eig.rs new file mode 100644 index 00000000..58172ec6 --- /dev/null +++ b/examples/truncated_eig.rs @@ -0,0 +1,23 @@ +extern crate ndarray; +extern crate ndarray_linalg; + +use ndarray::*; +use ndarray_linalg::*; + +fn main() { + let n = 10; + let v = random_unitary(n); + + // set eigenvalues in decreasing order + let t = Array1::linspace(n as f64, -(n as f64), n); + + println!("Generate spectrum: {:?}", &t); + + let t = Array2::from_diag(&t); + let a = v.dot(&t.dot(&v.t())); + + // calculate the truncated eigenproblem decomposition + for (val, _) in TruncatedEig::new(a, TruncatedOrder::Largest) { + println!("Found eigenvalue {}", val[0]); + } +} diff --git a/examples/truncated_svd.rs b/examples/truncated_svd.rs new file mode 100644 index 00000000..8cc164eb --- /dev/null +++ b/examples/truncated_svd.rs @@ -0,0 +1,22 @@ +extern crate ndarray; +extern crate ndarray_linalg; + +use ndarray::*; +use ndarray_linalg::*; + +fn main() { + let a = arr2(&[[3., 2., 2.], [2., 3., -2.]]); + + // calculate the truncated singular value decomposition for 2 singular values + let result = TruncatedSvd::new(a, TruncatedOrder::Largest).decompose(2).unwrap(); + + // acquire singular values, left-singular vectors and right-singular vectors + let (u, sigma, v_t) = result.values_vectors(); + println!("Result of the singular value decomposition A = UΣV^T:"); + println!(" === U ==="); + println!("{:?}", u); + println!(" === Σ ==="); + println!("{:?}", Array2::from_diag(&sigma)); + println!(" === V^T ==="); + println!("{:?}", v_t); +} diff --git a/src/eigh.rs b/src/eigh.rs index 5e2162b3..1fc7a031 100644 --- a/src/eigh.rs +++ b/src/eigh.rs @@ -42,6 +42,20 @@ where } } +impl EighInto for (ArrayBase, ArrayBase) +where + A: Scalar + Lapack, + S: DataMut, + S2: DataMut, +{ + type EigVal = Array1; + + fn eigh_into(mut self, uplo: UPLO) -> Result<(Self::EigVal, Self)> { + let (val, _) = self.eigh_inplace(uplo)?; + Ok((val, self)) + } +} + impl Eigh for ArrayBase where A: Scalar + Lapack, @@ -56,6 +70,21 @@ where } } +impl Eigh for (ArrayBase, ArrayBase) +where + A: Scalar + Lapack, + S: Data, + S2: Data, +{ + type EigVal = Array1; + type EigVec = (Array2, Array2); + + fn eigh(&self, uplo: UPLO) -> Result<(Self::EigVal, Self::EigVec)> { + let (a, b) = (self.0.to_owned(), self.1.to_owned()); + (a, b).eigh_into(uplo) + } +} + impl EighInplace for ArrayBase where A: Scalar + Lapack, @@ -75,6 +104,42 @@ where } } +impl EighInplace for (ArrayBase, ArrayBase) +where + A: Scalar + Lapack, + S: DataMut, + S2: DataMut, +{ + type EigVal = Array1; + + fn eigh_inplace(&mut self, uplo: UPLO) -> Result<(Self::EigVal, &mut Self)> { + let layout = self.0.square_layout()?; + // XXX Force layout to be Fortran (see #146) + match layout { + MatrixLayout::C(_) => self.0.swap_axes(0, 1), + MatrixLayout::F(_) => {} + } + + let layout = self.1.square_layout()?; + match layout { + MatrixLayout::C(_) => self.1.swap_axes(0, 1), + MatrixLayout::F(_) => {} + } + + let s = unsafe { + A::eigh_generalized( + true, + self.0.square_layout()?, + uplo, + self.0.as_allocated_mut()?, + self.1.as_allocated_mut()?, + )? + }; + + Ok((ArrayBase::from(s), self)) + } +} + /// Calculate eigenvalues without eigenvectors pub trait EigValsh { type EigVal; diff --git a/src/lapack/eigh.rs b/src/lapack/eigh.rs index 610564d2..11b40bb0 100644 --- a/src/lapack/eigh.rs +++ b/src/lapack/eigh.rs @@ -12,10 +12,17 @@ use super::{into_result, UPLO}; /// Wraps `*syev` for real and `*heev` for complex pub trait Eigh_: Scalar { unsafe fn eigh(calc_eigenvec: bool, l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result>; + unsafe fn eigh_generalized( + calc_eigenvec: bool, + l: MatrixLayout, + uplo: UPLO, + a: &mut [Self], + b: &mut [Self], + ) -> Result>; } macro_rules! impl_eigh { - ($scalar:ty, $ev:path) => { + ($scalar:ty, $ev:path, $evg:path) => { impl Eigh_ for $scalar { unsafe fn eigh(calc_v: bool, l: MatrixLayout, uplo: UPLO, mut a: &mut [Self]) -> Result> { let (n, _) = l.size(); @@ -24,11 +31,36 @@ macro_rules! impl_eigh { let info = $ev(l.lapacke_layout(), jobz, uplo as u8, n, &mut a, n, &mut w); into_result(info, w) } + + unsafe fn eigh_generalized( + calc_v: bool, + l: MatrixLayout, + uplo: UPLO, + mut a: &mut [Self], + mut b: &mut [Self], + ) -> Result> { + let (n, _) = l.size(); + let jobz = if calc_v { b'V' } else { b'N' }; + let mut w = vec![Self::Real::zero(); n as usize]; + let info = $evg( + l.lapacke_layout(), + 1, + jobz, + uplo as u8, + n, + &mut a, + n, + &mut b, + n, + &mut w, + ); + into_result(info, w) + } } }; } // impl_eigh! -impl_eigh!(f64, lapacke::dsyev); -impl_eigh!(f32, lapacke::ssyev); -impl_eigh!(c64, lapacke::zheev); -impl_eigh!(c32, lapacke::cheev); +impl_eigh!(f64, lapacke::dsyev, lapacke::dsygv); +impl_eigh!(f32, lapacke::ssyev, lapacke::ssygv); +impl_eigh!(c64, lapacke::zheev, lapacke::zhegv); +impl_eigh!(c32, lapacke::cheev, lapacke::chegv); diff --git a/src/lapack/svddc.rs b/src/lapack/svddc.rs index 9c59b7aa..22f770e3 100644 --- a/src/lapack/svddc.rs +++ b/src/lapack/svddc.rs @@ -3,10 +3,10 @@ use num_traits::Zero; use crate::error::*; use crate::layout::MatrixLayout; -use crate::types::*; use crate::svddc::UVTFlag; +use crate::types::*; -use super::{SVDOutput, into_result}; +use super::{into_result, SVDOutput}; pub trait SVDDC_: Scalar { unsafe fn svddc(l: MatrixLayout, jobz: UVTFlag, a: &mut [Self]) -> Result>; @@ -15,11 +15,7 @@ pub trait SVDDC_: Scalar { macro_rules! impl_svdd { ($scalar:ty, $gesdd:path) => { impl SVDDC_ for $scalar { - unsafe fn svddc( - l: MatrixLayout, - jobz: UVTFlag, - mut a: &mut [Self], - ) -> Result> { + unsafe fn svddc(l: MatrixLayout, jobz: UVTFlag, mut a: &mut [Self]) -> Result> { let (m, n) = l.size(); let k = m.min(n); let lda = l.lda(); @@ -51,11 +47,7 @@ macro_rules! impl_svdd { SVDOutput { s: s, u: if jobz == UVTFlag::None { None } else { Some(u) }, - vt: if jobz == UVTFlag::None { - None - } else { - Some(vt) - }, + vt: if jobz == UVTFlag::None { None } else { Some(vt) }, }, ) } diff --git a/src/lib.rs b/src/lib.rs index 198449cc..472cc348 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -36,6 +36,9 @@ //! - [Random matrix generators](generate/index.html) //! - [Scalar trait](types/trait.Scalar.html) +#[macro_use] +extern crate ndarray; + extern crate blas_src; extern crate lapack_src; @@ -50,6 +53,7 @@ pub mod inner; pub mod krylov; pub mod lapack; pub mod layout; +pub mod lobpcg; pub mod norm; pub mod operator; pub mod opnorm; @@ -70,6 +74,7 @@ pub use eigh::*; pub use generate::*; pub use inner::*; pub use layout::*; +pub use lobpcg::{TruncatedEig, TruncatedOrder, TruncatedSvd}; pub use norm::*; pub use operator::*; pub use opnorm::*; diff --git a/src/lobpcg/eig.rs b/src/lobpcg/eig.rs new file mode 100644 index 00000000..6f4d1ac4 --- /dev/null +++ b/src/lobpcg/eig.rs @@ -0,0 +1,184 @@ +use super::lobpcg::{lobpcg, LobpcgResult, Order}; +use crate::{Lapack, Scalar, generate}; +///! Implements truncated eigenvalue decomposition +/// +use ndarray::prelude::*; +use ndarray::stack; +use ndarray::ScalarOperand; +use num_traits::{Float, NumCast}; + +/// Truncated eigenproblem solver +/// +/// This struct wraps the LOBPCG algorithm and provides convenient builder-pattern access to +/// parameter like maximal iteration, precision and constraint matrix. Furthermore it allows +/// conversion into a iterative solver where each iteration step yields a new eigenvalue/vector +/// pair. +pub struct TruncatedEig { + order: Order, + problem: Array2, + pub constraints: Option>, + preconditioner: Option>, + precision: f32, + maxiter: usize, +} + +impl TruncatedEig { + pub fn new(problem: Array2, order: Order) -> TruncatedEig { + TruncatedEig { + precision: 1e-5, + maxiter: problem.len_of(Axis(0)) * 2, + preconditioner: None, + constraints: None, + order, + problem, + } + } + + pub fn precision(mut self, precision: f32) -> Self { + self.precision = precision; + + self + } + + pub fn maxiter(mut self, maxiter: usize) -> Self { + self.maxiter = maxiter; + + self + } + + pub fn orthogonal_to(mut self, constraints: Array2) -> Self { + self.constraints = Some(constraints); + + self + } + + pub fn precondition_with(mut self, preconditioner: Array2) -> Self { + self.preconditioner = Some(preconditioner); + + self + } + + // calculate the eigenvalues decompose + pub fn decompose(&self, num: usize) -> LobpcgResult { + let x: Array2 = generate::random((self.problem.len_of(Axis(0)), num)); + let x = x.mapv(|x| NumCast::from(x).unwrap()); + + if let Some(ref preconditioner) = self.preconditioner { + lobpcg( + |y| self.problem.dot(&y), + x, + |mut y| y.assign(&preconditioner.dot(&y)), + self.constraints.clone(), + self.precision, + self.maxiter, + self.order.clone(), + ) + } else { + lobpcg( + |y| self.problem.dot(&y), + x, + |_| {}, + self.constraints.clone(), + self.precision, + self.maxiter, + self.order.clone(), + ) + } + } +} + +impl IntoIterator for TruncatedEig { + type Item = (Array1, Array2); + type IntoIter = TruncatedEigIterator; + + fn into_iter(self) -> TruncatedEigIterator { + TruncatedEigIterator { + step_size: 1, + remaining: self.problem.len_of(Axis(0)), + eig: self, + } + } +} + +/// Truncate eigenproblem iterator +/// +/// This wraps a truncated eigenproblem and provides an iterator where each step yields a new +/// eigenvalue/vector pair. Useful for generating pairs until a certain condition is met. +pub struct TruncatedEigIterator { + step_size: usize, + remaining: usize, + eig: TruncatedEig, +} + +impl Iterator for TruncatedEigIterator { + type Item = (Array1, Array2); + + fn next(&mut self) -> Option { + if self.remaining == 0 { + return None; + } + + let step_size = usize::min(self.step_size, self.remaining); + let res = self.eig.decompose(step_size); + + match res { + LobpcgResult::Ok(vals, vecs, norms) | LobpcgResult::Err(vals, vecs, norms, _) => { + // abort if any eigenproblem did not converge + for r_norm in norms { + if r_norm > NumCast::from(0.1).unwrap() { + return None; + } + } + + // add the new eigenvector to the internal constrain matrix + let new_constraints = if let Some(ref constraints) = self.eig.constraints { + let eigvecs_arr = constraints + .gencolumns() + .into_iter() + .chain(vecs.gencolumns().into_iter()) + .map(|x| x.insert_axis(Axis(1))) + .collect::>(); + + stack(Axis(1), &eigvecs_arr).unwrap() + } else { + vecs.clone() + }; + + self.eig.constraints = Some(new_constraints); + self.remaining -= step_size; + + Some((vals, vecs)) + } + LobpcgResult::NoResult(_) => None, + } + } +} + +#[cfg(test)] +mod tests { + use super::Order; + use super::TruncatedEig; + use ndarray::{arr1, Array2}; + + #[test] + fn test_truncated_eig() { + let diag = arr1(&[ + 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12., 13., 14., 15., 16., 17., 18., 19., 20., + ]); + let a = Array2::from_diag(&diag); + + let teig = TruncatedEig::new(a, Order::Largest).precision(1e-5).maxiter(500); + + let res = teig.into_iter().take(3).flat_map(|x| x.0.to_vec()).collect::>(); + let ground_truth = vec![20., 19., 18.]; + + assert!( + ground_truth + .into_iter() + .zip(res.into_iter()) + .map(|(x, y)| (x - y) * (x - y)) + .sum::() + < 0.01 + ); + } +} diff --git a/src/lobpcg/lobpcg.rs b/src/lobpcg/lobpcg.rs new file mode 100644 index 00000000..b9f3e52c --- /dev/null +++ b/src/lobpcg/lobpcg.rs @@ -0,0 +1,573 @@ +///! Locally Optimal Block Preconditioned Conjugated +///! +///! This module implements the Locally Optimal Block Preconditioned Conjugated (LOBPCG) algorithm, +///which can be used as a solver for large symmetric positive definite eigenproblems. +use crate::error::{LinalgError, Result}; +use crate::{cholesky::*, close_l2, eigh::*, norm::*, triangular::*}; +use crate::{Lapack, Scalar}; +use ndarray::prelude::*; +use ndarray::{OwnedRepr, ScalarOperand, Data}; +use num_traits::{Float, NumCast}; + +/// Find largest or smallest eigenvalues +#[derive(Debug, Clone)] +pub enum Order { + Largest, + Smallest, +} + +/// The result of the eigensolver +/// +/// In the best case the eigensolver has converged with a result better than the given threshold, +/// then a `LobpcgResult::Ok` gives the eigenvalues, eigenvectors and norms. If an error ocurred +/// during the process, it is returned in `LobpcgResult::Err`, but the best result is still returned, +/// as it could be usable. If there is no result at all, then `LobpcgResult::NoResult` is returned. +/// This happens if the algorithm fails in an early stage, for example if the matrix `A` is not SPD +#[derive(Debug)] +pub enum LobpcgResult { + Ok(Array1, Array2, Vec), + Err(Array1, Array2, Vec, LinalgError), + NoResult(LinalgError), +} + +/// Solve full eigenvalue problem, sort by `order` and truncate to `size` +fn sorted_eig, A: Scalar + Lapack>( + a: ArrayBase, + b: Option>, + size: usize, + order: &Order, +) -> Result<(Array1, Array2)> { + let n = a.len_of(Axis(0)); + + let (vals, vecs) = match b { + Some(b) => (a, b).eigh(UPLO::Upper).map(|x| (x.0, (x.1).0))?, + _ => a.eigh(UPLO::Upper)?, + }; + + Ok(match order { + Order::Largest => ( + vals.slice_move(s![n-size..; -1]).mapv(|x| Scalar::from_real(x)), + vecs.slice_move(s![.., n-size..; -1]), + ), + Order::Smallest => ( + vals.slice_move(s![..size]).mapv(|x| Scalar::from_real(x)), + vecs.slice_move(s![.., ..size]), + ), + }) +} + +/// Masks a matrix with the given `matrix` +fn ndarray_mask(matrix: ArrayView2, mask: &[bool]) -> Array2 { + assert_eq!(mask.len(), matrix.ncols()); + + let indices = (0..mask.len()) + .zip(mask.into_iter()) + .filter(|(_, b)| **b) + .map(|(a, _)| a) + .collect::>(); + + matrix.select(Axis(1), &indices) +} + +/// Applies constraints ensuring that a matrix is orthogonal to it +/// +/// This functions takes a matrix `v` and constraint-matrix `y` and orthogonalize `v` to `y`. +fn apply_constraints( + mut v: ArrayViewMut, + cholesky_yy: &CholeskyFactorized>, + y: ArrayView2, +) { + let gram_yv = y.t().dot(&v); + + let u = gram_yv + .gencolumns() + .into_iter() + .map(|x| { + let res = cholesky_yy.solvec(&x).unwrap(); + + res.to_vec() + }) + .flatten() + .collect::>(); + + let rows = gram_yv.len_of(Axis(0)); + let u = Array2::from_shape_vec((rows, u.len() / rows), u).unwrap(); + + v -= &(y.dot(&u)); +} + +/// Orthonormalize `V` with Cholesky factorization +/// +/// This also returns the matrix `R` of the `QR` problem +fn orthonormalize(v: Array2) -> Result<(Array2, Array2)> { + let gram_vv = v.t().dot(&v); + let gram_vv_fac = gram_vv.cholesky(UPLO::Lower)?; + + close_l2( + &gram_vv, + &gram_vv_fac.dot(&gram_vv_fac.t()), + NumCast::from(1e-5).unwrap(), + ); + + let v_t = v.reversed_axes(); + let u = gram_vv_fac + .solve_triangular(UPLO::Lower, Diag::NonUnit, &v_t)? + .reversed_axes(); + + Ok((u, gram_vv_fac)) +} + +/// Eigenvalue solver for large symmetric positive definite (SPD) eigenproblems +/// +/// # Arguments +/// * `a` - An operator defining the problem, usually a sparse (sometimes also dense) matrix +/// multiplication. Also called the "stiffness matrix". +/// * `x` - Initial approximation of the k eigenvectors. If `a` has shape=(n,n), then `x` should +/// have shape=(n,k). +/// * `m` - Preconditioner to `a`, by default the identity matrix. Should approximate the inverse +/// of `a`. +/// * `y` - Constraints of (n,size_y), iterations are performed in the orthogonal complement of the +/// column-space of `y`. It must be full rank. +/// * `tol` - The tolerance values defines at which point the solver stops the optimization. The approximation +/// of a eigenvalue stops when then l2-norm of the residual is below this threshold. +/// * `maxiter` - The maximal number of iterations +/// * `order` - Whether to solve for the largest or lowest eigenvalues +/// +/// The function returns an `LobpcgResult` with the eigenvalue/eigenvector and achieved residual norm +/// for it. All iterations are tracked and the optimal solution returned. In case of an error a +/// special variant `LobpcgResult::NotConverged` additionally carries the error. This can happen when +/// the precision of the matrix is too low (switch then from `f32` to `f64` for example). +pub fn lobpcg< + A: Float + Scalar + Lapack + ScalarOperand + PartialOrd + Default, + F: Fn(ArrayView2) -> Array2, + G: Fn(ArrayViewMut2), +>( + a: F, + mut x: Array2, + m: G, + y: Option>, + tol: f32, + maxiter: usize, + order: Order, +) -> LobpcgResult { + // the initital approximation should be maximal square + // n is the dimensionality of the problem + let (n, size_x) = (x.nrows(), x.ncols()); + assert!(size_x <= n); + + /*let size_y = match y { + Some(ref y) => y.ncols(), + _ => 0, + }; + + if (n - size_y) < 5 * size_x { + panic!("Please use a different approach, the LOBPCG method only supports the calculation of a couple of eigenvectors!"); + }*/ + + // cap the number of iteration + let mut iter = usize::min(n * 10, maxiter); + let tol = NumCast::from(tol).unwrap(); + + // calculate cholesky factorization of YY' and apply constraints to initial guess + let cholesky_yy = y.as_ref().map(|y| { + let cholesky_yy = y.t().dot(y).factorizec(UPLO::Lower).unwrap(); + apply_constraints(x.view_mut(), &cholesky_yy, y.view()); + cholesky_yy + }); + + // orthonormalize the initial guess + let (x, _) = match orthonormalize(x) { + Ok(x) => x, + Err(err) => return LobpcgResult::NoResult(err), + }; + + // calculate AX and XAX for Rayleigh quotient + let ax = a(x.view()); + let xax = x.t().dot(&ax); + + // perform eigenvalue decomposition of XAX + let (mut lambda, eig_block) = match sorted_eig(xax.view(), None, size_x, &order) { + Ok(x) => x, + Err(err) => return LobpcgResult::NoResult(err), + }; + + // initiate approximation of the eigenvector + let mut x = x.dot(&eig_block); + let mut ax = ax.dot(&eig_block); + + // track residual below threshold + let mut activemask = vec![true; size_x]; + + // track residuals and best result + let mut residual_norms_history = Vec::new(); + let mut best_result = None; + + let mut previous_block_size = size_x; + + let mut ident: Array2 = Array2::eye(size_x); + let ident0: Array2 = Array2::eye(size_x); + let two: A = NumCast::from(2.0).unwrap(); + + let mut previous_p_ap: Option<(Array2, Array2)> = None; + let mut explicit_gram_flag = true; + + let final_norm = loop { + // calculate residual + let lambda_diag = Array2::from_diag(&lambda); + let lambda_x = x.dot(&lambda_diag); + + // calculate residual AX - lambdaX + let r = &ax - &lambda_x; + + // calculate L2 norm of error for every eigenvalue + let residual_norms = r.gencolumns().into_iter().map(|x| x.norm()).collect::>(); + residual_norms_history.push(residual_norms.clone()); + + // compare best result and update if we improved + let sum_rnorm: A::Real = residual_norms.iter().cloned().sum(); + if best_result + .as_ref() + .map(|x: &(_, _, Vec)| x.2.iter().cloned().sum::() > sum_rnorm) + .unwrap_or(true) + { + best_result = Some((lambda.clone(), x.clone(), residual_norms.clone())); + } + + // disable eigenvalues which are below the tolerance threshold + activemask = residual_norms + .iter() + .zip(activemask.iter()) + .map(|(x, a)| *x > tol && *a) + .collect(); + + // resize identity block if necessary + let current_block_size = activemask.iter().filter(|x| **x).count(); + if current_block_size != previous_block_size { + previous_block_size = current_block_size; + ident = Array2::eye(current_block_size); + } + + // if we are below the threshold for all eigenvalue or exceeded the number of iteration, + // abort + if current_block_size == 0 || iter == 0 { + break Ok(residual_norms); + } + + // select active eigenvalues, apply pre-conditioner, orthogonalize to Y and orthonormalize + let mut active_block_r = ndarray_mask(r.view(), &activemask); + // apply preconditioner + m(active_block_r.view_mut()); + // apply constraints to the preconditioned residuals + if let (Some(ref y), Some(ref cholesky_yy)) = (&y, &cholesky_yy) { + apply_constraints(active_block_r.view_mut(), cholesky_yy, y.view()); + } + // orthogonalize the preconditioned residual to x + active_block_r -= &x.dot(&x.t().dot(&active_block_r)); + + let (r, _) = match orthonormalize(active_block_r) { + Ok(x) => x, + Err(err) => break Err(err), + }; + + let ar = a(r.view()); + + // check whether `A` is of type `f32` or `f64` + let max_rnorm_float = if A::epsilon() > NumCast::from(1e-8).unwrap() { + NumCast::from(1.0).unwrap() + } else { + NumCast::from(1.0e-8).unwrap() + }; + + // if we are once below the max_rnorm, enable explicit gram flag + let max_norm = residual_norms.into_iter().fold(A::Real::neg_infinity(), A::Real::max); + explicit_gram_flag = max_norm <= max_rnorm_float || explicit_gram_flag; + + // perform the Rayleigh Ritz procedure + let xar = x.t().dot(&ar); + let mut rar = r.t().dot(&ar); + + // for small residuals calculate covariance matrices explicitely, otherwise approximate + // them such that X is orthogonal and uncorrelated to the residual R and use eigenvalues of + // previous decomposition + let (xax, xx, rr, xr) = if explicit_gram_flag { + rar = (&rar + &rar.t()) / two; + let xax = x.t().dot(&ax); + + ((&xax + &xax.t()) / two, x.t().dot(&x), r.t().dot(&r), x.t().dot(&r)) + } else { + ( + lambda_diag, + ident0.clone(), + ident.clone(), + Array2::zeros((size_x, current_block_size)), + ) + }; + + // mask and orthonormalize P and AP + let mut p_ap = previous_p_ap + .as_ref() + .and_then(|(p, ap)| { + let active_p = ndarray_mask(p.view(), &activemask); + let active_ap = ndarray_mask(ap.view(), &activemask); + + orthonormalize(active_p).map(|x| (active_ap, x)).ok() + }) + .and_then(|(active_ap, (active_p, p_r))| { + // orthonormalize AP with R^{-1} of A + let active_ap = active_ap.reversed_axes(); + p_r.solve_triangular(UPLO::Lower, Diag::NonUnit, &active_ap) + .map(|active_ap| (active_p, active_ap.reversed_axes())) + .ok() + }); + + // compute symmetric gram matrices and calculate solution of eigenproblem + // + // first try to compute the eigenvalue decomposition of the span{R, X, P}, + // if this fails (or the algorithm was restarted), then just use span{R, X} + let result = p_ap.as_ref() + .ok_or(LinalgError::Lapack { return_code: 1 }) + .and_then(|(active_p, active_ap)| { + let xap = x.t().dot(active_ap); + let rap = r.t().dot(active_ap); + let pap = active_p.t().dot(active_ap); + let xp = x.t().dot(active_p); + let rp = r.t().dot(active_p); + let (pap, pp) = if explicit_gram_flag { + ((&pap + &pap.t()) / two, active_p.t().dot(active_p)) + } else { + (pap, ident.clone()) + }; + + sorted_eig( + stack![ + Axis(0), + stack![Axis(1), xax, xar, xap], + stack![Axis(1), xar.t(), rar, rap], + stack![Axis(1), xap.t(), rap.t(), pap] + ], + Some(stack![ + Axis(0), + stack![Axis(1), xx, xr, xp], + stack![Axis(1), xr.t(), rr, rp], + stack![Axis(1), xp.t(), rp.t(), pp] + ]), + size_x, + &order + ) + }) + .or_else(|_| { + p_ap = None; + + sorted_eig( + stack![Axis(0), stack![Axis(1), xax, xar], stack![Axis(1), xar.t(), rar]], + Some(stack![Axis(0), stack![Axis(1), xx, xr], stack![Axis(1), xr.t(), rr]]), + size_x, + &order + ) + }); + + + // update eigenvalues and eigenvectors (lambda is also used in the next iteration) + let eig_vecs; + match result { + Ok((x, y)) => { + lambda = x; + eig_vecs = y; + }, + Err(x) => break Err(x) + } + + // approximate eigenvector X and conjugate vectors P with solution of eigenproblem + let (p, ap, tau) = if let Some((active_p, active_ap)) = p_ap { + // tau are eigenvalues to basis of X + let tau = eig_vecs.slice(s![..size_x, ..]); + // alpha are eigenvalues to basis of R + let alpha = eig_vecs.slice(s![size_x..size_x + current_block_size, ..]); + // gamma are eigenvalues to basis of P + let gamma = eig_vecs.slice(s![size_x + current_block_size.., ..]); + + // update AP and P in span{R, P} as linear combination + let updated_p = r.dot(&alpha) + active_p.dot(&gamma); + let updated_ap = ar.dot(&alpha) + active_ap.dot(&gamma); + + (updated_p, updated_ap, tau) + } else { + // tau are eigenvalues to basis of X + let tau = eig_vecs.slice(s![..size_x, ..]); + // alpha are eigenvalues to basis of R + let alpha = eig_vecs.slice(s![size_x.., ..]); + + // update AP and P as linear combination of the residual matrix R + let updated_p = r.dot(&alpha); + let updated_ap = ar.dot(&alpha); + + (updated_p, updated_ap, tau) + }; + + // update approximation of X as linear combinations of span{X, P, R} + x = x.dot(&tau) + &p; + ax = ax.dot(&tau) + ≈ + + previous_p_ap = Some((p, ap)); + + iter -= 1; + }; + + // retrieve best result and convert norm into `A` + let (vals, vecs, rnorm) = best_result.unwrap(); + let rnorm = rnorm.into_iter().map(|x| Scalar::from_real(x)).collect(); + + match final_norm { + Ok(_) => LobpcgResult::Ok(vals, vecs, rnorm), + Err(err) => LobpcgResult::Err(vals, vecs, rnorm, err), + } +} + +#[cfg(test)] +mod tests { + use super::lobpcg; + use super::ndarray_mask; + use super::orthonormalize; + use super::sorted_eig; + use super::LobpcgResult; + use super::Order; + use crate::close_l2; + use crate::qr::*; + use crate::generate; + use ndarray::prelude::*; + + /// Test the `sorted_eigen` function + #[test] + fn test_sorted_eigen() { + let matrix: Array2 = generate::random((10, 10)) * 10.0; + let matrix = matrix.t().dot(&matrix); + + // return all eigenvectors with largest first + let (vals, vecs) = sorted_eig(matrix.view(), None, 10, &Order::Largest).unwrap(); + + // calculate V * A * V' and compare to original matrix + let diag = Array2::from_diag(&vals); + let rec = (vecs.dot(&diag)).dot(&vecs.t()); + + close_l2(&matrix, &rec, 1e-5); + } + + /// Test the masking function + #[test] + fn test_masking() { + let matrix: Array2 = generate::random((10, 5)) * 10.0; + let masked_matrix = ndarray_mask(matrix.view(), &[true, true, false, true, false]); + close_l2(&masked_matrix.slice(s![.., 2]), &matrix.slice(s![.., 3]), 1e-12); + } + + /// Test orthonormalization of a random matrix + #[test] + fn test_orthonormalize() { + let matrix: Array2 = generate::random((10, 10)) * 10.0; + + let (n, l) = orthonormalize(matrix.clone()).unwrap(); + + // check for orthogonality + let identity = n.dot(&n.t()); + close_l2(&identity, &Array2::eye(10), 1e-2); + + // compare returned factorization with QR decomposition + let (_, r) = matrix.qr().unwrap(); + close_l2(&r.mapv(|x| x.abs()), &l.t().mapv(|x| x.abs()), 1e-2); + } + + fn assert_symmetric(a: &Array2) { + close_l2(a, &a.t(), 1e-5); + } + + fn check_eigenvalues(a: &Array2, order: Order, num: usize, ground_truth_eigvals: &[f64]) { + assert_symmetric(a); + + let n = a.len_of(Axis(0)); + let x: Array2 = generate::random((n, num)); + + let result = lobpcg(|y| a.dot(&y), x, |_| {}, None, 1e-5, n * 2, order); + match result { + LobpcgResult::Ok(vals, _, r_norms) | LobpcgResult::Err(vals, _, r_norms, _) => { + // check convergence + for (i, norm) in r_norms.into_iter().enumerate() { + if norm > 1e-5 { + println!("==== Assertion Failed ===="); + println!("The {}th eigenvalue estimation did not converge!", i); + panic!("Too large deviation of residual norm: {} > 0.01", norm); + } + } + + // check correct order of eigenvalues + if ground_truth_eigvals.len() == num { + close_l2(&Array1::from(ground_truth_eigvals.to_vec()), &vals, num as f64 * 5e-4) + } + } + LobpcgResult::NoResult(err) => panic!("Did not converge: {:?}", err), + } + } + + /// Test the eigensolver with a identity matrix problem and a random initial solution + #[test] + fn test_eigsolver_diag() { + let diag = arr1(&[ + 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12., 13., 14., 15., 16., 17., 18., 19., 20., + ]); + let a = Array2::from_diag(&diag); + + check_eigenvalues(&a, Order::Largest, 3, &[20., 19., 18.]); + check_eigenvalues(&a, Order::Smallest, 3, &[1., 2., 3.]); + } + + /// Test the eigensolver with matrix of constructed eigenvalues + #[test] + fn test_eigsolver_constructed() { + let n = 50; + let tmp = generate::random((n, n)); + //let (v, _) = tmp.qr_square().unwrap(); + let (v, _) = orthonormalize(tmp).unwrap(); + + // set eigenvalues in decreasing order + let t = Array2::from_diag(&Array1::linspace(n as f64, -(n as f64), n)); + let a = v.dot(&t.dot(&v.t())); + + // find five largest eigenvalues + check_eigenvalues(&a, Order::Largest, 5, &[50.0, 48.0, 46.0, 44.0, 42.0]); + check_eigenvalues(&a, Order::Smallest, 5, &[-50.0, -48.0, -46.0, -44.0, -42.0]); + } + + #[test] + fn test_eigsolver_constrained() { + let diag = arr1(&[1., 2., 3., 4., 5., 6., 7., 8., 9., 10.]); + let a = Array2::from_diag(&diag); + let x: Array2 = generate::random((10, 1)); + let y: Array2 = arr2(&[ + [1.0, 0., 0., 0., 0., 0., 0., 0., 0., 0.], + [0., 1.0, 0., 0., 0., 0., 0., 0., 0., 0.], + ]) + .reversed_axes(); + + let result = lobpcg(|y| a.dot(&y), x, |_| {}, Some(y), 1e-10, 50, Order::Smallest); + match result { + LobpcgResult::Ok(vals, vecs, r_norms) | LobpcgResult::Err(vals, vecs, r_norms, _) => { + // check convergence + for (i, norm) in r_norms.into_iter().enumerate() { + if norm > 0.01 { + println!("==== Assertion Failed ===="); + println!("The {}th eigenvalue estimation did not converge!", i); + panic!("Too large deviation of residual norm: {} > 0.01", norm); + } + } + + // should be the third eigenvalue + close_l2(&vals, &Array1::from(vec![3.0]), 1e-10); + close_l2( + &vecs.column(0).mapv(|x| x.abs()), + &arr1(&[0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + 1e-5, + ); + } + LobpcgResult::NoResult(err) => panic!("Did not converge: {:?}", err), + } + } +} diff --git a/src/lobpcg/mod.rs b/src/lobpcg/mod.rs new file mode 100644 index 00000000..0c0bce47 --- /dev/null +++ b/src/lobpcg/mod.rs @@ -0,0 +1,7 @@ +mod eig; +mod lobpcg; +mod svd; + +pub use eig::TruncatedEig; +pub use lobpcg::{lobpcg, LobpcgResult, Order as TruncatedOrder}; +pub use svd::TruncatedSvd; diff --git a/src/lobpcg/svd.rs b/src/lobpcg/svd.rs new file mode 100644 index 00000000..bc85f06f --- /dev/null +++ b/src/lobpcg/svd.rs @@ -0,0 +1,226 @@ +///! Truncated singular value decomposition +///! +///! This module computes the k largest/smallest singular values/vectors for a dense matrix. +use super::lobpcg::{lobpcg, LobpcgResult, Order}; +use crate::error::Result; +use crate::{Lapack, Scalar, generate}; +use ndarray::prelude::*; +use ndarray::ScalarOperand; +use num_traits::{Float, NumCast}; +use std::ops::DivAssign; + +/// The result of a eigenvalue decomposition, not yet transformed into singular values/vectors +/// +/// Provides methods for either calculating just the singular values with reduced cost or the +/// vectors with additional cost of matrix multiplication. +#[derive(Debug)] +pub struct TruncatedSvdResult { + eigvals: Array1, + eigvecs: Array2, + problem: Array2, + ngm: bool, +} + +impl + 'static + MagnitudeCorrection> TruncatedSvdResult { + /// Returns singular values ordered by magnitude with indices. + fn singular_values_with_indices(&self) -> (Array1, Vec) { + // numerate eigenvalues + let mut a = self.eigvals.iter().enumerate().collect::>(); + + // sort by magnitude + a.sort_by(|(_, x), (_, y)| x.partial_cmp(&y).unwrap().reverse()); + + // calculate cut-off magnitude (borrowed from scipy) + let cutoff = A::epsilon() * // float precision + A::correction() * // correction term (see trait below) + *a[0].1; // max eigenvalue + + // filter low singular values away + let (values, indices): (Vec, Vec) = a + .into_iter() + .filter(|(_, x)| *x > &cutoff) + .map(|(a, b)| (b.sqrt(), a)) + .unzip(); + + (Array1::from(values), indices) + } + + /// Returns singular values ordered by magnitude + pub fn values(&self) -> Array1 { + let (values, _) = self.singular_values_with_indices(); + + values + } + + /// Returns singular values, left-singular vectors and right-singular vectors + pub fn values_vectors(&self) -> (Array2, Array1, Array2) { + let (values, indices) = self.singular_values_with_indices(); + + // branch n > m (for A is [n x m]) + let (u, v) = if self.ngm { + let vlarge = self.eigvecs.select(Axis(1), &indices); + let mut ularge = self.problem.dot(&vlarge); + + ularge + .gencolumns_mut() + .into_iter() + .zip(values.iter()) + .for_each(|(mut a, b)| a.mapv_inplace(|x| x / *b)); + + (ularge, vlarge) + } else { + let ularge = self.eigvecs.select(Axis(1), &indices); + + let mut vlarge = self.problem.t().dot(&ularge); + vlarge + .gencolumns_mut() + .into_iter() + .zip(values.iter()) + .for_each(|(mut a, b)| a.mapv_inplace(|x| x / *b)); + + (ularge, vlarge) + }; + + (u, values, v.reversed_axes()) + } +} + +/// Truncated singular value decomposition +/// +/// Wraps the LOBPCG algorithm and provides convenient builder-pattern access to +/// parameter like maximal iteration, precision and constraint matrix. +pub struct TruncatedSvd { + order: Order, + problem: Array2, + precision: f32, + maxiter: usize, +} + +impl TruncatedSvd { + pub fn new(problem: Array2, order: Order) -> TruncatedSvd { + TruncatedSvd { + precision: 1e-5, + maxiter: problem.len_of(Axis(0)) * 2, + order, + problem, + } + } + + pub fn precision(mut self, precision: f32) -> Self { + self.precision = precision; + + self + } + + pub fn maxiter(mut self, maxiter: usize) -> Self { + self.maxiter = maxiter; + + self + } + + // calculate the eigenvalue decomposition + pub fn decompose(self, num: usize) -> Result> { + if num < 1 { + panic!("The number of singular values to compute should be larger than zero!"); + } + + let (n, m) = (self.problem.nrows(), self.problem.ncols()); + + // generate initial matrix + let x: Array2 = generate::random((usize::min(n, m), num)); + let x = x.mapv(|x| NumCast::from(x).unwrap()); + + // square precision because the SVD squares the eigenvalue as well + let precision = self.precision * self.precision; + + // use problem definition with less operations required + let res = if n > m { + lobpcg( + |y| self.problem.t().dot(&self.problem.dot(&y)), + x, + |_| {}, + None, + precision, + self.maxiter, + self.order.clone(), + ) + } else { + lobpcg( + |y| self.problem.dot(&self.problem.t().dot(&y)), + x, + |_| {}, + None, + precision, + self.maxiter, + self.order.clone(), + ) + }; + + // convert into TruncatedSvdResult + match res { + LobpcgResult::Ok(vals, vecs, _) | LobpcgResult::Err(vals, vecs, _, _) => Ok(TruncatedSvdResult { + problem: self.problem.clone(), + eigvals: vals, + eigvecs: vecs, + ngm: n > m, + }), + LobpcgResult::NoResult(err) => Err(err), + } + } +} + +pub trait MagnitudeCorrection { + fn correction() -> Self; +} + +impl MagnitudeCorrection for f32 { + fn correction() -> Self { + 1.0e3 + } +} + +impl MagnitudeCorrection for f64 { + fn correction() -> Self { + 1.0e6 + } +} + +#[cfg(test)] +mod tests { + use super::Order; + use super::TruncatedSvd; + use crate::{close_l2, generate}; + + use ndarray::{arr1, arr2, Array2}; + + #[test] + fn test_truncated_svd() { + let a = arr2(&[[3., 2., 2.], [2., 3., -2.]]); + + let res = TruncatedSvd::new(a, Order::Largest) + .precision(1e-5) + .maxiter(10) + .decompose(2) + .unwrap(); + + let (_, sigma, _) = res.values_vectors(); + + close_l2(&sigma, &arr1(&[5.0, 3.0]), 1e-5); + } + + #[test] + fn test_truncated_svd_random() { + let a: Array2 = generate::random((50, 10)); + + let res = TruncatedSvd::new(a.clone(), Order::Largest) + .precision(1e-5) + .maxiter(10) + .decompose(10) + .unwrap(); + + let (u, sigma, v_t) = res.values_vectors(); + let reconstructed = u.dot(&Array2::from_diag(&sigma).dot(&v_t)); + + close_l2(&a, &reconstructed, 1e-5); + } +} diff --git a/tests/svddc.rs b/tests/svddc.rs index 1dbbf32b..2c9204c8 100644 --- a/tests/svddc.rs +++ b/tests/svddc.rs @@ -14,7 +14,7 @@ fn test(a: &Array2, flag: UVTFlag) { assert!(u.is_none()); assert!(vt.is_none()); return; - }, + } }; let u: Array2<_> = u.unwrap(); let vt: Array2<_> = vt.unwrap(); @@ -53,7 +53,7 @@ macro_rules! test_svd_impl { let a = random(($n, $m).f()); test(&a, UVTFlag::Full); } - + #[test] fn []() { let a = random(($n, $m).f());