NumbaPro的@vectorize装饰器似乎是利用多核处理器进行数值计算的一种巧妙方法。不幸的是,以下相当小的例子会产生错误:
import numpy as np
from scipy.integrate import odeint
from numbapro import vectorize, float64, int64, jit
@vectorize([float64[:](float64[:], float64, float64, int64, float64, float64[:], float64)], target='parallel')
def heat_equation(x, t, a, p, h, dxdt, pi):
for i in xrange(p-1):
dxdt[i] = a * (x[i-1] - 2 * x[i] + x[i+1]) / h / h
dxdt[0] = 2*pi*np.cos(2*pi*t)
dxdt[p-1] = 0
return dxdt
if __name__ == '__main__':
p = 200
h = 1. / (p-1)
a = 0.125
x = np.linspace(0, 1, p)
y0 = np.zeros(p)
dxdt = np.zeros(p, dtype=np.float64)
pi = np.pi
for i in xrange(p):
y0[i] = 0
timeVector = np.linspace(0, 10, 100)
solVector = odeint(heat_equation, y0, timeVector, args=(a, p, h, dxdt, pi))
print solVector[-1, p/2]
上面的代码使用@jit装饰器工作得很好,但尝试@vectorize会出现以下错误:
ValueError: format number 1 of "array(float64, 1d, A)" is not recognized
显然,装饰器参数存在问题,但类型签名对我来说是正确的。是否有一些我不遵守的额外限制?
编辑:修改代码以避免在装饰函数中使用numpy.zeros和numpy.pi,根据下面的Bakuriu的有用评论并相应地调整收到的错误。
答案 0 :(得分:0)
用@jit修饰的代码通常不能直接移植到@vectorize。 @vectorize将装饰函数转换为NumPy ufunc核心(http://docs.scipy.org/doc/numpy/reference/ufuncs.html)。对@vectorize的额外限制是核心函数(正在修饰的函数)的所有参数都必须是标量。如果需要数组参数,请使用@guvectorize。有关示例,请参阅http://numba.pydata.org/numba-doc/dev/ufuncs.html#generalized-ufuncs。