使用 List Comprehension 和/或 map 避免嵌套 for 循环

Question

几天来，我一直在苦苦思考如何优化（不仅让它看起来更漂亮）3 个 嵌套循环，其中包含一个 条件 和一个函数调用。我现在拥有的是以下内容：

def build_prolongation_operator(p,qs):
    '''
    p: dimension of the coarse basis
    q: dimension of the fine basis

    The prolongation operator describes the relationship between
    the coarse and fine bases:    
    V_coarse = np.dot(V_fine, I)
    '''

    q = sum(qs)

    I = np.zeros([q, p])

    for i in range(0, q):
        for j in range(0, p):
            for k in range(0, qs[j]):
                # if BV i is a child of j, we set I[i, j] = 1
                if i == f_map(j, k, qs):
                    I[i, j] = 1
                    break

    return I

f_map 在哪里：

def f_map(i, j, q):
    '''
    Mapping which returns the index k of the fine basis vector which
    corresponds to the jth child of the ith coarse basis vector.    
    '''

    if j < 0 or j > q[i]:
        print('ERROR in f_map')
        return None

    result = j

    for k in range(0, i):
        result += q[k]

    return result

在分析我的整个代码时，我发现 build_prolongation_operator 被调用了 45 次，f_map 大约 850 万次！！

图片如下：

我曾尝试对列表理解和地图做同样的事情，但没有任何运气。

这是build_prolongation_operator 期望的输入示例：

p = 10
qs = randint(3, size=p)

Answer 1

我不知道基数和延长运算符，但您应该关注算法本身。在优化方面，这几乎总是合理的建议。

这可能是症结所在——如果不是，它可以帮助您入门：f_map 计算不依赖于 i，但您正在为每个 i 值重新计算它。由于i 的范围从零到qs 中的值之和，因此您将通过缓存结果来节省大量的重新计算；谷歌“python memoize”，它实际上会自己写。修复此问题，您可能已经完成，无需任何微优化。

您将需要足够的空间来存储 max(p) * max(qs[j]) 值，但从您报告的调用次数来看，这应该不是太大的障碍。

Answer 2

试试看是否可行，

for j in range(0,p):
    for k in range(0, qs[j]):
        # if BV i is a child of j, we set I[i,j] = 1
        val = f_map(j,k,qs)
        if I[val, j] == 0:
            I[val, j] = 1

Answer 3

一方面，您确实不需要p 作为函数的参数：len(qs) 只需要调用一次，而且非常便宜。如果您的输入始终是一个 numpy 数组（并且在这种情况下没有理由不应该这样），qs.size 也可以。

让我们从重写f_map 开始。那里的循环只是qs 的累积总和（但从零开始），您可以预先计算一次（或每次调用外部函数至少一次）。

def f_map(i, j, cumsum_q):
    return j + cumsum_q[i]

cumsum_q 将在 build_prolongation_operator 中定义为

cumsum_q = np.roll(np.cumsum(qs), 1)
cumsum_q[0] = 0

我相信您会体会到在f_map 中拥有与在build_prolongation_operator 中相同的一组变量名的用处。为了更简单，我们可以完全删除 f_map 并在您的条件下使用它所代表的表达式：

if i == k + cumsum_q[j]:
    I[i, j] = 1

k 上的循环意味着“如果i 是k + cumsum[j] for any k”，则将元素设置为1。如果我们将条件重写为i - cumsum_q[j] == k，你可以看到我们根本不需要k 的循环。 i - cumsum_q[j] 将等于 some k 在[0, qs[j]) 范围内，如果它是非负数并且严格小于qs[j]。你可以检查一下

if i >= cumsum_q[j] and i - cumsum_q[j] < qs[j]:
    I[i, j] = 1

这将您的循环减少到矩阵的每个元素一次迭代。没有比这更好的了：

def build_prolongation_operator_optimized(qs):
    '''
    The prolongation operator describes the relationship between
    the coarse and fine bases:    
    V_coarse = np.dot(V_fine, I)
    '''
    qs = np.asanyarray(qs)
    p = qs.size
    cumsum_q = np.roll(np.cumsum(qs), 1)
    q = cumsum_q[0]
    cumsum_q[0] = 0

    I = np.zeros([q, p])

    for i in range(0, q):
        for j in range(0, p):
            # if BV i is a child of j, we set I[i, j] = 1
            if 0 <= i - cumsum_q[j] < qs[j]:
                I[i, j] = 1
    return I

