diff options
Diffstat (limited to 'numpy/lib/function_base.py')
| -rw-r--r-- | numpy/lib/function_base.py | 140 |
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): """ |