如何向量化Numpy数组的PyOLS运算？优化三维数组循环

我有一个三维Numpy数组y_{n×m×p}和一维Numpy数组x_{p×1}，需要对y的每个切片y[i,j,:]与x执行带截距的OLS回归，提取第一个回归系数。当前用嵌套循环实现，但数组规模极大时效率极低，怎么通过向量化操作优化？

原示例代码 #

import numpy as np
def PyOLS(xvec, yvec):
   n = xvec.shape[0]
   X = np.c_[xvec, np.ones(n)]
   betas = np.dot(np.linalg.inv(np.dot(X.T, X)), np.dot(X.T, yvec))
   return betas   

n = 20
m = 20
p = 7
x = np.random.rand(n, m, p)
y = np.random.rand(p)

res = np.zeros((n, m))

for i in range(n):
    for j in range(m):
       cx = x[i, j, :]
       b = PyOLS(cx, y)
       res[i, j] = b[0]

向量化优化思路 #

先拆解原PyOLS的计算逻辑：带截距的OLS回归中，第一个系数的解析公式可以直接推导出来，不用依赖矩阵求逆，同时利用Numpy的轴运算一次性完成所有切片的统计量计算。

推导后的第一个回归系数公式：
[
b_0 = \frac{p \cdot \sum(x_i y_i) - \sum(x_i) \cdot \sum(y_i)}{p \cdot \sum(x_i^2) - (\sum(x_i))^2}
]
其中p是一维数组的长度，所有求和操作针对每个y[i,j,:]切片执行。

向量化实现代码 #

import numpy as np

n = 20
m = 20
p = 7
x = np.random.rand(n, m, p)
y = np.random.rand(p)

# 沿第三维(axis=2)批量计算所有切片的统计量
sum_x = x.sum(axis=2)
sum_x2 = (x ** 2).sum(axis=2)
sum_xy = (x * y).sum(axis=2)
sum_y = y.sum()

# 代入公式批量计算结果
numerator = p * sum_xy - sum_x * sum_y
denominator = p * sum_x2 - sum_x ** 2
res_vectorized = numerator / denominator

# 验证与原循环结果一致（可选）
# print(np.allclose(res, res_vectorized))  # 输出True

优化优势 #

• 速度提升：Numpy的向量化操作基于C底层实现，彻底避免Python循环的开销，数组规模越大，提升越明显；
• 数值稳定性：用解析公式替代矩阵求逆，减少数值计算误差；
• 代码简洁：无需嵌套循环，逻辑更清晰。

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