用于评估相同长度的1d numpy阵列上的1-d函数数组的高效算法

时间:2015-02-11 16:51:13

标签: python performance numpy scientific-computing numba

我有一个由k个不同函数组成的(大)长度为N的数组,以及一个长度为n的abcissa数组。我想评估abcissa中的函数返回一个长度为N的纵坐标数组,而且关键的是,我需要非常快速地完成它。

我在调用np.where时尝试了以下循环,这太慢了:

创建一些假数据来说明问题:

def trivial_functional(i): return lambda x : i*x
k = 250
func_table = [trivial_functional(j) for j in range(k)]
func_table = np.array(func_table) # possibly unnecessary

我们有250个不同功能的表格。现在我创建一个包含这些函数的许多重复条目的大型数组,以及一组应该评估这些函数的相同长度的点。

Npts = 1e6
abcissa_array = np.random.random(Npts)
function_indices = np.random.random_integers(0,len(func_table)-1,Npts)
func_array = func_table[function_indices]

最后,循环遍历数据使用的每个函数,并在相关点集上进行评估:

desired_output = np.zeros(Npts)
for func_index in set(function_indices):
    idx = np.where(function_indices==func_index)[0]
    desired_output[idx] = func_table[func_index](abcissa_array[idx])

这个循环在我的笔记本电脑上需要大约0.35秒,这是我的代码中最大的瓶颈。

有没有人看到如何避免对np.where的盲查询调用?是否巧妙地使用了numba来加速这种循环?

3 个答案:

答案 0 :(得分:4)

这与你的(优秀的!)自我答案几乎完全相同,但有点不那么严格。我的机器上的速度似乎稍微快一点 - 基于粗略的test约为30毫秒。

def apply_indexed_fast(array, func_indices, func_table):
    func_argsort = func_indices.argsort()
    func_ranges = list(np.searchsorted(func_indices[func_argsort], range(len(func_table))))
    func_ranges.append(None)
    out = np.zeros_like(array)
    for f, start, end in zip(func_table, func_ranges, func_ranges[1:]):
        ix = func_argsort[start:end]
        out[ix] = f(array[ix])
    return out

与您的一样,这会将一系列argsort索引拆分为多个块,每个块对应func_table中的一个函数。然后,它使用每个块为其相应的函数选择输入和输出索引。要确定块边界,它使用np.searchsorted而不是np.unique - 其中searchsorted(a, b)可以被认为是二进制搜索算法,它返回{{1}中第一个值的索引等于或大于a中给定值的值。

然后zip函数简单地并行迭代它的参数,从每个参数中返回一个项目,一起收集在一个元组中,并将它们串联到一个列表中。 (因此b会返回zip([1, 2, 3], ['a', 'b', 'c'], ['b', 'c', 'd'])。)这与[(1, 'a', 'b'), (2, 'b', 'c'), (3, 'c', 'd')]语句的内置功能一致,即可解压缩"那些元组,允许一种简洁但富有表现力的方式来并行迭代多个序列。

在这种情况下,我使用它来迭代for中的函数以及func_tables的两个不同步副本。这样可以确保func_ranges变量中func_ranges项的项始终比end变量中的项目提前一步。通过将start附加到None,我确保优雅地处理最后一个块 - func_ranges在其任何一个参数用完项目时停止,这会切断最终值序列。方便的是,zip值也可以作为开放式切片索引!

执行相同操作的另一个技巧需要更多行,但内存开销较低,尤其是与None等效itertoolsizip一起使用时:

zip

然而,这些低开销的基于发生器的方法有时可能比香草列表慢一点。另请注意,在Python 3中,range_iter_a = iter(func_ranges) # create generators that iterate over the range_iter_b = iter(func_ranges) # values in `func_ranges` without making copies next(range_iter_b, None) # advance the second generator by one for f, start, end in itertools.izip(func_table, range_iter_a, range_iter_b): ... 的行为更像zip

答案 1 :(得分:2)

感谢hpaulj建议采用 groupby 方法。这个操作有许多固定的例程,例如Pandas DataFrames,但它们都带有数据结构初始化的开销成本,这只是一次性的,但如果仅用于一次计算则成本很高。

这是我的纯粹numpy解决方案,比我正在使用的原始 循环快13倍。 结果摘要是我使用 np.argsort np.unique 以及一些花哨的索引体操。

首先我们对函数索引进行排序,然后找到每个新索引开始的排序数组的元素

idx_funcsort = np.argsort(function_indices)
unique_funcs, unique_func_indices = np.unique(function_indices[idx_funcsort], return_index=True)

现在不再需要盲查,因为我们确切地知道排序数组的哪个切片对应于每个唯一函数。所以我们仍然遍历每个被调用的函数,但是没有调用 where:

for func_index in range(len(unique_funcs)-1):
    idx_func = idx_funcsort[unique_func_indices[func_index]:unique_func_indices[func_index+1]]
    func = func_table[unique_funcs[func_index]]
    desired_output[idx_func] = func(abcissa_array[idx_func])

这涵盖了除最终索引之外的所有索引,由于Python索引约定,我们需要单独调用这些索引:

func_index = len(unique_funcs)-1
idx_func = idx_funcsort[unique_func_indices[func_index]:]
func = func_table[unique_funcs[func_index]]
desired_output[idx_func] = func(abcissa_array[idx_func])

这给 where 循环(簿记健全性检查)提供了相同的结果,但是这个循环的运行时间是0.027秒,比我原来的计算速度提高了13倍。

答案 2 :(得分:0)

这是函数式编程的一个很好的例子,在Python中有点模仿。

现在,如果您想将您的功能应用于一组点,我建议numpy ufunc框架,这将允许您创建极快的矢量化你的功能版本。