python - 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 导致减少 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

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

就我得到的时间而言(以秒为单位，20 次迭代):

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

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

你的第三个问题:“这个脚本看起来像是一种模拟具有离散跳跃的 ODE 的合理方法吗？从数学上讲，这等同于求解具有 delta 函数的 ODE。”我真的不知道。对不起!