From c2da1aa07677b88cefa2311d44e3a84ab2dfd1cb Mon Sep 17 00:00:00 2001
From: Pierre Martinon <pierrecmartinon@gmail.com>
Date: Mon, 29 Apr 2024 18:42:02 +0200
Subject: [PATCH] added CI test for continuation

---
 test/Project.toml          |   1 +
 test/suite/continuation.jl | 119 ++++++++++++++++
 test/test_continuation.jl  | 278 ++++++++++++++++++-------------------
 3 files changed, 258 insertions(+), 140 deletions(-)
 create mode 100644 test/suite/continuation.jl

diff --git a/test/Project.toml b/test/Project.toml
index a2cd2f8..5533bba 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,3 +1,4 @@
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/suite/continuation.jl b/test/suite/continuation.jl
new file mode 100644
index 0000000..f33fea6
--- /dev/null
+++ b/test/suite/continuation.jl
@@ -0,0 +1,119 @@
+using CTDirect
+using Statistics
+
+println("Test: discrete continuation")
+
+# continuation on final time with parametric ocp definition 
+function ocp_T(T)
+    @def ocp begin
+        t ∈ [ 0, T ], time
+        x ∈ R², state
+        u ∈ R, control
+        q = x₁
+        v = x₂
+        q(0) == 0
+        v(0) == 0
+        q(T) == 1
+        v(T) == 0
+        ẋ(t) == [ v(t), u(t) ]
+        ∫(u(t)^2) → min
+    end
+    return ocp
+end
+@testset verbose = true showtiming = true ":parametric_ocp :warm_start" begin
+    init = OptimalControlInit()
+    obj_list = []
+    iter_list = []
+    for T=1:5
+        ocp = ocp_T(T) 
+        sol = solve(ocp, print_level=0, init=init)
+        init = OptimalControlInit(sol)
+        push!(obj_list, sol.objective)
+        push!(iter_list, sol.iterations)
+    end
+    @test last(obj_list) ≈ 0.096038 rtol=1e-2
+    @test mean(iter_list) ≈ 2.8 rtol=1e-2
+end
+
+
+# continuation with parametric definition of the ocp
+relu(x) = max(0, x)
+μ = 10
+p_relu(x) = log(abs(1 + exp(μ*x)))/μ
+f(x) = 1-x
+m(x) = (p_relu∘f)(x)
+T = 2
+function myocp(ρ)
+    @def ocp begin
+        τ ∈ R, variable
+        s ∈ [ 0, 1 ], time
+        x ∈ R², state
+        u ∈ R², control
+        x₁(0) == 0
+        x₂(0) == 1
+        x₁(1) == 1
+        ẋ(s) == [τ*(u₁(s)+2), (T-τ)*u₂(s)]
+        -1 ≤ u₁(s) ≤ 1
+        -1 ≤ u₂(s) ≤ 1
+        0 ≤ τ ≤ T
+        -(x₂(1)-2)^3 - ∫( ρ * ( τ*m(x₁(s))^2 + (T-τ)*m(x₂(s))^2 ) ) → min
+    end
+    return ocp
+end
+@testset verbose = true showtiming = true ":parametric_ocp :warm_start" begin
+    init = OptimalControlInit()
+    obj_list = []
+    iter_list = []
+    for ρ in [0.1, 5, 10, 30, 100]
+        ocp = myocp(ρ)
+        sol = solve(ocp, print_level=0, init=init)
+        init = OptimalControlInit(sol)
+        push!(obj_list, sol.objective)
+        push!(iter_list, sol.iterations)
+    end
+    @test last(obj_list) ≈ -148.022367 rtol=1e-2
+    @test mean(iter_list) ≈ 43.8 rtol=1e-2
+end
+
+# goddard
+Cd = 310
+Tmax = 3.5
+β = 500
+b = 2
+function F00(x)
+    r, v, m = x
+    D = Cd * v^2 * exp(-β*(r - 1))
+    return [ v, -D/m - 1/r^2, 0 ]
+end
+function F01(x)
+    r, v, m = x
+    return [ 0, Tmax/m, -b*Tmax ]
+end
+ocp = Model(variable=true)
+state!(ocp, 3)
+control!(ocp, 1)
+variable!(ocp, 1)
+time!(ocp, 0, Index(1))
+constraint!(ocp, :initial, [1,0,1], :initial_constraint)
+constraint!(ocp, :final, Index(3), 0.6, :final_constraint)
+constraint!(ocp, :state, 1:2:3, [1,0.6], [1.2,1], :state_box)
+constraint!(ocp, :control, Index(1), 0, 1, :control_box)
+constraint!(ocp, :variable, Index(1), 0.01, Inf, :variable_box)
+constraint!(ocp, :state, Index(2), 0, Inf, :speed_limit)
+objective!(ocp, :mayer, (x0, xf, v) -> xf[1], :max)
+dynamics!(ocp, (x, u, v) -> F00(x) + u*F01(x) )
+sol0 = solve(ocp, print_level=0)
+
+@testset verbose = true showtiming = true ":global_variable :warm_start" begin
+    sol = sol0
+    obj_list = []
+    iter_list = []
+    for Tmax_local=3.5:-0.5:1
+        global Tmax = Tmax_local 
+        sol = solve(ocp, print_level=0, init=sol)
+        push!(obj_list, sol.objective)
+        push!(iter_list, sol.iterations)
+    end
+    @test last(obj_list) ≈ 1.00359 rtol=1e-2
+    @test mean(iter_list) ≈ 17 rtol=1e-2
+end
diff --git a/test/test_continuation.jl b/test/test_continuation.jl
index 8a93a71..e004c1f 100644
--- a/test/test_continuation.jl
+++ b/test/test_continuation.jl
@@ -9,122 +9,45 @@ test2 = true
 test3 = true
 test4 = true
 
