从三个1D阵列创建一个numpy 3D坐标数组

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!