如何将RK4函数内的for循环替换为索引方式实现等效功能

你期望的完全消除迭代循环、仅靠全局索引批量计算所有时间步的需求无法实现。RK4是步进式递推算法，第i+1步的解完全依赖第i步的计算结果，不存在跳过前置步骤直接批量计算所有时间步的数学可能性。但你可以通过优化替换Python原生for循环，获得数十到上百倍的性能提升，等效满足你的性能需求。

这是当前Python生态下求解常微分方程数值解的最优实践，只需要给原有函数加一行装饰器，不需要改动任何业务逻辑，编译后的循环效率接近C语言：
首先安装依赖：
pip install numba
优化后代码示例：

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

# 加njit装饰器，自动编译函数，消除Python循环开销
@njit
def RK4(t, f, y0):
    y = np.zeros((len(t), len(y0)))
    y[0] = y0
    h = t[1] - t[0]
    # 此处for循环会被编译为机器码，运行速度远快于原生Python循环
    for i in range(len(t)-1):
        k1 = h * f(t[i], y[i])
        k2 = h * f(t[i] + h/2, y[i] + k1/2)
        k3 = h * f(t[i] + h/2, y[i] + k2/2)
        k4 = h * f(t[i] + h, y[i] + k3)    
        y[i+1] = y[i] + (k1 + 2*k2 + 2*k3 + k4) / 6
    return y

# SIR函数也要加njit装饰器才能被编译后的RK4调用
@njit
def func(t, y):
    b, g = 0.143, 0.071
    S = -b * y[0] * y[1]
    I = b * y[0] * y[1] - g * y[1]
    R = g * y[1]
    return np.array([S, I, R])

# 测试代码
if __name__ == "__main__":
    t = np.linspace(0, 100, 1000)
    y0 = np.array([0.99, 0.01, 0.0])
    sol = RK4(t, func, y0)
    print(sol)

如果你需要同时计算多组不同初始值/参数的SIR模型解，可以把初始值和参数按批次维度堆叠，用索引向量化同时计算所有批次的结果，仅保留时间步的循环，大幅提升批量计算的效率：

def RK4_batch(t, f, y0_batch):
    # y0_batch维度: (批次数量, 状态量数量)
    n_batch, n_state = y0_batch.shape
    y = np.zeros((len(t), n_batch, n_state))
    y[0] = y0_batch
    h = t[1] - t[0]
    for i in range(len(t)-1):
        # 所有批次的k1、k2等同时计算，用numpy广播机制向量化批次维度
        k1 = h * f(t[i], y[i])
        k2 = h * f(t[i] + h/2, y[i] + k1/2)
        k3 = h * f(t[i] + h/2, y[i] + k2/2)
        k4 = h * f(t[i] + h, y[i] + k3)
        y[i+1] = y[i] + (k1 + 2*k2 + 2*k3 + k4) / 6
    return y

该方案适用于参数扫描、不确定度分析等需要同时运行数百上千组SIR模拟的场景，比逐次循环调用单组模拟效率高很多。

两种优化方案的计算结果和原有原生循环的实现完全一致，没有精度损失。

内容的提问来源于stack exchange，提问作者Rahul Niranjan