在 NumPy 数组的每个单元格处对函数进行有效评估

Posted 2023-02-19

【中文标题】在 NumPy 数组的每个单元格处对函数进行有效评估

【英文标题】：Efficient evaluation of a function at every cell of a NumPy array

【发布时间】：2011-12-03 19:48:16

【问题描述】：

给定一个NumPy数组A，将same函数f应用于每个细胞？

假设我们将分配给 A(i,j) f(A(i,j))。

函数 f 没有二进制输出，因此掩码操作无济于事。

“明显的”双循环迭代（通过每个单元）是最优解吗？

【问题讨论】：

numpy.apply_over_axes

【参考方案1】：

以上所有答案都比较好，但是如果您需要使用自定义函数进行映射，并且您有numpy.ndarray，则需要保留数组的形状。

我只比较了两个，但它会保留ndarray 的形状。我使用了包含 100 万个条目的数组进行比较。这里我使用平方函数。我正在介绍 n 维数组的一般情况。对于二维，只需将 iter 设置为 2D。

import numpy, time

def A(e):
    return e * e

def timeit():
    y = numpy.arange(1000000)
    now = time.time()
    numpy.array([A(x) for x in y.reshape(-1)]).reshape(y.shape)        
    print(time.time() - now)
    now = time.time()
    numpy.fromiter((A(x) for x in y.reshape(-1)), y.dtype).reshape(y.shape)
    print(time.time() - now)
    now = time.time()
    numpy.square(y)  
    print(time.time() - now)

输出

>>> timeit()
1.162431240081787    # list comprehension and then building numpy array
1.0775556564331055   # from numpy.fromiter
0.002948284149169922 # using inbuilt function

在这里您可以清楚地看到numpy.fromiter 用户方功能，使用您的任何选择。如果您的函数依赖于数组索引i, j，则迭代数组的大小，如for ind in range(arr.size)，使用numpy.unravel_index 根据您的一维索引和数组形状numpy.unravel_index 获取i, j, ..

这个答案的灵感来自我对其他问题here的回答

【参考方案2】：

当二维数组（或 nd 数组）是 C 或 F 连续时，将函数映射到二维数组的任务实际上与将函数映射到一维数组的任务相同- 我们只需要这样看待它，例如通过np.ravel(A,'K')。

已经讨论了一维数组的可能解决方案，例如here。

但是，当二维数组的内存不连续时，情况会稍微复杂一些，因为如果轴以错误的顺序处理，我们希望避免可能的缓存未命中。

Numpy 已经有一台机器可以以最佳顺序处理轴。使用这种机器的一种可能性是np.vectorize。然而，numpy 在np.vectorize 上的文档指出它“主要是为了方便，而不是为了性能”——一个慢的python 函数仍然是一个慢的python 函数，带有整个相关的开销！另一个问题是其巨大的内存消耗 - 例如，请参见 SO-post。

当一个人想要一个 C 函数的性能但使用 numpy 的机制时，一个好的解决方案是使用 numba 来创建 ufunc，例如：

# runtime generated C-function as ufunc
import numba as nb
@nb.vectorize(target="cpu")
def nb_vf(x):
    return x+2*x*x+4*x*x*x

它很容易击败np.vectorize，但同样的功能也可以作为 numpy-array 乘法/加法执行，即

# numpy-functionality
def f(x):
    return x+2*x*x+4*x*x*x

# python-function as ufunc
import numpy as np
vf=np.vectorize(f)
vf.__name__="vf"

有关时间测量代码，请参阅此答案的附录：

Numba 的版本（绿色）比 python 函数（即np.vectorize）快大约 100 倍，这并不奇怪。但它也比 numpy-functionality 快 10 倍左右，因为 numbas 版本不需要中间数组，因此更有效地使用缓存。

---

虽然 numba 的 ufunc 方法在可用性和性能之间取得了很好的平衡，但它仍然不是我们能做的最好的。然而，对于任何任务来说，没有灵丹妙药或最适合的方法 - 人们必须了解限制是什么以及如何减轻这些限制。

例如，对于超越函数（例如exp、sin、cos），numba 并没有提供任何优于 numpy 的 np.exp 的优势（没有创建临时数组 - 速度的主要来源 -向上）。但是，我的 Anaconda 安装使用英特尔的 VML 来处理向量 bigger than 8192 - 如果内存不连续，它就无法做到这一点。所以最好将元素复制到一个连续的内存中以便能够使用英特尔的 VML：

import numba as nb
@nb.vectorize(target="cpu")
def nb_vexp(x):
    return np.exp(x)

