Sys.command "ocaml -version";;

#thread ;;
#require "jupyter.notebook" ;;
#require "jupyter.archimedes" ;;

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))
    |]
;;

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
;;

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 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;;

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;;

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
;;

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;;

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
;;

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;;

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 ();;