Skip to content

Commit

Permalink
add solution_label
Browse files Browse the repository at this point in the history
  • Loading branch information
@ocots
ocots committed Jun 20, 2024
1 parent fad9005 commit 424e131
Show file tree
Hide file tree
Showing 2 changed files with 50 additions and 28 deletions.
74 changes: 48 additions & 26 deletions src/plot.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,8 +37,15 @@ Update the plot `p` with the i-th component of a vectorial function of time `f(t
- The argument `s` can be `:state`, `:control` or `:costate`.
- `time` can be `:default` or `:normalized`.
"""
function __plot_time!(p::Union{Plots.Plot, Plots.Subplot}, sol::OptimalControlSolution, s::Symbol, i::Integer, time::Symbol;
t_label, label::String, kwargs...)
function __plot_time!(
p::Union{Plots.Plot, Plots.Subplot},
sol::OptimalControlSolution,
s::Symbol,
i::Integer,
time::Symbol;
t_label,
label::String,
kwargs...)

# t_label depends if time is normalized or not
t_label = @match time begin
Expand Down Expand Up @@ -165,8 +172,11 @@ $(TYPEDSIGNATURES)
Initial plot.
"""
function __initial_plot(sol::OptimalControlSolution; layout::Symbol=:split,
control::Symbol=:components, kwargs...)
function __initial_plot(
sol::OptimalControlSolution;
layout::Symbol=:split,
control::Symbol=:components,
kwargs...)

# parameters
n = sol.state_dimension
Expand Down Expand Up @@ -280,15 +290,20 @@ Plot the optimal control solution `sol` using the layout `layout`.
"""
function Plots.plot!(
p::Plots.Plot,
sol::OptimalControlSolution;
sol::OptimalControlSolution;
layout::Symbol=:split,
control::Symbol=:components,
time::Symbol=:default,
solution_label::String="",
state_style=(),
control_style=(),
costate_style=(),
kwargs...)

if solution_label!=""
solution_label = " " * solution_label
end

#
n = sol.state_dimension
m = sol.control_dimension
Expand All @@ -302,42 +317,42 @@ function Plots.plot!(

if layout==:group

__plot_time!(p[1], sol, n, :state, time; t_label=t_label, labels=x_labels, title="state", lims=:auto, series_attr..., state_style...)
__plot_time!(p[2], sol, n, :costate, time; t_label=t_label, labels="p".*x_labels, title="costate", lims=:auto, series_attr..., costate_style...)
__plot_time!(p[1], sol, n, :state, time; t_label=t_label, labels=x_labels.*solution_label, title="state", lims=:auto, series_attr..., state_style...)
__plot_time!(p[2], sol, n, :costate, time; t_label=t_label, labels="p".*x_labels.*solution_label, title="costate", lims=:auto, series_attr..., costate_style...)
@match control begin
:components => begin
__plot_time!(p[3], sol, m, :control, time; t_label=t_label, labels=u_labels, title="control", lims=:auto, series_attr..., control_style...)
__plot_time!(p[3], sol, m, :control, time; t_label=t_label, labels=u_labels.*solution_label, title="control", lims=:auto, series_attr..., control_style...)
end
:norm => begin
__plot_time!(p[3], sol, :control_norm, -1, time; t_label=t_label, label="‖"*u_label*"‖", title="control norm", lims=:auto, series_attr..., control_style...)
__plot_time!(p[3], sol, :control_norm, -1, time; t_label=t_label, label="‖"*u_label*"‖".*solution_label, title="control norm", lims=:auto, series_attr..., control_style...)
end
:all => begin
__plot_time!(p[3], sol, m, :control, time; t_label=t_label, labels=u_labels, title="control", lims=:auto, series_attr..., control_style...)
__plot_time!(p[4], sol, :control_norm, -1, time; t_label=t_label, label="‖"*u_label*"‖", title="control norm", lims=:auto, series_attr..., control_style...)
__plot_time!(p[3], sol, m, :control, time; t_label=t_label, labels=u_labels.*solution_label, title="control", lims=:auto, series_attr..., control_style...)
__plot_time!(p[4], sol, :control_norm, -1, time; t_label=t_label, label="‖"*u_label*"‖".*solution_label, title="control norm", lims=:auto, series_attr..., control_style...)
end
_ => throw(IncorrectArgument("No such choice for control. Use :components, :norm or :all"))
end

elseif layout==:split

for i ∈ 1:n
__plot_time!(p[i], sol, :state, i, time; t_label=t_label, label=x_labels[i], series_attr..., state_style...)
__plot_time!(p[i+n], sol, :costate, i, time; t_label=t_label, label="p"*x_labels[i], series_attr..., costate_style...)
__plot_time!(p[i], sol, :state, i, time; t_label=t_label, label=x_labels[i]*solution_label, series_attr..., state_style...)
__plot_time!(p[i+n], sol, :costate, i, time; t_label=t_label, label="p"*x_labels[i]*solution_label, series_attr..., costate_style...)
end
@match control begin
:components => begin
for i ∈ 1:m
__plot_time!(p[i+2*n], sol, :control, i, time; t_label=t_label, label=u_labels[i], series_attr..., control_style...)
__plot_time!(p[i+2*n], sol, :control, i, time; t_label=t_label, label=u_labels[i]*solution_label, series_attr..., control_style...)
end
end
:norm => begin
__plot_time!(p[2*n+1], sol, :control_norm, -1, time; t_label=t_label, label="‖"*u_label*"‖", series_attr..., control_style...)
__plot_time!(p[2*n+1], sol, :control_norm, -1, time; t_label=t_label, label="‖"*u_label*"‖"*solution_label, series_attr..., control_style...)
end
:all => begin
for i ∈ 1:m
__plot_time!(p[i+2*n], sol, :control, i, time; t_label=t_label, label=u_labels[i], series_attr..., control_style...)
__plot_time!(p[i+2*n], sol, :control, i, time; t_label=t_label, label=u_labels[i]*solution_label, series_attr..., control_style...)
end
__plot_time!(p[2*n+m+1], sol, :control_norm, -1, time; t_label=t_label, label="‖"*u_label*"‖", series_attr..., control_style...)
__plot_time!(p[2*n+m+1], sol, :control_norm, -1, time; t_label=t_label, label="‖"*u_label*"‖"*solution_label, series_attr..., control_style...)
end
_ => throw(IncorrectArgument("No such choice for control. Use :components, :norm or :all"))
end
Expand Down Expand Up @@ -377,16 +392,17 @@ Plot the optimal control solution `sol` using the layout `layout`.
function Plots.plot(sol::OptimalControlSolution;
layout::Symbol=:split,
control::Symbol=:components,
time::Symbol=:default,
time::Symbol=:default,
size=__size_plot(sol, control),
solution_label::String="",
state_style=(),
control_style=(),
costate_style=(),
kwargs...)
#
p = __initial_plot(sol; layout=layout, control=control, size=size, kwargs...)
#
return plot!(p, sol; layout=layout, control=control, time=time,
return plot!(p, sol; layout=layout, control=control, time=time, solution_label=solution_label,
state_style=state_style, control_style=control_style, costate_style=costate_style, kwargs...)
end

Expand All @@ -402,9 +418,11 @@ corresponding respectively to the argument `xx` and the argument `yy`.
- If `xx` is `:time`, then, a label is added to the plot.
- The argument `yy` can be `:state`, `:control` or `:costate`.
"""
@recipe function f(sol::OptimalControlSolution,
@recipe function f(
sol::OptimalControlSolution,
xx::Union{Symbol,Tuple{Symbol,Integer}},
yy::Union{Symbol,Tuple{Symbol,Integer}}, time::Symbol=:default)
yy::Union{Symbol,Tuple{Symbol,Integer}},
time::Symbol=:default)

#
x = __get_data_plot(sol, xx, time=time)
Expand All @@ -416,7 +434,11 @@ corresponding respectively to the argument `xx` and the argument `yy`.
return x, y
end

function recipe_label(sol::OptimalControlSolution, xx::Union{Symbol,Tuple{Symbol,Integer}}, yy::Union{Symbol,Tuple{Symbol,Integer}})
function recipe_label(
sol::OptimalControlSolution,
xx::Union{Symbol,Tuple{Symbol,Integer}},
yy::Union{Symbol,Tuple{Symbol,Integer}})

#
label = false
#
Expand All @@ -428,11 +450,11 @@ function recipe_label(sol::OptimalControlSolution, xx::Union{Symbol,Tuple{Symbol
end

label = @match s begin
:state => sol.state_components_names[i]
:control => sol.control_components_names[i]
:costate => "p"*sol.state_components_names[i]
:state => sol.state_components_names[i]
:control => sol.control_components_names[i]
:costate => "p"*sol.state_components_names[i]
:control_norm => "‖"*sol.control_name*"‖"
_ => error("Internal error, no such choice for label")
_ => error("Internal error, no such choice for label")
end

end
Expand Down
4 changes: 2 additions & 2 deletions test/test_plot_manual.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ sol.message = "ceci est un test"
sol.success = true

#
plt = plot(sol, layout=layout, control=control_plt, size=size, flip=true, linewidth=5)
plt = plot(sol, layout=layout, control=control_plt, size=size, flip=true, linewidth=5, solution_label="sol1")
#plot(sol, layout=:group)
#ps=plot(sol, :time, (:state, 1))
#plot!(ps, sol, :time, (:control, 1))
Expand Down Expand Up @@ -101,7 +101,7 @@ sol.message = "ceci est un test"
sol.success = true

if do_plot_2
plot!(plt, sol, layout=layout, size=size, control=control_plt)
plot!(plt, sol, layout=layout, size=size, control=control_plt, solution_label="sol2")
else
plt
end
Expand Down

0 comments on commit 424e131

Please sign in to comment.