RK4 in Scala

August 27, 2009

I’ve been looking at Scala recently and I’m impressed. It’s got a nice mix of imperative and functional styles that means that it can be really useful and easy to code in.

As an example I thought I’d knock up an implementation of RK4 in it.

The work-horse of RK4 is implemented as follows:

def rk4Iter(f: (Double, Double) => Double, 
            t : Double, 
            y : Double, 
            h : Double) = {
  val k1 = f(t, y)
  val k2 = f(t + h / 2, y + h * k1 / 2)
  val k3 = f(t + h / 2, y + h * k2 / 2)
  val k4 = f(t + h, y + h * k3)
  y + h * (k1 + k2 * 2 + k3 * 2 + k4) / 6
}

One thing I think we can all agree on is that having functions as first-class objects makes this incredibly succinct and easy to follow. I’m just calling the function f a whole bunch of times. Coming from a Java background this is pleasantly less noisy (compare with f.calculate(t, y) etc).

Another subtlety is the lack of specifying a return type. Scala is smart enough to infer the return type đŸ™‚

This method is intended to be called from a loop that performs n steps of size h:

def rk4(f: (Double, Double) => Double, 
            t0 : Double, 
            y0 : Double, 
            n : Int, 
            h : Double) 
    : Double = {
  n match {
    case 0 => y0
    case _ => 
      val y1 = rk4Iter(f, t0, y0, h)
      rk4(f, t0 + h, y1, n - 1, h)
  }
}

For those coming from a Java background you’ll notice that match looks a lot like switch. You can think of it as switch on steroids (I’ll leave off explaining more as it’s a blog post in its own right!). It is sufficient to say that case _ is similar to default.

In functional programming the preferred way to do things is using recursion rather than iteration (when written properly the compiler can identify tail-recursive functions and make them just as fast as loops). We do this because it makes them easier to mathematically reason about them. Suppose the function works for:

  • the 0th case
  • the nth case if it works for the n -1th case

then the function works for all positive integers.

A recursive style also reduces the scope for off-by-one errors.

And now for an example that ties it together:

rk4(
  (t, y) => 4 * Math.exp(0.8 * t) - 0.5 * y, 
  0, 2, 10, 0.1)

I think anyone coming from a Java background will agree that passing around functions like this is lovely: goodbye interfaces and lots of curly braces, hello Functional Programming.