使用元素求幂加速嵌套 for 循环

Question

我正在编写一个大型代码，我发现自己需要加快其中的特定部分。我创建了一个MWE，如下所示：

import numpy as np
import time

def random_data(N):
    # Generate some random data.
    return np.random.uniform(0., 10., N).tolist()

# Lists that contain all the data.
list1 = [random_data(10) for _ in range(1000)]
list2 = [random_data(1000), random_data(1000)]

# Start taking the time.
tik = time.time()

list4 = []
# Loop through all elements in list1.
for elem in list1:

    list3 = []
    # Loop through elements in list2.
    for elem2 in zip(*list2):

        A = np.exp(-0.5*((elem[0]-elem2[0])/elem[3])**2)
        B = np.exp(-0.5*((elem[1]-elem2[1])/elem[3])**2)
        list3.append(A*B)

    # Sum elements in list3 and append result to list4.
    sum_list3 = sum(list3) if sum(list3)>0. else 1e-06
    list4.append(sum_list3)

# Print the elapsed time.
print time.time()-tik

list1 和 list2 的奇怪格式是因为这就是这段代码接收它们的方式。

花费大部分时间的明显部分是递归计算 A 和 B 术语。

有没有什么方法可以加快这段代码的速度而不必并行化它（我之前尝试过它，它给了我a lot of troubles）？我愿意使用任何软件包，numpy、scipy 等。

---

添加

这是应用 abarnert 优化的结果，也是 Jaime 建议只进行一次幂运算的结果。优化后的功能在我的系统上平均快了约 60 倍。

import numpy as np
import timeit

def random_data(N):
    return np.random.uniform(0., 10., N).tolist()

# Lists that contain all the data.
list1 = [random_data(10) for _ in range(1000)]
list2 = [random_data(1000), random_data(1000)]

array1 = np.array(list1)
array2 = np.array(zip(*list2))


# Old non-optimezed function.
def func1():
    list4 = []
    # Process all elements in list1.
    for elem in list1:
        # Process all elements in list2.
        list3 = []
        for elem2 in zip(*list2):
            A = np.exp(-0.5*((elem[0]-elem2[0])/elem[3])**2)
            B = np.exp(-0.5*((elem[1]-elem2[1])/elem[3])**2)
            list3.append(A*B)
        # Sum elements in list3 and append result to list4.
        sum_list3 = sum(list3) if sum(list3)>0. else 1e-06
        list4.append(sum_list3)

# New optimized function.
def func2():
    list4 = []
    # Process all elements in list1.
    for elem in array1:

        # Broadcast over elements in array2.
        A = -0.5*((elem[0]-array2[:,0])/elem[3])**2
        B = -0.5*((elem[1]-array2[:,1])/elem[3])**2
        array3 = np.exp(A+B)

        # Sum elements in array3 and append result to list4.
        sum_list3 = max(array3.sum(), 1e-10)
        list4.append(sum_list3)


# Get time for both functions.
func1_time = timeit.timeit(func1, number=10)
func2_time = timeit.timeit(func2, number=10)

# Print hom many times faster func2 is versus func1.
print func1_time/func2_time

Answer 1

您希望逐渐将其从使用列表和循环转换为使用数组和广播，首先抓住最简单和/或最关键的部分，直到速度足够快。

第一步是不要一遍又一遍地这样做zip(*list2)（特别是如果这是Python 2.x）。当我们这样做时，我们不妨将它存储在一个数组中，并对list1 做同样的事情——你现在仍然可以遍历它们。所以：

array1 = np.array(list1)
array2 = np.array(zip(*list2))
# …
for elem in array1:
    # …
    for elem2 in array2:

这不会加快速度——在我的机器上，它需要我们从 14.1 秒到 12.9 秒——但它给了我们一个开始工作的地方。

你还应该去掉sum(list3)的双重计算：

sum_list3 = sum(list3)
sum_list3 = sum_list3 if sum_list3>0. else 1e-06

同时，您希望value <= 0 转到1e-6，但不理会0 < value < 1e-6，这有点奇怪。这真的是故意的吗？如果没有，你可以解决这个问题，同时简化代码，只需这样做：

sum_list3 = max(array3.sum(), 1e-06)

现在，让我们广播A 和B 的计算：

# Broadcast over elements in list2.
A = np.exp(-0.5*((elem[0]-array2[:,0])/elem[3])**2)
B = np.exp(-0.5*((elem[1]-array2[:, 1])/elem[3])**2)
array3 = A*B

# Sum elements in list3 and append result to list4.
sum_list3 = max(array3.sum(), 1e-06)

list4.append(sum_list3)

这让我们从 12.9 秒缩短到 0.12 秒。您还可以通过array1 进行广播，并用预先分配的数组替换list4 等等，更进一步，但这可能已经足够快了。