Improved RK4 Implementation :: 21 Apr 2009

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

Output graph