numpy - 如何通过固定索引将 4D numpy 数组转换为 2D

Question

我有一个 4D numpy 数组 A形状(N,N,N,N)我想转换为二维矩阵 M形状(N,N)通过固定索引对。例如

M[i,j] = A[i,j,i,j]

这应该如何在 numpy 中完成，避免 for 循环？
编辑:
我随后将访问 M 的元素使用 numpy.ix_ 提供的索引数组所以以类似的方式访问 4D 数组的元素也是一个解决方案。

Answer 1

这是一种解决方法:

i, j = np.arange(N), np.arange(N)
j_idx, i_idx = np.meshgrid(i, j)
M = A[i_idx, j_idx, i_idx, j_idx]

用途 meshgrid预先生成索引模式，然后 fancy indexing数组 A获取 M .正如@hpaulj 所建议的，您可以指定 sparse = True在 np.meshgrid()获得可广播的一维数组而不是完整的二维索引数组以节省一些空间。
您也可以使用 np.ix_() 执行此操作还有:

ixgrid = np.ix_(i, j)
M = A[ixgrid + ixgrid]

自 ixgrid是一个二元组， ixgrid + ixgrid生成索引所需的 4 元组 A .