Runge-Kutta method: 1st, 2nd and 4th Order

from __future__ import division
import matplotlib.pyplot as plt

def rungeKutta(func, tspan, steps, y0, order):

	dy = [0] * (steps + 1)
	t  = [0] * (steps + 1)

	dy[0] = y  = y0
	t[0]  = ti = tspan[0] 
	h = (tspan[1] - tspan[0]) / float(steps)
	
	if order not in [1, 2, 4]:

		print('Error! Order integer must be == 1, 2 or 4!')

		return

	# FIRST ORDER RUNGE-KUTTA

	for i in xrange(1, steps + 1): # ITERATE 
		
		if order == 1:

			k1 = func(ti, y)

			t[i]  = ti = tspan[0] + i * h
			dy[i] = y  = y + k1 * h
			
	## SECOND ORDER RUNGE-KUTTA

		if order == 2:

			k1 = func(ti, y) * h
			k2 = func(ti, y + k1 / 2) * h

			t[i]  = ti = tspan[0] + i * h
			dy[i] = y  = y + k2

	# FOURTH ORDER RUNGE-KUTTA

		if order == 4:

			k1 = func(ti, y) * h
			k2 = func(ti + h / 2, y + k1 / 2) * h
			k3 = func(ti + h / 2, y + k2 / 2) * h
			k4 = func(ti + h, y + k3) * h

			t[i]  = ti = tspan[0] + i * h
			dy[i] = y  = y + (k1 + 2 * (k2 + k3) + k4) / 6

	return(dy, t)




## USAGE EXAMPLE ##

function = lambda t, x : -2 * t ** 3 + 12 * t ** 2 - 20 * t + 8.5 # ODE function x(t)

tspan = [0, 4] # TIME SPAN
steps = 100    # NO. OF TIME STEPS

y0 = 1 # Define initial value

# Call 1st order Runge-Kutta (Euler's) function
y1, t = rungeKutta(function, tspan, steps, y0, 1) 

# Call 2nd order Runge-Kutta function
y2, t = rungeKutta(function, tspan, steps, y0, 2) 

# Call 4th order Runge-Kutta function
y4, t = rungeKutta(function, tspan, steps, y0, 4) 

plt.plot(t, y1, '--') # Plot the results
plt.plot(t, y2, '*')
plt.plot(t, y4, 's')

plt.legend(('First rder', 
	    'Second order', 
	    'Fourth order'))

plt.grid('on')
plt.show()