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

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

使用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

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

【问题讨论】：

没有jit 是否可以工作？您的 rhs 接受参数、求解器变化的 t 以及您提供的 X（通过额外的参数元组）。
是的，没有jit 它可以正常工作。额外参数元组是什么意思？
我也有类似的问题。

标签： python scipy ode integrate odeint

【解决方案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)

【讨论】：