Add proptest for lobpcg eig and found problematic test case

Author: YuhanLiin

Cargo.toml

@@ -25,8 +25,3 @@ rand_xoshiro = { version = "0.6" }
 [features]
 default = ["iterative"]
 iterative = ["rand"]
-
-[[test]]
-name = "lobpcg"
-path = "tests/lobpcg.rs"
-required-features = ["iterative"]

src/lobpcg/mod.rs

@@ -45,8 +45,10 @@ where
 /// as it could be of value. If there is no result at all, then the second field is `None`.
 /// This happens if the algorithm fails in an early stage, for example if the matrix `A` is not SPD
 pub type LobpcgResult<A> = std::result::Result<Lobpcg<A>, (LinalgError, Option<Lobpcg<A>>)>;
+
+#[derive(Debug, Clone, PartialEq)]
 pub struct Lobpcg<A> {
-    eigvals: Array1<A>,
-    eigvecs: Array2<A>,
+    pub eigvals: Array1<A>,
+    pub eigvecs: Array2<A>,
     rnorm: Vec<A>,
 }

src/lobpcg/svd.rs

@@ -268,8 +268,7 @@ mod tests {
     use super::TruncatedSvd;
 
     use approx::assert_abs_diff_eq;
-    use ndarray::{arr1, arr2, s, Array1, Array2, NdFloat};
-    use ndarray_rand::{rand_distr::StandardNormal, RandomExt};
+    use ndarray::{arr1, arr2, Array2, NdFloat};
     use rand::distributions::{Distribution, Standard};
     use rand::SeedableRng;
     use rand_xoshiro::Xoshiro256Plus;

tests/cholesky.rs

@@ -6,18 +6,6 @@ use ndarray_linalg_rs::{cholesky::*, triangular::*};
 
 mod common;
 
-prop_compose! {
-    fn hpd_arr()
-        (arr in common::square_arr()) -> Array2<f64> {
-        let dim = arr.nrows();
-        let mut mul = arr.t().dot(&arr);
-        for i in 0..dim {
-            mul[(i, i)] += 1.0;
-        }
-        mul
-    }
-}
-
 fn run_cholesky_test(orig: Array2<f64>) {
     let chol = orig.cholesky().unwrap();
     assert_abs_diff_eq!(chol.dot(&chol.t()), orig, epsilon = 1e-7);
@@ -69,17 +57,17 @@ fn run_invc_test(a: Array2<f64>) {
 proptest! {
     #![proptest_config(ProptestConfig::with_cases(1000))]
     #[test]
-    fn cholesky_test(arr in hpd_arr()) {
+    fn cholesky_test(arr in common::hpd_arr()) {
         run_cholesky_test(arr)
     }
 
     #[test]
-    fn solvec_test((a, x) in common::system_of_arr(hpd_arr())) {
+    fn solvec_test((a, x) in common::system_of_arr(common::hpd_arr())) {
         run_solvec_test(a, x)
     }
 
     #[test]
-    fn invc_test(arr in hpd_arr()) {
+    fn invc_test(arr in common::hpd_arr()) {
         run_invc_test(arr)
     }
 }

tests/common.rs

@@ -2,6 +2,7 @@
 
 use std::ops::RangeInclusive;
 
+use approx::assert_abs_diff_eq;
 use ndarray::prelude::*;
 use proptest::prelude::*;
 use proptest_derive::Arbitrary;
@@ -49,6 +50,18 @@ prop_compose! {
     }
 }
 
+prop_compose! {
+    pub fn hpd_arr()
+        (arr in square_arr()) -> Array2<f64> {
+        let dim = arr.nrows();
+        let mut mul = arr.t().dot(&arr);
+        for i in 0..dim {
+            mul[(i, i)] += 1.0;
+        }
+        mul
+    }
+}
+
 prop_compose! {
     pub fn rect_arr()(rows in DIM_RANGE, cols in DIM_RANGE)
         (arr in matrix(rows, cols)) -> Array2<f64> {
@@ -93,3 +106,12 @@ pub fn system_of_arr(
         )
     })
 }
+
+pub fn check_eigh(arr: &Array2<f64>, vals: &Array1<f64>, vecs: &Array2<f64>) {
+    // Original array multiplied with eigenvec should equal eigenval times eigenvec
+    for (i, v) in vecs.axis_iter(Axis(1)).enumerate() {
+        let av = arr.dot(&v);
+        let ev = v.mapv(|x| vals[i] * x);
+        assert_abs_diff_eq!(av, ev, epsilon = 1e-5);
+    }
+}

tests/eigh.rs

@@ -6,15 +6,6 @@ use ndarray_linalg_rs::eigh::*;
 
 mod common;
 
-fn check_eigh(arr: &Array2<f64>, vals: &Array1<f64>, vecs: &Array2<f64>) {
-    // Original array multiplied with eigenvec should equal eigenval times eigenvec
-    for (i, v) in vecs.axis_iter(Axis(1)).enumerate() {
-        let av = arr.dot(&v);
-        let ev = v.mapv(|x| vals[i] * x);
-        assert_abs_diff_eq!(av, ev, epsilon = 1e-5);
-    }
-}
-
 fn run_eigh_test(arr: Array2<f64>) {
     let n = arr.nrows();
     let d = arr.eigh().unwrap();
@@ -23,7 +14,7 @@ fn run_eigh_test(arr: Array2<f64>) {
     // Eigenvecs should be orthogonal
     let s = vecs.t().dot(&vecs);
     assert_abs_diff_eq!(s, Array2::eye(n), epsilon = 1e-5);
-    check_eigh(&arr, &vals, &vecs);
+    common::check_eigh(&arr, &vals, &vecs);
 
     let (evals, evecs) = arr.clone().eigh_into().unwrap();
     assert_abs_diff_eq!(evals, vals);
@@ -33,12 +24,12 @@ fn run_eigh_test(arr: Array2<f64>) {
 
     // Check if ascending eigen is actually sorted and valid
     let (vals, vecs) = d.clone().sort_eig_asc();
-    check_eigh(&arr, &vals, &vecs);
+    common::check_eigh(&arr, &vals, &vecs);
     assert!(vals.windows(2).into_iter().all(|w| w[0] <= w[1]));
 
     // Check if descending eigen is actually sorted and valid
     let (vals, vecs) = d.sort_eig_desc();
-    check_eigh(&arr, &vals, &vecs);
+    common::check_eigh(&arr, &vals, &vecs);
     assert!(vals.windows(2).into_iter().all(|w| w[0] >= w[1]));
 }
 

tests/lobpcg.rs

@@ -2,11 +2,12 @@ use approx::assert_abs_diff_eq;
 use ndarray::prelude::*;
 use ndarray_rand::{rand_distr::StandardNormal, RandomExt};
 use proptest::prelude::*;
-
-use ndarray_linalg_rs::lobpcg::*;
 use rand::SeedableRng;
 use rand_xoshiro::Xoshiro256Plus;
 
+use ndarray_linalg_rs::eigh::*;
+use ndarray_linalg_rs::lobpcg::*;
+
 mod common;
 
 /// Eigenvalue structure in high dimensions
@@ -63,3 +64,44 @@ fn test_marchenko_pastur() {
         assert_abs_diff_eq!(mp_law, empirical, epsilon = 0.05);
     }
 }
+
+fn run_lobpcg_eig_test(arr: Array2<f64>, num: usize, ordering: Order) {
+    let (eigvals, _) = arr.eigh().unwrap().sort_eig(ordering == Order::Largest);
+    let res = TruncatedEig::new_with_rng(arr.clone(), ordering, Xoshiro256Plus::seed_from_u64(42))
+        .precision(1e-3)
+        .decompose(num)
+        .unwrap_or_else(|e| e.1.unwrap());
+
+    assert_abs_diff_eq!(eigvals.slice(s![..num]), res.eigvals, epsilon = 1e-5);
+    common::check_eigh(&arr, &res.eigvals, &res.eigvecs);
+}
+
+fn generate_order() -> impl Strategy<Value = Order> {
+    prop_oneof![Just(Order::Largest), Just(Order::Smallest)]
+}
+
+prop_compose! {
+    pub fn hpd_arr_num()(arr in common::hpd_arr())
+        (num in (1..arr.ncols()), arr in Just(arr)) -> (Array2<f64>, usize) {
+        (arr, num)
+    }
+}
+
+proptest! {
+    #![proptest_config(ProptestConfig::with_cases(1000))]
+    #[test]
+    fn lobpcg_eig_test((arr, num) in hpd_arr_num(), ordering in generate_order()) {
+        run_lobpcg_eig_test(arr, num, ordering);
+    }
+}
+
+#[test]
+fn problematic_eig_matrix() {
+    let arr = array![
+        [1.0, 0.0, 0.0, 0.0],
+        [0.0, 1.0, 0.0, 0.0],
+        [0.0, 0.0, 7854.796948298437, 2495.5155877621937],
+        [0.0, 0.0, 2495.5155877621937, 5995.696530257453]
+    ];
+    run_lobpcg_eig_test(arr, 3, Order::Largest);
+}