-#################################################
-# goddard max final altitude
-println("Test: discrete continuation")
-
-Cd = 310
-Tmax = 3.5
-β = 500
-b = 2
-function F0(x)
-    r, v, m = x
-    D = Cd * v^2 * exp(-β*(r - 1))
-    return [ v, -D/m - 1/r^2, 0 ]
-end
-function F1(x)
-    r, v, m = x
-    return [ 0, Tmax/m, -b*Tmax ]
-end
-
-ocp = Model(variable=true)
-state!(ocp, 3)
-control!(ocp, 1)
-variable!(ocp, 1)
-time!(ocp, 0, Index(1))
-constraint!(ocp, :initial, [1,0,1], :initial_constraint)
-constraint!(ocp, :final, Index(3), 0.6, :final_constraint)
-constraint!(ocp, :state, 1:2:3, [1,0.6], [1.2,1], :state_box)
-constraint!(ocp, :control, Index(1), 0, 1, :control_box)
-constraint!(ocp, :variable, Index(1), 0.01, Inf, :variable_box)
-constraint!(ocp, :state, Index(2), 0, Inf, :speed_limit)
-objective!(ocp, :mayer, (x0, xf, v) -> xf[1], :max)
-dynamics!(ocp, (x, u, v) -> F0(x) + u*F1(x) )
-
-# solve unconstrained problem
-if (test1 || test2)
-    sol0 = solve(ocp, print_level=0)
-    @printf("Objective for goddard reference solution %.6f",  sol0.objective)
-    sol = sol0
-end
-
+# continuation on fixed final time
+# NB. time! can be called only once, so we redefine the ocp
 if test1
-    # default init
-    print("\nSolve goddard for different speed limits, default initial guess\nvmax ")
-    iter_list = []
-    for vmax=0.15:-0.01:0.05
-        print(vmax," ")
-        remove_constraint!(ocp, :speed_limit)
-        constraint!(ocp, :state, Index(2), 0, vmax, :speed_limit)
-        global sol = solve(ocp, print_level=0)   
-        #@printf("vmax %.2f objective %.6f iterations %d\n",vmax,sol.objective, sol.iterations)
-        push!(iter_list, sol.iterations)
+    println("\nDiscrete continuation on final time")
+    # parametric ocp
+    function ocp_T(T)
+        @def ocp begin
+            t ∈ [ 0, T ], time
+            x ∈ R², state
+            u ∈ R, control
+            q = x₁
+            v = x₂
+            q(0) == 0
+            v(0) == 0
+            q(T) == 1
+            v(T) == 0
+            ẋ(t) == [ v(t), u(t) ]
+            ∫(u(t)^2) → min
+        end
+        return ocp
     end
-    @printf("\nAverage iterations %d\n", mean(iter_list))
 
-    # warm start
-    print("\nDiscrete continuation on speed limit, with warm start\nvmax ")
-    vmax_list = []
-    obj_list = []
+    # continuation on final time
+    init1 = OptimalControlInit()
     iter_list = []
