Runge-Kutta in Swift

Apple’s Scene Kit has a built in physics engine. However, we discovered it wasn’t accurate for particular cases. We decided to replace the physics engine with a Runge-Kutta method.

We first did a simple fixed time step (rk4FixedStep below) to advance the solution from one time to another time. This worked fine. However, in the case we were concerned with, having a fixed time step would introduce too large of an error.

To get around that, we introduce a variable time step Runge-Kutta based on solving both the fourth and fifth order Runge-Kutta and by looking at the error between the two orders, we adjust the time step so the accuracy is within user specified limit (rkVariableStep below).

For example, let’s do a simple pendulum of length L under uniform gravity g. The independent variable is time t and the dependent variable is the angle the pendulum makes with downward vertical, q. Applying Netwon’s second law of motion leads to an ordinary second order differential equation q”(t) = – g sin(q) / L. The prime denotes differentiation with respect to time. Runge-Kutta requires systems of first order differential equations, not second order.

We introduce a new variable w(t) = q‘(t) leading to two coupled equations: q‘(t) = w(t) and w‘(t) = – g sin(q) / L. That is x = t, y[0] = q and y[1] = w in the Swift routines that follow.

class RungeKutta {
    
    func rk4(y: [Double], dydx: [Double], x: Double, h: Double, f: ([Double], Double) -> [Double] ) -> [Double] {
        // given y[] and dydx[] at x, use 4th order Runge-Kutta to advance solution over an interval h
        // f returns derivatives at x, i.e. y'(x) = f(y, x)

        let n = y.count  // number of first order ODE
        
        let hh = 0.5 * h    // half interval
        let xh = x + hh
        var yt = [Double](repeating: 0.0, count: n) // temporary storage
        
        for i in 0..<n { yt[i] = y[i] + hh * dydx[i] }  // first step
        var dyt = f(yt, xh)                             // get dydx at half time step
        for i in 0..<n { yt[i] = y[i] + hh * dyt[i] }   // second step
        var dym = f(yt, xh)                             // get updated dydx at half time step
        for i in 0..<n {
            yt[i] = y[i] + h * dym[i]                   // third step
            dym[i] = dyt[i] + dym[i]
        }
        dyt = f(yt, x + h)                              // fourth step
        for i in 0..<n { yt[i] = y[i] + h * (dydx[i] + dyt[i] + 2.0 * dym[i]) / 6.0 }   // do final step w/ appropriate weights
        return yt   // return updated y[] at x+h
    }
    
    func rk4FixedStep(ystart: [Double], x1: Double, x2: Double, nstep: Int, f: ([Double], Double) -> [Double] ) -> ( x: [Double], y: [[Double]] ) {
        // given initial values (ystart[]) and interval to integrate (x1 to x2) use fourth order Runge-Kutta to advance solution in equal steps of (x2-x1)/nstep. y'(x) = f(y, x)
        // we return the path at each step, [x], [y(x)]
        var pathX = [Double]()
        var pathY = [[Double]]()
        
        pathX.append(x1)
        var v = ystart
        pathY.append(v)
        var x = x1
        let h = (x2 - x1) / Double(nstep)
        for _ in 0..<nstep {
            let dv = f(v, x)
            assert(n == dv.count, "number of variables != number of derivatives in f rk4FixedStep")
            v = rk4(y: v, dydx: dv, x: x, h: h, f: f)
            if x + h == x { print("stepsize underflow in rkFixedStep"); exit(1) }
            x += h
            pathX.append(x)
            pathY.append(v)
        }
        return ( pathX, pathY )
    }
    
