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

Question

I have a (16000000,5) numpy array, and I want to apply this function on each row.

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]

vectorizing would operate slow.

I tried going like this

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

but I get Memory Error.

What is the fastest way to do this?

Answer 1

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

Answer 2

You're better off preallocate and splitting the operations into separate lines, you don't gain anything here in terms of readability or speed by using 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