Skip to content

Commit

Permalink
fixed 2 of the 3 type piracy in solution.jl, the remaining constructo…
Browse files Browse the repository at this point in the history
…r could be moved to ctbase
  • Loading branch information
PierreMartinon committed Oct 17, 2024
1 parent 6f1fbe0 commit afe8e68
Show file tree
Hide file tree
Showing 3 changed files with 170 additions and 105 deletions.
3 changes: 2 additions & 1 deletion ext/CTSolveExtMadNLP.jl
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,8 @@ function CTDirect.solve_docp(
return docp_solution
end

function CTBase.OptimalControlSolution(docp, docp_solution::MadNLP.MadNLPExecutionStats)

function CTBase.OptimalControlSolution(docp::CTDirect.DOCP, docp_solution::MadNLP.MadNLPExecutionStats)

# adjust objective sign for maximization problems
if is_min(docp.ocp)
Expand Down
143 changes: 142 additions & 1 deletion src/ctbase.jl
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
################################################################
################################################################
# A integrer dans CTBase vu que c'est la sauvegarde d'une solution OCP
# NB. copie de CTDirectExt.jl qui definit l'extension avec les weakdeps JLD2 / JSON3
# copie de CTDirectExt.jl qui definit l'extension avec les weakdeps JLD2 / JSON3

module CTDirectExt

Expand Down Expand Up @@ -66,3 +68,142 @@ function CTDirect.import_ocp_solution(ocp::OptimalControlModel; filename_prefix
end

end



################################################################
################################################################
# cf solution.jl: a mettre dans CTBase ? (aqua indique un type piracy car on ne passe pas d'argument avec des types specifiques a DOCP). Ce constructeur pourrait etre utilise pour toute solution sous forme discrete, pas seulement les methodes directes.
"""
$(TYPEDSIGNATURES)
Build OCP functional solution from DOCP vector solution (given as raw variables and multipliers plus some optional infos)
"""
function CTBase.OptimalControlSolution(
ocp::OptimalControlModel,
T,
X,
U,
v,
P;
objective = 0,
iterations = 0,
constraints_violation = 0,
message = "No msg",
stopping = nothing,
success = nothing,
constraints_types = (nothing, nothing, nothing, nothing, nothing),
constraints_mult = (nothing, nothing, nothing, nothing, nothing),
box_multipliers = (nothing, nothing, nothing, nothing, nothing, nothing),
)
dim_x = state_dimension(ocp)
dim_u = control_dimension(ocp)
dim_v = variable_dimension(ocp)

# check that time grid is strictly increasing
# if not proceed with list of indexes as time grid
if !issorted(T, lt = <=)
println(
"WARNING: time grid at solution is not strictly increasing, replacing with list of indices...",
)
println(T)
dim_NLP_steps = length(T) - 1
T = LinRange(0, dim_NLP_steps, dim_NLP_steps + 1)
end

# variables: remove additional state for lagrange cost
x = ctinterpolate(T, matrix2vec(X[:, 1:dim_x], 1))
p = ctinterpolate(T[1:(end - 1)], matrix2vec(P[:, 1:dim_x], 1))
u = ctinterpolate(T, matrix2vec(U[:, 1:dim_u], 1))

# force scalar output when dimension is 1
fx = (dim_x == 1) ? deepcopy(t -> x(t)[1]) : deepcopy(t -> x(t))
fu = (dim_u == 1) ? deepcopy(t -> u(t)[1]) : deepcopy(t -> u(t))
fp = (dim_x == 1) ? deepcopy(t -> p(t)[1]) : deepcopy(t -> p(t))
var = (dim_v == 1) ? v[1] : v

# misc infos
infos = Dict{Symbol, Any}()
infos[:constraints_violation] = constraints_violation

# nonlinear constraints and multipliers
control_constraints = t -> ctinterpolate(T, matrix2vec(constraints_types[1], 1))(t)
mult_control_constraints = t -> ctinterpolate(T, matrix2vec(constraints_mult[1], 1))(t)
state_constraints = t -> ctinterpolate(T, matrix2vec(constraints_types[2], 1))(t)
mult_state_constraints = t -> ctinterpolate(T, matrix2vec(constraints_mult[2], 1))(t)
mixed_constraints = t -> ctinterpolate(T, matrix2vec(constraints_types[3], 1))(t)
mult_mixed_constraints = t -> ctinterpolate(T, matrix2vec(constraints_mult[3], 1))(t)

# boundary and variable constraints
boundary_constraints = constraints_types[4]
mult_boundary_constraints = constraints_mult[4]
variable_constraints = constraints_types[5]
mult_variable_constraints = constraints_mult[5]

# box constraints multipliers
mult_state_box_lower = t -> ctinterpolate(T, matrix2vec(box_multipliers[1][:, 1:dim_x], 1))(t)
mult_state_box_upper = t -> ctinterpolate(T, matrix2vec(box_multipliers[2][:, 1:dim_x], 1))
mult_control_box_lower = t -> ctinterpolate(T, matrix2vec(box_multipliers[3][:, 1:dim_u], 1))(t)
mult_control_box_upper = t -> ctinterpolate(T, matrix2vec(box_multipliers[4][:, 1:dim_u], 1))
mult_variable_box_lower, mult_variable_box_upper = box_multipliers[5], box_multipliers[6]

# build and return solution
if is_variable_dependent(ocp)
return OptimalControlSolution(
ocp;
state = fx,
control = fu,
objective = objective,
costate = fp,
time_grid = T,
variable = var,
iterations = iterations,
stopping = stopping,
message = message,
success = success,
infos = infos,
control_constraints = control_constraints,
state_constraints = state_constraints,
mixed_constraints = mixed_constraints,
boundary_constraints = boundary_constraints,
variable_constraints = variable_constraints,
mult_control_constraints = mult_control_constraints,
mult_state_constraints = mult_state_constraints,
mult_mixed_constraints = mult_mixed_constraints,
mult_boundary_constraints = mult_boundary_constraints,
mult_variable_constraints = mult_variable_constraints,
mult_state_box_lower = mult_state_box_lower,
mult_state_box_upper = mult_state_box_upper,
mult_control_box_lower = mult_control_box_lower,
mult_control_box_upper = mult_control_box_upper,
mult_variable_box_lower = mult_variable_box_lower,
mult_variable_box_upper = mult_variable_box_upper,
)
else
return OptimalControlSolution(
ocp;
state = fx,
control = fu,
objective = objective,
costate = fp,
time_grid = T,
iterations = iterations,
stopping = stopping,
message = message,
success = success,
infos = infos,
control_constraints = control_constraints,
state_constraints = state_constraints,
mixed_constraints = mixed_constraints,
boundary_constraints = boundary_constraints,
mult_control_constraints = mult_control_constraints,
mult_state_constraints = mult_state_constraints,
mult_mixed_constraints = mult_mixed_constraints,
mult_boundary_constraints = mult_boundary_constraints,
mult_state_box_lower = mult_state_box_lower,
mult_state_box_upper = mult_state_box_upper,
mult_control_box_lower = mult_control_box_lower,
mult_control_box_upper = mult_control_box_upper,
)
end
end
129 changes: 26 additions & 103 deletions src/solution.jl
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ $(TYPEDSIGNATURES)
Build OCP functional solution from DOCP discrete solution (given as a SolverCore.GenericExecutionStats)
"""
function CTBase.OptimalControlSolution(docp, docp_solution)
function CTBase.OptimalControlSolution(docp::DOCP, docp_solution)

# retrieve data (could pass some status info too (get_status ?))
if docp.is_maximization
Expand Down Expand Up @@ -34,7 +34,7 @@ $(TYPEDSIGNATURES)
Build OCP functional solution from the DOCP discrete solution, given as a vector. Costate will be retrieved from dual variables (multipliers) if available.
"""
function CTBase.OptimalControlSolution(
docp;
docp::DOCP;
primal = Vector(),
dual = nothing,
objective = nothing,
Expand Down Expand Up @@ -138,9 +138,9 @@ function parse_DOCP_solution_primal(docp, solution; mult_LB = nothing, mult_UB =
end

box_multipliers = (
(mult_state_box_lower, mult_state_box_upper),
(mult_control_box_lower, mult_control_box_upper),
(mult_variable_box_lower, mult_variable_box_upper),
mult_state_box_lower, mult_state_box_upper,
mult_control_box_lower, mult_control_box_upper,
mult_variable_box_lower, mult_variable_box_upper
)

return X, U, v, box_multipliers
Expand Down Expand Up @@ -254,6 +254,8 @@ function parse_DOCP_solution_dual(docp, multipliers, constraints)
return P, constraints_types, constraints_mult
end


# +++ move this one to CTBase ? (aqua flags a type piracy since we don't have arguments DOCP specific types)
"""
$(TYPEDSIGNATURES)
Expand All @@ -274,7 +276,7 @@ function CTBase.OptimalControlSolution(
success = nothing,
constraints_types = (nothing, nothing, nothing, nothing, nothing),
constraints_mult = (nothing, nothing, nothing, nothing, nothing),
box_multipliers = ((nothing, nothing), (nothing, nothing), (nothing, nothing)),
box_multipliers = (nothing, nothing, nothing, nothing, nothing, nothing),
)
dim_x = state_dimension(ocp)
dim_u = control_dimension(ocp)
Expand Down Expand Up @@ -306,29 +308,26 @@ function CTBase.OptimalControlSolution(
infos = Dict{Symbol, Any}()
infos[:constraints_violation] = constraints_violation

# +++ put interpolations here directly and reuse vectors ?
# nonlinear constraints and multipliers
(
control_constraints,
state_constraints,
mixed_constraints,
boundary_constraints,
variable_constraints,
mult_control_constraints,
mult_state_constraints,
mult_mixed_constraints,
mult_boundary_constraints,
mult_variable_constraints,
) = set_constraints_and_multipliers(T, constraints_types, constraints_mult)
control_constraints = t -> ctinterpolate(T, matrix2vec(constraints_types[1], 1))(t)
mult_control_constraints = t -> ctinterpolate(T, matrix2vec(constraints_mult[1], 1))(t)
state_constraints = t -> ctinterpolate(T, matrix2vec(constraints_types[2], 1))(t)
mult_state_constraints = t -> ctinterpolate(T, matrix2vec(constraints_mult[2], 1))(t)
mixed_constraints = t -> ctinterpolate(T, matrix2vec(constraints_types[3], 1))(t)
mult_mixed_constraints = t -> ctinterpolate(T, matrix2vec(constraints_mult[3], 1))(t)

# boundary and variable constraints
boundary_constraints = constraints_types[4]
mult_boundary_constraints = constraints_mult[4]
variable_constraints = constraints_types[5]
mult_variable_constraints = constraints_mult[5]

# box constraints multipliers
(
mult_state_box_lower,
mult_state_box_upper,
mult_control_box_lower,
mult_control_box_upper,
mult_variable_box_lower,
mult_variable_box_upper,
) = set_box_multipliers(T, box_multipliers, dim_x, dim_u)
mult_state_box_lower = t -> ctinterpolate(T, matrix2vec(box_multipliers[1][:, 1:dim_x], 1))(t)
mult_state_box_upper = t -> ctinterpolate(T, matrix2vec(box_multipliers[2][:, 1:dim_x], 1))
mult_control_box_lower = t -> ctinterpolate(T, matrix2vec(box_multipliers[3][:, 1:dim_u], 1))(t)
mult_control_box_upper = t -> ctinterpolate(T, matrix2vec(box_multipliers[4][:, 1:dim_u], 1))
mult_variable_box_lower, mult_variable_box_upper = box_multipliers[5], box_multipliers[6]

# build and return solution
if is_variable_dependent(ocp)
Expand Down Expand Up @@ -390,79 +389,3 @@ function CTBase.OptimalControlSolution(
)
end
end

"""
$(TYPEDSIGNATURES)
Process data related to constraints for solution building
"""
function set_constraints_and_multipliers(T, constraints_types, constraints_mult)

# control, state, mixed constraints
control_constraints = t -> ctinterpolate(T, matrix2vec(constraints_types[1], 1))(t)
mult_control_constraints = t -> ctinterpolate(T, matrix2vec(constraints_mult[1], 1))(t)

state_constraints = t -> ctinterpolate(T, matrix2vec(constraints_types[2], 1))(t)
mult_state_constraints = t -> ctinterpolate(T, matrix2vec(constraints_mult[2], 1))(t)

mixed_constraints = t -> ctinterpolate(T, matrix2vec(constraints_types[3], 1))(t)
mult_mixed_constraints = t -> ctinterpolate(T, matrix2vec(constraints_mult[3], 1))(t)

# boundary and variable constraints
boundary_constraints = constraints_types[4]
mult_boundary_constraints = constraints_mult[4]
variable_constraints = constraints_types[5]
mult_variable_constraints = constraints_mult[5]

return (
control_constraints,
state_constraints,
mixed_constraints,
boundary_constraints,
variable_constraints,
mult_control_constraints,
mult_state_constraints,
mult_mixed_constraints,
mult_boundary_constraints,
mult_variable_constraints,
)
end

"""
$(TYPEDSIGNATURES)
Process data related to box constraints for solution building
"""
function set_box_multipliers(T, box_multipliers, dim_x, dim_u)

# state box
mult_state_box_lower, mult_state_box_upper = set_box_block(T, box_multipliers[1], dim_x)
# control box
mult_control_box_lower, mult_control_box_upper = set_box_block(T, box_multipliers[2], dim_u)
# variable box
mult_variable_box_lower, mult_variable_box_upper = box_multipliers[3]

return (
mult_state_box_lower,
mult_state_box_upper,
mult_control_box_lower,
mult_control_box_upper,
mult_variable_box_lower,
mult_variable_box_upper,
)
end

"""
$(TYPEDSIGNATURES)
Process data related to a box type for solution building
"""
# +++ integrate above ?
function set_box_block(T, mults, dim)
mult_l, mult_u = mults
if !isnothing(mult_l) && !isnothing(mult_u)
m_l = ctinterpolate(T, matrix2vec(mult_l[:, 1:dim], 1))
m_u = ctinterpolate(T, matrix2vec(mult_u[:, 1:dim], 1))
end
return t -> m_l(t), t -> m_u(t)
end

0 comments on commit afe8e68

Please sign in to comment.