在 numpy 中将 2D 数组广播到 4D 数组

Question

我有一个形状 (48, 7) 的二维数组 x 和一个形状 (48, 7, 48, 7) 的 4 维数组 T。当我乘以 x * T 时，python 广播维度，但不是我预期的方式（实际上，我不明白它是如何广播的）。以下循环将实现我想要的：

for i in range(48):
    for j in range(7):
        Tx[i, j, :, :] = x[i, j] * T[i, j, :, :]

其中 Tx 是形状为 (48, 7, 48, 7) 的数组。我的问题是，有没有办法使用广播实现相同的结果？

Answer 1

我找到了解决方案。 Python 从最右边的维度广播并向左传播 (source)。通过转置前两个维度和后两个维度：

T = np.transpose(T, (2,3,0,1))

然后它将按照我预期的方式进行广播。之后，结果数组可以再次转置以恢复原始形状：

Tx = x*T
Tx = np.transpose(Tx, (2,3,0,1))

Answer 2

广播对齐尾随维度。换句话说，x * Tx 是这样做的：

for i in range(48):
    for j in range(7):
        Tx[:, :, i, j] = x[i, j] * T[:, :, i, j]

要使前导维度对齐，请将单位维度添加到 x:

Tx = x[..., None, None] * T

或者，您可以使用 np.einsum 明确指定维度：

Tx = np.einsum('ij,ij...->ij...', x, T)

Broadcasting 2D array to 4D array in numpy