generic irk discretization

benchmark/prof.jl

@@ -9,7 +9,7 @@ using Profile
 using PProf
 using JET
 
-#include("../test/problems/goddard.jl")
+include("../test/problems/goddard.jl")
 include("../test/problems/simple_integrator.jl")
 
 # local version of mayer cost
@@ -19,18 +19,22 @@ function local_mayer(obj, x0, xf, v)
 end
 
 function init(;grid_size, disc_method)
-    #prob = goddard_all()
+    prob = goddard_all()
     #prob = goddard()
-    prob = simple_integrator()
+    #prob = simple_integrator()
     ocp = prob[:ocp]
-    docp = CTDirect.DOCP(ocp, grid_size=grid_size, time_grid=CTDirect.__time_grid(), disc_method=string(disc_method))
+    docp = CTDirect.DOCP(ocp, grid_size=grid_size, time_grid=CTDirect.__time_grid(), disc_method=disc_method)
     xu = CTDirect.DOCP_initial_guess(docp)
     return prob, docp, xu
 end
 
 
-function test_unit(;test_get=false, test_dyn=false, test_unit_cons=false, test_mayer=false, test_obj=true, test_block=false, test_cons=true, test_trans=true, test_solve=true, warntype=false, jet=false, profile=false, grid_size=100, disc_method=:trapeze)
-
+function test_unit(;test_get=false, test_obj=true, test_cons=true, test_trans=true, test_solve=true, warntype=false, jet=false, profile=false, grid_size=100, disc_method=:trapeze)
+
+    if profile
+        Profile.Allocs.clear()
+    end
+
     # define problem and variables
     prob, docp, xu = init(grid_size=grid_size, disc_method=disc_method)
     disc = docp.discretization
@@ -56,52 +60,7 @@ function test_unit(;test_get=false, test_dyn=false, test_unit_cons=false, test_m
         end
     end
 
-    f = similar(xu, docp.dim_NLP_x)
-
-    if test_dyn
-        print("dynamics_ext"); @btime $docp.dynamics_ext($f, $t, $x, $u, $v)
-        warntype && @code_warntype docp.dynamics_ext(f, t, x, u, v)
-        Profile.clear_malloc_data()
-        docp.dynamics_ext(f, t, x, u, v)
-    end
-
-    if test_unit_cons
-        print("u cons"); @btime $docp.control_constraints[2]($c, $t, $u, $v)
-        print("x cons"); @btime $docp.state_constraints[2]($c, $t, $x, $v)
-        print("xu cons"); @btime $docp.mixed_constraints[2]($c, $t, $x, $u, $v)
-        if warntype
-            @code_warntype docp.control_constraints[2](c, t, u, v)
-            @code_warntype docp.state_constraints[2](c, t, x, v)
-            @code_warntype docp.mixed_constraints[2](c, t, x, u, v)
-        end
-    end
-
-    # objective
-    if test_mayer
-        n = docp.dim_OCP_x
-        nx = docp.dim_NLP_x
-        m = docp.dim_NLP_u
-        N = docp.dim_NLP_steps
-        x0 = CTDirect.get_OCP_state_at_time_step(xu, docp, 1)
-        xf = CTDirect.get_OCP_state_at_time_step(xu, docp, N+1)
-        v = CTDirect.get_OCP_variable(xu, docp)
-        obj = similar(xu,1)
-
-        # local mayer
-        println("")
-        print("Local Mayer: views for x0/xf and scalar v"); @btime local_mayer($obj, (@view $xu[1:$n]), (@view $xu[($nx + $m) * $N + 1: ($nx + $m) * $N + $n]), $xu[end]) # OK
-        print("Local Mayer: param scal/vec getters"); @btime local_mayer($obj, $x0, $xf, $v) # OK
-        print("OCP Mayer: param scal/vec getters"); @btime $docp.ocp.mayer($obj, $x0, $xf, $v) # 3 allocs (112)
-
-        warntype && @code_warntype docp.ocp.mayer(obj, x0, xf, v)
-        jet && display(@report_opt docp.ocp.mayer(obj, x0, xf, v))
-        if profile
-            Profile.Allocs.@profile sample_rate=1.0 docp.ocp.mayer(obj, x0, xf, v)
-            results = Profile.Allocs.fetch()
-            PProf.Allocs.pprof()
-        end
-    end
-
+    # DOCP_objective
     if test_obj
         print("Objective"); @btime CTDirect.DOCP_objective($xu, $docp)
         warntype && @code_warntype CTDirect.DOCP_objective(xu, docp)
@@ -113,22 +72,6 @@ function test_unit(;test_get=false, test_dyn=false, test_unit_cons=false, test_m
         end
     end
 
-    if test_block
-        times = CTDirect.get_time_grid(xu, docp) # type OK
-        i = 1
-        v = CTDirect.get_OCP_variable(xu, docp)
-        work = CTDirect.setWorkArray(docp, xu, times, v)
-        print("Constraints block")
-        @btime CTDirect.setConstraintBlock!($docp, $c, $xu, $v, $times, $i, $work)
-        warntype && @code_warntype CTDirect.setConstraintBlock!(docp, c, xu, v, times, i, work)
-        jet && display(@report_opt CTDirect.setConstraintBlock!(docp, c, xu, v, times, i, work))
-        if profile
-            Profile.Allocs.@profile sample_rate=1.0 CTDirect.setConstraintBlock!(docp, c, xu, v, times, i, work)
-            results = Profile.Allocs.fetch()
-            PProf.Allocs.pprof()
-        end        
-    end
-
     # DOCP_constraints
     if test_cons
         print("Constraints"); @btime CTDirect.DOCP_constraints!($c, $xu, $docp)
@@ -144,16 +87,16 @@ function test_unit(;test_get=false, test_dyn=false, test_unit_cons=false, test_m
 
     # transcription
     if test_trans
-        print("Transcription"); @btime direct_transcription($prob.ocp, grid_size=$grid_size)
+        print("Transcription"); @btime direct_transcription($prob.ocp, grid_size=$grid_size, disc_method=$disc_method)
     end
 
     # solve
     if test_solve
-        sol = direct_solve(prob.ocp, display=false, grid_size=grid_size)
+        sol = direct_solve(prob.ocp, display=false, grid_size=grid_size, disc_method=disc_method)
         if !isapprox(sol.objective, prob.obj, rtol=1e-2)
             error("objective mismatch: ", sol.objective, " vs ", prob.obj)
         end
-        print("Solve"); @btime direct_solve($prob.ocp, display=false, grid_size=$grid_size)
+        print("Solve"); @btime direct_solve($prob.ocp, display=false, grid_size=$grid_size, disc_method=$disc_method)
     end
 
 end

src/CTDirect.jl

@@ -17,6 +17,7 @@ include("default.jl")
 include("problem.jl") # rename as docp.jl ?
 include("midpoint.jl")
 include("trapeze.jl")
+include("irk.jl")
 include("solution.jl")
 include("solve.jl")
 

src/default.jl

@@ -6,7 +6,7 @@ $(TYPEDSIGNATURES)
 Used to set the default discretization method.
 The default value is `trapeze`.
 """
-__disc_method() = "trapeze"
+__disc_method() = :trapeze
 
 """
 $(TYPEDSIGNATURES)