通过 for 循环广播 (N,2,2) 数组和 (N,2) 数组

Question

我正在尝试将一些代码从 MATLAB 迁移到 python，但我无法理解为什么以下代码不起作用。

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()

当我执行此代码时出现错误。

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

我不明白为什么这会给我一个一维数组而不是 (ngrid*ngrid,2) 数组。

我尝试重新创建的 MATLAB 代码是

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

任何关于理论和如何更正我的代码的指导都会很有帮助。

谢谢

Answer 1

在 numpy 中，乘法符号总是逐元素的（就像 Matlab 中的 .*）。对于矩阵乘法，使用 @.

此外，您可以使用 ** 在 Python 中求幂。

因此，如果将循环更改为下面的代码，它就可以工作。

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

通过 for 循环广播 (N,2,2) 数组和 (N,2) 数组

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

python

arrays

matlab

broadcasting