快速python numpy的功能在哪里？

Question

我在多个for循环中多次使用numpy的where函数，但是它变得太慢了。有什么方法可以更快地执行此功能？我读到您应该尝试内联进行循环，并在for循环之前为函数创建局部变量，但是似乎没有什么可以提高速度的（<1％）。在len(UNIQ_IDS)〜800emiss_data和obj_data有numpy的ndarrays具有形状=（2600,5200）。我用import profile得到的瓶颈在哪里手柄，并且where在for循环是一个大的。

import numpy as np
max = np.max
where = np.where
MAX_EMISS = [max(emiss_data[where(obj_data == i)]) for i in UNIQ_IDS)]

Answer 1

事实证明，在这种情况下，纯Python循环比NumPy索引（或对np.where的调用）要快得多。

考虑以下替代方案：

import numpy as np
import collections
import itertools as IT

shape = (2600,5200)
# shape = (26,52)
emiss_data = np.random.random(shape)
obj_data = np.random.random_integers(1, 800, size=shape)
UNIQ_IDS = np.unique(obj_data)

def using_where():
    max = np.max
    where = np.where
    MAX_EMISS = [max(emiss_data[where(obj_data == i)]) for i in UNIQ_IDS]
    return MAX_EMISS

def using_index():
    max = np.max
    MAX_EMISS = [max(emiss_data[obj_data == i]) for i in UNIQ_IDS]
    return MAX_EMISS

def using_max():
    MAX_EMISS = [(emiss_data[obj_data == i]).max() for i in UNIQ_IDS]
    return MAX_EMISS

def using_loop():
    result = collections.defaultdict(list)
    for val, idx in IT.izip(emiss_data.ravel(), obj_data.ravel()):
        result[idx].append(val)
    return [max(result[idx]) for idx in UNIQ_IDS]

def using_sort():
    uind = np.digitize(obj_data.ravel(), UNIQ_IDS) - 1
    vals = uind.argsort()
    count = np.bincount(uind)
    start = 0
    end = 0
    out = np.empty(count.shape[0])
    for ind, x in np.ndenumerate(count):
        end += x
        out[ind] = np.max(np.take(emiss_data, vals[start:end]))
        start += x
    return out

def using_split():
    uind = np.digitize(obj_data.ravel(), UNIQ_IDS) - 1
    vals = uind.argsort()
    count = np.bincount(uind)
    return [np.take(emiss_data, item).max()
            for item in np.split(vals, count.cumsum())[:-1]]

for func in (using_index, using_max, using_loop, using_sort, using_split):
    assert using_where() == func()

以下是基准测试shape = (2600,5200)：

In [57]: %timeit using_loop()
1 loops, best of 3: 9.15 s per loop

In [90]: %timeit using_sort()
1 loops, best of 3: 9.33 s per loop

In [91]: %timeit using_split()
1 loops, best of 3: 9.33 s per loop

In [61]: %timeit using_index()
1 loops, best of 3: 63.2 s per loop

In [62]: %timeit using_max()
1 loops, best of 3: 64.4 s per loop

In [58]: %timeit using_where()
1 loops, best of 3: 112 s per loop

因此，using_loop（纯Python）的运行速度比快11倍以上using_where。

我不完全确定为什么纯Python比NumPy更快。我的猜测是纯Python版本一次通过两个数组（是，是双关语）。它利用了这样一个事实，即尽管进行了所有花哨的索引编制，但实际上我们只想访问每个值一次。因此，它避免了必须确定每个值所属的组的问题emiss_data。但这只是模糊的推测。我不知道在进行基准测试之前会更快。