Merge pull request #150 from rust-ndarray/krylov

Orthogonalizer trait

Author: termoshtt

src/krylov/mgs.rs

@@ -0,0 +1,108 @@
+//! Modified Gram-Schmit orthogonalizer
+
+use super::*;
+use crate::{generate::*, inner::*, norm::Norm};
+
+/// Iterative orthogonalizer using modified Gram-Schmit procedure
+///
+/// ```rust
+/// # use ndarray::*;
+/// # use ndarray_linalg::{krylov::*, *};
+/// let mut mgs = MGS::new(3);
+/// let coef = mgs.append(array![0.0, 1.0, 0.0], 1e-9).unwrap();
+/// close_l2(&coef, &array![1.0], 1e-9);
+///
+/// let coef = mgs.append(array![1.0, 1.0, 0.0], 1e-9).unwrap();
+/// close_l2(&coef, &array![1.0, 1.0], 1e-9);
+///
+/// // Fail if the vector is linearly dependent
+/// assert!(mgs.append(array![1.0, 2.0, 0.0], 1e-9).is_err());
+///
+/// // You can get coefficients of dependent vector
+/// if let Err(coef) = mgs.append(array![1.0, 2.0, 0.0], 1e-9) {
+///     close_l2(&coef, &array![2.0, 1.0, 0.0], 1e-9);
+/// }
+/// ```
+#[derive(Debug, Clone)]
+pub struct MGS<A> {
+    /// Dimension of base space
+    dimension: usize,
+    /// Basis of spanned space
+    q: Vec<Array1<A>>,
+}
+
+impl<A: Scalar> MGS<A> {
+    /// Create an empty orthogonalizer
+    pub fn new(dimension: usize) -> Self {
+        Self {
+            dimension,
+            q: Vec::new(),
+        }
+    }
+}
+
+impl<A: Scalar + Lapack> Orthogonalizer for MGS<A> {
+    type Elem = A;
+
+    fn dim(&self) -> usize {
+        self.dimension
+    }
+
+    fn len(&self) -> usize {
+        self.q.len()
+    }
+
+    fn orthogonalize<S>(&self, a: &mut ArrayBase<S, Ix1>) -> Array1<A>
+    where
+        A: Lapack,
+        S: DataMut<Elem = A>,
+    {
+        assert_eq!(a.len(), self.dim());
+        let mut coef = Array1::zeros(self.len() + 1);
+        for i in 0..self.len() {
+            let q = &self.q[i];
+            let c = q.inner(&a);
+            azip!(mut a (&mut *a), q (q) in { *a = *a - c * q } );
+            coef[i] = c;
+        }
+        let nrm = a.norm_l2();
+        coef[self.len()] = A::from_real(nrm);
+        coef
+    }
+
+    fn append<S>(&mut self, a: ArrayBase<S, Ix1>, rtol: A::Real) -> Result<Array1<A>, Array1<A>>
+    where
+        A: Lapack,
+        S: Data<Elem = A>,
+    {
+        let mut a = a.into_owned();
+        let coef = self.orthogonalize(&mut a);
+        let nrm = coef[coef.len() - 1].re();
+        if nrm < rtol {
+            // Linearly dependent
+            return Err(coef);
+        }
+        azip!(mut a in { *a = *a / A::from_real(nrm) });
+        self.q.push(a);
+        Ok(coef)
+    }
+
+    fn get_q(&self) -> Q<A> {
+        hstack(&self.q).unwrap()
+    }
+}
+
+/// Online QR decomposition of vectors using modified Gram-Schmit algorithm
+pub fn mgs<A, S>(
+    iter: impl Iterator<Item = ArrayBase<S, Ix1>>,
+    dim: usize,
+    rtol: A::Real,
+    strategy: Strategy,
+) -> (Q<A>, R<A>)
+where
+    A: Scalar + Lapack,
+    S: Data<Elem = A>,
+{
+    let mgs = MGS::new(dim);
+    qr(iter, mgs, rtol, strategy)
+}

src/krylov/mod.rs

