Runge-Kutta methods for ODE integration in OCaml¶
I want to implement and illustrate the Runge-Kutta method (actually, different variants), in the OCaml programming language.
The Runge-Kutta methods are a family of numerical iterative algorithms to approximate solutions of Ordinary Differential Equations. I will simply implement them, for the mathematical descriptions, I let the interested reader refer to the Wikipedia page, or any good book or course on numerical integration of ODE.
- I will start with the order 1 method, then the order 2 and the most famous order 4.
- They will be compared on different ODE.
Preliminary¶
Sys.command "ocaml -version";;
#thread ;;
#require "jupyter.notebook" ;;
#require "jupyter.archimedes" ;;
I don't want to try to find a reference implementation of an Ordinary Differential Equations solver in OCaml. I'm sure there is, just don't care.
I will use as a first example the one included in the scipy (Python) documentation for this odeint function.
$$\theta''(t) + b \theta'(t) + c \sin(\theta(t)) = 0.$$
If $\omega(t) := \theta'(t)$, this gives $$ \begin{cases} \theta'(t) = \omega(t) \\ \omega'(t) = -b \omega(t) - c \sin(\theta(t)) \end{cases} $$
Vectorially, if $y(t) = [\theta(t), \omega(t)]$, then the equation is $y' = f(t, y)$ where $f(t, y) = [y_2(t), -b y_2(t) - c \sin(y_1(t))]$.
We assume the values of $b$ and $c$ to be known, and the starting point to be also fixed:
let b = 0.25 ;;
let c = 5.0 ;;
let y0 = [| 3.14156 -. 0.1; 0.0 |];;
let f_pend (y : float array) (_ : float) : float array =
[|
y.(1);
(-.b) *. y.(1) +. (-.c) *. sin(y.(0))
|]
;;
Runge-Kutta method of order 1, or the Euler method¶
The approximation is computed using this update: $$y_{n+1} = y_n + (t_{n+1} - t_n) f(y_n, t_n).$$
The math behind this formula are the following: if $g$ is a solution to the ODE, and so far the approximation is correct, $y_n \simeq g(t_n)$, then a small step $h = t_{n+1} - t_n$ satisfy $g(t_n + h) \simeq g(t_n) + h g'(t_n) \simeq y_n + h f(g(t_n), t_n) + \simeq y_n + h f(y_n, t_n)$.
let rungekutta1 (f : float array -> float -> float array) (y0 : float array) (t : float array) =
let d = Array.length y0 in
let n = Array.length t in
let y = Array.make_matrix n d 0. in
for j = 0 to d-1 do
y.(0).(j) <- y0.(j);
done;
for i = 0 to n-2 do
let h = t.(i+1) -. t.(i) in
let fyt = f y.(i) t.(i) in
for j = 0 to d-1 do
y.(i+1).(j) <- y.(i).(j) +. h *. fyt.(j);
done;
done;
y
;;
We have to redefine ourselfves most of the useful functions:
let linspace (a : float) (b : float) (n : int) =
let t = Array.make n a in
let h = (b -. a) /. (float_of_int n) in
for i = 0 to n-1 do
t.(i) <- a +. (float_of_int i) *. h;
done;
t
;;
let t = linspace 0. 10. 101 ;;
let sol = rungekutta1 f_pend y0 t ;;
let column sol i =
Array.map (fun x -> x.(i)) sol
;;
Let see a first plot!
let vp = A.init ~w:800. ~h:360. ["jupyter"] in
A.Axes.box vp ;
A.Viewport.xlabel vp "Time t";
A.Viewport.title vp "Solution to the pendulum ODE with Runge-Kutta 1";
A.set_color vp A.Color.red ;
A.Viewport.text vp 7. 3. "o theta(t)" ;
A.Array.xy ~style:(`Linesmarkers "o") vp t (column sol 0);
A.set_color vp A.Color.blue ;
A.Viewport.text vp 7. 2.5 "+ omega(t)" ;
A.Array.xy ~style:(`Linesmarkers "+") vp t (column sol 1);
A.close vp;;
With the same number of points, the Euler method (i.e. the Runge-Kutta method of order 1) is less precise than the reference solve method. With more points, it can give a satisfactory approximation of the solution:
let t = linspace 0. 10. 1001 ;;
let sol = rungekutta1 f_pend y0 t ;;
let vp = A.init ~w:800. ~h:360. ["jupyter"] in
A.Axes.box vp ;
A.Viewport.xlabel vp "Time t";
A.Viewport.title vp "Solution to the pendulum ODE with Runge-Kutta 1";
A.set_color vp A.Color.red ;
A.Viewport.text vp 8.5 3. "o theta(t)" ;
A.Array.xy ~style:(`Linesmarkers "o") vp t (column sol 0);
A.set_color vp A.Color.blue ;
A.Viewport.text vp 8.5 2.5 "+ omega(t)" ;
A.Array.xy vp ~style:(`Linesmarkers "+") t (column sol 1);
A.close vp;;
let t = linspace 0. 10. 10001 ;;
let sol = rungekutta1 f_pend y0 t ;;
let vp = A.init ~w:800. ~h:360. ["jupyter"] in
A.Axes.box vp ;
A.Viewport.xlabel vp "Time t";
A.Viewport.title vp "Solution to the pendulum ODE with Runge-Kutta 1";
A.set_color vp A.Color.red ;
A.Viewport.text vp 8.5 3. "o theta(t)" ;
A.Array.xy ~style:(`Linesmarkers "o") vp t (column sol 0);
A.set_color vp A.Color.blue ;
A.Viewport.text vp 8.5 2.5 "+ omega(t)" ;
A.Array.xy vp ~style:(`Linesmarkers "+") t (column sol 1);
A.close vp;;
Runge-Kutta method of order 2¶
The order 2 Runge-Method uses this update: $$ y_{n+1} = y_n + h f(t + \frac{h}{2}, y_n + \frac{h}{2} f(t, y_n)),$$ if $h = t_{n+1} - t_n$.
Again, we need some basic functions as OCaml Array are quite poor.
With Julia arrays or NumPy Python arrays, we can write h * f(y, t) to multiply each entry of f(y, t) by a number h.
let aplus a k =
Array.map ( ( +. ) k) a
;;
let ( ++. ) = aplus ;;
let aaplus a b =
Array.map2 ( +. ) a b
;;
let ( +++. ) = aaplus ;;
let amul a k =
Array.map ( ( *. ) k ) a
;;
let ( **. ) = amul ;;
let rungekutta2 (f : float array -> float -> float array) (y0 : float array) (t : float array) =
let d = Array.length y0 in
let n = Array.length t in
let y = Array.make_matrix n d 0. in
for j = 0 to d-1 do
y.(0).(j) <- y0.(j);
done;
for i = 0 to n-2 do
let h = t.(i+1) -. t.(i) in
let fyt = f y.(i) t.(i) in
let fyt2 = f (y.(i) +++. (fyt **. (h /. 2.))) (t.(i) +. (h /. 2.)) in
for j = 0 to d-1 do
y.(i+1).(j) <- y.(i).(j) +. h *. fyt2.(j);
done;
done;
y
;;
For our simple ODE example, this method is already quite efficient.
let t = linspace 0. 10. 101 ;;
let sol = rungekutta2 f_pend y0 t ;;
let vp = A.init ~w:800. ~h:360. ["jupyter"] in
A.Axes.box vp ;
A.Viewport.xlabel vp "Time t";
A.Viewport.title vp "Solution to the pendulum ODE with Runge-Kutta 2";
A.set_color vp A.Color.red ;
A.Viewport.text vp 8.5 3. "o theta(t)" ;
A.Array.xy ~style:(`Linesmarkers "o") vp t (column sol 0);
A.set_color vp A.Color.blue ;
A.Viewport.text vp 8.5 2.5 "+ omega(t)" ;
A.Array.xy vp ~style:(`Linesmarkers "+") t (column sol 1);
A.close vp;;
Runge-Kutta method of order 4, "RK4"¶
The order 4 Runge-Method uses this update: $$ y_{n+1} = y_n + \frac{h}{6} (k_1 + 2 k_2 + 2 k_3 + k_4),$$ if $h = t_{n+1} - t_n$, and $$\begin{cases} k_1 &= f(y_n, t_n), \\ k_2 &= f(y_n + \frac{h}{2} k_1, t_n + \frac{h}{2}), \\ k_3 &= f(y_n + \frac{h}{2} k_2, t_n + \frac{h}{2}), \\ k_4 &= f(y_n + h k_3, t_n + h). \end{cases}$$
let rungekutta4 (f : float array -> float -> float array) (y0 : float array) (t : float array) =
let d = Array.length y0 in
let n = Array.length t in
let y = Array.make_matrix n d 0. in
for j = 0 to d-1 do
y.(0).(j) <- y0.(j);
done;
for i = 0 to n-2 do
let h = t.(i+1) -. t.(i) in
let k1 = f y.(i) t.(i) in
let k2 = f (y.(i) +++. (k1 **. (h /. 2.))) (t.(i) +. (h /. 2.)) in
let k3 = f (y.(i) +++. (k2 **. (h /. 2.))) (t.(i) +. (h /. 2.)) in
let k4 = f (y.(i) +++. (k3 **. h)) (t.(i) +. h) in
for j = 0 to d-1 do
y.(i+1).(j) <- y.(i).(j) +. (h/.6.) *. (k1.(j) +. 2.*.k2.(j) +. 2.*.k3.(j) +. k4.(j));
done;
done;
y
;;
For our simple ODE example, this method is even more efficient.
let t = linspace 0. 10. 31 ;;
let sol = rungekutta4 f_pend y0 t ;;
let vp = A.init ~w:800. ~h:360. ["jupyter"] in
A.Axes.box vp ;
A.Viewport.xlabel vp "Time t";
A.Viewport.title vp "Solution to the pendulum ODE with Runge-Kutta 4 (31 points)";
A.set_color vp A.Color.red ;
A.Viewport.text vp 8.5 3. "o theta(t)" ;
A.Array.xy ~style:(`Linesmarkers "o") vp t (column sol 0);
A.set_color vp A.Color.blue ;
A.Viewport.text vp 8.5 2.5 "+ omega(t)" ;
A.Array.xy vp ~style:(`Linesmarkers "+") t (column sol 1);
A.close vp;;
let t = linspace 0. 10. 101 ;;
let sol = rungekutta4 f_pend y0 t ;;
let vp = A.init ~w:800. ~h:360. ["jupyter"] in
A.Axes.box vp ;
A.Viewport.xlabel vp "Time t";
A.Viewport.title vp "Solution to the pendulum ODE with Runge-Kutta 4 (101 points)";
A.set_color vp A.Color.red ;
A.Viewport.text vp 8.5 3. "o theta(t)" ;
A.Array.xy ~style:(`Linesmarkers "o") vp t (column sol 0);
A.set_color vp A.Color.blue ;
A.Viewport.text vp 8.5 2.5 "+ omega(t)" ;
A.Array.xy vp ~style:(`Linesmarkers "+") t (column sol 1);
A.close vp;;
Comparisons¶
let methods = [|rungekutta1; rungekutta2; rungekutta4|] ;;
let names = [|"Runge-Kutta 1"; "Runge-Kutta 2"; "Runge-Kutta 4"|] ;;
let markers = [|"o"; "s"; ">"|] ;;
let colors = [|A.Color.red; A.Color.blue; A.Color.green|] ;;
let test_1 ?(n=101) () =
let t = linspace 0. 10. n in
let vp = A.init ~w:800. ~h:360. ["jupyter"] in
A.Axes.box vp ;
A.Viewport.xlabel vp "Time t";
A.Viewport.title vp (Format.sprintf "Solution to the pendulum ODE with different methods (%i points)" n);
for i = 0 to (Array.length methods) - 1 do
let sol = methods.(i) f_pend y0 t in
A.set_color vp colors.(i);
A.Viewport.text vp 6.5 (2. -. 2.*.(float_of_int i)) (Format.sprintf "%s %s" markers.(i) names.(i)) ;
A.Array.xy ~style:(`Linesmarkers markers.(i)) vp t (column sol 0);
done;
A.close vp;
;;
test_1 ~n:10 ();;
test_1 ~n:30 ();;
test_1 ~n:100 ();;
Conclusion¶
That's it for today, folks! See my other notebooks, available on GitHub.