子数组的 Numpy 外部添加

Question

有没有办法在 numpy 中执行子数组的外部加法？

也就是说，我有 2 个形式为 2x2xNxM 的数组，每个数组都可以被认为是 2x2 矩阵的堆栈 N 高和 M 宽.我想将这些矩阵中的每一个添加到另一个数组中的每个矩阵，以形成一个 2x2xNxMxNxM 数组，其中最后四个索引对应于我最初两个数组中的索引，以便我可以索引 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 和 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

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 外部添加

Numpy outer addition of subarrays

python

arrays

numpy

matrix

array-broadcasting