Solving system of coupled differential equations using Runge-Kutta in python

This python code can solve one non- coupled differential equation:

import numpy as np
import matplotlib.pyplot as plt
import numba
import time
start_time = time.clock()
@numba.jit()
# A sample differential equation "dy / dx = (x - y**2)/2"
def dydx(x, y): return ((x - y**2)/2)
# Finds value of y for a given x using step size h
# and initial value y0 at x0.
def rungeKutta(x0, y0, x, h): # Count number of iterations using step size or # step height h n = (int)((x - x0)/h) # Iterate for number of iterations y = y0 for i in range(1, n + 1): "Apply Runge Kutta Formulas to find next value of y" k1 = h * dydx(x0, y) k2 = h * dydx(x0 + 0.5 * h, y + 0.5 * k1) k3 = h * dydx(x0 + 0.5 * h, y + 0.5 * k2) k4 = h * dydx(x0 + h, y + k3) # Update next value of y y = y + (1.0 / 6.0)*(k1 + 2 * k2 + 2 * k3 + k4) # Update next value of x x0 = x0 + h return y
def dplot(start,end,steps): Y=list() for x in np.linspace(start,end,steps): Y.append(rungeKutta(x0, y, x , h)) plt.plot(np.linspace(start,end,steps),Y) print("Execution time:",time.clock() - start_time, "seconds") plt.show()
start,end = 0, 10
steps = end* 100
x0 = 0
y = 1
h = 0.002
dplot(start,end,steps)

This code can solve this differential equation:

 dydx= (x - y**2)/2

Now I have a system of coupled differential equations:

 dydt= (x - y**2)/2 dxdt= x*3 + 3y

How can I implement these two as a system of coupled differential equations in the above code? Is there any more generalized way for system of n-number of coupled differential equations?

1 Answer

With the help of others, I got to this:

import numpy as np
from math import sqrt
import matplotlib.pyplot as plt
import numba
import time
start_time = time.clock()
a=1
b=1
c=1
d=1
# Equations:
@numba.jit()
#du/dt=V(u,t)
def V(u,t): x, y, vx, vy = u return np.array([vy,vx,a*x+b*y,c*x+d*y])
def rk4(f, u0, t0, tf , n): t = np.linspace(t0, tf, n+1) u = np.array((n+1)*[u0]) h = t[1]-t[0] for i in range(n): k1 = h * f(u[i], t[i]) k2 = h * f(u[i] + 0.5 * k1, t[i] + 0.5*h) k3 = h * f(u[i] + 0.5 * k2, t[i] + 0.5*h) k4 = h * f(u[i] + k3, t[i] + h) u[i+1] = u[i] + (k1 + 2*(k2 + k3 ) + k4) / 6 return u, t
u, t = rk4(V, np.array([1., 0., 1. , 0.]) , 0. , 10. , 100000)
x,y, vx,vy = u.T
# plt.plot(t, x, t,y)
plt.semilogy(t, x, t,y)
plt.grid('on')
print("Execution time:",time.clock() - start_time, "seconds")
plt.show()