Python中高效计算稀疏矩阵希尔伯特-施密特内积的最优方法

我需要计算两个矩阵的希尔伯特-施密特内积（Hilbert-Schmidt Inner Product）。该内积的定义为矩阵A的厄米共轭与B的乘积的迹，即：
$$Tr(A^\dagger B) = \sum_{ij} (A_{ij}^* B_{ij})$$
其中$A_{ij}^*$是$A_{ij}$的共轭复数，i、j分别为行、列索引。

从公式可以看出，理论上我们只需要对两个矩阵共有的非零元素进行运算即可，不需要计算完整的矩阵乘积（非对角元素的计算属于冗余操作）。但实际测试中，我的手动实现反而比直接计算完整乘积再取迹的方法慢很多，想知道针对稀疏矩阵有没有更优的计算方式。

现有实现代码 #

手动筛选非零元素的实现 #

def hilbert_schmidt_inner_product(mat1, mat2):
    # 找到每个矩阵的非零元素索引对
    mat1_ij = set([tuple(x) for x in np.array(list(zip(*mat1.nonzero())))])
    mat2_ij = set([tuple(x) for x in np.array(list(zip(*mat2.nonzero())))])
    
    # 找出两个矩阵都有非零元素的索引对
    common_ij = np.array(list(mat1_ij & mat2_ij))
    
    # 提取对应位置的元素
    mat1_survied = np.array(mat1[common_ij[:,0], common_ij[:,1]])[0]
    mat2_survied = np.array(mat2[common_ij[:,0], common_ij[:,1]])[0]
    
    # 计算共轭点积得到结果
    trace = np.dot(mat1_survied.conj(), mat2_survied)
    return trace

直接计算矩阵乘积取迹的实现 #

import numpy as np
sum((mat1.conj().T@mat2).diagonal())

基准测试 #

测试代码：

import numpy as np
from scipy.sparse import rand

Dimension = 2**12

A = rand(Dimension, Dimension, density=0.001, format='csr')
B = rand(Dimension, Dimension, density=0.001, format='csr')

测试结果：

%timeit hilbert_schmidt_inner_product(A,B)
49.2 ms ± 3.13 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit sum((A.conj().T@B).diagonal())
1.48 ms ± 32 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit np.einsum('ij,ij->', A.conj().todense(), B.todense())
53.9 ms ± 2.74 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

可以看到，手动筛选非零元素的方法比直接计算乘积取迹慢了一个数量级，明明前者理论上应该更高效，想请教有没有针对稀疏矩阵的更优实现方式？

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