diff --git a/src/plot.jl b/src/plot.jl index 1ded14b0..5f5861ac 100644 --- a/src/plot.jl +++ b/src/plot.jl @@ -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 @@ -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 @@ -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 @@ -302,18 +317,18 @@ 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 @@ -321,23 +336,23 @@ function Plots.plot!( 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 @@ -377,8 +392,9 @@ 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=(), @@ -386,7 +402,7 @@ function Plots.plot(sol::OptimalControlSolution; # 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 @@ -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) @@ -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 # @@ -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 diff --git a/test/test_plot_manual.jl b/test/test_plot_manual.jl index 827ec926..e1467443 100644 --- a/test/test_plot_manual.jl +++ b/test/test_plot_manual.jl @@ -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)) @@ -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