Fix polyfit for 2D case and add test for same. Fixes ticket 697.

1 files changed, 68 insertions, 41 deletions

diff --git a/numpy/lib/polynomial.py b/numpy/lib/polynomial.py
index da1908d93..e33bc4e57 100644
--- a/numpy/lib/polynomial.py
+++ b/numpy/lib/polynomial.py
@@ -182,44 +182,71 @@ def polyder(p, m=1):
 def polyfit(x, y, deg, rcond=None, full=False):
     """Least squares polynomial fit.
 
-    Required arguments
-
-        x -- vector of sample points
-        y -- vector or 2D array of values to fit
-        deg -- degree of the fitting polynomial
-
-    Keyword arguments
+    Do a best fit polynomial of degree 'deg' of 'x' to 'y'.  Return value is a
+    vector of polynomial coefficients [pk ... p1 p0].  Eg, for n=2
 
-        rcond -- relative condition number of the fit (default len(x)*eps)
-        full -- return full diagnostic output (default False)
+        p2*x0^2 +  p1*x0 + p0 = y1
+        p2*x1^2 +  p1*x1 + p0 = y1
+        p2*x2^2 +  p1*x2 + p0 = y2
+        .....
+        p2*xk^2 +  p1*xk + p0 = yk
+
+    Parameters
+    ----------
+
+    x : array_like
+        1D vector of sample points.
+    y : array_like
+        1D vector or 2D array of values to fit. The values should run down the
+        columes in the 2D case.
+    deg : integer
+        Degree of the fitting polynomial
+    rcond: {None, float}, optional
+        Relative condition number of the fit. Singular values smaller than this
+        relative to the largest singular value will be ignored. The defaul value
+        is len(x)*eps, where eps is the relative precision of the float type,
+        about 2e-16 in most cases.
+    full : {False, boolean}, optional
+        Switch determining nature of return value. When it is False just the
+        coefficients are returned, when True diagnostic information from the
+        singular value decomposition is also returned.
 
     Returns
+    -------
 
-        full == False -- coefficients
-        full == True -- coefficients, residuals, rank, singular values, rcond.
+    coefficients, [residuals, rank, singular_values, rcond] : variable
+        When full=False, only the coefficients are returned, running down the
+        appropriate colume when y is a 2D array. When full=True, the rank of
+        the scaled Vandermonde matrix, it's effective rank in light of the rcond
+        value, its singular values, and the specified value of rcond are also
+        returned.
 
     Warns
+    -----
 
-        RankWarning -- if rank is reduced and not full output
-
-    Do a best fit polynomial of degree 'deg' of 'x' to 'y'.  Return value is a
-    vector of polynomial coefficients [pk ... p1 p0].  Eg, for n=2
+    RankWarning : if rank is reduced and not full output
+        The warnings can be turned off by:
+        >>> import numpy as np
+        >>> import warnings
+        >>> warnings.simplefilter('ignore',np.RankWarning)
 
-      p2*x0^2 +  p1*x0 + p0 = y1
-      p2*x1^2 +  p1*x1 + p0 = y1
-      p2*x2^2 +  p1*x2 + p0 = y2
-      .....
-      p2*xk^2 +  p1*xk + p0 = yk
 
+    See Also
+    --------
+    
+    polyval : computes polynomial values.
 
-    Method: if X is a the Vandermonde Matrix computed from x (see
+    Notes
+    -----
+    
+    If X is a the Vandermonde Matrix computed from x (see
     http://mathworld.wolfram.com/VandermondeMatrix.html), then the
     polynomial least squares solution is given by the 'p' in
 
       X*p = y
 
-    where X is a len(x) x N+1 matrix, p is a N+1 length vector, and y
-    is a len(x) x 1 vector
+    where X.shape is a matrix of dimensions (len(x), deg + 1), p is a vector of
+    dimensions (deg + 1, 1), and y is a vector of dimensions (len(x), 1).
 
     This equation can be solved as
 
@@ -227,7 +254,7 @@ def polyfit(x, y, deg, rcond=None, full=False):
 
     where XT is the transpose of X and -1 denotes the inverse. However, this
     method is susceptible to rounding errors and generally the singular value
-    decomposition is preferred and that is the method used here. The singular
+    decomposition of the matrix X is preferred and that is what is done here. The singular
     value method takes a paramenter, 'rcond', which sets a limit on the
     relative size of the smallest singular value to be used in solving the
     equation. This may result in lowering the rank of the Vandermonde matrix,
@@ -235,27 +262,24 @@ def polyfit(x, y, deg, rcond=None, full=False):
     a fit of lower degree or replace x by x - x.mean(), both of which will
     generally improve the condition number. The routine already normalizes the
     vector x by its maximum absolute value to help in this regard. The rcond
-    parameter may also be set to a value smaller than its default, but this may
-    result in bad fits. The current default value of rcond is len(x)*eps, where
+    parameter can be set to a value smaller than its default, but the
+    resulting fit may be spurious. The current default value of rcond is len(x)*eps, where
     eps is the relative precision of the floating type being used, generally
     around 1e-7 and 2e-16 for IEEE single and double precision respectively.
     This value of rcond is fairly conservative but works pretty well when x -
     x.mean() is used in place of x.
 
-    The warnings can be turned off by:
-
-    >>> import numpy
-    >>> import warnings
-    >>> warnings.simplefilter('ignore',numpy.RankWarning)
 
     DISCLAIMER: Power series fits are full of pitfalls for the unwary once the
     degree of the fit becomes large or the interval of sample points is badly
-    centered. The basic problem is that the powers x**n are generally a poor
-    basis for the functions on the sample interval with the result that the
-    Vandermonde matrix is ill conditioned and computation of the polynomial
-    values is sensitive to coefficient error. The quality of the resulting fit
-    should be checked against the data whenever the condition number is large,
-    as the quality of polynomial fits *can not* be taken for granted. If all
+    centered. The problem is that the powers x**n are generally a poor
+    basis for the polynomial functions on the sample interval, resulting in a
+    Vandermonde matrix is ill conditioned and coefficients sensitive to rounding
+    erros. The computation of the polynomial
+    values will also sensitive to rounding errors. Consequently, the quality of
+    the polynomial fit
+    should be checked against the data whenever the condition number is large.
+    The quality of polynomial fits *can not* be taken for granted. If all
     you want to do is draw a smooth curve through the y values and polyfit is
     not doing the job, try centering the sample range or look into
     scipy.interpolate, which includes some nice spline fitting functions that
@@ -266,8 +290,6 @@ def polyfit(x, y, deg, rcond=None, full=False):
     but note that the k's and n's in the superscripts and subscripts
     on that page.  The linear algebra is correct, however.
 
-    See also polyval
-
     """
     order = int(deg) + 1
     x = NX.asarray(x) + 0.0
@@ -276,7 +298,9 @@ def polyfit(x, y, deg, rcond=None, full=False):
     # check arguments.
     if deg < 0 :
         raise ValueError, "expected deg >= 0"
-    if x.ndim != 1 or x.size == 0:
+    if x.ndim != 1:
+        raise TypeError, "expected 1D vector for x"
+    if x.size == 0:
         raise TypeError, "expected non-empty vector for x"
     if y.ndim < 1 or y.ndim > 2 :
         raise TypeError, "expected 1D or 2D array for y"
@@ -307,7 +331,10 @@ def polyfit(x, y, deg, rcond=None, full=False):
 
     # scale returned coefficients
     if scale != 0 :
-        c /= vander([scale], order)[0]
+        if c.ndim == 1 :
+            c /= vander([scale], order)[0]
+        else :
+            c /= vander([scale], order).T
 
     if full :
         return c, resids, rank, s, rcond