Change the QR implementation, inspired by LAPACK

src/linalg/qr.rs

@@ -3,15 +3,12 @@ use num::Zero;
 use serde::{Deserialize, Serialize};
 
 use crate::allocator::{Allocator, Reallocator};
-use crate::base::{DefaultAllocator, Matrix, MatrixMN, MatrixN, Unit, VectorN};
+use crate::base::{DefaultAllocator, Matrix, MatrixMN, MatrixN, VectorN};
 use crate::constraint::{SameNumberOfRows, ShapeConstraint};
 use crate::dimension::{Dim, DimMin, DimMinimum, U1};
 use crate::storage::{Storage, StorageMut};
 use simba::scalar::ComplexField;
 
-use crate::geometry::Reflection;
-use crate::linalg::householder;
-
 /// The QR decomposition of a general matrix.
 #[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
 #[cfg_attr(
@@ -34,7 +31,7 @@ where
     DefaultAllocator: Allocator<N, R, C> + Allocator<N, DimMinimum<R, C>>,
 {
     qr: MatrixMN<N, R, C>,
-    diag: VectorN<N, DimMinimum<R, C>>,
+    tau: VectorN<N, DimMinimum<R, C>>,
 }
 
 impl<N: ComplexField, R: DimMin<C>, C: Dim> Copy for QR<N, R, C>
@@ -47,30 +44,64 @@ where
 
 impl<N: ComplexField, R: DimMin<C>, C: Dim> QR<N, R, C>
 where
