MAINT: polynomial: Add an N-d vander implementation used under the hood of the vander2d and vander3d functions

The generalization is not exposed in the public API yet, but it could be if the need arises.

The shape / dtype conversion logic is left as is for now, even if it might be broken.

7 files changed, 86 insertions, 41 deletions

diff --git a/numpy/polynomial/chebyshev.py b/numpy/polynomial/chebyshev.py
index 093eb0048..0cd9c4d23 100644
--- a/numpy/polynomial/chebyshev.py
+++ b/numpy/polynomial/chebyshev.py
@@ -1468,7 +1468,7 @@ def chebvander2d(x, y, deg):
     .. versionadded:: 1.7.0
 
     """
-    return pu._vander2d(chebvander, x, y, deg)
+    return pu._vander_nd_flat((chebvander, chebvander), (x, y), deg)
 
 
 def chebvander3d(x, y, z, deg):
@@ -1522,7 +1522,7 @@ def chebvander3d(x, y, z, deg):
     .. versionadded:: 1.7.0
 
     """
-    return pu._vander3d(chebvander, x, y, z, deg)
+    return pu._vander_nd_flat((chebvander, chebvander, chebvander), (x, y, z), deg)
 
 
 def chebfit(x, y, deg, rcond=None, full=False, w=None):
diff --git a/numpy/polynomial/hermite.py b/numpy/polynomial/hermite.py
index 0011fa3b7..9b1aea239 100644
--- a/numpy/polynomial/hermite.py
+++ b/numpy/polynomial/hermite.py
@@ -1193,7 +1193,7 @@ def hermvander2d(x, y, deg):
     .. versionadded:: 1.7.0
 
     """
-    return pu._vander2d(hermvander, x, y, deg)
+    return pu._vander_nd_flat((hermvander, hermvander), (x, y), deg)
 
 
 def hermvander3d(x, y, z, deg):
@@ -1247,7 +1247,7 @@ def hermvander3d(x, y, z, deg):
     .. versionadded:: 1.7.0
 
     """
-    return pu._vander3d(hermvander, x, y, z, deg)
+    return pu._vander_nd_flat((hermvander, hermvander, hermvander), (x, y, z), deg)
 
 
 def hermfit(x, y, deg, rcond=None, full=False, w=None):
diff --git a/numpy/polynomial/hermite_e.py b/numpy/polynomial/hermite_e.py
index b1cc2d3ab..c5a0a05a2 100644
--- a/numpy/polynomial/hermite_e.py
+++ b/numpy/polynomial/hermite_e.py
@@ -1186,7 +1186,7 @@ def hermevander2d(x, y, deg):
     .. versionadded:: 1.7.0
 
     """
-    return pu._vander2d(hermevander, x, y, deg)
+    return pu._vander_nd_flat((hermevander, hermevander), (x, y), deg)
 
 
 def hermevander3d(x, y, z, deg):
@@ -1240,7 +1240,7 @@ def hermevander3d(x, y, z, deg):
     .. versionadded:: 1.7.0
 
     """
-    return pu._vander3d(hermevander, x, y, z, deg)
+    return pu._vander_nd_flat((hermevander, hermevander, hermevander), (x, y, z), deg)
 
 
 def hermefit(x, y, deg, rcond=None, full=False, w=None):
diff --git a/numpy/polynomial/laguerre.py b/numpy/polynomial/laguerre.py
index 7e7e45ca1..538a1d449 100644
--- a/numpy/polynomial/laguerre.py
+++ b/numpy/polynomial/laguerre.py
@@ -1193,7 +1193,7 @@ def lagvander2d(x, y, deg):
     .. versionadded:: 1.7.0
 
     """
-    return pu._vander2d(lagvander, x, y, deg)
+    return pu._vander_nd_flat((lagvander, lagvander), (x, y), deg)
 
 
 def lagvander3d(x, y, z, deg):
@@ -1247,7 +1247,7 @@ def lagvander3d(x, y, z, deg):
     .. versionadded:: 1.7.0
 
     """
