From 7260d71a5e1560631d41c9bd0e77c8b7066768d7 Mon Sep 17 00:00:00 2001
From: Olivier Cots <olivier.cots@irit.fr>
Date: Sat, 8 Jun 2024 17:14:21 +0200
Subject: [PATCH] default plot size

---
 src/CTBase.jl            |  3 +--
 src/plot.jl              | 28 +++++++++++++++++++++++-----
 src/print.jl             | 14 +++++++-------
 test/test_plot_manual.jl |  4 +++-
 4 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/src/CTBase.jl b/src/CTBase.jl
index 44d30887..3eea616f 100644
--- a/src/CTBase.jl
+++ b/src/CTBase.jl
@@ -19,8 +19,7 @@ using Interpolations: linear_interpolation, Line, Interpolations # for default i
 using MLStyle # pattern matching
 using Parameters # @with_kw: to have default values in struct
 using Plots
-import Plots: plot, plot! # import instead of using to overload the plot and plot! functions
-using Printf # to print an OptimalControlModel
+using Printf # to print an Opt imalControlModel
 using DataStructures # OrderedDict for aliases
 using Unicode # unicode primitives
 using PrettyTables # to print a table
diff --git a/src/plot.jl b/src/plot.jl
index 2518b233..5f6295a1 100644
--- a/src/plot.jl
+++ b/src/plot.jl
@@ -48,7 +48,7 @@ function __plot_time!(p::Union{Plots.Plot, Plots.Subplot}, sol::OptimalControlSo
     end
 
     # reset ylims: ylims=:auto
-    CTBase.plot!(p, sol, :time, (s, i), time; ylims=:auto, xlabel=t_label, label=label, kwargs...) # use simple plot
+    Plots.plot!(p, sol, :time, (s, i), time; ylims=:auto, xlabel=t_label, label=label, kwargs...) # use simple plot
 
     # change ylims if the gap between min and max is less than a tol
     tol  = 1e-3
@@ -262,7 +262,7 @@ Plot the optimal control solution `sol` using the layout `layout`.
 - `time` can be `:default` or `:normalized`.
 - The keyword arguments `state_style`, `control_style` and `costate_style` are passed to the `plot` function of the `Plots` package. The `state_style` is passed to the plot of the state, the `control_style` is passed to the plot of the control and the `costate_style` is passed to the plot of the costate.
 """
-function CTBase.plot!(p::Plots.Plot, sol::OptimalControlSolution; layout::Symbol=:split,
+function Plots.plot!(p::Plots.Plot, sol::OptimalControlSolution; layout::Symbol=:split,
     control::Symbol=:components, time::Symbol=:default,
     state_style=(), control_style=(), costate_style=(), kwargs...)
 
@@ -326,6 +326,17 @@ function CTBase.plot!(p::Plots.Plot, sol::OptimalControlSolution; layout::Symbol
 
 end
 
+function __size_plot(sol::OptimalControlSolution, control::Symbol)
+    n = sol.state_dimension
+    #m = sol.control_dimension
+    m = @match control begin
+        :components => sol.control_dimension
+        :norm => 1
+        :all => sol.control_dimension + 1
+        _ => throw(IncorrectArgument("No such choice for control. Use :components, :norm or :all"))
+    end
+    return (600, 140*(n+m))
+end
 
 """
 $(TYPEDSIGNATURES)
@@ -337,10 +348,17 @@ Plot the optimal control solution `sol` using the layout `layout`.
 - The argument `layout` can be `:group` or `:split` (default).
 - The keyword arguments `state_style`, `control_style` and `costate_style` are passed to the `plot` function of the `Plots` package. The `state_style` is passed to the plot of the state, the `control_style` is passed to the plot of the control and the `costate_style` is passed to the plot of the costate.
 """
-function CTBase.plot(sol::OptimalControlSolution; layout::Symbol=:split, 
-    control::Symbol=:components, time::Symbol=:default, state_style=(), control_style=(), costate_style=(), kwargs...)
+function Plots.plot(sol::OptimalControlSolution; 
+    layout::Symbol=:split, 
+    control::Symbol=:components, 
+    time::Symbol=:default, 
+    size=__size_plot(sol, control),
+    state_style=(), 
+    control_style=(), 
+    costate_style=(),
+    kwargs...)
     #
-    p = __initial_plot(sol; layout=layout, kwargs...)
+    p = __initial_plot(sol; layout=layout, control=control, size=size, kwargs...)
     #
     return plot!(p, sol; layout=layout, control=control, time=time,
         state_style=state_style, control_style=control_style, costate_style=costate_style, kwargs...)
diff --git a/src/print.jl b/src/print.jl
index da0697b9..947a10c6 100644
--- a/src/print.jl
+++ b/src/print.jl
@@ -198,18 +198,18 @@ function Base.show(io::IO, ::MIME"text/plain", ocp::OptimalControlModel{<: TimeD
     #is_variable_dependent(ocp) && push!(header, "variable")
     push!(header, "variable")
     push!(header, "dynamics*", "objective*", "constraints")
-    data = hcat(__is_time_not_set(ocp) ? "❌" : "✅",
-        __is_state_not_set(ocp) ? "❌" : "✅", 
-        __is_control_not_set(ocp) ? "❌" : "✅")
+    data = hcat(__is_time_not_set(ocp)      ? "❌" : "✅",
+                __is_state_not_set(ocp)     ? "❌" : "✅", 
+                __is_control_not_set(ocp)   ? "❌" : "✅")
     #is_variable_dependent(ocp) && 
     begin
         (data = hcat(data, 
-        __is_variable_not_set(ocp) ? "❌" : "✅")) 
+                __is_variable_not_set(ocp)  ? "❌" : "✅")) 
     end
     data = hcat(data, 
-        __is_dynamics_not_set(ocp) ? "❌" : "✅",
-        __is_objective_not_set(ocp) ? "❌" : "✅",
-        isempty(ocp.constraints) ? "❌" : "✅")
+                __is_dynamics_not_set(ocp)  ? "❌" : "✅",
+                __is_objective_not_set(ocp) ? "❌" : "✅",
+                isempty(ocp.constraints)    ? "❌" : "✅")
     println("")
     pretty_table(data, header=header, header_crayon=crayon"yellow")
     nothing
diff --git a/test/test_plot_manual.jl b/test/test_plot_manual.jl
index 7d2cdad4..fe552380 100644
--- a/test/test_plot_manual.jl
+++ b/test/test_plot_manual.jl
@@ -1,4 +1,5 @@
 using CTBase
+#using Plots
 
 layout = :split
 size   = (900, 600)
@@ -51,11 +52,12 @@ sol.message = "ceci est un test"
 sol.success = true
 
 #
-plt = plot(sol, layout=layout, size=size, control=control_plt)
+plt = plot(sol, layout=layout, control=control_plt) #, size=size)
 #plot(sol, layout=:group)
 #ps=plot(sol, :time, (:state, 1))
 #plot!(ps, sol, :time, (:control, 1))
 
+return 
 # ----------------------------------------
 # SOL 2
 n=2