diff options
Diffstat (limited to 'numpy/polynomial/hermite_e.py')
| -rw-r--r-- | numpy/polynomial/hermite_e.py | 39 |
1 files changed, 34 insertions, 5 deletions
diff --git a/numpy/polynomial/hermite_e.py b/numpy/polynomial/hermite_e.py index 394f256a5..da441af83 100644 --- a/numpy/polynomial/hermite_e.py +++ b/numpy/polynomial/hermite_e.py @@ -1385,8 +1385,14 @@ def hermefit(x, y, deg, rcond=None, full=False, w=None): y-coordinates of the sample points. Several data sets of sample points sharing the same x-coordinates can be fitted at once by passing in a 2D-array that contains one dataset per column. - deg : int - Degree of the fitting polynomial + 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 rcond : float, optional Relative condition number of the fit. Singular values smaller than this relative to the largest singular value will be ignored. The @@ -1483,12 +1489,14 @@ def hermefit(x, y, deg, rcond=None, full=False, w=None): array([ 1.01690445, 1.99951418, 2.99948696]) """ - order = int(deg) + 1 x = np.asarray(x) + 0.0 y = np.asarray(y) + 0.0 + deg = np.asarray([deg,], dtype=int).flatten() # check arguments. - if deg < 0: + if deg.size < 1: + raise TypeError("expected deg to be one or more integers") + if deg.min() < 0: raise ValueError("expected deg >= 0") if x.ndim != 1: raise TypeError("expected 1D vector for x") @@ -1499,8 +1507,20 @@ 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") + if deg.size == 1: + restricted_fit = False + lmax = deg[0] + order = lmax + 1 + else: + restricted_fit = True + lmax = deg.max() + order = deg.size + # set up the least squares matrices in transposed form - lhs = hermevander(x, deg).T + van = hermevander(x, lmax) + if restricted_fit: + van = van[:, deg] + lhs = van.T rhs = y.T if w is not None: w = np.asarray(w) + 0.0 @@ -1528,6 +1548,15 @@ def hermefit(x, y, deg, rcond=None, full=False, w=None): c, resids, rank, s = la.lstsq(lhs.T/scl, rhs.T, rcond) c = (c.T/scl).T + # Expand c to include non-fitted coefficients which are set to zero + if restricted_fit: + if c.ndim == 2: + cc = np.zeros((lmax+1, c.shape[1]), dtype=c.dtype) + else: + cc = np.zeros(lmax+1, dtype=c.dtype) + cc[deg] = c + c = cc + # warn on rank reduction if rank != order and not full: msg = "The fit may be poorly conditioned" |