Efficient iteration over kron products

I would like to get a more compact, and efficient, algorithm for what I am doing. My task is to take the kron product of a matrix with the identity multiple times.

Right now I am doing something that, while it works, I cannot scale it up.

import scipy.sparse as sp
from scipy.sparse import csr_matrix
A=[[0., 0.],
   [1., 0.]]
Id=sp.identity(2)
A_1=csr_matrix(sp.kron(A,sp.kron(Id,sp.kron(Id,Id))))
A_2=csr_matrix(sp.kron(Id,sp.kron(A,sp.kron(Id,Id))))
A_3=csr_matrix(sp.kron(Id,sp.kron(Id,sp.kron(A,Id))))
A_4=csr_matrix(sp.kron(Id,sp.kron(Id,sp.kron(Id,A))))

Notice how for each A_i, the A matrix is moved one space to the right. While this works, at some point if I need to do it for the A_20, it would require writing down an extremely long sequence of sp.kron.

That is why I would like to get something like:

n_q=20
A_list=[]
for i in range(i,n_q):
    A_i=csr_matrix(sp.kron(Id,...sp.kron(A...,sp.kron(Id,Id))))
    A_list.append(A_i)

Does anybody know if something like this is possible?

Solution

Here's an trial idea.

Define a function to chain some kron calls from a list of matrices:

In [50]: def foo(alist):
    ...:     M = alist[0]
    ...:     for a in alist[1:]:
    ...:         M = sparse.kron(M,a)
    ...:     return M
    ...:

Your two matrices:

In [51]: A=sparse.csr_matrix([[0., 0.],
    ...:    [1., 0.]])
    ...: Id=sparse.identity(2)

Make a list, in this example, 3 id, and 2 after A:

In [52]: alist = [Id]*3+[A]+[Id]*2

In [53]: alist
Out[53]: 
[<2x2 sparse matrix of type '<class 'numpy.float64'>'
    with 2 stored elements (1 diagonals) in DIAgonal format>,
 <2x2 sparse matrix of type '<class 'numpy.float64'>'
    with 2 stored elements (1 diagonals) in DIAgonal format>,
 <2x2 sparse matrix of type '<class 'numpy.float64'>'
    with 2 stored elements (1 diagonals) in DIAgonal format>,
 <2x2 sparse matrix of type '<class 'numpy.float64'>'
    with 1 stored elements in Compressed Sparse Row format>,
 <2x2 sparse matrix of type '<class 'numpy.float64'>'
    with 2 stored elements (1 diagonals) in DIAgonal format>,
 <2x2 sparse matrix of type '<class 'numpy.float64'>'
    with 2 stored elements (1 diagonals) in DIAgonal format>]

And pass it to foo:

In [54]: M = foo(alist)

In [55]: M
Out[55]: 
<64x64 sparse matrix of type '<class 'numpy.float64'>'
    with 512 stored elements (blocksize = 2x2) in Block Sparse Row format>

Then it's just a matter of iterating over the possible lists of Id/A mixes, with different numbers of before and after.

It saves typing if not speed.

In numpy the main way to make 'iteration' more efficient is to move it from python to compiled code, that is, by using the builtin whole-array methods. Or by using compilers like numba.

Scipy sparse works on top of numpy, with most of the matrix attributes stored as arrays. There's limited amount of compiled code (cython?), mainly for matrix multiplication (a key sparse calculation), and conversion between formats.

You have to iteration levels, the different A_i sequences, and kron chaining within each.

I don't know enough of how kron works to offhand suggest improvements on that. I suppose you could precompute some chained sequences, say for 2 and 4 I_d in a row. So in effect the above example could be written as

sparse.kron(Id_3, sparse.kron(A, Id_2))

In fact let's try that

Chain 2 and 3 Id matrices:

In [68]: Id2 = sparse.kron(Id,Id); Id3 = sparse.kron(Id,Id2)

In [69]: Id3
Out[69]: 
<8x8 sparse matrix of type '<class 'numpy.float64'>'
    with 32 stored elements (blocksize = 4x4) in Block Sparse Row format>

In [70]: Id3.A
Out[70]: 
array([[1., 0., 0., 0., 0., 0., 0., 0.],
       [0., 1., 0., 0., 0., 0., 0., 0.],
       [0., 0., 1., 0., 0., 0., 0., 0.],
       [0., 0., 0., 1., 0., 0., 0., 0.],
       [0., 0., 0., 0., 1., 0., 0., 0.],
       [0., 0., 0., 0., 0., 1., 0., 0.],
       [0., 0., 0., 0., 0., 0., 1., 0.],
       [0., 0., 0., 0., 0., 0., 0., 1.]])

Actually we don't need to use kron to get that. It's just a eye matrix for a desired size.

Now kron 3 matrices:

In [71]: M1 = sparse.kron(Id3,sparse.kron(A,Id2))

In [72]: M1
Out[72]: 
<64x64 sparse matrix of type '<class 'numpy.float64'>'
    with 512 stored elements in COOrdinate format>

and compare that the result of the full blown iteration:

In [73]: np.allclose(M1.A, foo(alist).A)
Out[73]: True

That's what I had in mind when first suggesting examining the kron sequence for patterns.

In fact using eye directly:

In [92]: M2 = sparse.kron(sparse.eye(8),sparse.kron(A,sparse.eye(4)))

In [93]: M2
Out[93]: 
<64x64 sparse matrix of type '<class 'numpy.float64'>'
    with 32 stored elements in COOrdinate format>

In [94]: np.allclose(M2.A, foo(alist).A)
Out[94]: True