如何以无Python循环方式构建(...,N,N)形状的NumPy数组

问题场景 #

给定一个返回N×N子数组的函数，我们需要基于任意形状（比如示例中的(M,M)）的输入数组pi，构建形状为(...,N,N)的输出数组，要求完全避免嵌套Python循环以保证效率。以N=2为例，期望输出数组y满足y[i,j,k,l] = h0(pi[i,j])[k,l]，初始尝试的代码如下：

import numpy as np

a = 1.
b = 2 + 3.j
def h0(pi):
    return np.array([
        [a,        b * pi ],
        [b * pi,   a      ]
    ])

原代码问题 #

直接按上述方式构建数组会出现维度不匹配：常量a无法与带前导维度的pi数组正确广播，生成的数组形状不符合预期，无法对应每个pi[i,j]对应的2×2子矩阵。

解决方案 #

将与pi无关的常量a通过ones_like广播至与pi相同的形状，再通过转置调整维度顺序，最终得到目标形状的数组。代码实现如下：

import numpy as np

a = 1.
b = 2 + 3.j
def h0(pi):
    one = np.ones_like(pi)
    return np.array([
        [a * one, b * pi ],
        [b * pi,  a * one]
    ]).transpose((2,3,0,1))

当输入pi为二维(M,M)数组时，输出数组形状为(M,M,2,2)，完全符合需求，且全程依赖NumPy向量化操作，彻底规避了Python循环的性能开销。

内容的提问来源于stack exchange，提问作者q_mat