ENH: Allow specification of terms to fit in chebfit

The argument `deg` is enhanced to allow an
array_like argument to past which specifies
which terms to include in the fit.

The returned coef array is exapnded to have
entries of 0 for all coefficients which were
not included in the fit.

1 files changed, 34 insertions, 5 deletions

diff --git a/numpy/polynomial/chebyshev.py b/numpy/polynomial/chebyshev.py
index de5cbb734..9db613b78 100644
--- a/numpy/polynomial/chebyshev.py
+++ b/numpy/polynomial/chebyshev.py
@@ -1617,8 +1617,14 @@ def chebfit(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 series
+    deg : int or array_like
+        Degree of the fitting series. 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
@@ -1710,12 +1716,14 @@ def chebfit(x, y, deg, rcond=None, full=False, w=None):
     --------
 
     """
-    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")
@@ -1726,8 +1734,20 @@ def chebfit(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 = chebvander(x, deg).T
+    van = chebvander(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
@@ -1755,6 +1775,15 @@ def chebfit(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"