ENH: Allow specification of terms to fit in legfit

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, 39 insertions, 8 deletions

diff --git a/numpy/polynomial/legendre.py b/numpy/polynomial/legendre.py
index c91cb72ec..2035ba6e9 100644
--- a/numpy/polynomial/legendre.py
+++ b/numpy/polynomial/legendre.py
@@ -1418,8 +1418,14 @@ def legfit(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
@@ -1440,9 +1446,11 @@ def legfit(x, y, deg, rcond=None, full=False, w=None):
     Returns
     -------
     coef : ndarray, shape (M,) or (M, K)
-        Legendre coefficients ordered from low to high. If `y` was 2-D,
-        the coefficients for the data in column k  of `y` are in column
-        `k`.
+        Legendre coefficients ordered from low to high. If `y` was
+        2-D, the coefficients for the data in column k of `y` are in
+        column `k`. If `deg` is specified as a list, coefficients for
+        terms not included in the fit are set equal to zero in the
+        returned `coef`.
 
     [residuals, rank, singular_values, rcond] : list
         These values are only returned if `full` = True
@@ -1511,12 +1519,14 @@ def legfit(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")
@@ -1527,8 +1537,20 @@ def legfit(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 = legvander(x, deg).T
+    van = legvander(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
@@ -1556,6 +1578,15 @@ def legfit(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"