ENH: Make derivatives and integrals work on multidimensional array.

An axis keyword was added to the function signatures of xxxder and
xxxint, where xxx is any of poly, cheb, leg, lag, herm, herme. The
evaluation method for the Chebeshev series was also changed to avoid
using z_series and to more closely resemble the other implementations.

At some point the z_series will be removed from the chebyshev module
and only used for trigonometric series.

1 files changed, 98 insertions, 62 deletions

diff --git a/numpy/polynomial/polynomial.py b/numpy/polynomial/polynomial.py
index bf9c1a216..87acfa616 100644
--- a/numpy/polynomial/polynomial.py
+++ b/numpy/polynomial/polynomial.py
@@ -414,8 +414,8 @@ def polypow(cs, pow, maxpower=None) :
     Parameters
     ----------
     cs : array_like
-        1d array of chebyshev series coefficients ordered from low to
-        high.
+        1d array of array of series coefficients ordered from low to
+        high degree.
     pow : integer
         Power to which the series will be raised
     maxpower : integer, optional
@@ -425,11 +425,11 @@ def polypow(cs, pow, maxpower=None) :
     Returns
     -------
     coef : ndarray
-        Chebyshev series of power.
+        Power series of power.
 
     See Also
     --------
-    chebadd, chebsub, chebmul, chebdiv
+    polyadd, polysub, polymul, polydiv
 
     Examples
     --------
@@ -455,31 +455,37 @@ def polypow(cs, pow, maxpower=None) :
         return prd
 
 
-def polyder(cs, m=1, scl=1):
+def polyder(c, m=1, scl=1, axis=0):
     """
     Differentiate a polynomial.
 
-    Returns the polynomial `cs` differentiated `m` times.  At each
-    iteration the result is multiplied by `scl` (the scaling factor is for
-    use in a linear change of variable).  The argument `cs` is the sequence
-    of coefficients from lowest order term to highest, e.g., [1,2,3]
-    represents the polynomial ``1 + 2*x + 3*x**2``.
+    Returns the polynomial coefficients `c` differentiated `m` times along
+    `axis`.  At each iteration the result is multiplied by `scl` (the
+    scaling factor is for use in a linear change of variable).  The
+    argument `c` is an array of coefficients from low to high degree along
+    each axis, e.g., [1,2,3] represents the polynomial ``1 + 2*x + 3*x**2``
+    while [[1,2],[1,2]] represents ``1 + 1*x + 2*y + 2*x*y`` if axis=0 is
+    ``x`` and axis=1 is ``y``.
 
     Parameters
     ----------
-    cs: array_like
-        1-d array of polynomial coefficients ordered from low to high.
+    c: array_like
+        Array of polynomial coefficients. If c is multidimensional the
+        different axis correspond to different variables with the degree
+        in each axis given by the corresponding index.
     m : int, optional
         Number of derivatives taken, must be non-negative. (Default: 1)
     scl : scalar, optional
         Each differentiation is multiplied by `scl`.  The end result is
         multiplication by ``scl**m``.  This is for use in a linear change
         of variable. (Default: 1)
+    axis : int, optional
+        Axis over which the derivative is taken. (Default: 0).
 
     Returns
     -------
     der : ndarray
-        Polynomial of the derivative.
+        Polynomial coefficients of the derivative.
 
     See Also
     --------
@@ -488,55 +494,71 @@ def polyder(cs, m=1, scl=1):
     Examples
     --------
     >>> from numpy import polynomial as P
-    >>> cs = (1,2,3,4) # 1 + 2x + 3x**2 + 4x**3
-    >>> P.polyder(cs) # (d/dx)(cs) = 2 + 6x + 12x**2
+    >>> c = (1,2,3,4) # 1 + 2x + 3x**2 + 4x**3
+    >>> P.polyder(c) # (d/dx)(c) = 2 + 6x + 12x**2
     array([  2.,   6.,  12.])
-    >>> P.polyder(cs,3) # (d**3/dx**3)(cs) = 24
+    >>> P.polyder(c,3) # (d**3/dx**3)(c) = 24
     array([ 24.])
-    >>> P.polyder(cs,scl=-1) # (d/d(-x))(cs) = -2 - 6x - 12x**2
+    >>> P.polyder(c,scl=-1) # (d/d(-x))(c) = -2 - 6x - 12x**2
     array([ -2.,  -6., -12.])
-    >>> P.polyder(cs,2,-1) # (d**2/d(-x)**2)(cs) = 6 + 24x
+    >>> P.polyder(c,2,-1) # (d**2/d(-x)**2)(c) = 6 + 24x
     array([  6.,  24.])
 
     """
