adding plotting recipes

Author: simon-lc

Project.toml

@@ -7,12 +7,11 @@ version = "0.1.5"
 Altro = "5dcf52e5-e2fb-48e0-b826-96f46d2e3e73"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-GeometryBasics = "5c1252a2-5f33-56bf-86c9-59e7332b4326"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-MeshCat = "283c5d60-a78f-5afe-a0af-af636b173e11"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 RobotDynamics = "38ceca67-d8d3-44e8-9852-78a5596522e1"
 Rotations = "6038ab10-8711-5258-84ad-4b1120ba62dc"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

examples/intro_example.jl

@@ -75,5 +75,6 @@ prob = GameProblem(N,dt,x0,model,opts,game_obj,game_con)
 @time newton_solve!(prob)
 
 # Visualize the Results
-Algames.plot_traj!(prob.model, prob.pdtraj.pr)
-Algames.plot_violation!(prob.stats)
+using Plots
+plot(prob.model, prob.pdtraj.pr)
+plot(prob.stats)

src/Algames.jl

@@ -6,11 +6,10 @@ using Altro
 using BenchmarkTools
 using ForwardDiff
 using LinearAlgebra
-using MeshCat
-using GeometryBasics
 using Parameters
 using Printf
 using Random
+using RecipesBase
 using RobotDynamics
 using Rotations
 using SparseArrays
@@ -162,8 +161,9 @@ export
 
 # Plots
 export
-    plot_traj!,
-    plot_violation!
+    recipe_traj,
+    recipe_violation
+
 
 # Utils
 include("utils.jl")

src/plots/solver_plots.jl

@@ -2,60 +2,119 @@
 # Trajectory Plot
 ################################################################################
 
-function plot_traj!(model::AbstractGameModel, traj::Traj; plt=plot())
-    plot!(plt, legend=false, aspect_ratio=:equal)
+# function plot_traj!(model::AbstractGameModel, traj::Traj; plt=plot())
+#     plot!(plt, legend=false, aspect_ratio=:equal)
+#     N = length(traj)
+#     for i = 1:model.p
+#         xi = [Algames.state(traj[k])[model.pz[i][1]] for k=1:N]
+#         yi = [Algames.state(traj[k])[model.pz[i][2]] for k=1:N]
+#         plot!(xi, yi, label=false)
+#         scatter!(xi, yi)
+#     end
+#     display(plt)
+#     return nothing
+# end
+
+@recipe function recipe_traj(model::AbstractGameModel, traj::Traj)
     N = length(traj)
+    x = []
+    y = []
     for i = 1:model.p
         xi = [Algames.state(traj[k])[model.pz[i][1]] for k=1:N]
         yi = [Algames.state(traj[k])[model.pz[i][2]] for k=1:N]
-        plot!(xi, yi, label=false)
-        scatter!(xi, yi)
+        push!(x, xi)
+        push!(y, yi)
     end
-    display(plt)
-    return nothing
+    aspect_ratio  --> :equal
+	xguide        --> "X"
+	yguide        --> "Y"
+	title         --> "Trajectories"
+    legend        --> :false
+	seriestype    --> [fill(:scatter, model.p)... fill(:path, model.p)...]
+    return [x, x], [y, y]
 end
 
+
 ################################################################################
 # Constraint Violation Plot
 ################################################################################
 
