Question

有没有办法在numpy中执行相当于外部添加子阵列的方法？

也就是说，我有2个2x2xNxM形式的数组，每个数组可能被认为是2x2矩阵N高和M宽的堆栈。我想将这些矩阵中的每一个添加到另一个数组的每个矩阵中，以形成一个2x2xNxMxNxM数组，其中最后四个索引对应于我最初的两个数组中的索引，以便我可以索引{{1} }。

如果这些是标量数组，那将是微不足道的，我所要做的就是：

output[:,:,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由矩阵组成的情况。当然，我可以使用嵌套的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

Answer 2

与插入额外轴的标量相同，也适用于更高维数组（这称为broadcasting）：

import numpy as np

a1 = np.random.randn(2, 2, 3, 4)
a2 = np.random.randn(2, 2, 3, 4)
added = a1[..., np.newaxis, np.newaxis] + a2[..., np.newaxis, np.newaxis, :, :]

print(added.shape)  # (2, 2, 3, 4, 3, 4)

Numpy外部添加子阵列

2 个答案: