基于四阶Runge-Kutta法求解二维数组多谐振子的实现疑问

背景 #

我已经用Python的四阶龙格-库塔（RK4）方法求解了单个非线性摆的微分方程，还成功用该方法求解过由三个联立微分方程构成的自旋系统，但这些案例都是针对单个粒子的。现在需要求解多个无相互作用粒子的相同微分方程，重复编写单粒子代码太繁琐，也没法灵活调整粒子数量。我对NumPy不熟悉，也没时间系统学习它。

单粒子RK4实现代码 #

微分方程定义 #

import numpy as np

def f1(t, x, y): return y

def f2(t, x, y): return -k*np.sin(x)

参数与初始值 #

k = 1.0                      # 系统参数
t, x, y = 0, 8.0*np.pi/9.0, 0       # 初始值（t：时间/秒，x：摆角/弧度，y：角速度/弧度/秒）
h = 0.01                     # 时间步长

RK4循环与数据存储 #

T, X, Y = [t], [x], [y]          # 存储时间、摆角、角速度的列表

# 迭代计算
for i in range(2000):
    a1 = h * f1(t, x, y)                      
    b1 = h * f2(t, x, y)                      
    a2 = h * f1(t + 0.5*h, x + 0.5*a1, y + 0.5*b1)  
    b2 = h * f2(t + 0.5*h, x + 0.5*a1, y + 0.5*b1)  
    a3 = h * f1(t + 0.5*h, x + 0.5*a2, y + 0.5*b2)  
    b3 = h * f2(t + 0.5*h, x + 0.5*a2, y + 0.5*b2)  
    a4 = h * f1(t + h, x + a3, y + b3)              
    b4 = h * f2(t + h, x + a3, y + b3)              
    x = x + (1/6)*(a1 + 2*a2 + 2*a3 + a4)         # 更新t+h时刻的摆角
    y = y + (1/6)*(b1 + 2*b2 + 2*b3 + b4)         # 更新t+h时刻的角速度
    t = t + h                               # 更新当前时间
    T.append(t)
    X.append(x)
    Y.append(y)

多粒子实现的问题 #

现在需要处理**1D数组（2个粒子）或2D数组（4个粒子）**的情况，我尝试写了下面的代码，但出现很多错误，不知道循环部分该怎么补充：

# 错误的尝试代码
def f1(t, x[i][j], y[i][j]): return y[i][j]

def f2(t, x[i][j], y[i][j]): return -k*sin(x[i][j])

k = 1.0                      # 参数
t, x[i][j], y[i][j] = 0, 8.0*pi/9.0, 0       # 初始值
h = 0.01                     # t的步长

请问怎么用基础NumPy技术实现多粒子的RK4求解？

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