用于在相同长度的一维 numpy 数组上评估一维函数数组的高效算法答案

【问题标题】：Efficient algorithm for evaluating a 1-d array of functions on a same-length 1d numpy array用于在相同长度的一维 numpy 数组上评估一维函数数组的高效算法
【发布时间】：2015-02-11 16:51:12
【问题描述】：

我有一个（大）长度为 N 的 k 个不同函数数组，以及一个长度为 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 可以加快这个循环？

【问题讨论】：

可能想将此发布到 codereview.stackexchange.com
如果您跳过对np.where 的调用并使用布尔索引（即idx = function_indices == func_index），您会更快，其他一切都保持不变。
@Jaime - 此更改实际上将运行时间减慢到 0.5 秒
“where”行和最后一个函数评估行之间的时间间隔是多少？
所以是对where 的反复调用会害死你。您需要某种排序或 groupby 来组织索引一次，然后在循环中提供快速访问。

标签： python performance numpy scientific-computing numba

【解决方案1】：

这与您的（优秀的！）自我回答几乎相同，但少了一些繁琐。它在我的机器上似乎也稍微快了一点——根据粗略的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)可以被认为是一种二分搜索算法，它返回a中第一个值的索引等于或大于给定b 中的一个或多个值。

然后 zip 函数简单地并行迭代它的参数，从每个参数中返回一个项目，收集到一个元组中，然后将它们串在一起形成一个列表。（所以zip([1, 2, 3], ['a', 'b', 'c'], ['b', 'c', 'd']) 返回[(1, 'a', 'b'), (2, 'b', 'c'), (3, 'c', 'd')]。）这与for 语句内置的“解包”这些元组的能力一起，允许以简洁但富有表现力的方式并行迭代多个序列。

在这种情况下，我使用它来迭代 func_tables 中的函数以及 func_ranges 的两个不同步副本。这可确保end 变量中func_ranges 中的项目始终比start 变量中的项目领先一步。通过将None 附加到func_ranges，我确保最终块得到妥善处理——zip 在其任何一个参数用完项目时停止，从而切断序列中的最终值。方便的是，None 值也可用作开放式切片索引！

做同样事情的另一个技巧需要多几行，但内存开销较低，尤其是与 itertools 等效的 zip、izip 一起使用时：

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):
    ...

但是，这些基于生成器的低开销方法有时会比普通列表慢一些。另外，请注意，在 Python 3 中，zip 的行为更像izip。

【讨论】：

非常好。我可以通过我自己的独立测试来确认你的时间安排：你的解决方案比我的快 20%，你的语法更精简，我们的代码甚至在一些有害的边缘情况下也同意。这太棒了！
@aph，很高兴听到测试结果一致——我本来打算跑几个，但不得不突然离开。我也会为以后的访问者添加几句解释。
如果您不介意充实一些内容，那就太好了。特别是你对 zip 的使用，我觉得它并不直观，但看起来很优雅，所以我想在这样的实际应用程序中学习它的语法。
再次被撕毁，但刚刚找到时间发布。如果您还有其他问题，请告诉我！
嗯，这超出了“使命召唤”，senderle。但是，哇，一百万感谢您抽出宝贵的时间。这对我在几个主题上非常具有指导意义，包括关于 itertools 和生成器的奖励位。多年来，我一直在编写自己的自制版本的 groupby 计算（尽管不是在 python 中），并且以前从未见过这种计算做得这么好，所以这对我来说是一个真正令人大开眼界的 SO 答案。如果我可以投多个赞成票，我会的。干杯！

【解决方案2】：

感谢 hpaulj 提出采用 groupby 方法的建议。有很多用于此操作的固定例程，例如 Pandas DataFrames，但它们都带有数据结构初始化的开销成本，这是一次性的，但如果仅用于单次计算，成本可能会很高.

这是我的纯 numpy 解决方案，它比我使用的原始 where 循环快 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 倍。

【讨论】：

一个小建议。 x[a:] 等价于 x[a:None] - 所以不是中断最后一次调用，你不能将 unique_func_indices 转换为一个普通列表并附加一个 None 值吗？

【解决方案3】：

这是函数式编程在某种程度上被 Python 模拟的一个很好的例子。

现在，如果您想将您的函数应用于一组点，我建议您使用 numpy 的 ufunc 框架，它可以让您创建极快的函数矢量化版本。

【讨论】：