python – 聚合numpy函数

2023年5月29日 236次阅读

我有一个强调的操作,我需要优化：

 np.sum(a**2, axis=1)**.5   # where a is a 2 dimensional ndarray

该操作由三个函数组成,需要三次迭代’a’.在一个函数下聚合所有操作并沿轴1仅应用该函数会更有效.不幸的是,numpy的apply_along_axis函数不是一个选项,因为性能大约是x1000差.

有没有办法聚合几个numpy操作,所以它只需要在数组上循环一次？

最佳答案 使用浮点数组时,可以使用
np.einsum –

np.sqrt(np.einsum('ij,ij->i',a,a))

运行时测试 –

In [34]: a = np.random.rand(1000,1000)

In [35]: np.allclose(np.sum(a**2, axis=1)**.5,np.sqrt(np.einsum('ij,ij->i',a,a)))
Out[35]: True

In [36]: %timeit np.sum(a**2, axis=1)**.5
100 loops, best of 3: 7.57 ms per loop

In [37]: %timeit np.sqrt(np.einsum('ij,ij->i',a,a))
1000 loops, best of 3: 1.52 ms per loop
点赞