Save coefficients in Arnoldi struct

Author: termoshtt

src/krylov/arnoldi.rs

@@ -13,6 +13,8 @@ where
     a: F,
     v: ArrayBase<S, Ix1>,
     ortho: Ortho,
+    /// Coefficients
+    h: Vec<Array1<A>>,
 }
 
 impl<A, S, F, Ortho> Arnoldi<A, S, F, Ortho>
@@ -29,28 +31,33 @@ where
         let norm = v.norm_l2();
         azip!(mut v(&mut v) in { *v = v.div_real(norm) });
         ortho.append(v.view());
-        Arnoldi { a, v, ortho }
+        Arnoldi {
+            a,
+            v,
+            ortho,
+            h: Vec::new(),
+        }
+    }
+
+    /// Dimension of Krylov subspace
+    pub fn dim(&self) -> usize {
+        self.ortho.len()
     }
 
     /// Iterate until convergent
-    pub fn complete(self) -> (Q<A>, H<A>) {
+    pub fn complete(mut self) -> (Q<A>, H<A>) {
+        for _ in &mut self {} // execute iteration until convergent
         let q = self.ortho.get_q();
-        let hs: Vec<Array1<A>> = self.collect();
-        let n = hs.len();
+        let n = self.dim();
         let mut h = Array2::zeros((n, n).f());
-        for (i, hc) in hs.iter().enumerate() {
+        for (i, hc) in self.h.iter().enumerate() {
             let m = std::cmp::max(n, i + 1);
             for j in 0..m {
                 h[(j, i)] = hc[j];
             }
         }
         (q, h)
     }
-
-    /// Dimension of Krylov subspace
-    pub fn krylov_dimension(&self) -> usize {
-        self.ortho.len()
-    }
 }
 
 impl<A, S, F, Ortho> Iterator for Arnoldi<A, S, F, Ortho>
@@ -69,57 +76,45 @@ where
         if result.is_dependent() {
             None
         } else {
-            Some(result.into_coeff())
+            let coef = result.into_coeff();
+            self.h.push(coef.clone());
+            Some(coef)
         }
     }
 }
 
-pub fn arnoldi_householder<A, S1, S2>(
-    a: ArrayBase<S1, Ix2>,
-    v: ArrayBase<S2, Ix1>,
-    tol: A::Real,
-) -> impl Iterator<Item = Array1<A>>
+/// Interpret a matrix as a linear operator
+pub fn mul_mat<A, S1, S2>(a: ArrayBase<S1, Ix2>) -> impl Fn(&mut ArrayBase<S2, Ix1>)
 where
-    A: Scalar + Lapack,
+    A: Scalar,
     S1: Data<Elem = A>,
     S2: DataMut<Elem = A>,
 {
     let (n, m) = a.dim();
     assert_eq!(n, m, "Input matrix must be square");
-    assert_eq!(m, v.len(), "Input matrix and vector sizes mismach");
-
-    let householder = Householder::new(n, tol);
-    Arnoldi::new(
-        move |x| {
-            let ax = a.dot(x);
-            azip!(mut x(x), ax in { *x = ax });
-        },
-        v,
-        householder,
-    )
+    move |x| {
+        assert_eq!(m, x.len(), "Input matrix and vector sizes mismatch");
+        let ax = a.dot(x);
+        azip!(mut x(x), ax in { *x = ax });
+    }
 }
 
-pub fn arnoldi_mgs<A, S1, S2>(
-    a: ArrayBase<S1, Ix2>,
-    v: ArrayBase<S2, Ix1>,
-    tol: A::Real,
-) -> impl Iterator<Item = Array1<A>>
+pub fn arnoldi_householder<A, S1, S2>(a: ArrayBase<S1, Ix2>, v: ArrayBase<S2, Ix1>, tol: A::Real) -> (Q<A>, H<A>)
 where
     A: Scalar + Lapack,
     S1: Data<Elem = A>,
     S2: DataMut<Elem = A>,
 {
-    let (n, m) = a.dim();
-    assert_eq!(n, m, "Input matrix must be square");
-    assert_eq!(m, v.len(), "Input matrix and vector sizes mismach");
+    let householder = Householder::new(v.len(), tol);
+    Arnoldi::new(mul_mat(a), v, householder).complete()
+}
 
-    let mgs = MGS::new(n, tol);
-    Arnoldi::new(
-        move |x| {
-            let ax = a.dot(x);
-            azip!(mut x(x), ax in { *x = ax });
-        },
-        v,
-        mgs,
-    )
+pub fn arnoldi_mgs<A, S1, S2>(a: ArrayBase<S1, Ix2>, v: ArrayBase<S2, Ix1>, tol: A::Real) -> (Q<A>, H<A>)
+where
+    A: Scalar + Lapack,
+    S1: Data<Elem = A>,
+    S2: DataMut<Elem = A>,
+{
+    let mgs = MGS::new(v.len(), tol);
+    Arnoldi::new(mul_mat(a), v, mgs).complete()
 }