在numpy数组上映射函数的最有效方法

Question

在numpy数组上映射函数最有效的方法是什么？我在当前项目中的做法如下：

import numpy as np 

x = np.array([1, 2, 3, 4, 5])

# Obtain array of square of each element in x
squarer = lambda t: t ** 2
squares = np.array([squarer(xi) for xi in x])

然而，这看起来可能是非常低效的，因为我使用列表理解来构造新的数组作为Python列表，然后再将其转换回Numpy数组。

我们能做得更好吗？

Answer 1

正如@user2357112所指出的，应用函数的“直接”方法总是在Numpy数组上映射函数的最快和最简单的方法：

import numpy as np
x = np.array([1, 2, 3, 4, 5])
f = lambda x: x ** 2
squares = f(x)

通常避免使用np.vectorize，因为它的性能不好，并且存在（或曾经存在）许多问题。如果您正在处理其他数据类型，则可能需要研究下面所示的其他方法。

下面是一些比较映射函数的三种方法的简单测试，本例使用Python3.6和Numpy1.15.4。首先，测试的设置功能：

import timeit
import numpy as np

f = lambda x: x ** 2
vf = np.vectorize(f)

def test_array(x, n):
    t = timeit.timeit(
        'np.array([f(xi) for xi in x])',
        'from __main__ import np, x, f', number=n)
    print('array: {0:.3f}'.format(t))

def test_fromiter(x, n):
    t = timeit.timeit(
        'np.fromiter((f(xi) for xi in x), x.dtype, count=len(x))',
        'from __main__ import np, x, f', number=n)
    print('fromiter: {0:.3f}'.format(t))

def test_direct(x, n):
    t = timeit.timeit(
        'f(x)',
        'from __main__ import x, f', number=n)
    print('direct: {0:.3f}'.format(t))

def test_vectorized(x, n):
    t = timeit.timeit(
        'vf(x)',
        'from __main__ import x, vf', number=n)
    print('vectorized: {0:.3f}'.format(t))

用五个元素进行测试（从最快到最慢排序）：

x = np.array([1, 2, 3, 4, 5])
n = 100000
test_direct(x, n)      # 0.265
test_fromiter(x, n)    # 0.479
test_array(x, n)       # 0.865
test_vectorized(x, n)  # 2.906

有100个元素：

x = np.arange(100)
n = 10000
test_direct(x, n)      # 0.030
test_array(x, n)       # 0.501
test_vectorized(x, n)  # 0.670
test_fromiter(x, n)    # 0.883

以及1000个或更多阵列元素：

x = np.arange(1000)
n = 1000
test_direct(x, n)      # 0.007
test_fromiter(x, n)    # 0.479
test_array(x, n)       # 0.516
test_vectorized(x, n)  # 0.945

不同版本的Python/NumPy和编译器优化将有不同的结果，因此请针对您的环境执行类似的测试。

Answer 2

使用numpy.vectorize怎么样。

import numpy as np
x = np.array([1, 2, 3, 4, 5])
squarer = lambda t: t ** 2
vfunc = np.vectorize(squarer)
vfunc(x)
# Output : array([ 1,  4,  9, 16, 25])

Answer 3

我已经测试了所有建议的方法，加上np.array（map（f， x））和PERPLANT（我的一个小项目）。

信息#1：如果您可以使用numpy的本机函数，请这样做。

如果您尝试矢量化的函数已经矢量化（如原始帖子中的x**2示例），则使用该函数比使用其他函数快得多（注意日志比例）：

如果您真的需要矢量化，那么使用哪种变体其实并不重要。

代码再现的情节：

import numpy as np
import perfplot
import math

def f(x):
    # return math.sqrt(x)
    return np.sqrt(x)

vf = np.vectorize(f)

def array_for(x):
    return np.array([f(xi) for xi in x])

def array_map(x):
    return np.array(list(map(f, x)))

def fromiter(x):
    return np.fromiter((f(xi) for xi in x), x.dtype)

def vectorize(x):
    return np.vectorize(f)(x)

def vectorize_without_init(x):
    return vf(x)

perfplot.show(
    setup=np.random.rand,
    n_range=[2 ** k for k in range(20)],
    kernels=[f, array_for, array_map, fromiter,
             vectorize, vectorize_without_init],
    xlabel="len(x)",
)