## 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` = q and `y` = 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 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 and q' = w = y
func derivs(y: [Double], x: Double) -> [Double] {
return [ y, - g * sin(y) / L]  // i.e. [ q' = w, w' = - g sin(q) / L ]
}

// q is q and q is q' = w
var q = [Double](repeating: 0.0, count: 2)
q = - 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]), q'(\(t1)) = \(resultRK.y[resultRK.x.count-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
q = -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]), q'(\(t1)) = \(resultRK.y[resultRK.x.count-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)")```