Эффективное вычисление разреженного массива и произведения Кронекера

Я хочу вычислить произведение Кронекера длинного списка небольших матриц (скажем, матриц Паули). Я пытался определить маленькие матрицы, используя scipy.sparse.csr_array, и использовать scipy.sparse.kron для выполнения произведения Кронекера. Однако производительность не идеальна.

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)

=== Обновление 2 ===

Нашел простое решение: мне следует использовать ключевое слово format='csr' при вызове sparse.kron. Поэтому я определяю

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

а потом

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

даст результат всего за 0,005 секунды.

=== Обновление 1 ===

Следуя комментарию ниже, я разделил вычисление на два этапа: сначала вычислил произведение Кронекера матриц Паули 3 * 4, затем вычислил произведение Кронекера этого результата с последними 3 матрицами Паули:

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)

или сначала вычислите произведение Кронекера sigma_x, sigma_y, sigma_z, затем вычислите произведение Кронекера пяти из этих промежуточных матриц,

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

Производительность увеличивается примерно до 4–9 секунд.

Время выполнения дает что-то около 11 секунд. Те же вычисления с использованием Mathematica занимают всего около 0,01 секунды.

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]]

Интересно, как улучшить производительность кода Python (надеюсь, до чего-то подобного в Mathematica)?

🤔 А знаете ли вы, что...
С Python можно создавать кросс-платформенные приложения для Windows, macOS и Linux.