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

Question

我正在尝试对允许离散跳跃的 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（这是一个 python 函数，顺便说一下，它已经用 Numba 编译了）。

由于我是 Numba 的新手，我有几个问题

我该怎么做才能让 Numba 理解输入字段 fdot 是一个函数？
对函数 fdot 使用 JIT "only" 导致减少 50%。我应该期待更多吗？或者这是正常的吗？
这个脚本看起来像是模拟具有离散跳跃的 ODE 的合理方法吗？从数学上讲，这等同于用 delta 函数求解 ODE。

Numba 版本为 0.17

Answer 1

您认为 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

我更喜欢这个版本，因为它为你分配了return值，所以它更容易使用。不过，这与您的实际问题略有不同。

就我得到的时间而言（以秒为单位，迭代 20 次）：

6.90394687653（仅适用于 numba 中的 fdot）
0.0584900379181（对于版本 1）
0.0640540122986（对于版本 2 - 即速度稍慢但更易于使用）

因此，它大约快 100 倍 - 加速循环会产生很大的不同！

你的第三个问题："Does this script look like a reasonable way to simulate an ODE with discrete jumps? Mathematically this is equivalent at solving an ODE with delta functions."我真的不知道。抱歉！

Answer 2

最后一点：

在目前的形式中，它甚至不是一个有效的实现行为良好的 ODE。它过早停止了一步，最后 "regular" 一步应该是朝向noiter*dt，不考虑时间余数 T-noiter*dt.

请注意 range(N) 生成数字 0,1,…,N-1。相等， res=zeros(N) 生成一个包含 N 个条目的数组，从 res[0] 到 res[N-1].
切换不应该依赖于离散化，即步长长度。为此，一个更准确的穿越时间开关条件应通过插值确定（线性或 reverse quadratic)，然后修改后的系统或新系统重新启动新的初始条件。要保留所需的网格，请使用短第一步.

def solve_my_ODE(res, fdot, x0, T, dt):
    noiter = int(T / dt)
    dt = T/noiter          #adapt the timestep 
    res = zeros(noiter+1)
    res[0] = x0
    for i in range(noiter):
        res[i + 1] = res[i] + dt * fdot(res[i], i * dt)
        if res[i + 1] >= 2:
            h = (2-res[i])/(res[i+1]-res[i]) # precautions against zero division ?
            res[i + 1] = 0 + (1-h)*dt * fdot(0, (i+h)*dt)
    return res

看来最终精度要优于1e-4。这里 dt=1e-5 计算使用 100 000 步以及同样多的功能评估。

将经典的 Runge-Kutta 方法与 h=0.05 结合使用将导致误差略大于 1e-5 (dt**4=6.25e-6)，即，具有与欧拉方法误差相当的大小。但是，现在这只需要 T/dt=20 个步骤，总共 80 功能评价。请注意，切换时间也需要准确 O(dt**4) 不污染全局错误顺序。

因此，如果速度是 objective，那么调查是有利可图的高阶方法。

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

Numba's jit fails to compile function that has another function as input

python

jit

ode

numba