Merge pull request #2947 from charris/fix-complex-polynomial-fit

Fix complex polynomial fit

12 files changed, 73 insertions, 12 deletions

diff --git a/numpy/polynomial/chebyshev.py b/numpy/polynomial/chebyshev.py
index 00cac2527..afeafcb68 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 issubclass(lhs.dtype.type, np.complexfloating):
+        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..2beb848ae 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 issubclass(lhs.dtype.type, np.complexfloating):
+        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..c183a5a86 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 issubclass(lhs.dtype.type, np.complexfloating):
+        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..f0c1268bf 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 issubclass(lhs.dtype.type, np.complexfloating):
+        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..0efd13ffa 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 issubclass(lhs.dtype.type, np.complexfloating):
+        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..8d7251b19 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 issubclass(lhs.dtype.type, np.complexfloating):
+        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/tests/test_chebyshev.py b/numpy/polynomial/tests/test_chebyshev.py
index 331cc17bc..95da83631 100644
--- a/numpy/polynomial/tests/test_chebyshev.py
+++ b/numpy/polynomial/tests/test_chebyshev.py
@@ -435,6 +435,10 @@ class TestFitting(TestCase):
         #
         wcoef2d = cheb.chebfit(x, np.array([yw,yw]).T, 3, w=w)
         assert_almost_equal(wcoef2d, np.array([coef3,coef3]).T)
+        # test scaling with complex values x points whose square
+        # is zero when summed.
+        x = [1, 1j, -1, -1j]
+        assert_almost_equal(cheb.chebfit(x, x, 1), [0, 1])
 
 class TestGauss(TestCase):
 
diff --git a/numpy/polynomial/tests/test_hermite.py b/numpy/polynomial/tests/test_hermite.py
index b5649a693..a0ef3b515 100644
--- a/numpy/polynomial/tests/test_hermite.py
+++ b/numpy/polynomial/tests/test_hermite.py
@@ -425,6 +425,10 @@ class TestFitting(TestCase):
         #
         wcoef2d = herm.hermfit(x, np.array([yw,yw]).T, 3, w=w)
         assert_almost_equal(wcoef2d, np.array([coef3,coef3]).T)
+        # test scaling with complex values x points whose square
+        # is zero when summed.
+        x = [1, 1j, -1, -1j]
+        assert_almost_equal(herm.hermfit(x, x, 1), [0, .5])
 
 class TestGauss(TestCase):
 
diff --git a/numpy/polynomial/tests/test_hermite_e.py b/numpy/polynomial/tests/test_hermite_e.py
index 018fe8595..f6bfe5e5e 100644
--- a/numpy/polynomial/tests/test_hermite_e.py
+++ b/numpy/polynomial/tests/test_hermite_e.py
@@ -422,6 +422,10 @@ class TestFitting(TestCase):
         #
         wcoef2d = herme.hermefit(x, np.array([yw,yw]).T, 3, w=w)
         assert_almost_equal(wcoef2d, np.array([coef3,coef3]).T)
+        # test scaling with complex values x points whose square
+        # is zero when summed.
+        x = [1, 1j, -1, -1j]
+        assert_almost_equal(herme.hermefit(x, x, 1), [0, 1])
 
 class TestGauss(TestCase):
 
diff --git a/numpy/polynomial/tests/test_laguerre.py b/numpy/polynomial/tests/test_laguerre.py
index 14fafe37d..b6e0376a2 100644
--- a/numpy/polynomial/tests/test_laguerre.py
+++ b/numpy/polynomial/tests/test_laguerre.py
@@ -420,6 +420,10 @@ class TestFitting(TestCase):
         #
         wcoef2d = lag.lagfit(x, np.array([yw,yw]).T, 3, w=w)
         assert_almost_equal(wcoef2d, np.array([coef3,coef3]).T)
+        # test scaling with complex values x points whose square
+        # is zero when summed.
+        x = [1, 1j, -1, -1j]
+        assert_almost_equal(lag.lagfit(x, x, 1), [1, -1])
 
 class TestGauss(TestCase):
 
diff --git a/numpy/polynomial/tests/test_legendre.py b/numpy/polynomial/tests/test_legendre.py
index cdfaa96f1..3180db907 100644
--- a/numpy/polynomial/tests/test_legendre.py
+++ b/numpy/polynomial/tests/test_legendre.py
@@ -423,6 +423,10 @@ class TestFitting(TestCase):
         #
         wcoef2d = leg.legfit(x, np.array([yw,yw]).T, 3, w=w)
         assert_almost_equal(wcoef2d, np.array([coef3,coef3]).T)
+        # test scaling with complex values x points whose square
+        # is zero when summed.
+        x = [1, 1j, -1, -1j]
+        assert_almost_equal(leg.legfit(x, x, 1), [0, 1])
 
 class TestGauss(TestCase):
 
diff --git a/numpy/polynomial/tests/test_polynomial.py b/numpy/polynomial/tests/test_polynomial.py
index bae711cbf..bd09c07f6 100644
--- a/numpy/polynomial/tests/test_polynomial.py
+++ b/numpy/polynomial/tests/test_polynomial.py
@@ -440,6 +440,11 @@ class TestMisc(TestCase) :
         #
         wcoef2d = poly.polyfit(x, np.array([yw,yw]).T, 3, w=w)
         assert_almost_equal(wcoef2d, np.array([coef3,coef3]).T)
+        # test scaling with complex values x points whose square
+        # is zero when summed.
+        x = [1, 1j, -1, -1j]
+        assert_almost_equal(poly.polyfit(x, x, 1), [0, 1])
+
 
     def test_polytrim(self) :
         coef = [2, -1, 1, 0]