在 numpy 中,有没有一种方法可以执行相当于子数组外部添加的操作?
也就是说,我有 2 个 2x2xNxM
形式的数组,每个数组都可以被视为 2x2
矩阵 N
高的堆栈和 M
宽。我想将这些矩阵中的每一个添加到另一个数组中的每个矩阵,以形成一个 2x2xNxMxNxM 数组,其中最后四个索引对应于我最初的两个数组中的索引,以便我可以索引 输出[:,:,x1,y1,x2,y2] == a1[:,:,x1,y1] + a2[:,:,x2,y2]
。
如果这些是标量数组,那就很简单了,我所要做的就是:
A, B = a.ravel(), b.ravel()
four_D = (a[...:np.newaxis] + b).reshape(*a1.shape, *a2.shape)
for (x1, y1, x2, y2), added in np.ndenumerate(four_D):
assert added == a1[x1,y1] + a2[x2,y2]
但是,这对于 a
和 b
由矩阵组成的情况不起作用。当然,我可以使用嵌套的 for 循环,但我的数据集将相当大,并且我希望在多个数据集上运行它。
有没有有效的方法来做到这一点?
最佳答案
扩展数组以具有更多维度,然后利用 broadcasting
-
output = a1[...,None,None] + a2[...,None,None,:,:]
示例运行 -
In [38]: # Setup input arrays
...: N = 3
...: M = 4
...: a1 = np.random.rand(2,2,N,M)
...: a2 = np.random.rand(2,2,N,M)
...:
...: output = np.zeros((2,2,N,M,N,M))
...: for x1 in range(N):
...: for x2 in range(N):
...: for y1 in range(M):
...: for y2 in range(M):
...: output[:,:,x1,y1,x2,y2] = a1[:,:,x1,y1] + a2[:,:,x2,y2]
...:
...: output1 = a1[...,None,None] + a2[...,None,None,:,:]
...:
...: print np.allclose(output, output1)
True
关于python - Numpy 子数组的外部添加,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/47627934/