ENH: Add weights and covariance estimate to standard polyfit.

1 files changed, 47 insertions, 18 deletions

diff --git a/numpy/lib/polynomial.py b/numpy/lib/polynomial.py
index 229fe020d..fad06f4df 100644
--- a/numpy/lib/polynomial.py
+++ b/numpy/lib/polynomial.py
@@ -10,11 +10,11 @@ import re
 import warnings
 import numpy.core.numeric as NX
 
-from numpy.core import isscalar, abs, finfo, atleast_1d, hstack
+from numpy.core import isscalar, abs, finfo, atleast_1d, hstack, dot
 from numpy.lib.twodim_base import diag, vander
 from numpy.lib.function_base import trim_zeros, sort_complex
 from numpy.lib.type_check import iscomplex, real, imag
-from numpy.linalg import eigvals, lstsq
+from numpy.linalg import eigvals, lstsq, inv
 
 class RankWarning(UserWarning):
     """
@@ -391,7 +391,7 @@ def polyder(p, m=1):
         val = poly1d(val)
     return val
 
-def polyfit(x, y, deg, rcond=None, full=False):
+def polyfit(x, y, deg, rcond=None, full=False, w=None, cov=False):
     """
     Least squares polynomial fit.
 
@@ -419,6 +419,11 @@ def polyfit(x, y, deg, rcond=None, full=False):
         False (the default) just the coefficients are returned, when True
         diagnostic information from the singular value decomposition is also
         returned.
+    w : array_like, shape (M,), optional
+        weights to apply to the y-coordinates of the sample points.
+    cov : bool, optional
+        Return the estimate and the covariance matrix of the estimate
+        If full is True, then cov is not returned.
 
     Returns
     -------
@@ -431,6 +436,12 @@ def polyfit(x, y, deg, rcond=None, full=False):
         Vandermonde coefficient matrix, its singular values, and the specified
         value of `rcond`. For more details, see `linalg.lstsq`.
 
+    V : ndaray, shape (M,M) or (M,M,K) : present only if `full` = False and `cov`=True
+        The covariance matrix of the polynomial coefficient estimates.  The diagonal
+        of this matrix are the variance estimates for each coefficient.  If y is a 2-d
+        array, then the covariance matrix for the `k`-th data set are in ``V[:,:,k]``
+
+
     Warns
     -----
     RankWarning
@@ -545,34 +556,52 @@ def polyfit(x, y, deg, rcond=None, full=False):
     if rcond is None :
         rcond = len(x)*finfo(x.dtype).eps
 
-    # scale x to improve condition number
-    scale = abs(x).max()
-    if scale != 0 :
-        x /= scale
+    # set up least squares equation for powers of x
+    lhs = vander(x, order)
+    rhs = y
+
+    # apply weighting
+    if w is not None:
+        w = NX.asarray(w) + 0.0
+        if w.ndim != 1:
+            raise TypeError, "expected a 1-d array for weights"
+        if w.shape[0] != y.shape[0] :
+            raise TypeError, "expected w and y to have the same length"
+        lhs *= w[:, NX.newaxis]
+        if rhs.ndim == 2:
+            rhs *= w[:, NX.newaxis]
+        else:
+            rhs *= w
 
-    # solve least squares equation for powers of x
-    v = vander(x, order)
-    c, resids, rank, s = lstsq(v, y, rcond)
+    # scale lhs to improve condition number and solve
+    scale = NX.sqrt((lhs*lhs).sum(axis=0))
+    lhs /= scale
+    c, resids, rank, s = lstsq(lhs, rhs, rcond)
+    c = (c.T/scale).T  # broadcast scale coefficients
 
     # warn on rank reduction, which indicates an ill conditioned matrix
     if rank != order and not full:
         msg = "Polyfit may be poorly conditioned"
         warnings.warn(msg, RankWarning)
 
-    # scale returned coefficients
-    if scale != 0 :
-        if c.ndim == 1 :
-            c /= vander([scale], order)[0]
-        else :
-            c /= vander([scale], order).T
-
     if full :
         return c, resids, rank, s, rcond
+    elif cov :
+        Vbase = inv(dot(lhs.T,lhs))
+        Vbase /= NX.outer(scale, scale)
+        # Some literature ignores the extra -2.0 factor in the denominator, but
+        #  it is included here because the covariance of Multivariate Student-T
+        #  (which is implied by a Bayesian uncertainty analysis) includes it.
+        #  Plus, it gives a slightly more conservative estimate of uncertainty.
+        fac = resids / (len(x) - order - 2.0)
+        if y.ndim == 1:
+            return c, Vbase * fac
+        else:
+            return c, Vbase[:,:,NX.newaxis] * fac
     else :
         return c
 
 
-
 def polyval(p, x):
     """
     Evaluate a polynomial at specific values.