ENH: Modify the various polynomial series so that multidimensional coefficient arrays can be used. Add functions for evaluation of 2D and 3D polynomial series evaluated either on a specified set of points or on a cartesian product of 1D points.

The new functions have names polyval2d, polygrid2d, polyval3d, and
polygrid3d, where 'poly' can be replaced by any of 'leg', 'cheb', 'lag',
'herm', or 'herme'. These additional functions should cover the common
multidimensional cases and provide examples for anyone who wants to go to
higher dimensions.

1 files changed, 230 insertions, 22 deletions

diff --git a/numpy/polynomial/polynomial.py b/numpy/polynomial/polynomial.py
index a0716b9a5..a0c9f903a 100644
--- a/numpy/polynomial/polynomial.py
+++ b/numpy/polynomial/polynomial.py
@@ -22,6 +22,10 @@ Arithmetic
 - `polydiv` -- divide one polynomial by another.
 - `polypow` -- raise a polynomial to an positive integer power
 - `polyval` -- evaluate a polynomial at given points.
+- `polyval2d` -- evaluate a 2D polynomial at given points.
+- `polyval3d` -- evaluate a 3D polynomial at given points.
+- `polygrid2d` -- evaluate a 2D polynomial on a Cartesian product.
+- `polygrid3d` -- evaluate a 3D polynomial on a Cartesian product.
 
 Calculus
 --------
@@ -50,8 +54,9 @@ from __future__ import division
 
 __all__ = ['polyzero', 'polyone', 'polyx', 'polydomain', 'polyline',
     'polyadd', 'polysub', 'polymulx', 'polymul', 'polydiv', 'polypow',
-    'polyval', 'polyder', 'polyint', 'polyfromroots', 'polyvander',
-    'polyfit', 'polytrim', 'polyroots', 'Polynomial']
+    'polyval', 'polyval2d', 'polyval3d', 'polygrid2d', 'polygrid3d',
+    'polyder', 'polyint', 'polyfromroots', 'polyvander', 'polyfit',
+    'polytrim', 'polyroots', 'Polynomial']
 
 import numpy as np
 import numpy.linalg as la
@@ -615,52 +620,255 @@ def polyint(cs, m=1, k=[], lbnd=0, scl=1):
     return cs
 
 
-def polyval(x, cs):
+def polyval(x, c, tensor=True):
     """
     Evaluate a polynomial.
 
-    If `cs` is of length `n`, this function returns :
+    If `c` is of length ``n + 1``, this function returns the value:
 
-    ``p(x) = cs[0] + cs[1]*x + ... + cs[n-1]*x**(n-1)``
+    ``p(x) = c[0] + c[1]*x + ... + c[n]*x**(n)``
 
-    If x is a sequence or array then p(x) will have the same shape as x.
-    If r is a ring_like object that supports multiplication and addition
-    by the values in `cs`, then an object of the same type is returned.
+    If `x` is a sequence or array and `c` is 1 dimensional, then ``p(x)``
+    will have the same shape as `x`.  If `x` is a algebra_like object that
+    supports multiplication and addition with itself and the values in `c`,
+    then an object of the same type is returned.
+
+    In the case where c is multidimensional, the shape of the result
+    depends on the value of `tensor`. If tensor is true the shape of the
+    return will be ``c.shape[1:] + x.shape``, where the shape of a scalar
+    is the empty tuple. If tensor is false the shape is ``c.shape[1:]`` if
+    `x` is broadcast compatible with that.
+
+    If there are trailing zeros in the coefficients they still take part in
+    the evaluation, so they should be avoided if efficiency is a concern.
 
     Parameters
     ----------
-    x : array_like, ring_like
+    x : array_like, algebra_like
         If x is a list or tuple, it is converted to an ndarray. Otherwise
-        it must support addition and multiplication with itself and the
-        elements of `cs`.
-    cs : array_like
-        1-d array of Chebyshev coefficients ordered from low to high.
+        it is left unchanged and if it isn't an ndarray it is treated as a
+        scalar. In either case, `x` or any element of an ndarray must
+        support addition and multiplication with itself and the elements of
+        `c`.
+    c : array_like
+        Array of coefficients ordered so that the coefficients for terms of
+        degree n are contained in ``c[n]``. If `c` is multidimesional the
+        remaining indices enumerate multiple polynomials. In the two
+        dimensional case the coefficients may be thought of as stored in
+        the columns of `c`.
+    tensor : boolean, optional
+        If true, the coefficient array shape is extended with ones on the
+        right, one for each dimension of `x`. Scalars are treated as having
+        dimension 0 for this action. The effect is that every column of
+        coefficients in `c` is evaluated for every value in `x`. If False,
+        the `x` are broadcast over the columns of `c` in the usual way.
+        This gives some flexibility in evaluations in the multidimensional
+        case. The default value it ``True``.
 
     Returns
     -------
-    values : ndarray
-        The return array has the same shape as `x`.
+    values : ndarray, algebra_like
+        The shape of the return value is described above.
 
     See Also
     --------
-    polyfit
+    polyval2d, polygrid2d, polyval3d, polygrid3d
 
     Notes
     -----
     The evaluation uses Horner's method.
 
+    Examples
+    --------
+    >>> from numpy.polynomial.polynomial import polyval
+    >>> polyval(1, [1,2,3])
+    6.0
+    >>> a = np.arange(4).reshape(2,2)
+    >>> a
+    array([[[  0,   1],
+            [  2,   3]],
+    >>> polyval(a, [1,2,3])
+    array([[  1.,   6.],
+           [ 17.,  34.]])
+    >>> c = np.arange(4).reshape(2,2)
+    >>> c
+    array([[[  0,   1],
+            [  2,   3]],
+    >>> polyval([1,2], c, tensor=True)
+    array([[ 2.,  4.],
+           [ 4.,  7.]])
+    >>> polyval([1,2], c, tensor=False)
+    array([ 2.,  7.])
+
     """
