如何加快第二个循环依赖于第一个循环的嵌套循环？

Question

考虑以下代码：

n = 20000


def f(i, j):

    return (i+1j*j)/(i-1j*j+1)  # a sample function, not necessary this form


lst = []
for i in range(n):
    for j in range(i, n):
        lst.append((i, j, f(i, j)))

由于循环非常大，我想对其进行矢量化或加速。阅读其他帖子，似乎itertools.product 可以加快循环，但在我的情况下，第二个循环取决于第一个，看来我不能简单地使用它。那么如何加速呢？

例如，我可以使用 4 个处理器。

Answer 1

一般情况下，你可以用np.triu_indices：

i, j = np.triu_indices(n)
np.stack([i, j, f(i, j)])

这可能会阻塞您的系统（因为 i 和 j 对于 n = 20000 将有 200M 个元素）在这种情况下您将需要 itertools

Answer 2

建议：

1. 不要动态增加列表。请参阅：Is there a way to circumvent Python list.append() becoming progressively slower in a loop as the list grows? 和 Create an empty list in python with certain size。
2. 分配一个所需大小的空白 numpy 数组，而不是列表，然后使用您的函数填充它。
3. 将输入范围划分为子集，并创建子进程来处理每个子集。例如，进程 1 处理 0-10000，进程 2 处理 10001-20000。主进程等待子进程退出，然后合并结果。

Answer 3

你可以使用：

lst = list( map ( 
  lambda item: (item[0],item[1], f(item[0],item[1] )), 
  filter( lambda item: item[0] >= item[1] ,  itertools.product(range(n),range(n))) 
))

但是，由于itertools.product 使用嵌套循环，我看不出它如何加速循环。

Answer 4

我将展示一个较小的 n 值的示例：

n = 4

首先创建一个包含函数结果的完整数组：

arr = np.fromfunction(f, (n,n), dtype='complex')

到目前为止，结果（精度降低）是：

array([[ 0.     +0.j, -0.5    +0.5j    , -0.8    +0.4j    , -0.9    +0.3j    ],
       [ 0.5    +0.j,  0.2    +0.6j    , -0.25   +0.75j   , -0.53846+0.69231j],
       [ 0.66667+0.j,  0.5    +0.5j    ,  0.15385+0.76923j, -0.16667+0.83333j],
       [ 0.75   +0.j,  0.64706+0.41176j,  0.4    +0.7j    ,  0.12   +0.84j   ]])

然后，要生成您的预期结果，运行：

result = [ [i,j, arr[i,j]] for i, j in zip(*np.triu_indices(n)) ]

结果是一个列表列表，其中包含：

[[0, 0, 0j]],
 [0, 1, (-0.5    +0.5j)]],
 [0, 2, (-0.8    +0.4j)],
 [0, 3, (-0.89999+0.3j)],
 [1, 1, ( 0.2    +0.6j)],
 [1, 2, (-0.25   +0.75j)],
 [1, 3, (-0.53846+0.69231j)],
 [2, 2, ( 0.15385+0.76923j)],
 [2, 3, (-0.16667+0.83333j)],
 [3, 3, ( 0.12   +0.84j)]]

（我也降低了精度）。

如果您遇到有关可用内存的问题，请不要创建 任何临时数组，但改为运行：

result = [ [i, j, f(i,j)] for i, j in zip(*np.triu_indices(n)) ]

但是这个变种会比使用 np.fromfunction 慢很多。

Answer 5

您应该能够利用numba 的功能来完成此类工作：

import numpy as np
from numba import jit, prange

n = 10000
@jit(nopython=True)
def f(i, j):
    return (i+1j*j)/(i-1j*j+1)  # a sample function, not necessary this form

def loops_lst(n):
    lst = []
    for i in range(n):
        for j in range(i, n):
            lst.append((i, j, f(i, j)))
    return np.asarray(lst)
res1 = loops_lst(n)

def loops_np(n): # solution by @Daniel F
    i, j = np.triu_indices(n)
    out = np.stack([i, j, f(i, j)])
    return out.T
res2 = loops_np(n)

@jit(nopython=True,parallel=True)
def loops_nb(n):
    dim = np.sum(np.arange(1, n+1))
    out_idx = np.empty((dim,2),dtype=np.int64)
    out_f = np.empty((dim),dtype=np.complex64)
    for i in prange(n):
        for j in range(i, n):
            idx = int((n*(n-1))/2 - ((n-i)*(n-i-1))/2+j)
            out_idx[idx,0] = i
            out_idx[idx,1] = j
            out_f[idx] = f(i, j)
    return out_idx, out_f

res3_ = loops_nb(n)
res3 = np.hstack((res3_[0], res3_[1][:,None]))

这里是完整性检查：

np.allclose(res1,res2)
>>> True
np.allclose(res1,res3)
>>> True

下面是n=10000的时间安排：

# numpy triu_indices
%timeit res2 = loops_np(n)
2.32 s ± 58.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# numba
# parallel = False (and simple range instead of numba.prange)
%timeit res3_ = loops_nb(n);res3 = np.hstack((res3_[0], res3_[1][:,None]))
2.13 s ± 47.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# parallel = True
%timeit res3_ = loops_nb(n);res3 = np.hstack((res3_[0], res3_[1][:,None]))
1.46 s ± 71.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

所以如果有可能njit 你的真正功能f，你一定要看看numba 解决方案。