Updated 2012-05-08: There was a pretty serious mistake in my original implementation: it didn't work for more than one dependent variable.
This is a Python implementation of RK4. It can be used to solve systems of first order ordinary differential equations (and, by extension, higher order ordinary differential equations).
For an implementation specialized for second order ordinary differential equations, see my older post.
Here is the main function:
def rk4(x, h, y, f):
k1 = h * f(x, y)
k2 = h * f(x + 0.5*h, y + 0.5*k1)
k3 = h * f(x + 0.5*h, y + 0.5*k2)
k4 = h * f(x + h, y + k3)
return x + h, y + (k1 + 2*(k2 + k3) + k4)/6.0
"That's all?" you may ask. Yes, it is.
Here is an example, using numpy for
numerical arrays. Note that numpy overloads the + and * operations
on arrays so that we don't have to iterate over each element by hand.
import numpy
def damped_spring(t, state):
pos, vel = state
stiffness = 1
damping = 0.05
return numpy.array([vel, -stiffness*pos - damping*vel])
t = 0
dt = 1.0/40
state = numpy.array([5, 0])
print('%10f %10f' % (t, state[0]))
while t < 100:
t, state = rk4(t, dt, state, damped_spring)
print('%10f %10f' % (t, state[0]))
If you run this program, it will output a line for each timestep indicating the time that has passed and the current position of the spring. If you plot it, you'll get something like this:
sh$ python example.py > output.dat
sh$ gnuplot
gnuplot> plot 'output.dat' using 1:2 with lines
