Использование 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 и друзей). Он также знает некоторую алгебру и оптимизирует вычисление.

Кроме того, Anano может скомпилировать для GPU (это было моей причиной, чтобы вначале попробовать numba.jit). Однако использование GPU позволило повысить производительность для огромных систем (возможно, миллион уравнений) из-за накладных расходов.

Answer 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)