import numpy as np def meshgrid(*xi,**kwargs): """ Return coordinate matrices from one 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 ---------- 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 Returns ------- 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. See Also -------- index_tricks.mgrid : Construct a multi-dimensional "meshgrid" using indexing notation. index_tricks.ogrid : Construct an open multi-dimensional "meshgrid" using indexing notation. Examples -------- >>> x = np.linspace(0,1,3) # coordinates along x axis >>> y = np.linspace(0,1,2) # coordinates along y axis >>> xv, yv = meshgrid(x,y) # extend x and y for a 2D xy grid >>> 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(x,y,sparse=True,indexing='ij') # change to matrix indexing [array([[ 0. ], [ 0.5], [ 1. ]]), array([[ 0., 1.]])] >>> meshgrid(x,y,indexing='ij') [array([[ 0. , 0. ], [ 0.5, 0.5], [ 1. , 1. ]]), array([[ 0., 1.], [ 0., 1.], [ 0., 1.]])] >>> meshgrid(0,1,5) # just a 3D point [array([[[0]]]), array([[[1]]]), array([[[5]]])] >>> map(np.squeeze,meshgrid(0,1,5)) # just a 3D point [array(0), array(1), array(5)] >>> meshgrid(3) array([3]) >>> meshgrid(y) # 1D grid; y is just returned 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 = meshgrid(x, y, sparse=True) >>> z = np.sin(xx**2+yy**2)/(xx**2+yy**2) """ copy = kwargs.get('copy',True) args = np.atleast_1d(*xi) if not isinstance(args, list): if args.size>0: return args.copy() if copy else args else: raise TypeError('meshgrid() take 1 or more arguments (0 given)') sparse = kwargs.get('sparse',False) indexing = kwargs.get('indexing','xy') # 'ij' ndim = len(args) 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) def ndgrid(*args,**kwargs): """ Same as calling meshgrid with indexing='ij' (see meshgrid for documentation). """ kwargs['indexing'] = 'ij' return meshgrid(*args,**kwargs) if __name__=='__main__': import doctest doctest.testmod()