Efficient scipy sparse array and kronecker product computation

Question

I want to compute the Kronecker product of a long list of small matrices (say, Pauli matrices). I tried to define the small matrices using scipy.sparse.csr_array, and use scipy.sparse.kron to perform the Kronecker product. However, the performance isn't ideal.


import numpy as np
from scipy import sparse
import functools
import time


sigma_x = sparse.csr_array(np.array([[0, 1], [1, 0]]))
sigma_y = sparse.csr_array(np.array([[0, -1j], [1j, 0]]))
sigma_z = sparse.csr_array(np.array([[1., 0], [0, -1.]]))

sigma = [sigma_x, sigma_y, sigma_z]
arguments = [sigma[i % 3] for i in range(3 * 5)]


start_time = time.time()

result = functools.reduce(sparse.kron, arguments)

end_time = time.time()
execution_time = end_time - start_time
print(execution_time)

=== Update 2 ===

Found a simple fix: I should use format='csr' keyword when calling sparse.kron. So I define

def kron_csr(x, y):
    return sparse.kron(x, y, format='csr')

and then

arguments = [sigma[i % 3] for i in range(3 * 5)]
result = functools.reduce(kron_csr, arguments)

would yield the result in just 0.005 second.

=== Update 1 ===

Following a comment below, I split the computation into two steps, first compute the Kronecker product of 3 * 4 Pauli matrices, then compute the Kronecker product of that result with the last 3 Pauli matrices,

start_time = time.time()
arguments = [sigma[i % 3] for i in range(3 * 4)] # reduces to 3 * 4
first = functools.reduce(sparse.kron, arguments)
second = functools.reduce(sparse.kron, [sigma_x, sigma_y, sigma_z])
result = functools.reduce(sparse.kron, [first, second])

end_time = time.time()
execution_time = end_time - start_time
print(execution_time)

or first compute the Kronecker product of sigma_x, sigma_y, sigma_z, then compute the Kronecker product of five of this intermediate matrices,

start_time = time.time()
result_1 = functools.reduce(sparse.kron, sigma)
result = functools.reduce(sparse.kron, [result_1 for i in range(5)])
end_time = time.time()
execution_time = end_time - start_time

The performance improves to around 4~9 seconds.

The execution time gives something like 11 seconds. The same computatoin using Mathematica only takes around 0.01 second,

Sigma = {SparseArray[( {
      {0, 1},
      {1, 0}
     } )], SparseArray[( {
      {0, -I},
      {I, 0}
     } )], SparseArray[( {
      {1, 0},
      {0, -1}
     } )]};

((KroneckerProduct @@ 
     Join @@ Table[{Sigma[[1]], Sigma[[2]], Sigma[[3]]}, {i, 5}]) // 
   Timing)[[1]]

I wonder how to improve the python code performence (hopefully to something like that in Mathematica)?

Why sparse? Have you comoared it to regukar dense numpy arrays? — hpaulj
– hpaulj, Commented Jun 17, 2024 at 2:54
@hpaulj Using np.array takes around the same time, 11 seconds. So sparse or not in python doesn't make much of a difference in this example. But in mathematica Sparse drastically improve performance. — Lelouch
– Lelouch, Commented Jun 17, 2024 at 3:08
functools.reduce(sparse.kron, arguments[:i]), the array increases by 4x for each increase in i, and has a consistent sparsity of 0.5. Evaluation time tends to increase by 3.5x for each step. By 15, you are getting a result in the Gb sizes. (` 2**30 ~ 1G` ) — hpaulj
– hpaulj, Commented Jun 17, 2024 at 5:08
@user24714692 oh what do you mean subproblem? Do you mean first generate the Kronecker product of 3 * 4 Pauli matrices, and then compute its product with 3 Pauli matrices? I tried this but the the time consumption stays roughly the same. — Lelouch
– Lelouch, Commented Jun 17, 2024 at 7:17
@user24714692 thanks for the tips, I added the attempt to the question, with different ways of breaking up the problem into smaller pieces. The performance does improve to around 5 seconds. — Lelouch
– Lelouch, Commented Jun 17, 2024 at 8:13

Aicody · Accepted Answer · 2024-06-17 12:48:12Z

0

You can solve the subproblems by which you can make it somewhat efficient:

import numpy as np
from scipy import sparse
import functools
import time


def _kron():
    sigma_x = sparse.csr_matrix(np.array(((0, 1), (1, 0))))
    sigma_y = sparse.csr_matrix(np.array(((0, -1j), (1j, 0))))
    sigma_z = sparse.csr_matrix(np.array(((1, 0), (0, -1))))

    sigma = (sigma_x, sigma_y, sigma_z)

    kron_csr = lambda x, y: sparse.kron(x, y, format='csr')

    def _product(sigma, repeat):
        arguments = (sigma[i % 3] for i in range(3 * repeat))
        return functools.reduce(kron_csr, arguments)

    start_time = time.time()
    result = _product(sigma, 5)
    end_time = time.time()
    execution_time = end_time - start_time
    return f"Execution time: {execution_time} seconds"

print(_kron())

answered Jun 17, 2024 at 12:48

Aicody

4,9063 gold badges8 silver badges30 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Efficient scipy sparse array and kronecker product computation

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related