Repository: mars-project/mars

Issue:
[TENSOR] QR decomposition for Short-and-Fat matrix
**Is your feature request related to a problem? Please describe.**
`mt.linalg.qr` now only works for tall and skinny matrix, short and fat matrix is easy to implement in theory, something should be metioned is that 'sfqr' method cannot be used for `mt.linalg.svd`.

**Describe the solution you'd like**
Short-and-Fat matrix A can be split like [A1, A2, A3, ...], we can choose A1 to be decomposed as a product A1 = Q1R1 of an orthogonal matrix Q1 and an upper triangular matrix R1, matrix Q1 as A's orthogonal matrix, then compute R ([R1, R2, R3, ...]) matrix by matrix Q1 to make sure A2=Q1R2, A3=Q1R3...



Discussion:

Resolving patch:
diff --git a/mars/tensor/expressions/linalg/core.py b/mars/tensor/expressions/linalg/core.py
--- a/mars/tensor/expressions/linalg/core.py
+++ b/mars/tensor/expressions/linalg/core.py
@@ -18,15 +18,71 @@
 
 from ....compat import izip
 from ..utils import decide_chunk_sizes
-from ..core import TensorOperandMixin
 
 
-class TSQR(TensorOperandMixin):
+class SFQR(object):
     __slots__ = ()
 
     @classmethod
-    def _is_svd(cls):
-        return False
+    def tile(cls, op):
+        """
+        Short-and-Fat QR
+
+        Q [R_1 R_2 ...] = [A_1 A_2 ...]
+        """
+        from .qr import TensorQR
+        from .dot import TensorDot
+        from ..base import TensorTranspose
+
+        a = op.input
+
+        tinyq, tinyr = np.linalg.qr(np.ones((1, 1), dtype=a.dtype))
+        q_dtype, r_dtype = tinyq.dtype, tinyr.dtype
+
+        rechunk_size = dict()
+        if a.chunk_shape[0] != 1:
+            rechunk_size[0] = a.shape[0]
+
+        if len(a.chunks) > 1 and a.chunks[0].shape[0] > a.chunks[0].shape[1]:
+            rechunk_size[1] = a.shape[0]
+
+        if rechunk_size:
+            new_chunks = decide_chunk_sizes(a.shape, rechunk_size, a.dtype.itemsize)
+            a = a.rechunk(new_chunks).single_tiles()
+
+        # A_1's QR decomposition
+        r_chunks = []
+        first_chunk = a.chunks[0]
+        x, y = first_chunk.shape
+        q_shape = first_chunk.shape if x > y else (x, x)
+        r_shape = first_chunk.shape if x < y else (y, y)
+        qr_op = TensorQR()
+        q_chunk, r_chunk = qr_op.new_chunks([first_chunk], (q_shape, r_shape),
+                                            index=(0, 0),
+                                            kws=[{'side': 'q', 'dtype': q_dtype},
+                                                 {'side': 'r', 'dtype': r_dtype}])
+        # q is an orthogonal matrix, so q.T and inverse of q is equal
+        trans_op = TensorTranspose()
+        q_transpose = trans_op.new_chunk([q_chunk], q_chunk.shape)
+        r_chunks.append(r_chunk)
+
+        r_rest = [TensorDot().new_chunk([q_transpose, c], (q_transpose.shape[0], c.shape[1]),
+                                        index=c.index) for c in a.chunks[1:]]
+        r_chunks.extend(r_rest)
+
+        q, r = op.outputs
+        new_op = op.copy()
+        q_nsplits = ((q_chunk.shape[0],), (q_chunk.shape[1],))
+        r_nsplits = ((1,), (c.shape[1] for c in r_chunks))
+        kws = [
+            {'chunks': [q_chunk], 'nsplits': q_nsplits, 'dtype': q.dtype},
+            {'chunks': r_chunks, 'nsplits': r_nsplits, 'dtype': r.dtype}
+        ]
+        return new_op.new_tensors(op.inputs, [q.shape, r.shape], kws=kws)
+
+
+class TSQR(object):
+    __slots__ = ()
 
     @classmethod
     def tile(cls, op):
@@ -36,7 +92,7 @@ def tile(cls, op):
         from .qr import TensorQR
         from .svd import TensorSVD
 
-        calc_svd = cls._is_svd()
+        calc_svd = getattr(op, '_is_svd', lambda: None)() or False
 
         a = op.input
 
@@ -90,6 +146,7 @@ def tile(cls, op):
         if not calc_svd:
             q, r = op.outputs
             new_op = op.copy()
