Improved RK4 Implementation

Damped harmonic motion

This is a Python implementation of the RK4 numerical integrator that works with differential functions of all orders. That is, any function in the form F(x, y, y’, y”, …, y^{(n)}).

If you’re new to numerical integration or even RK4 integration, please read my other post first. It’s easier to understand because it’s a less generalized function.

def rk4(t0, h, s0, f):
	"""RK4 implementation.
	t = current value of the independent variable
	h = amount to increase the independent variable (step size)
	s0 = initial state as a list. ex.: [initial_position, initial_velocity]
	f = function(state, t) to integrate"""
	r = range(len(s0))
	s1 = s0 + [f(t0, s0)]
	s2 = [s0[i] + 0.5*s1[i+1]*h for i in r]
	s2 += [f(t0+0.5*h, s2)]
	s3 = [s0[i] + 0.5*s2[i+1]*h for i in r]
	s3 += [f(t0+0.5*h, s3)]
	s4 = [s0[i] + s3[i+1]*h for i in r]
	s4 += [f(t0+h, s4)]
	return t+h, [s0[i] + (s1[i+1] + 2*(s2[i+1]+s3[i+1]) + s4[i+1])*h/6.0 for i in r]

An example usage of this function is:

def damped_spring_accel(t, state):
	x = state[0]
	v = state[1]
	stiffness = 1
	damping = 0.005
	return -stiffness*x - damping*v
 
t = 0
h = 1/40.0
s = [50, 5]
 
while t<100:
	t, s = rk4(t, h, s, damped_spring_accel)
print "Final state: t=%.2f x=%.2f v=%.2f"%(t,s[0],s[1])

On a side note, I recently found out about Sage (sagemath.org and sagenb.org). Its plotting capabilities and convenient mathematical notation are especially useful. Plus, it uses Python!

This entry was posted in Mathematics, Physics, Programming. Bookmark the permalink.

Leave a Reply

Your email address will not be published. Required fields are marked *

*

You may use these HTML tags and attributes: <a href="" title=""> <abbr title=""> <acronym title=""> <b> <blockquote cite=""> <cite> <code> <del datetime=""> <em> <i> <q cite=""> <strike> <strong> <pre lang="" line="" escaped="">