Numpy meshgrid in 3D

Numpy's meshgrid is very useful for converting two vectors to a coordinate grid. What is the easiest way to extend this to three dimensions? So given three vectors x, y, and z, construct 3x3D arrays (instead of 2x2D arrays) which can be used as coordinates.

Solution 1:

Numpy (as of 1.8 I think) now supports higher that 2D generation of position grids with meshgrid. One important addition which really helped me is the ability to chose the indexing order (either xy or ij for Cartesian or matrix indexing respectively), which I verified with the following example:

import numpy as np

x_ = np.linspace(0., 1., 10)
y_ = np.linspace(1., 2., 20)
z_ = np.linspace(3., 4., 30)

x, y, z = np.meshgrid(x_, y_, z_, indexing='ij')

assert np.all(x[:,0,0] == x_)
assert np.all(y[0,:,0] == y_)
assert np.all(z[0,0,:] == z_)

Solution 2:

Here is the source code of meshgrid:

def meshgrid(x,y):
    """
    Return coordinate matrices from two coordinate vectors.

    Parameters
    ----------
    x, y : ndarray
        Two 1-D arrays representing the x and y coordinates of a grid.

    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`.

    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, 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]])

    `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)

    """
    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)

    y = y.reshape(numRows,1)
    Y = y.repeat(numCols, axis=1)
    return X, Y

It is fairly simple to understand. I extended the pattern to an arbitrary number of dimensions, but this code is by no means optimized (and not thoroughly error-checked either), but you get what you pay for. Hope it helps:

def meshgrid2(*arrs):
    arrs = tuple(reversed(arrs))  #edit
    lens = map(len, arrs)
    dim = len(arrs)

    sz = 1
    for s in lens:
        sz*=s

    ans = []    
    for i, arr in enumerate(arrs):
        slc = [1]*dim
        slc[i] = lens[i]
        arr2 = asarray(arr).reshape(slc)
        for j, sz in enumerate(lens):
            if j!=i:
                arr2 = arr2.repeat(sz, axis=j) 
        ans.append(arr2)

    return tuple(ans)