[Python] How to improve efficiency of my python function involving sparse matrices?

I want to implement a particular function in python involving sparse matrices and want to make it as fast as possible. I describe in detail what the problem is and how I implemented it so far.

Problem: I have N=1000 sparse matrices of size 1000x1000 (average sparsity, i.e. number of non-zero entries over all entries, is 0.0001). For given vectors a and b (of size 1000) I want to compute c[j] = a @ B[j] @ b for each j=0,...,999 and the output should be the numpy array c.

My implementation so far: I can precompute all the matrices and I decided to safe them as scipy.sparse matrices in csr_matrix format (I read that this is the best format when I just want to calculate matrix vector products) in a Tuple. To calculate my desired function I do

def func(a: np.ndarray, b: np.ndarray, sparse_matrices: Tuple[csr_matrix]) -> np.ndarray:
"""
a and b are np.arrays of length N
sparse_matrices is precomputed tuple of csr_matrices
"""
return np.array([a.dot(sparse_matrices[k].dot(b.transpose())) for k in range(len(sparse_matrices))])


Question: Is this the most efficient it can get, or do you see room for improvement regarding speed, e.g. change the structure how I safe the matrices or how I compute all the vector-matrix-vector products? Also, do you see potential for parallelising this computation? If so, do you have a suggestion what libraries/functions I should look into?

Thanks in advance!

Edit: With c[j] = a @ B[j] @ b I mean that I want to compute the matrix-vector product of B[j] and b and then the inner product with a. (Mathematically a.transposed() * B * b)

If it is of help to anyone. Here is a small benchmark program where I create some random sparse matrices and evaluate it on some random vector.

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

# My function
def func(a: np.ndarray, b: np.ndarray, sparse_matrices: Tuple[csr_matrix]) -> np.ndarray:
"""
a and b are np.arrays of length N
sparse_matrices is precomputed tuple of csr_matrices
"""
return np.array([a.dot(sparse_matrices[k].dot(b.transpose())) for k in range(len(sparse_matrices))])

# Create random tuple of sparse matrices
N = 1000
n = 100
sparse_matrices = []
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)]
B = csr_matrix(coo_matrix((data, (rows, cols)), shape=(N, N)))
sparse_matrices.append(B)
sparse_matrices = tuple(sparse_matrices)

# random vector to evaluate function
u = np.random.uniform(-1, 1, N)

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


>>> Speed : 0.005247831344604492

Continue reading...