Chris@87: """
Chris@87: Functions to operate on polynomials.
Chris@87: 
Chris@87: """
Chris@87: from __future__ import division, absolute_import, print_function
Chris@87: 
Chris@87: __all__ = ['poly', 'roots', 'polyint', 'polyder', 'polyadd',
Chris@87:            'polysub', 'polymul', 'polydiv', 'polyval', 'poly1d',
Chris@87:            'polyfit', 'RankWarning']
Chris@87: 
Chris@87: import re
Chris@87: import warnings
Chris@87: import numpy.core.numeric as NX
Chris@87: 
Chris@87: from numpy.core import isscalar, abs, finfo, atleast_1d, hstack, dot
Chris@87: from numpy.lib.twodim_base import diag, vander
Chris@87: from numpy.lib.function_base import trim_zeros, sort_complex
Chris@87: from numpy.lib.type_check import iscomplex, real, imag
Chris@87: from numpy.linalg import eigvals, lstsq, inv
Chris@87: 
Chris@87: class RankWarning(UserWarning):
Chris@87:     """
Chris@87:     Issued by `polyfit` when the Vandermonde matrix is rank deficient.
Chris@87: 
Chris@87:     For more information, a way to suppress the warning, and an example of
Chris@87:     `RankWarning` being issued, see `polyfit`.
Chris@87: 
Chris@87:     """
Chris@87:     pass
Chris@87: 
Chris@87: def poly(seq_of_zeros):
Chris@87:     """
Chris@87:     Find the coefficients of a polynomial with the given sequence of roots.
Chris@87: 
Chris@87:     Returns the coefficients of the polynomial whose leading coefficient
Chris@87:     is one for the given sequence of zeros (multiple roots must be included
Chris@87:     in the sequence as many times as their multiplicity; see Examples).
Chris@87:     A square matrix (or array, which will be treated as a matrix) can also
Chris@87:     be given, in which case the coefficients of the characteristic polynomial
Chris@87:     of the matrix are returned.
Chris@87: 
Chris@87:     Parameters
Chris@87:     ----------
Chris@87:     seq_of_zeros : array_like, shape (N,) or (N, N)
Chris@87:         A sequence of polynomial roots, or a square array or matrix object.
Chris@87: 
Chris@87:     Returns
Chris@87:     -------
Chris@87:     c : ndarray
Chris@87:         1D array of polynomial coefficients from highest to lowest degree:
Chris@87: 
Chris@87:         ``c[0] * x**(N) + c[1] * x**(N-1) + ... + c[N-1] * x + c[N]``
Chris@87:         where c[0] always equals 1.
Chris@87: 
Chris@87:     Raises
Chris@87:     ------
Chris@87:     ValueError
Chris@87:         If input is the wrong shape (the input must be a 1-D or square
Chris@87:         2-D array).
Chris@87: 
Chris@87:     See Also
Chris@87:     --------
Chris@87:     polyval : Evaluate a polynomial at a point.
Chris@87:     roots : Return the roots of a polynomial.
Chris@87:     polyfit : Least squares polynomial fit.
Chris@87:     poly1d : A one-dimensional polynomial class.
Chris@87: 
Chris@87:     Notes
Chris@87:     -----
Chris@87:     Specifying the roots of a polynomial still leaves one degree of
Chris@87:     freedom, typically represented by an undetermined leading
Chris@87:     coefficient. [1]_ In the case of this function, that coefficient -
Chris@87:     the first one in the returned array - is always taken as one. (If
Chris@87:     for some reason you have one other point, the only automatic way
Chris@87:     presently to leverage that information is to use ``polyfit``.)
Chris@87: 
Chris@87:     The characteristic polynomial, :math:`p_a(t)`, of an `n`-by-`n`
Chris@87:     matrix **A** is given by
Chris@87: 
Chris@87:         :math:`p_a(t) = \\mathrm{det}(t\\, \\mathbf{I} - \\mathbf{A})`,
Chris@87: 
Chris@87:     where **I** is the `n`-by-`n` identity matrix. [2]_
Chris@87: 
Chris@87:     References
Chris@87:     ----------
Chris@87:     .. [1] M. Sullivan and M. Sullivan, III, "Algebra and Trignometry,
Chris@87:        Enhanced With Graphing Utilities," Prentice-Hall, pg. 318, 1996.
Chris@87: 
Chris@87:     .. [2] G. Strang, "Linear Algebra and Its Applications, 2nd Edition,"
Chris@87:        Academic Press, pg. 182, 1980.
Chris@87: 
Chris@87:     Examples
Chris@87:     --------
Chris@87:     Given a sequence of a polynomial's zeros:
Chris@87: 
Chris@87:     >>> np.poly((0, 0, 0)) # Multiple root example
Chris@87:     array([1, 0, 0, 0])
Chris@87: 
Chris@87:     The line above represents z**3 + 0*z**2 + 0*z + 0.
Chris@87: 
Chris@87:     >>> np.poly((-1./2, 0, 1./2))
Chris@87:     array([ 1.  ,  0.  , -0.25,  0.  ])
Chris@87: 
Chris@87:     The line above represents z**3 - z/4
Chris@87: 
Chris@87:     >>> np.poly((np.random.random(1.)[0], 0, np.random.random(1.)[0]))
Chris@87:     array([ 1.        , -0.77086955,  0.08618131,  0.        ]) #random
Chris@87: 
Chris@87:     Given a square array object:
Chris@87: 
Chris@87:     >>> P = np.array([[0, 1./3], [-1./2, 0]])
Chris@87:     >>> np.poly(P)
Chris@87:     array([ 1.        ,  0.        ,  0.16666667])
Chris@87: 
Chris@87:     Or a square matrix object:
Chris@87: 
Chris@87:     >>> np.poly(np.matrix(P))
Chris@87:     array([ 1.        ,  0.        ,  0.16666667])
Chris@87: 
Chris@87:     Note how in all cases the leading coefficient is always 1.
Chris@87: 
Chris@87:     """
Chris@87:     seq_of_zeros = atleast_1d(seq_of_zeros)
Chris@87:     sh = seq_of_zeros.shape
Chris@87:     if len(sh) == 2 and sh[0] == sh[1] and sh[0] != 0:
Chris@87:         seq_of_zeros = eigvals(seq_of_zeros)
Chris@87:     elif len(sh) == 1:
Chris@87:         pass
Chris@87:     else:
Chris@87:         raise ValueError("input must be 1d or non-empty square 2d array.")
Chris@87: 
Chris@87:     if len(seq_of_zeros) == 0:
Chris@87:         return 1.0
Chris@87: 
Chris@87:     a = [1]
Chris@87:     for k in range(len(seq_of_zeros)):
Chris@87:         a = NX.convolve(a, [1, -seq_of_zeros[k]], mode='full')
Chris@87: 
Chris@87:     if issubclass(a.dtype.type, NX.complexfloating):
Chris@87:         # if complex roots are all complex conjugates, the roots are real.
Chris@87:         roots = NX.asarray(seq_of_zeros, complex)
Chris@87:         pos_roots = sort_complex(NX.compress(roots.imag > 0, roots))
Chris@87:         neg_roots = NX.conjugate(sort_complex(
Chris@87:                                         NX.compress(roots.imag < 0, roots)))
Chris@87:         if (len(pos_roots) == len(neg_roots) and
Chris@87:                 NX.alltrue(neg_roots == pos_roots)):
Chris@87:             a = a.real.copy()
Chris@87: 
Chris@87:     return a
Chris@87: 
Chris@87: def roots(p):
Chris@87:     """
Chris@87:     Return the roots of a polynomial with coefficients given in p.
Chris@87: 
Chris@87:     The values in the rank-1 array `p` are coefficients of a polynomial.
Chris@87:     If the length of `p` is n+1 then the polynomial is described by::
Chris@87: 
Chris@87:       p[0] * x**n + p[1] * x**(n-1) + ... + p[n-1]*x + p[n]
Chris@87: 
Chris@87:     Parameters
Chris@87:     ----------
Chris@87:     p : array_like
Chris@87:         Rank-1 array of polynomial coefficients.
Chris@87: 
Chris@87:     Returns
Chris@87:     -------
Chris@87:     out : ndarray
Chris@87:         An array containing the complex roots of the polynomial.
Chris@87: 
Chris@87:     Raises
Chris@87:     ------
Chris@87:     ValueError
Chris@87:         When `p` cannot be converted to a rank-1 array.
Chris@87: 
Chris@87:     See also
Chris@87:     --------
Chris@87:     poly : Find the coefficients of a polynomial with a given sequence
Chris@87:            of roots.
Chris@87:     polyval : Evaluate a polynomial at a point.
Chris@87:     polyfit : Least squares polynomial fit.
Chris@87:     poly1d : A one-dimensional polynomial class.
Chris@87: 
Chris@87:     Notes
Chris@87:     -----
Chris@87:     The algorithm relies on computing the eigenvalues of the
Chris@87:     companion matrix [1]_.
Chris@87: 
Chris@87:     References
Chris@87:     ----------
Chris@87:     .. [1] R. A. Horn & C. R. Johnson, *Matrix Analysis*.  Cambridge, UK:
Chris@87:         Cambridge University Press, 1999, pp. 146-7.
Chris@87: 
Chris@87:     Examples
Chris@87:     --------
Chris@87:     >>> coeff = [3.2, 2, 1]
Chris@87:     >>> np.roots(coeff)
Chris@87:     array([-0.3125+0.46351241j, -0.3125-0.46351241j])
Chris@87: 
Chris@87:     """
Chris@87:     # If input is scalar, this makes it an array
Chris@87:     p = atleast_1d(p)
Chris@87:     if len(p.shape) != 1:
Chris@87:         raise ValueError("Input must be a rank-1 array.")
Chris@87: 
Chris@87:     # find non-zero array entries
Chris@87:     non_zero = NX.nonzero(NX.ravel(p))[0]
Chris@87: 
Chris@87:     # Return an empty array if polynomial is all zeros
Chris@87:     if len(non_zero) == 0:
Chris@87:         return NX.array([])
Chris@87: 
Chris@87:     # find the number of trailing zeros -- this is the number of roots at 0.
Chris@87:     trailing_zeros = len(p) - non_zero[-1] - 1
Chris@87: 
Chris@87:     # strip leading and trailing zeros
Chris@87:     p = p[int(non_zero[0]):int(non_zero[-1])+1]
Chris@87: 
Chris@87:     # casting: if incoming array isn't floating point, make it floating point.
Chris@87:     if not issubclass(p.dtype.type, (NX.floating, NX.complexfloating)):
Chris@87:         p = p.astype(float)
Chris@87: 
Chris@87:     N = len(p)
Chris@87:     if N > 1:
Chris@87:         # build companion matrix and find its eigenvalues (the roots)
Chris@87:         A = diag(NX.ones((N-2,), p.dtype), -1)
Chris@87:         A[0,:] = -p[1:] / p[0]
Chris@87:         roots = eigvals(A)
Chris@87:     else:
Chris@87:         roots = NX.array([])
Chris@87: 
Chris@87:     # tack any zeros onto the back of the array
Chris@87:     roots = hstack((roots, NX.zeros(trailing_zeros, roots.dtype)))
Chris@87:     return roots
Chris@87: 
Chris@87: def polyint(p, m=1, k=None):
Chris@87:     """
Chris@87:     Return an antiderivative (indefinite integral) of a polynomial.
Chris@87: 
Chris@87:     The returned order `m` antiderivative `P` of polynomial `p` satisfies
Chris@87:     :math:`\\frac{d^m}{dx^m}P(x) = p(x)` and is defined up to `m - 1`
Chris@87:     integration constants `k`. The constants determine the low-order
Chris@87:     polynomial part
Chris@87: 
Chris@87:     .. math:: \\frac{k_{m-1}}{0!} x^0 + \\ldots + \\frac{k_0}{(m-1)!}x^{m-1}
Chris@87: 
Chris@87:     of `P` so that :math:`P^{(j)}(0) = k_{m-j-1}`.
Chris@87: 
Chris@87:     Parameters
Chris@87:     ----------
Chris@87:     p : {array_like, poly1d}
Chris@87:         Polynomial to differentiate.
Chris@87:         A sequence is interpreted as polynomial coefficients, see `poly1d`.
Chris@87:     m : int, optional
Chris@87:         Order of the antiderivative. (Default: 1)
Chris@87:     k : {None, list of `m` scalars, scalar}, optional
Chris@87:         Integration constants. They are given in the order of integration:
Chris@87:         those corresponding to highest-order terms come first.
Chris@87: 
Chris@87:         If ``None`` (default), all constants are assumed to be zero.
Chris@87:         If `m = 1`, a single scalar can be given instead of a list.
Chris@87: 
Chris@87:     See Also
Chris@87:     --------
Chris@87:     polyder : derivative of a polynomial
Chris@87:     poly1d.integ : equivalent method
Chris@87: 
Chris@87:     Examples
Chris@87:     --------
Chris@87:     The defining property of the antiderivative:
Chris@87: 
Chris@87:     >>> p = np.poly1d([1,1,1])
Chris@87:     >>> P = np.polyint(p)
Chris@87:     >>> P
Chris@87:     poly1d([ 0.33333333,  0.5       ,  1.        ,  0.        ])
Chris@87:     >>> np.polyder(P) == p
Chris@87:     True
Chris@87: 
Chris@87:     The integration constants default to zero, but can be specified:
Chris@87: 
Chris@87:     >>> P = np.polyint(p, 3)
Chris@87:     >>> P(0)
Chris@87:     0.0
Chris@87:     >>> np.polyder(P)(0)
Chris@87:     0.0
Chris@87:     >>> np.polyder(P, 2)(0)
Chris@87:     0.0
Chris@87:     >>> P = np.polyint(p, 3, k=[6,5,3])
Chris@87:     >>> P
Chris@87:     poly1d([ 0.01666667,  0.04166667,  0.16666667,  3. ,  5. ,  3. ])
Chris@87: 
Chris@87:     Note that 3 = 6 / 2!, and that the constants are given in the order of
Chris@87:     integrations. Constant of the highest-order polynomial term comes first:
Chris@87: 
Chris@87:     >>> np.polyder(P, 2)(0)
Chris@87:     6.0
Chris@87:     >>> np.polyder(P, 1)(0)
Chris@87:     5.0
Chris@87:     >>> P(0)
Chris@87:     3.0
Chris@87: 
Chris@87:     """
Chris@87:     m = int(m)
Chris@87:     if m < 0:
Chris@87:         raise ValueError("Order of integral must be positive (see polyder)")
Chris@87:     if k is None:
Chris@87:         k = NX.zeros(m, float)
Chris@87:     k = atleast_1d(k)
Chris@87:     if len(k) == 1 and m > 1:
Chris@87:         k = k[0]*NX.ones(m, float)
Chris@87:     if len(k) < m:
Chris@87:         raise ValueError(
Chris@87:               "k must be a scalar or a rank-1 array of length 1 or >m.")
Chris@87: 
Chris@87:     truepoly = isinstance(p, poly1d)
Chris@87:     p = NX.asarray(p)
Chris@87:     if m == 0:
Chris@87:         if truepoly:
Chris@87:             return poly1d(p)
Chris@87:         return p
Chris@87:     else:
Chris@87:         # Note: this must work also with object and integer arrays
Chris@87:         y = NX.concatenate((p.__truediv__(NX.arange(len(p), 0, -1)), [k[0]]))
Chris@87:         val = polyint(y, m - 1, k=k[1:])
Chris@87:         if truepoly:
Chris@87:             return poly1d(val)
Chris@87:         return val
Chris@87: 
Chris@87: def polyder(p, m=1):
Chris@87:     """
Chris@87:     Return the derivative of the specified order of a polynomial.
Chris@87: 
Chris@87:     Parameters
Chris@87:     ----------
Chris@87:     p : poly1d or sequence
Chris@87:         Polynomial to differentiate.
Chris@87:         A sequence is interpreted as polynomial coefficients, see `poly1d`.
Chris@87:     m : int, optional
Chris@87:         Order of differentiation (default: 1)
Chris@87: 
Chris@87:     Returns
Chris@87:     -------
Chris@87:     der : poly1d
Chris@87:         A new polynomial representing the derivative.
Chris@87: 
Chris@87:     See Also
Chris@87:     --------
Chris@87:     polyint : Anti-derivative of a polynomial.
Chris@87:     poly1d : Class for one-dimensional polynomials.
Chris@87: 
Chris@87:     Examples
Chris@87:     --------
Chris@87:     The derivative of the polynomial :math:`x^3 + x^2 + x^1 + 1` is:
Chris@87: 
Chris@87:     >>> p = np.poly1d([1,1,1,1])
Chris@87:     >>> p2 = np.polyder(p)
Chris@87:     >>> p2
Chris@87:     poly1d([3, 2, 1])
Chris@87: 
Chris@87:     which evaluates to:
Chris@87: 
Chris@87:     >>> p2(2.)
Chris@87:     17.0
Chris@87: 
Chris@87:     We can verify this, approximating the derivative with
Chris@87:     ``(f(x + h) - f(x))/h``:
Chris@87: 
Chris@87:     >>> (p(2. + 0.001) - p(2.)) / 0.001
Chris@87:     17.007000999997857
Chris@87: 
Chris@87:     The fourth-order derivative of a 3rd-order polynomial is zero:
Chris@87: 
Chris@87:     >>> np.polyder(p, 2)
Chris@87:     poly1d([6, 2])
Chris@87:     >>> np.polyder(p, 3)
Chris@87:     poly1d([6])
Chris@87:     >>> np.polyder(p, 4)
Chris@87:     poly1d([ 0.])
Chris@87: 
Chris@87:     """
Chris@87:     m = int(m)
Chris@87:     if m < 0:
Chris@87:         raise ValueError("Order of derivative must be positive (see polyint)")
Chris@87: 
Chris@87:     truepoly = isinstance(p, poly1d)
Chris@87:     p = NX.asarray(p)
Chris@87:     n = len(p) - 1
Chris@87:     y = p[:-1] * NX.arange(n, 0, -1)
Chris@87:     if m == 0:
Chris@87:         val = p
Chris@87:     else:
Chris@87:         val = polyder(y, m - 1)
Chris@87:     if truepoly:
Chris@87:         val = poly1d(val)
Chris@87:     return val
Chris@87: 
Chris@87: def polyfit(x, y, deg, rcond=None, full=False, w=None, cov=False):
Chris@87:     """
Chris@87:     Least squares polynomial fit.
Chris@87: 
Chris@87:     Fit a polynomial ``p(x) = p[0] * x**deg + ... + p[deg]`` of degree `deg`
Chris@87:     to points `(x, y)`. Returns a vector of coefficients `p` that minimises
Chris@87:     the squared error.
Chris@87: 
Chris@87:     Parameters
Chris@87:     ----------
Chris@87:     x : array_like, shape (M,)
Chris@87:         x-coordinates of the M sample points ``(x[i], y[i])``.
Chris@87:     y : array_like, shape (M,) or (M, K)
Chris@87:         y-coordinates of the sample points. Several data sets of sample
Chris@87:         points sharing the same x-coordinates can be fitted at once by
Chris@87:         passing in a 2D-array that contains one dataset per column.
Chris@87:     deg : int
Chris@87:         Degree of the fitting polynomial
Chris@87:     rcond : float, optional
Chris@87:         Relative condition number of the fit. Singular values smaller than
Chris@87:         this relative to the largest singular value will be ignored. The
Chris@87:         default value is len(x)*eps, where eps is the relative precision of
Chris@87:         the float type, about 2e-16 in most cases.
Chris@87:     full : bool, optional
Chris@87:         Switch determining nature of return value. When it is False (the
Chris@87:         default) just the coefficients are returned, when True diagnostic
Chris@87:         information from the singular value decomposition is also returned.
Chris@87:     w : array_like, shape (M,), optional
Chris@87:         weights to apply to the y-coordinates of the sample points.
Chris@87:     cov : bool, optional
Chris@87:         Return the estimate and the covariance matrix of the estimate
Chris@87:         If full is True, then cov is not returned.
Chris@87: 
Chris@87:     Returns
Chris@87:     -------
Chris@87:     p : ndarray, shape (M,) or (M, K)
Chris@87:         Polynomial coefficients, highest power first.  If `y` was 2-D, the
Chris@87:         coefficients for `k`-th data set are in ``p[:,k]``.
Chris@87: 
Chris@87:     residuals, rank, singular_values, rcond :
Chris@87:         Present only if `full` = True.  Residuals of the least-squares fit,
Chris@87:         the effective rank of the scaled Vandermonde coefficient matrix,
Chris@87:         its singular values, and the specified value of `rcond`. For more
Chris@87:         details, see `linalg.lstsq`.
Chris@87: 
Chris@87:     V : ndarray, shape (M,M) or (M,M,K)
Chris@87:         Present only if `full` = False and `cov`=True.  The covariance
Chris@87:         matrix of the polynomial coefficient estimates.  The diagonal of
Chris@87:         this matrix are the variance estimates for each coefficient.  If y
Chris@87:         is a 2-D array, then the covariance matrix for the `k`-th data set
Chris@87:         are in ``V[:,:,k]``
Chris@87: 
Chris@87: 
Chris@87:     Warns
Chris@87:     -----
Chris@87:     RankWarning
Chris@87:         The rank of the coefficient matrix in the least-squares fit is
Chris@87:         deficient. The warning is only raised if `full` = False.
Chris@87: 
Chris@87:         The warnings can be turned off by
Chris@87: 
Chris@87:         >>> import warnings
Chris@87:         >>> warnings.simplefilter('ignore', np.RankWarning)
Chris@87: 
Chris@87:     See Also
Chris@87:     --------
Chris@87:     polyval : Computes polynomial values.
Chris@87:     linalg.lstsq : Computes a least-squares fit.
Chris@87:     scipy.interpolate.UnivariateSpline : Computes spline fits.
Chris@87: 
Chris@87:     Notes
Chris@87:     -----
Chris@87:     The solution minimizes the squared error
Chris@87: 
Chris@87:     .. math ::
Chris@87:         E = \\sum_{j=0}^k |p(x_j) - y_j|^2
Chris@87: 
Chris@87:     in the equations::
Chris@87: 
Chris@87:         x[0]**n * p[0] + ... + x[0] * p[n-1] + p[n] = y[0]
Chris@87:         x[1]**n * p[0] + ... + x[1] * p[n-1] + p[n] = y[1]
Chris@87:         ...
Chris@87:         x[k]**n * p[0] + ... + x[k] * p[n-1] + p[n] = y[k]
Chris@87: 
Chris@87:     The coefficient matrix of the coefficients `p` is a Vandermonde matrix.
Chris@87: 
Chris@87:     `polyfit` issues a `RankWarning` when the least-squares fit is badly
Chris@87:     conditioned. This implies that the best fit is not well-defined due
Chris@87:     to numerical error. The results may be improved by lowering the polynomial
Chris@87:     degree or by replacing `x` by `x` - `x`.mean(). The `rcond` parameter
Chris@87:     can also be set to a value smaller than its default, but the resulting
Chris@87:     fit may be spurious: including contributions from the small singular
Chris@87:     values can add numerical noise to the result.
Chris@87: 
Chris@87:     Note that fitting polynomial coefficients is inherently badly conditioned
Chris@87:     when the degree of the polynomial is large or the interval of sample points
Chris@87:     is badly centered. The quality of the fit should always be checked in these
Chris@87:     cases. When polynomial fits are not satisfactory, splines may be a good
Chris@87:     alternative.
Chris@87: 
Chris@87:     References
Chris@87:     ----------
Chris@87:     .. [1] Wikipedia, "Curve fitting",
Chris@87:            http://en.wikipedia.org/wiki/Curve_fitting
Chris@87:     .. [2] Wikipedia, "Polynomial interpolation",
Chris@87:            http://en.wikipedia.org/wiki/Polynomial_interpolation
Chris@87: 
Chris@87:     Examples
Chris@87:     --------
Chris@87:     >>> x = np.array([0.0, 1.0, 2.0, 3.0,  4.0,  5.0])
Chris@87:     >>> y = np.array([0.0, 0.8, 0.9, 0.1, -0.8, -1.0])
Chris@87:     >>> z = np.polyfit(x, y, 3)
Chris@87:     >>> z
Chris@87:     array([ 0.08703704, -0.81349206,  1.69312169, -0.03968254])
Chris@87: 
Chris@87:     It is convenient to use `poly1d` objects for dealing with polynomials:
Chris@87: 
Chris@87:     >>> p = np.poly1d(z)
Chris@87:     >>> p(0.5)
Chris@87:     0.6143849206349179
Chris@87:     >>> p(3.5)
Chris@87:     -0.34732142857143039
Chris@87:     >>> p(10)
Chris@87:     22.579365079365115
Chris@87: 
Chris@87:     High-order polynomials may oscillate wildly:
Chris@87: 
Chris@87:     >>> p30 = np.poly1d(np.polyfit(x, y, 30))
Chris@87:     /... RankWarning: Polyfit may be poorly conditioned...
Chris@87:     >>> p30(4)
Chris@87:     -0.80000000000000204
Chris@87:     >>> p30(5)
Chris@87:     -0.99999999999999445
Chris@87:     >>> p30(4.5)
Chris@87:     -0.10547061179440398
Chris@87: 
Chris@87:     Illustration:
Chris@87: 
Chris@87:     >>> import matplotlib.pyplot as plt
Chris@87:     >>> xp = np.linspace(-2, 6, 100)
Chris@87:     >>> _ = plt.plot(x, y, '.', xp, p(xp), '-', xp, p30(xp), '--')
Chris@87:     >>> plt.ylim(-2,2)
Chris@87:     (-2, 2)
Chris@87:     >>> plt.show()
Chris@87: 
Chris@87:     """
Chris@87:     order = int(deg) + 1
Chris@87:     x = NX.asarray(x) + 0.0
Chris@87:     y = NX.asarray(y) + 0.0
Chris@87: 
Chris@87:     # check arguments.
Chris@87:     if deg < 0:
Chris@87:         raise ValueError("expected deg >= 0")
Chris@87:     if x.ndim != 1:
Chris@87:         raise TypeError("expected 1D vector for x")
Chris@87:     if x.size == 0:
Chris@87:         raise TypeError("expected non-empty vector for x")
Chris@87:     if y.ndim < 1 or y.ndim > 2:
Chris@87:         raise TypeError("expected 1D or 2D array for y")
Chris@87:     if x.shape[0] != y.shape[0]:
Chris@87:         raise TypeError("expected x and y to have same length")
Chris@87: 
Chris@87:     # set rcond
Chris@87:     if rcond is None:
Chris@87:         rcond = len(x)*finfo(x.dtype).eps
Chris@87: 
Chris@87:     # set up least squares equation for powers of x
Chris@87:     lhs = vander(x, order)
Chris@87:     rhs = y
Chris@87: 
Chris@87:     # apply weighting
Chris@87:     if w is not None:
Chris@87:         w = NX.asarray(w) + 0.0
Chris@87:         if w.ndim != 1:
Chris@87:             raise TypeError("expected a 1-d array for weights")
Chris@87:         if w.shape[0] != y.shape[0]:
Chris@87:             raise TypeError("expected w and y to have the same length")
Chris@87:         lhs *= w[:, NX.newaxis]
Chris@87:         if rhs.ndim == 2:
Chris@87:             rhs *= w[:, NX.newaxis]
Chris@87:         else:
Chris@87:             rhs *= w
Chris@87: 
Chris@87:     # scale lhs to improve condition number and solve
Chris@87:     scale = NX.sqrt((lhs*lhs).sum(axis=0))
Chris@87:     lhs /= scale
Chris@87:     c, resids, rank, s = lstsq(lhs, rhs, rcond)
Chris@87:     c = (c.T/scale).T  # broadcast scale coefficients
Chris@87: 
Chris@87:     # warn on rank reduction, which indicates an ill conditioned matrix
Chris@87:     if rank != order and not full:
Chris@87:         msg = "Polyfit may be poorly conditioned"
Chris@87:         warnings.warn(msg, RankWarning)
Chris@87: 
Chris@87:     if full:
Chris@87:         return c, resids, rank, s, rcond
Chris@87:     elif cov:
Chris@87:         Vbase = inv(dot(lhs.T, lhs))
Chris@87:         Vbase /= NX.outer(scale, scale)
Chris@87:         # Some literature ignores the extra -2.0 factor in the denominator, but
Chris@87:         #  it is included here because the covariance of Multivariate Student-T
Chris@87:         #  (which is implied by a Bayesian uncertainty analysis) includes it.
Chris@87:         #  Plus, it gives a slightly more conservative estimate of uncertainty.
Chris@87:         fac = resids / (len(x) - order - 2.0)
Chris@87:         if y.ndim == 1:
Chris@87:             return c, Vbase * fac
Chris@87:         else:
Chris@87:             return c, Vbase[:,:, NX.newaxis] * fac
Chris@87:     else:
Chris@87:         return c
Chris@87: 
Chris@87: 
Chris@87: def polyval(p, x):
Chris@87:     """
Chris@87:     Evaluate a polynomial at specific values.
Chris@87: 
Chris@87:     If `p` is of length N, this function returns the value:
Chris@87: 
Chris@87:         ``p[0]*x**(N-1) + p[1]*x**(N-2) + ... + p[N-2]*x + p[N-1]``
Chris@87: 
Chris@87:     If `x` is a sequence, then `p(x)` is returned for each element of `x`.
Chris@87:     If `x` is another polynomial then the composite polynomial `p(x(t))`
Chris@87:     is returned.
Chris@87: 
Chris@87:     Parameters
Chris@87:     ----------
Chris@87:     p : array_like or poly1d object
Chris@87:        1D array of polynomial coefficients (including coefficients equal
Chris@87:        to zero) from highest degree to the constant term, or an
Chris@87:        instance of poly1d.
Chris@87:     x : array_like or poly1d object
Chris@87:        A number, a 1D array of numbers, or an instance of poly1d, "at"
Chris@87:        which to evaluate `p`.
Chris@87: 
Chris@87:     Returns
Chris@87:     -------
Chris@87:     values : ndarray or poly1d
Chris@87:        If `x` is a poly1d instance, the result is the composition of the two
Chris@87:        polynomials, i.e., `x` is "substituted" in `p` and the simplified
Chris@87:        result is returned. In addition, the type of `x` - array_like or
Chris@87:        poly1d - governs the type of the output: `x` array_like => `values`
Chris@87:        array_like, `x` a poly1d object => `values` is also.
Chris@87: 
Chris@87:     See Also
Chris@87:     --------
Chris@87:     poly1d: A polynomial class.
Chris@87: 
Chris@87:     Notes
Chris@87:     -----
Chris@87:     Horner's scheme [1]_ is used to evaluate the polynomial. Even so,
Chris@87:     for polynomials of high degree the values may be inaccurate due to
Chris@87:     rounding errors. Use carefully.
Chris@87: 
Chris@87:     References
Chris@87:     ----------
Chris@87:     .. [1] I. N. Bronshtein, K. A. Semendyayev, and K. A. Hirsch (Eng.
Chris@87:        trans. Ed.), *Handbook of Mathematics*, New York, Van Nostrand
Chris@87:        Reinhold Co., 1985, pg. 720.
Chris@87: 
Chris@87:     Examples
Chris@87:     --------
Chris@87:     >>> np.polyval([3,0,1], 5)  # 3 * 5**2 + 0 * 5**1 + 1
Chris@87:     76
Chris@87:     >>> np.polyval([3,0,1], np.poly1d(5))
Chris@87:     poly1d([ 76.])
Chris@87:     >>> np.polyval(np.poly1d([3,0,1]), 5)
Chris@87:     76
Chris@87:     >>> np.polyval(np.poly1d([3,0,1]), np.poly1d(5))
Chris@87:     poly1d([ 76.])
Chris@87: 
Chris@87:     """
Chris@87:     p = NX.asarray(p)
Chris@87:     if isinstance(x, poly1d):
Chris@87:         y = 0
Chris@87:     else:
Chris@87:         x = NX.asarray(x)
Chris@87:         y = NX.zeros_like(x)
Chris@87:     for i in range(len(p)):
Chris@87:         y = x * y + p[i]
Chris@87:     return y
Chris@87: 
Chris@87: def polyadd(a1, a2):
Chris@87:     """
Chris@87:     Find the sum of two polynomials.
Chris@87: 
Chris@87:     Returns the polynomial resulting from the sum of two input polynomials.
Chris@87:     Each input must be either a poly1d object or a 1D sequence of polynomial
Chris@87:     coefficients, from highest to lowest degree.
Chris@87: 
Chris@87:     Parameters
Chris@87:     ----------
Chris@87:     a1, a2 : array_like or poly1d object
Chris@87:         Input polynomials.
Chris@87: 
Chris@87:     Returns
Chris@87:     -------
Chris@87:     out : ndarray or poly1d object
Chris@87:         The sum of the inputs. If either input is a poly1d object, then the
Chris@87:         output is also a poly1d object. Otherwise, it is a 1D array of
Chris@87:         polynomial coefficients from highest to lowest degree.
Chris@87: 
Chris@87:     See Also
Chris@87:     --------
Chris@87:     poly1d : A one-dimensional polynomial class.
Chris@87:     poly, polyadd, polyder, polydiv, polyfit, polyint, polysub, polyval
Chris@87: 
Chris@87:     Examples
Chris@87:     --------
Chris@87:     >>> np.polyadd([1, 2], [9, 5, 4])
Chris@87:     array([9, 6, 6])
Chris@87: 
Chris@87:     Using poly1d objects:
Chris@87: 
Chris@87:     >>> p1 = np.poly1d([1, 2])
Chris@87:     >>> p2 = np.poly1d([9, 5, 4])
Chris@87:     >>> print p1
Chris@87:     1 x + 2
Chris@87:     >>> print p2
Chris@87:        2
Chris@87:     9 x + 5 x + 4
Chris@87:     >>> print np.polyadd(p1, p2)
Chris@87:        2
Chris@87:     9 x + 6 x + 6
Chris@87: 
Chris@87:     """
Chris@87:     truepoly = (isinstance(a1, poly1d) or isinstance(a2, poly1d))
Chris@87:     a1 = atleast_1d(a1)
Chris@87:     a2 = atleast_1d(a2)
Chris@87:     diff = len(a2) - len(a1)
Chris@87:     if diff == 0:
Chris@87:         val = a1 + a2
Chris@87:     elif diff > 0:
Chris@87:         zr = NX.zeros(diff, a1.dtype)
Chris@87:         val = NX.concatenate((zr, a1)) + a2
Chris@87:     else:
Chris@87:         zr = NX.zeros(abs(diff), a2.dtype)
Chris@87:         val = a1 + NX.concatenate((zr, a2))
Chris@87:     if truepoly:
Chris@87:         val = poly1d(val)
Chris@87:     return val
Chris@87: 
Chris@87: def polysub(a1, a2):
Chris@87:     """
Chris@87:     Difference (subtraction) of two polynomials.
Chris@87: 
Chris@87:     Given two polynomials `a1` and `a2`, returns ``a1 - a2``.
Chris@87:     `a1` and `a2` can be either array_like sequences of the polynomials'
Chris@87:     coefficients (including coefficients equal to zero), or `poly1d` objects.
Chris@87: 
Chris@87:     Parameters
Chris@87:     ----------
Chris@87:     a1, a2 : array_like or poly1d
Chris@87:         Minuend and subtrahend polynomials, respectively.
Chris@87: 
Chris@87:     Returns
Chris@87:     -------
Chris@87:     out : ndarray or poly1d
Chris@87:         Array or `poly1d` object of the difference polynomial's coefficients.
Chris@87: 
Chris@87:     See Also
Chris@87:     --------
Chris@87:     polyval, polydiv, polymul, polyadd
Chris@87: 
Chris@87:     Examples
Chris@87:     --------
Chris@87:     .. math:: (2 x^2 + 10 x - 2) - (3 x^2 + 10 x -4) = (-x^2 + 2)
Chris@87: 
Chris@87:     >>> np.polysub([2, 10, -2], [3, 10, -4])
Chris@87:     array([-1,  0,  2])
Chris@87: 
Chris@87:     """
Chris@87:     truepoly = (isinstance(a1, poly1d) or isinstance(a2, poly1d))
Chris@87:     a1 = atleast_1d(a1)
Chris@87:     a2 = atleast_1d(a2)
Chris@87:     diff = len(a2) - len(a1)
Chris@87:     if diff == 0:
Chris@87:         val = a1 - a2
Chris@87:     elif diff > 0:
Chris@87:         zr = NX.zeros(diff, a1.dtype)
Chris@87:         val = NX.concatenate((zr, a1)) - a2
Chris@87:     else:
Chris@87:         zr = NX.zeros(abs(diff), a2.dtype)
Chris@87:         val = a1 - NX.concatenate((zr, a2))
Chris@87:     if truepoly:
Chris@87:         val = poly1d(val)
Chris@87:     return val
Chris@87: 
Chris@87: 
Chris@87: def polymul(a1, a2):
Chris@87:     """
Chris@87:     Find the product of two polynomials.
Chris@87: 
Chris@87:     Finds the polynomial resulting from the multiplication of the two input
Chris@87:     polynomials. Each input must be either a poly1d object or a 1D sequence
Chris@87:     of polynomial coefficients, from highest to lowest degree.
Chris@87: 
Chris@87:     Parameters
Chris@87:     ----------
Chris@87:     a1, a2 : array_like or poly1d object
Chris@87:         Input polynomials.
Chris@87: 
Chris@87:     Returns
Chris@87:     -------
Chris@87:     out : ndarray or poly1d object
Chris@87:         The polynomial resulting from the multiplication of the inputs. If
Chris@87:         either inputs is a poly1d object, then the output is also a poly1d
Chris@87:         object. Otherwise, it is a 1D array of polynomial coefficients from
Chris@87:         highest to lowest degree.
Chris@87: 
Chris@87:     See Also
Chris@87:     --------
Chris@87:     poly1d : A one-dimensional polynomial class.
Chris@87:     poly, polyadd, polyder, polydiv, polyfit, polyint, polysub,
Chris@87:     polyval
Chris@87:     convolve : Array convolution. Same output as polymul, but has parameter
Chris@87:                for overlap mode.
Chris@87: 
Chris@87:     Examples
Chris@87:     --------
Chris@87:     >>> np.polymul([1, 2, 3], [9, 5, 1])
Chris@87:     array([ 9, 23, 38, 17,  3])
Chris@87: 
Chris@87:     Using poly1d objects:
Chris@87: 
Chris@87:     >>> p1 = np.poly1d([1, 2, 3])
Chris@87:     >>> p2 = np.poly1d([9, 5, 1])
Chris@87:     >>> print p1
Chris@87:        2
Chris@87:     1 x + 2 x + 3
Chris@87:     >>> print p2
Chris@87:        2
Chris@87:     9 x + 5 x + 1
Chris@87:     >>> print np.polymul(p1, p2)
Chris@87:        4      3      2
Chris@87:     9 x + 23 x + 38 x + 17 x + 3
Chris@87: 
Chris@87:     """
Chris@87:     truepoly = (isinstance(a1, poly1d) or isinstance(a2, poly1d))
Chris@87:     a1, a2 = poly1d(a1), poly1d(a2)
Chris@87:     val = NX.convolve(a1, a2)
Chris@87:     if truepoly:
Chris@87:         val = poly1d(val)
Chris@87:     return val
Chris@87: 
Chris@87: def polydiv(u, v):
Chris@87:     """
Chris@87:     Returns the quotient and remainder of polynomial division.
Chris@87: 
Chris@87:     The input arrays are the coefficients (including any coefficients
Chris@87:     equal to zero) of the "numerator" (dividend) and "denominator"
Chris@87:     (divisor) polynomials, respectively.
Chris@87: 
Chris@87:     Parameters
Chris@87:     ----------
Chris@87:     u : array_like or poly1d
Chris@87:         Dividend polynomial's coefficients.
Chris@87: 
Chris@87:     v : array_like or poly1d
Chris@87:         Divisor polynomial's coefficients.
Chris@87: 
Chris@87:     Returns
Chris@87:     -------
Chris@87:     q : ndarray
Chris@87:         Coefficients, including those equal to zero, of the quotient.
Chris@87:     r : ndarray
Chris@87:         Coefficients, including those equal to zero, of the remainder.
Chris@87: 
Chris@87:     See Also
Chris@87:     --------
Chris@87:     poly, polyadd, polyder, polydiv, polyfit, polyint, polymul, polysub,
Chris@87:     polyval
Chris@87: 
Chris@87:     Notes
Chris@87:     -----
Chris@87:     Both `u` and `v` must be 0-d or 1-d (ndim = 0 or 1), but `u.ndim` need
Chris@87:     not equal `v.ndim`. In other words, all four possible combinations -
Chris@87:     ``u.ndim = v.ndim = 0``, ``u.ndim = v.ndim = 1``,
Chris@87:     ``u.ndim = 1, v.ndim = 0``, and ``u.ndim = 0, v.ndim = 1`` - work.
Chris@87: 
Chris@87:     Examples
Chris@87:     --------
Chris@87:     .. math:: \\frac{3x^2 + 5x + 2}{2x + 1} = 1.5x + 1.75, remainder 0.25
Chris@87: 
Chris@87:     >>> x = np.array([3.0, 5.0, 2.0])
Chris@87:     >>> y = np.array([2.0, 1.0])
Chris@87:     >>> np.polydiv(x, y)
Chris@87:     (array([ 1.5 ,  1.75]), array([ 0.25]))
Chris@87: 
Chris@87:     """
Chris@87:     truepoly = (isinstance(u, poly1d) or isinstance(u, poly1d))
Chris@87:     u = atleast_1d(u) + 0.0
Chris@87:     v = atleast_1d(v) + 0.0
Chris@87:     # w has the common type
Chris@87:     w = u[0] + v[0]
Chris@87:     m = len(u) - 1
Chris@87:     n = len(v) - 1
Chris@87:     scale = 1. / v[0]
Chris@87:     q = NX.zeros((max(m - n + 1, 1),), w.dtype)
Chris@87:     r = u.copy()
Chris@87:     for k in range(0, m-n+1):
Chris@87:         d = scale * r[k]
Chris@87:         q[k] = d
Chris@87:         r[k:k+n+1] -= d*v
Chris@87:     while NX.allclose(r[0], 0, rtol=1e-14) and (r.shape[-1] > 1):
Chris@87:         r = r[1:]
Chris@87:     if truepoly:
Chris@87:         return poly1d(q), poly1d(r)
Chris@87:     return q, r
Chris@87: 
Chris@87: _poly_mat = re.compile(r"[*][*]([0-9]*)")
Chris@87: def _raise_power(astr, wrap=70):
Chris@87:     n = 0
Chris@87:     line1 = ''
Chris@87:     line2 = ''
Chris@87:     output = ' '
Chris@87:     while True:
Chris@87:         mat = _poly_mat.search(astr, n)
Chris@87:         if mat is None:
Chris@87:             break
Chris@87:         span = mat.span()
Chris@87:         power = mat.groups()[0]
Chris@87:         partstr = astr[n:span[0]]
Chris@87:         n = span[1]
Chris@87:         toadd2 = partstr + ' '*(len(power)-1)
Chris@87:         toadd1 = ' '*(len(partstr)-1) + power
Chris@87:         if ((len(line2) + len(toadd2) > wrap) or
Chris@87:                 (len(line1) + len(toadd1) > wrap)):
Chris@87:             output += line1 + "\n" + line2 + "\n "
Chris@87:             line1 = toadd1
Chris@87:             line2 = toadd2
Chris@87:         else:
Chris@87:             line2 += partstr + ' '*(len(power)-1)
Chris@87:             line1 += ' '*(len(partstr)-1) + power
Chris@87:     output += line1 + "\n" + line2
Chris@87:     return output + astr[n:]
Chris@87: 
Chris@87: 
Chris@87: class poly1d(object):
Chris@87:     """
Chris@87:     A one-dimensional polynomial class.
Chris@87: 
Chris@87:     A convenience class, used to encapsulate "natural" operations on
Chris@87:     polynomials so that said operations may take on their customary
Chris@87:     form in code (see Examples).
Chris@87: 
Chris@87:     Parameters
Chris@87:     ----------
Chris@87:     c_or_r : array_like
Chris@87:         The polynomial's coefficients, in decreasing powers, or if
Chris@87:         the value of the second parameter is True, the polynomial's
Chris@87:         roots (values where the polynomial evaluates to 0).  For example,
Chris@87:         ``poly1d([1, 2, 3])`` returns an object that represents
Chris@87:         :math:`x^2 + 2x + 3`, whereas ``poly1d([1, 2, 3], True)`` returns
Chris@87:         one that represents :math:`(x-1)(x-2)(x-3) = x^3 - 6x^2 + 11x -6`.
Chris@87:     r : bool, optional
Chris@87:         If True, `c_or_r` specifies the polynomial's roots; the default
Chris@87:         is False.
Chris@87:     variable : str, optional
Chris@87:         Changes the variable used when printing `p` from `x` to `variable`
Chris@87:         (see Examples).
Chris@87: 
Chris@87:     Examples
Chris@87:     --------
Chris@87:     Construct the polynomial :math:`x^2 + 2x + 3`:
Chris@87: 
Chris@87:     >>> p = np.poly1d([1, 2, 3])
Chris@87:     >>> print np.poly1d(p)
Chris@87:        2
Chris@87:     1 x + 2 x + 3
Chris@87: 
Chris@87:     Evaluate the polynomial at :math:`x = 0.5`:
Chris@87: 
Chris@87:     >>> p(0.5)
Chris@87:     4.25
Chris@87: 
Chris@87:     Find the roots:
Chris@87: 
Chris@87:     >>> p.r
Chris@87:     array([-1.+1.41421356j, -1.-1.41421356j])
Chris@87:     >>> p(p.r)
Chris@87:     array([ -4.44089210e-16+0.j,  -4.44089210e-16+0.j])
Chris@87: 
Chris@87:     These numbers in the previous line represent (0, 0) to machine precision
Chris@87: 
Chris@87:     Show the coefficients:
Chris@87: 
Chris@87:     >>> p.c
Chris@87:     array([1, 2, 3])
Chris@87: 
Chris@87:     Display the order (the leading zero-coefficients are removed):
Chris@87: 
Chris@87:     >>> p.order
Chris@87:     2
Chris@87: 
Chris@87:     Show the coefficient of the k-th power in the polynomial
Chris@87:     (which is equivalent to ``p.c[-(i+1)]``):
Chris@87: 
Chris@87:     >>> p[1]
Chris@87:     2
Chris@87: 
Chris@87:     Polynomials can be added, subtracted, multiplied, and divided
Chris@87:     (returns quotient and remainder):
Chris@87: 
Chris@87:     >>> p * p
Chris@87:     poly1d([ 1,  4, 10, 12,  9])
Chris@87: 
Chris@87:     >>> (p**3 + 4) / p
Chris@87:     (poly1d([  1.,   4.,  10.,  12.,   9.]), poly1d([ 4.]))
Chris@87: 
Chris@87:     ``asarray(p)`` gives the coefficient array, so polynomials can be
Chris@87:     used in all functions that accept arrays:
Chris@87: 
Chris@87:     >>> p**2 # square of polynomial
Chris@87:     poly1d([ 1,  4, 10, 12,  9])
Chris@87: 
Chris@87:     >>> np.square(p) # square of individual coefficients
Chris@87:     array([1, 4, 9])
Chris@87: 
Chris@87:     The variable used in the string representation of `p` can be modified,
Chris@87:     using the `variable` parameter:
Chris@87: 
Chris@87:     >>> p = np.poly1d([1,2,3], variable='z')
Chris@87:     >>> print p
Chris@87:        2
Chris@87:     1 z + 2 z + 3
Chris@87: 
Chris@87:     Construct a polynomial from its roots:
Chris@87: 
Chris@87:     >>> np.poly1d([1, 2], True)
Chris@87:     poly1d([ 1, -3,  2])
Chris@87: 
Chris@87:     This is the same polynomial as obtained by:
Chris@87: 
Chris@87:     >>> np.poly1d([1, -1]) * np.poly1d([1, -2])
Chris@87:     poly1d([ 1, -3,  2])
Chris@87: 
Chris@87:     """
Chris@87:     coeffs = None
Chris@87:     order = None
Chris@87:     variable = None
Chris@87:     __hash__ = None
Chris@87: 
Chris@87:     def __init__(self, c_or_r, r=0, variable=None):
Chris@87:         if isinstance(c_or_r, poly1d):
Chris@87:             for key in c_or_r.__dict__.keys():
Chris@87:                 self.__dict__[key] = c_or_r.__dict__[key]
Chris@87:             if variable is not None:
Chris@87:                 self.__dict__['variable'] = variable
Chris@87:             return
Chris@87:         if r:
Chris@87:             c_or_r = poly(c_or_r)
Chris@87:         c_or_r = atleast_1d(c_or_r)
Chris@87:         if len(c_or_r.shape) > 1:
Chris@87:             raise ValueError("Polynomial must be 1d only.")
Chris@87:         c_or_r = trim_zeros(c_or_r, trim='f')
Chris@87:         if len(c_or_r) == 0:
Chris@87:             c_or_r = NX.array([0.])
Chris@87:         self.__dict__['coeffs'] = c_or_r
Chris@87:         self.__dict__['order'] = len(c_or_r) - 1
Chris@87:         if variable is None:
Chris@87:             variable = 'x'
Chris@87:         self.__dict__['variable'] = variable
Chris@87: 
Chris@87:     def __array__(self, t=None):
Chris@87:         if t:
Chris@87:             return NX.asarray(self.coeffs, t)
Chris@87:         else:
Chris@87:             return NX.asarray(self.coeffs)
Chris@87: 
Chris@87:     def __repr__(self):
Chris@87:         vals = repr(self.coeffs)
Chris@87:         vals = vals[6:-1]
Chris@87:         return "poly1d(%s)" % vals
Chris@87: 
Chris@87:     def __len__(self):
Chris@87:         return self.order
Chris@87: 
Chris@87:     def __str__(self):
Chris@87:         thestr = "0"
Chris@87:         var = self.variable
Chris@87: 
Chris@87:         # Remove leading zeros
Chris@87:         coeffs = self.coeffs[NX.logical_or.accumulate(self.coeffs != 0)]
Chris@87:         N = len(coeffs)-1
Chris@87: 
Chris@87:         def fmt_float(q):
Chris@87:             s = '%.4g' % q
Chris@87:             if s.endswith('.0000'):
Chris@87:                 s = s[:-5]
Chris@87:             return s
Chris@87: 
Chris@87:         for k in range(len(coeffs)):
Chris@87:             if not iscomplex(coeffs[k]):
Chris@87:                 coefstr = fmt_float(real(coeffs[k]))
Chris@87:             elif real(coeffs[k]) == 0:
Chris@87:                 coefstr = '%sj' % fmt_float(imag(coeffs[k]))
Chris@87:             else:
Chris@87:                 coefstr = '(%s + %sj)' % (fmt_float(real(coeffs[k])),
Chris@87:                                           fmt_float(imag(coeffs[k])))
Chris@87: 
Chris@87:             power = (N-k)
Chris@87:             if power == 0:
Chris@87:                 if coefstr != '0':
Chris@87:                     newstr = '%s' % (coefstr,)
Chris@87:                 else:
Chris@87:                     if k == 0:
Chris@87:                         newstr = '0'
Chris@87:                     else:
Chris@87:                         newstr = ''
Chris@87:             elif power == 1:
Chris@87:                 if coefstr == '0':
Chris@87:                     newstr = ''
Chris@87:                 elif coefstr == 'b':
Chris@87:                     newstr = var
Chris@87:                 else:
Chris@87:                     newstr = '%s %s' % (coefstr, var)
Chris@87:             else:
Chris@87:                 if coefstr == '0':
Chris@87:                     newstr = ''
Chris@87:                 elif coefstr == 'b':
Chris@87:                     newstr = '%s**%d' % (var, power,)
Chris@87:                 else:
Chris@87:                     newstr = '%s %s**%d' % (coefstr, var, power)
Chris@87: 
Chris@87:             if k > 0:
Chris@87:                 if newstr != '':
Chris@87:                     if newstr.startswith('-'):
Chris@87:                         thestr = "%s - %s" % (thestr, newstr[1:])
Chris@87:                     else:
Chris@87:                         thestr = "%s + %s" % (thestr, newstr)
Chris@87:             else:
Chris@87:                 thestr = newstr
Chris@87:         return _raise_power(thestr)
Chris@87: 
Chris@87:     def __call__(self, val):
Chris@87:         return polyval(self.coeffs, val)
Chris@87: 
Chris@87:     def __neg__(self):
Chris@87:         return poly1d(-self.coeffs)
Chris@87: 
Chris@87:     def __pos__(self):
Chris@87:         return self
Chris@87: 
Chris@87:     def __mul__(self, other):
Chris@87:         if isscalar(other):
Chris@87:             return poly1d(self.coeffs * other)
Chris@87:         else:
Chris@87:             other = poly1d(other)
Chris@87:             return poly1d(polymul(self.coeffs, other.coeffs))
Chris@87: 
Chris@87:     def __rmul__(self, other):
Chris@87:         if isscalar(other):
Chris@87:             return poly1d(other * self.coeffs)
Chris@87:         else:
Chris@87:             other = poly1d(other)
Chris@87:             return poly1d(polymul(self.coeffs, other.coeffs))
Chris@87: 
Chris@87:     def __add__(self, other):
Chris@87:         other = poly1d(other)
Chris@87:         return poly1d(polyadd(self.coeffs, other.coeffs))
Chris@87: 
Chris@87:     def __radd__(self, other):
Chris@87:         other = poly1d(other)
Chris@87:         return poly1d(polyadd(self.coeffs, other.coeffs))
Chris@87: 
Chris@87:     def __pow__(self, val):
Chris@87:         if not isscalar(val) or int(val) != val or val < 0:
Chris@87:             raise ValueError("Power to non-negative integers only.")
Chris@87:         res = [1]
Chris@87:         for _ in range(val):
Chris@87:             res = polymul(self.coeffs, res)
Chris@87:         return poly1d(res)
Chris@87: 
Chris@87:     def __sub__(self, other):
Chris@87:         other = poly1d(other)
Chris@87:         return poly1d(polysub(self.coeffs, other.coeffs))
Chris@87: 
Chris@87:     def __rsub__(self, other):
Chris@87:         other = poly1d(other)
Chris@87:         return poly1d(polysub(other.coeffs, self.coeffs))
Chris@87: 
Chris@87:     def __div__(self, other):
Chris@87:         if isscalar(other):
Chris@87:             return poly1d(self.coeffs/other)
Chris@87:         else:
Chris@87:             other = poly1d(other)
Chris@87:             return polydiv(self, other)
Chris@87: 
Chris@87:     __truediv__ = __div__
Chris@87: 
Chris@87:     def __rdiv__(self, other):
Chris@87:         if isscalar(other):
Chris@87:             return poly1d(other/self.coeffs)
Chris@87:         else:
Chris@87:             other = poly1d(other)
Chris@87:             return polydiv(other, self)
Chris@87: 
Chris@87:     __rtruediv__ = __rdiv__
Chris@87: 
Chris@87:     def __eq__(self, other):
Chris@87:         if self.coeffs.shape != other.coeffs.shape:
Chris@87:             return False
Chris@87:         return (self.coeffs == other.coeffs).all()
Chris@87: 
Chris@87:     def __ne__(self, other):
Chris@87:         return not self.__eq__(other)
Chris@87: 
Chris@87:     def __setattr__(self, key, val):
Chris@87:         raise ValueError("Attributes cannot be changed this way.")
Chris@87: 
Chris@87:     def __getattr__(self, key):
Chris@87:         if key in ['r', 'roots']:
Chris@87:             return roots(self.coeffs)
Chris@87:         elif key in ['c', 'coef', 'coefficients']:
Chris@87:             return self.coeffs
Chris@87:         elif key in ['o']:
Chris@87:             return self.order
Chris@87:         else:
Chris@87:             try:
Chris@87:                 return self.__dict__[key]
Chris@87:             except KeyError:
Chris@87:                 raise AttributeError(
Chris@87:                     "'%s' has no attribute '%s'" % (self.__class__, key))
Chris@87: 
Chris@87:     def __getitem__(self, val):
Chris@87:         ind = self.order - val
Chris@87:         if val > self.order:
Chris@87:             return 0
Chris@87:         if val < 0:
Chris@87:             return 0
Chris@87:         return self.coeffs[ind]
Chris@87: 
Chris@87:     def __setitem__(self, key, val):
Chris@87:         ind = self.order - key
Chris@87:         if key < 0:
Chris@87:             raise ValueError("Does not support negative powers.")
Chris@87:         if key > self.order:
Chris@87:             zr = NX.zeros(key-self.order, self.coeffs.dtype)
Chris@87:             self.__dict__['coeffs'] = NX.concatenate((zr, self.coeffs))
Chris@87:             self.__dict__['order'] = key
Chris@87:             ind = 0
Chris@87:         self.__dict__['coeffs'][ind] = val
Chris@87:         return
Chris@87: 
Chris@87:     def __iter__(self):
Chris@87:         return iter(self.coeffs)
Chris@87: 
Chris@87:     def integ(self, m=1, k=0):
Chris@87:         """
Chris@87:         Return an antiderivative (indefinite integral) of this polynomial.
Chris@87: 
Chris@87:         Refer to `polyint` for full documentation.
Chris@87: 
Chris@87:         See Also
Chris@87:         --------
Chris@87:         polyint : equivalent function
Chris@87: 
Chris@87:         """
Chris@87:         return poly1d(polyint(self.coeffs, m=m, k=k))
Chris@87: 
Chris@87:     def deriv(self, m=1):
Chris@87:         """
Chris@87:         Return a derivative of this polynomial.
Chris@87: 
Chris@87:         Refer to `polyder` for full documentation.
Chris@87: 
Chris@87:         See Also
Chris@87:         --------
Chris@87:         polyder : equivalent function
Chris@87: 
Chris@87:         """
Chris@87:         return poly1d(polyder(self.coeffs, m=m))
Chris@87: 
Chris@87: # Stuff to do on module import
Chris@87: 
Chris@87: warnings.simplefilter('always', RankWarning)