如何删除嵌套的for循环？

Question

我有以下嵌套循环：

sum_tot = 0.0

for i in range(len(N)-1):
    for j in range(len(N)-1):
        sum_tot = sum_tot + N[i]**2*N[j]**2*W[i]*W[j]*x_i[j][-1] / (N[j]**2 - x0**2) *(z_i[i][j] - z_j[i][j])*x_j[i][-1] / (N[i]**2 - x0**2)

它基本上是一个具有双重求和的数学函数。每个总和都达到 N 的长度。我一直在试图弄清楚是否有一种方法可以在不使用嵌套 for 循环的情况下编写它以减少计算时间。我尝试使用列表理解，但计算时间即使不一样也很相似。有没有办法将这个表达式写成矩阵来避免循环？

Answer 1

请注意，鉴于您当前的循环，范围将停止在 N-2：range 上升到但不包括其参数。你可能想写 for i in range(len(N)).

减少求和也很困难：它实际花费的时间取决于计算的项数，因此如果您以不同的方式编写它仍然涉及相同数量的项，则需要同样长的时间。然而，O(n^2) 并不完全糟糕：它看起来是您在这种情况下可以做的最好的事情，除非您找到问题的数学简化。

您可以考虑检查以收集以更简洁的方式写出总和的方法。

Answer 2

@Kraigolas 提出了有效的观点。但是，无论哪种方式，让我们在一个虚拟的、双重嵌套的操作上尝试一些基准测试。（提示：Numba 可能会帮助您加快速度）

请注意，我会特别避免使用 numpy 数组，因为范围之间的所有叉积都会立即存储在内存中。如果这是一个很大的范围，您可能运行内存不足。

嵌套 for 循环

n = 5000
s1 = 0

for i in range(n):
    for j in range(n):
        s1 += (i/2) + (j/3)
        
print(s1)

#2.26 s ± 101 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

列表理解

n = 5000
s2 = 0

s2 = sum([i/2+j/3 for i in range(n) for j in range(n)])
print(s2)

#3.2 s ± 307 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Itertools 产品

from itertools import product

n = 5000
s3 = 0

for i,j in product(range(n),repeat=2):
    s3 += (i/2) + (j/3)
print(s3)

#2.35 s ± 186 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Note: When using Numba, you would want to run the code at least once before, because the first time it compiles the code and therefore the speed is slow. The real speedup comes second run onwards.

Numba njit (SIMD)

from numba import njit

n=5000

@njit
def f(n):
    s = 0
    for i in range(n):
        for j in range(n):
            s += (i/2) + (j/3)
    return s

s4 = f(n)

#29.4 ms ± 1.85 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Numba njit 与 prange 平行

@Tim 的出色建议，已添加到基准测试中

@njit(parallel=True)
def f(n):
    s = 0
    for i in prange(n):
        for j in prange(n):
            s += (i/2) + (j/3)
    return s

s5 = f(n)

#21.8 ms ± 4.81 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Numba 正如预期的那样显着提升。也许试试看？

Answer 3

要将其转换为矩阵计算，我建议先合并一些项。

如果这些对象不是 numpy 数组，最好将它们转换为 numpy 数组，因为它们支持逐元素操作。

要转换，只需执行

import numpy
N = numpy.array(N)
w = numpy.array(w)
x_i = numpy.array(x_i)
x_j = numpy.array(x_j)
z_i = numpy.array(z_i)
z_j = numpy.array(z_j)

然后，


common_terms = N**2*w/(N**2-x0**2)
i_terms = common_terms*x_j[:,-1]
j_terms = common_terms*x_i[:,-1]
i_j_matrix = z_i - z_j
sum_output = (i_terms.reshape((1,-1)) @ i_j_matrix @ j_terms.reshape((-1,1)))[0,0]

如何删除嵌套的for循环？

How to remove nested for loop?

python

nested-loops

python-2.7

python-3.x

嵌套 for 循环

列表理解

Itertools 产品

Numba njit (SIMD)

Numba njit 与 prange 平行