优化NumPy中大规模矩阵S1的构建效率

需要构建尺寸为num×num的矩阵S1，其中num=d²（通常d极大，比如d=81时num=6561）。当前实现通过嵌套循环计算每个元素：

S1 = np.zeros((num, num), dtype=np.complex128)
for i in range(num):
    for j in range(num):
        S1[i][j] = np.trace(np.dot(b1[i], np.dot(b1[j], state)))

其中：

• b1是长度为d²的d×d稀疏厄米矩阵列表，满足希尔伯特-施密特内积正交基，生成函数如下：

def basis(n):
    list2=[]
    list3=[]
    list4=[]
    list5=[]
    list6=[]

    for i in range(0,n):
        for j in range(i+1,n):
            m1=np.zeros((n,n),dtype=np.complex128)
            m1[i,j]=1j
            m1[j,i]=1j
            list2.append(m1/np.sqrt(2))
            
            m2=np.zeros((n,n),dtype=np.complex128)
            m2[i,j]=-1
            m2[j,i]=1
            list3.append(m2/np.sqrt(2))
            
    for i in range(0,n-1):
        m3=np.zeros(n,dtype=np.complex128)
        m3[i]=1j
        m3[i+1]=-1j
        list4.append(m3)
   
        org=np.linalg.qr(np.array(list4).T)[0]
        l1=org[:,i]
        list5.append(l1)
 
        l2=np.diag(list5[i])
        list6.append(l2)
    
    return [np.identity(n)/np.sqrt(n)]+[mat*(-1j) for mat in (list2+list3+list6)]

• state为d×d稀疏密度矩阵。

当前嵌套循环需执行num²次运算，当d=81时速度极慢，希望利用稀疏性优化代码。

1. 数学简化：利用希尔伯特-施密特内积特性 #

由于b1中的矩阵是厄米矩阵（b1[i]^† = b1[i]），state也是厄米矩阵，原计算式可简化为：

S1[i][j] = trace(b1[i] · b1[j] · state) = ⟨b1[i], state · b1[j]⟩_HS

其中⟨A,B⟩_HS是希尔伯特-施密特内积，等价于两个矩阵对应元素相乘后求和——稀疏矩阵的非零元素占比极低，这一步能大幅减少无效运算。

2. 改用稀疏矩阵库操作 #

放弃numpy密集矩阵，改用scipy.sparse的稀疏矩阵类型（如csr_matrix）存储b1和state，矩阵乘法、元素相乘等操作会自动跳过零元素，从底层减少计算量。

3. 优化循环逻辑+并行加速 #

具体实现步骤 #

（1）转换为稀疏矩阵

先把所有矩阵转为稀疏格式，避免密集矩阵的内存浪费和无效计算：

from scipy.sparse import csr_matrix

# 转换state为稀疏矩阵
state_sparse = csr_matrix(state)
# 转换b1列表为稀疏矩阵列表
b1_sparse = [csr_matrix(mat) for mat in b1]

（2）预计算中间结果

提前计算所有state · b1[j]的乘积，避免循环中重复计算：

temp_list = [state_sparse @ mat for mat in b1_sparse]

（3）稀疏内积填充S1

用稀疏矩阵的元素相乘+求和替代原有的trace运算，再通过并行计算加速行遍历：

import numpy as np
from joblib import Parallel, delayed

num = len(b1_sparse)

def compute_row(i):
    bi = b1_sparse[i]
    row = np.zeros(num, dtype=np.complex128)
    for j in range(num):
        # 希尔伯特-施密特内积：仅计算非零元素的乘积和
        row[j] = (bi.multiply(temp_list[j])).sum()
    return row

# 并行计算所有行，n_jobs设为CPU核心数
S1_rows = Parallel(n_jobs=-1)(delayed(compute_row)(i) for i in range(num))
S1 = np.array(S1_rows)

4. 优化基矩阵生成（可选） #

原basis函数生成密集矩阵，可修改为直接生成稀疏矩阵，减少内存占用和后续转换开销：

from scipy.sparse import csr_matrix, identity

def basis_sparse(n):
    list2=[]
    list3=[]
    list6=[]

    for i in range(0,n):
        for j in range(i+1,n):
            # 生成m1的稀疏矩阵
            data = [1j, 1j]
            rows = [i, j]
            cols = [j, i]
            m1 = csr_matrix((data, (rows, cols)), shape=(n,n), dtype=np.complex128)
            list2.append(m1 / np.sqrt(2))
            
            # 生成m2的稀疏矩阵
            data = [-1, 1]
            rows = [i, j]
            cols = [j, i]
            m2 = csr_matrix((data, (rows, cols)), shape=(n,n), dtype=np.complex128)
            list3.append(m2 / np.sqrt(2))
            
    list4=[]
    list5=[]
    for i in range(0,n-1):
        m3=np.zeros(n,dtype=np.complex128)
        m3[i]=1j
        m3[i+1]=-1j
        list4.append(m3)
   
        org=np.linalg.qr(np.array(list4).T)[0]
        l1=org[:,i]
        list5.append(l1)
 
        # 生成对角稀疏矩阵l2
        data = l1
        rows = np.arange(n)
        cols = np.arange(n)
        l2 = csr_matrix((data, (rows, cols)), shape=(n,n), dtype=np.complex128)
        list6.append(l2)
    
    # 单位矩阵转为稀疏矩阵
    id_mat = identity(n, dtype=np.complex128) / np.sqrt(n)
    return [id_mat] + [mat*(-1j) for mat in (list2+list3+list6)]

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