Coding 4th-order Runge-Kutta method for multiple harmonic oscillators in a 2D array. Stuck in a single oscillator

Question

I solved numerically, in Python, the general (nonlinear) pengulum using 4th-order Runge-Kutta method previously. And here is a question related to this, Attempt to solve the nonlinear pendulum 2nd order differential equation using 4th order Runge-Kutta method, not getting expected result, it contains details about the code and processes.

I have also tried to solve a 'spin' system, made of three simultaneous differential equations, using the same method. And it was also successful.

But, both of these questions dealt with a single particle.

Now what if I have is a number of particles. So I need to solve the same differential equations for each of the particles (no interactions). Maybe I can try it by writing the same steps for all the particles, but that's cumbursome. Also, in that case I can't change the particle number as may be required, I have to rewrite the whole thing.

I am not that familier with NumPy, also I can't wait to learn it in detail first.

Let's say for the single particle/oscillator, we have the following functions:

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

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

Then there is the initial values:

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

And the RK4 loop:

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)

Now what to do if, let's say, we have two particles/oscillator in an 1D array, or, say, 2x2=4 oscillators in a 2D array?
How can we achive our task by using elementary numpy techniques?

I thought of something as follows:

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

But several errors are being shown, and I know I am missing a lot of stuff.
That's why I can't proceed to the loop part at all.

What things need to be added?

Answer 1

Rather than dealing with multi-dimensional arrays, you can use a single 1d numpy array of size 2N (where N is the number of oscillators).

Here you store POSITIONS in y[0], y[2], y[4], ... and velocities in y[1], y[3], y[5], ...

This has the advantage that you can keep the same fully-vectorised Runge-Kutta routine.

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] / m         # "true" harmonic oscillator; use np.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()