如何优化涉及稀疏矩阵的Python函数计算效率？

我现在想在Python里实现一个涉及稀疏矩阵的特定函数，并且希望它尽可能快。下面我会详细说明问题是什么，以及目前我是怎么实现的。

问题描述 #

我有固定的N=1000个稀疏矩阵（维度和元素都固定），统称为B，每个矩阵大小是1000x1000（平均稀疏度，也就是非零元素占总元素的比例，是0.0001）。对于给定的向量u（大小为1000），我需要为每个j=0,...,999计算c[j] = u @ B[j] @ u，输出是numpy数组c。也就是说，稀疏矩阵B[j]（存储在元组B中）是固定的，u是函数的输入。

现有实现 #

我预先计算了所有矩阵，并把它们作为全局变量在程序中使用。我决定用scipy.sparse的csr_matrix格式存储它们（我看到资料说这种格式最适合计算矩阵-向量乘积），放在一个元组里。计算函数的代码如下：

# 预先计算矩阵，作为全局固定变量
B = []
for j in range(1000):
   sparse_matrix = ...compute it... # 大小为1000x1000的稀疏csr_matrix
   B.append(sparse_matrix)
B = tuple(B)

def func(u: np.ndarray) -> np.ndarray:
    """
    u是长度为N=1000的np.array
    """
    return np.array([u.dot(B[k].dot(u.transpose())) for k in range(len(B))])

我的疑问 #

这已经是能达到的最高效率了吗？还是说在速度方面还有优化空间？比如改变矩阵的存储结构，或者改变计算所有向量-矩阵-向量乘积的方式？另外，这个计算有没有并行化的潜力？如果有的话，有没有推荐的库或函数可以参考？

提前感谢大家的帮助！

补充说明 #

这里的c[j] = u @ B[j] @ u，我的意思是先计算B[j]和u的矩阵-向量乘积，再和u做内积。（数学上就是u.transposed() * B * u）

给大家提供一个小基准测试程序，我在里面创建了一些随机稀疏矩阵，并在随机向量上测试：

import numpy as np
from random import randint
from scipy.sparse import coo_matrix, csr_matrix
from time import time

# 创建随机稀疏矩阵元组
N = 1000
n = 100
B = []
for i in range(N):
    data = np.random.uniform(-1, 1, n).tolist()
    rows = [randint(0, N-1) for _ in range(n)]
    cols = [randint(0, N-1) for _ in range(n)]
    sparse_matrix = csr_matrix(coo_matrix((data, (rows, cols)), shape=(N, N)))
    B.append(sparse_matrix)
B = tuple(B)

# 我的函数
def func(u: np.ndarray) -> np.ndarray:
    """
    u是长度为N=1000的np.array
    """
    return np.array([u.dot(B[k].dot(u.transpose())) for k in range(len(B))])

# 用于测试的随机向量
u = np.random.uniform(-1, 1, N)

START = time()
func(u)
END = time()
print(f"Speed : {END - START}")

运行结果：

>>> Speed : 0.005256175994873047

备注：内容来源于stack exchange，提问作者User341562