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”, …,
).
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!
