如何向量化实现numpy.linalg.lstsq操作以优化多频相位解包裹

我正在使用Python3和NumPy实现多频相位解包裹算法。现有7张单通道（灰度）图像，每张形状为(1080, 1920)，沿第三轴堆叠后得到形状为(1080, 1920, 7)的数组。存在一个固定的7×7位移矩阵A，每个像素对应形状为(1,7)的不同强度数组r。为通过最小化L2范数||r - Au||求解每个像素的u，可执行如下示例代码：

# 示例代码
A = np.random.randn(7, 7)
r = np.random.randn(7, 1)

# 针对单个像素求解
u = np.linalg.lstsq(a=A, b=r, rcond=None)

目前通过两层循环实现，但效率低下，求高效的NumPy实现方案。

核心思路：利用固定矩阵的伪逆 + 批量广播运算 #

由于矩阵A是固定不变的，无需对每个像素重复执行最小二乘求解。最小二乘问题||r - Au||的解等价于u = A⁺r，其中A⁺是A的Moore-Penrose伪逆。预先计算一次伪逆后，即可通过批量矩阵乘法一次性处理所有像素。

具体代码实现 #

1. 预先计算A的伪逆：

import numpy as np

# 固定的7×7矩阵
A = np.random.randn(7, 7)
# 计算伪逆（仅需执行一次）
A_pinv = np.linalg.pinv(A)

2. 批量处理所有像素：
   假设堆叠后的图像数组为img_stack（形状(1080, 1920, 7)），直接利用NumPy的广播机制完成批量运算：

# 假设img_stack是堆叠后的图像数组，形状(1080, 1920, 7)
u_stack = img_stack @ A_pinv
# 结果u_stack的形状为(1080, 1920, 7)，每个位置对应像素的解u

备选方案：批量调用lstsq #

如果偏好直接使用np.linalg.lstsq，可将所有像素的r向量拼接成二维矩阵，一次性求解后再恢复形状：

# 将图像数组展平为(N, 7)，N=1080*1920为总像素数
r_flat = img_stack.reshape(-1, 7).T
# 批量求解所有像素的u
u_flat, _, _, _ = np.linalg.lstsq(A, r_flat, rcond=None)
# 将结果恢复为原图像形状
u_stack = u_flat.T.reshape(1080, 1920, 7)

效率说明 #

两种方案均避免了Python层面的循环，完全依赖NumPy底层的BLAS/LAPACK优化实现，运算速度比两层循环快数十到上百倍，尤其适合大尺寸图像的处理。

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