将 numba.jit 与 scipy.integrate.ode 一起使用

Posted

技术标签:

【中文标题】将 numba.jit 与 scipy.integrate.ode 一起使用【英文标题】:Using numba.jit with scipy.integrate.ode 【发布时间】:2015-12-21 01:42:44 【问题描述】:

使用numba.jitscipy.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

任何想法如何克服这个问题?

【问题讨论】:

没有jit 是否可以工作?您的 rhs 接受参数、求解器变化的 t 以及您提供的 X(通过额外的参数元组)。 是的,没有jit 它可以正常工作。额外参数元组是什么意思? 我也有类似的问题。 【参考方案1】:

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

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

【讨论】:

【参考方案2】:

您可以使用包装函数,但我认为它不会提高小型 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)

【讨论】:

以上是关于将 numba.jit 与 scipy.integrate.ode 一起使用的主要内容,如果未能解决你的问题,请参考以下文章

调用 numba jit 函数时,cProfile 会增加大量开销

numba - guvectorize 几乎比 jit 快

numba初体验

使用带有 numba njit 功能的字典

分析 numba 包装函数的覆盖率

numba安装和使用