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\prime, y\prime\!\prime, \dots, 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!