Merge pull request #13 from termoshtt/qr

QR decomposition

Author: termoshtt

src/hermite.rs

@@ -8,6 +8,7 @@ use matrix::Matrix;
 use square::SquareMatrix;
 use error::LinalgError;
 use eigh::ImplEigh;
+use qr::ImplQR;
 use svd::ImplSVD;
 use norm::ImplNorm;
 use solve::ImplSolve;
@@ -21,7 +22,7 @@ pub trait HermiteMatrix: SquareMatrix + Matrix {
 }
 
 impl<A> HermiteMatrix for Array<A, (Ix, Ix)>
-    where A: ImplSVD + ImplNorm + ImplSolve + ImplEigh + LinalgScalar + Float
+    where A: ImplQR + ImplSVD + ImplNorm + ImplSolve + ImplEigh + LinalgScalar + Float
 {
     fn eigh(self) -> Result<(Self::Vector, Self), LinalgError> {
         try!(self.check_square());

src/lib.rs

@@ -9,6 +9,7 @@ pub mod matrix;
 pub mod square;
 pub mod hermite;
 
+pub mod qr;
 pub mod svd;
 pub mod eigh;
 pub mod norm;

src/matrix.rs

@@ -1,9 +1,11 @@
 //! Define trait for general matrix
 
+use std::cmp::min;
 use ndarray::prelude::*;
 use ndarray::LinalgScalar;
 
 use error::LapackError;
+use qr::ImplQR;
 use svd::ImplSVD;
 use norm::ImplNorm;
 
@@ -21,11 +23,11 @@ pub trait Matrix: Sized {
     fn norm_f(&self) -> Self::Scalar;
     /// singular-value decomposition (SVD)
     fn svd(self) -> Result<(Self, Self::Vector, Self), LapackError>;
-    // fn qr(self) -> (Self, Self);
+    fn qr(self) -> Result<(Self, Self), LapackError>;
 }
 
 impl<A> Matrix for Array<A, (Ix, Ix)>
-    where A: ImplSVD + ImplNorm + LinalgScalar
+    where A: ImplQR + ImplSVD + ImplNorm + LinalgScalar
 {
     type Scalar = A;
     type Vector = Array<A, Ix>;
@@ -74,4 +76,39 @@ impl<A> Matrix for Array<A, (Ix, Ix)>
             Ok((ua, sv, va))
         }
     }
+    fn qr(self) -> Result<(Self, Self), LapackError> {
+        let (n, m) = self.size();
+        let strides = self.strides();
+        let k = min(n, m);
+        let (q, r) = if strides[0] < strides[1] {
+            try!(ImplQR::qr(m, n, self.clone().into_raw_vec()))
+        } else {
+            try!(ImplQR::lq(n, m, self.clone().into_raw_vec()))
+        };
+        let (qa, ra) = if strides[0] < strides[1] {
+            (Array::from_vec(q).into_shape((m, n)).unwrap().reversed_axes(),
+             Array::from_vec(r).into_shape((m, n)).unwrap().reversed_axes())
+        } else {
+            (Array::from_vec(q).into_shape((n, m)).unwrap(),
+             Array::from_vec(r).into_shape((n, m)).unwrap())
+        };
+        let qm = if m > k {
+            let (qsl, _) = qa.view().split_at(Axis(1), k);
+            qsl.to_owned()
+        } else {
+            qa
+        };
+        let mut rm = if n > k {
+            let (rsl, _) = ra.view().split_at(Axis(0), k);
+            rsl.to_owned()
+        } else {
+            ra
+        };
+        for ((i, j), val) in rm.indexed_iter_mut() {
+            if i > j {
+                *val = A::zero();
+            }
+        }
+        Ok((qm, rm))
+    }
 }

src/qr.rs