-    cnt = int(m)
+    c = np.array(c, ndmin=1, copy=1)
+    if c.dtype.char in '?bBhHiIlLqQpP':
+        c = c.astype(np.double)
+    cnt, iaxis = [int(t) for t in [m, axis]]
 
     if cnt != m:
         raise ValueError("The order of derivation must be integer")
     if cnt < 0:
         raise ValueError("The order of derivation must be non-negative")
+    if iaxis != axis:
+        raise ValueError("The axis must be integer")
+    if not -c.ndim <= iaxis < c.ndim:
+        raise ValueError("The axis is out of range")
+    if iaxis < 0:
+        iaxis += c.ndim
 
-    # cs is a trimmed copy
-    [cs] = pu.as_series([cs])
     if cnt == 0:
-        return cs
-    elif cnt >= len(cs):
-        return cs[:1]*0
+        return c
+
+    c = np.rollaxis(c, iaxis)
+    n = len(c)
+    if cnt >= n:
+        c = c[:1]*0
     else :
-        n = len(cs)
-        d = np.arange(n)*scl
         for i in range(cnt):
-            cs[i:] *= d[:n-i]
-        return cs[i+1:].copy()
+            n = n - 1
+            c *= scl
+            der = np.empty((n,) + c.shape[1:], dtype=c.dtype)
+            for j in range(n, 0, -1):
+                der[j - 1] = j*c[j]
+            c = der
+    c = np.rollaxis(c, 0, iaxis + 1)
+    return c
 
 
-def polyint(cs, m=1, k=[], lbnd=0, scl=1):
+def polyint(c, m=1, k=[], lbnd=0, scl=1, axis=0):
     """
     Integrate a polynomial.
 
-    Returns the polynomial `cs`, integrated `m` times from `lbnd` to `x`.
-    At each iteration the resulting series is **multiplied** by `scl` and
-    an integration constant, `k`, is added.  The scaling factor is for use
-    in a linear change of variable.  ("Buyer beware": note that, depending
-    on what one is doing, one may want `scl` to be the reciprocal of what
-    one might expect; for more information, see the Notes section below.)
-    The argument `cs` is a sequence of coefficients, from lowest order
-    term to highest, e.g., [1,2,3] represents the polynomial
-    ``1 + 2*x + 3*x**2``.
+    Returns the polynomial coefficients `c` integrated `m` times from
+    `lbnd` along `axis`.  At each iteration the resulting series is
+    **multiplied** by `scl` and an integration constant, `k`, is added.
+    The scaling factor is for use in a linear change of variable.  ("Buyer
+    beware": note that, depending on what one is doing, one may want `scl`
+    to be the reciprocal of what one might expect; for more information,
+    see the Notes section below.) The argument `c` is an array of
+    coefficients, from low to high degree along each axis, e.g., [1,2,3]
+    represents the polynomial ``1 + 2*x + 3*x**2`` while [[1,2],[1,2]]
+    represents ``1 + 1*x + 2*y + 2*x*y`` if axis=0 is ``x`` and axis=1 is
+    ``y``.
 
     Parameters
     ----------
-    cs : array_like
+    c : array_like
         1-d array of polynomial coefficients, ordered from low to high.
     m : int, optional
         Order of integration, must be positive. (Default: 1)
@@ -551,11 +573,13 @@ def polyint(cs, m=1, k=[], lbnd=0, scl=1):
     scl : scalar, optional
         Following each integration the result is *multiplied* by `scl`
         before the integration constant is added. (Default: 1)
+    axis : int, optional
+        Axis over which the integral is taken. (Default: 0).
 
     Returns
     -------
     S : ndarray
-        Coefficients of the integral.
+        Coefficient array of the integral.
 
     Raises
     ------
@@ -577,23 +601,26 @@ def polyint(cs, m=1, k=[], lbnd=0, scl=1):
     Examples
     --------
     >>> from numpy import polynomial as P
-    >>> cs = (1,2,3)
-    >>> P.polyint(cs) # should return array([0, 1, 1, 1])
+    >>> c = (1,2,3)
+    >>> P.polyint(c) # should return array([0, 1, 1, 1])
     array([ 0.,  1.,  1.,  1.])
-    >>> P.polyint(cs,3) # should return array([0, 0, 0, 1/6, 1/12, 1/20])
+    >>> P.polyint(c,3) # should return array([0, 0, 0, 1/6, 1/12, 1/20])
     array([ 0.        ,  0.        ,  0.        ,  0.16666667,  0.08333333,
             0.05      ])
-    >>> P.polyint(cs,k=3) # should return array([3, 1, 1, 1])
+    >>> P.polyint(c,k=3) # should return array([3, 1, 1, 1])
     array([ 3.,  1.,  1.,  1.])
-    >>> P.polyint(cs,lbnd=-2) # should return array([6, 1, 1, 1])
+    >>> P.polyint(c,lbnd=-2) # should return array([6, 1, 1, 1])
     array([ 6.,  1.,  1.,  1.])
-    >>> P.polyint(cs,scl=-2) # should return array([0, -2, -2, -2])
+    >>> P.polyint(c,scl=-2) # should return array([0, -2, -2, -2])
     array([ 0., -2., -2., -2.])
 
     """
