Ниже приведена улучшенная версия вашего кода:
import numpy as np
from scipy.sparse import csr_matrix
def vec(A):
m, n = A.shape[0], A.shape[1]
return A.reshape(m*n, order='F')
def commutation_matrix_sp(A):
m, n = A.shape[0], A.shape[1]
row = np.arange(m*n)
col = row.reshape((m, n), order='F').ravel()
data = np.ones(m*n, dtype=np.int8)
K = csr_matrix((data, (row, col)), shape=(m*n, m*n))
return K
Тест:
A = np.random.rand(500, 500)
K = commutation_matrix_sp(A)
print(f'{K.data.nbytes/2**20:.2f} MB')
# 0.24 MB
print(np.all(K @ vec(A) == vec(A.T)))
# True