python - 加速嵌套 python 数组的操作

Question

我正在使用嵌套数组。我需要对该数组的每个元素应用相当简单但成本高昂的算术运算。

下面是MWE，其中“第二个 block ”是占用大部分时间的 block (它将运行数千次)。

考虑到获取a、b、c的真实方式，第一个和第二个 block 需要分开，因为第一个 block 仅处理一次 时间成本非常高。

我不确定如何提高应用于嵌套数组的每个元素的操作的性能。当然，numpy 会更快地完成此操作，但我不太熟悉数组上的广播操作。

import numpy as np
import time


# Generate some random data.
N = 100
x, y, z = [np.random.uniform(0., 10., N) for _ in range(3)]

# Grid of values in 2 dimensions.
M = 200
p_lst, q_lst = np.linspace(0., 50., M), np.linspace(0., 25., M)

# Define empty nested list to be filled below.
# The shape is given by the length of the lists defined above.
abc_lst = [[[] for _ in p_lst] for _ in q_lst]

# First block. This needs to be separated from the block below.
# Fill nested list with values.
for i, p in enumerate(p_lst):
    for j, q in enumerate(q_lst):

        # a,b,c are obtained via some complicated function of p,q.
        # This is just for the purpose of this example.
        a, b, c = 1.*p, 1.*q, p+q

        # Store in nested list.
        abc_lst[i][j] = [a, b, c]

# Second block <-- THIS IS THE BOTTLENECK
tik = time.time()
# Apply operation on nested list.
lst = []
for i in range(len(p_lst)):
    for j in range(len(q_lst)):

        # Extract a,b,c values from nested list.
        a, b, c = abc_lst[i][j]

        # Apply operation. This is the *actual* operation
        # I need to apply.
        d = sum(abs(a*x + y*b + c*z))

        # Store value.
        lst.append(d)

print time.time() - tik

Answer 1

我在 this question 中找到了答案，使用np.outer()功能。

只需要对第一个 block 进行一些重新排列，第二个 block 的运行速度就会快很多倍。

# First block. Store a,b,c separately.
a_lst, b_lst, c_lst = [], [], []
for i, p in enumerate(p_lst):
    for j, q in enumerate(q_lst):

        # a,b,c are obtained via some complicated function of p,q.
        # This is just for the purpose of this example.
        a_lst.append(1.*p)
        b_lst.append(1.*q)
        c_lst.append(p+q)

# As arrays.
a_lst, b_lst, c_lst = np.asarray(a_lst), np.asarray(b_lst), np.asarray(c_lst)

# Second block.
# Apply operation on nested list using np.outer.
lst = np.sum(abs(np.outer(a_lst, x) + np.outer(b_lst, y) + np.outer(c_lst, z)), axis=1)