BUG: numpy.ma.polyfit masks NaNs incorrectly

This fixes the incorrect handing of masked NaNs by
``np.ma.polyfit``. Instead of passing the mask into
``np.polyfit`` by setting the weight of the masked
points to zero, the subset of elements of which are
to be fitted are passed instead.

Closes #5591

1 files changed, 15 insertions, 0 deletions

diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py
index 8902d67dc..9e1e8ba38 100644
--- a/numpy/ma/tests/test_extras.py
+++ b/numpy/ma/tests/test_extras.py
@@ -734,6 +734,21 @@ class TestPolynomial(TestCase):
         for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)):
             assert_almost_equal(a, a_)
 
+    def test_polyfit_with_masked_NaNs(self):
+        x = np.random.rand(10)
+        y = np.random.rand(20).reshape(-1, 2)
+
+        x[0] = np.nan
+        y[-1,-1] = np.nan
+        x = x.view(MaskedArray)
+        y = y.view(MaskedArray)
+        x[0] = masked
+        y[-1,-1] = masked
+
+        (C, R, K, S, D) = polyfit(x, y, 3, full=True)
+        (c, r, k, s, d) = np.polyfit(x[1:-1], y[1:-1,:], 3, full=True)
+        for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)):
+            assert_almost_equal(a, a_)
 
 class TestArraySetOps(TestCase):