使用numba.jit和scipy.integrate.ode

时间:2015-09-23 16:18:24

标签: python scipy ode integrate odeint

使用numba.jit来加速scipy.integrate odeint的右侧计算效果很好:

from scipy.integrate import ode, odeint
from numba import jit

@jit
def rhs(t, X):
    return 1

X = odeint(rhs, 0, np.linspace(0, 1, 11))

但是使用integrate.ode这样:

solver = ode(rhs)
solver.set_initial_value(0, 0)
while solver.successful() and solver.t < 1:
    solver.integrate(solver.t + 0.1)

使用装饰器@jit生成以下错误:

capi_return is NULL
Call-back cb_f_in_dvode__user__routines failed.
Traceback (most recent call last):
  File "sandbox/numba_cubic.py", line 15, in <module>
    solver.integrate(solver.t + 0.1)
  File "/home/pgermann/Software/anaconda3/lib/python3.4/site-packages/scipy/integrate/_ode.py", line 393, in integrate
    self.f_params, self.jac_params)
  File "/home/pgermann/Software/anaconda3/lib/python3.4/site-packages/scipy/integrate/_ode.py", line 848, in run
    y1, t, istate = self.runner(*args)
TypeError: not enough arguments: expected 2, got 1

任何想法如何克服这个?

2 个答案:

答案 0 :(得分:0)

我不知道原因或解决方案,但在这种情况下,Theano帮助加快了计算速度。 Theano基本上编译了numpy表达式,因此只有当你能将rhs写成多维数组的表达式时才会有用(而jit知道for和朋友)。它也知道一些代数并优化计算。

除了Theano可以为GPU编译(这是我首先尝试numba.jit的原因)。然而,由于开销,使用GPU只能提高大型系统(可能是一百万个方程式)的性能。

答案 1 :(得分:0)

您可以使用包装函数,但我认为它不会提高小rhs函数的性能。

@jit(nopython=True)
def rhs(t, X):
    return 1

def wrapper(t, X):
    return rhs(t, X)

solver = ode(wrapper)
solver.set_initial_value(0, 0)
while solver.successful() and solver.t < 1:
solver.integrate(solver.t + 0.1)