-    return pu._vander3d(lagvander, x, y, z, deg)
+    return pu._vander_nd_flat((lagvander, lagvander, lagvander), (x, y, z), deg)
 
 
 def lagfit(x, y, deg, rcond=None, full=False, w=None):
diff --git a/numpy/polynomial/legendre.py b/numpy/polynomial/legendre.py
index 281982d0b..c11824761 100644
--- a/numpy/polynomial/legendre.py
+++ b/numpy/polynomial/legendre.py
@@ -1229,7 +1229,7 @@ def legvander2d(x, y, deg):
     .. versionadded:: 1.7.0
 
     """
-    return pu._vander2d(legvander, x, y, deg)
+    return pu._vander_nd_flat((legvander, legvander), (x, y), deg)
 
 
 def legvander3d(x, y, z, deg):
@@ -1283,7 +1283,7 @@ def legvander3d(x, y, z, deg):
     .. versionadded:: 1.7.0
 
     """
-    return pu._vander3d(legvander, x, y, z, deg)
+    return pu._vander_nd_flat((legvander, legvander, legvander), (x, y, z), deg)
 
 
 def legfit(x, y, deg, rcond=None, full=False, w=None):
diff --git a/numpy/polynomial/polynomial.py b/numpy/polynomial/polynomial.py
index 3f0a902cf..315ea1495 100644
--- a/numpy/polynomial/polynomial.py
+++ b/numpy/polynomial/polynomial.py
@@ -1133,7 +1133,7 @@ def polyvander2d(x, y, deg):
     polyvander, polyvander3d, polyval2d, polyval3d
 
     """
-    return pu._vander2d(polyvander, x, y, deg)
+    return pu._vander_nd_flat((polyvander, polyvander), (x, y), deg)
 
 
 def polyvander3d(x, y, z, deg):
@@ -1187,7 +1187,7 @@ def polyvander3d(x, y, z, deg):
     .. versionadded:: 1.7.0
 
     """
-    return pu._vander3d(polyvander, x, y, z, deg)
+    return pu._vander_nd_flat((polyvander, polyvander, polyvander), (x, y, z), deg)
 
 
 def polyfit(x, y, deg, rcond=None, full=False, w=None):
diff --git a/numpy/polynomial/polyutils.py b/numpy/polynomial/polyutils.py
index 35b24d1ab..5dcfa7a7a 100644
--- a/numpy/polynomial/polyutils.py
+++ b/numpy/polynomial/polyutils.py
@@ -46,6 +46,7 @@ Functions
 from __future__ import division, absolute_import, print_function
 
 import operator
+import functools
 import warnings
 
 import numpy as np
@@ -415,45 +416,89 @@ def mapdomain(x, old, new):
     return off + scl*x
 
 
-def _vander2d(vander_f, x, y, deg):
-    """
-    Helper function used to implement the ``<type>vander2d`` functions.
+def _nth_slice(i, ndim):
+    sl = [np.newaxis] * ndim
+    sl[i] = slice(None)
+    return tuple(sl)
+
+
+def _vander_nd(vander_fs, points, degrees):
+    r"""
+    A generalization of the Vandermonde matrix for N dimensions
+
+    The result is built by combining the results of 1d Vandermonde matrices,
+
+    .. math::
+        W[i_0, \ldots, i_M, j_0, \ldots, j_N] = \prod_{k=0}^N{V_k(x_k)[i_0, \ldots, i_M, j_k]}
+
+    where
+
+    .. math::
+        N &= \texttt{len(points)} = \texttt{len(degrees)} = \texttt{len(vander\_fs)} \\
+        M &= \texttt{points[k].ndim} \\
+        V_k &= \texttt{vander\_fs[k]} \\
+        x_k &= \texttt{points[k]} \\
+        0 \le j_k &\le \texttt{degrees[k]}
+
+    Expanding the one-dimensional :math:`V_k` functions gives:
+
+    .. math::
+        W[i_0, \ldots, i_M, j_0, \ldots, j_N] = \prod_{k=0}^N{B_{k, j_k}(x_k[i_0, \ldots, i_M])}
+
+    where :math:`B_{k,m}` is the m'th basis of the polynomial construction used along
+    dimension :math:`k`. For a regular polynomial, :math:`B_{k, m}(x) = P_m(x) = x^m`.
 
     Parameters
     ----------
