在一个太大的数组中加倍并添加

时间:2013-06-07 14:16:35

标签: python numpy

我有一个形状为(M,N)的阵列A现在我想做手术

R = (A[:,newaxis,:] * A[newaxis,:,:]).sum(2)

应该产生一个(MxM)数组。现在问题是数组很大,我得到一个内存错误,因为MxMxN数组不适合内存。

完成这项工作的最佳策略是什么? C?地图()?还是有特殊功能吗?

谢谢你, 大卫

1 个答案:

答案 0 :(得分:7)

我不确定你的数组有多大,但以下是等价的:

R = np.einsum('ij,kj',A,A)

并且速度可以更快,并且内存密集程度更低:

In [7]: A = np.random.random(size=(500,400))

In [8]: %timeit R = (A[:,np.newaxis,:] * A[np.newaxis,:,:]).sum(2)
1 loops, best of 3: 1.21 s per loop

In [9]: %timeit R = np.einsum('ij,kj',A,A)
10 loops, best of 3: 54 ms per loop

如果我将A的大小增加到(500,4000)np.einsum会在大约2秒内完成计算,而原始配方会因为大小而使我的机器停止运转它必须创建的临时数组。

<强>更新

正如@Jaime在评论中指出的那样,np.dot(A,A.T)也是问题的等效表述,甚至可能比np.einsum解决方案更快。完全赞同他指出这一点,但如果他不把它作为一个正式的解决方案发布,我想把它拉出来作为主答案。