-    for vmax=0.15:-0.01:0.05
-        print(vmax," ")
-        remove_constraint!(ocp, :speed_limit)
-        constraint!(ocp, :state, Index(2), 0, vmax, :speed_limit)
-        global sol = solve(ocp, print_level=0, init=sol)  
-        #@printf("vmax %.2f objective %.6f iterations %d\n",vmax,sol.objective, sol.iterations)
-        push!(vmax_list, vmax)
-        push!(obj_list, sol.objective)
-        push!(iter_list, sol.iterations)
-    end
-    @printf("\nAverage iterations %d\n", mean(iter_list))
-
-    # plot obj(vmax)
-    pobj = plot(vmax_list, obj_list, label="r(tf)", xlabel="Speed limit (vmax)", ylabel="Maximal altitude r(tf)",seriestype=:scatter)
-
-    # plot multiple solutions
-    plot(sol0)
-    p = plot!(sol)
-
-    display(plot(pobj, p, layout=2))
-end
-
-if test2
-    print("\nDiscrete continuation on maximal thrust\nTmax ")
-    # reset speed constraint
-    remove_constraint!(ocp, :speed_limit)
-    vmax = 0.1
-    constraint!(ocp, :state, Index(2), 0, vmax, :speed_limit)
-    sol2 = solve(ocp, print_level=0)
-    sol0 = sol2
-
-    # continuation on Tmax (using default init is slower)
-    Tmax_list = []
-    obj_list = []
-    for Tmax_local=3.5:-0.5:1
-        global Tmax = Tmax_local 
-        print(Tmax," ")   
-        global sol2 = solve(ocp, print_level=0, init=sol2)
-        push!(Tmax_list, Tmax)
-        push!(obj_list, sol2.objective)
-        #print(Tmax, " ", sol2.objective, " ", sol2.iterations)
+    for T=1:5
+        ocp1 = ocp_T(T) 
+        sol1 = solve(ocp1, print_level=0, init=init1)
+        global init1 = OptimalControlInit(sol1)
+        @printf("T %.2f objective %.6f iterations %d\n", T, sol1.objective, sol1.iterations)
+        push!(iter_list, sol1.iterations)
     end
-    
-    # plot obj(vmax)
-    pobj = plot(Tmax_list, obj_list, label="r(tf)", xlabel="Maximal thrust (Tmax)", ylabel="Maximal altitude r(tf)",seriestype=:scatter)
-
-    # plot multiple solutions
-    plot(sol0)
-    p = plot!(sol2)
-    
-    display(plot(pobj, p, layout=2, reuse=false))
-
+    @printf("Average iterations %d\n", mean(iter_list))
 end
 
 
 # continuation with parametric definition of the ocp
-if test3
-    println("\nDiscrete contiuation on parametric ocp")
+if test2
+    println("\nDiscrete continuation on parametric ocp")
     # definition of the parametric OCP
     relu(x) = max(0, x)
     μ = 10
@@ -152,45 +75,120 @@ if test3
     end
 
     # continuation on rho
-    init3 = OptimalControlInit()
+    init2 = OptimalControlInit()
+    iter_list = []
     ρs = [0.1, 5, 10, 30, 100]
     for ρ in ρs
-        ocp3 = myocp(ρ)
-        sol3 = solve(ocp3, print_level=0, init=init3)
-        global init3 = OptimalControlInit(sol3)
-        @printf("Rho %.2f objective %.6f iterations %d\n", ρ, sol3.objective, sol3.iterations)
+        ocp2 = myocp(ρ)
+        sol2 = solve(ocp2, print_level=0, init=init2)
+        global init2 = OptimalControlInit(sol2)
+        @printf("Rho %.2f objective %.6f iterations %d\n", ρ, sol2.objective, sol2.iterations)
+        push!(iter_list, sol2.iterations)
     end
+    @printf("Average iterations %d\n", mean(iter_list))
 end
 
-# continuation on fixed final time
-# NB. time! can be called only once, so we redefine the ocp
-if (test4)
-    println("\nDiscrete continuation on final time (double integrator)")
-    # parametric ocp
-    function ocp_T(T)
-        @def ocp begin
-            t ∈ [ 0, T ], time
-            x ∈ R², state
-            u ∈ R, control
-            q = x₁
-            v = x₂
-            q(0) == 0
-            v(0) == 0
-            q(T) == 1
-            v(T) == 0
-            ẋ(t) == [ v(t), u(t) ]
-            ∫(u(t)^2) → min
-        end
-        return ocp
+
+# goddard max final altitude
+if (test3 || test4)
+    Cd = 310
+    Tmax = 3.5
+    β = 500
+    b = 2
+    function F0(x)
+        r, v, m = x
+        D = Cd * v^2 * exp(-β*(r - 1))
+        return [ v, -D/m - 1/r^2, 0 ]
+    end
+    function F1(x)
+        r, v, m = x
+        return [ 0, Tmax/m, -b*Tmax ]
     end
 
