Runge-Kutta 4 For Solving Systems Of ODEs Python

I wrote code for Runge-Kutta 4 for solving system of ODEs. It works fine for 1-D ODE but when I try to solve x'' + kx = 0 I have a problem trying to define a vectorial function: L

Solution 1:

You are on the right path, but when applying time-integration methods such as RK to vector valued ODEs, one essentially does the exact same thing as in the scalar case, just with vectors.

Thus, you skip the for j in range(len(X_0)) loop and associated indexation and you make sure that you pass initial values as vectors (numpy arrays).

Also cleaned up the indexation for t a little and stored the solution in a list.

import numpy as np

def ode_RK4(f, X_0, dt, T):    
    N_t = int(round(T/dt))
    # Initial conditions
    usol = [X_0]
    u = np.copy(X_0)
    
    tt = np.linspace(0, N_t*dt, N_t + 1)
    # RK4
    for t in tt[:-1]:
        u1 = f(u + 0.5*dt* f(u, t), t + 0.5*dt)
        u2 = f(u + 0.5*dt*u1, t + 0.5*dt)
        u3 = f(u + dt*u2, t + dt)
        u = u + (1/6)*dt*( f(u, t) + 2*u1 + 2*u2 + u3)
        usol.append(u)
    return usol, tt

def demo_exp():
    import matplotlib.pyplot as plt
    
    def f(u,t):
        return np.asarray([u])

    u, t = ode_RK4(f, np.array([1]) , 0.1, 1.5)
    
    plt.plot(t, u, "b*", t, np.exp(t), "r-")
    plt.show()
    
def demo_osci():
    import matplotlib.pyplot as plt
    
    def f(u,t, omega=2):
        u, v = u 
        return np.asarray([v, -omega**2*u])
    
    u, t = ode_RK4(f, np.array([2,0]), 0.1, 2)
    
    u1 = [a[0] for a in u]
    
    for i in [1]:
        plt.plot(t, u1, "b*")
    plt.show()

Solution 2:

The model is this: enter image description here

From the Langtangen’s book Programming for Computations - Python.