如何使用RK4方法求解含三对角M矩阵的大规模耦合微分方程组

1 统一向量化方程组形式 #

你遇到的所有耦合线性微分方程组都可以统一写成向量形式：

$\frac{d\boldsymbol{x}}{dt} = M \cdot \boldsymbol{x}$
其中$\boldsymbol{x}$是长度为$n$的状态向量（对应你示例里的$x_0,x_1,x_2$），$M$是$n\times n$的系数矩阵（你的示例对应3阶矩阵，实际业务中为100阶三对角矩阵）。

你的3维示例对应的M矩阵为：
$$
M = \begin{bmatrix}
0 & -2 & 0 \
-2 & -8 & -2\sqrt{2} \
0 & -2\sqrt{2} & -14
\end{bmatrix}
$$
这种向量化形式不受维度限制，n为任意值都可以用同一套逻辑处理，无需单独定义每个分量的微分方程。

2 向量化适配RK4迭代逻辑 #

RK4的核心迭代公式天然支持向量运算，仅需将原标量计算替换为向量计算即可，无需修改核心逻辑，通用迭代步骤如下：
已知当前时刻状态向量$\boldsymbol{x}_t$，时间步长$h$：

• 计算斜率向量$k_1 = h \cdot (M \cdot \boldsymbol{x}_t)$
• 计算斜率向量$k_2 = h \cdot (M \cdot (\boldsymbol{x}_t + k_1/2))$
• 计算斜率向量$k_3 = h \cdot (M \cdot (\boldsymbol{x}_t + k_2/2))$
• 计算斜率向量$k_4 = h \cdot (M \cdot (\boldsymbol{x}_t + k_3))$
• 下一时刻状态更新为$\boldsymbol{x}_{t+h} = \boldsymbol{x}_t + \frac{k_1 + 2k_2 + 2k_3 + k_4}{6}$

3 三对角矩阵优化（针对100维场景） #

因为业务中M为三对角矩阵，仅存在主对角线、上对角线、下对角线三个非零对角，无需存储完整的n×n矩阵，也无需做O(n²)复杂度的全矩阵乘法，仅用O(n)复杂度即可完成矩阵乘向量计算，实现代码如下：

def tridiag_multiply(lower_diag, main_diag, upper_diag, x):
    """
    三对角矩阵乘向量，仅需存储三个对角数组
    lower_diag: 下对角线，长度n-1
    main_diag: 主对角线，长度n
    upper_diag: 上对角线，长度n-1
    x: 状态向量，长度n
    """
    n = len(x)
    res = np.zeros(n)
    res[0] = main_diag[0] * x[0] + upper_diag[0] * x[1]
    for i in range(1, n-1):
        res[i] = lower_diag[i-1] * x[i-1] + main_diag[i] * x[i] + upper_diag[i] * x[i+1]
    res[-1] = lower_diag[-1] * x[-2] + main_diag[-1] * x[-1]
    return res

将RK4迭代步骤中的矩阵乘法替换为上述函数，可大幅提升100维场景的计算效率。

4 可运行示例代码（Python） #

以你给出的3维场景为例，完整实现代码如下，可直接扩展到100维场景：

import numpy as np

# 3维示例系数矩阵，可替换为100维三对角矩阵或对应的三个对角数组
M = np.array([
    [0, -2, 0],
    [-2, -8, -2*np.sqrt(2)],
    [0, -2*np.sqrt(2), -14]
])
# t=0初始条件
init_x = np.array([-0.00076896, -0.01033249, -0.06899846])
# 时间配置，可根据精度需求调整步长h
h = 1e-4
total_time = 1.0
t_arr = np.arange(0, total_time, h)
x_result = np.zeros((len(t_arr), len(init_x)))
x_result[0] = init_x

# RK4迭代求解
for step in range(1, len(t_arr)):
    x_current = x_result[step-1]
    k1 = h * (M @ x_current)
    k2 = h * (M @ (x_current + k1/2))
    k3 = h * (M @ (x_current + k2/2))
    k4 = h * (M @ (x_current + k3))
    x_result[step] = x_current + (k1 + 2*k2 + 2*k3 + k4)/6

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