Skip to main content
 首页 » 编程设计

python之MATLAB 的 bsxfun 是最好的吗 Python 的 numpy.einsum

2025年05月04日10zhengyun_ustc

我有一个非常大的乘法和求和运算,我需要尽可能高效地实现它。到目前为止,我发现的最佳方法是 MATLAB 中的 bsxfun,我将问题表述为:

L = 10000; 
x = rand(4,1,L+1); 
A_k = rand(4,4,L); 
tic 
for k = 2:L 
    i = 2:k; 
    x(:,1,k+1) = x(:,1,k+1)+sum(sum(bsxfun(@times,A_k(:,:,2:k),x(:,1,k+1-i)),2),3); 
end 
toc 

请注意,L 实际上会更大。有没有更快的方法?奇怪的是,我需要先将单例维度添加到 x,然后再对它进行 sum,但我无法让它正常工作。

它仍然比我尝试过的任何其他方法快得多,但对于我们的应用程序来说还不够。我听说 Python 函数 numpy.einsum 可能更高效,但我想在考虑移植我的代码之前先在这里问一下。

我正在使用 MATLAB R2017b。

请您参考如下方法:

我相信你的两个总结都可以去掉,但我暂时只去掉了比较容易的那个。第二个维度的求和是微不足道的,因为它只影响 A_k 数组:

B_k = sum(A_k,2); 
for k = 2:L 
    i = 2:k; 
    x(:,1,k+1) = x(:,1,k+1) + sum(bsxfun(@times,B_k(:,1,2:k),x(:,1,k+1-i)),3); 
end 

通过这一单一更改,我的笔记本电脑上的运行时间从约 8 秒减少到约 2.5 秒。

通过将 times+sum 转换为矩阵向量乘积,也可以删除第二个求和。它需要一些单例摆弄来获得正确的维度,但是如果您定义一个辅助数组 B_k 并且第二个维度被反转,您可以生成剩余的总和作为 ~x*C_k 和这个辅助数组 C_k,发出或接收一些 reshape 的调用。


所以在仔细观察之后,我意识到我最初的评估过于乐观:你在剩余的任期中在两个维度上都有乘法,所以它不是一个简单的矩阵乘积。无论如何,我们可以将该项重写为矩阵乘积的对角线。这意味着我们正在计算一堆不必要的矩阵元素,但这似乎仍然比 bsxfun 方法稍微快一点,我们也可以摆脱你讨厌的单例维度:

L = 10000; 
x = rand(4,L+1); 
A_k = rand(4,4,L); 
B_k = squeeze(sum(A_k,2)).'; 
 
tic 
for k = 2:L 
    ii = 1:k-1; 
    x(:,k+1) = x(:,k+1) + diag(x(:,ii)*B_k(k+1-ii,:)); 
end 
toc 

这在我的笔记本电脑上运行大约 2.2 秒,比之前获得的大约 2.5 秒快一些。