4

I have two matrix arrays A and B such with identical shape: A.shape = B.shape = (M,N,P)

I would like to compute the Kronecker product along the axis 0, so that:

KP[ii,:,:] = A[ii,:,:]⊗B[ii,:,:]

Is there a way of doing this in numpy without using for loops?

Thanks!

Example:

A = np.array([ [[1,0],
                [0,1]],
               [[1,0],
                [0,1]]
            ])

B = np.array([ [[1,0],
                [0,-1]],
               [[0,1],
                [1,0]]
            ])

KP = np.array( [
                [[1,0,0,0],
                 [0,-1,0,0],
                 [0,0,1,0],
                 [0,0,0,-1]],
                [[0,1,0,0],
                 [1,0,0,0],
                 [0,0,0,1],
                 [0,0,1,0]]
               ] )

which would be equivalent to:

KP= np.zeros( (A.shape[0],
               A.shape[1]**2,
               A.shape[2]**2) )

for ii in range(A.shape[0]):

    KP[ii,:,:] = np.kron(A[ii,:,:],B[ii,:,:])
blackbody
  • 55
  • 4

2 Answers2

5

you can use einsum which with a bit of practice is quite intuitive or go the classic reshape-and-broadcast route

A = np.array([ [[1,0],
                [0,1]],
               [[1,0],
                [0,1]]
            ])

B = np.array([ [[1,0],
                [0,-1]],
               [[0,1],
                [1,0]]
            ])

i,j,k = A.shape
i,l,m = B.shape
np.einsum("ijk,ilm->ijlkm",A,B).reshape(i,j*l,k*m)

# array([[[ 1,  0,  0,  0],
#         [ 0, -1,  0,  0],
#         [ 0,  0,  1,  0],
#         [ 0,  0,  0, -1]],
# 
#        [[ 0,  1,  0,  0],
#         [ 1,  0,  0,  0],
#         [ 0,  0,  0,  1],
#         [ 0,  0,  1,  0]]])

equivalent non-einsum expression:

(A[:,:,None,:,None]*B[:,None,:,None,:]).reshape(i,j*l,k*m)
Paul Panzer
  • 51,835
  • 3
  • 54
  • 99
  • Thanks, this definitely works! I'm a bit confused about how the non-einsum expression works in this case though. – blackbody Jul 29 '19 at 19:56
  • @blackbody I've added trailing colon to the B index. It is redundant but perhaps makes it more readable? – Paul Panzer Jul 29 '19 at 20:15
0

in case someone finds this useful as well (maybe it's not that efficient though).

def give_kr_prod(matrices):
    #matrices list of 2 (or more in principle) matrices
    while len(matrices) != 1:
        sm, smf=[],[]
        for ind in range(len(matrices)):
            sm.append(matrices[ind])
            if ind%2==1 and ind>0:
                smf.append(np.kron(*sm))
                sm=[]
        matrices = smf
    return matrices[0]

matrices = np.random.randn(8,2,2)