-    DefaultAllocator: Allocator<N, R, C> + Allocator<N, R> + Allocator<N, DimMinimum<R, C>>,
+    DefaultAllocator:
+        Allocator<N, R, C> + Allocator<N, R> + Allocator<N, C> + Allocator<N, DimMinimum<R, C>>,
 {
-    /// Computes the QR decomposition using householder reflections.
+    /// Computes the QR decomposition using Householder reflections.
     pub fn new(mut matrix: MatrixMN<N, R, C>) -> Self {
         let (nrows, ncols) = matrix.data.shape();
         let min_nrows_ncols = nrows.min(ncols);
 
-        let mut diag = unsafe { MatrixMN::new_uninitialized_generic(min_nrows_ncols, U1) };
+        let mut tau = unsafe { MatrixMN::new_uninitialized_generic(min_nrows_ncols, U1) };
 
         if min_nrows_ncols.value() == 0 {
-            return QR {
+            return QR { qr: matrix, tau };
-                qr: matrix,
+        }
-                diag: diag,
+
+        let mut work = unsafe { MatrixMN::new_uninitialized_generic(ncols, U1) };
+        for icol in 0..min_nrows_ncols.value() {
+            let (mut left, mut right) = matrix.columns_range_pair_mut(icol, icol + 1..);
+            let mut axis = left.rows_range_mut(icol..);
+            // Compute the scaled Housholder vector
+            let (beta, tau_i) = {
+                let alpha = unsafe { *axis.vget_unchecked(0) };
+                let xnorm = axis.rows_range(1..).norm();
+                if xnorm.is_zero() && alpha.imaginary().is_zero() {
+                    (alpha, N::zero())
+                } else {
+                    let a_r = alpha.real();
+                    let a_i = alpha.imaginary();
+                    // FIXME: Use LAPACK's ?LAPY3 once we have F::max
+                    let reflection_norm = (a_r * a_r + a_i * a_i + xnorm * xnorm).sqrt();
+                    // FIXME: Use reflection_norm.copysign(a_r)
+                    let beta = -reflection_norm.abs() * a_r.signum();
+                    // FIXME: Check for tiny beta and recompute, cf. LAPACK's ?LARFG
+                    let tau = (N::from_real(beta) - alpha).unscale(beta);
+                    let tmp = alpha - N::from_real(beta);
+                    axis.rows_range_mut(1..).apply(|x| x / tmp);
+                    (N::from_real(beta), tau)
+                }
             };
+            unsafe {
+                *tau.vget_unchecked_mut(icol) = tau_i;
+            }
+            if !tau_i.is_zero() {
+                // apply the Householder reflection to the remaining columns
+                unsafe {
+                    *axis.vget_unchecked_mut(0) = N::one();
+                }
+                let mut work = work.rows_range_mut(icol + 1..);
+                work.gemv_ad(N::one(), &right.rows_range(icol..), &axis, N::zero());
+                right
+                    .rows_range_mut(icol..)
+                    .gerc(-tau_i.conjugate(), &axis, &work, N::one());
+            }
+            unsafe {
+                *axis.vget_unchecked_mut(0) = beta;
+            }
         }
 
-        for ite in 0..min_nrows_ncols.value() {
+        QR { qr: matrix, tau }
-            householder::clear_column_unchecked(&mut matrix, &mut diag[ite], ite, 0, None);
-        }
-
-        QR {
-            qr: matrix,
-            diag: diag,
-        }
     }
 
     /// Retrieves the upper trapezoidal submatrix `R` of this decomposition.
@@ -80,9 +111,7 @@ where
         DefaultAllocator: Allocator<N, DimMinimum<R, C>, C>,
     {
         let (nrows, ncols) = self.qr.data.shape();
-        let mut res = self.qr.rows_generic(0, nrows.min(ncols)).upper_triangle();
+        self.qr.rows_generic(0, nrows.min(ncols)).upper_triangle()
-        res.set_partial_diagonal(self.diag.iter().map(|e| N::from_real(e.modulus())));
-        res
     }
 
     /// Retrieves the upper trapezoidal submatrix `R` of this decomposition.
@@ -96,32 +125,59 @@ where
         let (nrows, ncols) = self.qr.data.shape();
         let mut res = self.qr.resize_generic(nrows.min(ncols), ncols, N::zero());
         res.fill_lower_triangle(N::zero(), 1);
-        res.set_partial_diagonal(self.diag.iter().map(|e| N::from_real(e.modulus())));
         res
     }
 
+    /// Compute the first `ncols` columns of `Q`.
+    ///
+    /// Panics if `ncols` is bigger than the number of rows of the original matrix.
+    pub fn q_columns<K: Dim>(&self, ncols: K) -> MatrixMN<N, R, K>
+    where
+        DefaultAllocator: Allocator<N, R, K> + Allocator<N, K>,
+    {
+        let (q_nrows, q_ncols) = self.qr.data.shape();
+        assert!(
+            ncols.value() <= q_nrows.value(),
+            "k must be less than the number of columns"
+        );
+        let mut a = MatrixMN::<N, R, K>::identity_generic(q_nrows, ncols);
+        let mut work = unsafe { VectorN::<N, K>::new_uninitialized_generic(ncols, U1) };
+        let k = q_nrows.min(q_ncols);
+        let ncols_to_copy = k.value().min(q_ncols.value());
+        a.slice_range_mut(.., ..ncols_to_copy)
+            .copy_from(&self.qr.slice_range(.., ..ncols_to_copy));
+        for i in (0..k.value()).rev() {
+            let tau_i = unsafe { *self.tau.vget_unchecked(i) };
+            if i < q_ncols.value() {
+                unsafe {
+                    *a.get_unchecked_mut((i, i)) = N::one();
+                }
+                let (left, mut right) = a.columns_range_pair_mut(i, i + 1..);
+                let axis = left.rows_range(i..);
+                let mut work = work.rows_range_mut(i + 1..);
+                work.gemv_ad(N::one(), &right.rows_range(i..), &axis, N::zero());
+                right
+                    .rows_range_mut(i..)
+                    .gerc(-tau_i, &axis, &work, N::one());
+            }
+            if i < q_nrows.value() {
+                a.slice_range_mut(i + 1.., i).apply(|x| x * (-tau_i));
+            }
+            unsafe {
+                *a.get_unchecked_mut((i, i)) = N::one() - tau_i;
+            }
+            a.slice_range_mut(0..i, i).fill(N::zero());
+        }
+        a
+    }
+
     /// Computes the orthogonal matrix `Q` of this decomposition.
     pub fn q(&self) -> MatrixMN<N, R, DimMinimum<R, C>>
     where
-        DefaultAllocator: Allocator<N, R, DimMinimum<R, C>>,
+        DefaultAllocator: Allocator<N, R, DimMinimum<R, C>> + Allocator<N, DimMinimum<R, C>>,
     {
         let (nrows, ncols) = self.qr.data.shape();
-
+        self.q_columns::<DimMinimum<R, C>>(nrows.min(ncols))
-        // NOTE: we could build the identity matrix and call q_mul on it.
-        // Instead we don't so that we take in account the matrix sparseness.
-        let mut res = Matrix::identity_generic(nrows, nrows.min(ncols));
-        let dim = self.diag.len();
-
-        for i in (0..dim).rev() {
-            let axis = self.qr.slice_range(i.., i);
-            // FIXME: sometimes, the axis might have a zero magnitude.
-            let refl = Reflection::new(Unit::new_unchecked(axis), N::zero());
-
-            let mut res_rows = res.slice_range_mut(i.., i..);
-            refl.reflect_with_sign(&mut res_rows, self.diag[i].signum());
-        }
-
-        res
     }
 
     /// Unpacks this decomposition into its two matrix factors.
@@ -133,8 +189,9 @@ where
     )
     where
         DimMinimum<R, C>: DimMin<C, Output = DimMinimum<R, C>>,
