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

时间:2014-06-06 17:33:41

标签: python numpy jit anaconda numba

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

1 个答案:

答案 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