python – Numba的jit无法编译具有另一个函数作为输入的函数

2024-04-02 10:40:16

我试图在数字上解决一个允许离散跳跃的ODE.我正在使用Euler方法,并希望Numba的jit可以帮助我加快这个过程(现在脚本需要300秒才能运行,我需要它运行200次).

这是我简化的第一次尝试：

import numpy as np
from numba import jit

dt = 1e-5
T = 1
x0 = 1
noiter = int(T / dt)
res = np.zeros(noiter)

def fdot(x, t):
    return -x + t / (x + 1) ** 2

def solve_my_ODE(res, fdot, x0, T, dt):
    res[0] = x0
    noiter = int(T / dt)
    for i in range(noiter - 1):
        res[i + 1] = res[i] + dt * fdot(res[i], i * dt)
        if res[i + 1] >= 2:
            res[i + 1] -= 2
    return res

%timeit fdot(x0, T)
%timeit solve_my_ODE(res, fdot, x0, T, dt)
    ->The slowest run took 8.38 times longer than the fastest. This could mean that an intermediate result is being cached 
    ->1000000 loops, best of 3: 465 ns per loop
    ->10 loops, best of 3: 122 ms per loop

@jit(nopython=True)
def fdot(x, t):
    return -x + t / (x + 1) ** 2
%timeit fdot(x0, T)
%timeit solve_my_ODE(res, fdot, x0, T, dt)
    ->The slowest run took 106695.67 times longer than the fastest. This could mean that an intermediate result is being cached 
    ->1000000 loops, best of 3: 240 ns per loop
    ->10 loops, best of 3: 99.3 ms per loop

@jit(nopython=True)
def solve_my_ODE(res, fdot, x0, T, dt):
    res[0] = x0
    noiter = int(T / dt)
    for i in range(noiter - 1):
        res[i + 1] = res[i] + dt * fdot(res[i], i * dt)
        if res[i + 1] >= 2:
            res[i + 1] -= 2
    return res
%timeit fdot(x0, T)
%timeit solve_my_ODE(res, fdot, x0, T, dt)
    ->The slowest run took 10.21 times longer than the fastest. This could mean that an intermediate result is being cached 
    ->1000000 loops, best of 3: 274 ns per loop
    ->TypingError                               Traceback (most recent call last)
ipython-input-10-27199e82c72c> in <module>()
  1 get_ipython().magic('timeit fdot(x0, T)')
----> 2 get_ipython().magic('timeit solve_my_ODE(res, fdot, x0, T, dt)')

(...)


TypingError: Failed at nopython (nopython frontend)
Undeclared pyobject(float64, float64)
File "<ipython-input-9-112bd04325a4>", line 6

我不明白为什么我得到这个错误.我怀疑numba无法识别输入字段fdot(这是一个已经用Numba编译的python函数).

由于我对Numba这么新,我有几个问题

>我能做些什么让Numba了解输入字段fdot是一个函数？
>在函数fdot上使用JIT“仅”导致减少50％.我应该期待更多吗？或者这是正常的吗？
>这个脚本看起来像是一种合理的方法来模拟具有离散跳跃的ODE吗？在数学上,这相当于用delta函数求解ODE.

Numba版本是0.17

解决方法:

你认为numba不认识fdot是一个numba编译函数你是对的.我不认为你可以把它识别为函数参数,但你可以使用这种方法(使用变量捕获,因此在构建函数时已知fdot)来构建ODE求解器：

def make_solver(f):
    @jit(nopython=True)
    def solve_my_ODE(res, x0, T, dt):
        res[0] = x0
        noiter = int(T / dt)
        for i in range(noiter - 1):
            res[i + 1] = res[i] + dt * f(res[i], i * dt)
            if res[i + 1] >= 2:
                res[i + 1] -= 2
        return res
    return solve_my_ODE

fdot_solver = make_solver(fdot) # call this for each function you 
      # want to make an ODE solver for

这是一个替代版本,不需要您将res传递给它.只有循环加速,但因为那是慢速位,这是唯一重要的位.

def make_solver_2(f):
    @jit
    def solve_my_ODE(x0, T, dt):
        # this bit ISN'T in no python mode
        noiter = int(T / dt)
        res = np.zeros(noiter)
        res[0] = x0
        # but the loop is nopython (so fast)
        for i in range(noiter - 1):
            res[i + 1] = res[i] + dt * f(res[i], i * dt)
            if res[i + 1] >= 2:
                res[i + 1] -= 2
        return res
    return solve_my_ODE

我更喜欢这个版本,因为它为你分配了返回值,所以它更容易使用.尽管如此,这可能会从你真正的问题中略微转移.

在时间方面我得到(在几秒钟内,20次迭代)：

> 6.90394687653(仅适用于numba的fdot)
> 0.0584900379181(版本1)
> 0.0640540122986(对于版本2 – 即它稍微慢一点但更容易使用)

因此,它大约快了100倍 – 加速循环产生了很大的不同！

第三个问题：“这个脚本看起来像是一种合理的方法来模拟具有离散跳跃的ODE吗？在数学上,这相当于解决具有delta函数的ODE.”我真的不知道.抱歉!

码农公寓