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 的有用评论,并相应地调整收到的错误。