Merge pull request #184 from bytesnake/master

Add LOBPCG solver for large symmetric positive definite eigenproblems

Author: termoshtt

Cargo.toml

@@ -52,6 +52,11 @@ optional = true
 paste = "0.1.9"
 criterion = "0.3.1"
 
+[[bench]]
+name = "truncated_eig"
+harness = false
+
 [[bench]]
 name = "eigh"
 harness = false
+

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 [<teig $n>](c: &mut Criterion) {
+                c.bench_function(&format!("truncated_eig{}", $n), |b| {
+                    let a: Array2<f64> = 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<f64> = 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);

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]);
+    }
+}

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);
+}

src/eigh.rs

@@ -42,6 +42,20 @@ where
     }
 }
 
+impl<A, S, S2> EighInto for (ArrayBase<S, Ix2>, ArrayBase<S2, Ix2>)
+where
+    A: Scalar + Lapack,
+    S: DataMut<Elem = A>,
+    S2: DataMut<Elem = A>,
+{
+    type EigVal = Array1<A::Real>;
+
+    fn eigh_into(mut self, uplo: UPLO) -> Result<(Self::EigVal, Self)> {
+        let (val, _) = self.eigh_inplace(uplo)?;
+        Ok((val, self))
+    }
+}
+
 impl<A, S> Eigh for ArrayBase<S, Ix2>
 where
     A: Scalar + Lapack,
@@ -56,6 +70,21 @@ where
     }
 }
 
+impl<A, S, S2> Eigh for (ArrayBase<S, Ix2>, ArrayBase<S2, Ix2>)
+where
+    A: Scalar + Lapack,
+    S: Data<Elem = A>,
+    S2: Data<Elem = A>,
+{
+    type EigVal = Array1<A::Real>;
+    type EigVec = (Array2<A>, Array2<A>);
+
+    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<A, S> EighInplace for ArrayBase<S, Ix2>
 where
     A: Scalar + Lapack,
@@ -75,6 +104,42 @@ where
     }
 }
 
+impl<A, S, S2> EighInplace for (ArrayBase<S, Ix2>, ArrayBase<S2, Ix2>)
+where
+    A: Scalar + Lapack,
+    S: DataMut<Elem = A>,
+    S2: DataMut<Elem = A>,
+{
+    type EigVal = Array1<A::Real>;
+
+    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;

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<Vec<Self::Real>>;
+    unsafe fn eigh_generalized(
+        calc_eigenvec: bool,
+        l: MatrixLayout,
+        uplo: UPLO,
+        a: &mut [Self],
+        b: &mut [Self],
+    ) -> Result<Vec<Self::Real>>;
 }
 
 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<Vec<Self::Real>> {
                 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<Vec<Self::Real>> {
+                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);

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<SVDOutput<Self>>;
@@ -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<SVDOutput<Self>> {
+            unsafe fn svddc(l: MatrixLayout, jobz: UVTFlag, mut a: &mut [Self]) -> Result<SVDOutput<Self>> {
                 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) },
                     },
                 )
             }

src/lib.rs

@@ -37,6 +37,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;
 
@@ -52,6 +55,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;
@@ -73,6 +77,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::*;