-    vander_f : function(array_like, int) -> ndarray
-        The 1d vander function, such as ``polyvander``
-    x, y, deg :
-        See the ``<type>vander2d`` functions for more detail
+    vander_fs : Sequence[function(array_like, int) -> ndarray]
+        The 1d vander function to use for each axis, such as ``polyvander``
+    points : Sequence[array_like]
+        Arrays of point coordinates, all of the same shape. The dtypes
+        will be converted to either float64 or complex128 depending on
+        whether any of the elements are complex. Scalars are converted to
+        1-D arrays.
+        This must be the same length as `vander_fs`.
+    degrees : Sequence[int]
+        The maximum degree (inclusive) to use for each axis.
+        This must be the same length as `vander_fs`.
+
+    Returns
+    -------
+    vander_nd : ndarray
+        An array of shape ``points[0].shape + tuple(d + 1 for d in degrees)``.
     """
-    degx, degy = deg
-    x, y = np.array((x, y), copy=False) + 0.0
+    n_dims = len(vander_fs)
+    if n_dims != len(points):
+        raise ValueError(
+            "Expected {} dimensions of sample points, got {}".format(n_dims, len(points)))
+    if n_dims != len(degrees):
+        raise ValueError(
+            "Expected {} dimensions of degrees, got {}".format(n_dims, len(degrees)))
+    if n_dims == 0:
+        raise ValueError("Unable to guess a dtype or shape when no points are given")
+
+    # convert to the same shape and type
+    points = tuple(np.array(tuple(points), copy=False) + 0.0)
 
-    vx = vander_f(x, degx)
-    vy = vander_f(y, degy)
-    v = vx[..., None]*vy[..., None,:]
-    return v.reshape(v.shape[:-2] + (-1,))
+    # produce the vandermonde matrix for each dimension, placing the last
+    # axis of each in an independent trailing axis of the output
+    vander_arrays = (
+        vander_fs[i](points[i], degrees[i])[(...,) + _nth_slice(i, n_dims)]
+        for i in range(n_dims)
+    )
 
+    # we checked this wasn't empty already, so no `initial` needed
+    return functools.reduce(operator.mul, vander_arrays)
 
-def _vander3d(vander_f, x, y, z, deg):
+
+def _vander_nd_flat(vander_fs, points, degrees):
     """
-    Helper function used to implement the ``<type>vander3d`` functions.
+    Like `_vander_nd`, but flattens the last ``len(degrees)`` axes into a single axis
 
-    Parameters
-    ----------
-    vander_f : function(array_like, int) -> ndarray
-        The 1d vander function, such as ``polyvander``
-    x, y, z, deg :
-        See the ``<type>vander3d`` functions for more detail
+    Used to implement the public ``<type>vander<n>d`` functions.
     """
-    degx, degy, degz = deg
-    x, y, z = np.array((x, y, z), copy=False) + 0.0
-
-    vx = vander_f(x, degx)
-    vy = vander_f(y, degy)
-    vz = vander_f(z, degz)
-    v = vx[..., None, None]*vy[..., None,:, None]*vz[..., None, None,:]
-    return v.reshape(v.shape[:-3] + (-1,))
+    v = _vander_nd(vander_fs, points, degrees)
+    return v.reshape(v.shape[:-len(degrees)] + (-1,))
 
 
 def _fromroots(line_f, mul_f, roots):