使用 NumPy 替换 Python 中的嵌套循环答案

【问题标题】：Replacing nested loops in Python using NumPy使用 NumPy 替换 Python 中的嵌套循环
【发布时间】：2021-04-07 17:39:27
【问题描述】：

所以我有这 6 个嵌套循环，它们的目的只是将数组 X 和 Y 在不同的索引上相乘和相加，得到数组 Z。

import numpy as np

dim_a = 5
dim_b = 9
Z = np.zeros((dim_a,dim_b,dim_b,dim_a))
X = np.arange(2025).reshape(dim_a,dim_b,dim_b,dim_a)
Y = np.arange(2025).reshape(dim_a,dim_b,dim_b,dim_a)

for i in range(0, dim_a):
    for j in range(0,dim_a):
        for a in range(0, dim_b):
            for b in range(0, dim_b):
                for m in range(0,dim_a):
                    for e in range(0,dim_b):
                        Z[i,a,b,j] += X[m,e,b,j] * Y[m,e,a,i] * 2

有没有办法使用 NumPy 只用几行代码来编写它？这些嵌套循环的计算量是巨大的。我感觉 NumPy 可以显着优化它。

【问题讨论】：

标签： python arrays python-3.x numpy

【解决方案1】：

你当然可以用np.einsum：

Z[i,a,b,j] += X[m,e,b,j] * Y[m,e,a,i] * 2

翻译成

Z = np.einsum('mebj,meai->iabj', X,Y) * 2

【讨论】：

这太棒了！我检查了计算时间：6 nested loops: 0.14539456367492676np.einsum: 0.000209808349609375
也和np.einsum('mebj,meai,...->iabj', X, Y, 2)一样。虽然，不确定这是否更快。