使用 Numba 求解 ODE

Question

我正在尝试使用 Numba 加快我的 ODE 求解器的速度，但以下代码会引发键入错误：

    import numpy as np
    import matplotlib.pyplot as plt
    from numba import njit

    @njit
    def pend(t, y, b, c):
        theta, omega = y
        dydt = np.array([omega, -b*omega - c*np.sin(theta)])
        return dydt

    @njit
    def rungeStep(f, t, y0, tau, params):
        k1 = tau * f(t, y0, *params)
        k2 = tau * f(t, y0 + k1 / 2, *params)
        k3 = tau * f(t, y0 + k2 / 2, *params)
        k4 = tau * f(t, y0 + k3, *params)
        return (k1 + 2 * k2 + 2 * k3 + k4) / 6
    @njit
    def integrate(f, t0, y0, tEnd, h, params):
        ys = y0.copy()
        t = np.array(t0)
        while t0 <= tEnd:
            y0 += rungeStep(f, t0, y0[0], h, params)
            t0 += h
            ys = np.concatenate((ys, y0), axis=0)
            t = np.append(t, t0)
        return t, ys.T

    args = (0.25, 5)
    y0 = np.array([[np.pi - 0.1, 0.0]])
    t, y = integrate(pend, 0, y0, 10, 1, args)

这导致：

    TypingError: Failed in nopython mode pipeline (step: nopython frontend)
    Cannot unify array(int64, 0d, C) and array(int64, 1d, C) for 't.2', defined at <ipython-input-56-38d2ea70b889> (6)
    
    File "<ipython-input-56-38d2ea70b889>", line 6:
    def inagrate(f, t0, y0, tEnd, h, params):
        <source elided>
        while t0 <= tEnd:
            y0 += rungeStep(f, t0, y0[0], h, params)
            ^
    
    During: typing of assignment at <ipython-input-56-38d2ea70b889> (6)
    
    File "<ipython-input-56-38d2ea70b889>", line 6:
    def inagrate(f, t0, y0, tEnd, h, params):
        <source elided>
        while t0 <= tEnd:
            y0 += rungeStep(f, t0, y0[0], h, params)
            ^

没有 njit-decorator 它工作正常。有人可以帮我吗？

Answer 1

如果没有 JIT，应该会出现同样的错误，但也许广播规则足够灵活。你的 y0 是一个二维数组，你传递给 rungeStep （Heun-Kutta 方法）的是一个一维数组，它的 return 值也是。所以你不能立即使用那个结果来更新 y0，你必须更新 y0[0]。但是接下来的问题是为什么首先要引入这种复杂性。

Numpy 数组被优化为固定大小的对象，追加到它们需要在每个实例中分配内存和复制，这比追加到列表慢，在列表中重新分配是通过（增长的）预留完成的。 Numba-JIT 在将 numpy 数组列表转换为 2d numpy 数组时似乎有问题。有效的是手动将数组降级为简单列表。经过一些反复试验，我得到了以下代码运行（仅重复有更改的部分）：

@njit
def rungeStep(f, t, y0, tau, params):
    k1 = tau * f(t, y0, *params)
    k2 = tau * f(t + tau/2, y0 + k1 / 2, *params)
    k3 = tau * f(t + tau/2, y0 + k2 / 2, *params)
    k4 = tau * f(t + tau, y0 + k3, *params)
    return (k1 + 2 * k2 + 2 * k3 + k4) / 6
@njit
def integrate(f, t0, y0, tEnd, h, params):
    ys = [list(y0)]
    t = [t0]
    while t0 <= tEnd:
        y0 += rungeStep(f, t0, y0, h, params)
        t0 += h
        ys.append(list(y0))
        t.append(t0)
    return np.array(t), np.array(ys).T

args = (0.25, 5.0)
y0 = np.array([np.pi - 0.1, 0.0])
t, y = integrate(pend, 0, y0, 10, 1e-1, args)

请注意，您的 RK4 实现有一些遗漏，这里并不重要，但时间相关的 ODE 示例的错误来源。

绘制结果给出合理的图形

使用 Numba 求解 ODE

Solve ODEs with Numba

python

numpy

numeric

ode

numba