Merge pull request #192 from rgommers/meshgrid3d

Meshgrid enhancements (>2-D, sparse grids, matrix indexing)

1 files changed, 106 insertions, 33 deletions

diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index a0781ebf9..1edfc4c6b 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -3247,62 +3247,135 @@ def add_newdoc(place, obj, doc):
         pass
 
 
-# From matplotlib
-def meshgrid(x,y):
+# Based on scitools meshgrid
+def meshgrid(*xi, **kwargs):
     """
-    Return coordinate matrices from two coordinate vectors.
+    Return coordinate matrices from two or more coordinate vectors.
+
+    Make N-D coordinate arrays for vectorized evaluations of
+    N-D scalar/vector fields over N-D grids, given
+    one-dimensional coordinate arrays x1, x2,..., xn.
 
     Parameters
     ----------
-    x, y : ndarray
-        Two 1-D arrays representing the x and y coordinates of a grid.
+    x1, x2,..., xn : array_like
+        1-D arrays representing the coordinates of a grid.
+    indexing : {'xy', 'ij'}, optional
+        Cartesian ('xy', default) or matrix ('ij') indexing of output.
+        See Notes for more details.
+    sparse : bool, optional
+         If True a sparse grid is returned in order to conserve memory.
+         Default is False.
+    copy : bool, optional
+        If False, a view into the original arrays are returned in
+        order to conserve memory.  Default is True.  Please note that
+        ``sparse=False, copy=False`` will likely return non-contiguous arrays.
+        Furthermore, more than one element of a broadcast array may refer to
+        a single memory location.  If you need to write to the arrays, make
+        copies first.
 
     Returns
     -------
-    X, Y : ndarray
-        For vectors `x`, `y` with lengths ``Nx=len(x)`` and ``Ny=len(y)``,
-        return `X`, `Y` where `X` and `Y` are ``(Ny, Nx)`` shaped arrays
-        with the elements of `x` and y repeated to fill the matrix along
-        the first dimension for `x`, the second for `y`.
+    X1, X2,..., XN : ndarray
+        For vectors `x1`, `x2`,..., 'xn' with lengths ``Ni=len(xi)`` ,
+        return ``(N1, N2, N3,...Nn)`` shaped arrays if indexing='ij'
+        or ``(N2, N1, N3,...Nn)`` shaped arrays if indexing='xy'
+        with the elements of `xi` repeated to fill the matrix along
+        the first dimension for `x1`, the second for `x2` and so on.
+
+    Notes
+    -----
+    This function supports both indexing conventions through the indexing keyword
+    argument.  Giving the string 'ij' returns a meshgrid with matrix indexing,
+    while 'xy' returns a meshgrid with Cartesian indexing.  In the 2-D case
+    with inputs of length M and N, the outputs are of shape (N, M) for 'xy'
+    indexing and (M, N) for 'ij' indexing.  In the 3-D case with inputs of
+    length M, N and P, outputs are of shape (N, M, P) for 'xy' indexing and (M,
+    N, P) for 'ij' indexing.  The difference is illustrated by the following
+    code snippet::
+
+        xv, yv = meshgrid(x, y, sparse=False, indexing='ij')
+        for i in range(nx):
+            for j in range(ny):
+                # treat xv[i,j], yv[i,j]
+
+        xv, yv = meshgrid(x, y, sparse=False, indexing='xy')
+        for i in range(nx):
+            for j in range(ny):
+                # treat xv[j,i], yv[j,i]
 
     See Also
     --------
     index_tricks.mgrid : Construct a multi-dimensional "meshgrid"
-                         using indexing notation.
+                     using indexing notation.
     index_tricks.ogrid : Construct an open multi-dimensional "meshgrid"
-                         using indexing notation.
+                     using indexing notation.
 
     Examples
     --------
-    >>> X, Y = np.meshgrid([1,2,3], [4,5,6,7])
-    >>> X
-    array([[1, 2, 3],
-           [1, 2, 3],
-           [1, 2, 3],
-           [1, 2, 3]])
-    >>> Y
-    array([[4, 4, 4],
-           [5, 5, 5],
-           [6, 6, 6],
-           [7, 7, 7]])
+    >>> nx, ny = (3, 2)
+    >>> x = np.linspace(0, 1, nx)
+    >>> y = np.linspace(0, 1, ny)
+    >>> xv, yv = meshgrid(x, y)
+    >>> xv
+    array([[ 0. ,  0.5,  1. ],
+           [ 0. ,  0.5,  1. ]])
+    >>> yv
+    array([[ 0.,  0.,  0.],
+           [ 1.,  1.,  1.]])
+    >>> xv, yv = meshgrid(x, y, sparse=True)  # make sparse output arrays
+    >>> xv
+    array([[ 0. ,  0.5,  1. ]])
+    >>> yv
+    array([[ 0.],
+           [ 1.]])
 
     `meshgrid` is very useful to evaluate functions on a grid.
 
     >>> x = np.arange(-5, 5, 0.1)
     >>> y = np.arange(-5, 5, 0.1)
-    >>> xx, yy = np.meshgrid(x, y)
-    >>> z = np.sin(xx**2+yy**2)/(xx**2+yy**2)
+    >>> xx, yy = meshgrid(x, y, sparse=True)
+    >>> z = np.sin(xx**2 + yy**2) / (xx**2 + yy**2)
+    >>> h = plt.contourf(x,y,z)
 
     """
-    x = asarray(x)
-    y = asarray(y)
-    numRows, numCols = len(y), len(x)  # yes, reversed
-    x = x.reshape(1,numCols)
-    X = x.repeat(numRows, axis=0)
+    if len(xi) < 2:
+        msg = 'meshgrid() takes 2 or more arguments (%d given)' % int(len(xi) > 0)
+        raise ValueError(msg)
+
+    args = np.atleast_1d(*xi)
+    ndim = len(args)
+
+    copy_ = kwargs.get('copy', True)
+    sparse = kwargs.get('sparse', False)
+    indexing = kwargs.get('indexing', 'xy')
+    if not indexing in ['xy', 'ij']:
+        raise ValueError("Valid values for `indexing` are 'xy' and 'ij'.")
+
+    s0 = (1,) * ndim
+    output = [x.reshape(s0[:i] + (-1,) + s0[i + 1::]) for i, x in enumerate(args)]
+
+    shape = [x.size for x in output]
+
+    if indexing == 'xy':
+        # switch first and second axis
+        output[0].shape = (1, -1) + (1,)*(ndim - 2)
+        output[1].shape = (-1, 1) + (1,)*(ndim - 2)
+        shape[0], shape[1] = shape[1], shape[0]
+
+    if sparse:
+        if copy_:
+            return [x.copy() for x in output]
+        else:
+            return output
+    else:
+        # Return the full N-D matrix (not only the 1-D vector)
+        if copy_:
+            mult_fact = np.ones(shape, dtype=int)
+            return [x * mult_fact for x in output]
+        else:
+            return np.broadcast_arrays(*output)
 
-    y = y.reshape(numRows,1)
-    Y = y.repeat(numCols, axis=1)
-    return X, Y
 
 def delete(arr, obj, axis=None):
     """