Merge pull request #155 from rust-ndarray/arnoldi2

Arnoldi iterator

Author: termoshtt

src/krylov/arnoldi.rs

@@ -0,0 +1,135 @@
+//! Arnoldi iteration
+
+use super::*;
+use crate::norm::Norm;
+use num_traits::One;
+use std::iter::*;
+
+/// Execute Arnoldi iteration as Rust iterator
+///
+/// - [Arnoldi iteration - Wikipedia](https://en.wikipedia.org/wiki/Arnoldi_iteration)
+///
+pub struct Arnoldi<A, S, F, Ortho>
+where
+    A: Scalar,
+    S: DataMut<Elem = A>,
+    F: Fn(&mut ArrayBase<S, Ix1>),
+    Ortho: Orthogonalizer<Elem = A>,
+{
+    a: F,
+    /// Next vector (normalized `|v|=1`)
+    v: ArrayBase<S, Ix1>,
+    /// Orthogonalizer
+    ortho: Ortho,
+    /// Coefficients to be composed into H-matrix
+    h: Vec<Array1<A>>,
+}
+
+impl<A, S, F, Ortho> Arnoldi<A, S, F, Ortho>
+where
+    A: Scalar + Lapack,
+    S: DataMut<Elem = A>,
+    F: Fn(&mut ArrayBase<S, Ix1>),
+    Ortho: Orthogonalizer<Elem = A>,
+{
+    /// Create an Arnoldi iterator from any linear operator `a`
+    pub fn new(a: F, mut v: ArrayBase<S, Ix1>, mut ortho: Ortho) -> Self {
+        assert_eq!(ortho.len(), 0);
+        assert!(ortho.tolerance() < One::one());
+        // normalize before append because |v| may be smaller than ortho.tolerance()
+        let norm = v.norm_l2();
+        azip!(mut v(&mut v) in { *v = v.div_real(norm) });
+        ortho.append(v.view());
+        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(mut self) -> (Q<A>, H<A>) {
+        for _ in &mut self {} // execute iteration until convergent
+        let q = self.ortho.get_q();
+        let n = self.h.len();
+        let mut h = Array2::zeros((n, n).f());
+        for (i, hc) in self.h.iter().enumerate() {
+            let m = std::cmp::min(n, i + 2);
+            for j in 0..m {
+                h[(j, i)] = hc[j];
+            }
+        }
+        (q, h)
+    }
+}
+
+impl<A, S, F, Ortho> Iterator for Arnoldi<A, S, F, Ortho>
+where
+    A: Scalar + Lapack,
+    S: DataMut<Elem = A>,
+    F: Fn(&mut ArrayBase<S, Ix1>),
+    Ortho: Orthogonalizer<Elem = A>,
+{
+    type Item = Array1<A>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        (self.a)(&mut self.v);
+        let result = self.ortho.div_append(&mut self.v);
+        let norm = self.v.norm_l2();
+        azip!(mut v(&mut self.v) in { *v = v.div_real(norm) });
+        match result {
+            AppendResult::Added(coef) => {
+                self.h.push(coef.clone());
+                Some(coef)
+            }
+            AppendResult::Dependent(coef) => {
+                self.h.push(coef);
+                None
+            }
+        }
+    }
+}
+
+/// 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,
+    S1: Data<Elem = A>,
+    S2: DataMut<Elem = A>,
+{
+    let (n, m) = a.dim();
+    assert_eq!(n, m, "Input matrix must be square");
+    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 });
+    }
+}
+
+/// Utility to execute Arnoldi iteration with Householder reflection
+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 householder = Householder::new(v.len(), tol);
+    Arnoldi::new(mul_mat(a), v, householder).complete()
+}
+
+/// Utility to execute Arnoldi iteration with modified Gram-Schmit orthogonalizer
+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()
+}

src/krylov/householder.rs

@@ -8,17 +8,17 @@ use crate::{inner::*, norm::*};
 use num_traits::One;
 
 /// Calc a reflactor `w` from a vector `x`
-pub fn calc_reflector<A, S>(x: &mut ArrayBase<S, Ix1>) -> A
+pub fn calc_reflector<A, S>(x: &mut ArrayBase<S, Ix1>)
 where
     A: Scalar + Lapack,
     S: DataMut<Elem = A>,
 {
+    assert!(x.len() > 0);
     let norm = x.norm_l2();
     let alpha = -x[0].mul_real(norm / x[0].abs());
     x[0] -= alpha;
     let inv_rev_norm = A::Real::one() / x.norm_l2();
     azip!(mut a(x) in { *a = a.mul_real(inv_rev_norm)});
-    alpha
 }
 
 /// Take a reflection `P = I - 2ww^T`
