diff options
Diffstat (limited to 'numpy/ma/extras.py')
| -rw-r--r-- | numpy/ma/extras.py | 38 |
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__) ################################################################################ |