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

问题描述

我正在尝试将一些代码从 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,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

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

谢谢

解决方法

在 numpy 中，乘法符号始终是按元素计算的（如 Matlab 中的 .*）。对于矩阵乘法，使用 @。

此外，您可以在 Python 中使用 ** 进行取幂。

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

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