使用 NumPy 向量化 python 循环

Question

由于我的数学建模的原因，我有一个 vector_U 函数

import numpy as np
def vector_U(U_0, t, func, dt):
    res = np.empty((len(t), 4))
    res[0] = U_0
    for i in range(1, len(t)):
        res[i] = res[i-1]+runge_kutta(res[i-1], t[i-1], func, dt)
    return res

其中 func 是微分方程的函数

f([x, y, vx, vy]) = [vx, vy, ax, ay]

，runge_kutta 是

def runge_kutta(x, t, func, dt):
    k1 = dt*func(x, t)
    k2 = dt*func(x+k1/2, t+dt/2)
    k3 = dt*func(x+k2/2, t+dt/2)
    k4 = dt*func(x+k3, t+dt)
    return (k1+2*k2+2*k3+k4)/6

vector_U 函数的目的是获取起始条件 U_0 和 t=np.linspace(0, t_final, t_final/dt) 的 4-向量，并返回每个 dt 步骤的 4-向量数组。然后 x = res[0] 和 y = res[1] 列用于动画。

虽然它工作得很好，但对于一个大的 t （例如

len(t)==50000

），该函数显然会变得很慢，因为它使用了循环。由于我正在尝试（而不是 KISS，是的）使程序尽可能矢量化以使其运行得更快，NumPy 是否有任何方法来矢量化类似的东西？我指的不是具体的决定，而是矢量化循环的方式，因为我认为我不完全理解矢量化方法。

我考虑过使用 runge_kuttas 的 np.cumsum...但是对于每个 runge_kutta 我仍然需要 U 的先前值。

Answer 1

对于这样的代码，Numba 是加快速度的好方法。这是快速指南：https://numba.pydata.org/numba-doc/latest/user/5minguide.html

您基本上可以保持代码原样，添加装饰器（或包装函数）并获得相当好的加速，有时与使用 NumPy 矢量化代码获得的速度相当。