diff --git a/Project.toml b/Project.toml index 492e9be2..c7cdc9e1 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "CTDirect" uuid = "790bbbee-bee9-49ee-8912-a9de031322d5" authors = ["Olivier Cots "] -version = "0.12.0" +version = "0.12.1" [deps] ADNLPModels = "54578032-b7ea-4c30-94aa-7cbd1cce6c9a" diff --git a/benchmark/benchmark.jl b/benchmark/benchmark.jl index fa29fb49..a8cafa8d 100644 --- a/benchmark/benchmark.jl +++ b/benchmark/benchmark.jl @@ -21,6 +21,7 @@ function bench(; precompile = true, display = false, verbose = true, + discretization = :trapeze ) ####################################################### @@ -66,6 +67,7 @@ function bench(; linear_solver = linear_solver, max_iter = 0, display = display, + discretization = discretization ) t_precomp += t end @@ -84,6 +86,7 @@ function bench(; linear_solver = linear_solver, grid_size = grid_size, tol = tol, + discretization = discretization ) if !isnothing(problem[:obj]) && !isapprox(sol.objective, problem[:obj], rtol = 5e-2) error( @@ -115,29 +118,29 @@ function bench(; end # +++ put repeat directly in bench() -function bench_average(; repeat = 2, verbose = false, kwargs...) +function bench_average(; repeat = 4, verbose = false, kwargs...) # execute series of benchmark runs t_list = [] for i = 1:repeat t = bench(; verbose = verbose, kwargs...) append!(t_list, t) - @printf("Run %d / %d: time (s) = %6.2f\n", i, repeat, t) + verbose && @printf("Run %d / %d: time (s) = %6.2f\n", i, repeat, t) end # print / return average total time avg_time = sum(t_list) / length(t_list) - @printf("Average time (s): %6.2f\n", avg_time) + verbose && @printf("Average time (s): %6.2f\n", avg_time) return avg_time end -function bench_series(; grid_size_list = [250, 500, 1000, 2500, 5000, 10000], kwargs...) +function bench_series(; grid_size_list = [250, 500, 1000, 2500, 5000], kwargs...) println(grid_size_list) t_list = [] for grid_size in grid_size_list t = bench_average(; grid_size = grid_size, kwargs...) append!(t_list, t) - @printf("Grid size %d: time (s) = %6.2f\n\n", grid_size, t) + @printf("Grid size %d: time (s) = %6.1f\n", grid_size, t) end - return t_list + #return t_list end diff --git a/benchmark/prof.jl b/benchmark/prof.jl index 01297ae2..692bff6e 100644 --- a/benchmark/prof.jl +++ b/benchmark/prof.jl @@ -1,6 +1,10 @@ +# +++ TODO: make function with bools as args ? # Profiling -include("../test/deps.jl") +using CTDirect +using CTBase +using LinearAlgebra +using NLPModelsIpopt using BenchmarkTools #using Traceur #using Profile @@ -25,20 +29,31 @@ println("Load problem ", prob[:name]) if precompile println("Precompilation") - direct_solve(ocp, grid_size = grid_size, display = false, max_iter = 2) - CTDirect.DOCP_objective(CTDirect.DOCP_initial_guess(docp), docp) - CTDirect.DOCP_constraints!( - zeros(docp.dim_NLP_constraints), - CTDirect.DOCP_initial_guess(docp), - docp, - ) + if test_solve + direct_solve(ocp, grid_size = grid_size, display = false, max_iter = 2) + end + if test == :objective + CTDirect.DOCP_objective(CTDirect.DOCP_initial_guess(docp), docp) + else + CTDirect.DOCP_constraints!(zeros(docp.dim_NLP_constraints), CTDirect.DOCP_initial_guess(docp), docp) + end +end + +if test == :objective + println("Timed objective") + #@timev CTDirect.DOCP_objective(CTDirect.DOCP_initial_guess(docp), docp) + @btime CTDirect.DOCP_objective(CTDirect.DOCP_initial_guess(docp), docp) +else + println("Timed constraints") + #@timev CTDirect.DOCP_constraints!(zeros(docp.dim_NLP_constraints), CTDirect.DOCP_initial_guess(docp), docp) + @btime CTDirect.DOCP_constraints!(zeros(docp.dim_NLP_constraints), CTDirect.DOCP_initial_guess(docp), docp) end # full solve if test_solve - println("Timed solve") - @timev sol = direct_solve(ocp, grid_size = grid_size, display = false) - @btime sol = direct_solve(ocp, grid_size = grid_size, display = false) + println("Timed full solve") + #@timev sol = direct_solve(ocp, grid_size = grid_size, display=false) + @btime sol = direct_solve(ocp, grid_size = grid_size, display=false) end if test_code_warntype @@ -71,12 +86,3 @@ if test_jet end end -#= -# ipopt statistics -if test_ipopt - docp = directTranscription(ocp, grid_size=200, init=(state=[1,0.2,0.5], control=0.5)) - dsol = solve(docp, print_level=5, tol=1e-12, print_timing_statistics="yes") - println("\n Discrete solution") - println(dsol) -end -=# diff --git a/src/CTDirect.jl b/src/CTDirect.jl index 74e17ab4..367257d4 100644 --- a/src/CTDirect.jl +++ b/src/CTDirect.jl @@ -13,7 +13,9 @@ const matrix2vec = CTBase.matrix2vec # includes include("utils.jl") include("default.jl") -include("problem.jl") +include("problem.jl") # rename as docp.jl ? +include("midpoint.jl") +include("trapeze.jl") include("solution.jl") include("solve.jl") diff --git a/src/default.jl b/src/default.jl index e8472d33..ab8e9c63 100644 --- a/src/default.jl +++ b/src/default.jl @@ -3,6 +3,14 @@ """ $(TYPEDSIGNATURES) +Used to set the default discretization method. +The default value is `trapeze`. +""" +__discretization() = "trapeze" + +""" +$(TYPEDSIGNATURES) + Used to set the default grid size. The default value is `250`. """ diff --git a/src/gauss.jl b/src/gauss.jl new file mode 100644 index 00000000..954b3a61 --- /dev/null +++ b/src/gauss.jl @@ -0,0 +1,105 @@ +#= Functions for implicit gauss-legendre 2 discretization scheme +Internal layout for NLP variables: +[X_0,U_0,K_0, X_1,U_1,K_1 .., X_N-1,U_N-1,K_N-1, XN, V] +with the conventions U_i constant per step, K_i := [K_i1 K_i2] + +NB. try both control_stage and control_step versions ! +Q. if taking controls at stages, how to define the control at 'step' (for path constraints or tf). Unless we evaluate path constraints at stages, but then we have to determine the state at stages (a possibility is to take the state argument as for the stage dynamics) +=# + +# Later adjust this one for generic IRK scheme ! (use stage value) +# define several tags for different methods +# use intermediate abstract type ImplicitRungeKuttaTag + +struct GaussLegendre2Tag <: DiscretizationTag + stage::Int + additional_controls::Int + butcher_a + butcher_b + butcher_c + GaussLegendre2Tag() = new(2, 0, [0.25 (0.25 - sqrt(3)/6); (0.25 + sqrt(3)/6) 0.25], [0.5, 0.5], [(0.5 - sqrt(3)/6), (0.5 + sqrt(3)/6)]) +end + + +""" +$(TYPEDSIGNATURES) + +Retrieve state and control variables at given time step from the NLP variables. +""" +function get_variables_at_time_step(xu, docp, i, tag::GuassLegendre2Tag) + + # block: [X_i U_i1 U_i2 K_i1 K_i2] + nx = docp.dim_NLP_x + n = docp.dim_OCP_x + m = docp.dim_NLP_u + N = docp.dim_NLP_steps + offset = (nx*3 + m) * i + + # retrieve scalar/vector OCP state (w/o lagrange state) + if n == 1 + xi = xu[offset + 1] + else + xi = xu[(offset + 1):(offset + n)] + end + if docp.has_lagrange + xli = xu[offset + nx] + else + xli = nothing # dummy. use xu type ? + end + + # retrieve scalar/vector controls + if i < N + offset_u = offset + else + offset_u = (nx*3 + m) * (i-1) + end + if m == 1 + ui = xu[offset_u + nx + 1] + else + ui = xu[(offset_u + nx + 1):(offset_u + nx + m)] + end + + # retrieve vector stage variable (except at final time) + if i < N + ki = (xu[(offset + nx + m + 1):(offset + nx + m + nx)], + xu[(offset + nx*2 + m + 1):(offset + nx*2 + m + nx)]) + else + ki = nothing + end + + return xi, ui, xli, ki +end + + +# internal NLP version for solution parsing +# could be fused with one above if +# - using extended dynamics that include lagrange cost +# - scalar case is handled at OCP level +function get_NLP_variables_at_time_step(xu, docp, i, tag::GaussLegendre2Tag) + + +++ + nx = docp.dim_NLP_x + m = docp.dim_NLP_u + N = docp.dim_NLP_steps + offset = (nx*2 + m) * i + + # state + xi = xu[(offset + 1):(offset + nx)] + # control + if i < N + offset_u = offset + else + offset_u = (nx*2 + m) * (i-1) + end + ui = xu[(offset_u + nx + 1):(offset_u + nx + m)] + # stage + if i < N + ki = xu[(offset + nx + m + 1):(offset + nx + m + nx) ] + else + ki = nothing + end + + return xi, ui, ki +end + + diff --git a/src/midpoint.jl b/src/midpoint.jl new file mode 100644 index 00000000..3f31fa3f --- /dev/null +++ b/src/midpoint.jl @@ -0,0 +1,239 @@ +#= Functions for implicit midpoint discretization scheme +Internal layout for NLP variables: +[X_0,U_0,K_0, X_1,U_1,K_1 .., X_N-1,U_N-1,K_N-1, XN, V] +with the convention u([t_i,t_i+1[) = U_i and u(tf) = U_N-1 +=# + +# +++ TODO: use args + +struct MidpointTag <: DiscretizationTag + stage::Int + additional_controls::Int + MidpointTag() = new(1, 0) +end + + +""" +$(TYPEDSIGNATURES) + +Retrieve state and control variables at given time step from the NLP variables. +""" +function get_variables_at_time_step(xu, docp::DOCP{MidpointTag}, i) + + nx = docp.dim_NLP_x + n = docp.dim_OCP_x + m = docp.dim_NLP_u + N = docp.dim_NLP_steps + offset = (nx*(1+docp.discretization.stage) + m) * i + + # retrieve scalar/vector OCP state (w/o lagrange state) + if n == 1 + xi = xu[offset + 1] + else + xi = xu[(offset + 1):(offset + n)] + end + if docp.has_lagrange + xli = xu[offset + nx] + else + xli = nothing # dummy. use xu type ? + end + + # retrieve scalar/vector control (convention u(tf) = U_N-1) + if i < N + offset_u = offset + else + offset_u = (nx*2 + m) * (i-1) + end + if m == 1 + ui = xu[offset_u + nx + 1] + else + ui = xu[(offset_u + nx + 1):(offset_u + nx + m)] + end + + # retrieve vector stage variable (except at final time) + if i < N + ki = xu[(offset + nx + m + 1):(offset + nx + m + nx) ] + else + ki = nothing + end + + return xi, ui, xli, ki +end + + +# internal NLP version for solution parsing +# could be fused with one above if +# - using extended dynamics that include lagrange cost +# - scalar case is handled at OCP level +function get_NLP_variables_at_time_step(xu, docp, i, tag::MidpointTag) + + nx = docp.dim_NLP_x + m = docp.dim_NLP_u + N = docp.dim_NLP_steps + offset = (nx*2 + m) * i + + # state + xi = xu[(offset + 1):(offset + nx)] + # control + if i < N + offset_u = offset + else + offset_u = (nx*2 + m) * (i-1) + end + ui = xu[(offset_u + nx + 1):(offset_u + nx + m)] + # stage + if i < N + ki = xu[(offset + nx + m + 1):(offset + nx + m + nx) ] + else + ki = nothing + end + + return xi, ui, ki +end + + +function set_variables_at_time_step!(xu, x_init, u_init, docp, i, tag::MidpointTag) + + nx = docp.dim_NLP_x + n = docp.dim_OCP_x + m = docp.dim_NLP_u + N = docp.dim_NLP_steps + offset = (nx*2 + m) * i + + # NB. only set the actual state variables from the OCP + # - skip the possible additional state for lagrange cost + # - skip internal discretization variables (K_i) + if !isnothing(x_init) + xu[(offset + 1):(offset + n)] .= x_init + end + if (i < N) && !isnothing(u_init) + xu[(offset + nx + 1):(offset + nx + m)] .= u_init + end +end + + +# trivial version for now... +# +++multiple dispatch here seems to cause more allocations ! +# +++? use abstract type for all Args ? +""" +$(TYPEDSIGNATURES) + +Useful values at a time step: time, state, control, dynamics... +""" +struct ArgsAtTimeStep_Midpoint + time::Any + state::Any + control::Any + lagrange_state::Any + stage_k::Any + next_time::Any + next_state::Any + next_lagrange_state::Any + + function ArgsAtTimeStep_Midpoint(xu, docp::DOCP{MidpointTag}, v, time_grid, i::Int) + + tag = docp.discretization + + # variables + ti = time_grid[i+1] + xi, ui, xli, ki = get_variables_at_time_step(xu, docp, i) + + if i == docp.dim_NLP_steps + return new(ti, xi, ui, xli, ki, tag) + else + tip1 = time_grid[i+2] + xip1, uip1, xlip1 = get_variables_at_time_step(xu, docp, i+1) + return new(ti, xi, ui, xli, ki, tip1, xip1, xlip1) + end + end +end +function initArgs(xu, docp::DOCP{MidpointTag}, time_grid) + v = Float64[] + docp.has_variable && (v = get_optim_variable(xu, docp)) + args = ArgsAtTimeStep_Midpoint(xu, docp, v, time_grid, 0) + return args, v +end +function updateArgs(args, xu, docp::DOCP{MidpointTag}, v, time_grid, i::Int) + return ArgsAtTimeStep_Midpoint(xu, docp, v, time_grid, i+1) +end + + +""" +$(TYPEDSIGNATURES) + +Set the constraints corresponding to the state equation +""" +function setStateEquation!(docp::DOCP{MidpointTag}, c, index::Int, args, v, i) + + ocp = docp.ocp + + # +++ later use butcher table in struct ? + + # variables + ti = args.time + xi = args.state + ui = args.control + xli = args.lagrange_state + ki = args.stage_k + tip1 = args.next_time + xip1 = args.next_state + xlip1 = args.next_lagrange_state + hi = tip1 - ti + + # midpoint rule + @. c[index:(index + docp.dim_OCP_x - 1)] = + xip1 - (xi + hi * ki[1:docp.dim_OCP_x]) + # +++ just define extended dynamics ! + if docp.has_lagrange + c[index + docp.dim_OCP_x] = xlip1 - (xli + hi * ki[end]) + end + index += docp.dim_NLP_x + + # stage equation at mid-step + t_s = 0.5 * (ti + tip1) + x_s = 0.5 * (xi + xip1) + c[index:(index + docp.dim_OCP_x - 1)] .= + ki[1:docp.dim_OCP_x] .- ocp.dynamics(t_s, x_s, ui, v) + # +++ just define extended dynamics ! + if docp.has_lagrange + c[index + docp.dim_OCP_x] = ki[end] - ocp.lagrange(t_s, x_s, ui, v) + end + index += docp.dim_NLP_x + + return index +end + + +""" +$(TYPEDSIGNATURES) + +Set the path constraints at given time step +""" +function setPathConstraints!(docp::DOCP{MidpointTag}, c, index::Int, args, v, i::Int) + + ocp = docp.ocp + ti = args.time + xi = args.state + ui = args.control + + # NB. using .= below *doubles* the allocations oO ?? + # pure control constraints + if docp.dim_u_cons > 0 + c[index:(index + docp.dim_u_cons - 1)] = docp.control_constraints[2](ti, ui, v) + index += docp.dim_u_cons + end + + # pure state constraints + if docp.dim_x_cons > 0 + c[index:(index + docp.dim_x_cons - 1)] = docp.state_constraints[2](ti, xi, v) + index += docp.dim_x_cons + end + + # mixed state / control constraints + if docp.dim_mixed_cons > 0 + c[index:(index + docp.dim_mixed_cons - 1)] = docp.mixed_constraints[2](ti, xi, ui, v) + index += docp.dim_mixed_cons + end + + return index +end diff --git a/src/problem.jl b/src/problem.jl index 18588e44..ea54e91c 100644 --- a/src/problem.jl +++ b/src/problem.jl @@ -1,5 +1,14 @@ -# Internal layout for NLP variables: -# [X0,U0, X1,U1, .., XN,UN,V] +# todo: add more discretization schemes +# use a scheme struct to allow multiple dispatch (cf solve) +# this struct can contains useful info about the scheme +# (name, order, stage, properties etc) +# later we can add an option for control discretization: step or stage +# (meaningful only for schemes with more than 1 stage, at least I will be able to compare the two ! further options may include CVP -control vector parametrization-, and maybe even pseudo-spectral ?) + + +# generic discretization struct +# NB. can we mutualize common fields at the abstract level ? +abstract type DiscretizationTag end """ $(TYPEDSIGNATURES) @@ -8,10 +17,9 @@ Struct for discretized optimal control problem DOCP Contains: - a copy of the original OCP -- a NLP formulation of the DOCP -- data required to link the two problems +- data required to link the OCP with the discretized DOCP """ -struct DOCP +struct DOCP{DiscretizationTag} ## OCP ocp::OptimalControlModel @@ -58,8 +66,11 @@ struct DOCP con_l::Vector{Float64} con_u::Vector{Float64} + # discretization scheme + discretization::DiscretizationTag + # constructor - function DOCP(ocp::OptimalControlModel, grid_size::Integer, time_grid) + function DOCP(ocp::OptimalControlModel, grid_size::Integer, time_grid, discretization::DiscretizationTag) # time grid if time_grid == nothing @@ -76,9 +87,10 @@ struct DOCP #println("INFO: normalizing given time grid...") t0 = time_grid[1] tf = time_grid[end] - time_grid = (time_grid .- t0) ./ (tf - t0) + NLP_normalized_time_grid = (time_grid .- t0) ./ (tf - t0) + else + NLP_normalized_time_grid = time_grid end - NLP_normalized_time_grid = time_grid dim_NLP_steps = length(time_grid) - 1 end @@ -105,9 +117,10 @@ struct DOCP dim_OCP_x = ocp.state_dimension N = dim_NLP_steps + dim_stage = discretization.stage - # NLP unknown (state + control + variable) - dim_NLP_variables = (N + 1) * (dim_NLP_x + dim_NLP_u) + dim_NLP_v + # NLP unknown (state + control + variable [+ stage]) + dim_NLP_variables = (N + 1) * dim_NLP_x + (N + discretization.additional_controls) * dim_NLP_u + dim_NLP_v + N * dim_NLP_x * dim_stage # NLP constraints # parse NLP constraints (and initialize dimensions) @@ -130,15 +143,16 @@ struct DOCP dim_mixed_cons = dim_mixed_constraints(ocp) dim_boundary_cons = dim_boundary_constraints(ocp) - # lagrange to mayer transformation + # constraints (dynamics, stage, path, boundary, variable) dim_NLP_constraints = - N * (dim_NLP_x + dim_path_cons) + dim_path_cons + dim_boundary_cons + dim_v_cons + N * (dim_NLP_x + (dim_NLP_x * dim_stage) + dim_path_cons) + dim_path_cons + dim_boundary_cons + dim_v_cons if has_lagrange + # add initial condition for lagrange state dim_NLP_constraints += 1 end # call constructor with const fields - docp = new( + docp = new{typeof(discretization)}( ocp, control_constraints, state_constraints, @@ -175,12 +189,14 @@ struct DOCP Inf * ones(dim_NLP_variables), zeros(dim_NLP_constraints), zeros(dim_NLP_constraints), + discretization ) return docp end end + """ $(TYPEDSIGNATURES) @@ -191,12 +207,14 @@ function is_solvable(ocp) return solvable end + """ $(TYPEDSIGNATURES) Build upper and lower bounds vectors for the DOCP nonlinear constraints. """ function constraints_bounds!(docp::DOCP) + lb = docp.con_l ub = docp.con_u @@ -204,6 +222,8 @@ function constraints_bounds!(docp::DOCP) for i = 0:(docp.dim_NLP_steps - 1) # skip (ie leave 0) for equality dynamics constraint index = index + docp.dim_NLP_x + # skip (ie leave 0) for equality stage constraint (ki) + index = index + docp.dim_NLP_x * docp.discretization.stage # path constraints index = setPathBounds!(docp, index, lb, ub) end @@ -217,6 +237,7 @@ function constraints_bounds!(docp::DOCP) return lb, ub end + """ $(TYPEDSIGNATURES) @@ -227,69 +248,55 @@ function variables_bounds!(docp::DOCP) var_l = docp.var_l var_u = docp.var_u ocp = docp.ocp + tag = docp.discretization - # NB. keep offset for each block since they are optional ! + # first we build full ordered sets of bounds, then set them in NLP - # build ordered bounds vectors for state and control - x_lb = -Inf * ones(docp.dim_OCP_x) - x_ub = Inf * ones(docp.dim_OCP_x) - for j = 1:(docp.dim_x_box) - indice = docp.state_box[2][j] - x_lb[indice] = docp.state_box[1][j] - x_ub[indice] = docp.state_box[3][j] - end - u_lb = -Inf * ones(docp.dim_NLP_u) - u_ub = Inf * ones(docp.dim_NLP_u) - for j = 1:(docp.dim_u_box) - indice = docp.control_box[2][j] - u_lb[indice] = docp.control_box[1][j] - u_ub[indice] = docp.control_box[3][j] - end - - # apply bounds for NLP variables + # state / control box + x_lb, x_ub = build_bounds(docp.dim_OCP_x, docp.dim_x_box, docp.state_box) + u_lb, u_ub = build_bounds(docp.dim_NLP_u, docp.dim_u_box, docp.control_box) for i = 0:N - set_variables_at_time_step!(var_l, x_lb, u_lb, docp, i) - set_variables_at_time_step!(var_u, x_ub, u_ub, docp, i) + set_variables_at_time_step!(var_l, x_lb, u_lb, docp, i, tag) + set_variables_at_time_step!(var_u, x_ub, u_ub, docp, i, tag) end # variable box - offset = (N + 1) * (docp.dim_NLP_x + docp.dim_NLP_u) - if docp.dim_v_box > 0 - for j = 1:(docp.dim_v_box) - indice = docp.variable_box[2][j] - var_l[offset + indice] = docp.variable_box[1][j] - var_u[offset + indice] = docp.variable_box[3][j] - end + if docp.has_variable + v_lb, v_ub = build_bounds(docp.dim_NLP_v, docp.dim_v_box, docp.variable_box) + set_optim_variable!(var_l, v_lb, docp) + set_optim_variable!(var_u, v_ub, docp) end return var_l, var_u end + """ $(TYPEDSIGNATURES) Compute the objective for the DOCP problem. """ -# DOCP objective function DOCP_objective(xu, docp::DOCP) + obj = 0.0 N = docp.dim_NLP_steps ocp = docp.ocp + # optimization variables + v = Float64[] + docp.has_variable && (v = get_optim_variable(xu, docp)) + + # final state is always needed since lagrange cost is there + xf, uf, xlf = get_variables_at_time_step(xu, docp, N) + # mayer cost if docp.has_mayer - v = get_variable(xu, docp) - t0 = get_initial_time(xu, docp) - tf = get_final_time(xu, docp) x0, u0, xl0 = get_variables_at_time_step(xu, docp, 0) - xf, uf, xlf = get_variables_at_time_step(xu, docp, N) - #obj = obj + ocp.mayer(t0, tf, x0, xf, v) obj = obj + ocp.mayer(x0, xf, v) end # lagrange cost if docp.has_lagrange - xf, uf, xlf = get_variables_at_time_step(xu, docp, N) obj = obj + xlf end @@ -301,154 +308,70 @@ function DOCP_objective(xu, docp::DOCP) return obj end + """ $(TYPEDSIGNATURES) Compute the constraints C for the DOCP problem (modeled as LB <= C(X) <= UB). """ function DOCP_constraints!(c, xu, docp::DOCP) - """ - compute the constraints for the NLP : - - discretization of the dynamics via the trapeze method - - boundary conditions - inputs - ocp :: ocp model - xu :: - return - c :: - """ - - # could use a single v, however v is set only in Args constructor, not update, so not much is wasted - v = get_variable(xu, docp) - # main loop on time steps - args_i = ArgsAtTimeStep(xu, docp, 0, v) - args_ip1 = ArgsAtTimeStep(xu, docp, 1, v) + # +++ todo: help AD by avoid passing the whole xu to inner functions + # + # args = initArgs(xu) + # for i=1:N + # setStateEquation!(docp, c, index, args[i]) + # setPathConstraints!(docp, c, index, args[i]) + # setPathConstraints!(docp, c, index, args[N+1]) + # setPointConstraints!(docp, xu, index) + # + # where args is a N size array of discretization-dependent + # tuples (or structs if we need the mutable part ?), + # with each one containing every scalar and vector needed + # to evaluate both setStateEquation and setPathConstraints + # + # initArgs will avoid unnecessary recomputations + # get_time_grid (for the t_i) + # get_optim_variables (for v) + # for i=0:N-1 + # getters for xi, ui, ki, xi+1 (some redundancy here) + # compute needed values and save in args[i] + # + # note: maybe reuse values from one iteration to the other + # (eg x_i+1 in general, f_i+1 for trapeze) + # but this would require a mutable struct (tuple are non mutable) + # the struct could be tailored to each discretization method + # test both for instance on midpoint: mutable vs non mutable + + # initialization + # +++ use and pass a single tuple (args, v, time_grid) instead ? + time_grid = get_time_grid(xu, docp) + args, v = initArgs(xu, docp, time_grid) + # main loop on time steps index = 1 # counter for the constraints for i = 0:(docp.dim_NLP_steps - 1) # state equation - index = setStateEquation!(docp, c, index, (args_i, args_ip1)) + index = setStateEquation!(docp, c, index, args, v, i) # path constraints - index = setPathConstraints!(docp, c, index, args_i, v) + index = setPathConstraints!(docp, c, index, args, v, i) + # update + args = updateArgs(args, xu, docp, v, time_grid, i) - # smart update for next iteration - if i < docp.dim_NLP_steps - 1 - args_i = args_ip1 - args_ip1 = ArgsAtTimeStep(xu, docp, i + 2, v) - end end # path constraints at final time - args_0 = ArgsAtTimeStep(xu, docp, 0, v) - args_f = ArgsAtTimeStep(xu, docp, docp.dim_NLP_steps, v) - index = setPathConstraints!(docp, c, index, args_f, v) + index = setPathConstraints!(docp, c, index, args, v, docp.dim_NLP_steps) # boundary conditions and variable constraints - index = setPointConstraints!(docp, c, index, args_0, args_f, v) + index = setPointConstraints!(docp, c, index, xu, v) # needed even for inplace version, AD error otherwise # may be because actual return would be index above ? return c end -# +++ later use abstract interface, based on Args variants ? -# this struct and the related functions will change according to discretization scheme -# put this part in trapeze.jl file -""" -$(TYPEDSIGNATURES) - -Useful values at a time step: time, state, control, dynamics... -""" -struct ArgsAtTimeStep - time::Any - state::Any - control::Any - dynamics::Any - lagrange_state::Any - lagrange_cost::Any - - function ArgsAtTimeStep(xu, docp::DOCP, i::Int, v) - - # variables - ti = get_time_at_time_step(xu, docp, i) - #xi = get_state_at_time_step(xu, docp, i) - #ui = get_control_at_time_step(xu, docp, i) - xi, ui, xli = get_variables_at_time_step(xu, docp, i) - - # dynamics and lagrange cost - fi = docp.ocp.dynamics(ti, xi, ui, v) - - if docp.has_lagrange - #xli = get_lagrange_cost_at_time_step(xu, docp, i) - li = docp.ocp.lagrange(ti, xi, ui, v) - args = new(ti, xi, ui, fi, xli, li) - else - args = new(ti, xi, ui, fi) - end - - return args - end -end - -""" -$(TYPEDSIGNATURES) - -Set the constraints corresponding to the state equation -""" -function setStateEquation!(docp::DOCP, c, index::Int, args_trapeze) - ocp = docp.ocp - args_i = args_trapeze[1] - args_ip1 = args_trapeze[2] - hi = args_ip1.time - args_i.time - - # trapeze rule - c[index:(index + docp.dim_OCP_x - 1)] .= - args_ip1.state .- (args_i.state .+ 0.5 * hi * (args_i.dynamics .+ args_ip1.dynamics)) - - if docp.has_lagrange - c[index + docp.dim_OCP_x] = - args_ip1.lagrange_state - - (args_i.lagrange_state + 0.5 * hi * (args_i.lagrange_cost + args_ip1.lagrange_cost)) - end - - index = index + docp.dim_NLP_x - return index -end - -""" -$(TYPEDSIGNATURES) - -Set the path constraints at given time step -""" -function setPathConstraints!(docp::DOCP, c, index::Int, args::ArgsAtTimeStep, v) - ocp = docp.ocp - - ti = args.time - xi = args.state - ui = args.control - - # pure control constraints - if docp.dim_u_cons > 0 - c[index:(index + docp.dim_u_cons - 1)] = docp.control_constraints[2](ti, ui, v) - index = index + docp.dim_u_cons - end - - # pure state constraints - if docp.dim_x_cons > 0 - c[index:(index + docp.dim_x_cons - 1)] = docp.state_constraints[2](ti, xi, v) - index = index + docp.dim_x_cons - end - - # mixed state / control constraints - if docp.dim_mixed_cons > 0 - c[index:(index + docp.dim_mixed_cons - 1)] = docp.mixed_constraints[2](ti, xi, ui, v) - index = index + docp.dim_mixed_cons - end - - return index -end """ $(TYPEDSIGNATURES) @@ -456,6 +379,7 @@ $(TYPEDSIGNATURES) Set bounds for the path constraints at given time step """ function setPathBounds!(docp::DOCP, index::Int, lb, ub) + ocp = docp.ocp # pure control constraints @@ -482,23 +406,18 @@ function setPathBounds!(docp::DOCP, index::Int, lb, ub) return index end + """ $(TYPEDSIGNATURES) Set the boundary and variable constraints """ -function setPointConstraints!( - docp::DOCP, - c, - index::Int, - args_0::ArgsAtTimeStep, - args_f::ArgsAtTimeStep, - v, -) +function setPointConstraints!(docp::DOCP, c, index::Int, xu, v) + ocp = docp.ocp - x0 = args_0.state - xf = args_f.state + x0, u0, xl0 = get_variables_at_time_step(xu, docp, 0) + xf, = get_variables_at_time_step(xu, docp, docp.dim_NLP_steps) # boundary constraints if docp.dim_boundary_cons > 0 @@ -514,19 +433,21 @@ function setPointConstraints!( # null initial condition for lagrangian cost state if docp.has_lagrange - c[index] = args_0.lagrange_state + c[index] = xl0 index = index + 1 end return index end + """ $(TYPEDSIGNATURES) Set bounds for the boundary and variable constraints """ function setPointBounds!(docp::DOCP, index::Int, lb, ub) + ocp = docp.ocp # boundary constraints @@ -553,6 +474,51 @@ function setPointBounds!(docp::DOCP, index::Int, lb, ub) return index end + +""" +$(TYPEDSIGNATURES) + +Build initial guess for discretized problem +""" +function DOCP_initial_guess(docp::DOCP, init::OptimalControlInit = OptimalControlInit()) + + # default initialization (internal variables such as lagrange cost, k_i for RK schemes) will keep these default values + NLP_X = 0.1 * ones(docp.dim_NLP_variables) + + # set variables if provided (needed first in case of free times !) + if !isnothing(init.variable_init) + set_optim_variable!(NLP_X, init.variable_init, docp) + end + + # set state / control variables if provided + time_grid = get_time_grid(NLP_X, docp) + for i = 0:(docp.dim_NLP_steps) + ti = time_grid[i+1] + set_variables_at_time_step!(NLP_X, init.state_init(ti), init.control_init(ti), docp, i, docp.discretization) + end + + return NLP_X +end + +#= OLD + + # recompute value of constraints at solution + # NB. the constraint formulation is LB <= C <= UB + constraints = zeros(docp.dim_NLP_constraints) + DOCP_constraints!(constraints, solution, docp) + # set constraint violation if needed + # is not saved in OCP solution currently... + if constraints_violation==nothing + constraints_check = zeros(docp.dim_NLP_constraints) + DOCP_constraints_check!(constraints_check, constraints, docp) + println("Recomputed constraints violation ", norm(constraints_check, Inf)) + variables_check = zeros(docp.dim_NLP_variables) + DOCP_variables_check!(variables_check, solution, docp) + println("Recomputed variable bounds violation ", norm(variables_check, Inf)) + constraints_violation = norm(append!(variables_check, constraints_check), Inf) + + end + """ $(TYPEDSIGNATURES) @@ -560,7 +526,7 @@ Check the nonlinear constraints violation for the DOCP problem. """ function DOCP_constraints_check!(cb, constraints, docp) - # +++ todo add a single utils function check_bounds(v,lb,ub) that returns the error vector + # todo add a single utils function check_bounds(v,lb,ub) that returns the error vector ? # check constraints vs bounds # by construction only one of the two can be active @@ -593,3 +559,4 @@ function DOCP_variables_check!(vb, variables, docp) end return nothing end +=# diff --git a/src/solution.jl b/src/solution.jl index 06f76f0b..8ce71eac 100644 --- a/src/solution.jl +++ b/src/solution.jl @@ -46,11 +46,7 @@ function CTBase.OptimalControlSolution( ) # time grid - N = docp.dim_NLP_steps - T = zeros(N + 1) - for i = 1:(N + 1) - T[i] = get_unnormalized_time(primal, docp, docp.NLP_normalized_time_grid[i]) - end + T = get_time_grid(primal, docp) # recover primal variables X, U, v, box_multipliers = @@ -100,9 +96,11 @@ Recover OCP primal variables from DOCP solution function parse_DOCP_solution_primal(docp, solution; mult_LB = nothing, mult_UB = nothing) # state and control variables + tag = docp.discretization N = docp.dim_NLP_steps X = zeros(N + 1, docp.dim_NLP_x) U = zeros(N + 1, docp.dim_NLP_u) + v = Float64[] # multipliers for box constraints if isnothing(mult_LB) || length(mult_LB) == 0 @@ -119,20 +117,22 @@ function parse_DOCP_solution_primal(docp, solution; mult_LB = nothing, mult_UB = mult_variable_box_upper = zeros(N + 1, docp.dim_NLP_v) # retrieve optimization variables - v = get_variable(solution, docp) - mult_variable_box_lower = get_variable(mult_LB, docp) - mult_variable_box_upper = get_variable(mult_UB, docp) + if docp.has_variable + v = get_optim_variable(solution, docp) + mult_variable_box_lower = get_optim_variable(mult_LB, docp) + mult_variable_box_upper = get_optim_variable(mult_UB, docp) + end # loop over time steps for i = 1:(N + 1) # state and control variables at current step - X[i, :], U[i, :] = get_NLP_variables_at_time_step(solution, docp, i - 1) + X[i, :], U[i, :] = get_NLP_variables_at_time_step(solution, docp, i - 1, tag) # box multipliers mult_state_box_lower[i, :], mult_control_box_lower[i, :] = - get_NLP_variables_at_time_step(mult_LB, docp, i - 1) + get_NLP_variables_at_time_step(mult_LB, docp, i - 1, tag) mult_state_box_upper[i, :], mult_control_box_upper[i, :] = - get_NLP_variables_at_time_step(mult_UB, docp, i - 1) + get_NLP_variables_at_time_step(mult_UB, docp, i - 1, tag) end box_multipliers = ( @@ -187,8 +187,12 @@ function parse_DOCP_solution_dual(docp, multipliers, constraints) # state equation multiplier for costate (except last step) if i < N + 1 P[i, :] = multipliers[i_m:(i_m + docp.dim_NLP_x - 1)] - i_c += docp.dim_NLP_x # skip dynamics constraints + # skip dynamics constraints + i_c += docp.dim_NLP_x i_m += docp.dim_NLP_x + # skip stage constraints + i_c += docp.dim_NLP_x * docp.discretization.stage + i_m += docp.dim_NLP_x * docp.discretization.stage end # path constraints and multipliers @@ -451,7 +455,7 @@ $(TYPEDSIGNATURES) Process data related to a box type for solution building """ -# +++ integrate above +# +++ integrate above ? function set_box_block(T, mults, dim) mult_l, mult_u = mults if !isnothing(mult_l) && !isnothing(mult_u) @@ -461,22 +465,3 @@ function set_box_block(T, mults, dim) return t -> m_l(t), t -> m_u(t) end -#= OLD - - # recompute value of constraints at solution - # NB. the constraint formulation is LB <= C <= UB - constraints = zeros(docp.dim_NLP_constraints) - DOCP_constraints!(constraints, solution, docp) - # set constraint violation if needed - # is not saved in OCP solution currently... - if constraints_violation==nothing - constraints_check = zeros(docp.dim_NLP_constraints) - DOCP_constraints_check!(constraints_check, constraints, docp) - println("Recomputed constraints violation ", norm(constraints_check, Inf)) - variables_check = zeros(docp.dim_NLP_variables) - DOCP_variables_check!(variables_check, solution, docp) - println("Recomputed variable bounds violation ", norm(variables_check, Inf)) - constraints_violation = norm(append!(variables_check, constraints_check), Inf) - - end -=# diff --git a/src/solve.jl b/src/solve.jl index 5ff9355d..7cd458c3 100644 --- a/src/solve.jl +++ b/src/solve.jl @@ -24,10 +24,21 @@ function direct_transcription( init = CTBase.__ocp_init(), grid_size = __grid_size(), time_grid = __time_grid(), + discretization = __discretization() ) # build DOCP - docp = DOCP(ocp, grid_size, time_grid) + discretization = string(discretization) + if discretization == "midpoint" + disc_tag = CTDirect.MidpointTag() + elseif discretization == "trapeze" + disc_tag = CTDirect.TrapezeTag() + elseif discretization == "gausslegendre2" + disc_tag = CTDirect.GaussLegendre2Tag() + else + error("Unknown discretization method:", discretization) + end + docp = DOCP(ocp, grid_size, time_grid, disc_tag) # set bounds in DOCP variables_bounds!(docp) @@ -88,6 +99,7 @@ function direct_solve( init = CTBase.__ocp_init(), grid_size::Int = CTDirect.__grid_size(), time_grid = CTDirect.__time_grid(), + discretization = __discretization(), kwargs..., ) method = getFullDescription(description, available_methods()) @@ -100,6 +112,7 @@ function direct_solve( init = init, grid_size = grid_size, time_grid = time_grid, + discretization = discretization ) # solve DOCP diff --git a/src/trapeze.jl b/src/trapeze.jl new file mode 100644 index 00000000..ed95deb5 --- /dev/null +++ b/src/trapeze.jl @@ -0,0 +1,202 @@ +#= Functions for trapeze discretization scheme +Internal layout for NLP variables: +[X_0,U_0, X_1,U_1, .., X_N,U_N, V] +=# + + +struct TrapezeTag <: DiscretizationTag + stage::Int + additional_controls::Int + # add control at tf + TrapezeTag() = new(0, 1) +end + + +""" +$(TYPEDSIGNATURES) + +Retrieve state and control variables at given time step from the NLP variables. +""" +function get_variables_at_time_step(xu, docp::DOCP{TrapezeTag}, i) + + nx = docp.dim_NLP_x + n = docp.dim_OCP_x + m = docp.dim_NLP_u + offset = (nx + m) * i + + # retrieve scalar/vector OCP state (w/o lagrange state) + if n == 1 + xi = xu[offset + 1] + else + xi = xu[(offset + 1):(offset + n)] + end + if docp.has_lagrange + xli = xu[offset + nx] + else + xli = nothing # dummy. use xu type ? + end + + # retrieve scalar/vector control + if m == 1 + ui = xu[offset + nx + 1] + else + ui = xu[(offset + nx + 1):(offset + nx + m)] + end + + return xi, ui, xli +end + + +# internal NLP version for solution parsing +# could be fused with one above if +# - using extended dynamics that include lagrange cost +# - scalar case is handled at OCP level +function get_NLP_variables_at_time_step(xu, docp, i, tag::TrapezeTag) + + nx = docp.dim_NLP_x + m = docp.dim_NLP_u + offset = (nx + m) * i + + # state + xi = xu[(offset + 1):(offset + nx)] + # control + ui = xu[(offset + nx + 1):(offset + nx + m)] + + return xi, ui +end + + +function set_variables_at_time_step!(xu, x_init, u_init, docp, i, tag::TrapezeTag) + + nx = docp.dim_NLP_x + n = docp.dim_OCP_x + m = docp.dim_NLP_u + N = docp.dim_NLP_steps + offset = (nx + m) * i + + # NB. only set the actual state variables from the OCP + # - skip the possible additional state for lagrange cost + if !isnothing(x_init) + xu[(offset + 1):(offset + n)] .= x_init + end + if !isnothing(u_init) + xu[(offset + nx + 1):(offset + nx + m)] .= u_init + end +end + + +# ? use abstract type for args ? +""" +$(TYPEDSIGNATURES) + +Useful values at a time step: time, state, control, dynamics... +""" +struct ArgsAtTimeStep_Trapeze + time::Any + state::Any + control::Any + dynamics::Any + lagrange_state::Any + lagrange_cost::Any + + function ArgsAtTimeStep_Trapeze(xu, docp::DOCP{TrapezeTag}, v, time_grid, i::Int) + + # variables + ti = time_grid[i+1] + xi, ui, xli = get_variables_at_time_step(xu, docp, i) + + # dynamics and lagrange cost + fi = docp.ocp.dynamics(ti, xi, ui, v) + + if docp.has_lagrange + li = docp.ocp.lagrange(ti, xi, ui, v) + args = new(ti, xi, ui, fi, xli, li) + else + args = new(ti, xi, ui, fi) + end + + return args + end +end +# +++multiple dispatch here seems to cause more allocations ! +function initArgs(xu, docp::DOCP{TrapezeTag}, time_grid) + # optimization variables + v = Float64[] + docp.has_variable && (v = get_optim_variable(xu, docp)) + args_i = ArgsAtTimeStep_Trapeze(xu, docp, v, time_grid, 0) + args_ip1 = ArgsAtTimeStep_Trapeze(xu, docp, v, time_grid, 1) + return (args_i, args_ip1), v +end +function updateArgs(args, xu, docp::DOCP{TrapezeTag}, v, time_grid, i) + args_i, args_ip1 = args + if i < docp.dim_NLP_steps - 1 + # are we allocating more than one args here ? + return (args_ip1, ArgsAtTimeStep_Trapeze(xu, docp, v, time_grid, i+2)) + else + return (args_ip1, args_ip1) + end +end + + +""" +$(TYPEDSIGNATURES) + +Set the constraints corresponding to the state equation +""" +function setStateEquation!(docp::DOCP{TrapezeTag}, c, index::Int, args, v, i) + + # NB. arguments v,i are unused here but present for unified call + ocp = docp.ocp + args_i, args_ip1 = args + hi = args_ip1.time - args_i.time + + # trapeze rule (NB. @. allocates more ...) + c[index:(index + docp.dim_OCP_x - 1)] .= + args_ip1.state .- (args_i.state .+ 0.5 * hi * (args_i.dynamics .+ args_ip1.dynamics)) + # +++ just define extended dynamics ! + if docp.has_lagrange + c[index + docp.dim_OCP_x] = + args_ip1.lagrange_state - + (args_i.lagrange_state + 0.5 * hi * (args_i.lagrange_cost + args_ip1.lagrange_cost)) + end + index += docp.dim_NLP_x + + return index +end + + +""" +$(TYPEDSIGNATURES) + +Set the path constraints at given time step +""" +function setPathConstraints!(docp::DOCP{TrapezeTag}, c, index::Int, args, v, i::Int) + + # note: i is unused but passed for call compatibility + ocp = docp.ocp + args_i, args_ip1 = args + ti = args_i.time + xi = args_i.state + ui = args_i.control + + # NB. using .= below *doubles* the allocations oO + # pure control constraints + if docp.dim_u_cons > 0 + c[index:(index + docp.dim_u_cons - 1)] = docp.control_constraints[2](ti, ui, v) + index += docp.dim_u_cons + end + + # pure state constraints + if docp.dim_x_cons > 0 + c[index:(index + docp.dim_x_cons - 1)] = docp.state_constraints[2](ti, xi, v) + index += docp.dim_x_cons + end + + # mixed state / control constraints + if docp.dim_mixed_cons > 0 + c[index:(index + docp.dim_mixed_cons - 1)] = docp.mixed_constraints[2](ti, xi, ui, v) + index += docp.dim_mixed_cons + end + + return index +end diff --git a/src/utils.jl b/src/utils.jl index be5aa3bb..e2844322 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1,146 +1,29 @@ +#= Functions that are not dependent on the discretization scheme +Internal layout for NLP variables: +- Trapeze: [X_0,U_0, X_1,U_1, .., X_N,U_N, V] +- Midpoint: [X_0,U_0,K_0, X_1,U_1,K_1 .., X_N-1,U_N-1,K_N-1, XN, V] +with the convention u([t_i,t_i+1[) = U_i and u(tf) = U_N-1 +=# + + """ $(TYPEDSIGNATURES) Retrieve optimization variables from the NLP variables. -Internal layout: [X0,U0, X1,U1, .., XN,UN,V] """ -function get_variable(xu, docp) - if docp.has_variable - nx = docp.dim_NLP_x - m = docp.dim_NLP_u - N = docp.dim_NLP_steps - offset = (nx + m) * (N + 1) +function get_optim_variable(xu, docp) + if docp.has_variable if docp.dim_NLP_v == 1 - return xu[offset + 1] - #return xu[end] + return xu[end] else - #return xu[end-docp.dim_NLP_v+1:end] - return xu[(offset + 1):(offset + docp.dim_NLP_v)] + return xu[(end - docp.dim_NLP_v + 1):end] end else - return Float64[] - end -end - -""" -$(TYPEDSIGNATURES) - -Retrieve a single optimization variable (no dim check). -Internal layout: [X0,U0, X1,U1, .., XN,UN,V] -""" -function get_single_variable(xu, docp, i::Int) - if docp.has_variable - #return xu[end-docp.dim_NLP_v+i] - nx = docp.dim_NLP_x - m = docp.dim_NLP_u - N = docp.dim_NLP_steps - offset = (nx + m) * (N + 1) - return xu[offset + i] - else - error("Tring to access variable in variable independent problem") - end -end - -""" -$(TYPEDSIGNATURES) - -Retrieve state and control variables at given time step from the NLP variables. -Internal layout: [X0,U0, X1,U1, .., XN,UN,V] -""" -# +++ this one would depend on the discretization scheme -function get_variables_at_time_step(xu, docp, i) - nx = docp.dim_NLP_x - n = docp.dim_OCP_x - m = docp.dim_NLP_u - offset = (nx + m) * i - - # retrieve scalar/vector OCP state (w/o lagrange state) - if n == 1 - xi = xu[offset + 1] - else - xi = xu[(offset + 1):(offset + n)] - end - # NB. meaningful ONLY if has_lagrange is true ! - # add check without performance cost ? - if docp.has_lagrange - xli = xu[offset + nx] - else - xli = 0.0 - end - - # retrieve scalar/vector control - if m == 1 - ui = xu[offset + nx + 1] - else - ui = xu[(offset + nx + 1):(offset + nx + m)] - end - - return xi, ui, xli -end - -# internal NLP version for solution parsing -function get_NLP_variables_at_time_step(xu, docp, i) - nx = docp.dim_NLP_x - m = docp.dim_NLP_u - offset = (nx + m) * i - - xi = xu[(offset + 1):(offset + nx)] - ui = xu[(offset + nx + 1):(offset + nx + m)] - - return xi, ui -end - -#= -""" -$(TYPEDSIGNATURES) - -Retrieve scalar/vector state variables at given time step from the NLP variables -""" -function get_state_at_time_step(xu, docp, i) - nx = docp.dim_NLP_x - n = docp.dim_OCP_x - if n == 1 - return xu[i*nx + 1] - else - return xu[i*nx + 1 : i*nx + n] + error("Problem is not variable dependent") end end -""" -$(TYPEDSIGNATURES) - -Retrieve the additional state variable corresponding to the lagrange (running) cost at given time step from the NLP variables -""" -function get_lagrange_cost_at_time_step(xu, docp, i) - nx = docp.dim_NLP_x - return xu[(i+1)*nx] -end - -# internal NLP version for solution parsing -function get_NLP_state_at_time_step(xu, docp, i) - nx = docp.dim_NLP_x - return xu[i*nx + 1 : (i+1)*nx] -end - -""" -$(TYPEDSIGNATURES) - -Retrieve scalar/vector control variables at given time step from the NLP variables -""" -function get_control_at_time_step(xu, docp, i) - nx = docp.dim_NLP_x - m = docp.dim_NLP_u - N = docp.dim_NLP_steps - offset = (N+1)*nx - - if docp.dim_NLP_u == 1 - return xu[offset + i*m + 1] - else - return xu[offset + i*m + 1 : offset + (i+1)*m] - end -end -=# """ $(TYPEDSIGNATURES) @@ -149,12 +32,13 @@ Retrieve initial time for OCP (may be fixed or variable) """ function get_initial_time(xu, docp) if docp.has_free_t0 - return get_single_variable(xu, docp, Base.to_index(docp.ocp.initial_time)) + return get_optim_variable(xu, docp)[docp.ocp.initial_time] else return docp.ocp.initial_time end end + """ $(TYPEDSIGNATURES) @@ -162,110 +46,54 @@ Retrieve final time for OCP (may be fixed or variable) """ function get_final_time(xu, docp) if docp.has_free_tf - return get_single_variable(xu, docp, Base.to_index(docp.ocp.final_time)) + return get_optim_variable(xu, docp)[docp.ocp.final_time] else return docp.ocp.final_time end end + """ $(TYPEDSIGNATURES) -Get actual (un-normalized) time value +Get full (un-normalized) time grid """ -function get_unnormalized_time(xu, docp, t_normalized) +function get_time_grid(xu, docp) t0 = get_initial_time(xu, docp) tf = get_final_time(xu, docp) - return t0 + t_normalized * (tf - t0) -end - -""" -$(TYPEDSIGNATURES) - -Get actual (un-normalized) time at give time step -""" -function get_time_at_time_step(xu, docp, i) - return get_unnormalized_time(xu, docp, docp.NLP_normalized_time_grid[i + 1]) + return @. t0 + docp.NLP_normalized_time_grid * (tf - t0) end -#= -""" -$(TYPEDSIGNATURES) - -Set state variables at given time step in the NLP variables (for initial guess) -""" -function set_state_at_time_step!(xu, x_init, docp, i) - nx = docp.dim_NLP_x - n = docp.dim_OCP_x - # NB. only set first the actual state variables from the OCP (not the possible additional state for lagrange cost) - xu[i*nx + 1 : i*nx + n] .= x_init -end - -""" -$(TYPEDSIGNATURES) - -Set control variables at given time step in the NLP variables (for initial guess) -""" -function set_control_at_time_step!(xu, u_init, docp, i) - nx = docp.dim_NLP_x - m = docp.dim_NLP_u - N = docp.dim_NLP_steps - offset = (N+1)*nx - xu[offset + i*m + 1 : offset + i*m + m] .= u_init -end -=# - -function set_variables_at_time_step!(xu, x_init, u_init, docp, i) - nx = docp.dim_NLP_x - n = docp.dim_OCP_x - m = docp.dim_NLP_u - N = docp.dim_NLP_steps - offset = (nx + m) * i - - # NB. only set first the actual state variables from the OCP (not the possible additional state for lagrange cost) - if !isnothing(x_init) - xu[(offset + 1):(offset + n)] .= x_init - end - if !isnothing(u_init) - xu[(offset + nx + 1):(offset + nx + m)] .= u_init - end -end """ $(TYPEDSIGNATURES) Set optimization variables in the NLP variables (for initial guess) """ -function set_variable!(xu, v_init, docp) +function set_optim_variable!(xu, v_init, docp) xu[(end - docp.dim_NLP_v + 1):end] .= v_init end + """ $(TYPEDSIGNATURES) -Build initial guess for discretized problem +Build full, ordered sets of bounds for state, control or optimization variables """ -function DOCP_initial_guess(docp, init::OptimalControlInit = OptimalControlInit()) - - # default initialization - # note: internal variables (lagrange cost, k_i for RK schemes) will keep these default values - NLP_X = 0.1 * ones(docp.dim_NLP_variables) +function build_bounds(dim_var, dim_box, box_triplet) - # set variables if provided - # (needed first in case of free times !) - if !isnothing(init.variable_init) - set_variable!(NLP_X, init.variable_init, docp) + x_lb = -Inf * ones(dim_var) + x_ub = Inf * ones(dim_var) + for j = 1:(dim_box) + indice = box_triplet[2][j] + x_lb[indice] = box_triplet[1][j] + x_ub[indice] = box_triplet[3][j] end - # set state / control variables if provided - for i = 0:(docp.dim_NLP_steps) - ti = get_time_at_time_step(NLP_X, docp, i) - set_variables_at_time_step!(NLP_X, init.state_init(ti), init.control_init(ti), docp, i) - end - - return NLP_X + return x_lb, x_ub end + # placeholders (see CTDirectExt) function export_ocp_solution end function import_ocp_solution end diff --git a/test/Project.toml b/test/Project.toml index a3c049f0..a3fcce23 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -6,6 +6,7 @@ JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" NLPModelsIpopt = "f4238b75-b362-5c4c-b852-0801c9a21d71" +SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/test/deps.jl b/test/deps.jl index bea978a1..225ce119 100644 --- a/test/deps.jl +++ b/test/deps.jl @@ -5,6 +5,7 @@ using LinearAlgebra using NLPModelsIpopt using MadNLP using HSL +using SplitApplyCombine using JLD2 using JSON3 diff --git a/test/problems/simple_integrator.jl b/test/problems/simple_integrator.jl index 75271969..25160215 100644 --- a/test/problems/simple_integrator.jl +++ b/test/problems/simple_integrator.jl @@ -12,5 +12,5 @@ function simple_integrator() dynamics!(ocp, (x, u) -> -x - u[1] + u[2]) objective!(ocp, :lagrange, (x, u) -> (u[1] + u[2])^2) - return ((ocp = ocp, obj = nothing, name = "simple_integrator", init = nothing)) + return ((ocp = ocp, obj = 3.13e-1, name = "simple_integrator", init = nothing)) end diff --git a/test/suite/test_grid.jl b/test/suite/test_discretization.jl similarity index 63% rename from test/suite/test_grid.jl rename to test/suite/test_discretization.jl index 77454aa2..24961a2d 100644 --- a/test/suite/test_grid.jl +++ b/test/suite/test_discretization.jl @@ -1,4 +1,7 @@ -println("Test: grid options") +println("Test: discretization options") + + +normalize_grid(t) = return (t .- t[1]) ./ (t[end] - t[1]) # 1. simple integrator min energy (dual control for test) if !isdefined(Main, :simple_integrator) @@ -17,7 +20,7 @@ end @testset verbose = true showtiming = true ":non_uniform_grid" begin time_grid = [0, 0.1, 0.3, 0.6, 0.98, 0.99, 1] sol = direct_solve(ocp, time_grid = time_grid, display = false) - @test sol.objective ≈ 0.309 rtol = 1e-2 + @test sol.time_grid ≈ time_grid end # 2. integrator free times @@ -33,8 +36,9 @@ sol0 = direct_solve(ocp, display = false) end @testset verbose = true showtiming = true ":max_t0 :non_uniform_grid" begin - sol = direct_solve(ocp, time_grid = [0, 0.1, 0.6, 0.95, 1], display = false) - @test sol.objective ≈ 7.48 rtol = 1e-2 + time_grid = [0, 0.1, 0.6, 0.95, 1] + sol = direct_solve(ocp, time_grid = time_grid, display = false) + @test normalize_grid(sol.time_grid) ≈ time_grid end # 3. parametric ocp (T=2) with explicit / non-uniform grid @@ -50,6 +54,23 @@ sol0 = direct_solve(ocp, display = false) end @testset verbose = true showtiming = true ":non_uniform_grid" begin - sol = direct_solve(ocp, time_grid = [0, 0.3, 1, 1.9, 2], display = false) - @test sol.objective ≈ 2.43 rtol = 1e-2 + time_grid = [0, 0.3, 1, 1.9, 2] + sol = direct_solve(ocp, time_grid = time_grid, display = false) + @test sol.time_grid ≈ time_grid end + +# implicit midpoint scheme +@testset verbose = true showtiming = true ":implicit_midpoint" begin + ocp = simple_integrator().ocp + sol_t = direct_solve(ocp, display = false) + sol_m = direct_solve(ocp, display = false, discretization = "midpoint") + @test sol_m.objective ≈ sol_t.objective rtol = 1e-2 +end + +@testset verbose = true showtiming = true ":implicit_midpoint" begin + ocp = double_integrator_freet0tf().ocp + sol_t = direct_solve(ocp, display = false) + sol_m = direct_solve(ocp, display = false, discretization = :midpoint) + @test sol_m.objective ≈ sol_t.objective rtol = 1e-2 +end + diff --git a/test/suite/test_initial_guess.jl b/test/suite/test_initial_guess.jl index a9dd1549..9e73f94e 100644 --- a/test/suite/test_initial_guess.jl +++ b/test/suite/test_initial_guess.jl @@ -3,17 +3,6 @@ println("Test: initial guess options") # use 0 iterations to check initial guess, >0 to check cv maxiter = 0 -# test functions -function check_xf(sol, xf) - return xf == sol.state(sol.time_grid[end]) -end -function check_uf(sol, uf) - return uf == sol.control(sol.time_grid[end]) -end -function check_v(sol, v) - return v == sol.variable -end - # reference solution prob = double_integrator_a() ocp = prob.ocp @@ -39,29 +28,40 @@ u_vec = [0, 0.3, 0.1] # 1.a default initial guess @testset verbose = true showtiming = true ":default_init_no_arg" begin sol = direct_solve(ocp, display = false, max_iter = maxiter) - @test(check_xf(sol, [0.1, 0.1]) && check_uf(sol, 0.1) && check_v(sol, 0.1)) + T = sol.time_grid + @test isapprox(sol.state.(T), (t->[0.1,0.1]).(T), rtol = 1e-2) + @test isapprox(sol.control.(T), (t->0.1).(T), rtol = 1e-2) + @test isapprox(sol.variable, 0.1, rtol = 1e-2) end @testset verbose = true showtiming = true ":default_init_()" begin sol = direct_solve(ocp, display = false, init = (), max_iter = maxiter) - @test(check_xf(sol, [0.1, 0.1]) && check_uf(sol, 0.1) && check_v(sol, 0.1)) + T = sol.time_grid + @test isapprox(sol.state.(T), (t->[0.1,0.1]).(T), rtol = 1e-2) + @test isapprox(sol.control.(T), (t->0.1).(T), rtol = 1e-2) + @test isapprox(sol.variable, 0.1, rtol = 1e-2) end @testset verbose = true showtiming = true ":default_init_nothing" begin sol = direct_solve(ocp, display = false, init = nothing, max_iter = maxiter) - @test(check_xf(sol, [0.1, 0.1]) && check_uf(sol, 0.1) && check_v(sol, 0.1)) + T = sol.time_grid + @test isapprox(sol.state.(T), (t->[0.1,0.1]).(T), rtol = 1e-2) + @test isapprox(sol.control.(T), (t->0.1).(T), rtol = 1e-2) + @test sol.variable == 0.1 end # 1.b constant initial guess @testset verbose = true showtiming = true ":constant_x" begin sol = direct_solve(ocp, display = false, init = (state = x_const,), max_iter = maxiter) - @test(check_xf(sol, x_const)) + T = sol.time_grid + @test isapprox(sol.state.(T), (t->x_const).(T), rtol = 1e-2) end @testset verbose = true showtiming = true ":constant_u" begin sol = direct_solve(ocp, display = false, init = (control = u_const,), max_iter = maxiter) - @test(check_uf(sol, u_const)) + T = sol.time_grid + @test isapprox(sol.control.(T), (t->u_const).(T), rtol = 1e-2) end @testset verbose = true showtiming = true ":constant_v" begin sol = direct_solve(ocp, display = false, init = (variable = v_const,), max_iter = maxiter) - @test(check_v(sol, v_const)) + @test sol.variable == v_const end @testset verbose = true showtiming = true ":constant_xu" begin sol = direct_solve( @@ -70,7 +70,9 @@ end init = (state = x_const, control = u_const), max_iter = maxiter, ) - @test(check_xf(sol, x_const) && check_uf(sol, u_const)) + T = sol.time_grid + @test isapprox(sol.state.(T), (t->x_const).(T), rtol = 1e-2) + @test isapprox(sol.control.(T), (t->u_const).(T), rtol = 1e-2) end @testset verbose = true showtiming = true ":constant_xv" begin sol = direct_solve( @@ -79,7 +81,9 @@ end init = (state = x_const, variable = v_const), max_iter = maxiter, ) - @test(check_xf(sol, x_const) && check_v(sol, v_const)) + T = sol.time_grid + @test isapprox(sol.state.(T), (t->x_const).(T), rtol = 1e-2) + @test sol.variable == v_const end @testset verbose = true showtiming = true ":constant_uv" begin sol = direct_solve( @@ -88,7 +92,9 @@ end init = (control = u_const, variable = v_const), max_iter = maxiter, ) - @test(check_uf(sol, u_const) && check_v(sol, v_const)) + T = sol.time_grid + @test isapprox(sol.control.(T), (t->u_const).(T), rtol = 1e-2) + @test sol.variable == v_const end @testset verbose = true showtiming = true ":constant_xuv" begin sol = direct_solve( @@ -97,17 +103,22 @@ end init = (state = x_const, control = u_const, variable = v_const), max_iter = maxiter, ) - @test(check_xf(sol, x_const) && check_uf(sol, u_const) && check_v(sol, v_const)) + T = sol.time_grid + @test isapprox(sol.state.(T), (t->x_const).(T), rtol = 1e-2) + @test isapprox(sol.control.(T), (t->u_const).(T), rtol = 1e-2) + @test sol.variable == v_const end # 1. functional initial guess @testset verbose = true showtiming = true ":functional_x" begin sol = direct_solve(ocp, display = false, init = (state = x_func,), max_iter = maxiter) - @test(check_xf(sol, x_func(sol.time_grid[end]))) + T = sol.time_grid + @test isapprox(sol.state.(T), x_func.(T), rtol = 1e-2) end @testset verbose = true showtiming = true ":functional_u" begin sol = direct_solve(ocp, display = false, init = (control = u_func,), max_iter = maxiter) - @test(check_uf(sol, u_func(sol.time_grid[end]))) + T = sol.time_grid + @test isapprox(sol.control.(T), u_func.(T), rtol = 1e-2) end @testset verbose = true showtiming = true ":functional_xu" begin sol = direct_solve( @@ -116,7 +127,9 @@ end init = (state = x_func, control = u_func), max_iter = maxiter, ) - @test(check_xf(sol, x_func(sol.time_grid[end])) && check_uf(sol, u_func(sol.time_grid[end]))) + T = sol.time_grid + @test isapprox(sol.state.(T), x_func.(T), rtol = 1e-2) + @test isapprox(sol.control.(T), u_func.(T), rtol = 1e-2) end # 1.d interpolated initial guess @@ -128,7 +141,9 @@ t_vec = [0, 0.1, v_const] init = (time = t_vec, state = x_vec, control = u_vec, variable = v_const), max_iter = maxiter, ) - @test(check_xf(sol, x_vec[end]) && check_uf(sol, u_vec[end]) && check_v(sol, v_const)) + @test isapprox(sol.state.(t_vec), x_vec, rtol = 1e-2) + @test isapprox(sol.control.(t_vec), u_vec, rtol = 1e-2) + @test sol.variable == v_const end t_matrix = [0 0.1 v_const] @testset verbose = true showtiming = true ":matrix_t :vector_xu :constant_v" begin @@ -138,7 +153,9 @@ t_matrix = [0 0.1 v_const] init = (time = t_matrix, state = x_vec, control = u_vec, variable = v_const), max_iter = maxiter, ) - @test(check_xf(sol, x_vec[end]) && check_uf(sol, u_vec[end]) && check_v(sol, v_const)) + @test isapprox(sol.state.(flatten(t_matrix)), x_vec, rtol = 1e-2) + @test isapprox(sol.control.(flatten(t_matrix)), u_vec, rtol = 1e-2) + @test sol.variable == v_const end @testset verbose = true showtiming = true ":matrix_x :vector_tu :constant_v" begin sol = direct_solve( @@ -147,7 +164,9 @@ end init = (time = t_vec, state = x_matrix, control = u_vec, variable = v_const), max_iter = maxiter, ) - @test(check_xf(sol, x_vec[end]) && check_uf(sol, u_vec[end]) && check_v(sol, v_const)) + @test isapprox(stack(sol.state.(t_matrix),dims=1), x_matrix, rtol = 1e-2) + @test isapprox(sol.control.(t_vec), u_vec, rtol = 1e-2) + @test sol.variable == v_const end # 1.e mixed initial guess @@ -158,21 +177,19 @@ end init = (time = t_vec, state = x_vec, control = u_func, variable = v_const), max_iter = maxiter, ) - @test( - check_xf(sol, x_vec[end]) && - check_uf(sol, u_func(sol.time_grid[end])) && - check_v(sol, v_const) - ) + T = sol.time_grid + @test isapprox(sol.state.(t_vec), x_vec, rtol = 1e-2) + @test isapprox(sol.control.(T), u_func.(T), rtol = 1e-2) + @test sol.variable == v_const end # 1.f warm start @testset verbose = true showtiming = true ":warm_start" begin sol = direct_solve(ocp, display = false, init = sol0, max_iter = maxiter) - @test( - check_xf(sol, sol.state(sol.time_grid[end])) && - check_uf(sol, sol.control(sol.time_grid[end])) && - check_v(sol, sol.variable) - ) + T = sol.time_grid + @test isapprox(sol.state.(T), sol0.state.(T), rtol = 1e-2) + @test isapprox(sol.control.(T), sol0.control.(T), rtol = 1e-2) + @test sol.variable == sol0.variable end ################################################# @@ -188,20 +205,18 @@ tag = CTDirect.IpoptTag() ) dsol = CTDirect.solve_docp(tag, docp, nlp, display = false, max_iter = maxiter) sol = OptimalControlSolution(docp, dsol) - @test( - check_xf(sol, x_vec[end]) && - check_uf(sol, u_func(sol.time_grid[end])) && - check_v(sol, v_const) - ) + T = sol.time_grid + @test isapprox(sol.state.(t_vec), x_vec, rtol = 1e-2) + @test isapprox(sol.control.(T), u_func.(T), rtol = 1e-2) + @test sol.variable == v_const end # warm start @testset verbose = true showtiming = true ":docp_warm_start" begin set_initial_guess(docp, nlp, sol0) dsol = CTDirect.solve_docp(tag, docp, nlp, display = false, max_iter = maxiter) sol = OptimalControlSolution(docp, dsol) - @test( - check_xf(sol, sol.state(sol.time_grid[end])) && - check_uf(sol, sol.control(sol.time_grid[end])) && - check_v(sol, sol.variable) - ) + T = sol.time_grid + @test isapprox(sol.state.(T), sol0.state.(T), rtol = 1e-2) + @test isapprox(sol.control.(T), sol0.control.(T), rtol = 1e-2) + @test sol.variable == sol0.variable end diff --git a/test/suite/test_misc.jl b/test/suite/test_misc.jl index 4a97cb6f..7d473156 100644 --- a/test/suite/test_misc.jl +++ b/test/suite/test_misc.jl @@ -1,13 +1,13 @@ println("Test: misc") -@testset verbose = true ":default :direct" begin +@testset verbose = true ":default_direct" begin @test CTDirect.__grid_size() isa Integer @test isnothing(CTDirect.__time_grid()) @test CTDirect.__tolerance() < 1 @test CTDirect.__max_iterations() isa Integer end -@testset verbose = true ":default :ipopt" begin +@testset verbose = true ":default_ipopt" begin @test CTDirect.__ipopt_print_level() isa Integer @test CTDirect.__ipopt_print_level() ≤ 12 @test CTDirect.__ipopt_print_level() ≥ 0 diff --git a/test/suite/test_nlp.jl b/test/suite/test_nlp.jl index 884aa806..f42c1440 100644 --- a/test/suite/test_nlp.jl +++ b/test/suite/test_nlp.jl @@ -11,7 +11,7 @@ ocp = simple_integrator().ocp @test (:adnlp, :madnlp) in available_methods() end -@testset verbose = true showtiming = true ":solve_docp :ipopt" begin +@testset verbose = true showtiming = true ":solve_docp" begin docp, nlp = direct_transcription(ocp) tag = CTDirect.IpoptTag() dsol = CTDirect.solve_docp(tag, docp, nlp, display = false) @@ -23,6 +23,18 @@ end @test sol.objective ≈ 0.313 rtol = 1e-2 end +@testset verbose = true showtiming = true ":solve_docp :midpoint" begin + docp, nlp = direct_transcription(ocp, discretization = :midpoint) + tag = CTDirect.IpoptTag() + dsol = CTDirect.solve_docp(tag, docp, nlp, display = false) + sol = OptimalControlSolution(docp, dsol) + @test sol.objective ≈ 0.313 rtol = 1e-2 + sol = OptimalControlSolution(docp, primal = dsol.solution) + @test sol.objective ≈ 0.313 rtol = 1e-2 + sol = OptimalControlSolution(docp, primal = dsol.solution, dual = dsol.multipliers) + @test sol.objective ≈ 0.313 rtol = 1e-2 +end + @testset verbose = true showtiming = true ":solve_docp :madnlp" begin docp, nlp = direct_transcription(ocp) tag = CTDirect.MadNLPTag()