can choose plot by desc

Author: ocots

ext/default.jl

@@ -79,24 +79,35 @@ function __size_plot(
             end
         end
         if control == :all && nb_plots == 4
-            return (600, 560)
+            return (600, 420)
         else
             return (600, 280)
         end
     else
         n = CTModels.state_dimension(sol)
+        m = CTModels.control_dimension(sol)
         l = @match control begin
-            :components => CTModels.control_dimension(sol)
+            :components => m
             :norm => 1
-            :all => CTModels.control_dimension(sol) + 1
+            :all => m + 1
             _ => throw(
                 CTBase.IncorrectArgument(
                     "No such choice for control. Use :components, :norm or :all"
                 ),
             )
         end
         nc = model === nothing ? 0 : CTModels.dim_path_constraints_nl(model)
-        return (600, 140 * (n + l + nc))
+        nb_lines = 0
+        nb_lines += (:state ∈ description || :costate ∈ description) ? n : 0
+        nb_lines += :control ∈ description ? l : 0
+        nb_lines += (:cons ∈ description || :dual ∈ description) ? nc : 0
+        if nb_lines==1
+            return (600, 280)
+        elseif nb_lines==2
+            return (600, 420)
+        else 
+            return (600, 140 * nb_lines)
+        end
     end
 
 end

ext/plot.jl

@@ -231,7 +231,7 @@ function __initial_plot(
                 :control ∈ description && push!(plots, Plots.plot()) # control
                 :control ∈ description && push!(plots, Plots.plot()) # control norm
                 if length(plots) == 4
-                    return Plots.plot(plots...; layout=(2, 2), bottommargin=5mm, kwargs...)   
+                    return Plots.plot(plots...; layout=(2, 2), kwargs...)   
                 else
                     return Plots.plot(plots...; layout=(1, length(plots)), bottommargin=5mm, kwargs...)
                 end
@@ -256,6 +256,9 @@ function __initial_plot(
             :costate ∈ description && push!(costate_plots, PlotLeaf())
         end
 
+        println("state_plots: ", state_plots)
+        println("costate_plots: ", costate_plots)
+
         # create the control plots
         if :control ∈ description
             l = m
@@ -284,10 +287,6 @@ function __initial_plot(
             end
         end
 
-        # starting root
-        root = EmptyPlot()
-        nblines = 0
-
         # assemble the state and costate plots
         node_x = isempty(state_plots)   ? EmptyPlot() : PlotNode(:column, state_plots)
         node_p = isempty(costate_plots) ? EmptyPlot() : PlotNode(:column, costate_plots)
@@ -301,6 +300,8 @@ function __initial_plot(
         node_u = isempty(control_plots) ? EmptyPlot() : PlotNode(:column, control_plots)
 
         # create the root node
+        root = EmptyPlot()
+        nblines = 0
         if ( !(node_xp isa EmptyPlot) && !(node_u isa EmptyPlot) )
             nblines = n + l
             a = __height(round(n / nblines; digits=2))
@@ -309,9 +310,12 @@ function __initial_plot(
                 b
             ]
             root = PlotNode(lay, [node_xp, node_u])
-        else
-            root = node_xp isa EmptyPlot ? node_u : node_xp
-            nblines = node_xp isa EmptyPlot ? l : n
+        elseif !(node_xp isa EmptyPlot)
+            root = node_xp
+            nblines = n
+        elseif !(node_u isa EmptyPlot)
+            root = node_u
+            nblines = l
         end
 
         # Add the path constraints and their dual variables (in two columns as for state and costate) plots 
@@ -336,8 +340,9 @@ function __initial_plot(
             end
 
             # update the root node
-            if !(node_cocp isa EmptyPlot)    
-                c = __height(round(nc / (nblines + nc); digits=2))
+            if !(node_cocp isa EmptyPlot)
+                nblines += nc  
+                c = __height(round(nc / nblines; digits=2))
                 @eval lay = @layout [
                     a
                     $c
@@ -347,7 +352,7 @@ function __initial_plot(
         end
 
         # plot
-        return __plot_tree(root; kwargs...)
+        return nblines==1 ? __plot_tree(root; bottommargin=5mm, kwargs...) : __plot_tree(root; kwargs...)
 
     else
         throw(CTBase.IncorrectArgument("No such choice for layout. Use :group or :split"))

test/extras/plot_duals.jl

@@ -62,10 +62,19 @@ end;
 sol = SOL(ocp, t0, tf);
 
 plt = plot(sol, ocp)
+plt = plot(sol, ocp, :state)
+plt = plot(sol, ocp, :state, :costate)
+plt = plot(sol, ocp, :state, :control)
+plt = plot(sol, ocp, :state, :control, :constraint)
+plt = plot(sol, ocp, :costate)
+plt = plot(sol, ocp, :control)
+plt = plot(sol, ocp, :constraint)
+plt = plot(sol, ocp, :dual)
+plt = plot(sol, ocp, :constraint, :dual)
 
-mixed_con_dual = CTModels.dual(sol, ocp, :mixed_con)
-plot(range(t0, tf; length=101), mixed_con_dual)
+# mixed_con_dual = CTModels.dual(sol, ocp, :mixed_con)
+# plot(range(t0, tf; length=101), mixed_con_dual)
 
-eq2_dual = CTModels.dual(sol, ocp, :eq2)
-plot(range(t0, tf; length=101), t -> eq2_dual(t)[1]; label="eq2_dual 1")
-plot!(range(t0, tf; length=101), t -> eq2_dual(t)[2]; label="eq2_dual 2")
+# eq2_dual = CTModels.dual(sol, ocp, :eq2)
+# plot(range(t0, tf; length=101), t -> eq2_dual(t)[1]; label="eq2_dual 1")
+# plot!(range(t0, tf; length=101), t -> eq2_dual(t)[2]; label="eq2_dual 2")

test/extras/plot_manual.jl

@@ -129,7 +129,9 @@ end;
 sol = get_solution();
 
 # layout = :group
-plot(sol; layout=:group)
+plot(sol; layout=:group, control=:components)
+plot(sol; layout=:group, control=:norm)
+plot(sol; layout=:group, control=:all)
 plot(sol, :state; layout=:group)
 plot(sol, :costate; layout=:group)
 plot(sol, :control; layout=:group)
@@ -147,6 +149,7 @@ plot(sol, :control; layout=:split, control=:norm)
 plot(sol, :state, :control; layout=:split)
 plot(sol, :control; layout=:split, control=:all)
 plot(sol, :state, :control; layout=:split, control=:all)
+plot(sol, :state, :costate; layout=:split)
 nothing
 # plt = plot(sol; size=(700, 450), time=:normalise)