使用NumbaPro的@vectorize装饰器时出错

Question

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

显然，装饰器的参数有问题，但我觉得类型签名是正确的。我是不是忽略了什么额外的限制？

编辑：根据Bakuriu下面的有用评论，我修改了代码，避免在装饰的函数中使用numpy.zeros和numpy.pi，并相应地调整了收到的错误。

Answer 1

用@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。