现在您已经知道每个单元格的公式，您可以让 numpy 使用广播在基本上一行中为您计算整个矩阵：

def build_prolongation_operator_numpy(qs):
    qs = np.asanyarray(qs)
    cumsum_q = np.roll(np.cumsum(qs), 1)
    q = cumsum_q[0]
    cumsum_q[0] = 0
    i_ = np.arange(q).reshape(-1, 1)  # Make this a column vector
    return (i_ >= cumsum_q) & (i_ - cumsum_q < qs)

我运行了一个小型演示，以确保 (A) 建议的解决方案获得与原始解决方案相同的结果，并且 (B) 工作得更快：

In [1]: p = 10
In [2]: q = np.random.randint(3, size=p)

In [3]: ops = (
...     build_prolongation_operator(p, qs),
...     build_prolongation_operator_optimized(qs),
...     build_prolongation_operator_numpy(qs),
...     build_prolongation_operator_RaunaqJain(p, qs),
...     build_prolongation_operator_gboffi(p, qs),
... )

In [4]: np.array([[(op1 == op2).all() for op1 in ops] for op2 in ops])
Out[4]: 
array([[ True,  True,  True,  True,  True],
       [ True,  True,  True,  True,  True],
       [ True,  True,  True,  True,  True],
       [ True,  True,  True,  True,  True],
       [ True,  True,  True,  True,  True]])

In [5]: %timeit build_prolongation_operator(p, qs)
321 µs ± 890 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [6]: %timeit build_prolongation_operator_optimized(qs)
75.1 µs ± 1.85 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [7]: %timeit build_prolongation_operator_numpy(qs)
24.8 µs ± 77.7 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [8]: %timeit build_prolongation_operator_RaunaqJain(p, qs)
28.5 µs ± 1.55 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [9]: %timeit build_prolongation_operator_gboffi(p, qs)
31.8 µs ± 772 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [10]: %timeit build_prolongation_operator_gboffi2(p, qs)
26.6 µs ± 768 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

如您所见，最快的选项是完全矢量化的选项，但 @RaunaqJain's 和 @gboffi's 选项紧随其后。

注意

我的矢量化解决方案创建了一个布尔数组。如果您不希望使用 I.astype(...) 转换为所需的 dtype，或将其视为 np.uint8 数组：I.view(dtype=np.uint8)。

Answer 4

这是Raunaq Jain在their answer中提出的优化循环

for j in range(0,p):
    for k in range(0, qs[j]):
        # if BV i is a child of j, we set I[i,j] = 1
            val = f_map(j,k,qs)
            if I[val, j] == 0:
                I[val, j] = 1

这里是 f_map 函数，我在其中编辑了参数的名称以反映调用者使用的名称

def f_map(j,k,qs):
    if k < 0 or k > qs[j]:
        print('ERROR in f_map')
        return None
    result = k
    for i in range(0, j):
        result += qs[i]
    return result

我们首先注意到它始终是0 ≤ k < qs[j]，因为在k 上定义了循环，因此我们可以安全地删除健全性检查并写入

def f_map(j,k,qs):
    result = k
    for i in range(0, j):
        result += q[i]
    return result

现在，这是 sum 内置的文档字符串

签名：sum(iterable, start=0, /)
文档字符串：
返回一个“开始”值（默认值：0）加上一个可迭代的数字的总和

当iterable为空时，返回起始值。
此函数专门用于数值，可能会拒绝非数值类型。
类型：builtin_function_or_method

显然我们可以写

def f_map(j,k,qs):
    return sum(qs[:j], k)

很明显，我们也可以不使用函数调用

for j in range(0,p):
    for k in range(0, qs[j]):
        # if BV i is a child of j, we set I[i,j] = 1
            val = sum(qs[:j], k)
            if I[val, j] == 0:
                I[val, j] = 1

调用内置函数应该比函数调用和循环更有效，不是吗？

---

针对疯狂物理学家的言论

我们可以预先计算qs 的部分和以进一步加快速度

sqs = [sum(qs[:i]) for i in range(len(qs))] # there are faster ways...
...
for j in range(0,p):
    for k in range(0, qs[j]):
        # if BV i is a child of j, we set I[i,j] = 1
            val = k+sqs[j]
            if I[val, j] == 0:
                I[val, j] = 1