Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Solve second order differential equation using the Euler and the Runge-Kutta methods
#!/usr/bin/env python
"""
Find the solution for the second order differential equation
u'' = -u
with u(0) = 10 and u'(0) = -5
using the Euler and the Runge-Kutta methods.
This works by splitting the problem into 2 first order differential equations
u' = v
v' = f(t,u)
with u(0) = 10 and v(0) = -5
"""
from math import cos, sin
def f(t, u):
return -u
def exact(u0, du0, t):
# analytical solution
return u0 * cos(t) + du0 * sin(t)
def iterate(func, u, v, tmax, n):
dt = tmax/(n-1)
t = 0.0
for i in range(n):
u,v = func(u,v,t,dt)
t += dt
return u
def euler_iter(u, v, t, dt):
v_new = v + dt * f(t, u)
u_new = u + dt * v
return u_new, v_new
def rk_iter(u, v, t, dt):
k1 = f(t,u)
k2 = f(t+dt*0.5,u+k1*0.5*dt)
k3 = f(t+dt*0.5,u+k2*0.5*dt)
k4 = f(t+dt,u+k3*dt)
v += dt * (k1+2*k2+2*k3+k4)/6
# v doesn't explicitly depend on other variables
k1 = k2 = k3 = k4 = v
u += dt * (k1+2*k2+2*k3+k4)/6
return u,v
euler = lambda u, v, tmax, n: iterate(euler_iter, u, v, tmax, n)
runge_kutta = lambda u, v, tmax, n: iterate(rk_iter, u, v, tmax, n)
def plot_result(u, v, tmax, n):
dt = tmax/(n-1)
t = 0.0
allt = []
error_euler = []
error_rk = []
r_exact = []
r_euler = []
r_rk = []
u0 = u_euler = u_rk = u
v0 = v_euler = v_rk = v
for i in range(n):
u = exact(u0, v0, t)
u_euler, v_euler = euler_iter(u_euler, v_euler, t, dt)
u_rk, v_rk = rk_iter(u_rk, v_rk, t, dt)
allt.append(t)
error_euler.append(abs(u_euler-u))
error_rk.append(abs(u_rk-u))
r_exact.append(u)
r_euler.append(u_euler)
r_rk.append(u_rk)
t += dt
_plot("error.png", "Error", "time t", "error e", allt, error_euler, error_rk)
#_plot("result.png", "Result", "time t", "u(t)", allt, r_euler, r_rk, r_exact)
def _plot(out, title, xlabel, ylabel, allt, euler, rk, exact=None):
import matplotlib.pyplot as plt
plt.title(title)
plt.ylabel(ylabel)
plt.xlabel(xlabel)
plt.plot(allt, euler, 'b-', label="Euler")
plt.plot(allt, rk, 'r--', label="Runge-Kutta")
if exact:
plt.plot(allt, exact, 'g.', label='Exact')
plt.legend(loc=4)
plt.grid(True)
plt.savefig(out, dpi=None, facecolor='w', edgecolor='w',
orientation='portrait', papertype=None, format=None,
transparent=False)
u0 = 10
du0 = v0 = -5
tmax = 10.0
n = 2000
print "t=", tmax
print "euler =", euler(u0, v0, tmax, n)
print "runge_kutta=", runge_kutta(u0, v0, tmax, n)
print "exact=", exact(u0, v0, tmax)
plot_result(u0, v0, tmax*2, n*2)
@houkensjtu

This comment has been minimized.

Copy link

@houkensjtu houkensjtu commented May 24, 2016

Thank you for this inspiring script.
Your elegant usage of the function "iterate" is sweet and I really didn't see that kind of functional programming in numerical codes before.
Is this kind of programming trick quite common in Python's numerical calculation??

@Nazek42

This comment has been minimized.

Copy link

@Nazek42 Nazek42 commented Feb 11, 2017

Suggestion: replace lines 59 and 60 with:

euler = functools.partial(iterate, euler_iter)
runge_kutta = functools.partial(iterate, rk_iter)

You would also need to import functools.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment