added: separate function for each RungeKutta implementation

Author: franckgaga

src/model/solver.jl

@@ -39,54 +39,70 @@ RungeKutta(order::Int=4; supersample::Int=1) = RungeKutta(order, supersample)
 
 "Get the `f!` and `h!` functions for the explicit Runge-Kutta solvers."
 function get_solver_functions(NT::DataType, solver::RungeKutta, fc!, hc!, Ts, _ , nx, _ , _ )
-    order = solver.order
-    Ts_inner = Ts/solver.supersample
     Nc = nx + 2
+    f! = if solver.order==4
+        get_rk4_function(NT, solver, fc!, Ts, nx, Nc)
+    elseif solver.order==1
+        get_euler_function(NT, solver, fc!, Ts, nx, Nc)
+    else
+        throw(ArgumentError("only 1st and 4th order Runge-Kutta is supported."))
+    end
+    h! = hc!
+    return f!, h!
+end
+
+"Get the f! function for the 4th order explicit Runge-Kutta solver."
+function get_rk4_function(NT, solver, fc!, Ts, nx, Nc)
+    Ts_inner = Ts/solver.supersample
     xcur_cache::DiffCache{Vector{NT}, Vector{NT}} = DiffCache(zeros(NT, nx), Nc)
     k1_cache::DiffCache{Vector{NT}, Vector{NT}}   = DiffCache(zeros(NT, nx), Nc)
     k2_cache::DiffCache{Vector{NT}, Vector{NT}}   = DiffCache(zeros(NT, nx), Nc)
     k3_cache::DiffCache{Vector{NT}, Vector{NT}}   = DiffCache(zeros(NT, nx), Nc)
     k4_cache::DiffCache{Vector{NT}, Vector{NT}}   = DiffCache(zeros(NT, nx), Nc)
-    if order==1
-        f! = function euler_solver!(xnext, x, u, d, p)
-            CT = promote_type(eltype(x), eltype(u), eltype(d))
-            xcur = get_tmp(xcur_cache, CT)
-            k1   = get_tmp(k1_cache, CT)
-            xterm = xnext
-            @. xcur = x
-            for i=1:solver.supersample
-                fc!(k1, xcur, u, d, p)
-                @. xcur = xcur + k1 * Ts_inner
-            end
-            @. xnext = xcur
-            return nothing
+    f! = function rk4_solver!(xnext, x, u, d, p)
+        CT = promote_type(eltype(x), eltype(u), eltype(d))
+        xcur = get_tmp(xcur_cache, CT)
+        k1   = get_tmp(k1_cache, CT)
+        k2   = get_tmp(k2_cache, CT)
+        k3   = get_tmp(k3_cache, CT)
+        k4   = get_tmp(k4_cache, CT)
+        xterm = xnext
+        @. xcur = x
+        for i=1:solver.supersample
+            fc!(k1, xcur, u, d, p)
+            @. xterm = xcur + k1 * Ts_inner/2
+            fc!(k2, xterm, u, d, p)
+            @. xterm = xcur + k2 * Ts_inner/2
+            fc!(k3, xterm, u, d, p)
+            @. xterm = xcur + k3 * Ts_inner
+            fc!(k4, xterm, u, d, p)
+            @. xcur = xcur + (k1 + 2k2 + 2k3 + k4)*Ts_inner/6
         end
-    elseif order==4
-        f! = function rk4_solver!(xnext, x, u, d, p)
-            CT = promote_type(eltype(x), eltype(u), eltype(d))
-            xcur = get_tmp(xcur_cache, CT)
-            k1   = get_tmp(k1_cache, CT)
-            k2   = get_tmp(k2_cache, CT)
-            k3   = get_tmp(k3_cache, CT)
-            k4   = get_tmp(k4_cache, CT)
-            xterm = xnext
-            @. xcur = x
-            for i=1:solver.supersample
-                fc!(k1, xcur, u, d, p)
-                @. xterm = xcur + k1 * Ts_inner/2
-                fc!(k2, xterm, u, d, p)
-                @. xterm = xcur + k2 * Ts_inner/2
-                fc!(k3, xterm, u, d, p)
-                @. xterm = xcur + k3 * Ts_inner
-                fc!(k4, xterm, u, d, p)
-                @. xcur = xcur + (k1 + 2k2 + 2k3 + k4)*Ts_inner/6
-            end
-            @. xnext = xcur
-            return nothing
+        @. xnext = xcur
+        return nothing
+    end
+    return f!
+end
+
+"Get the f! function for the explicit Euler solver."
+function get_euler_function(NT, solver, fc!, Ts, nx, Nc)
+    Ts_inner = Ts/solver.supersample
+    xcur_cache::DiffCache{Vector{NT}, Vector{NT}} = DiffCache(zeros(NT, nx), Nc)
+    k_cache::DiffCache{Vector{NT}, Vector{NT}}    = DiffCache(zeros(NT, nx), Nc)
+    f! = function euler_solver!(xnext, x, u, d, p)
+        CT = promote_type(eltype(x), eltype(u), eltype(d))
+        xcur = get_tmp(xcur_cache, CT)
+        k    = get_tmp(k_cache, CT)
+        xterm = xnext
+        @. xcur = x
+        for i=1:solver.supersample
+            fc!(k, xcur, u, d, p)
+            @. xcur = xcur + k * Ts_inner
         end
+        @. xnext = xcur
+        return nothing
     end
-    h! = hc!
-    return f!, h!
+    return f!
 end
 
 """