python - Numpy 子数组的外部添加

Question

在 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 循环，但我的数据集将相当大，并且我希望在多个数据集上运行它。

有没有有效的方法来做到这一点？

Answer 1

扩展数组以具有更多维度，然后利用 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