+            # unify_nsplits will get nsplits by chunk shape if it was set to (1,)
             q_nsplits = ((c.shape[0] for c in stage3_q_chunks), (1,))
             r_nsplits = ((stage2_r_chunk.shape[0],), (stage2_r_chunk.shape[1],))
             kws = [
diff --git a/mars/tensor/expressions/linalg/qr.py b/mars/tensor/expressions/linalg/qr.py
--- a/mars/tensor/expressions/linalg/qr.py
+++ b/mars/tensor/expressions/linalg/qr.py
@@ -20,10 +20,11 @@
 from .... import operands
 from ...core import ExecutableTuple
 from ..datasource import tensor as astensor
-from .core import TSQR
+from ..core import TensorOperandMixin
+from .core import SFQR, TSQR
 
 
-class TensorQR(operands.QR, TSQR):
+class TensorQR(operands.QR, TensorOperandMixin):
     def __init__(self, method=None, dtype=None, **kw):
         super(TensorQR, self).__init__(_method=method, _dtype=dtype, **kw)
 
@@ -68,8 +69,9 @@ def tile(cls, op):
             ]
             return new_op.new_tensors(op.inputs, [q_shape, r_shape], kws=kws)
         elif op.method == 'tsqr':
-            return super(TensorQR, cls).tile(op)
-        # TODO(hks): support sfqr(short-and-fat qr)
+            return TSQR.tile(op)
+        elif op.method == 'sfqr':
+            return SFQR.tile(op)
         else:
             raise NotImplementedError('Only tsqr method supported for now')
 
@@ -85,7 +87,7 @@ def qr(a, method='tsqr'):
     ----------
     a : array_like, shape (M, N)
         Matrix to be factored.
-    method: {'tsqr'}, optional
+    method: {'tsqr', 'sfqr'}, optional
         method to calculate qr factorization, tsqr as default
 
         TSQR is presented in:
@@ -96,6 +98,10 @@ def qr(a, method='tsqr'):
             IEEE International Conference on Big Data, 2013.
             http://arxiv.org/abs/1301.1071
 
+        FSQR is a QR decomposition for fat and short matrix:
+            A = [A1, A2, A3, ...], A1 may be decomposed as A1 = Q1 * R1,
+            for A = Q * R, Q = Q1, R = [R1, R2, R3, ...] where A2 = Q1 * R2, A3 = Q1 * R3, ...
+
     Returns
     -------
     q : Tensor of float or complex, optional
diff --git a/mars/tensor/expressions/linalg/svd.py b/mars/tensor/expressions/linalg/svd.py
--- a/mars/tensor/expressions/linalg/svd.py
+++ b/mars/tensor/expressions/linalg/svd.py
@@ -20,10 +20,11 @@
 from .... import operands
 from ...core import ExecutableTuple
 from ..datasource import tensor as astensor
+from ..core import TensorOperandMixin
 from .core import TSQR
 
 
-class TensorSVD(operands.SVD, TSQR):
+class TensorSVD(operands.SVD, TensorOperandMixin):
     def __init__(self, method=None, dtype=None, **kw):
         super(TensorSVD, self).__init__(_method=method, _dtype=dtype, **kw)
 
@@ -44,9 +45,17 @@ def __call__(self, a):
 
         tiny_U, tiny_s, tiny_V = np.linalg.svd(np.ones((1, 1), dtype=a.dtype))
 