def np_copy_exp(x):
    copy = np.ravel(x, 'K')
    return np.exp(copy).reshape(x.shape) 

为了比较的公平，我关闭了 VML 的并行化（见附录中的代码）：

正如我们所看到的，一旦 VML 启动，复制的开销就会得到补偿。然而，一旦数据变得对于 L3 缓存来说太大，优势就会变得微乎其微，因为任务再次成为内存带宽限制。

另一方面，numba 也可以使用英特尔的 SVML，如 this post 中所述：

from llvmlite import binding
# set before import
binding.set_option('SVML', '-vector-library=SVML')

import numba as nb

@nb.vectorize(target="cpu")
def nb_vexp_svml(x):
    return np.exp(x)

并使用具有并行化效果的 VML：

numba 的版本开销较小，但对于某些大小，即使有额外的复制开销，VML 也能胜过 SVML - 这并不奇怪，因为 numba 的 ufunc 没有并行化。

---

列表：

A.多项式函数的比较：

import perfplot
perfplot.show(
    setup=lambda n: np.random.rand(n,n)[::2,::2],
    n_range=[2**k for k in range(0,12)],
    kernels=[
        f,
        vf, 
        nb_vf
        ],
    logx=True,
    logy=True,
    xlabel='len(x)'
    ) 

B. exp的对比：

import perfplot
import numexpr as ne # using ne is the easiest way to set vml_num_threads
ne.set_vml_num_threads(1)
perfplot.show(
    setup=lambda n: np.random.rand(n,n)[::2,::2],
    n_range=[2**k for k in range(0,12)],
    kernels=[
        nb_vexp, 
        np.exp,
        np_copy_exp,
        ],
    logx=True,
    logy=True,
    xlabel='len(x)',
    )

【参考方案3】：

我相信我找到了更好的解决方案。将函数更改为 python 通用函数的想法（参见documentation），可以在后台进行并行计算。

可以用C编写自己定制的ufunc，当然效率更高，或者调用np.frompyfunc，这是内置的工厂方法。经测试，这个比np.vectorize效率高：

f = lambda x, y: x * y
f_arr = np.frompyfunc(f, 2, 1)
vf = np.vectorize(f)
arr = np.linspace(0, 1, 10000)

%timeit f_arr(arr, arr) # 307ms
%timeit f_arr(arr, arr) # 450ms

我还测试了更大的样本，并且改进是成比例的。其他方法性能对比见this post

【参考方案4】：

您可以只 vectorize 该函数，然后在每次需要时将其直接应用于 Numpy 数组：

import numpy as np

def f(x):
    return x * x + 3 * x - 2 if x > 0 else x * 5 + 8

f = np.vectorize(f)  # or use a different name if you want to keep the original f

result_array = f(A)  # if A is your Numpy array

在矢量化时直接指定显式输出类型可能会更好：

f = np.vectorize(f, otypes=[np.float])

【讨论】：

恐怕矢量化函数不能比“手动”双循环迭代和遍历所有数组元素更快。特别是，因为它将结果存储到 新 创建的变量（而不是直接存储到初始输入）。非常感谢您的回复:)

@Peter：啊，现在我看到您在原始问题中提到将结果分配回前一个数组。很抱歉我在第一次阅读时错过了它。是的，在那种情况下，双循环必须更快。但是您是否也在阵列的扁平视图上尝试了一个循环？这可能会稍微快一些，因为您节省了一点循环开销，并且 Numpy 需要在每次迭代时少做一次乘法和加法（用于计算数据偏移量）。此外，它适用于任意尺寸的数组。不过，在非常小的阵列上可能会更慢。

注意vectorize 函数描述中给出的警告：提供vectorize 函数主要是为了方便，而不是为了性能。该实现本质上是一个 for 循环。 所以这很可能根本不会加速这个过程。

注意vectorize如何判断返回类型。这产生了错误。 frompyfunc 快一点，但返回一个 dtype 对象数组。两者都提要标量，而不是行或列。

@Gabriel 只需将np.vectorize 扔到我的函数（使用 RK45）上，我的速度就会提高约 20 倍。

【参考方案5】：

如果您使用数字和f(A(i,j)) = f(A(j,i))，您可以使用scipy.spatial.distance.cdist 将f 定义为A(i) 和A(j) 之间的距离。

【参考方案6】：

类似的问题是：Mapping a NumPy array in place。 如果你可以为你的 f() 找到一个ufunc，那么你应该使用 out 参数。