Skip to content

Commit

Permalink
fix merge
Browse files Browse the repository at this point in the history
  • Loading branch information
ocots committed Aug 25, 2024
1 parent c1a00f9 commit 1b43237
Show file tree
Hide file tree
Showing 7 changed files with 225 additions and 88 deletions.
11 changes: 7 additions & 4 deletions ext/CTDirectExt.jl
Original file line number Diff line number Diff line change
Expand Up @@ -32,8 +32,11 @@ $(TYPEDSIGNATURES)
Export OCP solution in JSON format
"""
function CTDirect.export_ocp_solution(sol::OptimalControlSolution; filename_prefix="solution")
# +++ redo this, start with basics
function CTDirect.export_ocp_solution(
sol::OptimalControlSolution;
filename_prefix = "solution",
)
# +++ redo this, start with basics
open(filename_prefix * ".json", "w") do io
#JSON3.pretty(io, CTDirect.OCPDiscreteSolution(sol))
end
Expand All @@ -45,8 +48,8 @@ $(TYPEDSIGNATURES)
Read OCP solution in JSON format
"""
function CTDirect.import_ocp_solution(filename_prefix="solution")
# +++ add constructor from json blob
function CTDirect.import_ocp_solution(filename_prefix = "solution")
# +++ add constructor from json blob
json_string = read(filename_prefix * ".json", String)
return OptimalControlSolution(JSON3.read(json_string))
end
Expand Down
146 changes: 117 additions & 29 deletions src/solution.jl
Original file line number Diff line number Diff line change
Expand Up @@ -262,11 +262,23 @@ $(TYPEDSIGNATURES)
Build OCP functional solution from DOCP vector solution (given as raw variables and multipliers plus some optional infos)
"""
function CTBase.OptimalControlSolution(docp, 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)))
function CTBase.OptimalControlSolution(
docp,
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)),
)

ocp = docp.ocp
dim_x = state_dimension(ocp)
Expand Down Expand Up @@ -300,29 +312,86 @@ function CTBase.OptimalControlSolution(docp, T, X, U, v, P;

# +++ 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,
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)
# 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,
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)

# build and return solution
if docp.has_variable
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)
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)
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
Expand Down Expand Up @@ -351,7 +420,18 @@ function set_constraints_and_multipliers(T, constraints_types, constraints_mult)
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)
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


Expand All @@ -365,11 +445,19 @@ 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)
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)
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

"""
Expand All @@ -381,10 +469,10 @@ Process data related to a box type for solution building
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))
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)
return t -> m_l(t), t -> m_u(t)
end


Expand Down
8 changes: 7 additions & 1 deletion src/utils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -270,7 +270,13 @@ function DOCP_initial_guess(docp, init::OptimalControlInit = OptimalControlInit(
# 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)
set_variables_at_time_step!(
NLP_X,
init.state_init(ti),
init.control_init(ti),
docp,
i,
)
end

return NLP_X
Expand Down
81 changes: 46 additions & 35 deletions test/suite/test_constraints.jl
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,12 @@ if check_constraint_mult
# variable constraints / box
println("\nVariable constraints: ", sol.variable_constraints)
println("multipliers: ", sol.mult_variable_constraints)
println("\nVariable box multipliers LB: ", sol.mult_variable_box_lower, " UB ", sol.mult_variable_box_upper)
println(
"\nVariable box multipliers LB: ",
sol.mult_variable_box_lower,
" UB ",
sol.mult_variable_box_upper,
)

# PATH CONSTRAINTS
# state box
Expand All @@ -55,9 +60,9 @@ if check_constraint_mult
u = sol.control.(T)
u_box_lb = flatten(sol.mult_control_box_lower.(T))
u_box_ub = flatten(sol.mult_control_box_upper.(T))
p4_a = plot(T, u, label="u")
p4_b = plot(T, [u_box_lb u_box_ub], label=["LB" "UB"])
p4 = plot(p4_a, p4_b, layout=(2,1))
p4_a = plot(T, u, label = "u")
p4_b = plot(T, [u_box_lb u_box_ub], label = ["LB" "UB"])
p4 = plot(p4_a, p4_b, layout = (2, 1))

p_box = plot(
p1,
Expand All @@ -74,47 +79,53 @@ if check_constraint_mult
# control constraints
c_u = flatten(sol.control_constraints.(T))
m_c_u = flatten(sol.mult_control_constraints.(T))
p5_a = plot(T, c_u, label="c_u")
p5_b = plot(T, m_c_u, label="mul")
p5 = plot(p5_a, p5_b, layout=(2,1))
p5_a = plot(T, c_u, label = "c_u")
p5_b = plot(T, m_c_u, label = "mul")
p5 = plot(p5_a, p5_b, layout = (2, 1))

# state constraints
c_x = flatten(sol.state_constraints.(T))
m_c_x = flatten(sol.mult_state_constraints.(T))
p6_a = plot(T, c_x, label="c_x")
p6_b = plot(T, m_c_x, label="mul")
p6 = plot(p6_a, p6_b, layout=(2,1))
p6_a = plot(T, c_x, label = "c_x")
p6_b = plot(T, m_c_x, label = "mul")
p6 = plot(p6_a, p6_b, layout = (2, 1))

# mixed constraints
c_xu = flatten(sol.mixed_constraints.(T))
m_c_xu = flatten(sol.mult_mixed_constraints.(T))
p7_a = plot(T, c_xu, label="c_xu")
p7_b = plot(T, m_c_xu, label="mul")
p7 = plot(p7_a, p7_b, layout=(2,1))

p_cons = plot(p5, p6, p7, layout=(1,3), title=["control cons" "" "state cons" "" "mixed cons" ""])
p7_a = plot(T, c_xu, label = "c_xu")
p7_b = plot(T, m_c_xu, label = "mul")
p7 = plot(p7_a, p7_b, layout = (2, 1))

p_cons = plot(
p5,
p6,
p7,
layout = (1, 3),
title = ["control cons" "" "state cons" "" "mixed cons" ""],
)

else

# box constraints
@testset verbose = true showtiming = true ":goddard :box_constraints" begin
ocp = goddard()
sol = direct_solve(ocp.ocp, display=false)
@test sol.objective ocp.obj rtol=1e-2
end

# functional constraints
@testset verbose = true showtiming = true ":goddard :functional_constraints" begin
ocp = goddard(functional_constraints=true)
sol = direct_solve(ocp.ocp, display=false, init=ocp.init)
@test sol.objective ocp.obj rtol=1e-2
end
# box constraints
@testset verbose = true showtiming = true ":goddard :box_constraints" begin
ocp = goddard()
sol = direct_solve(ocp.ocp, display = false)
@test sol.objective ocp.obj rtol = 1e-2
end

# functional constraints
@testset verbose = true showtiming = true ":goddard :functional_constraints" begin
ocp = goddard(functional_constraints = true)
sol = direct_solve(ocp.ocp, display = false, init = ocp.init)
@test sol.objective ocp.obj rtol = 1e-2
end

# all constraints
@testset verbose = true showtiming = true ":goddard :all_constraints" begin
ocp = goddard_all()
sol = direct_solve(ocp.ocp, display = false, init = ocp.init)
@test sol.objective ocp.obj rtol = 1e-2
end

# all constraints
@testset verbose = true showtiming = true ":goddard :all_constraints" begin
ocp = goddard_all()
sol = direct_solve(ocp.ocp, display=false, init=ocp.init)
@test sol.objective ocp.obj rtol=1e-2
end

end
Loading

0 comments on commit 1b43237

Please sign in to comment.