Question

我正在使用numpy运行python，我有一个循环，它被剥离，看起来像这样：

result = np.zeros(bins)
for i in xrange(bins):
    result[f(i)] += source[i]

这里，结果和源都是numpy数组，f是一组稍微复杂的算术运算。例如，f的一个简化示例可能看起来像

f = lambda x: min(int(width*pow(x, 3)), bins-1)

虽然f在其论证中通常不是单调的。

这个循环目前是我程序中的瓶颈。我设法将其他所有内容都矢量化，但我目前在这里难以理解如何做到这一点。如何对这个循环进行矢量化？

Answer 1

向量化f 主要思想是用基于矢量的numpy函数替换标量运算。例如，如果最初我们有

def f(x):
    return min(int(width*pow(x, 3)), bins-1)

然后我们可以使用

def fvec(x):
    return np.minimum((width*np.power(x, 3)).astype(int), bins-1)

一些Python标量函数和NumPy之间存在自然的对应关系矢量化函数：

| pow   | np.power   |
| min   | np.minimum |
| max   | np.maximum |
| floor | np.floor   |
| log   | np.log     |
| <     | np.less    |
| >     | np.greater |

矢量化函数采用输入数组并返回相同形状的数组。但是，还有其他结构可能不那么明显。比如说矢量化等效x if condition else y为np.where(condition, x, y)。

不幸的是，一般来说没有简单的捷径。翻译自矢量化函数的标量函数可能需要许多函数中的任何一个提供NumPy功能，以及广播和高级等NumPy概念索引。

例如，此时很容易替换

for i in range(bins):
    result[f(i)] += source[i]

使用integer-array indexed assignment：

result[fvec(np.arange(bins))] += source

但如果fvec(np.arange(bins))重复值，则会产生错误的结果。而是使用 np.bincount，因为当source表示相同的bin时，这会正确累积多个fvec(np.arange(bins))值：

result = np.bincount(fvec(np.arange(bins)), weights=source, minlength=bins)

import numpy as np
import pandas as pd

bins = 1000
width = 1.5
source = np.random.random(bins)

def fvec(x):
    return np.minimum((width*np.power(x, 3)).astype(int), bins-1)

def f(x):
    return min(int(width*pow(x, 3)), bins-1)

def orig():
    result = np.zeros(bins)
    for i in range(bins):
        result[f(i)] += source[i]
    return result

def alt():
    result = np.bincount(fvec(np.arange(bins)), weights=source, minlength=bins)
    return result

assert np.allclose(orig(), alt())

对于上面bins=1000的示例，alt比orig（在我的机器上）快约62倍：

In [194]: %timeit orig()
1000 loops, best of 3: 1.37 ms per loop

In [195]: %timeit alt()
10000 loops, best of 3: 21.8 µs per loop

alt超过orig的速度优势将随着orig for-loop所需的迭代次数的增加而增加 - 也就是{ {1}}增加。

使用可变数组索引进行矢量化

1 个答案: