diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/ma/extras.py | 38 | ||||
| -rw-r--r-- | numpy/ma/tests/test_extras.py | 10 |
2 files changed, 20 insertions, 28 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__) ################################################################################ diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py index 082d36f95..be995d617 100644 --- a/numpy/ma/tests/test_extras.py +++ b/numpy/ma/tests/test_extras.py @@ -630,7 +630,15 @@ class TestPolynomial(TestCase): (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_) - + # + w = np.random.rand(10) + 1 + xs = x[1:-1] + ys = y[1:-1] + ws = w[1:-1] + (C, R, K, S, D) = polyfit(x, y, 3, full=True, w=w) + (c, r, k, s, d) = np.polyfit(xs, ys, 3, full=True, w=ws) + for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)): + assert_almost_equal(a, a_) class TestArraySetOps(TestCase): |