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 numpy#5591

Author: simongibbons

numpy/ma/extras.py

@@ -1921,12 +1921,12 @@ def polyfit(x, y, deg, rcond=None, full=False, w=None, cov=False):
         m = mask_or(m, getmask(w))
 
     if m is not nomask:
+        not_m = ~m
         if w is not None:
-            w = ~m*w
-        else:
-            w = ~m
-
-    return np.polyfit(x, y, deg, rcond, full, w, cov)
+            w = w[not_m]
+        return np.polyfit(x[not_m], y[not_m], deg, rcond, full, w, cov)
+    else:
+        return np.polyfit(x, y, deg, rcond, full, w, cov)
 
 polyfit.__doc__ = ma.doc_note(np.polyfit.__doc__, polyfit.__doc__)
 

numpy/ma/tests/test_extras.py

@@ -734,6 +734,21 @@ def test_polyfit(self):
         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):