-    # cs is a trimmed copy
-    [cs] = pu.as_series([cs])
-    if isinstance(x, tuple) or isinstance(x, list) :
+    c = np.array(c, ndmin=1, copy=0)
+    if c.dtype.char not in 'efdgFDGO':
+        c = c + 0.0
+    if isinstance(x, (tuple, list)):
         x = np.asarray(x)
+    if isinstance(x, np.ndarray) and tensor:
+       c = c.reshape(c.shape + (1,)*x.ndim)
 
-    c0 = cs[-1] + x*0
-    for i in range(2, len(cs) + 1) :
-        c0 = cs[-i] + c0*x
+    c0 = c[-1] + x*0
+    for i in range(2, len(c) + 1) :
+        c0 = c[-i] + c0*x
     return c0
 
 
+def polyval2d(x, y, c):
+    """
+    Evaluate 2D polynomials at points (x,y).
+
+    This function returns the values:
+
+    ``p(x,y) = \sum_{i,j} c[i,j] * x^i * y^j``
+
+    Parameters
+    ----------
+    x,y : array_like, algebra_like
+        The two dimensional polynomial is evaluated at the points
+        ``(x,y)``, where `x` and `y` must have the same shape. If `x` or
+        `y` is a list or tuple, it is first converted to an ndarray,
+        otherwise it is left unchanged and if it isn't an ndarray it is
+        treated as a scalar. See `polyval` for explanation of algebra_like.
+    c : array_like
+        Array of coefficients ordered so that the coefficients for terms of
+        degree i,j are contained in ``c[i,j]``. If `c` has dimension
+        greater than 2 the remaining indices enumerate multiple sets of
+        coefficients.
+
+    Returns
+    -------
+    values : ndarray, algebra_like
+        The values of the two dimensional polynomial at points formed with
+        pairs of corresponding values from `x` and `y`.
+
+    See Also
+    --------
+    polyval, polygrid2d, polyval3d, polygrid3d
+
+    """
+    return polyval(y, polyval(x, c), False)
+
+
+def polygrid2d(x, y, c):
+    """
+    Evaluate 2D polynomials on the Cartesion product of x,y.
+
+    This function returns the values:
+
+    ``p(a,b) = \sum_{i,j} c[i,j] * a^i * b^j``
+
+    where the points ``(a,b)`` consist of all pairs of points formed by
+    taking ``a`` from `x` and ``b`` from `y`. The resulting points form a
+    grid with `x` in the first dimension and `y` in the second.
+
+    Parameters
+    ----------
+    x,y : array_like, algebra_like
+        The two dimensional polynomial is evaluated at the points in the
+        Cartesian product of `x` and `y`.  If `x` or `y` is a list or
+        tuple, it is first converted to an ndarray, otherwise it is left
+        unchanged and if it isn't an ndarray it is treated as a scalar. See
+        `polyval` for explanation of algebra_like.
+    c : array_like
+        Array of coefficients ordered so that the coefficients for terms of
+        degree i,j are contained in ``c[i,j]``. If `c` has dimension
+        greater than 2 the remaining indices enumerate multiple sets of
+        coefficients.
+
+    Returns
+    -------
+    values : ndarray, algebra_like
+        The values of the two dimensional polynomial at points in the Cartesion
+        product of `x` and `y`.
+
+    See Also
+    --------
+    polyval, polyval2d, polyval3d, polygrid3d
+
+    """
+    return polyval(y, polyval(x, c))
+
+
+def polyval3d(x, y, z, c):
+    """
+    Evaluate 3D polynomials at points (x,y,z).
+
+    This function returns the values:
+
+    ``p(x,y,z) = \sum_{i,j,k} c[i,j,k] * x^i * y^j * z^k``
+
+    Parameters
+    ----------
+    x,y,z : array_like, algebra_like
+        The three dimensional polynomial is evaluated at the points
+        ``(x,y,z)``, where `x`, `y`, and `z` must have the same shape.  If
+        any of `x`, `y`, or `z` is a list or tuple, it is first converted
+        to an ndarray, otherwise it is left unchanged and if it isn't an
+        ndarray it is  treated as a scalar. See `polyval` for an
+        explanation of algebra_like.
+    c : array_like
+        Array of coefficients ordered so that the coefficients for terms of
+        degree i, j are contained in ``c[i,j,k]``. If `c` has dimension
+        greater than 3 the remaining indices enumerate multiple sets of
+        coefficients.
+
+    Returns
+    -------
+    values : ndarray, algebra_like
+        The values of the multidimension polynomial on points formed with
+        triples of corresponding values from `x`, `y`, and `z`.
+
+    See Also
+    --------
+    polyval, polyval2d, polygrid2d, polygrid3d
+
+    """
+    return polyval(z, polyval2d(x, y, c), False)
+
+
+def polygrid3d(x, y, z, c):
+    """
+    Evaluate 3D polynomials on the Cartesion product of x,y,z.
+
+    This function returns the values:
+
+    ``p(a,b,c) = \sum_{i,j,k} c[i,j,k] * a^i * b^j * c^k``
+
+    where the points ``(a,b,c)`` consist of all triples formed by taking
+    ``a`` from `x`, ``b`` from `y`, and ``c`` from `z`. The resulting
+    points form a grid with `x` in the first dimension, `y` in the second,
+    and `z` in the third.
+
+    Parameters
+    ----------
+    x,y,z : array_like, algebra_like
+        The three dimensional polynomial is evaluated at the points
+        ``(x,y,z)``, where `x` and `y` must have the same shape. If `x` or
+        `y` is a list or tuple, it is first converted to an ndarray,
+        otherwise it is left unchanged and treated as a scalar. See
+        `polyval` for explanation of algebra_like.
+    c : array_like
+        Array of coefficients ordered so that the coefficients for terms of
+        degree i,j are contained in ``c[i,j,k]``. If `c` has dimension
+        greater than 3 the remaining indices enumerate multiple sets of
+        coefficients.
+
+    Returns
+    -------
+    values : ndarray, algebra_like
+        The values of the multidimensional polynomial on the cartesion
+        product of `x`, `y`, and `z`.
+
+    See Also
+    --------
+    polyval, polyval2d, polygrid2d, polyval3d
+
+    """
+    return polyval(z, polygrid2d(x, y, c))
+
+
 def polyvander(x, deg) :
     """Vandermonde matrix of given degree.