diff options
Diffstat (limited to 'numpy/polynomial/polynomial.py')
| -rw-r--r-- | numpy/polynomial/polynomial.py | 41 |
1 files changed, 18 insertions, 23 deletions
diff --git a/numpy/polynomial/polynomial.py b/numpy/polynomial/polynomial.py index 7c922c11b..5d05f5991 100644 --- a/numpy/polynomial/polynomial.py +++ b/numpy/polynomial/polynomial.py @@ -1217,14 +1217,11 @@ def polyfit(x, y, deg, rcond=None, full=False, w=None): sharing the same x-coordinates can be (independently) fit with one call to `polyfit` by passing in for `y` a 2-D array that contains one data set per column. - deg : int or array_like - Degree of the fitting polynomial. If `deg` is a single integer - all terms up to and including the `deg`'th term are included. - `deg` may alternatively be a list or array specifying which - terms in the Legendre expansion to include in the fit. - - .. versionchanged:: 1.11.0 - `deg` may be a list specifying which terms to fit + deg : int or 1-D array_like + Degree(s) of the fitting polynomials. If `deg` is a single integer + all terms up to and including the `deg`'th term are included in the + fit. For Numpy versions >= 1.11 a list of integers specifying the + degrees of the terms to include may be used instead. rcond : float, optional Relative condition number of the fit. Singular values smaller than `rcond`, relative to the largest singular value, will be @@ -1340,11 +1337,11 @@ def polyfit(x, y, deg, rcond=None, full=False, w=None): """ x = np.asarray(x) + 0.0 y = np.asarray(y) + 0.0 - deg = np.asarray([deg,], dtype=int).flatten() + deg = np.asarray(deg) # check arguments. - if deg.size < 1: - raise TypeError("expected deg to be one or more integers") + if deg.ndim > 1 or deg.dtype.kind not in 'iu' or deg.size == 0: + raise TypeError("deg must be an int or non-empty 1-D array of int") if deg.min() < 0: raise ValueError("expected deg >= 0") if x.ndim != 1: @@ -1356,19 +1353,17 @@ def polyfit(x, y, deg, rcond=None, full=False, w=None): if len(x) != len(y): raise TypeError("expected x and y to have same length") - if deg.size == 1: - restricted_fit = False - lmax = deg[0] + if deg.ndim == 0: + lmax = deg order = lmax + 1 + van = polyvander(x, lmax) else: - restricted_fit = True - lmax = deg.max() - order = deg.size + deg = np.sort(deg) + lmax = deg[-1] + order = len(deg) + van = polyvander(x, lmax)[:, deg] # set up the least squares matrices in transposed form - van = polyvander(x, lmax) - if restricted_fit: - van = van[:, deg] lhs = van.T rhs = y.T if w is not None: @@ -1398,11 +1393,11 @@ def polyfit(x, y, deg, rcond=None, full=False, w=None): c = (c.T/scl).T # Expand c to include non-fitted coefficients which are set to zero - if restricted_fit: + if deg.ndim == 1: if c.ndim == 2: - cc = np.zeros((lmax+1, c.shape[1]), dtype=c.dtype) + cc = np.zeros((lmax + 1, c.shape[1]), dtype=c.dtype) else: - cc = np.zeros(lmax+1, dtype=c.dtype) + cc = np.zeros(lmax + 1, dtype=c.dtype) cc[deg] = c c = cc |