Runge-Kutta with Haskell

Equation solving with class (and elegance, too!) 

This example in the R5RS document impressed me the first time I came across it, here's an implementation of the simple Runge-Kutta algorithm in Haskell:

instance (Num a) => Num [a] where
(+) = zipWith (+)
negate = fmap negate
fromInteger = undefined
(*) = undefined
abs = undefined
signum = undefined

integrate_system :: (Functor f, Floating a, Num (f a)) =>
(f a -> f a) -> f a -> a -> [f a]
integrate_system system_derivative initial_state h =
let next = runge_kutta4 system_derivative h
runge_kutta4 f h y =
let
(.*) n = fmap (*n)
shf = ((1/2).*)
k0 = (f y)
k1 = h .* (y + shf k0)
k2 = h .* (y + shf k1)
k3 = h .* (y + k2)
in
y + (h/6) .* (k0 + 2 .* k1 + 2 .* k2 + k3)
in
iterate next initial_state

We can then write code that solves, say, the simple equation y' = -k*y:


expEq k = integrate_system (\ (y:_) = -k*y) 

For some fun, use this to solve (with no claims whatsoever made regarding numerical stability) the Rossler and Lorenz systems:

rossler a b c (x:y:z:_) = [-(y+z), x+a*y, b + x*z - c*z]
lorenz s r b (x:y:z:_) = [s * (y - z), r * x - y - x * z, x * y - b * z]

There are, I suppose, two points worth noting here: firstly, the haskell code is extremely clear, and a simple translation of the algorithm. The second, is of course, that it is extremely generic - instead of lists, a Vector3 type with the relevant functions defined (fmap, negate, (+)) can directly use this code.

 

This can be used to create a lot more pretty pictures: create a pair of coupled oscillators, solve the wave equation in one dimension, and so on.