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

【问题标题】：Broadcasting (N,2,2) array and (N,2) array through a for loop通过 for 循环广播 (N,2,2) 数组和 (N,2) 数组
【发布时间】：2021-02-19 00:00:42
【问题描述】：

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

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

谢谢

【问题讨论】：

它给你一个一维数组，因为切片一个 numpy 数组的一维切片是一个一维数组。将此与不存在 1d 数组的 MATLAB 进行对比：一切都至少是 2d。
现在我来尝试提出一个优雅的（阅读：矢量化）解决方案为时已晚，但最坏的情况是您可以使用A = np.einsum('ja,jab -> jb', A_tmp, B.conj()) 之类的东西而不是整个循环。给出或采用 MATLAB 的主要运算符所隐含的转置。
我听说过这个“神奇的”einsum。感谢您的帮助，问题已解决。
我认为你不能用 matmul 解决这个问题，所以唯一的其他选择是A = (A_tmp[..., None] * B.conj()).sum(1)。它可能会比 einsum 更快，但会消耗更多内存。
谢谢，我在优化代码时会考虑到这一点

标签： python arrays matlab broadcasting

【解决方案1】：

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

另外，您可以在 Python 中使用 ** 进行幂运算。

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

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

【讨论】：