@@ -0,0 +1,113 @@
+//! Implement QR decomposition
+
+extern crate lapack;
+
+use std::cmp::min;
+use self::lapack::fortran::*;
+use num_traits::Zero;
+
+use error::LapackError;
+
+pub trait ImplQR: Sized {
+    fn qr(n: usize, m: usize, mut a: Vec<Self>) -> Result<(Vec<Self>, Vec<Self>), LapackError>;
+    fn lq(n: usize, m: usize, mut a: Vec<Self>) -> Result<(Vec<Self>, Vec<Self>), LapackError>;
+}
+
+macro_rules! impl_qr {
+    ($geqrf:path, $orgqr:path, $gelqf:path, $orglq:path) => {
+// XXX These codes are most same, but the argument of $orgqr and $orglq are different!
+fn qr(n: usize, m: usize, mut a: Vec<Self>) -> Result<(Vec<Self>, Vec<Self>), LapackError> {
+    let n = n as i32;
+    let m = m as i32;
+    let mut info = 0;
+    let k = min(m, n);
+    let lda = m;
+    let lw_default = 1000;
+    let mut tau = vec![Self::zero(); k as usize];
+    let mut work = vec![Self::zero(); lw_default];
+// estimate lwork
+    $geqrf(m, n, &mut a, lda, &mut tau, &mut work, -1, &mut info);
+    let lwork_r = work[0] as i32;
+    if lwork_r > lw_default as i32 {
+        work = vec![Self::zero(); lwork_r as usize];
+    }
+// calc R
+    $geqrf(m, n, &mut a, lda, &mut tau, &mut work, lwork_r, &mut info);
+    if info != 0 {
+        return Err(From::from(info));
+    }
+    let r = a.clone();
+// re-estimate lwork
+    $orgqr(m, k, k, &mut a, lda, &mut tau, &mut work, -1, &mut info);
+    let lwork_q = work[0] as i32;
+    if lwork_q > lwork_r {
+        work = vec![Self::zero(); lwork_q as usize];
+    }
+// calc Q
+    $orgqr(m,
+           k,
+           k,
+           &mut a,
+           lda,
+           &mut tau,
+           &mut work,
+           lwork_q,
+           &mut info);
+    if info == 0 {
+        Ok((a, r))
+    } else {
+        Err(From::from(info))
+    }
+}
+fn lq(n: usize, m: usize, mut a: Vec<Self>) -> Result<(Vec<Self>, Vec<Self>), LapackError> {
+    let n = n as i32;
+    let m = m as i32;
+    let mut info = 0;
+    let k = min(m, n);
+    let lda = m;
+    let lw_default = 1000;
+    let mut tau = vec![Self::zero(); k as usize];
+    let mut work = vec![Self::zero(); lw_default];
+// estimate lwork
+    $gelqf(m, n, &mut a, lda, &mut tau, &mut work, -1, &mut info);
+    let lwork_r = work[0] as i32;
+    if lwork_r > lw_default as i32 {
+        work = vec![Self::zero(); lwork_r as usize];
+    }
+// calc R
+    $gelqf(m, n, &mut a, lda, &mut tau, &mut work, lwork_r, &mut info);
+    if info != 0 {
+        return Err(From::from(info));
+    }
+    let r = a.clone();
+// re-estimate lwork
+    $orglq(k, n, k, &mut a, lda, &mut tau, &mut work, -1, &mut info);
+    let lwork_q = work[0] as i32;
+    if lwork_q > lwork_r {
+        work = vec![Self::zero(); lwork_q as usize];
+    }
+// calc Q
+    $orglq(k,
+           n,
+           k,
+           &mut a,
+           lda,
+           &mut tau,
+           &mut work,
+           lwork_q,
+           &mut info);
+    if info == 0 {
+        Ok((a, r))
+    } else {
+        Err(From::from(info))
+    }
+}
+}} // endmacro
+
+impl ImplQR for f64 {
+    impl_qr!(dgeqrf, dorgqr, dgelqf, dorglq);
+}
+
+impl ImplQR for f32 {
+    impl_qr!(sgeqrf, sorgqr, sgelqf, sorglq);
+}

src/square.rs

@@ -6,6 +6,7 @@ use num_traits::float::Float;
 
 use matrix::Matrix;
 use error::{LinalgError, NotSquareError};
+use qr::ImplQR;
 use svd::ImplSVD;
 use norm::ImplNorm;
 use solve::ImplSolve;
@@ -37,7 +38,7 @@ pub trait SquareMatrix: Matrix {
 }
 
 impl<A> SquareMatrix for Array<A, (Ix, Ix)>
-    where A: ImplSVD + ImplNorm + ImplSolve + LinalgScalar + Float
+    where A: ImplQR + ImplNorm + ImplSVD + ImplSolve + LinalgScalar + Float
 {
     fn inv(self) -> Result<Self, LinalgError> {
         try!(self.check_square());