Revamp plotting interface

Author: BatyLeo

Project.toml

@@ -7,7 +7,6 @@ authors = ["Members of JuliaDecisionFocusedLearning"]
 projects = ["docs", "test"]
 
 [deps]
-Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 ConstrainedShortestPaths = "b3798467-87dc-4d99-943d-35a1bd39e395"
 DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe"
@@ -22,12 +21,11 @@ InferOpt = "4846b161-c94e-4150-8dac-c7ae193c601f"
 Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
-JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
+JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Metalhead = "dbeba491-748d-5e0e-a39e-b530a07fa0cc"
 NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
-Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
@@ -37,8 +35,13 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
+[weakdeps]
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+
+[extensions]
+DFLBenchmarksPlotsExt = "Plots"
+
 [compat]
-Colors = "0.13.1"
 Combinatorics = "1.0.3"
 ConstrainedShortestPaths = "0.6.0"
 DataDeps = "0.7"
@@ -53,8 +56,8 @@ InferOpt = "0.7.0"
 Ipopt = "1.6"
 IterTools = "1.10.0"
 JSON = "1"
-JuMP = "1.22"
 LaTeXStrings = "1.4.0"
+JuMP = "1.22"
 LinearAlgebra = "1"
 Metalhead = "0.9.4"
 NPZ = "0.4"

docs/src/custom_benchmarks.md

@@ -69,13 +69,13 @@ generate_maximizer(bench::MyBenchmark)
 ### Optional methods
 
 ```julia
+generate_baseline_policies(bench::MyBenchmark) -> collection of callables
 is_minimization_problem(bench::MyBenchmark) -> Bool   # default: false (maximization)
 objective_value(bench::MyBenchmark, sample::DataSample, y) -> Real
 compute_gap(bench::MyBenchmark, dataset, model, maximizer) -> Float64
-plot_data(bench::MyBenchmark, sample::DataSample; kwargs...)
-plot_instance(bench::MyBenchmark, instance; kwargs...)
-plot_solution(bench::MyBenchmark, sample::DataSample, y; kwargs...)
-generate_baseline_policies(bench::MyBenchmark) -> collection of callables
+has_visualization(bench::MyBenchmark) -> Bool                            # default: false; return true when plot methods are implemented/available
+plot_instance(bench::MyBenchmark, sample::DataSample; kwargs...)
+plot_solution(bench::MyBenchmark, sample::DataSample; kwargs...)
 ```
 
 ---
@@ -148,6 +148,13 @@ generate_baseline_policies(bench::MyDynamicBenchmark)
 # Each callable performs a full episode rollout and returns the trajectory.
 ```
 
+### Optional visualization methods
+
+```julia
+plot_trajectory(bench::MyDynamicBenchmark, traj::Vector{DataSample}; kwargs...)
+animate_trajectory(bench::MyDynamicBenchmark, traj::Vector{DataSample}; kwargs...)
+```
+
 `generate_dataset` for dynamic benchmarks **requires** a `target_policy` kwarg, 
 there is no default. The `target_policy` must be a callable `(env) -> Vector{DataSample}`.
 

docs/src/tutorials/warcraft_tutorial.jl

@@ -32,8 +32,8 @@ y_true = sample.y
 # `context` is not used in this benchmark (no solver kwargs needed), so it is empty:
 isempty(sample.context)
 
-# For some benchmarks, we provide the following plotting method [`plot_data`](@ref) to visualize the data:
-plot_data(b, sample)
+# For some benchmarks, we provide the following plotting method [`plot_solution`](@ref) to visualize the data:
+plot_solution(b, sample)
 # We can see here the terrain image, the true terrain weights, and the true shortest path avoiding the high cost cells.
 
 # ## Building a pipeline
@@ -50,7 +50,7 @@ maximizer = generate_maximizer(b; dijkstra=true)
 # In the case o fthe Warcraft benchmark, the method has an additional keyword argument to chose the algorithm to use: Dijkstra's algorithm or Bellman-Ford algorithm.
 y = maximizer(θ)
 # As we can see, currently the pipeline predicts random noise as cell weights, and therefore the maximizer returns a straight line path.
-plot_data(b, DataSample(; x, θ, y))
+plot_solution(b, DataSample(; x, θ, y))
 # We can evaluate the current pipeline performance using the optimality gap metric:
 starting_gap = compute_gap(b, test_dataset, model, maximizer)
 
@@ -85,7 +85,7 @@ final_gap = compute_gap(b, test_dataset, model, maximizer)
 #
 θ = model(x)
 y = maximizer(θ)
-plot_data(b, DataSample(; x, θ, y))
+plot_solution(b, DataSample(; x, θ, y))
 
 using Test #src
 @test final_gap < starting_gap #src

docs/src/using_benchmarks.md

@@ -10,7 +10,7 @@ A benchmark bundles a problem family (an instance generator, a combinatorial sol
 Three abstract types cover the main settings:
 - **`AbstractBenchmark`**: static problems (one instance, one decision)
 - **`AbstractStochasticBenchmark{exogenous}`**: stochastic problems (type parameter indicates whether uncertainty is exogenous)
-- **`AbstractDynamicBenchmark`**: sequential / multi-stage problems
+- **`AbstractDynamicBenchmark{exogenous}`**: sequential / multi-stage problems
 
 The sections below explain what changes between these settings. For most purposes, start with a static benchmark to understand the core workflow.
 
@@ -180,10 +180,29 @@ rewards, samples = evaluate_policy!(pol, envs, n_episodes)
 
 ## Visualization
 
-Where implemented, benchmarks provide benchmark-specific plotting helpers:
-
+Plots is an **optional** dependency, load it with `using Plots` to unlock the plot functions. Not all benchmarks support visualization, call `has_visualization(bench)` to check.
 ```julia
