Runge-Kutta integration of some equations in Haskell

rk4.hs

#! /usr/bin/env nix-shell
#! nix-shell -i runghc -p 'ghc.withPackages (ps: [ ps.hmatrix ])'

import Numeric.LinearAlgebra

-- | Represents the coupled ODEs (Example (1)):
--
--   y[0] = position = p
--   y[1] = velocity = v
--
--   pdot = v
--   vdot = -k / (p^2)
--
--   Given time (unused), and y, returns ydot (the derivative of y
--   components wrt time).
ydot1 :: Double -> Vector Double -> Vector Double
ydot1 time y =
  let
    -- let's choose a k value
    k = 5.0
    -- current position and velocity
    p = y ! 0
    v = y ! 1
    -- derivatives of position and velocity wrt time
    pdot = v
    vdot = -k / (p * p)
    -- return the output vector
    ydot = vector [ pdot, vdot ]
  in
    ydot

-- | Another example ODE (Example (2)):
--
--   y[0] = position = p
--   y[1] = velocity = v
--
--   pdot = v
--   vdot = -9.81
--
--   (Chosen because it has an easy exact solution.)
ydot2 :: Double -> Vector Double -> Vector Double
ydot2 time y =
  let
    -- current position and velocity
    p = y ! 0
    v = y ! 1
    -- derivatives
    pdot = v
    vdot = -9.81
    -- return the output vector
    ydot = vector [ pdot, vdot ]
  in
    ydot

-- | Single step of RK4 integration.
--
--   From: https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods
rk4Step :: Double
        -> Double
        -> Vector Double
        -> (Double -> Vector Double -> Vector Double)
        -> Vector Double
rk4Step t h y f = 
  let
    h2 = h / 2.0
    h6 = h / 6.0

    k1 = f t y
    k2 = f (t + h2) (y + (scalar h2) * k1)
    k3 = f (t + h2) (y + (scalar h2) * k2)
    k4 = f (t + h)  (y + (scalar h)  * k3)

    y' = y + (scalar h6) * (k1 + 2*k2 + 2*k3 + k4)
  in
    y'

-- | RK4 integration until a particular time is reached.
rk4 :: Double
    -> Double
    -> Double
    -> Vector Double
    -> (Double -> Vector Double -> Vector Double)
    -> Vector Double
rk4 ti tf h yi f =
  let
    -- TODO: Remove explicit recursion.
    go t y = if t >= tf then y else go (t + h) (rk4Step t h y f)
  in
    go ti yi

main :: IO ()
main = do
  putStrLn "Initial Conditions:"
  putStrLn "  p = 10"
  putStrLn "  v =  0"
  let yi = vector [ 10.0, 0.0 ]

  putStrLn "Numerical solution of (1) when t = 10, (0.1 timestep):"
  let yf1 = rk4 0.0 10.0 0.1 yi ydot1
  putStrLn $ concat [ "  yf = ", show yf1 ]

  putStrLn "Numerical solution of (2) when t = 1, (0.005 timestep):"
  let yf2 = rk4 0.0 1.0 0.005 yi ydot2
  putStrLn $ concat [ "  yf = ", show yf2 ]

  putStrLn "Analytical solution of (2) when t=1:"
  let t = 1.0
  let p2 = (-9.81) * t * t / 2.0 + 10.0
  let v2 = (-9.81) * t
  let yf2' = vector [ p2, v2 ]
  putStrLn $ concat [ "  yf(exact) = ", show yf2' ]