在两个 Ndarrays 之间执行广播二进制操作的 Pythonic 方法

Pythonic way to perform broadcasting binary operations between two Ndarrays

标题说的是什么。我正在尝试执行以下任务:

1) 在两个二维 numpy 数组 A1A2 之间执行二元运算(如 +、-、*、/、>、<) 其中 A1.shape = (N1, N2_1)A2.shape = (N1, N2_2) 生成一个 3-d numpy 数组,其形状为 (N, N2_1, N2_2)

2) 在两个 3-d numpy 数组 A1A2 之间执行二元运算(如 +、-、*、/、>、<) 其中 A1.shape = (N1, N2, N3_1)A2.shape = (N1, N2, N3_2) 生成形状为 (N, N2, N3_1, N3_2)

的 4 维 numpy 数组

我发现自己在编写非常不符合 Python 规范的代码来完成第一个任务 1)。如果有人可以向我展示正确的 pythonic 代码示例来完成它,我将不胜感激。


我的尝试:

数据来源:

import numpy as np
n_row = 10000
n_col_a1 = 3
n_col_a2 = 4

a1 = np.tile(np.arange(n_col_a1), (n_row, 1))
a2 = np.tile(np.arange(n_col_a2), (n_row, 1))

我的 unpythonic 广播尝试:

X1 = np.broadcast_to(a1, (n_col_a2, *a1.shape))
X1 = np.moveaxis(X1, 0, -1)

X2 = np.broadcast_to(a2, (n_col_a1, *a2.shape))
X2 = np.moveaxis(X2, 0, -2)

result_a1_minus_a2 = X1 - X2
print(result_a1_minus_a2)

两个输入数组必须展开为:

(N1, N2_1)     => (N1, N2_1, 1)
(N1, N2_2)     => (N1, 1,    N2_2)
(N1, N2_1, N2_2)

例如

A1[:, :, None] * A2[:, None, :]

同样:

(N1, N2, N3_1)
(N1, N2, N3_2)
(N1, N2, N3_1, N3_2)

A1[...,None] * A2[:,:,None,:]