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

Question

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

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

这里，result 和 source 都是 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 因为当fvec(np.arange(bins)) 表示同一个 bin 时，这会正确累积多个 source 值：

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 所需迭代次数的增加而增加——即随着bins 的增加。