-        DefaultAllocator:
+        DefaultAllocator: Allocator<N, R, DimMinimum<R, C>>
-            Allocator<N, R, DimMinimum<R, C>> + Reallocator<N, R, C, DimMinimum<R, C>, C>,
+            + Allocator<N, DimMinimum<R, C>>
+            + Reallocator<N, R, C, DimMinimum<R, C>, C>,
     {
         (self.q(), self.unpack_r())
     }
@@ -145,19 +202,27 @@ where
     }
 
     /// Multiplies the provided matrix by the transpose of the `Q` matrix of this decomposition.
-    pub fn q_tr_mul<R2: Dim, C2: Dim, S2>(&self, rhs: &mut Matrix<N, R2, C2, S2>)
+    pub fn q_tr_mul<R2: Dim, C2: Dim, S2>(&mut self, rhs: &mut Matrix<N, R2, C2, S2>)
-    // FIXME: do we need a static constraint on the number of rows of rhs?
     where
         S2: StorageMut<N, R2, C2>,
+        ShapeConstraint: SameNumberOfRows<R2, R>,
+        DefaultAllocator: Allocator<N, C2>,
     {
-        let dim = self.diag.len();
+        let ncols = rhs.data.shape().1;
-
+        let mut work = unsafe { VectorN::<N, C2>::new_uninitialized_generic(ncols, U1) };
-        for i in 0..dim {
+        for i in 0..self.tau.len() {
-            let axis = self.qr.slice_range(i.., i);
+            let mut axis = self.qr.slice_range_mut(i.., i);
-            let refl = Reflection::new(Unit::new_unchecked(axis), N::zero());
+            let temp = unsafe { *axis.vget_unchecked(0) };
-
+            unsafe {
-            let mut rhs_rows = rhs.rows_range_mut(i..);
+                *axis.vget_unchecked_mut(0) = N::one();
-            refl.reflect_with_sign(&mut rhs_rows, self.diag[i].signum().conjugate());
+            }
+            let tau_i = unsafe { *self.tau.vget_unchecked(i) };
+            work.gemv_ad(N::one(), &rhs.rows_range(i..), &axis, N::zero());
+            rhs.rows_range_mut(i..)
+                .gerc(-tau_i.conjugate(), &axis, &work, N::one());
+            unsafe {
+                *axis.vget_unchecked_mut(0) = temp;
+            }
         }
     }
 }