    private func rkck(y: [Double], dydx: [Double], x: Double, h: Double, f: ([Double], Double) -> [Double] ) -> (yout: [Double], yerr: [Double]) {
        // use fifth order Cash-Karp Runge-Kutta method to advance over interval h returning yout and estimate of error in yerr

        let A2=0.2, A3=0.3, A4=0.6, A5=1.0, A6=0.875
        let B21 = 0.2, B31 = 3.0/40.0, B32 = 9.0/40.0, B41 = 0.3, B42 = -0.9, B43 = 1.2
        let B51 = -11.0/54.0,B52 = 2.5, B53 = -70.0/27.0, B54 = 35.0/27.0
        let B61 = 1631.0/55296.0, B62 = 175.0/512.0, B63 = 575.0/13824.0, B64 = 44275.0/110592.0, B65 = 253.0/4096.0
        let C1 = 37.0/378.0, C3 = 250.0/621.0, C4 = 125.0/594.0, C6 = 512.0/1771.0
        let DC1 = C1 - 2825.0/27648.0, DC3 = C3 - 18575.0/48384.0
        let DC4 = C4 - 13525.0/55296.0, DC5 = -277.0/14336.0, DC6 = C6 - 0.25
                
        let n = y.count  // number of first order ODE
        var ytemp = [Double](repeating: 0.0, count: n)
        var yerr = [Double](repeating: 0.0, count: n)
        for i in 0..<n { ytemp[i] = y[i] + h * (B21 * dydx[i]) }
        let ak2 = f(ytemp, x + A2 * h)
        for i in 0..<n { ytemp[i] = y[i] + h * (B31 * dydx[i] + B32 * ak2[i]) }
        let ak3 = f(ytemp, x + A3 * h)
        for i in 0..<n { ytemp[i] = y[i] + h * (B41 * dydx[i] + B42 * ak2[i] + B43 * ak3[i]) }
        let ak4 = f(ytemp, x + A4 * h)
        for i in 0..<n { ytemp[i] = y[i] + h * (B51 * dydx[i] + B52 * ak2[i] + B53 * ak3[i] + B54 * ak4[i]) }
        let ak5 = f(ytemp, x + A5 * h)
        for i in 0..<n { ytemp[i] = y[i] + h * (B61 * dydx[i] + B62 * ak2[i] + B63 * ak3[i] + B64 * ak4[i] + B65 * ak5[i]) }
        let ak6 = f(ytemp, x + A6 * h)
        for i in 0..<n { ytemp[i] = y[i] + h * (C1 * dydx[i] + C3 * ak3[i] + C4 * ak4[i] + C6 * ak6[i]) } // result?
        for i in 0..<n { yerr[i] = h * ( DC1 * dydx[i] + DC3 * ak3[i] + DC4 * ak4[i] + DC5 * ak5[i] + DC6 * ak6[i]) }
        return ( ytemp, yerr )
    }

    private func rkqs(y: [Double], dydx: [Double], x: Double, htry: Double, f: ([Double], Double) -> [Double], eps: Double, yscale: [Double]) -> (y: [Double], xnew: Double, hdid: Double, hnext: Double) {
        // use fifth-order Runge-Kutta step w/ monitoring of local error to adjust stepsize if necessary. eps determines the acceptable error in the step

        let SAFETY = 0.9, PGROW = -0.2, PSHRINK = -0.25, ERRCON = 1.89e-4
        var h = htry
        var errmax = 0.0
        
        let n = y.count
        var result = rkck(y: y, dydx: dydx, x: x, h: h, f: f)
        for i in 0..<n { errmax = max( errmax, abs(result.yerr[i]/yscale[i])) }
        errmax = errmax / eps

        while errmax > 1.0 { // ugh, stepsize too big, shrink and try again...
            let htemp = SAFETY * h * (pow(errmax, PSHRINK))
            h = max( abs(htemp), 0.1 * abs(h)) * ((h < 0.0) ? -1.0 : 1.0)
            if x + h == x { print("underflow on stepsize in rkqs"); exit(1) }
            result = rkck(y: y, dydx: dydx, x: x, h: h, f: f)
            errmax = 0.0
            for i in 0..<n {  errmax = max( errmax, abs(result.yerr[i]/yscale[i])) }
            errmax = errmax / eps
        }
        var hnext = 5.0 * h
        if errmax > ERRCON { hnext = SAFETY * h * (pow(errmax, PGROW)) }
        return ( result.yout, x+h, h, hnext )
    }

    func rkVariableStep(ystart: [Double], x1: Double, x2: Double, eps: Double, h1: Double, hmin: Double, f: ([Double], Double) -> [Double]) -> (x: [Double], y: [[Double]], hnext: Double, nok: Int, nbad: Int) {
  
        // given initial values (ystart[]) and interval to integrate (x1 to x2) use fourth and fifth order Runge-Kutta to advance solution in steps to within estimated accuracy of eps. start with h1 for step with min step hmin (can be zero).  y'(x) = f(y, x)    
        // we return the path at each step, [x], [y(x)]  
        let n = ystart.count
        
        let TINY = 1.0e-30
        var x = x1
        var h = abs(h1) * ((x2 - x1 < 0.0) ? -1.0 : 1.0)
        var nok = 0
        var nbad = 0
        
        var pathX = [Double]()
        var pathY = [[Double]]()
        var yscal = [Double](repeating: 0.0, count: n)
        
        pathX.append(x1)
        var y = ystart
        pathY.append(y)
        
        repeat {
            let dydx = f(y, x)
            assert(n == dydx.count, "number of variables != number of derivatives in f rkVariableStep")
            for i in 0..<n { yscal[i] = abs(y[i]) + abs(h * dydx[i])+TINY }
            if (x+h-x2)*(x+h-x1) > 0.0 { h = x2 - x }
            let result = rkqs(y: y, dydx: dydx, x: x, htry: h, f: f, eps: eps, yscale: yscal)
            if result.hdid == h { nok += 1 } else { nbad += 1 }
            x = result.xnew
            y = result.y
            if abs(result.hnext) < hmin {print("stepsize underflow in rkVariableStep"); exit(1)}
            pathX.append(x)
            pathY.append(y)
            h = result.hnext
        } while (x-x2) * (x2-x1) < 0.0
        return ( pathX, pathY, h, nok, nbad)
    }
}

