BUG: linalg: lstsq should always return real residuals (#937)

(cherry picked from commit 99f79e3)

numpy/linalg/linalg.py

@@ -1680,12 +1680,13 @@ def lstsq(a, b, rcond=-1):
     """
     Return the least-squares solution to a linear matrix equation.
 
-    Solves the equation `a x = b` by computing a vector `x` that minimizes
-    the norm `|| b - a x ||`.  The equation may be under-, well-, or over-
-    determined (i.e., the number of linearly independent rows of `a` can be
-    less than, equal to, or greater than its number of linearly independent
-    columns).  If `a` is square and of full rank, then `x` (but for round-off
-    error) is the "exact" solution of the equation.
+    Solves the equation `a x = b` by computing a vector `x` that
+    minimizes the Euclidean 2-norm `|| b - a x ||^2`.  The equation may
+    be under-, well-, or over- determined (i.e., the number of
+    linearly independent rows of `a` can be less than, equal to, or
+    greater than its number of linearly independent columns).  If `a`
+    is square and of full rank, then `x` (but for round-off error) is
+    the "exact" solution of the equation.
 
     Parameters
     ----------
@@ -1706,7 +1707,7 @@ def lstsq(a, b, rcond=-1):
         Least-squares solution.  The shape of `x` depends on the shape of
         `b`.
     residues : ndarray, shape (), (1,), or (K,)
-        Sums of residues; squared Euclidean norm for each column in
+        Sums of residues; squared Euclidean 2-norm for each column in
         ``b - a*x``.
         If the rank of `a` is < N or > M, this is an empty array.
         If `b` is 1-dimensional, this is a (1,) shape array.
@@ -1772,6 +1773,7 @@ def lstsq(a, b, rcond=-1):
     if m != b.shape[0]:
         raise LinAlgError, 'Incompatible dimensions'
     t, result_t = _commonType(a, b)
+    result_real_t = _realType(result_t)
     real_t = _linalgRealType(t)
     bstar = zeros((ldb, n_rhs), t)
     bstar[:b.shape[0],:n_rhs] = b.copy()
@@ -1811,16 +1813,27 @@ def lstsq(a, b, rcond=-1):
                                  0, work, lwork, iwork, 0)
     if results['info'] > 0:
         raise LinAlgError, 'SVD did not converge in Linear Least Squares'
-    resids = array([], t)
+    resids = array([], result_real_t)
     if is_1d:
         x = array(ravel(bstar)[:n], dtype=result_t, copy=True)
         if results['rank'] == n and m > n:
-            resids = array([sum((ravel(bstar)[n:])**2)], dtype=result_t)
+            if isComplexType(t):
+                resids = array([sum(abs(ravel(bstar)[n:])**2)],
+                               dtype=result_real_t)
+            else:
+                resids = array([sum((ravel(bstar)[n:])**2)],
+                               dtype=result_real_t)
     else:
         x = array(transpose(bstar)[:n,:], dtype=result_t, copy=True)
         if results['rank'] == n and m > n:
-            resids = sum((transpose(bstar)[n:,:])**2, axis=0).astype(result_t)
-    st = s[:min(n, m)].copy().astype(_realType(result_t))
+            if isComplexType(t):
+                resids = sum(abs(transpose(bstar)[n:,:])**2, axis=0).astype(
+                    result_real_t)
+            else:
+                resids = sum((transpose(bstar)[n:,:])**2, axis=0).astype(
+                    result_real_t)
+
+    st = s[:min(n, m)].copy().astype(result_real_t)
     return wrap(x), wrap(resids), results['rank'], st
 
 def norm(x, ord=None):

numpy/linalg/tests/test_linalg.py

@@ -33,6 +33,11 @@ def test_double(self):
         b = array([2., 1.], dtype=double)
         self.do(a, b)
 
+    def test_double_2(self):
+        a = array([[1.,2.], [3.,4.]], dtype=double)
+        b = array([[2., 1., 4.], [3., 4., 6.]], dtype=double)
+        self.do(a, b)
+
     def test_csingle(self):
         a = array([[1.+2j,2+3j], [3+4j,4+5j]], dtype=csingle)
         b = array([2.+1j, 1.+2j], dtype=csingle)
@@ -43,6 +48,11 @@ def test_cdouble(self):
         b = array([2.+1j, 1.+2j], dtype=cdouble)
         self.do(a, b)
 
+    def test_cdouble_2(self):
+        a = array([[1.+2j,2+3j], [3+4j,4+5j]], dtype=cdouble)
+        b = array([[2.+1j, 1.+2j, 1+3j], [1-2j, 1-3j, 1-6j]], dtype=cdouble)
+        self.do(a, b)
+
     def test_empty(self):
         a = atleast_2d(array([], dtype = double))
         b = atleast_2d(array([], dtype = double))
@@ -70,6 +80,58 @@ def test_matrix_a_and_b(self):
         self.do(a, b)
 
 
+class LinalgNonsquareTestCase:
+    def test_single_nsq_1(self):
+        a = array([[1.,2.,3.], [3.,4.,6.]], dtype=single)
+        b = array([2., 1.], dtype=single)
+        self.do(a, b)
+
+    def test_single_nsq_2(self):
+        a = array([[1.,2.], [3.,4.], [5.,6.]], dtype=single)
+        b = array([2., 1., 3.], dtype=single)
+        self.do(a, b)
+
+    def test_double_nsq_1(self):
+        a = array([[1.,2.,3.], [3.,4.,6.]], dtype=double)
+        b = array([2., 1.], dtype=double)
+        self.do(a, b)
+
+    def test_double_nsq_2(self):
+        a = array([[1.,2.], [3.,4.], [5.,6.]], dtype=double)
+        b = array([2., 1., 3.], dtype=double)
+        self.do(a, b)
+
+    def test_csingle_nsq_1(self):
+        a = array([[1.+1j,2.+2j,3.-3j], [3.-5j,4.+9j,6.+2j]], dtype=csingle)
+        b = array([2.+1j, 1.+2j], dtype=csingle)
+        self.do(a, b)
+
+    def test_csingle_nsq_2(self):
+        a = array([[1.+1j,2.+2j], [3.-3j,4.-9j], [5.-4j,6.+8j]], dtype=csingle)
+        b = array([2.+1j, 1.+2j, 3.-3j], dtype=csingle)
+        self.do(a, b)
+
+    def test_cdouble_nsq_1(self):
+        a = array([[1.+1j,2.+2j,3.-3j], [3.-5j,4.+9j,6.+2j]], dtype=cdouble)
+        b = array([2.+1j, 1.+2j], dtype=cdouble)
+        self.do(a, b)
+
+    def test_cdouble_nsq_2(self):
+        a = array([[1.+1j,2.+2j], [3.-3j,4.-9j], [5.-4j,6.+8j]], dtype=cdouble)
+        b = array([2.+1j, 1.+2j, 3.-3j], dtype=cdouble)
+        self.do(a, b)
+
+    def test_cdouble_nsq_1_2(self):
+        a = array([[1.+1j,2.+2j,3.-3j], [3.-5j,4.+9j,6.+2j]], dtype=cdouble)
+        b = array([[2.+1j, 1.+2j], [1-1j, 2-2j]], dtype=cdouble)
+        self.do(a, b)
+
+    def test_cdouble_nsq_2_2(self):
+        a = array([[1.+1j,2.+2j], [3.-3j,4.-9j], [5.-4j,6.+8j]], dtype=cdouble)
+        b = array([[2.+1j, 1.+2j], [1-1j, 2-2j], [1-1j, 2-2j]], dtype=cdouble)
+        self.do(a, b)
+
+
 class TestSolve(LinalgTestCase, TestCase):
     def do(self, a, b):
         x = linalg.solve(a, b)
@@ -153,13 +215,28 @@ def test_zero(self):
         assert_equal(type(linalg.slogdet([[0.0j]])[0]), cdouble)
         assert_equal(type(linalg.slogdet([[0.0j]])[1]), double)
 
-class TestLstsq(LinalgTestCase, TestCase):
+class TestLstsq(LinalgTestCase, LinalgNonsquareTestCase, TestCase):
     def do(self, a, b):
+        arr = np.asarray(a)
+        m, n = arr.shape
         u, s, vt = linalg.svd(a, 0)
         x, residuals, rank, sv = linalg.lstsq(a, b)
-        assert_almost_equal(b, dot(a, x))
-        assert_equal(rank, asarray(a).shape[0])
+        if m <= n:
+            assert_almost_equal(b, dot(a, x))
+            assert_equal(rank, m)
+        else:
+            assert_equal(rank, n)
         assert_almost_equal(sv, sv.__array_wrap__(s))
+        if rank == n and m > n:
+            expect_resids = (np.asarray(abs(np.dot(a, x) - b))**2).sum(axis=0)
+            expect_resids = np.asarray(expect_resids)
+            if len(np.asarray(b).shape) == 1:
+                expect_resids.shape = (1,)
+                assert_equal(residuals.shape, expect_resids.shape)
+        else:
+            expect_resids = type(x)([])
+        assert_almost_equal(residuals, expect_resids)
+        assert_(np.issubdtype(residuals.dtype, np.floating))
         assert imply(isinstance(b, matrix), isinstance(x, matrix))
         assert imply(isinstance(b, matrix), isinstance(residuals, matrix))