Search code examples
pythonarraysmatlabbroadcasting

Broadcasting (N,2,2) array and (N,2) array through a for loop

I am trying to migrate some code from MATLAB to python and I am having trouble understanding why the following code does not work.

import numpy as np

ngrid = 56
A = np.random.randint(10, size =(ngrid*ngrid,2))
A_tmp = A
B = np.random.randint(10,size =(ngrid*ngrid,2,2) )


for jj in range(ngrid*ngrid):
    A[jj,:] = A_tmp[jj,:]*B[jj,:,:].conj()

When I execute this code I get the error.

ValueError: could not broadcast input array from shape (2,2) into shape (2,)

I am not understanding why this is giving me a one dimensional array vs a (ngrid*ngrid,2) array.

The MATLAB code I am trying to recreate for is

for jj = 1:ngrid^2
    Psi0(jj, :) = Psi0_tmp(jj, :)*dia2adi(:,:, jj)';
end

Any guidance on theory and how to correct my code would be very helpful.

Thanks

Solution

  • In numpy, the multiplication symbol is always element-wise (like .* in Matlab). For matrix multiplication, use @.

    Also, you can do exponentiation in Python with **.

    Therefore, if you change the loop to the code below, it works.

    for jj in range(ngrid**2):
    A[jj,:] = A_tmp[jj,:] @ B[jj,:,:].conj()