MAINT: Make masked ma.polyfit match current polyfit.

1 files changed, 11 insertions, 27 deletions

diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
index bbc17d165..162890ee3 100644
--- a/numpy/ma/extras.py
+++ b/numpy/ma/extras.py
@@ -1845,14 +1845,14 @@ def vander(x, n=None):
 vander.__doc__ = ma.doc_note(np.vander.__doc__, vander.__doc__)
 
 
-def polyfit(x, y, deg, rcond=None, full=False):
+def polyfit(x, y, deg, rcond=None, full=False, w=None, cov=False):
     """
     Any masked values in x is propagated in y, and vice-versa.
     """
-    order = int(deg) + 1
     x = asarray(x)
-    mx = getmask(x)
     y = asarray(y)
+
+    mx = getmask(x)
     if y.ndim == 1:
         m = mask_or(mx, getmask(y))
     elif y.ndim == 2:
@@ -1864,31 +1864,15 @@ def polyfit(x, y, deg, rcond=None, full=False):
             m = mx
     else:
         raise TypeError("Expected a 1D or 2D array for y!")
+
     if m is not nomask:
-        x[m] = y[m] = masked
-    # Set rcond
-    if rcond is None :
-        rcond = len(x) * np.finfo(x.dtype).eps
-    # Scale x to improve condition number
-    scale = abs(x).max()
-    if scale != 0 :
-        x = x / scale
-    # solve least squares equation for powers of x
-    v = vander(x, order)
-    c, resids, rank, s = lstsq(v, y.filled(0), rcond)
-    # warn on rank reduction, which indicates an ill conditioned matrix
-    if rank != order and not full:
-        warnings.warn("Polyfit may be poorly conditioned", np.RankWarning)
-    # scale returned coefficients
-    if scale != 0 :
-        if c.ndim == 1 :
-            c /= np.vander([scale], order)[0]
-        else :
-            c /= np.vander([scale], order).T
-    if full :
-        return c, resids, rank, s, rcond
-    else :
-        return c
+        if w is not None:
+            w *= ~m
+        else:
+            w = ~m
+
+    return np.polyfit(x, y, deg, rcond, full, w, cov)
+
 polyfit.__doc__ = ma.doc_note(np.polyfit.__doc__, polyfit.__doc__)
 
 ################################################################################