BUG: gh-2790, fix column scaling in polynomial package least squares.

The columns should be scaled using their 2-norm, but in the complex case
that was being incorrectly computed as the square root of the sum of the
squared elements rather than as the square root of the  sum of their squared
real and imaginary parts.

6 files changed, 48 insertions, 12 deletions

diff --git a/numpy/polynomial/chebyshev.py b/numpy/polynomial/chebyshev.py
index 00cac2527..4d87c8470 100644
--- a/numpy/polynomial/chebyshev.py
+++ b/numpy/polynomial/chebyshev.py
@@ -1742,8 +1742,14 @@ def chebfit(x, y, deg, rcond=None, full=False, w=None):
     if rcond is None :
         rcond = len(x)*np.finfo(x.dtype).eps
 
-    # scale the design matrix and solve the least squares equation
-    scl = np.sqrt((lhs*lhs).sum(1))
+    # Determine the norms of the design matrix columns.
+    if lhs.dtype.char in np.typecodes['Complex']:
+        scl = np.sqrt((np.square(lhs.real) + np.square(lhs.imag)).sum(1))
+    else:
+        scl = np.sqrt(np.square(lhs).sum(1))
+    scl[scl == 0] = 1
+
+    # Solve the least squares problem.
     c, resids, rank, s = la.lstsq(lhs.T/scl, rhs.T, rcond)
     c = (c.T/scl).T
 
diff --git a/numpy/polynomial/hermite.py b/numpy/polynomial/hermite.py
index 0b637f40a..66940a28e 100644
--- a/numpy/polynomial/hermite.py
+++ b/numpy/polynomial/hermite.py
@@ -1519,8 +1519,14 @@ def hermfit(x, y, deg, rcond=None, full=False, w=None):
     if rcond is None :
         rcond = len(x)*np.finfo(x.dtype).eps
 
-    # scale the design matrix and solve the least squares equation
-    scl = np.sqrt((lhs*lhs).sum(1))
+    # Determine the norms of the design matrix columns.
+    if lhs.dtype.char in np.typecodes['Complex']:
+        scl = np.sqrt((np.square(lhs.real) + np.square(lhs.imag)).sum(1))
+    else:
+        scl = np.sqrt(np.square(lhs).sum(1))
+    scl[scl == 0] = 1
+
+    # Solve the least squares problem.
     c, resids, rank, s = la.lstsq(lhs.T/scl, rhs.T, rcond)
     c = (c.T/scl).T
 
diff --git a/numpy/polynomial/hermite_e.py b/numpy/polynomial/hermite_e.py
index c5abe03ca..9ff7c06d2 100644
--- a/numpy/polynomial/hermite_e.py
+++ b/numpy/polynomial/hermite_e.py
@@ -1515,8 +1515,14 @@ def hermefit(x, y, deg, rcond=None, full=False, w=None):
     if rcond is None :
         rcond = len(x)*np.finfo(x.dtype).eps
 
-    # scale the design matrix and solve the least squares equation
-    scl = np.sqrt((lhs*lhs).sum(1))
+    # Determine the norms of the design matrix columns.
+    if lhs.dtype.char in np.typecodes['Complex']:
+        scl = np.sqrt((np.square(lhs.real) + np.square(lhs.imag)).sum(1))
+    else:
+        scl = np.sqrt(np.square(lhs).sum(1))
+    scl[scl == 0] = 1
+
+    # Solve the least squares problem.
     c, resids, rank, s = la.lstsq(lhs.T/scl, rhs.T, rcond)
     c = (c.T/scl).T
 
diff --git a/numpy/polynomial/laguerre.py b/numpy/polynomial/laguerre.py
index 3533343b0..6e035ebd2 100644
--- a/numpy/polynomial/laguerre.py
+++ b/numpy/polynomial/laguerre.py
@@ -1518,8 +1518,14 @@ def lagfit(x, y, deg, rcond=None, full=False, w=None):
     if rcond is None :
         rcond = len(x)*np.finfo(x.dtype).eps
 
-    # scale the design matrix and solve the least squares equation
-    scl = np.sqrt((lhs*lhs).sum(1))
+    # Determine the norms of the design matrix columns.
+    if lhs.dtype.char in np.typecodes['Complex']:
+        scl = np.sqrt((np.square(lhs.real) + np.square(lhs.imag)).sum(1))
+    else:
+        scl = np.sqrt(np.square(lhs).sum(1))
+    scl[scl == 0] = 1
+
+    # Solve the least squares problem.
     c, resids, rank, s = la.lstsq(lhs.T/scl, rhs.T, rcond)
     c = (c.T/scl).T
 
diff --git a/numpy/polynomial/legendre.py b/numpy/polynomial/legendre.py
index 1216e29f4..e909ad9b1 100644
--- a/numpy/polynomial/legendre.py
+++ b/numpy/polynomial/legendre.py
@@ -1543,8 +1543,14 @@ def legfit(x, y, deg, rcond=None, full=False, w=None):
     if rcond is None :
         rcond = len(x)*np.finfo(x.dtype).eps
 
-    # scale the design matrix and solve the least squares equation
-    scl = np.sqrt((lhs*lhs).sum(1))
+    # Determine the norms of the design matrix columns.
+    if lhs.dtype.char in np.typecodes['Complex']:
+        scl = np.sqrt((np.square(lhs.real) + np.square(lhs.imag)).sum(1))
+    else:
+        scl = np.sqrt(np.square(lhs).sum(1))
+    scl[scl == 0] = 1
+
+    # Solve the least squares problem.
     c, resids, rank, s = la.lstsq(lhs.T/scl, rhs.T, rcond)
     c = (c.T/scl).T
 
diff --git a/numpy/polynomial/polynomial.py b/numpy/polynomial/polynomial.py
index 324bec9c0..4e24e9c1f 100644
--- a/numpy/polynomial/polynomial.py
+++ b/numpy/polynomial/polynomial.py
@@ -1365,8 +1365,14 @@ def polyfit(x, y, deg, rcond=None, full=False, w=None):
     if rcond is None :
         rcond = len(x)*np.finfo(x.dtype).eps
 
-    # scale the design matrix and solve the least squares equation
-    scl = np.sqrt((lhs*lhs).sum(1))
+    # Determine the norms of the design matrix columns.
+    if lhs.dtype.char in np.typecodes['Complex']:
+        scl = np.sqrt((np.square(lhs.real) + np.square(lhs.imag)).sum(1))
+    else:
+        scl = np.sqrt(np.square(lhs).sum(1))
+    scl[scl == 0] = 1
+
+    # Solve the least squares problem.
     c, resids, rank, s = la.lstsq(lhs.T/scl, rhs.T, rcond)
     c = (c.T/scl).T