-    # contiuation on final time
-    init4 = OptimalControlInit()
-    for T=1:5
-        ocp4 = ocp_T(T) 
-        sol4 = solve(ocp4, print_level=0, init=init4)
-        global init4 = OptimalControlInit(sol4)
-        @printf("T %.2f objective %.6f iterations %d\n", T, sol4.objective, sol4.iterations)
+    ocp = Model(variable=true)
+    state!(ocp, 3)
+    control!(ocp, 1)
+    variable!(ocp, 1)
+    time!(ocp, 0, Index(1))
+    constraint!(ocp, :initial, [1,0,1], :initial_constraint)
+    constraint!(ocp, :final, Index(3), 0.6, :final_constraint)
+    constraint!(ocp, :state, 1:2:3, [1,0.6], [1.2,1], :state_box)
+    constraint!(ocp, :control, Index(1), 0, 1, :control_box)
+    constraint!(ocp, :variable, Index(1), 0.01, Inf, :variable_box)
+    constraint!(ocp, :state, Index(2), 0, Inf, :speed_limit)
+    objective!(ocp, :mayer, (x0, xf, v) -> xf[1], :max)
+    dynamics!(ocp, (x, u, v) -> F0(x) + u*F1(x) )
+
+    # solve unconstrained problem
+    sol0 = solve(ocp, print_level=0)
+    @printf("\nObjective for goddard reference solution %.6f",  sol0.objective)
+end
+
+
+# using a global variable in ocp definition
+if test3
+    print("\nDiscrete continuation on maximal thrust\nTmax ")
+    # continuation on Tmax (using default init is slower)
+    Tmax_list = []
+    obj_list = []
+    iter_list = []
+    sol3 = sol0
+    for Tmax_local=3.5:-0.5:1
+        global Tmax = Tmax_local 
+        print(Tmax," ")   
+        global sol3 = solve(ocp, print_level=0, init=sol3)
+        push!(Tmax_list, Tmax)
+        push!(obj_list, sol3.objective)
+        push!(iter_list, sol3.iterations)
+    end
+    @printf("\nAverage iterations %d\n", mean(iter_list))
+
+    # plot obj(vmax)
+    pobj = plot(Tmax_list, obj_list, label="r(tf)", xlabel="Maximal thrust (Tmax)", ylabel="Maximal altitude r(tf)",seriestype=:scatter)
+    # plot multiple solutions
+    plot(sol0)
+    p = plot!(sol3)
+    display(plot(pobj, p, layout=2, reuse=false))
+end
+
+
+# manually edit a constraint in ocp
+if test4
+    # reset Tmax
+    global Tmax = 3.5
+    # default init
+    print("\nSolve goddard for different speed limits, default initial guess\nvmax ")
+    iter_list = []
+    for vmax=0.15:-0.01:0.05
+        print(vmax," ")
+        remove_constraint!(ocp, :speed_limit)
+        constraint!(ocp, :state, Index(2), 0, vmax, :speed_limit)
+        global sol = solve(ocp, print_level=0) 
+        push!(iter_list, sol.iterations)
     end
+    @printf("\nAverage iterations %d\n", mean(iter_list))
 
-end
\ No newline at end of file
+    # warm start
+    print("Discrete continuation on speed limit, with warm start\nvmax ")
+    vmax_list = []
+    obj_list = []
+    iter_list = []
+    sol = sol0
+    for vmax=0.15:-0.01:0.05
+        print(vmax," ")
+        remove_constraint!(ocp, :speed_limit)
+        constraint!(ocp, :state, Index(2), 0, vmax, :speed_limit)
+        global sol = solve(ocp, print_level=0, init=sol)
+        push!(vmax_list, vmax)
+        push!(obj_list, sol.objective)
+        push!(iter_list, sol.iterations)
+    end
+    @printf("\nAverage iterations %d\n", mean(iter_list))
+
+    # plot obj(vmax)
+    pobj = plot(vmax_list, obj_list, label="r(tf)", xlabel="Speed limit (vmax)", ylabel="Maximal altitude r(tf)",seriestype=:scatter)
+    # plot multiple solutions
+    plot(sol0)
+    p = plot!(sol)
+    display(plot(pobj, p, layout=2))
+end