编码二维阵列中多个谐振子的四阶龙格-库塔方法。陷入单一振荡器

Question

我之前使用 4 阶 Runge-Kutta 方法在 Python 中对一般（非线性）摆摆进行了数值求解。这里有一个与此相关的问题，尝试使用四阶龙格-库塔方法求解非线性摆二阶微分方程，没有得到预期结果，它包含有关代码和过程的详细信息。

我还尝试使用相同的方法求解由三个联立微分方程组成的“自旋”系统。而且也成功了。

但是，这两个问题都涉及单个粒子。

现在如果我有许多粒子怎么办。所以我需要为每个粒子求解相同的微分方程（无相互作用）。也许我可以尝试为所有粒子编写相同的步骤，但这很麻烦。另外，在这种情况下，我无法根据需要更改粒子数，我必须重写整个过程。

我不太熟悉NumPy，也迫不及待地想先详细学习它。

假设对于单粒子/振荡器，我们有以下功能：

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

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

然后是初始值：

k=1.0                      #parameter
t,x,y=0,8.0*pi/9.0,0       #initial values (t: second, x: radian, y: radian/second)
h=0.01                     #increment in t

还有 RK4 循环：

T,X,Y=[t],[x],[y]          #lists to store data

# Loop:
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)         # apprx. value of x1 for t+h
    y=y+(1/6)*(b1+2*b2+2*b3+b4)         # apprx. value of y1 for t+h
    t=t+h                               # current value of independent variable t
    T.append(t)
    X.append(x)
    Y.append(y)

现在，如果，假设，我们在一维数组中有两个粒子/振荡器，或者，比如说，在二维数组中有 2x2=4 个振荡器？
我们如何使用基本的 numpy 技术来完成我们的任务？

我想到了如下：

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                      #parameter
t,x[i][j],y[i][j]=0,8.0*pi/9.0,0       #initial values (t: second, x: radian, y: radian/second)
h=0.01                     #increment in t

但是显示了几个错误，我知道我错过了很多东西。
这就是为什么我根本无法继续循环部分。

需要添加哪些东西？

Answer 1

您可以使用大小为 2N 的单个 1d numpy 数组（其中 N 是振荡器的数量），而不是处理多维数组。

在这里，您将位置存储在

y[0], y[2], y[4], ...

中，将速度存储在

y[1], y[3], y[5], ...

中

这样做的优点是您可以保持相同的完全矢量化的龙格-库塔例程。

import numpy as np
import matplotlib.pyplot as plt
import math


# 4th-order explicit Runge-Kutta
def rk4( x, y, dx, f ):
    dy1 = dx * f( x           , y             )
    dy2 = dx * f( x + 0.5 * dx, y + 0.5 * dy1 )
    dy3 = dx * f( x + 0.5 * dx, y + 0.5 * dy2 )
    dy4 = dx * f( x +       dx, y +       dy3 )
    return x + dx, y + ( dy1 + 2.0 * dy2 + 2.0 * dy3 + dy4 ) / 6.0


# Equation parameters
N = 3                                  # number of systems
k = np.array ( [ 1.0, 4.0, 9.0 ] )     # stiffnesses
m = np.array ( [ 1.0, 1.0, 1.0 ] )     # masses
x0 = np.array( [ 0.5, 0.6, 0.7 ] )     # initial displacements
v0 = np.array( [ 0.0, 0.0, 0.0 ] )     # initial velocities


# Derivative function (outputs a numpy array)
def deriv( t, y ):
    f = np.zeros_like( y )
    f[0:2*N-1:2] = y[1:2*N: 2]         # derivative of positions
    f[1:2*N  :2] = -k * y[0:2*N-1:2]   # "true" harmonic oscillator; use math.sin() otherwise
    return f


# Initialise
t = 0
y = np.array( [ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 ] )
y[0:2*N-1:2] = x0                      # set positions in elements 0, 2, 4, ...
y[1:2*N  :2] = v0                      # set velocities in elements 1, 3, 5, ...
nt = 400
dt = 0.01

# Initialise for plotting
tvals=[]
xvals=[[] for i in range(N)]           # careful!!!


# Successive timesteps
tvals.append( t )
for i in range( N ): xvals[i].append( y[2*i] )
for i in range( nt + 1 ):
    t, y = rk4( t, y, dt, deriv )     # Runge-Kutta update
    tvals.append( t )
    for i in range( N ): xvals[i].append( y[2*i] )


# Plot
for i in range( N ): plt.plot( tvals, xvals[i] )
plt.show()