Numba的jit无法编译另一个函数作为inpu的函数

Question

我试图用数值方法求解一个允许离散跳跃的颂歌。我使用的是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函数，btw已经用numba编译过了）。在

因为我对麻巴这么陌生，我有几个问题要问

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

Numba版本是0.17

Answer 1

最后一点：

• 在当前表单中，它甚至不是 端庄的颂歌。它过早地停止了一个步骤，即最后一个“常规”步骤 应该朝向noiter*dt，并且不考虑时间 余数T-noiter*dt。在

  注意，range(N)生成的是0,1,…,N-1。同样地， res=zeros(N)生成一个包含N项的数组，从res[0]到 res[N-1]。

• 转换不应依赖于离散化，即步骤 长度。为了达到这个目的，一个更精确的穿越时间 应通过插值（线性或 反向二次方），然后用 新的初始条件。要保留所需的栅格，请使用short 第一步。

---

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.05will 导致错误略大于1e-5（dt**4=6.25e-6）， i、 与欧拉法的误差大小相当。 但是，现在这只需要T/dt=20个步骤，总共80 功能评估。注意，切换时间也需要 以O(dt**4)顺序精确，以防污染 全局错误顺序。在

  因此，如果以速度为目标，研究 高阶方法。

Answer 2

您认为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传递给它。只有循环是加速的，但是因为这是慢比特，所以这是唯一重要的比特。在

^{pr2}$

我更喜欢这个版本，因为它为您分配了返回值，所以使用起来更容易一些。不过，这有点偏离了你真正的问题。在

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

• 6.90394687653（仅适用于numba的fdot）
• 0.0584900379181（用于版本1）
• 0.0640540122986（对于版本2-即，它稍微慢一点，但使用起来更容易一些）

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

你的第三个问题：“这个脚本看起来像是一个合理的方法来模拟带有离散跳跃的ODE吗？从数学上讲，这相当于用delta函数求解ODE。对不起的！在