1 files changed, 27 insertions, 12 deletions

diff --git a/numpy/polynomial/legendre.py b/numpy/polynomial/legendre.py
index bc9b5c2e6..97c387359 100644
--- a/numpy/polynomial/legendre.py
+++ b/numpy/polynomial/legendre.py
@@ -1245,7 +1245,10 @@ def legvander(x, deg) :
         raise ValueError("deg must be non-negative")
 
     x = np.array(x, copy=0, ndmin=1) + 0.0
-    v = np.empty((ideg + 1,) + x.shape, dtype=x.dtype)
+    dims = (ideg + 1,) + x.shape
+    mask = x.flags.maskna
+    dtyp = x.dtype
+    v = np.empty(dims, dtype=dtyp, maskna=mask)
     # Use forward recursion to generate the entries. This is not as accurate
     # as reverse recursion in this application but it is more efficient.
     v[0] = x*0 + 1
@@ -1391,6 +1394,11 @@ def legfit(x, y, deg, rcond=None, full=False, w=None):
 
     where `n` is `deg`.
 
+    Since numpy version 1.7.0, legfit also supports NA. If any of the
+    elements of `x`, `y`, or `w` are NA, then the corresponding rows of the
+    linear least squares problem (see Notes) are set to 0. If `y` is 2-D,
+    then an NA in any row of `y` invalidates that whole row.
+
     Parameters
     ----------
     x : array_like, shape (M,)
@@ -1503,30 +1511,37 @@ def legfit(x, y, deg, rcond=None, full=False, w=None):
     if len(x) != len(y):
         raise TypeError("expected x and y to have same length")
 
-    # set up the least squares matrices
-    lhs = legvander(x, deg)
-    rhs = y
+    # set up the least squares matrices in transposed form
+    lhs = legvander(x, deg).T
+    rhs = y.T
     if w is not None:
         w = np.asarray(w) + 0.0
         if w.ndim != 1:
             raise TypeError("expected 1D vector for w")
         if len(x) != len(w):
             raise TypeError("expected x and w to have same length")
-        # apply weights
-        if rhs.ndim == 2:
-            lhs *= w[:, np.newaxis]
-            rhs *= w[:, np.newaxis]
+        # apply weights. Don't use inplace operations as they
+        # can cause problems with NA.
+        lhs = lhs * w
+        rhs = rhs * w
+
+    # deal with NA. Note that polyvander propagates NA from x
+    # into all columns, that is rows for transposed form.
+    if lhs.flags.maskna or rhs.flags.maskna:
+        if rhs.ndim == 1:
+            mask = np.isna(lhs[0]) | np.isna(rhs)
         else:
-            lhs *= w[:, np.newaxis]
-            rhs *= w
+            mask = np.isna(lhs[0]) | np.isna(rhs).any(0)
+        np.copyto(lhs, 0, where=mask)
+        np.copyto(rhs, 0, where=mask)
 
     # set rcond
     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(0))
-    c, resids, rank, s = la.lstsq(lhs/scl, rhs, rcond)
+    scl = np.sqrt((lhs*lhs).sum(1))
+    c, resids, rank, s = la.lstsq(lhs.T/scl, rhs.T, rcond)
     c = (c.T/scl).T
 
     # warn on rank reduction