-function plot_violation!(stats::Statistics; plt=plot(), lw::T=5.0) where {T}
-	plot!(plt,
-		size=(500,500),
-		layout=(1,1,))
-    iter = stats.iter
-    dyn = log.(10, max.(1e-10, [stats.dyn_vio[i].max for i=1:iter]))
+# function plot_violation!(stats::Statistics; plt=plot(), lw::T=5.0) where {T}
+# 	plot!(plt,
+# 		size=(500,500),
+# 		layout=(1,1,))
+#     iter = stats.iter
+#     dyn = log.(10, max.(1e-10, [stats.dyn_vio[i].max for i=1:iter]))
+#     con = log.(10, max.(1e-10, [stats.con_vio[i].max for i=1:iter]))
+#     sta = log.(10, max.(1e-10, [stats.sta_vio[i].max for i=1:iter]))
+#     opt = log.(10, max.(1e-10, [stats.opt_vio[i].max for i=1:iter]))
+# 	y_min = minimum([dyn; con; sta; opt])
+# 	y_max = maximum([dyn; con; sta; opt])
+# 	# Set up plot
+# 	plot!(plt[1,1],
+# 		legend=:bottomleft,
+# 		xlabel="Outer Loop Iterations",
+# 		ylabel="log(cons. vio.)",
+# 		title="Constraint Violation")
+# 	# Add curves
+# 	plot!(plt[1,1], dyn, linewidth=lw, label="dyn", legend=:bottomleft)
+# 	plot!(plt[1,1], con, linewidth=lw, label="con", legend=:bottomleft)
+# 	plot!(plt[1,1], sta, linewidth=lw, label="sta", legend=:bottomleft)
+# 	plot!(plt[1,1], opt, linewidth=lw, label="opt", legend=:bottomleft)
+# 	# Add rectangles
+# 	plot_epochs!(plt, y_min, y_max, stats.outer_iter)
+#
+#     display(plt)
+#     return nothing
+# end
+#
+# function plot_epochs!(plt, y_min::T, y_max::T, epochs::Vector{Int}) where {T}
+# 	rectangle(w, h, x, y) = Shape(x .+ [0,w,w,0], y .+ [0,0,h,h])
+# 	i_start = 1
+# 	i_end = -1
+# 	for k = 1:epochs[end]
+# 		i_end = findlast(x -> x==k, epochs )
+# 		plot!(rectangle(i_end-i_start,y_max-y_min,i_start,y_min), opacity=.1, label=false)
+# 		i_start = i_end + 1
+# 	end
+# 	return nothing
+# end
+
+@recipe function recipe_violation(stats::Statistics)
+	rectangle(w, h, x, y) = Shape(x .+ [0,w,w,0], y .+ [0,0,h,h])
+
+	iter = stats.iter
+	dyn = log.(10, max.(1e-10, [stats.dyn_vio[i].max for i=1:iter]))
     con = log.(10, max.(1e-10, [stats.con_vio[i].max for i=1:iter]))
     sta = log.(10, max.(1e-10, [stats.sta_vio[i].max for i=1:iter]))
     opt = log.(10, max.(1e-10, [stats.opt_vio[i].max for i=1:iter]))
 	y_min = minimum([dyn; con; sta; opt])
 	y_max = maximum([dyn; con; sta; opt])
-	# Set up plot
-	plot!(plt[1,1],
-		legend=:bottomleft,
-		xlabel="Outer Loop Iterations",
-		ylabel="log(cons. vio.)",
-		title="Constraint Violation")
-	# Add curves
-	plot!(plt[1,1], dyn, linewidth=lw, label="dyn", legend=:bottomleft)
-	plot!(plt[1,1], con, linewidth=lw, label="con", legend=:bottomleft)
-	plot!(plt[1,1], sta, linewidth=lw, label="sta", legend=:bottomleft)
-	plot!(plt[1,1], opt, linewidth=lw, label="opt", legend=:bottomleft)
-	# Add rectangles
-	plot_epochs!(plt, y_min, y_max, stats.outer_iter)
-
-    display(plt)
-    return nothing
-end
 
-function plot_epochs!(plt, y_min::T, y_max::T, epochs::Vector{Int}) where {T}
-	rectangle(w, h, x, y) = Shape(x .+ [0,w,w,0], y .+ [0,0,h,h])
+	epochs = stats.outer_iter
 	i_start = 1
 	i_end = -1
-	for k = 1:epochs[end]
+	epo = []
+	N_epochs = epochs[end]
+	for k = 1:N_epochs
 		i_end = findlast(x -> x==k, epochs )
-		plot!(rectangle(i_end-i_start,y_max-y_min,i_start,y_min), opacity=.1, label=false)
+		push!(epo, Vector(range(i_start, stop = i_end, length = iter)))
 		i_start = i_end + 1
 	end
-	return nothing
+
+	size        --> (500,500)
+	legend      --> :bottomleft
+	ylims       --> [y_min, y_max]
+	xguide      --> "Outer Loop Iterations"
+	yguide      --> "log(cons. vio.)"
+	title       --> "Constraint Violation"
+	linewidth   --> [zeros(N_epochs)... 5.0 5.0 5.0 5.0]
+	label       --> [fill("", N_epochs)... "dyn" "con" "sta" "opt" ]
+	fillrange   --> [[fill(y_min*ones(90), N_epochs)... dyn con sta opt ]
+					 [fill(zeros(iter), N_epochs)... zeros(iter) zeros(iter) zeros(iter) zeros(iter) ]]
+	fillalpha   --> [fill(0.1, N_epochs)... 0.0 0.0 0.0 0.0]
+
+	y = fill(y_max*ones(iter), N_epochs)
+	x = Vector(1:iter)
+	return  [epo..., x, x, x, x], [y..., dyn, con, sta, opt]
 end

test/Project.toml

@@ -3,10 +3,10 @@ Altro = "5dcf52e5-e2fb-48e0-b826-96f46d2e3e73"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-MeshCat = "283c5d60-a78f-5afe-a0af-af636b173e11"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 RobotDynamics = "38ceca67-d8d3-44e8-9852-78a5596522e1"
 Rotations = "6038ab10-8711-5258-84ad-4b1120ba62dc"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

test/runtests.jl

@@ -4,10 +4,10 @@ using Altro
 using BenchmarkTools
 using ForwardDiff
 using LinearAlgebra
-using MeshCat
 using Parameters
 using Printf
 using Random
+using RecipesBase
 using RobotDynamics
 using Rotations
 using SparseArrays