-    cnt = int(m)
+    c = np.array(c, ndmin=1, copy=1)
+    if c.dtype.char in '?bBhHiIlLqQpP':
+        c = c.astype(np.double)
     if not np.iterable(k):
         k = [k]
+    cnt, iaxis = [int(t) for t in [m, axis]]
 
     if cnt != m:
         raise ValueError("The order of integration must be integer")
@@ -601,25 +628,34 @@ def polyint(cs, m=1, k=[], lbnd=0, scl=1):
         raise ValueError("The order of integration must be non-negative")
     if len(k) > cnt :
         raise ValueError("Too many integration constants")
+    if iaxis != axis:
+        raise ValueError("The axis must be integer")
+    if not -c.ndim <= iaxis < c.ndim:
+        raise ValueError("The axis is out of range")
+    if iaxis < 0:
+        iaxis += c.ndim
+
 
-    # cs is a trimmed copy
-    [cs] = pu.as_series([cs])
     if cnt == 0:
-        return cs
+        return c
 
     k = list(k) + [0]*(cnt - len(k))
+    c = np.rollaxis(c, iaxis)
     for i in range(cnt):
-        n = len(cs)
-        cs *= scl
-        if n == 1 and cs[0] == 0:
-            cs[0] += k[i]
+        n = len(c)
+        c *= scl
+        if n == 1 and np.all(c[0] == 0):
+            c[0] += k[i]
         else:
-            tmp = np.empty(n + 1, dtype=cs.dtype)
-            tmp[0] = cs[0]*0
-            tmp[1:] = cs/np.arange(1, n + 1)
+            tmp = np.empty((n + 1,) + c.shape[1:], dtype=c.dtype)
+            tmp[0] = c[0]*0
+            tmp[1] = c[0]
+            for j in range(1, n):
+                tmp[j + 1] = c[j]/(j + 1)
             tmp[0] += k[i] - polyval(lbnd, tmp)
-            cs = tmp
-    return cs
+            c = tmp
+    c = np.rollaxis(c, 0, iaxis + 1)
+    return c
 
 
 def polyval(x, c, tensor=True):
@@ -704,12 +740,12 @@ def polyval(x, c, tensor=True):
 
     """
     c = np.array(c, ndmin=1, copy=0)
-    if c.dtype.char not in 'efdgFDGO':
-        c = c + 0.0
+    if c.dtype.char in '?bBhHiIlLqQpP':
+        c = c.astype(np.double)
     if isinstance(x, (tuple, list)):
         x = np.asarray(x)
     if isinstance(x, np.ndarray) and tensor:
-       c = c.reshape(c.shape + (1,)*x.ndim)
+        c = c.reshape(c.shape + (1,)*x.ndim)
 
     c0 = c[-1] + x*0
     for i in range(2, len(c) + 1) :