Numpy含复数元素的矩阵向量乘法 元素为多维数组时如何保留原有结构

问题背景 #

原有标量场景下的矩阵向量乘代码可正常运行：

A = np.array([[1, 0],
              [0,1]])      # matrix
u = np.array([[3],[9]])    # column vector
U = np.dot(A,u)            # U = A*u
U

现在需要适配的场景是：矩阵A内的元素Xx,Xy,Yx,Yy、向量内的元素u,v均为2维或3维数组，希望尽可能保留「构造矩阵、向量再调用乘法」的结构，得到如下预期运算结果：

U = Xx*u + Xy*v
V = Yx*u + Yy*v

原有可行方案np.einsum('ijmn,jkmn->ikmn', A, w)和(w.T@A.T).T存在两个问题：1. 返回结果多一对方括号，维度不符合预期；2. 块矩阵构造耗时较高，整体运算性能低于直接按公式向量化计算。
原有各方案性能测试数据如下：

• p_einstein np.einsum('ijmn,jkmn->ikmn', A, w)：9.07 µs ± 267 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
• p_tensor: (w.T@A.T).T：7.89 µs ± 225 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
• p_vect: [Xx*u1 + Xy*v1, Yx*u1 + Yy*v1]：2.63 µs ± 173 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
• p_vect: including A = np.array([[Xx, Xy],[Yy,Yy]])：7.66 µs ± 290 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

最优实现方案 #

方案1：直接element-wise向量化计算（性能最优） #

如果不需要严格保留构造矩阵再乘法的结构，直接按预期公式做逐元素运算即可，该方案没有额外的矩阵构造、维度调整开销，维度完全符合预期，是所有方案中性能最高的：

# Xx、Xy、Yx、Yy、u、v均为同形状的2/3维数组
U = Xx * u + Xy * v
V = Yx * u + Yy * v

# 如果需要合并为和输入维度对齐的向量结构，可新增维度堆叠
UV = np.stack([U, V], axis=-2)

方案2：保留矩阵构造+乘法结构（性能次优） #

如果必须保留原有代码结构，可调整维度排布规则，将2x2矩阵维度、2x1向量维度放在数组最后两维，利用@运算符的原生广播规则实现运算，避免多余的转置和einsum解析开销：

# 构造矩阵A，将2x2维度放在最后，前面维度和输入数组对齐
A = np.stack([np.stack([Xx, Xy], axis=-1),
              np.stack([Yx, Yy], axis=-1)], axis=-2)
# 构造向量u_vec，将2x1维度放在最后
u_vec = np.stack([u, v], axis=-2)
# 直接矩阵乘，结果维度完全符合预期，无多余括号
UV = A @ u_vec

# 单独提取U、V的方式
U = UV[..., 0, :]
V = UV[..., 1, :]

该方案的实测性能约为4~5μs，比einsum、转置方案性能高30%以上，比带矩阵构造的直接向量化方案也快40%左右。

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