@@ -170,13 +235,13 @@ where
     ///
     /// Returns `None` if `self` is not invertible.
     pub fn solve<R2: Dim, C2: Dim, S2>(
-        &self,
+        &mut self,
         b: &Matrix<N, R2, C2, S2>,
     ) -> Option<MatrixMN<N, R2, C2>>
     where
         S2: Storage<N, R2, C2>,
         ShapeConstraint: SameNumberOfRows<R2, D>,
-        DefaultAllocator: Allocator<N, R2, C2>,
+        DefaultAllocator: Allocator<N, R2, C2> + Allocator<N, C2>,
     {
         let mut res = b.clone_owned();
 
@@ -191,9 +256,10 @@ where
     ///
     /// If the decomposed matrix is not invertible, this returns `false` and its input `b` is
     /// overwritten with garbage.
-    pub fn solve_mut<R2: Dim, C2: Dim, S2>(&self, b: &mut Matrix<N, R2, C2, S2>) -> bool
+    pub fn solve_mut<R2: Dim, C2: Dim, S2>(&mut self, b: &mut Matrix<N, R2, C2, S2>) -> bool
     where
         S2: StorageMut<N, R2, C2>,
+        DefaultAllocator: Allocator<N, C2>,
         ShapeConstraint: SameNumberOfRows<R2, D>,
     {
         assert_eq!(
@@ -207,48 +273,13 @@ where
         );
 
         self.q_tr_mul(b);
-        self.solve_upper_triangular_mut(b)
+        self.qr.solve_upper_triangular_mut(b)
-    }
-
-    // FIXME: duplicate code from the `solve` module.
-    fn solve_upper_triangular_mut<R2: Dim, C2: Dim, S2>(
-        &self,
-        b: &mut Matrix<N, R2, C2, S2>,
-    ) -> bool
-    where
-        S2: StorageMut<N, R2, C2>,
-        ShapeConstraint: SameNumberOfRows<R2, D>,
-    {
-        let dim = self.qr.nrows();
-
-        for k in 0..b.ncols() {
-            let mut b = b.column_mut(k);
-            for i in (0..dim).rev() {
-                let coeff;
-
-                unsafe {
-                    let diag = self.diag.vget_unchecked(i).modulus();
-
-                    if diag.is_zero() {
-                        return false;
-                    }
-
-                    coeff = b.vget_unchecked(i).unscale(diag);
-                    *b.vget_unchecked_mut(i) = coeff;
-                }
-
-                b.rows_range_mut(..i)
-                    .axpy(-coeff, &self.qr.slice_range(..i, i), N::one());
-            }
-        }
-
-        true
     }
 
     /// Computes the inverse of the decomposed matrix.
     ///
     /// Returns `None` if the decomposed matrix is not invertible.
-    pub fn try_inverse(&self) -> Option<MatrixN<N, D>> {
+    pub fn try_inverse(&mut self) -> Option<MatrixN<N, D>> {
         assert!(
             self.qr.is_square(),
             "QR inverse: unable to compute the inverse of a non-square matrix."
@@ -271,33 +302,14 @@ where
             self.qr.is_square(),
             "QR: unable to test the invertibility of a non-square matrix."
         );
-
+        (0..self.qr.ncols()).all(|i| unsafe { !self.qr.get_unchecked((i, i)).is_zero() })
-        for i in 0..self.diag.len() {
-            if self.diag[i].is_zero() {
-                return false;
     }
 }
 
-        true
-    }
-
-    // /// Computes the determinant of the decomposed matrix.
-    // pub fn determinant(&self) -> N {
-    //     let dim = self.qr.nrows();
-    //     assert!(self.qr.is_square(), "QR determinant: unable to compute the determinant of a non-square matrix.");
-
-    //     let mut res = N::one();
-    //     for i in 0 .. dim {
-    //         res *= unsafe { *self.diag.vget_unchecked(i) };
-    //     }
-
-    //     res self.q_determinant()
-    // }
-}
-
 impl<N: ComplexField, R: DimMin<C>, C: Dim, S: Storage<N, R, C>> Matrix<N, R, C, S>
 where
-    DefaultAllocator: Allocator<N, R, C> + Allocator<N, R> + Allocator<N, DimMinimum<R, C>>,
+    DefaultAllocator:
+        Allocator<N, R, C> + Allocator<N, R> + Allocator<N, C> + Allocator<N, DimMinimum<R, C>>,
 {
     /// Computes the QR decomposition of this matrix.
     pub fn qr(self) -> QR<N, R, C> {

tests/linalg/qr.rs

@@ -1,6 +1,48 @@
-#![cfg(feature = "arbitrary")]
+use na::{Matrix2, Matrix4x2, U3, U4};
 
+#[test]
+fn simple_qr() {
+    #[rustfmt::skip]
+    let a = Matrix4x2::new(
+        -0.8943285241224914 ,  0.12787800716234649,
+        -0.37320804072796987,  0.21338804264385058,
+         0.                 , -0.2456767687354977 ,
+         0.2456767687354977 ,  0.                 );
+    let qr = a.qr();
+    // the reference values were generated by converting the input
+    // to the form `m * 2 ^ e` for integers m and e. This was then used to
+    // obtain the QR decomposition without rounding errors. The result was
+    // converted back to f64.
+    #[rustfmt::skip]
+    let r_ref = Matrix2::new(
+        0.99973237689865724, -0.19405501632841561,
+        0.                 , -0.2908383860381578);
+    assert_relative_eq!(qr.r(), r_ref);
 
+    #[rustfmt::skip]
+    let q_ref = Matrix4x2::new(
+        -0.89456793116659196,  0.15719172406996297, 
+        -0.3733079465583837 , -0.48461884587835711,
+         0.                 ,  0.8447191998351451, 
+         0.24574253511487697, -0.1639658791740342);
+    assert_relative_eq!(qr.q(), q_ref);
+}
+
+#[test]
+fn q_columns() {
+    let a = Matrix4x2::new(0., 1., 3., 3., 1., 1., 2., 1.);
+    let qr = a.qr();
+    assert!(qr.q_columns(U4).is_orthogonal(1.0e-15));
+}
+
+#[test]
+#[should_panic]
+fn q_columns_panic() {
+    Matrix2::<f64>::zeros().qr().q_columns(U3);
+}
+
+#[cfg(feature = "arbitrary")]
+mod quickcheck_test {
     macro_rules! gen_tests(
         ($module: ident, $scalar: ty) => {
             mod $module {
@@ -16,11 +58,8 @@ macro_rules! gen_tests(
                         let q  = qr.q();
                         let r  = qr.r();
 
-                    println!("m: {}", m);
+                        relative_eq!(m, &q * r, epsilon = 1.0e-9) &&
-                    println!("qr: {}", &q * &r);
+                        q.is_orthogonal(1.0e-15)
-
-                    relative_eq!(m, &q * r, epsilon = 1.0e-7) &&
-                    q.is_orthogonal(1.0e-7)
                     }
 
                     fn qr_static_5_3(m: Matrix5x3<$scalar>) -> bool {
@@ -29,8 +68,8 @@ macro_rules! gen_tests(
                         let q  = qr.q();
                         let r  = qr.r();
 
-                    relative_eq!(m, q * r, epsilon = 1.0e-7) &&
+                        relative_eq!(m, q * r, epsilon = 1.0e-8) &&
-                    q.is_orthogonal(1.0e-7)
+                        q.is_orthogonal(1.0e-15)
                     }
 
                     fn qr_static_3_5(m: Matrix3x5<$scalar>) -> bool {
@@ -39,8 +78,8 @@ macro_rules! gen_tests(
                         let q  = qr.q();
                         let r  = qr.r();
 
-                    relative_eq!(m, q * r, epsilon = 1.0e-7) &&
+                        relative_eq!(m, q * r, epsilon = 1.0e-9) &&
-                    q.is_orthogonal(1.0e-7)
+                        q.is_orthogonal(1.0e-15)
                     }
 
                     fn qr_static_square(m: Matrix4<$scalar>) -> bool {
@@ -49,10 +88,8 @@ macro_rules! gen_tests(
                         let q  = qr.q();
                         let r  = qr.r();
 
-                    println!("{}{}{}{}", q, r, q * r, m);
+                        relative_eq!(m, q * r, epsilon = 1.0e-9) &&
-
+                        q.is_orthogonal(1.0e-15)
-                    relative_eq!(m, q * r, epsilon = 1.0e-7) &&
-                    q.is_orthogonal(1.0e-7)
                     }
 
                     fn qr_solve(n: usize, nb: usize) -> bool {
@@ -61,7 +98,7 @@ macro_rules! gen_tests(
                             let nb = cmp::min(nb, 50); // To avoid slowing down the test too much.
                             let m  = DMatrix::<$scalar>::new_random(n, n).map(|e| e.0);
 
-                        let qr = m.clone().qr();
+                            let mut qr = m.clone().qr();
                             let b1 = DVector::<$scalar>::new_random(n).map(|e| e.0);
                             let b2 = DMatrix::<$scalar>::new_random(n, nb).map(|e| e.0);
 
@@ -69,8 +106,8 @@ macro_rules! gen_tests(
                                 let sol1 = qr.solve(&b1).unwrap();
                                 let sol2 = qr.solve(&b2).unwrap();
 
-                            return relative_eq!(&m * sol1, b1, epsilon = 1.0e-6) &&
+                                return relative_eq!(&m * sol1, b1, epsilon = 1.0e-8) &&
-                                relative_eq!(&m * sol2, b2, epsilon = 1.0e-6)
+                                    relative_eq!(&m * sol2, b2, epsilon = 1.0e-8)
                             }
                         }
 
@@ -79,7 +116,7 @@ macro_rules! gen_tests(
 
                     fn qr_solve_static(m: Matrix4<$scalar>) -> bool {
                         let m = m.map(|e| e.0);
-                     let qr = m.qr();
+                        let mut qr = m.qr();
                         let b1 = Vector4::<$scalar>::new_random().map(|e| e.0);
                         let b2 = Matrix4x3::<$scalar>::new_random().map(|e| e.0);
 
@@ -87,8 +124,8 @@ macro_rules! gen_tests(
                             let sol1 = qr.solve(&b1).unwrap();
                             let sol2 = qr.solve(&b2).unwrap();
 
-                         relative_eq!(m * sol1, b1, epsilon = 1.0e-6) &&
+                            relative_eq!(m * sol1, b1, epsilon = 1.0e-8) &&
-                         relative_eq!(m * sol2, b2, epsilon = 1.0e-6)
+                            relative_eq!(m * sol2, b2, epsilon = 1.0e-8)
                         }
                         else {
                             false
@@ -103,7 +140,7 @@ macro_rules! gen_tests(
                             let id1 = &m  * &m1;
                             let id2 = &m1 * &m;
 
-                        id1.is_identity(1.0e-5) && id2.is_identity(1.0e-5)
+                            id1.is_identity(1.0e-7) && id2.is_identity(1.0e-7)
                         }
                         else {
                             true
@@ -112,13 +149,13 @@ macro_rules! gen_tests(
 
                     fn qr_inverse_static(m: Matrix4<$scalar>) -> bool {
                         let m  = m.map(|e| e.0);
-                    let qr = m.qr();
+                        let mut qr = m.qr();
 
                         if let Some(m1) = qr.try_inverse() {
                             let id1 = &m  * &m1;
                             let id2 = &m1 * &m;
 
-                        id1.is_identity(1.0e-5) && id2.is_identity(1.0e-5)
+                            id1.is_identity(1.0e-7) && id2.is_identity(1.0e-7)
                         }
                         else {
                             true
@@ -131,3 +168,4 @@ macro_rules! gen_tests(
 
     gen_tests!(complex, RandComplex<f64>);
     gen_tests!(f64, RandScalar<f64>);
+}