To implement the pendulum, we’ll start at q = 0 and a velocity of 2 to the left, i.e. w = -2 / L (radians/s).

rk = RungeKutta()

let L = 1.0 // length of pendulum
let g = 9.8 // strength of gravity
 
// simple pendulum; q is angle from downward vertical then q = y[0] and q' = w = y[1]
func derivs(y: [Double], x: Double) -> [Double] {
        return [ y[1], - g * sin(y[0]) / L]  // i.e. [ q' = w, w' = - g sin(q) / L ]
 }

// q[0] is q and q[1] is q' = w
var q = [Double](repeating: 0.0, count: 2)
q[1] = - 2.0 / L

let t0 = 0.0  // starting time
let t1 = (1.0 / 60.0) * 10000 + t0  // 2.77 minutes of integration at 60 fps

let frames = (t1 - t0) * 60.0      // number of frames
// integrate from t0 to t1 with fixed step using fourth order Runge-Kutta
let resultRKFixed = rk.rk4FixedStep(ystart: q, x1: t0, x2: t1, nstep: frames * 20, f: derivs)    // take 20 time steps per frame
print("q(\(t1)) = \(resultRK.y[resultRK.x.count-1][0]), q'(\(t1)) = \(resultRK.y[resultRK.x.count-1][1])")    // final position and angular velocity

// just take one step with fourth order Runge-Kutta
let resultRKStep = rk.rk4(y: q, dydx: derivs(q, t), x: t1, h: (1.0 / 60.0) / 20.0, f: derivs )

// reset pendulum
q[0] = 0.0 
q[1] = -2.0 / L

// integrate from t0 to t1 with variable step using Runge-Kutta
let resultRK = rk.rkVariableStep(ystart: q, x1: t0, x2: t1, eps: 1.0e-16, h1: (1.0 / 60.0) / 10.0, hmin: 0.0, f: derivs)  // accuracy ~ double precision (brutal)
print("q(\(t1)) = \(resultRK.y[resultRK.x.count-1][0]), q'(\(t1)) = \(resultRK.y[resultRK.x.count-1][1])")    // final position and angular velocity

For a simple pendulum, there is an analytic solution for q available (with the same initial conditions)

q(t)=2 JacobiAmplitude[-t / L, g L]

q‘(t)=-2 JacobiDN[-t / L, g L] / L

in Mathematica-speak. In standard notation, where sn, cn and dn are the Jacobian elliptic functions then

q(t) = 2 atan2( sn(-t/L, sqrt(gL)), cn(-t/L, sqrt(gL)) )

q‘(t) = -2 dn(-t/L, sqrt(gL)) / L.

Using this, it’s easy to show both Runge-Kutta routines (as called above) return an error of less than 1.0e-12 over the 2.777 minutes of integration.

Routines to implement the analytic solution:

struct EllipticAndJacobian {
    
    private func rf(_ x: Double,_ y: Double,_ z: Double) -> Double { // WARNING lifted from Numerical Recipes
        let ERRTOL = 0.0025, TINY = 5.0 * Double.leastNormalMagnitude, BIG = 0.2 * Double.greatestFiniteMagnitude // for double precision
      //let ERRTOL = 0.08, TINY = 5.0 * Float.leastNormalMagnitude, BIG = 0.2 * Float.greatestFiniteMagnitude      // for single precision
        let THIRD = 1.0/3.0
        let C1 = 1.0/24.0, C2 = 0.1, C3 = 3.0/44.0, C4 = 1.0/14.0
        
        var alamb = 1.0, ave = 0.0
        var delx = 1.0, dely = 0.0, delz = 0.0
        var e2 = 0.0, e3 = 0.0
        var sqrtx = 0.0, sqrty = 0.0, sqrtz = 0.0
        
        if min(x,y,z) < 0.0 || min(x+y,x+z,y+z) < TINY || max(x,y,z) > BIG {
            print("invalid arguments in rf \(x) \(y) \(z)")
            exit(1)
        }
        
        var xt = x
        var yt = y
        var zt = z
        repeat {
            sqrtx = sqrt(xt)
            sqrty = sqrt(yt)
            sqrtz = sqrt(zt)
            alamb = sqrtx * (sqrty + sqrtz) + sqrty * sqrtz
            xt = 0.25 * (xt + alamb)
            yt = 0.25 * (yt + alamb)
            zt = 0.25 * (zt + alamb)
            ave = THIRD * (xt + yt + zt)
            delx = (ave - xt) / ave
            dely = (ave - yt) / ave
            delz = (ave - zt) / ave
        } while max(abs(delx),abs(dely),abs(delz)) > ERRTOL
        e2 = delx * dely - delz*delz
        e3 = delx * dely * delz
        return  (1.0 + (C1 * e2 - C2 - C3 * e3) * e2 + C4 * e3)/sqrt(ave)
    }
    
