init struct and constructors (constant init, warm start) (#48)

Author: PierreMartinon

src/CTDirect.jl

@@ -13,6 +13,7 @@ const __display = CTBase.__display
 const matrix2vec = CTBase.matrix2vec
 
 # includes
+include("init.jl")
 include("utils.jl")
 include("problem.jl")
 include("solution.jl")
@@ -21,5 +22,6 @@ include("solve.jl")
 # export functions only for user
 export solve
 export is_solvable
+export OptimalControlInit
 
 end

src/init.jl

@@ -0,0 +1,52 @@
+# NB. to be moved up to CTBase later
+
+# init struct (to be passed to solve)
+# functions of time for state and control variables
+# values for optimization variables (NB. free t0/tf in particular !)
+# later for indirect methods add costate and structure information
+
+# constructors
+# - default: keep empty for method-specific handling
+# - from constant vector (dim x + dim u + dim v)
+# - from existing solution
+mutable struct OptimalControlInit
+
+    state_dimension
+    control_dimension
+    variable_dimension
+    state_init
+    control_init
+    variable_init
+    info
+
+    function OptimalControlInit()
+        init = new()
+        init.info = :undefined
+        return init
+    end
+
+    function OptimalControlInit(x_init, u_init, v_init=Real[])
+        init = new()
+        init.info = :constant
+        init.state_dimension = length(x_init)
+        init.control_dimension = length(u_init)
+        init.variable_dimension = length(v_init)
+        #println("Init dims x u v: ",init.state_dimension," ",init.control_dimension," ",init.variable_dimension)
+        init.state_init = t -> x_init
+        init.control_init = t -> u_init
+        init.variable_init = v_init
+        return init
+    end
+
+    function OptimalControlInit(sol::OptimalControlSolution)
+        init = new()
+        init.info = :solution
+        init.state_dimension = sol.state_dimension
+        init.control_dimension = sol.control_dimension
+        init.variable_dimension = sol.variable_dimension
+        init.state_init = t-> sol.state(t)[1:sol.state_dimension]
+        init.control_init = sol.control
+        init.variable_init = sol.infos[:variable]  #+++ need proper variable field in ocp solution !
+        return init
+    end
+end

src/problem.jl

@@ -54,10 +54,12 @@ mutable struct CTDirect_data
 
     ## NLP
     # NLP problem
-    dim_NLP_state  # 'augmented state'
+    dim_NLP_state  # 'augmented state' with possible additional component for lagrange cost
     dim_NLP_constraints
     dim_NLP_variables
     dim_NLP_steps
+
+    # initialization
     NLP_init
 
     # NLP solution
@@ -68,7 +70,8 @@ mutable struct CTDirect_data
     NLP_iterations
     NLP_stats       # remove later ? type is https://juliasmoothoptimizers.github.io/SolverCore.jl/stable/reference/#SolverCore.GenericExecutionStats
 
-    function CTDirect_data(ocp::OptimalControlModel, N::Integer, init=nothing)
+
+    function CTDirect_data(ocp::OptimalControlModel, N::Integer, init::OptimalControlInit)       
 
         ctd = new()
 
@@ -224,12 +227,14 @@ end
 function variables_bounds(ctd)
 
     N = ctd.dim_NLP_steps
-    l_var = -Inf*ones(ctd.dim_NLP_variables)
-    u_var = Inf*ones(ctd.dim_NLP_variables)
+    l_var = -Inf * ones(ctd.dim_NLP_variables)
+    u_var = Inf * ones(ctd.dim_NLP_variables)
 
     # NLP variables layout: [X0, X1 .. XN, U0, U1 .. UN, V]
     # NB. keep offset for each block since blocks are optional !
 
+    # +++ we could use the setters here (build local vectors then call setter ?!)
+
     # state box
     offset = 0
     if ctd.has_state_box

src/solution.jl

@@ -30,7 +30,7 @@ function _OptimalControlSolution(ocp, ipopt_solution, ctd)
     sol.state = t -> x(t)
     sol.costate = t -> p(t)
     sol.control = t -> u(t)
-    #sol.variable = v +++no field variable?
+    #sol.variable = v +++need proper field variable
     sol.infos[:variable] = v
     sol.objective = ctd.NLP_objective
     sol.iterations = ctd.NLP_iterations

src/solve.jl

@@ -32,7 +32,7 @@ function solve(ocp::OptimalControlModel,
     print_level::Integer=__print_level_ipopt(),
     mu_strategy::String=__mu_strategy_ipopt(),
     display::Bool=__display(),
-    init=nothing,
+    init::OptimalControlInit=OptimalControlInit(),
     kwargs...)
 
     # get full description from partial

src/utils.jl

@@ -77,72 +77,81 @@ end
 function get_final_time(xu, ctd)
     if ctd.has_free_final_time
         v = get_variable(xu, ctd)
-        #println("v: ",v)
-        #println("ctd.final_time: ",ctd.final_time)
-        #println("v[ctd.final_time]: ", v[ctd.final_time])
         return v[ctd.final_time]
     else
         return ctd.final_time
     end
 end
 
+#=
+function get_time_step(xu, ctd, i)
+    N = ctd.dim_NLP_steps
+    @assert i <= N "trying to get time step for i > N"
+    t0 = get_initial_time(xu, ctd)
+    tf = get_final_time(xu, ctd)
+    return t0 + i * (tf - t0) / N 
+end
+=#
 
 ## Initialization for the NLP problem
 
 function set_state_at_time_step!(xu, x_init, ctd, i)
     nx = ctd.dim_NLP_state
+    n = ctd.state_dimension
     N = ctd.dim_NLP_steps
     @assert i <= N "trying to set init for x(t_i) with i > N"
-    xu[1+i*nx:(i+1)*nx] = x_init[1:nx]
+    # NB. only set first n components of state variable (nx = n+1 for lagrange cost)
+    if n == 1
+        xu[i*nx + 1] = x_init[]
+    else
+        xu[i*nx + 1 : i*nx + n] = x_init
+    end
 end
 
 function set_control_at_time_step!(xu, u_init, ctd, i)
     nx = ctd.dim_NLP_state
     m = ctd.control_dimension
     N = ctd.dim_NLP_steps
     @assert i <= N "trying to set init for u(t_i) with i > N"
-    xu[1+(N+1)*nx+i*m:m+(N+1)*nx+i*m] = u_init[1:m]
+    offset = (N+1)*nx
+    if m == 1
+        xu[offset + i*m + 1] = u_init[]
+    else        
+        xu[offset + i*m + 1 : offset + i*m + m] = u_init
+    end
 end
 
 function set_variable!(xu, v_init, ctd)
-    xu[end-ctd.variable_dimension+1:end] = v_init[1:ctd.variable_dimension]
+    if ctd.variable_dimension == 1
+        xu[end] = v_init[]
+    else
+        xu[end-ctd.variable_dimension+1 : end] = v_init
+    end
 end
 
 function initial_guess(ctd)
 
-    N = ctd.dim_NLP_steps
-    init = ctd.NLP_init
+    # default initialization
+    # note: internal variables (lagrange cost, k_i for RK schemes) will keep these default values 
+    xu0 = 0.1 * ones(ctd.dim_NLP_variables)
 
-    if init === nothing
-        # default initialization
-        xu0 = 0.1*ones(ctd.dim_NLP_variables)
-    else
-        # split state / control init values
-        if length(init) != (ctd.state_dimension + ctd.control_dimension + ctd.variable_dimension)
-            error("vector for initialization should be of size dim_x + dim_u + dim_v ie:",ctd.state_dimension+ctd.control_dimension+ctd.variable_dimension)
-        end
-        x_init = zeros(ctd.dim_NLP_state)
-        x_init[1:ctd.state_dimension] = init[1:ctd.state_dimension]
-        u_init = zeros(ctd.control_dimension)
-        u_init[1:ctd.control_dimension] = init[ctd.state_dimension+1:ctd.state_dimension+ctd.control_dimension]
-
-        # mayer -> lagrange additional state
-        if ctd.has_lagrange_cost
-            x_init[ctd.dim_NLP_state] = 0.1
-        end
+    init = ctd.NLP_init
+    if init.info != :undefined
+        N = ctd.dim_NLP_steps
+        t0 = get_initial_time(xu0, ctd)
+        tf = get_final_time(xu0, ctd)
+        h = (tf - t0) / N 
 
-        # set constant initialization for state / control variables
-        xu0 = zeros(ctd.dim_NLP_variables)
+        # set state / control variables
         for i in 0:N
-            set_state_at_time_step!(xu0, x_init, ctd, i)
-            set_control_at_time_step!(xu0, u_init, ctd, i)
+            ti = t0 + i * h
+            set_state_at_time_step!(xu0, init.state_init(ti), ctd, i)
+            set_control_at_time_step!(xu0, init.control_init(ti), ctd, i)
         end
 
         # set variables
         if (ctd.variable_dimension > 0)
-            v_init = zeros(ctd.variable_dimension)
-            v_init[1:ctd.variable_dimension] = init[ctd.state_dimension+ctd.control_dimension+1:ctd.state_dimension+ctd.control_dimension+ctd.variable_dimension]    
-            set_variable!(xu0, v_init, ctd)
+            set_variable!(xu0, init.variable_init, ctd)
         end
     end
 

test/suite/goddard_all_constraints.jl

@@ -50,10 +50,26 @@ function F1(x)
     return [ 0, Tmax/m, -b*Tmax ]
 end
 dynamics!(ocp, (x, u, v) -> F0(x) + u*F1(x) )
-init_constant = [1.05, 0.1, 0.8, 0.5, 0.1]
-sol = solve(ocp, grid_size=30, print_level=0, tol=1e-8, init=init_constant)
-#println(sol.objective)
-#plot(sol)
+sol1 = solve(ocp, grid_size=30, print_level=0, tol=1e-8)
 @testset verbose = true showtiming = true ":goddard :max_rf :all_constraints_types" begin
-    @test sol.objective ≈ 1.0125 rtol=1e-2
-end
+    @test sol1.objective ≈ 1.0125 rtol=1e-2
+end
+
+
+# with constant initial guess
+x_init = [1.05, 0.1, 0.8]
+u_init = 0.5
+v_init = 0.1
+init_constant = OptimalControlInit(x_init, u_init, v_init)
+sol2 = solve(ocp, grid_size=30, print_level=0, tol=1e-8, init=init_constant)
+@testset verbose = true showtiming = true ":goddard :max_rf :all_constraints_types :init_constant" begin
+    @test sol2.objective ≈ 1.0125 rtol=1e-2
+end
+
+
+# with initial guess from solution
+init_sol = OptimalControlInit(sol2)
+sol3 = solve(ocp, grid_size=30, print_level=0, tol=1e-8, init=init_sol)
+@testset verbose = true showtiming = true ":goddard :max_rf :all_constraints_types :init_sol" begin
+    @test sol3.objective ≈ 1.0125 rtol=1e-2
+end

test/suite/simple_integrator.jl

@@ -12,8 +12,23 @@ constraint!(ocp1, :initial, -1, :initial_constraint)
 constraint!(ocp1, :final, 0, :final_constraint)
 dynamics!(ocp1, (x, u) -> -x + u)
 objective!(ocp1, :lagrange, (x, u) -> u^2)
-init_constant = [-0.5, 0]
-sol1 = solve(ocp1, grid_size=100, print_level=0, tol=1e-12, init=init_constant)
+sol1 = solve(ocp1, grid_size=100, print_level=0, tol=1e-12)
 @testset verbose = true showtiming = true ":double_integrator :min_tf" begin
     @test sol1.objective ≈ 0.313 rtol=1e-2
 end
+
+# with initial guess (using both vector and scalar syntax, no optimization variables)
+x_init = [-0.5]
+u_init = 0
+init_constant = OptimalControlInit(x_init, u_init)
+sol2 = solve(ocp1, grid_size=100, print_level=0, tol=1e-12, init=init_constant)
+@testset verbose = true showtiming = true ":double_integrator :min_tf :init_constant" begin
+    @test sol2.objective ≈ 0.313 rtol=1e-2
+end
+
+# with initial guess from solution
+init_sol = OptimalControlInit(sol2)
+sol3 = solve(ocp1, grid_size=100, print_level=0, tol=1e-12, init=init_sol)
+@testset verbose = true showtiming = true ":double_integrator :min_tf :init_sol" begin
+    @test sol3.objective ≈ 0.313 rtol=1e-2
+end