如何删除嵌套的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)

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

Answer 1

请注意，给定当前循环，范围将在 N-2 处停止： range上升但不包括其参数。 您可能的意思是写for i in range(len(N)) 。

减少求和也很困难：实际花费的时间是基于计算的项数，所以如果你用不同的方式写它仍然涉及相同数量的项，它会花费同样长的时间。 但是， O(n^2)并不是很糟糕：除非您找到问题的数学简化，否则在这种情况下您似乎可以做到最好。

您可以考虑查看这篇文章以收集以更简洁的方式写出总和的方法。

Answer 2

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

请注意，我会避免 numpy arrays 特别是因为范围之间的所有叉积都将同时出现在 memory 中。 如果这是一个很大的范围，您可能会用完 memory。

嵌套 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)

注意：使用 Numba 时，您可能希望至少运行一次代码，因为第一次编译代码时速度很慢。 真正的加速来自第二次运行。

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 arrays，最好将它们转换为 numpy ZA3CBC3F9D06CE2F2C19CZE1，因为它们支持元素操作。

要转换，只需执行

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]