如何使用np.einsum加速大规模时间序列矩阵逐日期乘法运算？

你的问题核心是Python循环的巨大开销——1500万次循环调用np.dot会彻底抵消底层运算优化的优势。下面是具体的提速方案，包括np.einsum实现和更简洁的替代方法：

1. 先将数据转为三维Numpy数组 #

必须彻底摆脱pandas循环索引的性能损耗，把A、B转为形状匹配的三维数组：
假设你的多级索引结构是[date, row, col]，转换代码如下：

# 转换A为(k, n, n)的三维数组
A_arr = np.stack([group.values.reshape(n, n) for _, group in A.groupby(level=0)])
# 转换B为(k, n, m)的三维数组
B_arr = np.stack([group.values.reshape(n, m) for _, group in B.groupby(level=0)])

注意替换n和m为你实际的矩阵维度。

2. 用np.einsum实现批量运算 #

直接通过einsum指定维度运算规则，完全避免循环：

AB_arr = np.einsum('knn,knm->knm', A_arr, B_arr)

维度规则说明：

• knn：对应A的三个维度（日期数k、矩阵行数n、矩阵列数n）
• knm：对应B的三个维度（日期数k、矩阵行数n、矩阵列数m）
• ->knm：输出结果的维度，每个位置对应单个日期的A[n×n] * B[n×m]乘积

3. 更简洁的替代方案：np.matmul或@运算符 #

Numpy原生支持批量矩阵乘法，用matmul或@运算符能达到和einsum几乎一致的性能，代码更直观：

# 方法1：np.matmul
AB_arr = np.matmul(A_arr, B_arr)
# 方法2：@运算符（Python 3.5+支持）
AB_arr = A_arr @ B_arr

这两种方法内部都是优化过的向量化运算，完全规避了Python循环的上下文切换开销。

性能提升原理 #

原列表推导式中，每次循环都要触发Python到C的上下文切换，1500万次切换的开销远大于矩阵乘法本身。而向量化方法直接将整个运算交给Numpy的底层C实现，一次性处理所有数据，运算效率能提升100~1000倍（取决于硬件配置）。

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