Creating a numpy array of 3D coordinates from three 1D arrays

Question

Suppose I have three arbitrary 1D arrays, for example:

x_p = np.array((1.0, 2.0, 3.0, 4.0, 5.0))
y_p = np.array((2.0, 3.0, 4.0))
z_p = np.array((8.0, 9.0))

These three arrays represent sampling intervals in a 3D grid, and I want to construct a 1D array of three-dimensional vectors for all intersections, something like

points = np.array([[1.0, 2.0, 8.0],
                   [1.0, 2.0, 9.0],
                   [1.0, 3.0, 8.0],
                   ...
                   [5.0, 4.0, 9.0]])

Order doesn't actually matter for this. The obvious way to generate them:

npoints = len(x_p) * len(y_p) * len(z_p)
points = np.zeros((npoints, 3))
i = 0
for x in x_p:
    for y in y_p:
        for z in z_p:
            points[i, :] = (x, y, z)
            i += 1

So the question is... is there a faster way? I have looked but not found (possibly just failed to find the right Google keywords).

I am currently using this:

npoints = len(x_p) * len(y_p) * len(z_p)
points = np.zeros((npoints, 3))
i = 0
nz = len(z_p)
for x in x_p:
    for y in y_p:
        points[i:i+nz, 0] = x
        points[i:i+nz, 1] = y
        points[i:i+nz, 2] = z_p
        i += nz

but I feel like I am missing some clever fancy Numpy way?

Answer 1

To use numpy mesh grid on the above example the following will work:

np.vstack(np.meshgrid(x_p,y_p,z_p)).reshape(3,-1).T

Numpy meshgrid for grids of more then two dimensions require numpy 1.7. To circumvent this and pulling the relevant data from the source code.

def ndmesh(*xi,**kwargs):
    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)

    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]

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

Checking the result:

print np.vstack((ndmesh(x_p,y_p,z_p))).reshape(3,-1).T

[[ 1.  2.  8.]
 [ 1.  2.  9.]
 [ 1.  3.  8.]
 ....
 [ 5.  3.  9.]
 [ 5.  4.  8.]
 [ 5.  4.  9.]]

For the above example:

%timeit sol2()
10000 loops, best of 3: 56.1 us per loop

%timeit np.vstack((ndmesh(x_p,y_p,z_p))).reshape(3,-1).T
10000 loops, best of 3: 55.1 us per loop

For when each dimension is 100:

%timeit sol2()
1 loops, best of 3: 655 ms per loop
In [10]:

%timeit points = np.vstack((ndmesh(x_p,y_p,z_p))).reshape(3,-1).T
10 loops, best of 3: 21.8 ms per loop

Depending on what you want to do with the data, you can return a view:

%timeit np.vstack((ndmesh(x_p,y_p,z_p,copy=False))).reshape(3,-1).T
100 loops, best of 3: 8.16 ms per loop

Answer 2

For your specific example, mgrid is quite useful.:

In [1]: import numpy as np
In [2]: points = np.mgrid[1:6, 2:5, 8:10]
In [3]: points.reshape(3, -1).T
Out[3]:
array([[1, 2, 8],
       [1, 2, 9],
       [1, 3, 8],
       [1, 3, 9],
       [1, 4, 8],
       [1, 4, 9],
       [2, 2, 8],
       [2, 2, 9],
       [2, 3, 8],
       [2, 3, 9],
       [2, 4, 8],
       [2, 4, 9],
       [3, 2, 8],
       [3, 2, 9],
       [3, 3, 8],
       [3, 3, 9],
       [3, 4, 8],
       [3, 4, 9],
       [4, 2, 8],
       [4, 2, 9],
       [4, 3, 8],
       [4, 3, 9],
       [4, 4, 8],
       [4, 4, 9],
       [5, 2, 8],
       [5, 2, 9],
       [5, 3, 8],
       [5, 3, 9],
       [5, 4, 8],
       [5, 4, 9]])

More generally, if you have specific arrays that you want to do this with, you'd use meshgrid instead of mgrid. However, you'll need numpy 1.7 or later for it to work in more than two dimensions.

Answer 3

You can use itertools.product:

def sol1():
    points = np.array( list(product(x_p, y_p, z_p)) )

Another solution using iterators and np.take showed to be about 3X faster for your case:

from itertools import izip, product

def sol2():
    points = np.empty((len(x_p)*len(y_p)*len(z_p),3))

    xi,yi,zi = izip(*product( xrange(len(x_p)),
                              xrange(len(y_p)),
                              xrange(len(z_p))  ))

    points[:,0] = np.take(x_p,xi)
    points[:,1] = np.take(y_p,yi)
    points[:,2] = np.take(z_p,zi)
    return points

Timing results:

In [3]: timeit sol1()
10000 loops, best of 3: 126 µs per loop

In [4]: timeit sol2()
10000 loops, best of 3: 42.9 µs per loop

In [6]: timeit ops()
10000 loops, best of 3: 59 µs per loop

In [11]: timeit joekingtons() # with your permission Joe...
10000 loops, best of 3: 56.2 µs per loop

So, only my second solution and Joe Kington's solution would give you some performance gain...

Answer 4

For those who had to stay with numpy <1.7.x, here is a simple two-liner solution:

g_p=np.zeros((x_p.size, y_p.size, z_p.size))
array_you_want=array(zip(*[item.flatten() for item in \
                                 [g_p+x_p[...,np.newaxis,np.newaxis],\
                                  g_p+y_p[...,np.newaxis],\
                                  g_p+z_p]]))

Very easy to expand to even higher dimenision as well. Cheers!