rktest.py

import numpy as np
import matplotlib.pyplot as plt

# This function takes previous estimations of s(t) and generates a new RUNGE-KUTTA estimation based upon given timepoints


def NextApproximation(f, t, s, h):

    k1 = f(t, s)
    k2 = f(t + h/2, s + h*(k1/2))
    k3 = f(t + h/2, s + h*(k2/2))
    k4 = f(t + h, s + h*(k3/2))

    return s + (h/6)*(k1 + 2*k2 + 2*k3 + k4)


def rk4(f, t0, x0, tau, n):
    m = x0.size
    h = (tau - t0) / n
    # Time, the reshape makes it a nx1 matrix instead of an n vector
    time = np.linspace(t0, tau, n + 1).reshape(-1, 1)

    traj = np.zeros((n + 1, m))  # State trajectory
    traj[0, :] = x0

    for k in range(n):  # Iterate through time
        tprev = time[k]
        sprev = traj[k, :]

        snext = NextApproximation(f, tprev, sprev, h)

        traj[k + 1, :] = snext

    return np.hstack((time, traj))


def f(t, x, lam=0.25, mu=0.45):
    return np.array([
        -x[1],
        lam - mu*x[1] - np.square(x[0]) - x[1]*x[0]
    ])


t0 = 0
x0 = np.array([0.5 - 1e-6, 0])
tau = 160  # 160
n = 100_000

approx = rk4(f, t0, x0, tau, n)
plt.plot(approx[:, 0], approx[:, 1])
plt.savefig("test.png")