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

前端 未结 4 695
名媛妹妹
名媛妹妹 2020-12-03 14:15

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


        
4条回答
  •  [愿得一人]
    2020-12-03 14:59

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

提交回复
热议问题