基于einsum的多数组运算加速及免分配大矩阵C的优化咨询

问题核心 #

当前代码的最大痛点是硬生成了(1024,1024,1000)的超大矩阵C，直接导致内存爆炸，但我们最终只需要C的偏迹结果——完全没必要先搞出这个大矩阵，直接合并运算步骤绕开它就行。

优化方向 #

1. 先平均再运算：原逻辑是先做张量积再对m索引取平均，反过来先对A、B的m维度取平均，再处理后续运算，能直接砍掉一个超大中间量。
2. 偏迹数学简化：利用张量积的迹性质，把10量子比特的偏迹拆解成小矩阵的运算，不用先生成完整的10量子比特矩阵。

实操代码 #

第一步：先对m维度取平均 #

先把A和B在m维度上做平均，得到两个小很多的矩阵：

import numpy as np

# A、B原形状(32,32,1000,10000)，平均后变成(32,32,1000)
A_avg = A.mean(axis=3)
B_avg = B.mean(axis=3)

第二步：直接计算偏迹结果 #

原代码里C[:,:,n_i]是A_avg[:,:,n_i]和B_avg[:,:,n_i]的张量积，而我们要的是这个张量积矩阵的前2个量子比特偏迹。结合量子态偏迹的数学性质，直接拆解运算：

方式一：循环实现（直观易懂）

from qutip import Qobj

# 保留的系统维度：2个量子比特
dim_keep = [2, 2]

def trace_5q_to_2q(mat):
    # 把5量子比特矩阵(32×32)转成(4,8,4,8)，对后3个量子比特求迹得到2量子比特矩阵(4×4)
    mat_reshaped = mat.reshape(4, 8, 4, 8)
    return np.trace(mat_reshaped, axis1=1, axis2=3)

result_list = []
for n_i in range(len(t_list)):
    # 取第n_i个时刻的矩阵
    a_mat = A_avg[:, :, n_i]
    b_mat = B_avg[:, :, n_i]
    
    # 对A的后3个量子比特求偏迹
    a_pt = trace_5q_to_2q(a_mat)
    # 对B的全部5个量子比特求迹（得到标量）
    b_tr = np.trace(b_mat)
    
    # 张量积得到最终的2量子比特约化矩阵
    reduced_mat = np.kron(a_pt, [[b_tr]])
    # 转为Qobj格式（按需保留）
    result_list.append(Qobj(reduced_mat, dims=[dim_keep, dim_keep]))

方式二：向量化实现（更快）

用numpy的广播和einsum实现无循环运算，适合大规模数据：

from qutip import Qobj

dim_keep = [2, 2]

# 把A_avg转成(4,8,4,8,1000)，对应2量子比特+3量子比特的结构
A_reshaped = A_avg.reshape(4, 8, 4, 8, -1)
# 对后3量子比特求迹，得到每个时刻的2量子比特矩阵(4,4,1000)
A_pt = np.trace(A_reshaped, axis1=1, axis2=3)
# 计算每个时刻B_avg的迹，得到(1000,)的标量数组
B_tr = np.trace(B_avg, axis1=0, axis2=1)
# 广播张量积，得到所有时刻的约化矩阵(4,4,1000)
reduced_mats = np.einsum('abn,n->abn', A_pt, B_tr)

# 批量转成Qobj
result_list = [Qobj(reduced_mats[:, :, n], dims=[dim_keep, dim_keep]) for n in range(len(t_list))]

效果对比 #

• 内存：完全跳过了(1024,1024,1000)的超大矩阵，内存占用从几十GB直接降到几百MB级别。
• 速度：避免了大规模张量积的内存拷贝和冗余计算，运算效率提升至少一个数量级。

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