使用numba.jit和scipy.integrate.ode

Question

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

有什么想法可以克服这个问题吗？

- germannp

---

没有使用jit，它能正常工作吗？您的rhs接受两个参数，一个是求解器变化的t，另一个是您提供的X（通过额外的参数元组）。 - hpaulj

是的，没有“jit”它可以正常工作。您所说的额外参数元组是什么意思？ - germannp

我有类似的问题。 - LifeWorks

Answer 1

您可以使用包装函数，但我认为对于小的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)

- TheIdealis

Answer 2

我不知道原因或解决方案，但在这种情况下Theano对加速计算有很大帮助。Theano基本上会编译numpy表达式，因此仅在您可以将rhs编写为多维数组的表达式时才有帮助（而 jit 知道 for 和 friends）。它还知道一些代数并优化计算。

此外，Theano 可以编译GPU（这是我首次尝试使用 numba.jit 的原因）。然而，由于开销问题，使用GPU仅在大型系统（可能有一百万个方程式）中才能提高性能。

- germannp