    func sncndn(uu: Double, emmc: Double) -> (sn: Double, cn: Double, dn: Double) { // WARNING lifted from Numerical Recipes
      //let CA = 0.0003     // single precision
        let CA = 0.00000003 // double precision
        var emc = emmc
        var u = uu
        var bo = true
        var d = 1.0, a = 0.0, b = 0.0, c = 0.0
        var sn = 0.0, cn = 0.0,  dn = 0.0
        var l = 0
        var em = Array(repeating: 0.0, count: 13)
        var en = Array(repeating: 0.0, count: 13)
        
        if emc != 0.0 {
            bo = (emc < 0.0)
            if bo {
                d = 1.0 - emc
                emc = -emc / d
                d = sqrt(d)
                u *= d
            }
            a = 1.0
            dn = 1.0
            for i in 0..<13 {
                l = i
                em[i] = a
                emc = sqrt(emc)
                en[i] = emc
                c = 0.5 * (a + emc)
                if abs(a - emc) <= CA * a {break}
                emc *= a
                a = c
            }
            u *= c
            sn = sin(u)
            cn = cos(u)
            if sn != 0.0 {
                a = cn / sn
                c *= a
                for i in (0...l).reversed() {
                    b = em[i]
                    a *= c
                    c *= dn
                    dn = (en[i] + a) / (b + a)
                    a = c / b
                }
                a = 1.0 / sqrt(c * c + 1.0)
                sn = (sn < 0.0 ? -a : a)
                cn = c * sn
            }
            if bo {
                a = dn
                dn = cn
                cn = a
                sn /= d
            }
        } else {
            cn = 1.0 / cosh(u)
            dn = cn
            sn = tanh(u)
        }
        return (sn, cn, dn)
    }
    
    
    func ellf(phi: Double, k: Double) -> Double { // WARNING lifted from Numerical Recipes
        let s = sin(phi)
        let c = cos(phi)
        return s*rf(c * c, (1.0 - s * k) * (1.0 + s * k),1.0)
    }
 
}

class PendulumEOM {
    var length = 0.0
    var gravity = 9.8
    let ej = EllipticAndJacobian()
    
    init(length: Double, gravity: Double = 9.8) {
        self.length = length
        self.gravity = gravity
    }
    
    func advance(_ t: Double, phi0: Double, EnergyByMass: Double, omega0: Double) -> (phi: Double, phiDot: Double) {
        // t is total time, phi0 is initial angle to downward vertical (radians), EnergyByMass is total mechanical energy of pendulum divided by mass, 
        // omega0 is initial angular velocity (right hand rule for sign)
        // return is phi and d(phi)/dt at end of time
        let sOmega = (omega0 < 0.0 ? -1.0 : 1.0)
        let k = 2.0 * gravity * length / EnergyByMass
        let du = sOmega * t * sqrt(EnergyByMass / 2.0) / length
        let u0 = ej.ellf(phi: abs(phi0) / 2.0, k: sqrt(k)) * (phi0 < 0.0 ? -1.0 : 1.0)
        let jacob = ej.sncndn(uu: u0 + du, emmc: 1.0 - k)

        return ( 2.0 * atan2(jacob.sn, jacob.cn), sqrt(2.0 * EnergyByMass) * jacob.dn / length * sOmega )
    }
}

let pendulumEOM = PendulumEOM(length: 1.0, gravity: 9.8)
let result = pendulumEOM.advance(10000.0/60.0, phi0: 0.0, EnergyByMass: 2.0, omega0: -2.0 / 1.0)
print("q(\(10000.0/60.0)) = \(result.phi), q'(\(10000.0/60.0)) = \(result.phiDot)")

One Response

Add a Comment

Your email address will not be published. Required fields are marked *