结合使用numba.jit和scipy.integrate.ode

Question

使用numba.jit加快了右手端计算odeint从scipy.integrate正常工作：

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

任何想法如何克服这个？

Answer 1

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

此外，Theano 可以针对 GPU 进行编译（这是我numba.jit首先尝试的原因）。然而，由于开销，使用 GPU 只能提高大型系统（可能有一百万个方程）的性能。