如何在随机过程中向量化实现多轨迹下矩阵点积运算以优化性能

嗨，这个场景我太熟悉了！当处理大量批次的矩阵运算时，Python循环确实会拖慢速度，numpy的普通dot方法因为维度对齐问题没法直接用，不过咱们有几个高效的向量化方案能完美解决这个问题：

方法一：用np.einsum实现灵活的张量收缩 #

einsum可以说是处理这类自定义维度运算的神器，它通过字符串直接定义维度的运算关系，可读性和灵活性都拉满。针对你的场景，只需要一行代码：

C = np.einsum('ijk,ik->jk', B, A)

维度解释： #

• ijk对应B的三个维度：i是矩阵的行（2），j是矩阵的列（2），k是轨迹批次（Mmax）
• ik对应A的两个维度：i是向量的行（2），k是轨迹批次（Mmax）
• ->jk指定输出的维度：j对应结果的行（2），k对应轨迹批次（Mmax），正好和你要的C形状一致

这个操作会自动对每个批次k，计算B[:,:,k]和A[:,k]的点积，完全替代循环，而且是numpy内部优化的C级运算，速度比Python循环快几个数量级。

方法二：用np.matmul（或@运算符）配合维度扩展 #

如果你更习惯矩阵乘法的语法，可以给A扩展一个维度，让它和B的批次维度对齐，再用matmul自动广播：

# 给A增加最后一个维度，从(2, Mmax)变成(2, Mmax, 1)
C = np.matmul(B, A[..., np.newaxis]).squeeze(-1)
# 或者用@运算符，效果一样
C = (B @ A[..., np.newaxis]).squeeze(-1)

原理： #

• B的形状是(2, 2, Mmax)，A[..., np.newaxis]的形状是(2, Mmax, 1)
• matmul会自动将B的最后一个维度和A的倒数第二个维度对齐（都是Mmax），然后对每个批次执行2x2矩阵乘2x1向量的运算，得到(2, Mmax, 1)的结果
• 最后用squeeze(-1)去掉多余的最后一维，得到(2, Mmax)的C，和原循环结果完全一致

方法三：用np.tensordot做张量点积 #

如果需要更底层的张量运算控制，tensordot也能实现，不过可读性稍差一些：

C = np.tensordot(B, A, axes=([0], [0]))

这里指定B的第0维（2）和A的第0维（2）做内积，剩下的维度B[1]（2）和A[1]（Mmax）组合成输出的(2, Mmax)形状，效果和前两种方法一致。

验证结果一致性 #

为了确保向量化方法的正确性，你可以用小批量数据测试：

import numpy as np

Mmax = 10
# 生成测试数据
A = np.random.randn(2, Mmax).astype(np.complex64) + 1j * np.random.randn(2, Mmax).astype(np.complex64)
B = np.random.randn(2, 2, Mmax).astype(np.complex64) + 1j * np.random.randn(2, 2, Mmax).astype(np.complex64)

# 原循环方法
C_loop = np.zeros((2, Mmax), dtype=np.complex64)
for m in range(Mmax):
    C_loop[:, m] = B[:, :, m].dot(A[:, m])

# 验证einsum结果
C_einsum = np.einsum('ijk,ik->jk', B, A)
print("einsum和循环结果一致：", np.allclose(C_loop, C_einsum))

# 验证matmul结果
C_matmul = (B @ A[..., np.newaxis]).squeeze(-1)
print("matmul和循环结果一致：", np.allclose(C_loop, C_matmul))

运行后应该会输出两个True，说明结果完全匹配。

性能对比 #

当Mmax很大时（比如10^4甚至105），向量化方法的优势会非常明显：Python循环会因为解释器的开销变得很慢，而numpy的向量化运算直接调用底层优化的BLAS/LAPACK库，速度能提升几十到上百倍。

内容的提问来源于stack exchange，提问作者J.Agusti