二维数组的 Numpy einsum 外和

Question

我尝试搜索答案，但找不到我需要的答案。如果这是一个重复的问题，我们深表歉意。

假设我有一个形状为 (n, n*m) 的二维数组。我想要做的是这个数组的外部总和到它的转置，结果是一个形状为 (n*m, n*m) 的数组。例如，假设我有

A = array([[1., 1., 2., 2.],
           [1., 1., 2., 2.]])

我想对 A 和 A.T 进行外部求和，使得输出为：

>>> array([[2., 2., 3., 3.],
           [2., 2., 3., 3.],
           [3., 3., 4., 4.],
           [3., 3., 4., 4.]])

请注意，np.add.outer 不起作用，因为它会将输入分解为向量。我可以通过

实现类似的效果

np.tile(A, (2, 1)) + np.tile(A.T, (1, 2))

但是当 n 和 m 相当大时（n > 100 和 m > 1000），这似乎不合理。是否可以使用 einsum 来写这个总和？我就是想不通 einsum.

Answer 1

一种方法是

(A.reshape(-1,*A.shape).T+A)[:,0,:]

我认为 n>100 和 m>1000 会占用大量内存。

但这和

不一样吗

np.add.outer(A,A)[:,0,:].reshape(4,-1)

Answer 2

为了利用 broadcasting，我们需要将其分解为 3D，然后置换轴并添加 -

n = A.shape[0]
m = A.shape[1]//n
a = A.reshape(n,m,n) # reshape to 3D
out = (a[None,:,:,:] + a.transpose(1,2,0)[:,:,None,:]).reshape(n*m,-1)

样本运行验证-

In [359]: # Setup input array
     ...: np.random.seed(0)
     ...: n,m = 3,4
     ...: A = np.random.randint(1,10,(n,n*m))

In [360]: # Original soln
     ...: out0 = np.tile(A, (m, 1)) + np.tile(A.T, (1, m))

In [361]: # Posted soln
     ...: n = A.shape[0]
     ...: m = A.shape[1]//n
     ...: a = A.reshape(n,m,n)
     ...: out = (a[None,:,:,:] + a.transpose(1,2,0)[:,:,None,:]).reshape(n*m,-1)

In [362]: np.allclose(out0, out)
Out[362]: True

大 n、m -

的时间

In [363]: # Setup input array
     ...: np.random.seed(0)
     ...: n,m = 100,100
     ...: A = np.random.randint(1,10,(n,n*m))

In [364]: %timeit np.tile(A, (m, 1)) + np.tile(A.T, (1, m))
1 loop, best of 3: 407 ms per loop

In [365]: %%timeit
     ...: # Posted soln
     ...: n = A.shape[0]
     ...: m = A.shape[1]//n
     ...: a = A.reshape(n,m,n)
     ...: out = (a[None,:,:,:] + a.transpose(1,2,0)[:,:,None,:]).reshape(n*m,-1)
1 loop, best of 3: 219 ms per loop

numexpr

进一步提升性能

我们可以利用 multi-core with numexpr module 处理大数据并提高内存效率和性能 -

import numexpr as ne

n = A.shape[0]
m = A.shape[1]//n
a = A.reshape(n,m,n)
p1 = a[None,:,:,:]
p2 = a.transpose(1,2,0)[:,:,None,:]
out = ne.evaluate('p1+p2').reshape(n*m,-1)

具有相同大 n、m 设置的计时 -

In [367]: %%timeit
     ...: # Posted soln
     ...: n = A.shape[0]
     ...: m = A.shape[1]//n
     ...: a = A.reshape(n,m,n)
     ...: p1 = a[None,:,:,:]
     ...: p2 = a.transpose(1,2,0)[:,:,None,:]
     ...: out = ne.evaluate('p1+p2').reshape(n*m,-1)
10 loops, best of 3: 152 ms per loop

二维数组的 Numpy einsum 外和

Numpy einsum outer sum of 2d-arrays

python

numpy

numpy-einsum