numba-vectorize
参考文档： https://apachecn.github.io/numba-doc-zh/#/docs/17
vectorize装饰器:
允许python函数的标量入参使用 numpy 的 ufuncs。
import numba as nb
@nb.vectorize([nb.int32(nb.int32, nb.int32)], nopython=True)
def f(x, y):
return x + y
a = [1, 2, 3]
b = [4, 5, 6]
print("a + b: ", a + b) # a + b: [1, 2, 3, 4, 5, 6]
print("f(a, b): ", f(a, b)) # f(a, b): [5 7 9]
print("type(f(a, b)): ", type(f(a, b))) # type(f(a, b)): <class 'numpy.ndarray'>
c = [[1, 2, 3], [4, 5, 6]]
d = [[7, 8, 9], [10, 11, 12]]
print("f(c,d): ", f(c, d)) # [[8 10 12] [14 16 18]]
vectorize支持多个函数签名，多如果用多个函数签名的时候注意顺序，比如单精度浮点数在双精度浮点数之前。
@vectorize([int32(int32, int32),
int64(int64, int64),
float32(float32, float32),
float64(float64, float64)])
使用vectorize装饰之后，函数会自动获得numpy ufuncs的其他功能，比如缩小，累计或广播。ufunc的标准功能。
print(
"f.reduce(c) axis=0: %s, axis=1: %s" % (f.reduce(c), f.reduce(c, axis=1))
) # f.reduce(c) axis=0: [5 7 9], axis=1: [ 6 15]
guvectorize 装饰器
vectorize 是一次写一个元素的 ufunc, guvectorize 是允许编写输入任意数量数组元素的 ufunc，并获取和返回不同维度的数组。典型的例子是中值或卷积滤波器。
guvectorize 不返回结果，是将其作为函数参数，必须有函数填充。
@nb.guvectorize([(nb.int32[:], nb.int32, nb.int32[:])], "(n), ()->(n)", nopython=True)
def g(x, y, r):
for i in range(x.shape[0]):
r[i] = x[i] + y
a = [1, 2, 3, 4]
b = 10
print("g(a, b): ", g(a, b)) # g(a, b): [11 12 13 14]
底层的 python 代码是将一个给定的标量 y 添加到一维数组的所有元素中。
"(n),()->(n)" 表示输入一个n个元素的一维数组和一个标量(用符号()表示)并返回n个元素的一维数组。
1D 数组类型也可以接收标量参数，在上边例子中，第二个参数也可以声明为int[:] (), 在这种情况下，该值必须由 y[0] 读取。
我们可以检查一下能否支持不同维度的数组
# 2d + int
c = [[1, 2, 3], [4, 5, 6]]
print("g(c, 10): ", g(c, 10)) # g(c, 10): [[11 12 13] [14 15 16]]
# 2d + 1d
d = [10, 20]
print("g(c, d): ", g(c, d)) # g(c, d): [[11 12 13] [24 25 26]]
动态通用功能
如果给装饰器 vectorize 没有传递任何签名，python 函数将会构建动态通用函数或者DUDunc。
@nb.vectorize()
def d(x, y):
return x + y
print("d(1, 2)", d(1, 2)) # d(1, 2) 3
print(d.types) # ['ll->q']
print("d(1, 2.0)", d(1, 2.0)) # d(1, 2) 3.0
print(d.types) # ['ll->q', 'ld->d']
print("d([1, 2, 3], 10): ", d([1, 2, 3], 10)) # d([1, 2, 3], 10): [11 12 13]
print(d.types) # ['ll->q', 'ld->d']
在这种情况下，调用的顺序很重要，比如先传入浮点参数，那么任何带有整形的参数都会被转换成双精度浮点值。如： 先调用 d(1.0, 2.0) 得到3.0，再调用 d(1, 2) 也会得到3.0。
完整代码
import numba as nb
@nb.vectorize(
[nb.int32(nb.int32, nb.int32), nb.float64(nb.float64, nb.float64)], nopython=True
)
def f(x, y):
return x + y
a = [1.0, 2.0, 3.0]
b = [4, 5, 6]
# b = [4.0, 5.0, 6.0]
print("a + b: ", a + b) # a + b: [1, 2, 3, 4, 5, 6]
print("f(a, b): ", f(a, b)) # f(a, b): [5 7 9]
print("type(f(a, b)): ", type(f(a, b))) # type(f(a, b)): <class 'numpy.ndarray'>
c = [[1, 2, 3], [4, 5, 6]]
d = [[7, 8, 9], [10, 11, 12]]
print("f(c,d): ", f(c, d)) # [[8 10 12] [14 16 18]]
print(
"f.reduce(c) axis=0: %s, axis=1: %s" % (f.reduce(c), f.reduce(c, axis=1))
) # f.reduce(c) axis=0: [5 7 9], axis=1: [ 6 15]
print("-" * 20)
@nb.guvectorize([(nb.int32[:], nb.int32, nb.int32[:])], "(n), ()->(n)", nopython=True)
def g(x, y, r):
for i in range(x.shape[0]):
r[i] = x[i] + y
a = [1, 2, 3, 4]
b = 10
print("g(a, b): ", g(a, b)) # g(a, b): [11 12 13 14]
# 2d + int
c = [[1, 2, 3], [4, 5, 6]]
print("g(c, 10): ", g(c, 10)) # g(c, 10): [[11 12 13] [14 15 16]]
# 2d + 1d
d = [10, 20]
print("g(c, d): ", g(c, d)) # g(c, d): [[11 12 13] [24 25 26]]
print("-" * 20)
@nb.vectorize()
def d(x, y):
return x + y
print("d(1, 2)", d(1, 2)) # d(1, 2) 3
print(d.types) # ['ll->q']
print("d(1, 2.0)", d(1, 2.0)) # d(1, 2) 3.0
print(d.types) # ['ll->q', 'ld->d']
print("d([1, 2, 3], 10): ", d([1, 2, 3], 10)) # d([1, 2, 3], 10): [11 12 13]
print(d.types) # ['ll->q', 'ld->d']