@@ -50,12 +50,19 @@ pub struct Householder<A: Scalar> {
     ///
     /// The coefficient is copied into another array, and this does not contain
     v: Vec<Array1<A>>,
+
+    /// Tolerance
+    tol: A::Real,
 }
 
 impl<A: Scalar + Lapack> Householder<A> {
     /// Create a new orthogonalizer
-    pub fn new(dim: usize) -> Self {
-        Householder { dim, v: Vec::new() }
+    pub fn new(dim: usize, tol: A::Real) -> Self {
+        Householder {
+            dim,
+            v: Vec::new(),
+            tol,
+        }
     }
 
     /// Take a Reflection `P = I - 2ww^T`
@@ -92,12 +99,32 @@ impl<A: Scalar + Lapack> Householder<A> {
         }
     }
 
-    fn eval_residual<S>(&self, a: &ArrayBase<S, Ix1>) -> A::Real
+    /// Compose coefficients array using reflected vector
+    fn compose_coefficients<S>(&self, a: &ArrayBase<S, Ix1>) -> Coefficients<A>
     where
         S: Data<Elem = A>,
     {
-        let l = self.v.len();
-        a.slice(s![l..]).norm_l2()
+        let k = self.len();
+        let res = a.slice(s![k..]).norm_l2();
+        let mut c = Array1::zeros(k + 1);
+        azip!(mut c(c.slice_mut(s![..k])), a(a.slice(s![..k])) in { *c = a });
+        if k < a.len() {
+            let ak = a[k];
+            c[k] = -ak.mul_real(res / ak.abs());
+        } else {
+            c[k] = A::from_real(res);
+        }
+        c
+    }
+
+    /// Construct the residual vector from reflected vector
+    fn construct_residual<S>(&self, a: &mut ArrayBase<S, Ix1>)
+    where
+        S: DataMut<Elem = A>,
+    {
+        let k = self.len();
+        azip!(mut a( a.slice_mut(s![..k])) in { *a = A::zero() });
+        self.backward_reflection(a);
     }
 }
 
@@ -112,45 +139,61 @@ impl<A: Scalar + Lapack> Orthogonalizer for Householder<A> {
         self.v.len()
     }
 
+    fn tolerance(&self) -> A::Real {
+        self.tol
+    }
+
+    fn decompose<S>(&self, a: &mut ArrayBase<S, Ix1>) -> Array1<A>
+    where
+        S: DataMut<Elem = A>,
+    {
+        self.forward_reflection(a);
+        let coef = self.compose_coefficients(a);
+        self.construct_residual(a);
+        coef
+    }
+
     fn coeff<S>(&self, a: ArrayBase<S, Ix1>) -> Array1<A>
     where
         S: Data<Elem = A>,
     {
         let mut a = a.into_owned();
         self.forward_reflection(&mut a);
-        let res = self.eval_residual(&a);
-        let k = self.len();
-        let mut c = Array1::zeros(k + 1);
-        azip!(mut c(c.slice_mut(s![..k])), a(a.slice(s![..k])) in { *c = a });
-        c[k] = A::from_real(res);
-        c
+        self.compose_coefficients(&a)
     }
 
-    fn append<S>(&mut self, mut a: ArrayBase<S, Ix1>, rtol: A::Real) -> Result<Array1<A>, Array1<A>>
+    fn div_append<S>(&mut self, a: &mut ArrayBase<S, Ix1>) -> AppendResult<A>
     where
         S: DataMut<Elem = A>,
     {
         assert_eq!(a.len(), self.dim);
         let k = self.len();
-
-        self.forward_reflection(&mut a);
-        let mut coef = Array::zeros(k + 1);
-        for i in 0..k {
-            coef[i] = a[i];
-        }
-        if self.is_full() {
-            return Err(coef); // coef[k] must be zero in this case
+        self.forward_reflection(a);
+        let coef = self.compose_coefficients(a);
+        if coef[k].abs() < self.tol {
+            return AppendResult::Dependent(coef);
         }
+        calc_reflector(&mut a.slice_mut(s![k..]));
+        self.v.push(a.to_owned());
+        self.construct_residual(a);
+        AppendResult::Added(coef)
+    }
 
-        let alpha = calc_reflector(&mut a.slice_mut(s![k..]));
-        coef[k] = alpha;
-
-        if alpha.abs() < rtol {
-            // linearly dependent
-            return Err(coef);
+    fn append<S>(&mut self, a: ArrayBase<S, Ix1>) -> AppendResult<A>
+    where
+        S: Data<Elem = A>,
+    {
+        assert_eq!(a.len(), self.dim);
+        let mut a = a.into_owned();
+        let k = self.len();
+        self.forward_reflection(&mut a);
+        let coef = self.compose_coefficients(&a);
+        if coef[k].abs() < self.tol {
+            return AppendResult::Dependent(coef);
         }
-        self.v.push(a.into_owned());
-        Ok(coef)
+        calc_reflector(&mut a.slice_mut(s![k..]));
+        self.v.push(a.to_owned());
+        AppendResult::Added(coef)
     }
 
     fn get_q(&self) -> Q<A> {
@@ -175,8 +218,8 @@ where
     A: Scalar + Lapack,
     S: Data<Elem = A>,
 {
-    let h = Householder::new(dim);
-    qr(iter, h, rtol, strategy)
+    let h = Householder::new(dim, rtol);
+    qr(iter, h, strategy)
 }
 
 #[cfg(test)]