-plot_data(bench, sample)            # overview of a data sample
-plot_instance(bench, instance)      # raw problem instance
-plot_solution(bench, sample, y)     # overlay solution on instance
+using Plots
+
+bench = Argmax2DBenchmark()
+dataset = generate_dataset(bench, 10)
+sample = dataset[1]
+
+has_visualization(bench)           # true
+plot_instance(bench, sample)       # problem geometry only
+plot_solution(bench, sample)       # sample.y overlaid on the instance
+plot_solution(bench, sample, y)    # convenience 3-arg form: override y before plotting
+
+# Dynamic benchmarks only
+traj = generate_anticipative_solver(bench)(env)
+plot_trajectory(bench, traj)           # grid of epoch subplots
+anim = animate_trajectory(bench, traj; fps=2)
+gif(anim, "episode.gif")
 ```
+
+- `has_visualization(bench)`: returns `true` for benchmarks that implement plot support (if Plots is loaded).
+- `plot_instance(bench, sample; kwargs...)`: renders the problem geometry without any solution.
+- `plot_solution(bench, sample; kwargs...)`: renders `sample.y` overlaid on the instance.
+- `plot_solution(bench, sample, y; kwargs...)`: 3-arg convenience form that overrides `y` before plotting.
+- `plot_trajectory(bench, traj; kwargs...)`: dynamic benchmarks only; produces a grid of per-epoch subplots.
+- `animate_trajectory(bench, traj; kwargs...)`: dynamic benchmarks only, returns a `Plots.Animation` that can be saved with `gif(anim, "file.gif")`.

ext/DFLBenchmarksPlotsExt.jl

@@ -0,0 +1,29 @@
+module DFLBenchmarksPlotsExt
+
+using DecisionFocusedLearningBenchmarks
+using DocStringExtensions: TYPEDSIGNATURES
+using LaTeXStrings: @L_str
+using Plots
+import DecisionFocusedLearningBenchmarks:
+    has_visualization, plot_instance, plot_solution, plot_trajectory, animate_trajectory
+
+include("plots/argmax2d_plots.jl")
+include("plots/warcraft_plots.jl")
+include("plots/svs_plots.jl")
+include("plots/dvs_plots.jl")
+
+"""
+    plot_solution(bench::AbstractBenchmark, sample::DataSample, y; kwargs...)
+
+Reconstruct a new sample with `y` overridden and delegate to the 2-arg
+[`plot_solution`](@ref). Only available when `Plots` is loaded.
+"""
+function plot_solution(bench::AbstractBenchmark, sample::DataSample, y; kwargs...)
+    return plot_solution(
+        bench,
+        DataSample(; sample.context..., x=sample.x, θ=sample.θ, y=y, extra=sample.extra);
+        kwargs...,
+    )
+end
+
+end

ext/plots/argmax2d_plots.jl

@@ -0,0 +1,70 @@
+function _init_plot(title="")
+    pl = Plots.plot(;
+        aspect_ratio=:equal,
+        legend=:outerleft,
+        xlim=(-1.1, 1.1),
+        ylim=(-1.1, 1.1),
+        title=title,
+    )
+    return pl
+end
+
+function _plot_polytope!(pl, vertices)
+    return Plots.plot!(
+        pl,
+        vcat(map(first, vertices), first(vertices[1])),
+        vcat(map(last, vertices), last(vertices[1]));
+        fillrange=0,
+        fillcolor=:gray,
+        fillalpha=0.2,
+        linecolor=:black,
+        label=L"\mathrm{conv}(\mathcal{Y}(x))",
+    )
+end
+
+function _plot_objective!(pl, θ)
+    Plots.plot!(
+        pl, [0.0, θ[1]], [0.0, θ[2]]; color="#9558B2", arrow=true, lw=2, label=nothing
+    )
+    Plots.annotate!(pl, [-0.2 * θ[1]], [-0.2 * θ[2]], [L"\theta"])
+    return pl
+end
+
+function _plot_y!(pl, y)
+    return Plots.scatter!(
+        pl,
+        [y[1]],
+        [y[2]];
+        color="#CB3C33",
+        markersize=9,
+        markershape=:square,
+        label=L"f(\theta)",
+    )
+end
+
+function _plot_maximizer!(pl, θ, instance, maximizer)
+    ŷ = maximizer(θ; instance)
+    return _plot_y!(pl, ŷ)
+end
+
+has_visualization(::Argmax2DBenchmark) = true
+
+function plot_instance(::Argmax2DBenchmark, sample::DataSample)
+    pl = _init_plot()
+    _plot_polytope!(pl, sample.instance)
+    return pl
+end
+
+function plot_solution(::Argmax2DBenchmark, sample::DataSample)
+    pl = _init_plot()
+    _plot_polytope!(pl, sample.instance)
+    _plot_objective!(pl, sample.θ)
+    return _plot_y!(pl, sample.y)
+end
+
+function plot_solution(::Argmax2DBenchmark, sample::DataSample, y; θ=sample.θ)
+    pl = _init_plot()
+    _plot_polytope!(pl, sample.instance)
+    _plot_objective!(pl, θ)
+    return _plot_y!(pl, y)
+end