-        U_shape = a.shape
-        s_shape = (a.shape[1],)
-        V_shape = (a.shape[1],) * 2
+        # if a's shape is (6, 18), U's shape is (6, 6), s's shape is (6,), V's shape is (6, 18)
+        # if a's shape is (18, 6), U's shape is (18, 6), s's shape is (6,), V's shape is (6, 6)
+        x, y = a.shape
+        if x > y:
+            U_shape = (x, y)
+            s_shape = (y, )
+            V_shape = (y, y)
+        else:
+            U_shape = (x, x)
+            s_shape = (x, )
+            V_shape = (x, y)
         U, s, V = self.new_tensors([a], (U_shape, s_shape, V_shape),
                                    kws=[
                                        {'side': 'U', 'dtype': tiny_U.dtype},
@@ -83,7 +92,7 @@ def tile(cls, op):
             ]
             return new_op.new_tensors(op.inputs, [U_shape, s_shape, V_shape], kws=kws)
         elif op.method == 'tsqr':
-            return super(TensorSVD, cls).tile(op)
+            return TSQR.tile(op)
         else:
             raise NotImplementedError('Only tsqr method supported for now')
 


Test patch:
diff --git a/mars/tensor/execution/tests/test_linalg_execute.py b/mars/tensor/execution/tests/test_linalg_execute.py
--- a/mars/tensor/execution/tests/test_linalg_execute.py
+++ b/mars/tensor/execution/tests/test_linalg_execute.py
@@ -55,6 +55,30 @@ def testQRExecution(self):
         res = self.executor.execute_tensor(t, concat=True)[0]
         self.assertTrue(np.allclose(res, data))
 
+        # test for Short-and-Fat QR
+        data = np.random.randn(6, 18)
+
+        a = tensor(data, chunk_size=(6, 9))
+        q, r = qr(a, method='sfqr')
+        t = q.dot(r)
+
+        res = self.executor.execute_tensor(t, concat=True)[0]
+        self.assertTrue(np.allclose(res, data))
+
+        a = tensor(data, chunk_size=(3, 3))
+        q, r = qr(a, method='sfqr')
+        t = q.dot(r)
+
+        res = self.executor.execute_tensor(t, concat=True)[0]
+        self.assertTrue(np.allclose(res, data))
+
+        a = tensor(data, chunk_size=(6, 3))
+        q, r = qr(a, method='sfqr')
+        t = q.dot(r)
+
+        res = self.executor.execute_tensor(t, concat=True)[0]
+        self.assertTrue(np.allclose(res, data))
+
     def testSVDExecution(self):
         data = np.random.randn(18, 6) + 1j * np.random.randn(18, 6)
 
@@ -79,6 +103,15 @@ def testSVDExecution(self):
         res = self.executor.execute_tensor(t, concat=True)[0]
         self.assertTrue(np.allclose(res, data))
 
+        data = np.random.randn(6, 18) + 1j * np.random.randn(6, 18)
+
+        a = tensor(data)
+        U, s, V = svd(a)
+        t = U.dot(diag(s).dot(V))
+
+        res = self.executor.execute_tensor(t, concat=True)[0]
+        self.assertTrue(np.allclose(res, data))
+
     def testCholeskyExecution(self):
         data = np.array([[1, -2j], [2j, 5]])
 
diff --git a/mars/tensor/expressions/tests/test_linalg.py b/mars/tensor/expressions/tests/test_linalg.py
--- a/mars/tensor/expressions/tests/test_linalg.py
+++ b/mars/tensor/expressions/tests/test_linalg.py
@@ -35,6 +35,41 @@ def testQR(self):
         self.assertEqual(len(q.chunks), 3)
         self.assertEqual(len(r.chunks), 1)
 
+        # for Short-and-Fat QR
+        a = mt.random.rand(6, 18, chunk_size=(6, 6))
+        q, r = mt.linalg.qr(a, method='sfqr')
+
+        self.assertEqual(q.shape, (6, 6))
+        self.assertEqual(r.shape, (6, 18))
+
+        q.tiles()
+
+        self.assertEqual(len(q.chunks), 1)
+        self.assertEqual(len(r.chunks), 3)
+
+        # chunk width less than height
+        a = mt.random.rand(6, 9, chunk_size=(6, 3))
+        q, r = mt.linalg.qr(a, method='sfqr')
+
+        self.assertEqual(q.shape, (6, 6))
+        self.assertEqual(r.shape, (6, 9))
+
+        q.tiles()
+
+        self.assertEqual(len(q.chunks), 1)
+        self.assertEqual(len(r.chunks), 2)
+
+        a = mt.random.rand(9, 6, chunk_size=(9, 3))
+        q, r = mt.linalg.qr(a, method='sfqr')
+
+        self.assertEqual(q.shape, (9, 6))
+        self.assertEqual(r.shape, (6, 6))
+
+        q.tiles()
+
+        self.assertEqual(len(q.chunks), 1)
+        self.assertEqual(len(r.chunks), 1)
+
     def testNorm(self):
         data = np.random.rand(9, 6)
 
@@ -81,6 +116,13 @@ def testSVD(self):
         self.assertEqual(s.ndim, 1)
         self.assertEqual(len(s.chunks[0].index), 1)
 
+        a = mt.random.rand(6, 9, chunk_size=(6, 9))
+        U, s, V = mt.linalg.svd(a)
+
+        self.assertEqual(U.shape, (6, 6))
+        self.assertEqual(s.shape, (6,))
+        self.assertEqual(V.shape, (6, 9))
+
         rs = mt.random.RandomState(1)
         a = rs.rand(9, 6, chunk_size=(3, 6))
         U, s, V = mt.linalg.svd(a)