@@ -0,0 +1,134 @@
+//! Krylov subspace
+
+use crate::types::*;
+use ndarray::*;
+
+mod mgs;
+
+pub use mgs::{mgs, MGS};
+
+/// Q-matrix
+///
+/// - Maybe **NOT** square
+/// - Unitary for existing columns
+///
+pub type Q<A> = Array2<A>;
+
+/// R-matrix
+///
+/// - Maybe **NOT** square
+/// - Upper triangle
+///
+pub type R<A> = Array2<A>;
+
+/// Trait for creating orthogonal basis from iterator of arrays
+pub trait Orthogonalizer {
+    type Elem: Scalar;
+
+    /// Dimension of input array
+    fn dim(&self) -> usize;
+
+    /// Number of cached basis
+    fn len(&self) -> usize;
+
+    /// check if the basis spans entire space
+    fn is_full(&self) -> bool {
+        self.len() == self.dim()
+    }
+
+    fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    /// Orthogonalize given vector using current basis
+    ///
+    /// Panic
+    /// -------
+    /// - if the size of the input array mismatches to the dimension
+    ///
+    fn orthogonalize<S>(&self, a: &mut ArrayBase<S, Ix1>) -> Array1<Self::Elem>
+    where
+        S: DataMut<Elem = Self::Elem>;
+
+    /// Add new vector if the residual is larger than relative tolerance
+    ///
+    /// Returns
+    /// --------
+    /// Coefficients to the `i`-th Q-vector
+    ///
+    /// - The size of array must be `self.len() + 1`
+    /// - The last element is the residual norm of input vector
+    ///
+    /// Panic
+    /// -------
+    /// - if the size of the input array mismatches to the dimension
+    ///
+    fn append<S>(
+        &mut self,
+        a: ArrayBase<S, Ix1>,
+        rtol: <Self::Elem as Scalar>::Real,
+    ) -> Result<Array1<Self::Elem>, Array1<Self::Elem>>
+    where
+        S: DataMut<Elem = Self::Elem>;
+
+    /// Get Q-matrix of generated basis
+    fn get_q(&self) -> Q<Self::Elem>;
+}
+
+/// Strategy for linearly dependent vectors appearing in iterative QR decomposition
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum Strategy {
+    /// Terminate iteration if dependent vector comes
+    Terminate,
+
+    /// Skip dependent vector
+    Skip,
+
+    /// Orthogonalize dependent vector without adding to Q,
+    /// i.e. R must be non-square like following:
+    ///
+    /// ```text
+    /// x x x x x
+    /// 0 x x x x
+    /// 0 0 0 x x
+    /// 0 0 0 0 x
+    /// ```
+    Full,
+}
+
+/// Online QR decomposition using arbitrary orthogonalizer
+pub fn qr<A, S>(
+    iter: impl Iterator<Item = ArrayBase<S, Ix1>>,
+    mut ortho: impl Orthogonalizer<Elem = A>,
+    rtol: A::Real,
+    strategy: Strategy,
+) -> (Q<A>, R<A>)
+where
+    A: Scalar + Lapack,
+    S: Data<Elem = A>,
+{
+    assert_eq!(ortho.len(), 0);
+
+    let mut coefs = Vec::new();
+    for a in iter {
+        match ortho.append(a.into_owned(), rtol) {
+            Ok(coef) => coefs.push(coef),
+            Err(coef) => match strategy {
+                Strategy::Terminate => break,
+                Strategy::Skip => continue,
+                Strategy::Full => coefs.push(coef),
+            },
+        }
+    }
+    let n = ortho.len();
+    let m = coefs.len();
+    let mut r = Array2::zeros((n, m).f());
+    for j in 0..m {
+        for i in 0..n {
+            if i < coefs[j].len() {
+                r[(i, j)] = coefs[j][i];
+            }
+        }
+    }
+    (ortho.get_q(), r)
+}

src/lib.rs

@@ -44,9 +44,9 @@ pub mod eigh;
 pub mod error;
 pub mod generate;
 pub mod inner;
+pub mod krylov;
 pub mod lapack;
 pub mod layout;
-pub mod mgs;
 pub mod norm;
 pub mod operator;
 pub mod opnorm;