Element-wise tensor product for block matrices whos elements are 2x2 matrices

Question

I have a block matrix with its elements as 2x2 matrices numpy array, for example

X = np.array([[0, 1], [1, 0]], dtype=complex)
Y = np.array([[0, -1j], [1j, 0]], dtype=complex)
Z = np.array([[1, 0], [0, -1]], dtype=complex)

block_matrix = np.array([X,Y,Z])

I'm looking for a vectorised way(if it exists) where I could compute np.kron() without having to loop through each of the block matrix's elements(who themselves are 2x2 matrices again). Right now I have something like

def pl_rep_operation(matrix):

    op_sum = np.zeros((4,4), dtype=complex)
    tensored_seq = []
    for i in range(len(matrix)):
        tensored = np.kron(matrix[i], matrix[i].conj()) 
        op_sum += tensored
        tensored_seq.append(tensored)
    return op_sum, tensored_seq

where tensored_seq returns the original sequence with its block matrix element tensored and op_sum returns the element-wise sum of all tensored matrix elements. Output for example could be

op_sum, tensored_seq = pl_rep_operation(np.array([X,Y,Z]))
In[47]: op_sum
Out[47]: 
array([[ 1.+0.j,  0.+0.j,  0.+0.j,  2.+0.j],
       [ 0.+0.j, -1.+0.j,  0.+0.j,  0.+0.j],
       [ 0.+0.j,  0.+0.j, -1.+0.j,  0.+0.j],
       [ 2.+0.j,  0.+0.j,  0.+0.j,  1.+0.j]])

In[48]: tensored_seq
Out[48]: 
[array([[0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j],
        [0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j],
        [0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j],
        [1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j]]),
 array([[ 0.+0.j, -0.+0.j, -0.+0.j,  1.+0.j],
        [ 0.+0.j,  0.+0.j, -1.+0.j, -0.+0.j],
        [ 0.+0.j, -1.+0.j,  0.+0.j, -0.+0.j],
        [ 1.+0.j,  0.+0.j,  0.+0.j,  0.+0.j]]),
 array([[ 1.+0.j,  0.+0.j,  0.+0.j,  0.+0.j],
        [ 0.+0.j, -1.-0.j,  0.+0.j,  0.-0.j],
        [ 0.+0.j,  0.+0.j, -1.+0.j,  0.+0.j],
        [ 0.+0.j,  0.-0.j,  0.+0.j,  1.+0.j]])]

The elements of tensored_seq should be something like np.array([np.kron(X,X), np.kron(Y,Y), np.kron(Z,Z)]). I'm looking for some function np.func() or some way to vectorise this such that np.func(block_matrix, block_matrix) would return np.array([np.kron(X,X), np.kron(Y,Y), np.kron(Z,Z)]). Ideally, I want a vectorised way which also does

block_mat = np.array([[X, Y, Z], [X, Z, Y], [Z, Y, X]])
np.func(block_mat)

should return

np.array([[np.kron(X,X), np.kron(Y,Y), np.kron(Z,Z)],
          [np.kron(X,X), np.kron(Z,Z), np.kron(Y,Y)],
          [np.kron(Z,Z), np.kron(Y,Y), np.kron(X,X)]])

for example.

`np.kron` is not exceptional in its "vectoriztion" I just demonstrated in a recent answer that it does an `outer` followed by element rearrangement (reshape and transpose) — hpaulj, Mar 01 '20 at 16:14
If `X` is (2,2), then `block_matrix` is (3,2,2). The second version would be (3,3,2,2), with your psuedo-kron expanding to (3,3,4,4). Since you aren't doing things like `kron(X,Z)`, I think the most straight forward route is to just calculate the 3 krons, and assemble the target from them. This is too specialized of an operation for the base `numpy` code. — hpaulj, Mar 01 '20 at 16:45

score 0 · Answer 1 · answered Mar 01 '20 at 17:27

Working from my recent answer at

Why is numpy's kron so fast?

In [472]: X = np.array([[0, 1], [1, 0]], dtype=complex) 
     ...: Y = np.array([[0, -1j], [1j, 0]], dtype=complex) 
     ...: Z = np.array([[1, 0], [0, -1]], dtype=complex) 
     ...:  
     ...: block_matrix = np.array([X,Y,Z])                                                     
In [473]: block_matrix.shape                                                                   
Out[473]: (3, 2, 2)
In [474]: temp=block_matrix[:,:,:,None]*block_matrix.conj()[:,:,None,:]                        
In [475]: temp.shape                                                                           
Out[475]: (3, 2, 2, 2)                                                                         
In [477]: temp = block_matrix.ravel()                                                          
In [478]: temp = block_matrix.reshape(3,4)                                                     
In [479]: temp = temp[:,:,None]*temp.conj()[:,None,:]                                          
In [480]: temp.shape                                                                           
Out[480]: (3, 4, 4)
In [481]: nz = temp.shape                                                                      
In [482]: temp = temp.reshape(3,2,2,2,2)                                                       
In [483]: temp = temp.transpose(0,1,3,2,4).reshape(nz)                                         
In [484]: temp.shape                                                                           
Out[484]: (3, 4, 4)
In [485]: temp                                                                                 
Out[485]: 
array([[[ 0.+0.j,  0.+0.j,  0.+0.j,  1.+0.j],
        [ 0.+0.j,  0.+0.j,  1.+0.j,  0.+0.j],
        [ 0.+0.j,  1.+0.j,  0.+0.j,  0.+0.j],
        [ 1.+0.j,  0.+0.j,  0.+0.j,  0.+0.j]],

       [[ 0.+0.j, -0.+0.j, -0.+0.j,  1.+0.j],
        [ 0.+0.j,  0.+0.j, -1.+0.j, -0.+0.j],
        [ 0.+0.j, -1.+0.j,  0.+0.j, -0.+0.j],
        [ 1.+0.j,  0.+0.j,  0.+0.j,  0.+0.j]],

       [[ 1.+0.j,  0.+0.j,  0.+0.j,  0.+0.j],
        [ 0.+0.j, -1.-0.j,  0.+0.j,  0.-0.j],
        [ 0.+0.j,  0.+0.j, -1.+0.j,  0.+0.j],
        [ 0.+0.j,  0.-0.j,  0.+0.j,  1.+0.j]]])

Which matches your tensored_seq. If put together into a function, it should be faster than your 3 kron. But I don't know if it's enough to extend. The time penalty for 3 iterations on a relatively complex task is not great.

I would try to construct the last matrix with:

In [486]: res = np.array([[temp[0], temp[1], temp[2]], 
     ...:                 [temp[0], temp[2],...

assuming you want a (3,3,4,4) result. That's no more work than constructing

np.array([[X,Y,Z],[X,Z...])

Element-wise tensor product for block matrices whos elements are 2x2 matrices

1 Answers1