ENH: enhance meshgrid to generate 3D grids, sparse grids, matrix indexing.

1 files changed, 106 insertions, 34 deletions

diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index f254bbacf..86ae3ab08 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -6,7 +6,7 @@ __all__ = ['select', 'piecewise', 'trim_zeros', 'copy', 'iterable',
         'histogram', 'histogramdd', 'bincount', 'digitize', 'cov', 'corrcoef',
         'msort', 'median', 'sinc', 'hamming', 'hanning', 'bartlett',
         'blackman', 'kaiser', 'trapz', 'i0', 'add_newdoc', 'add_docstring',
-        'meshgrid', 'delete', 'insert', 'append', 'interp', 'add_newdoc_ufunc']
+        'meshgrid', 'ndgrid', 'delete', 'insert', 'append', 'interp', 'add_newdoc_ufunc']
 
 import warnings
 import types
@@ -3200,62 +3200,134 @@ 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' or 'ij' (optional)
+        cartesian ('xy', default) or matrix ('ij') indexing of output
+    sparse : True or False (default) (optional)
+         If True a sparse grid is returned in order to conserve memory.
+    copy : True (default) or False (optional)
+        If False a view into the original arrays are returned in order to
+        conserve memory. Please note that sparse=False, copy=False will likely 
+        return non-contiguous arrays. Furthermore, more than one element of a
+        broadcasted 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. 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)
+    >>> xx, yy = meshgrid(x, y, sparse=True)
     >>> z = np.sin(xx**2+yy**2)/(xx**2+yy**2)
-
+    
+    >>> import matplotlib.pyplot as plt
+    >>> 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)
+    copy_ = kwargs.get('copy', True)
+    args = np.atleast_1d(*xi)
+    ndim = len(args)
+    
+    if not isinstance(args, list) or ndim<2:
+        raise TypeError('meshgrid() takes 2 or more arguments (%d given)' % int(ndim>0))
+
+    sparse = kwargs.get('sparse', False)
+    indexing = kwargs.get('indexing', 'xy')
+
+    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]
 
-    y = y.reshape(numRows,1)
-    Y = y.repeat(numCols, axis=1)
-    return X, Y
+    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)
+
+
+def ndgrid(*args, **kwargs):
+    """
+    Same as calling meshgrid with indexing='ij' (see meshgrid for
+    documentation).
+    """
+    kwargs['indexing'] = 'ij'
+    return meshgrid(*args, **kwargs)
 
 def delete(arr, obj, axis=None):
     """