diff options
Diffstat (limited to 'numpy/polynomial/hermite_e.py')
| -rw-r--r-- | numpy/polynomial/hermite_e.py | 39 |
1 files changed, 27 insertions, 12 deletions
diff --git a/numpy/polynomial/hermite_e.py b/numpy/polynomial/hermite_e.py index 7a42d48a3..b81681a82 100644 --- a/numpy/polynomial/hermite_e.py +++ b/numpy/polynomial/hermite_e.py @@ -1214,7 +1214,10 @@ def hermevander(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) v[0] = x*0 + 1 if ideg > 0 : v[1] = x @@ -1358,6 +1361,11 @@ def hermefit(x, y, deg, rcond=None, full=False, w=None): where `n` is `deg`. + Since numpy version 1.7.0, hermefit 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,) @@ -1475,30 +1483,37 @@ def hermefit(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 = hermevander(x, deg) - rhs = y + # set up the least squares matrices in transposed form + lhs = hermevander(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 |