使用 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 方法）是一个一维数组，它的返回值也是如此。因此，您不能立即使用该结果更新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 示例来说，这是一个错误源。

绘制结果给出合理的图形