在numpy数组的每一行上应用一个函数

时间:2013-04-26 22:48:06

标签: python numpy

我有一个(16000000,5)numpy数组,我想在每一行上应用这个函数。

def f(row):
#returns a row of the same length.
    return [row[0]+0.5*row[1],row[2]+0.5*row[3],row[3]-0.5*row[2],row[4]-0.5*row[3],row[4]+1]

矢量化会很慢。

我试过这样的

np.column_stack((arr[:,0]+0.5*arr[:,1],arr[:,2]+0.5*arr[:,3],arr[:,3]-0.5*arr[:,2],arr[:,4]-0.5*arr[:,3],arr[:,4]+1))

但是我得到了内存错误。

最快的方法是什么?

2 个答案:

答案 0 :(得分:2)

In [104]: arr=np.random.rand(1000000,5)
In [105]: %timeit a=np.column_stack((arr[:,0]+0.5*arr[:,1],arr[:,2]+0.5*arr[:,3],arr[:,3]-0.5*arr[:,2],arr[:,4]-0.5*arr[:,3],arr[:,4]+1))
10 loops, best of 3: 86.3 ms per loop

In [106]: %timeit a2=map(f,arr)1 loops, best of 3: 10.2 s per loop


In [98]: a2=map(f,arr)

In [99]: %timeit a2=map(f,arr)
100 loops, best of 3: 10.5 ms per loop

In [100]: np.all(a==a2)
Out[100]: True

答案 1 :(得分:2)

你最好预先分配并将操作拆分成单独的行,你在使用column_stack时在可读性或速度方面没有获得任何东西。

result = np.zeros_like(arr)
result[:, 0] = arr[:, 0] + .5 * arr[:, 1]
result[:, 1] = arr[:, 2] + .5 * arr[:, 3]
result[:, 2] = arr[:, 3] - .5 * arr[:, 2]
result[:, 3] = arr[:, 4] - .5 * arr[:, 3]
result[:, 4] = arr[:, 4] + 1